qsym2/auxiliary/
molecule.rs

1//! Molecular structures.
2
3use std::collections::{HashMap, HashSet};
4use std::fmt;
5use std::fs;
6use std::io::{BufWriter, Write};
7use std::path::Path;
8use std::process;
9
10use anyhow;
11use itertools::Itertools;
12use log;
13use nalgebra::{DVector, Matrix3, Point3, Vector3};
14use ndarray::{Array2, ShapeBuilder};
15use ndarray_linalg::{Eigh, UPLO};
16use num_traits::ToPrimitive;
17use serde::{Deserialize, Serialize};
18
19use crate::auxiliary::atom::{Atom, AtomKind, ElementMap};
20use crate::auxiliary::geometry::{self, ImproperRotationKind, Transform};
21use crate::permutation::{permute_inplace, PermutableCollection, Permutation};
22
23#[cfg(test)]
24#[path = "sea_tests.rs"]
25mod sea_tests;
26
27#[cfg(test)]
28#[path = "molecule_tests.rs"]
29mod molecule_tests;
30
31// ==================
32// Struct definitions
33// ==================
34
35/// Structure containing the atoms constituting a molecule.
36#[derive(Clone, Debug, Serialize, Deserialize)]
37pub struct Molecule {
38    /// The atoms constituting this molecule.
39    pub atoms: Vec<Atom>,
40
41    /// Optional special atom to represent the electric fields applied to this molecule.
42    pub electric_atoms: Option<Vec<Atom>>,
43
44    /// Optional special atoms to represent the magnetic fields applied to this molecule.
45    pub magnetic_atoms: Option<Vec<Atom>>,
46
47    /// A threshold for approximate equality comparisons.
48    pub threshold: f64,
49}
50
51impl Molecule {
52    /// Parses an XYZ file to construct a molecule.
53    ///
54    /// # Arguments
55    ///
56    /// * `filename` - The XYZ file to be parsed.
57    ///
58    /// # Returns
59    ///
60    /// The parsed [`Molecule`] structure.
61    ///
62    /// # Panics
63    ///
64    /// Panics when unable to parse the provided XYZ file.
65    #[must_use]
66    pub fn from_xyz<P: AsRef<Path>>(filename: P, thresh: f64) -> Self {
67        let contents = fs::read_to_string(&filename).unwrap_or_else(|err| {
68            log::error!("Unable to read file {}.", filename.as_ref().display());
69            log::error!("{}", err);
70            process::exit(1);
71        });
72
73        let mut atoms: Vec<Atom> = vec![];
74        let emap = ElementMap::new();
75        let mut n_atoms = 0usize;
76        for (i, line) in contents.lines().enumerate() {
77            if i == 0 {
78                n_atoms = line.parse::<usize>().unwrap_or_else(|err| {
79                    log::error!(
80                        "Unable to read number of atoms in {}.",
81                        filename.as_ref().display()
82                    );
83                    log::error!("{}", err);
84                    process::exit(1);
85                });
86            } else if i == 1 {
87                continue;
88            } else {
89                atoms.push(
90                    Atom::from_xyz(line, &emap, thresh)
91                        .unwrap_or_else(|| panic!("Unable to parse {line} to give an atom.")),
92                );
93            }
94        }
95        assert_eq!(
96            atoms.len(),
97            n_atoms,
98            "Expected {} atoms, got {} instead.",
99            n_atoms,
100            atoms.len()
101        );
102        Molecule {
103            atoms,
104            electric_atoms: None,
105            magnetic_atoms: None,
106            threshold: thresh,
107        }
108    }
109
110    /// Writes the ordinary atoms in the molecule to an XYZ file.
111    ///
112    /// # Arguments
113    ///
114    /// * `filename` - The XYZ file to be written.
115    pub fn to_xyz<P: AsRef<Path>>(&self, filename: P) -> Result<(), anyhow::Error> {
116        let mut f = BufWriter::new(fs::File::create(filename)?);
117        writeln!(f, "{}", self.atoms.len())?;
118        writeln!(f, "")?;
119        for atom in self.atoms.iter() {
120            writeln!(f, "{}", atom.to_xyz())?;
121        }
122        Ok(())
123    }
124
125    /// Construct a molecule from an array of atoms.
126    ///
127    /// # Arguments
128    ///
129    /// * `all_atoms` - The atoms (of all types) constituting this molecule.
130    /// * `threshold` - A threshold for approximate equality comparisons.
131    ///
132    /// # Returns
133    ///
134    /// The constructed [`Molecule`] struct.
135    ///
136    /// # Panics
137    ///
138    /// Panics when the numbers of fictitious special atoms, if any, are invalid. It is expected
139    /// that, when present, there are two magnetic special atoms and/or one electric special atom.
140    #[must_use]
141    pub fn from_atoms(all_atoms: &[Atom], thresh: f64) -> Self {
142        let mut atoms: Vec<Atom> = all_atoms
143            .iter()
144            .filter(|atom| matches!(atom.kind, AtomKind::Ordinary))
145            .cloned()
146            .collect();
147        atoms.iter_mut().for_each(|atom| atom.threshold = thresh);
148
149        let mut magnetic_atoms_vec: Vec<Atom> = all_atoms
150            .iter()
151            .filter(|atom| matches!(atom.kind, AtomKind::Magnetic(_)))
152            .cloned()
153            .collect();
154        magnetic_atoms_vec
155            .iter_mut()
156            .for_each(|atom| atom.threshold = thresh);
157        let magnetic_atoms = if magnetic_atoms_vec.is_empty() {
158            None
159        } else {
160            Some(magnetic_atoms_vec)
161        };
162
163        let mut electric_atoms_vec: Vec<Atom> = all_atoms
164            .iter()
165            .filter(|atom| matches!(atom.kind, AtomKind::Electric(_)))
166            .cloned()
167            .collect();
168        electric_atoms_vec
169            .iter_mut()
170            .for_each(|atom| atom.threshold = thresh);
171        let electric_atoms = if electric_atoms_vec.is_empty() {
172            None
173        } else {
174            Some(electric_atoms_vec)
175        };
176
177        Molecule {
178            atoms,
179            electric_atoms,
180            magnetic_atoms,
181            threshold: thresh,
182        }
183    }
184
185    /// Constructs a new molecule containing only the ordinary atoms in this molecule.
186    #[must_use]
187    pub fn molecule_ordinary_atoms(&self) -> Self {
188        Self::from_atoms(&self.atoms, self.threshold)
189    }
190
191    /// Constructs a new molecule containing only the fictitious magnetic atoms in this molecule,
192    /// if any.
193    ///
194    /// # Returns
195    ///
196    /// Returns `None` if this molecule has no fictitious magnetic atoms.
197    #[must_use]
198    pub fn molecule_magnetic_atoms(&self) -> Option<Self> {
199        Some(Self::from_atoms(
200            self.magnetic_atoms.as_ref()?,
201            self.threshold,
202        ))
203    }
204
205    /// Constructs a new molecule containing only the fictitious electric atoms in this molecule,
206    /// if any.
207    ///
208    /// # Returns
209    ///
210    /// Returns `None` if this molecule has no fictitious electric atoms.
211    #[must_use]
212    pub fn molecule_electric_atoms(&self) -> Option<Self> {
213        Some(Self::from_atoms(
214            self.electric_atoms.as_ref()?,
215            self.threshold,
216        ))
217    }
218
219    /// Retrieves a vector of references to all atoms in this molecule, including special ones, if
220    /// any.
221    ///
222    /// # Returns
223    ///
224    /// All atoms in this molecule.
225    #[must_use]
226    pub fn get_all_atoms(&self) -> Vec<&Atom> {
227        let mut atoms: Vec<&Atom> = vec![];
228        for atom in &self.atoms {
229            atoms.push(atom);
230        }
231        if let Some(magnetic_atoms) = &self.magnetic_atoms {
232            for magnetic_atom in magnetic_atoms.iter() {
233                atoms.push(magnetic_atom);
234            }
235        }
236        if let Some(electric_atoms) = &self.electric_atoms {
237            for electric_atom in electric_atoms.iter() {
238                atoms.push(electric_atom);
239            }
240        }
241        atoms
242    }
243
244    /// Retrieves a vector of mutable references to all atoms in this molecule, including special
245    /// ones, if any.
246    ///
247    /// # Returns
248    ///
249    /// All atoms in this molecule.
250    #[must_use]
251    pub fn get_all_atoms_mut(&mut self) -> Vec<&mut Atom> {
252        let mut atoms: Vec<&mut Atom> = vec![];
253        for atom in &mut self.atoms {
254            atoms.push(atom);
255        }
256        if let Some(magnetic_atoms) = &mut self.magnetic_atoms {
257            for magnetic_atom in magnetic_atoms.iter_mut() {
258                atoms.push(magnetic_atom);
259            }
260        }
261        if let Some(electric_atoms) = &mut self.electric_atoms {
262            for electric_atom in electric_atoms.iter_mut() {
263                atoms.push(electric_atom);
264            }
265        }
266        atoms
267    }
268
269    /// Calculates the centre of mass of the molecule.
270    ///
271    /// This does not take into account fictitious special atoms.
272    ///
273    /// # Returns
274    ///
275    /// The centre of mass.
276    #[must_use]
277    pub fn calc_com(&self) -> Point3<f64> {
278        let atoms = &self.atoms;
279        let mut com: Point3<f64> = Point3::origin();
280        if atoms.is_empty() {
281            return com;
282        }
283        let mut tot_m: f64 = 0.0;
284        for atom in atoms.iter() {
285            let m: f64 = atom.atomic_mass;
286            com += atom.coordinates * m - Point3::origin();
287            tot_m += m;
288        }
289        com *= 1.0 / tot_m;
290        com
291    }
292
293    /// Calculates the inertia tensor of the molecule.
294    ///
295    /// This *does* take into account fictitious special atoms.
296    ///
297    /// # Arguments
298    ///
299    /// * `origin` - An origin about which the inertia tensor is evaluated.
300    ///
301    /// # Returns
302    ///
303    /// The inertia tensor as a $`3 \times 3`$ matrix.
304    #[must_use]
305    pub fn calc_inertia_tensor(&self, origin: &Point3<f64>) -> Matrix3<f64> {
306        let atoms = self.get_all_atoms();
307        let mut inertia_tensor = Matrix3::zeros();
308        for atom in &atoms {
309            let rel_coordinates: Vector3<f64> = atom.coordinates - origin;
310            for i in 0..3 {
311                for j in 0..=i {
312                    if i == j {
313                        inertia_tensor[(i, j)] += atom.atomic_mass
314                            * (rel_coordinates.norm_squared()
315                                - rel_coordinates[i] * rel_coordinates[j]);
316                    } else {
317                        inertia_tensor[(i, j)] -=
318                            atom.atomic_mass * rel_coordinates[i] * rel_coordinates[j];
319                        inertia_tensor[(j, i)] -=
320                            atom.atomic_mass * rel_coordinates[j] * rel_coordinates[i];
321                    }
322                }
323            }
324        }
325        log::debug!("Origin for inertia tensor:");
326        for component in origin.iter() {
327            log::debug!("  {component:+.14}");
328        }
329        log::debug!("Inertia tensor:\n{}", inertia_tensor);
330        inertia_tensor
331    }
332
333    /// Calculates the moments of inertia and the corresponding principal axes.
334    ///
335    /// This *does* take into account fictitious special atoms.
336    ///
337    /// # Returns
338    ///
339    /// * The moments of inertia in ascending order.
340    /// * The corresponding principal axes.
341    ///
342    /// # Panics
343    ///
344    /// Panics when any of the moments of inertia cannot be compared.
345    #[must_use]
346    pub fn calc_moi(&self) -> ([f64; 3], [Vector3<f64>; 3]) {
347        let inertia_tensor = Array2::from_shape_vec(
348            (3, 3).f(),
349            self.calc_inertia_tensor(&self.calc_com())
350                .into_iter()
351                .copied()
352                .collect::<Vec<_>>(),
353        )
354        .expect("Unable to construct the inertia tensor.");
355        let (eigvals, eigvecs) = inertia_tensor
356            .eigh(UPLO::Lower)
357            .expect("Unable to diagonalise the inertia tensor.");
358        let eigenvalues: Vec<f64> = eigvals.into_iter().collect::<Vec<_>>();
359        let eigenvectors = eigvecs
360            .columns()
361            .into_iter()
362            .map(|col| Vector3::from_iterator(col.into_owned()))
363            .collect::<Vec<_>>();
364        let mut eigen_tuple: Vec<(f64, _)> = eigenvalues
365            .iter()
366            .copied()
367            .zip(eigenvectors.iter().copied())
368            .collect();
369        eigen_tuple.sort_by(|(eigval0, _), (eigval1, _)| {
370            eigval0
371                .partial_cmp(eigval1)
372                .unwrap_or_else(|| panic!("{eigval0} and {eigval1} cannot be compared."))
373        });
374        let (sorted_eigenvalues, sorted_eigenvectors): (Vec<f64>, Vec<_>) =
375            eigen_tuple.into_iter().unzip();
376        let result = (
377            [
378                sorted_eigenvalues[0],
379                sorted_eigenvalues[1],
380                sorted_eigenvalues[2],
381            ],
382            [
383                geometry::get_standard_positive_pole(
384                    &Vector3::new(
385                        sorted_eigenvectors[0][(0, 0)],
386                        sorted_eigenvectors[0][(1, 0)],
387                        sorted_eigenvectors[0][(2, 0)],
388                    ),
389                    self.threshold,
390                ),
391                geometry::get_standard_positive_pole(
392                    &Vector3::new(
393                        sorted_eigenvectors[1][(0, 0)],
394                        sorted_eigenvectors[1][(1, 0)],
395                        sorted_eigenvectors[1][(2, 0)],
396                    ),
397                    self.threshold,
398                ),
399                geometry::get_standard_positive_pole(
400                    &Vector3::new(
401                        sorted_eigenvectors[2][(0, 0)],
402                        sorted_eigenvectors[2][(1, 0)],
403                        sorted_eigenvectors[2][(2, 0)],
404                    ),
405                    self.threshold,
406                ),
407            ],
408        );
409        result
410            .0
411            .iter()
412            .zip(result.1.iter())
413            .for_each(|(moi, axis)| {
414                log::debug!("Principal moment of inertia: {moi:.14}");
415                log::debug!("  -- Principal axis:\n{axis}");
416            });
417        result
418    }
419
420    /// Determines the interatomic distance matrix and the indices of symmetry-equivalent atoms.
421    ///
422    /// This *does* take into account fictitious special atoms.
423    ///
424    /// # Returns
425    ///
426    /// * The interatomic distance matrix where the distances in each column are sorted in ascending
427    /// order. Column $`j`$ contains the interatomic distances from atom $`j`$ to all other atoms
428    /// (both ordinary and fictitious) in the molecule. Also note that all atoms (both ordinary and
429    /// fictitious) are included here, so the matrix is square.
430    /// * A vector of vectors of symmetry-equivalent atom indices. Each inner vector contains
431    /// indices of atoms in one SEA group.
432    pub fn calc_interatomic_distance_matrix(&self) -> (Array2<f64>, Vec<Vec<usize>>) {
433        let all_atoms = &self.get_all_atoms();
434        let all_coords: Vec<_> = all_atoms.iter().map(|atm| atm.coordinates).collect();
435        let mut dist_columns: Vec<DVector<f64>> = vec![];
436        let mut sorted_dist_columns: Vec<DVector<f64>> = vec![];
437
438        // Determine indices of symmetry-equivalent atoms
439        let mut equiv_indicess: Vec<Vec<usize>> = vec![vec![0]];
440        for (j, coord_j) in all_coords.iter().enumerate() {
441            // column_j is the j-th column in the interatomic distance matrix. This column contains
442            // distances from ordinary atom j to all other atoms (both ordinary and fictitious) in
443            // the molecule.
444            let column_j = all_coords
445                .iter()
446                .map(|coord_i| (coord_j - coord_i).norm())
447                .collect_vec();
448            let mut sorted_column_j = column_j.clone();
449            dist_columns.push(DVector::from_vec(column_j));
450
451            sorted_column_j.sort_by(|a, b| {
452                a.partial_cmp(b).unwrap_or_else(|| {
453                    panic!("Mass-weighted interatomic distances {a} and {b} cannot be compared.")
454                })
455            });
456            let sorted_column_j_vec = DVector::from_vec(sorted_column_j);
457            if j == 0 {
458                sorted_dist_columns.push(sorted_column_j_vec);
459            } else {
460                let equiv_set_search = equiv_indicess.iter().position(|equiv_indices| {
461                    sorted_dist_columns[equiv_indices[0]].relative_eq(
462                        &sorted_column_j_vec,
463                        self.threshold,
464                        self.threshold,
465                    ) && match (&all_atoms[j].kind, &all_atoms[equiv_indices[0]].kind) {
466                        (AtomKind::Ordinary, AtomKind::Ordinary) => {
467                            all_atoms[j].atomic_number == all_atoms[equiv_indices[0]].atomic_number
468                        }
469                        (AtomKind::Magnetic(_), AtomKind::Magnetic(_))
470                        | (AtomKind::Electric(_), AtomKind::Electric(_)) => true,
471                        _ => false,
472                    }
473                });
474                sorted_dist_columns.push(sorted_column_j_vec);
475                if let Some(index) = equiv_set_search {
476                    equiv_indicess[index].push(j);
477                } else {
478                    equiv_indicess.push(vec![j]);
479                };
480            }
481        }
482
483        let dist_elements_f = dist_columns
484            .iter()
485            .flatten()
486            .cloned()
487            .collect::<Vec<_>>();
488        let n_atoms = all_atoms.len();
489        let dist_matrix =
490            Array2::<f64>::from_shape_vec((n_atoms, n_atoms).f(), dist_elements_f)
491                .expect("Unable to collect the interatomic distances into a square matrix.");
492
493        (dist_matrix, equiv_indicess)
494    }
495
496    /// Determines the sets of symmetry-equivalent atoms.
497    ///
498    /// This *does* take into account fictitious special atoms.
499    ///
500    /// # Returns
501    ///
502    /// * Copies of the atoms in the molecule, grouped into symmetry-equivalent
503    /// groups.
504    ///
505    /// # Panics
506    ///
507    /// Panics when the any of the mass-weighted interatomic distances cannot be compared.
508    #[must_use]
509    pub fn calc_sea_groups(&self) -> Vec<Vec<Atom>> {
510        let all_atoms = &self.get_all_atoms();
511        let (_, equiv_indicess) = self.calc_interatomic_distance_matrix();
512        let sea_groups: Vec<Vec<Atom>> = equiv_indicess
513            .iter()
514            .map(|equiv_indices| {
515                equiv_indices
516                    .iter()
517                    .map(|index| all_atoms[*index].clone())
518                    .collect()
519            })
520            .collect();
521
522        log::debug!("Number of SEA groups: {}", sea_groups.len());
523        sea_groups
524    }
525
526    /// Adds two fictitious magnetic atoms to represent the magnetic field.
527    ///
528    /// # Arguments
529    ///
530    /// * `magnetic_field` - The magnetic field vector. If zero or `None`, any magnetic
531    /// field present will be removed.
532    ///
533    /// # Panics
534    ///
535    /// Panics when the number of atoms cannot be represented as an `f64` value.
536    pub fn set_magnetic_field(&mut self, magnetic_field: Option<Vector3<f64>>) {
537        if let Some(b_vec) = magnetic_field {
538            if approx::relative_ne!(b_vec.norm(), 0.0) {
539                let com = self.calc_com();
540                let ave_mag = {
541                    let average_distance = self
542                        .atoms
543                        .iter()
544                        .fold(0.0, |acc, atom| acc + (atom.coordinates - com).magnitude())
545                        / self.atoms.len().to_f64().unwrap_or_else(|| {
546                            panic!("Unable to convert `{}` to `f64`.", self.atoms.len())
547                        });
548                    if average_distance > 0.0 {
549                        average_distance
550                    } else {
551                        0.5
552                    }
553                };
554                let b_vec_norm = b_vec.normalize() * ave_mag * 0.5;
555                self.magnetic_atoms = Some(vec![
556                    Atom::new_special(AtomKind::Magnetic(true), com + b_vec_norm, self.threshold)
557                        .expect("Unable to construct a special magnetic atom."),
558                    Atom::new_special(AtomKind::Magnetic(false), com - b_vec_norm, self.threshold)
559                        .expect("Unable to construct a special magnetic atom."),
560                ]);
561            } else {
562                self.magnetic_atoms = None;
563            }
564        } else {
565            self.magnetic_atoms = None;
566        }
567    }
568
569    /// Adds one fictitious electric atom to represent the electric field.
570    ///
571    /// # Arguments
572    ///
573    /// * `electric_field` - The electric field vector. If zero or `None`, any electric
574    /// field present will be removed.
575    ///
576    /// # Panics
577    ///
578    /// Panics when the number of atoms cannot be represented as an `f64` value.
579    pub fn set_electric_field(&mut self, electric_field: Option<Vector3<f64>>) {
580        if let Some(e_vec) = electric_field {
581            if approx::relative_ne!(e_vec.norm(), 0.0) {
582                let com = self.calc_com();
583                let ave_mag = {
584                    let average_distance = self
585                        .atoms
586                        .iter()
587                        .fold(0.0, |acc, atom| acc + (atom.coordinates - com).magnitude())
588                        / self.atoms.len().to_f64().unwrap_or_else(|| {
589                            panic!("Unable to convert `{}` to `f64`.", self.atoms.len())
590                        });
591                    if average_distance > 0.0 {
592                        average_distance
593                    } else {
594                        0.5
595                    }
596                };
597                let e_vec_norm = e_vec.normalize() * ave_mag * 0.5;
598                self.electric_atoms = Some(vec![Atom::new_special(
599                    AtomKind::Electric(true),
600                    com + e_vec_norm,
601                    self.threshold,
602                )
603                .expect("Unable to construct an electric special atom.")]);
604            } else {
605                self.electric_atoms = None;
606            }
607        } else {
608            self.electric_atoms = None;
609        }
610    }
611
612    /// Clones this molecule and adjusts all comparison thresholds to that specified by `thresh`.
613    ///
614    /// # Arguments
615    ///
616    /// * `thresh` - The new threshold to be assigned to the cloned molecule.
617    ///
618    /// # Returns
619    ///
620    /// A cloned copy of the molecule wit the adjusted threshold.
621    pub fn adjust_threshold(&self, thresh: f64) -> Self {
622        Self::from_atoms(
623            &self
624                .get_all_atoms()
625                .into_iter()
626                .cloned()
627                .collect::<Vec<_>>(),
628            thresh,
629        )
630    }
631
632    /// Reorientates the molecule in-place into a canonical alignment with the space-fixed axes of
633    /// the coordinate system.
634    ///
635    /// Fictitious special atoms are also moved during the reorientation.
636    ///
637    /// If the molecule has a unique principal axis, then this axis becomes aligned with the
638    /// $`z`$-axis and the other two degenerate axes become aligned with the $`x`$- and $`y`$-axes
639    /// of the coordinate system. If the molecule has no unique principal axes, then the axes are
640    /// aligned with $`x`$-, $`y`$-,  and $`z`$-axes in ascending order of moments of inertia.
641    ///
642    /// # Arguments
643    ///
644    /// * `moi_thresh` - Threshold for comparing moments of inertia.
645    pub fn reorientate_mut(&mut self, moi_thresh: f64) {
646        let (moi, principal_axes) = self.calc_moi();
647        let rotmat = if approx::relative_ne!(
648            moi[0],
649            moi[1],
650            max_relative = moi_thresh,
651            epsilon = moi_thresh
652        ) && approx::relative_eq!(
653            moi[1],
654            moi[2],
655            max_relative = moi_thresh,
656            epsilon = moi_thresh
657        ) {
658            // principal_axes[0] is unique.
659            Matrix3::from_columns(&[principal_axes[1], principal_axes[2], principal_axes[0]])
660                .transpose()
661        } else {
662            // principal_axes[2] is unique, or no unique axis, or isotropic.
663            Matrix3::from_columns(&[principal_axes[0], principal_axes[1], principal_axes[2]])
664                .transpose()
665        };
666        let com = self.calc_com();
667        self.recentre_mut();
668        self.transform_mut(&rotmat);
669        self.translate_mut(&(com - Point3::origin()));
670    }
671
672    /// Clones and reorientates the molecule into a canonical alignment with the space-fixed axes
673    /// of the coordinate system.
674    ///
675    /// Fictitious special atoms are also moved during the reorientation.
676    ///
677    /// If the molecule has a unique principal axis, then this axis becomes aligned with the
678    /// $`z`$-axis and the other two degenerate axes become aligned with the $`x`$- and $`y`$-axes
679    /// of the coordinate system. If the molecule has no unique principal axes, then the axes are
680    /// aligned with $`x`$-, $`y`$-,  and $`z`$-axes in ascending order of moments of inertia.
681    ///
682    /// # Arguments
683    ///
684    /// * `moi_thresh` - Threshold for comparing moments of inertia.
685    ///
686    /// # Returns
687    ///
688    /// A reoriented copy of the molecule.
689    pub fn reorientate(&self, moi_thresh: f64) -> Self {
690        let mut reoriented_mol = self.clone();
691        reoriented_mol.reorientate_mut(moi_thresh);
692        reoriented_mol
693    }
694}
695
696// =====================
697// Trait implementations
698// =====================
699
700impl fmt::Display for Molecule {
701    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
702        writeln!(f, "Molecule consisting of")?;
703        for atom in self.get_all_atoms().iter() {
704            writeln!(f, "  {atom}")?;
705        }
706        Ok(())
707    }
708}
709
710impl Transform for Molecule {
711    fn transform_mut(&mut self, mat: &Matrix3<f64>) {
712        for atom in &mut self.atoms {
713            atom.transform_mut(mat);
714        }
715        if let Some(ref mut mag_atoms) = self.magnetic_atoms {
716            for atom in mag_atoms.iter_mut() {
717                atom.transform_mut(mat);
718            }
719        }
720        if let Some(ref mut ele_atoms) = self.electric_atoms {
721            for atom in ele_atoms.iter_mut() {
722                atom.transform_mut(mat);
723            }
724        }
725    }
726
727    fn rotate_mut(&mut self, angle: f64, axis: &Vector3<f64>) {
728        for atom in &mut self.atoms {
729            atom.rotate_mut(angle, axis);
730        }
731        if let Some(ref mut mag_atoms) = self.magnetic_atoms {
732            for atom in mag_atoms.iter_mut() {
733                atom.rotate_mut(angle, axis);
734            }
735        }
736        if let Some(ref mut ele_atoms) = self.electric_atoms {
737            for atom in ele_atoms.iter_mut() {
738                atom.rotate_mut(angle, axis);
739            }
740        }
741    }
742
743    fn improper_rotate_mut(
744        &mut self,
745        angle: f64,
746        axis: &Vector3<f64>,
747        kind: &ImproperRotationKind,
748    ) {
749        for atom in &mut self.atoms {
750            atom.improper_rotate_mut(angle, axis, kind);
751        }
752        if let Some(ref mut mag_atoms) = self.magnetic_atoms {
753            for atom in mag_atoms.iter_mut() {
754                atom.improper_rotate_mut(angle, axis, kind);
755            }
756        }
757        if let Some(ref mut ele_atoms) = self.electric_atoms {
758            for atom in ele_atoms.iter_mut() {
759                atom.improper_rotate_mut(angle, axis, kind);
760            }
761        }
762    }
763
764    fn translate_mut(&mut self, tvec: &Vector3<f64>) {
765        for atom in &mut self.atoms {
766            atom.translate_mut(tvec);
767        }
768        if let Some(ref mut mag_atoms) = self.magnetic_atoms {
769            for atom in mag_atoms.iter_mut() {
770                atom.translate_mut(tvec);
771            }
772        }
773        if let Some(ref mut ele_atoms) = self.electric_atoms {
774            for atom in ele_atoms.iter_mut() {
775                atom.translate_mut(tvec);
776            }
777        }
778    }
779
780    fn recentre_mut(&mut self) {
781        let com = self.calc_com();
782        let tvec = -Vector3::new(com[0], com[1], com[2]);
783        self.translate_mut(&tvec);
784    }
785
786    fn reverse_time_mut(&mut self) {
787        if let Some(ref mut mag_atoms) = self.magnetic_atoms {
788            for atom in mag_atoms.iter_mut() {
789                atom.reverse_time_mut();
790            }
791        }
792    }
793
794    fn transform(&self, mat: &Matrix3<f64>) -> Self {
795        let mut transformed_mol = self.clone();
796        transformed_mol.transform_mut(mat);
797        transformed_mol
798    }
799
800    fn rotate(&self, angle: f64, axis: &Vector3<f64>) -> Self {
801        let mut rotated_mol = self.clone();
802        rotated_mol.rotate_mut(angle, axis);
803        rotated_mol
804    }
805
806    fn improper_rotate(
807        &self,
808        angle: f64,
809        axis: &Vector3<f64>,
810        kind: &ImproperRotationKind,
811    ) -> Self {
812        let mut improper_rotated_mol = self.clone();
813        improper_rotated_mol.improper_rotate_mut(angle, axis, kind);
814        improper_rotated_mol
815    }
816
817    fn translate(&self, tvec: &Vector3<f64>) -> Self {
818        let mut translated_mol = self.clone();
819        translated_mol.translate_mut(tvec);
820        translated_mol
821    }
822
823    fn recentre(&self) -> Self {
824        let mut recentred_mol = self.clone();
825        recentred_mol.recentre_mut();
826        recentred_mol
827    }
828
829    fn reverse_time(&self) -> Self {
830        let mut time_reversed_mol = self.clone();
831        time_reversed_mol.reverse_time_mut();
832        time_reversed_mol
833    }
834}
835
836impl PartialEq for Molecule {
837    fn eq(&self, other: &Self) -> bool {
838        if self.atoms.len() != other.atoms.len() {
839            return false;
840        };
841        let thresh = self
842            .atoms
843            .iter()
844            .chain(other.atoms.iter())
845            .fold(0.0_f64, |acc, atom| acc.max(atom.threshold));
846
847        let mut other_atoms_ref: HashSet<_> = other.atoms.iter().collect();
848        for s_atom in &self.atoms {
849            let o_atom = other
850                .atoms
851                .iter()
852                .find(|o_atm| (s_atom.coordinates - o_atm.coordinates).norm() < thresh);
853            match o_atom {
854                Some(atm) => {
855                    other_atoms_ref.remove(atm);
856                }
857                None => {
858                    break;
859                }
860            }
861        }
862        if !other_atoms_ref.is_empty() {
863            return false;
864        }
865
866        if let Some(self_mag_atoms) = &self.magnetic_atoms {
867            if let Some(other_mag_atoms) = &other.magnetic_atoms {
868                let mut other_mag_atoms_ref: HashSet<_> = other_mag_atoms.iter().collect();
869                for s_atom in self_mag_atoms.iter() {
870                    let o_atom = other_mag_atoms.iter().find(|o_atm| {
871                        (s_atom.coordinates - o_atm.coordinates).norm() < thresh
872                            && s_atom.kind == o_atm.kind
873                    });
874                    match o_atom {
875                        Some(atm) => {
876                            other_mag_atoms_ref.remove(atm);
877                        }
878                        None => {
879                            break;
880                        }
881                    }
882                }
883                if !other_mag_atoms_ref.is_empty() {
884                    return false;
885                }
886            } else {
887                return false;
888            }
889        } else if other.magnetic_atoms.is_some() {
890            return false;
891        };
892
893        if let Some(self_ele_atoms) = &self.electric_atoms {
894            if let Some(other_ele_atoms) = &other.electric_atoms {
895                let mut other_ele_atoms_ref: HashSet<_> = other_ele_atoms.iter().collect();
896                for s_atom in self_ele_atoms.iter() {
897                    let o_atom = other_ele_atoms.iter().find(|o_atm| {
898                        (s_atom.coordinates - o_atm.coordinates).norm() < thresh
899                            && s_atom.kind == o_atm.kind
900                    });
901                    match o_atom {
902                        Some(atm) => {
903                            other_ele_atoms_ref.remove(atm);
904                        }
905                        None => {
906                            break;
907                        }
908                    }
909                }
910                if !other_ele_atoms_ref.is_empty() {
911                    return false;
912                }
913            } else {
914                return false;
915            }
916        } else if other.electric_atoms.is_some() {
917            return false;
918        };
919        true
920    }
921}
922
923impl PermutableCollection for Molecule {
924    type Rank = usize;
925
926    /// Determines the permutation of *all* atoms to map `self` to `other`. Special fictitious
927    /// atoms are included after ordinary atoms, with magnetic atoms before electric atoms.
928    ///
929    /// # Arguments
930    ///
931    /// * `other` - Another molecule to be compared with `self`.
932    ///
933    /// # Returns
934    ///
935    /// Returns a permutation that permutes *all* atoms of `self` to give `other`, or `None` if no
936    /// such permutation exists.
937    fn get_perm_of(&self, other: &Self) -> Option<Permutation<Self::Rank>> {
938        let self_recentred = self.recentre();
939        let other_recentred = other.recentre();
940        let o_atoms: HashMap<Atom, usize> = other_recentred
941            .atoms
942            .into_iter()
943            .chain(
944                other_recentred
945                    .magnetic_atoms
946                    .unwrap_or_default()
947                    .into_iter(),
948            )
949            .chain(
950                other_recentred
951                    .electric_atoms
952                    .unwrap_or_default()
953                    .into_iter(),
954            )
955            .enumerate()
956            .map(|(i, atom)| (atom, i))
957            .collect();
958        let image_opt: Option<Vec<Self::Rank>> = self_recentred
959            .atoms
960            .iter()
961            .chain(self_recentred.magnetic_atoms.unwrap_or_default().iter())
962            .chain(self_recentred.electric_atoms.unwrap_or_default().iter())
963            .map(|s_atom| {
964                o_atoms
965                    .get(s_atom)
966                    .or_else(|| {
967                        let thresh = s_atom.threshold;
968                        log::debug!("Unable to retrieve matching original atom by hash. Falling back on distance comparisons with threshold {thresh:.3e}...");
969                        o_atoms.iter().find_map(|(o_atom, o_atom_idx)| {
970                            if s_atom.atomic_number == o_atom.atomic_number
971                                && s_atom.kind == o_atom.kind
972                                && (s_atom.coordinates - o_atom.coordinates).norm() < thresh
973                            {
974                                Some(o_atom_idx)
975                            } else {
976                                None
977                            }
978                        })
979                    })
980                    .copied()
981            })
982            .collect();
983        image_opt.and_then(|image| Permutation::from_image(image).ok())
984    }
985
986    /// Permutes *all* atoms in this molecule (including special fictitious atoms) and places them
987    /// in a new molecule to be returned.
988    ///
989    /// # Arguments
990    ///
991    /// * `perm` - A permutation for the atoms. Special fictitious atoms are included after
992    /// ordinary atoms, with magnetic atoms before electric atoms.
993    ///
994    /// # Returns
995    ///
996    /// A new molecule with the permuted atoms.
997    ///
998    /// # Panics
999    ///
1000    /// Panics if the rank of `perm` does not match the number of atoms in this molecule, or if the
1001    /// permutation results in atoms of different kind (*e.g.* ordinary and magnetic) are permuted
1002    /// into each other.
1003    fn permute(&self, perm: &Permutation<Self::Rank>) -> Result<Self, anyhow::Error> {
1004        let mut p_mol = self.clone();
1005        p_mol.permute_mut(perm)?;
1006        Ok(p_mol)
1007    }
1008
1009    /// Permutes in-place *all* atoms in this molecule (including special fictitious atoms).
1010    ///
1011    /// The in-place rearrangement implementation is taken from
1012    /// [here](https://stackoverflow.com/a/69774341/5112668).
1013    ///
1014    /// # Arguments
1015    ///
1016    /// * `perm` - A permutation for the atoms. Special fictitious atoms are included after
1017    /// ordinary atoms, with magnetic atoms before electric atoms.
1018    ///
1019    /// # Panics
1020    ///
1021    /// Panics if the rank of `perm` does not match the number of atoms in this molecule, or if the
1022    /// permutation results in atoms of different kind (*e.g.* ordinary and magnetic) are permuted
1023    /// into each other.
1024    fn permute_mut(&mut self, perm: &Permutation<Self::Rank>) -> Result<(), anyhow::Error> {
1025        let n_ordinary = self.atoms.len();
1026        let perm_ordinary = Permutation::from_image(perm.image()[0..n_ordinary].to_vec())?;
1027        permute_inplace(&mut self.atoms, &perm_ordinary);
1028
1029        let n_last = if let Some(mag_atoms) = self.magnetic_atoms.as_mut() {
1030            let n_magnetic = mag_atoms.len();
1031            let perm_magnetic = Permutation::from_image(
1032                perm.image()[n_ordinary..(n_ordinary + n_magnetic)]
1033                    .iter()
1034                    .map(|x| x - n_ordinary)
1035                    .collect::<Vec<_>>(),
1036            )?;
1037            permute_inplace(mag_atoms, &perm_magnetic);
1038            n_ordinary + n_magnetic
1039        } else {
1040            n_ordinary
1041        };
1042
1043        if let Some(elec_atoms) = self.electric_atoms.as_mut() {
1044            let n_electric = elec_atoms.len();
1045            let perm_electric = Permutation::from_image(
1046                perm.image()[n_last..(n_last + n_electric)]
1047                    .iter()
1048                    .map(|x| x - n_last)
1049                    .collect::<Vec<_>>(),
1050            )?;
1051            permute_inplace(elec_atoms, &perm_electric);
1052        }
1053        Ok(())
1054    }
1055}