1use std::fmt;
4
5use chematic_core::{AtomIdx, Molecule};
6
7use crate::coords::Coords3D;
8use crate::shape_descriptors::jacobi3;
9
10#[derive(Debug, PartialEq)]
15pub enum ConformerError {
16 AtomCountMismatch { expected: usize, got: usize },
17}
18
19impl fmt::Display for ConformerError {
20 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21 match self {
22 ConformerError::AtomCountMismatch { expected, got } => {
23 write!(f, "conformer has {got} atoms but molecule has {expected}")
24 }
25 }
26 }
27}
28
29impl std::error::Error for ConformerError {}
30
31pub struct ConformerEnsemble {
40 mol: Molecule,
41 conformers: Vec<Coords3D>,
42}
43
44impl ConformerEnsemble {
45 pub fn new(mol: Molecule) -> Self {
47 Self {
48 mol,
49 conformers: Vec::new(),
50 }
51 }
52
53 pub fn with_conformer(mol: Molecule, coords: Coords3D) -> Result<Self, ConformerError> {
57 let expected = mol.atom_count();
58 let got = coords.atom_count();
59 if got != expected {
60 return Err(ConformerError::AtomCountMismatch { expected, got });
61 }
62 Ok(Self {
63 mol,
64 conformers: vec![coords],
65 })
66 }
67
68 pub fn mol(&self) -> &Molecule {
70 &self.mol
71 }
72
73 pub fn conformer_count(&self) -> usize {
75 self.conformers.len()
76 }
77
78 pub fn add_conformer(&mut self, coords: Coords3D) -> Result<usize, ConformerError> {
83 let expected = self.mol.atom_count();
84 let got = coords.atom_count();
85 if got != expected {
86 return Err(ConformerError::AtomCountMismatch { expected, got });
87 }
88 let idx = self.conformers.len();
89 self.conformers.push(coords);
90 Ok(idx)
91 }
92
93 pub fn get_conformer(&self, idx: usize) -> Option<&Coords3D> {
95 self.conformers.get(idx)
96 }
97
98 pub fn get_conformer_mut(&mut self, idx: usize) -> Option<&mut Coords3D> {
100 self.conformers.get_mut(idx)
101 }
102
103 pub fn remove_conformer(&mut self, idx: usize) -> Option<Coords3D> {
108 if idx < self.conformers.len() {
109 Some(self.conformers.remove(idx))
110 } else {
111 None
112 }
113 }
114
115 pub fn conformer_rmsd_no_align(&self, a: usize, b: usize) -> Option<f64> {
119 let ca = self.conformers.get(a)?;
120 let cb = self.conformers.get(b)?;
121 let n = self.mol.atom_count();
122 if n == 0 {
123 return Some(0.0);
124 }
125 let sum_sq: f64 = (0..n)
126 .map(|i| {
127 let idx = AtomIdx(i as u32);
128 let pa = ca.get(idx);
129 let pb = cb.get(idx);
130 let dx = pa.x - pb.x;
131 let dy = pa.y - pb.y;
132 let dz = pa.z - pb.z;
133 dx * dx + dy * dy + dz * dz
134 })
135 .sum();
136 Some((sum_sq / n as f64).sqrt())
137 }
138
139 pub fn conformer_rmsd(&self, a: usize, b: usize) -> Option<f64> {
145 let ca = self.conformers.get(a)?;
146 let cb = self.conformers.get(b)?;
147 let n = self.mol.atom_count();
148 Some(kabsch_rmsd(ca, cb, n))
149 }
150
151 pub fn conformer_usr_descriptors(&self, idx: usize) -> Option<[f64; 12]> {
155 let c = self.conformers.get(idx)?;
156 let pts: Vec<[f64; 3]> = c.points.iter().map(|p| [p.x, p.y, p.z]).collect();
157 Some(crate::usr::usr_descriptors(&pts))
158 }
159
160 pub fn cluster_conformers_by_rms(&self, rms_threshold: f64) -> Vec<usize> {
181 let n = self.conformers.len();
182 if n == 0 {
183 return vec![];
184 }
185 if rms_threshold <= 0.0 {
186 return (0..n).collect();
187 }
188 let mut leaders: Vec<usize> = Vec::new();
189 'outer: for i in 0..n {
190 for &leader in &leaders {
191 let rmsd = self.conformer_rmsd(i, leader).unwrap_or(f64::INFINITY);
192 if rmsd < rms_threshold {
193 continue 'outer; }
195 }
196 leaders.push(i); }
198 leaders
199 }
200
201 pub fn is_duplicate(&self, coords: &Coords3D, rmsd_threshold: f64) -> bool {
207 if rmsd_threshold <= 0.0 {
208 return false;
209 }
210 let n = self.mol.atom_count();
211 self.conformers
212 .iter()
213 .any(|existing| kabsch_rmsd(coords, existing, n) < rmsd_threshold)
214 }
215
216 pub fn conformer_diversity_usr(&self) -> f64 {
222 let n = self.conformers.len();
223 if n < 2 {
224 return 0.0;
225 }
226 let descs: Vec<[f64; 12]> = (0..n)
227 .filter_map(|i| self.conformer_usr_descriptors(i))
228 .collect();
229 let mut total = 0.0;
230 let mut count = 0usize;
231 for i in 0..descs.len() {
232 for j in i + 1..descs.len() {
233 total += 1.0 - crate::usr::usr_similarity(&descs[i], &descs[j]);
234 count += 1;
235 }
236 }
237 if count == 0 {
238 0.0
239 } else {
240 total / count as f64
241 }
242 }
243}
244
245fn kabsch_rmsd(coords_a: &Coords3D, coords_b: &Coords3D, n: usize) -> f64 {
250 if n == 0 {
251 return 0.0;
252 }
253
254 let nf = n as f64;
255
256 let mut ca = [0.0f64; 3];
258 let mut cb = [0.0f64; 3];
259 for i in 0..n {
260 let idx = AtomIdx(i as u32);
261 let pa = coords_a.get(idx);
262 let pb = coords_b.get(idx);
263 ca[0] += pa.x;
264 ca[1] += pa.y;
265 ca[2] += pa.z;
266 cb[0] += pb.x;
267 cb[1] += pb.y;
268 cb[2] += pb.z;
269 }
270 for k in 0..3 {
271 ca[k] /= nf;
272 cb[k] /= nf;
273 }
274
275 let mut p = vec![[0.0f64; 3]; n];
277 let mut q = vec![[0.0f64; 3]; n];
278 for i in 0..n {
279 let idx = AtomIdx(i as u32);
280 let pa = coords_a.get(idx);
281 let pb = coords_b.get(idx);
282 p[i] = [pa.x - ca[0], pa.y - ca[1], pa.z - ca[2]];
283 q[i] = [pb.x - cb[0], pb.y - cb[1], pb.z - cb[2]];
284 }
285
286 let mut h = [[0.0f64; 3]; 3];
288 for i in 0..n {
289 for r in 0..3 {
290 for c in 0..3 {
291 h[r][c] += p[i][r] * q[i][c];
292 }
293 }
294 }
295
296 let mut hth = [[0.0f64; 3]; 3];
298 for r in 0..3 {
299 for c in 0..3 {
300 for k in 0..3 {
301 hth[r][c] += h[k][r] * h[k][c];
302 }
303 }
304 }
305
306 let (evals, v) = jacobi3(hth);
309
310 let mut hv = [[0.0f64; 3]; 3];
312 for r in 0..3 {
313 for c in 0..3 {
314 for k in 0..3 {
315 hv[r][c] += h[r][k] * v[k][c];
316 }
317 }
318 }
319 let mut u = [[0.0f64; 3]; 3];
320 for j in 0..3 {
321 let sigma = evals[j].max(0.0).sqrt();
322 for r in 0..3 {
323 u[r][j] = if sigma > 1e-10 { hv[r][j] / sigma } else { 0.0 };
324 }
325 }
326
327 let mut r_mat = [[0.0f64; 3]; 3];
329 for r in 0..3 {
330 for c in 0..3 {
331 for k in 0..3 {
332 r_mat[r][c] += u[r][k] * v[c][k];
333 }
334 }
335 }
336
337 let det = det3(r_mat);
339 let mut v_final = v;
340 if det < 0.0 {
341 for r in 0..3 {
342 v_final[r][0] *= -1.0;
343 }
344 r_mat = [[0.0f64; 3]; 3];
346 for r in 0..3 {
347 for c in 0..3 {
348 for k in 0..3 {
349 r_mat[r][c] += u[r][k] * v_final[c][k];
350 }
351 }
352 }
353 }
354
355 let mut sum_sq = 0.0f64;
357 for i in 0..n {
358 for row in 0..3 {
359 let rotated =
360 r_mat[row][0] * q[i][0] + r_mat[row][1] * q[i][1] + r_mat[row][2] * q[i][2];
361 let diff = p[i][row] - rotated;
362 sum_sq += diff * diff;
363 }
364 }
365 (sum_sq / nf).sqrt()
366}
367
368fn det3(m: [[f64; 3]; 3]) -> f64 {
369 m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
370 - m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
371 + m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0])
372}
373
374#[cfg(test)]
379mod tests {
380 use super::*;
381 use chematic_smiles::parse;
382
383 use crate::{coords::Point3, dg::generate_coords};
384
385 fn make_ensemble() -> ConformerEnsemble {
386 let mol = parse("CCC").unwrap();
387 let c = generate_coords(&mol);
388 ConformerEnsemble::with_conformer(mol, c).unwrap()
389 }
390
391 #[test]
394 fn new_has_zero_conformers() {
395 let mol = parse("C").unwrap();
396 let ens = ConformerEnsemble::new(mol);
397 assert_eq!(ens.conformer_count(), 0);
398 }
399
400 #[test]
401 fn with_conformer_has_one() {
402 let ens = make_ensemble();
403 assert_eq!(ens.conformer_count(), 1);
404 }
405
406 #[test]
407 fn add_conformer_increments_count() {
408 let mol = parse("CC").unwrap();
409 let c1 = generate_coords(&mol);
410 let c2 = generate_coords(&mol);
411 let mut ens = ConformerEnsemble::with_conformer(mol, c1).unwrap();
412 let idx = ens.add_conformer(c2).unwrap();
413 assert_eq!(idx, 1);
414 assert_eq!(ens.conformer_count(), 2);
415 }
416
417 #[test]
418 fn add_conformer_wrong_atom_count_errors() {
419 let mol = parse("CC").unwrap();
420 let wrong = Coords3D::new_zeroed(5);
421 let mut ens = ConformerEnsemble::new(mol);
422 let err = ens.add_conformer(wrong).unwrap_err();
423 assert!(matches!(
424 err,
425 ConformerError::AtomCountMismatch {
426 expected: 2,
427 got: 5
428 }
429 ));
430 }
431
432 #[test]
433 fn get_conformer_out_of_range_returns_none() {
434 let ens = make_ensemble();
435 assert!(ens.get_conformer(99).is_none());
436 }
437
438 #[test]
441 fn remove_conformer_decrements_count() {
442 let mut ens = make_ensemble();
443 let removed = ens.remove_conformer(0);
444 assert!(removed.is_some());
445 assert_eq!(ens.conformer_count(), 0);
446 }
447
448 #[test]
449 fn remove_conformer_shifts_indices() {
450 let mol = parse("C").unwrap();
451 let n = mol.atom_count();
452 let mut ens = ConformerEnsemble::new(mol);
453
454 for x in [1.0f64, 2.0, 3.0] {
456 let mut c = Coords3D::new_zeroed(n);
457 c.set(AtomIdx(0), Point3::new(x, 0.0, 0.0));
458 ens.add_conformer(c).unwrap();
459 }
460
461 ens.remove_conformer(0).unwrap();
463 assert_eq!(ens.conformer_count(), 2);
464 assert!((ens.get_conformer(0).unwrap().get(AtomIdx(0)).x - 2.0).abs() < 1e-10);
465 }
466
467 #[test]
468 fn remove_conformer_out_of_range_returns_none() {
469 let mut ens = make_ensemble();
470 assert!(ens.remove_conformer(99).is_none());
471 }
472
473 #[test]
476 fn rmsd_no_align_same_conformer_is_zero() {
477 let ens = make_ensemble();
478 let rmsd = ens.conformer_rmsd_no_align(0, 0).unwrap();
479 assert!(rmsd.abs() < 1e-10, "self-RMSD should be 0, got {rmsd}");
480 }
481
482 #[test]
483 fn rmsd_no_align_translated_is_nonzero() {
484 let mol = parse("CC").unwrap();
485 let n = mol.atom_count();
486 let mut c1 = Coords3D::new_zeroed(n);
487 let mut c2 = Coords3D::new_zeroed(n);
488 for i in 0..n {
489 c1.set(AtomIdx(i as u32), Point3::new(i as f64, 0.0, 0.0));
490 c2.set(AtomIdx(i as u32), Point3::new(i as f64 + 10.0, 0.0, 0.0));
491 }
492 let mut ens = ConformerEnsemble::with_conformer(mol, c1).unwrap();
493 ens.add_conformer(c2).unwrap();
494 let rmsd = ens.conformer_rmsd_no_align(0, 1).unwrap();
495 assert!(
496 rmsd > 0.0,
497 "translated conformers should have non-zero RMSD"
498 );
499 }
500
501 #[test]
502 fn kabsch_rmsd_same_conformer_is_zero() {
503 let ens = make_ensemble();
504 let rmsd = ens.conformer_rmsd(0, 0).unwrap();
505 assert!(
506 rmsd.abs() < 1e-8,
507 "Kabsch self-RMSD should be 0, got {rmsd}"
508 );
509 }
510
511 #[test]
512 fn kabsch_rmsd_pure_translation_is_zero() {
513 let mol = parse("CCC").unwrap();
515 let n = mol.atom_count();
516 let base = generate_coords(&mol);
517 let mut shifted = Coords3D::new_zeroed(n);
518 let offset = 5.0;
519 for i in 0..n {
520 let p = base.get(AtomIdx(i as u32));
521 shifted.set(
522 AtomIdx(i as u32),
523 Point3::new(p.x + offset, p.y + offset, p.z + offset),
524 );
525 }
526 let mut ens = ConformerEnsemble::with_conformer(mol, base).unwrap();
527 ens.add_conformer(shifted).unwrap();
528 let rmsd = ens.conformer_rmsd(0, 1).unwrap();
529 assert!(
530 rmsd < 1e-6,
531 "pure-translation Kabsch RMSD should be ~0, got {rmsd}"
532 );
533 }
534
535 #[test]
536 fn kabsch_rmsd_pure_rotation_is_zero() {
537 let mol = parse("CCC").unwrap();
539 let n = mol.atom_count();
540 let base = generate_coords(&mol);
541 let mut rotated = Coords3D::new_zeroed(n);
543 for i in 0..n {
544 let p = base.get(AtomIdx(i as u32));
545 rotated.set(AtomIdx(i as u32), Point3::new(-p.y, p.x, p.z));
546 }
547 let mut ens = ConformerEnsemble::with_conformer(mol, base).unwrap();
548 ens.add_conformer(rotated).unwrap();
549 let rmsd = ens.conformer_rmsd(0, 1).unwrap();
550 assert!(
551 rmsd < 1e-6,
552 "pure-rotation Kabsch RMSD should be ~0, got {rmsd}"
553 );
554 }
555
556 #[test]
557 fn kabsch_rmsd_different_conformers_nonzero() {
558 let mol = parse("CCC").unwrap();
559 let c1 = generate_coords(&mol);
560 let n = mol.atom_count();
561 let mut c2 = Coords3D::new_zeroed(n);
563 for i in 0..n {
564 let p = c1.get(AtomIdx(i as u32));
565 c2.set(AtomIdx(i as u32), Point3::new(-p.x, p.y, p.z));
566 }
567 let mut ens = ConformerEnsemble::with_conformer(mol, c1).unwrap();
568 ens.add_conformer(c2).unwrap();
569 let rmsd = ens.conformer_rmsd(0, 1).unwrap();
570 assert!(rmsd >= 0.0, "RMSD must be non-negative, got {rmsd}");
573 }
574
575 #[test]
576 fn kabsch_rmsd_out_of_range_returns_none() {
577 let ens = make_ensemble();
578 assert!(ens.conformer_rmsd(0, 99).is_none());
579 assert!(ens.conformer_rmsd(99, 0).is_none());
580 }
581
582 #[test]
583 fn rmsd_no_align_out_of_range_returns_none() {
584 let ens = make_ensemble();
585 assert!(ens.conformer_rmsd_no_align(0, 99).is_none());
586 }
587
588 #[test]
591 fn usr_descriptors_single_conformer() {
592 let ens = make_ensemble();
593 let d = ens.conformer_usr_descriptors(0);
594 assert!(d.is_some(), "valid index must return Some");
595 assert!(
596 d.unwrap().iter().all(|v| v.is_finite()),
597 "all USR values finite"
598 );
599 }
600
601 #[test]
602 fn usr_descriptors_out_of_range() {
603 let ens = make_ensemble();
604 assert!(ens.conformer_usr_descriptors(99).is_none());
605 }
606
607 #[test]
608 fn diversity_usr_single_conformer_is_zero() {
609 let ens = make_ensemble();
610 assert_eq!(
611 ens.conformer_diversity_usr(),
612 0.0,
613 "single conformer → diversity 0"
614 );
615 }
616
617 #[test]
618 fn diversity_usr_identical_conformers_is_zero() {
619 use crate::coords::Point3;
620 use chematic_core::{Atom, Element, MoleculeBuilder};
621
622 let mut b = MoleculeBuilder::new();
623 let a0 = b.add_atom(Atom::new(Element::C));
624 let a1 = b.add_atom(Atom::new(Element::C));
625 let mol = b.build();
626
627 let mut c = Coords3D::new_zeroed(2);
628 c.set(a0, Point3::new(0.0, 0.0, 0.0));
629 c.set(a1, Point3::new(1.5, 0.0, 0.0));
630
631 let mut ens = ConformerEnsemble::with_conformer(mol, c.clone()).unwrap();
632 ens.add_conformer(c).unwrap();
633
634 let div = ens.conformer_diversity_usr();
635 assert!(
636 div.abs() < 1e-9,
637 "identical conformers → diversity ~0, got {div}"
638 );
639 }
640
641 #[test]
642 fn diversity_usr_different_shapes_positive() {
643 use crate::coords::Point3;
644 use chematic_core::{Atom, Element, MoleculeBuilder};
645
646 let mut b = MoleculeBuilder::new();
648 let a0 = b.add_atom(Atom::new(Element::C));
649 let a1 = b.add_atom(Atom::new(Element::C));
650 let a2 = b.add_atom(Atom::new(Element::C));
651 let mol = b.build();
652
653 let mut c1 = Coords3D::new_zeroed(3);
654 c1.set(a0, Point3::new(0.0, 0.0, 0.0));
655 c1.set(a1, Point3::new(1.0, 0.0, 0.0));
656 c1.set(a2, Point3::new(2.0, 0.0, 0.0));
657
658 let mut c2 = Coords3D::new_zeroed(3);
659 c2.set(a0, Point3::new(0.0, 0.0, 0.0));
660 c2.set(a1, Point3::new(0.0, 10.0, 0.0));
661 c2.set(a2, Point3::new(0.0, 0.0, 10.0));
662
663 let mut ens = ConformerEnsemble::with_conformer(mol, c1).unwrap();
664 ens.add_conformer(c2).unwrap();
665
666 let div = ens.conformer_diversity_usr();
667 assert!(
668 div > 0.0 && div <= 1.0,
669 "diverse ensemble → diversity in (0,1], got {div}"
670 );
671 }
672
673 #[test]
676 fn cluster_empty_ensemble() {
677 let mol = parse("C").unwrap();
678 let ens = ConformerEnsemble::new(mol);
679 assert!(ens.cluster_conformers_by_rms(0.5).is_empty());
680 }
681
682 #[test]
683 fn cluster_single_conformer() {
684 let ens = make_ensemble();
685 assert_eq!(ens.cluster_conformers_by_rms(0.5), vec![0]);
686 }
687
688 #[test]
689 fn cluster_zero_threshold_keeps_all() {
690 let mol = parse("CCC").unwrap();
691 let c = generate_coords(&mol);
692 let mut ens = ConformerEnsemble::with_conformer(mol, c.clone()).unwrap();
693 ens.add_conformer(c.clone()).unwrap();
694 ens.add_conformer(c).unwrap();
695 let kept = ens.cluster_conformers_by_rms(0.0);
696 assert_eq!(kept, vec![0, 1, 2], "threshold ≤ 0 → keep all");
697 }
698
699 #[test]
700 fn cluster_negative_threshold_keeps_all() {
701 let mol = parse("CCC").unwrap();
702 let c = generate_coords(&mol);
703 let mut ens = ConformerEnsemble::with_conformer(mol, c.clone()).unwrap();
704 ens.add_conformer(c).unwrap();
705 assert_eq!(ens.cluster_conformers_by_rms(-1.0), vec![0, 1]);
706 }
707
708 #[test]
709 fn cluster_identical_conformers_keeps_first() {
710 let mol = parse("CCC").unwrap();
711 let c = generate_coords(&mol);
712 let mut ens = ConformerEnsemble::with_conformer(mol, c.clone()).unwrap();
713 ens.add_conformer(c.clone()).unwrap();
714 ens.add_conformer(c).unwrap();
715 let kept = ens.cluster_conformers_by_rms(0.01);
716 assert_eq!(
717 kept,
718 vec![0],
719 "three identical conformers → keep only first"
720 );
721 }
722
723 #[test]
724 fn cluster_distinct_conformers_keeps_all() {
725 let mol = parse("CCC").unwrap();
727 let n = mol.atom_count();
728 let mut ca = Coords3D::new_zeroed(n);
729 let mut cb = Coords3D::new_zeroed(n);
730 for i in 0..n {
731 ca.set(
732 chematic_core::AtomIdx(i as u32),
733 Point3 {
734 x: i as f64,
735 y: 0.0,
736 z: 0.0,
737 },
738 );
739 cb.set(
740 chematic_core::AtomIdx(i as u32),
741 Point3 {
742 x: 0.0,
743 y: i as f64 * 100.0,
744 z: 0.0,
745 },
746 );
747 }
748 let mut ens = ConformerEnsemble::with_conformer(mol, ca).unwrap();
749 ens.add_conformer(cb).unwrap();
750 let kept = ens.cluster_conformers_by_rms(0.5);
751 assert_eq!(kept, vec![0, 1], "two distinct conformers → both kept");
752 }
753
754 #[test]
755 fn cluster_ascending_order() {
756 let mol = parse("CCC").unwrap();
758 let c0 = generate_coords(&mol);
759 let mut ens = ConformerEnsemble::with_conformer(mol, c0.clone()).unwrap();
760 ens.add_conformer(c0).unwrap(); let n = ens.mol().atom_count();
762 let mut far = Coords3D::new_zeroed(n);
763 for i in 0..n {
764 far.set(
765 chematic_core::AtomIdx(i as u32),
766 Point3 {
767 x: 0.0,
768 y: i as f64 * 100.0,
769 z: 0.0,
770 },
771 );
772 }
773 ens.add_conformer(far).unwrap(); let kept = ens.cluster_conformers_by_rms(0.1);
775 for w in kept.windows(2) {
776 assert!(w[0] < w[1], "kept indices not ascending");
777 }
778 }
779}