qsym2/bindings/python/
symmetry_group_detection.rs

1//! Python bindings for QSym² symmetry-group detection.
2//!
3//! See [`crate::drivers::symmetry_group_detection`] for more information.
4
5use std::collections::HashMap;
6use std::fmt;
7use std::path::PathBuf;
8use std::sync::Arc;
9
10use anyhow::{self, format_err};
11use derive_builder::Builder;
12use nalgebra::{Point3, Vector3};
13use numpy::{PyArray1, ToPyArray};
14use pyo3::exceptions::PyRuntimeError;
15use pyo3::prelude::*;
16
17use crate::auxiliary::atom::{Atom, ElementMap};
18use crate::auxiliary::molecule::Molecule;
19use crate::drivers::symmetry_group_detection::{
20    SymmetryGroupDetectionDriver, SymmetryGroupDetectionParams,
21};
22use crate::drivers::QSym2Driver;
23#[allow(unused_imports)]
24use crate::io::QSym2FileType;
25use crate::symmetry::symmetry_core::Symmetry;
26use crate::symmetry::symmetry_element::{AntiunitaryKind, SymmetryElementKind};
27use crate::symmetry::symmetry_element_order::ElementOrder;
28
29// ===========================
30// Struct and enum definitions
31// ===========================
32
33// ----------
34// PyMolecule
35// ----------
36
37/// Python-exposed structure to marshall molecular structure information between Rust and Python.
38///
39/// # Constructor arguments
40///
41/// * `atoms` - The ordinary atoms in the molecule. Python type: `list[tuple[str, tuple[float,
42/// float, float]]]`.
43/// * `threshold` - Threshold for comparing molecules. Python type: `float`.
44/// * `magnetic_field` - An optional uniform external magnetic field. Python type:
45/// `Optional[tuple[float, float, float]]`.
46/// * `electric_field` - An optional uniform external electric field. Python type:
47/// `Optional[tuple[float, float, float]]`.
48#[pyclass]
49#[derive(Clone)]
50pub struct PyMolecule {
51    /// The ordinary atoms in the molecule.
52    ///
53    /// Python type: `list[tuple[str, tuple[float, float, float]]]`
54    #[pyo3(get)]
55    pub atoms: Vec<(String, [f64; 3])>,
56
57    /// An optional uniform external magnetic field.
58    ///
59    /// Python type: `Optional[tuple[float, float, float]]`
60    #[pyo3(get)]
61    pub magnetic_field: Option<[f64; 3]>,
62
63    /// An optional uniform external electric field.
64    ///
65    /// Python type: `Optional[tuple[float, float, float]]`
66    #[pyo3(get)]
67    pub electric_field: Option<[f64; 3]>,
68
69    /// Threshold for comparing molecules.
70    ///
71    /// Python type: `float`
72    #[pyo3(get)]
73    pub threshold: f64,
74}
75
76#[pymethods]
77impl PyMolecule {
78    /// Creates a new `PyMolecule` structure.
79    ///
80    /// # Arguments
81    ///
82    /// * `atoms` - The ordinary atoms in the molecule. Python type: `list[tuple[str, tuple[float,
83    /// float, float]]]`.
84    /// * `threshold` - Threshold for comparing molecules. Python type: `float`.
85    /// * `magnetic_field` - An optional uniform external magnetic field. Python type:
86    /// `Optional[tuple[float, float, float]]`.
87    /// * `electric_field` - An optional uniform external electric field. Python type:
88    /// `Optional[tuple[float, float, float]]`.
89    #[new]
90    #[pyo3(signature = (atoms, threshold, magnetic_field=None, electric_field=None))]
91    pub fn new(
92        atoms: Vec<(String, [f64; 3])>,
93        threshold: f64,
94        magnetic_field: Option<[f64; 3]>,
95        electric_field: Option<[f64; 3]>,
96    ) -> Self {
97        Self {
98            atoms,
99            threshold,
100            magnetic_field,
101            electric_field,
102        }
103    }
104}
105
106impl From<PyMolecule> for Molecule {
107    fn from(pymol: PyMolecule) -> Self {
108        let emap = ElementMap::new();
109        let mut mol = Self::from_atoms(
110            &pymol
111                .atoms
112                .iter()
113                .map(|(ele, r)| {
114                    Atom::new_ordinary(ele, Point3::new(r[0], r[1], r[2]), &emap, pymol.threshold)
115                })
116                .collect::<Vec<_>>(),
117            pymol.threshold,
118        );
119        mol.set_magnetic_field(pymol.magnetic_field.map(Vector3::from_iterator));
120        mol.set_electric_field(pymol.electric_field.map(Vector3::from_iterator));
121        mol
122    }
123}
124
125// ---------------------
126// PySymmetryElementKind
127// ---------------------
128
129/// Python-exposed structure to marshall symmetry element kind information one-way from Rust to
130/// Python.
131#[pyclass(eq, eq_int)]
132#[derive(Clone, Hash, PartialEq, Eq)]
133pub enum PySymmetryElementKind {
134    /// Variant denoting proper symmetry elements.
135    Proper,
136
137    /// Variant denoting time-reversed proper symmetry elements.
138    ProperTR,
139
140    /// Variant denoting improper symmetry elements (mirror-plane convention).
141    ImproperMirrorPlane,
142
143    /// Variant denoting time-reversed improper symmetry elements (mirror-plane convention).
144    ImproperMirrorPlaneTR,
145}
146
147impl fmt::Display for PySymmetryElementKind {
148    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
149        match self {
150            PySymmetryElementKind::Proper => write!(f, "Proper"),
151            PySymmetryElementKind::ProperTR => write!(f, "Time-reversed proper"),
152            PySymmetryElementKind::ImproperMirrorPlane => {
153                write!(f, "Improper (mirror-plane convention)")
154            }
155            PySymmetryElementKind::ImproperMirrorPlaneTR => {
156                write!(f, "Time-reversed improper (mirror-plane convention)")
157            }
158        }
159    }
160}
161
162impl TryFrom<&SymmetryElementKind> for PySymmetryElementKind {
163    type Error = anyhow::Error;
164
165    fn try_from(symkind: &SymmetryElementKind) -> Result<Self, Self::Error> {
166        match symkind {
167            SymmetryElementKind::Proper(None) => Ok(Self::Proper),
168            SymmetryElementKind::Proper(Some(AntiunitaryKind::TimeReversal)) => Ok(Self::ProperTR),
169            SymmetryElementKind::ImproperMirrorPlane(None) => Ok(Self::ImproperMirrorPlane),
170            SymmetryElementKind::ImproperMirrorPlane(Some(AntiunitaryKind::TimeReversal)) => {
171                Ok(Self::ImproperMirrorPlaneTR)
172            }
173            _ => Err(format_err!(
174                "Symmetry element kind `{symkind}` is not yet supported in Python."
175            )),
176        }
177    }
178}
179
180// ----------
181// PySymmetry
182// ----------
183
184/// Python-exposed structure to marshall symmetry information one-way from Rust to Python.
185#[pyclass]
186#[derive(Clone, Builder)]
187pub struct PySymmetry {
188    /// The name of the symmetry group.
189    #[pyo3(get)]
190    group_name: String,
191
192    /// The symmetry elements.
193    ///
194    /// Python type: `dict[PySymmetryElementKind, dict[int, list[numpy.1darray[float]]]]`
195    elements: HashMap<PySymmetryElementKind, HashMap<i32, Vec<Arc<Py<PyArray1<f64>>>>>>,
196
197    /// The symmetry generators.
198    ///
199    /// Python type: `dict[PySymmetryElementKind, dict[int, list[numpy.1darray[float]]]]`
200    generators: HashMap<PySymmetryElementKind, HashMap<i32, Vec<Arc<Py<PyArray1<f64>>>>>>,
201}
202
203impl PySymmetry {
204    fn builder() -> PySymmetryBuilder {
205        PySymmetryBuilder::default()
206    }
207}
208
209#[pymethods]
210impl PySymmetry {
211    /// Returns a boolean indicating if the group is infinite.
212    pub fn is_infinite(&self) -> bool {
213        self.elements
214            .values()
215            .any(|kind_elements| kind_elements.contains_key(&-1))
216            || self
217                .generators
218                .values()
219                .any(|kind_generators| kind_generators.contains_key(&-1))
220    }
221
222    /// Returns symmetry elements of all *finite* orders of a given kind.
223    ///
224    /// # Arguments
225    ///
226    /// * `kind` - The symmetry element kind. Python type: `PySymmetryElementKind`.
227    ///
228    /// # Returns
229    ///
230    /// A hashmap where the keys are integers indicating the orders of the elements and the values
231    /// are vectors of one-dimensional arrays, each of which gives the axis of a symmetry element.
232    /// If the order value is `-1`, then the associated elements have infinite order.
233    ///
234    /// Python type: `dict[int, list[numpy.1darray[float]]]`.
235    pub fn get_elements_of_kind(
236        &self,
237        kind: &PySymmetryElementKind,
238    ) -> PyResult<HashMap<i32, Vec<Py<PyArray1<f64>>>>> {
239        self.elements
240            .get(kind)
241            .cloned()
242            .and_then(|elements| {
243                elements
244                    .iter()
245                    .map(|(order, axes_arc)| {
246                        let axes_opt = axes_arc
247                            .iter()
248                            .map(|axis_arc| Arc::into_inner(axis_arc.clone()))
249                            .collect::<Option<Vec<_>>>();
250                        axes_opt.map(|axes| (order.clone(), axes))
251                    })
252                    .collect::<Option<HashMap<_, _>>>()
253            })
254            .ok_or(PyRuntimeError::new_err(format!(
255                "Elements of kind `{kind}` not found."
256            )))
257    }
258
259    /// Returns symmetry generators of *finite*  and *infinite* orders of a given kind.
260    ///
261    /// # Arguments
262    ///
263    /// * `kind` - The symmetry generator kind. Python type: `PySymmetryElementKind`.
264    ///
265    /// # Returns
266    ///
267    /// A hashmap where the keys are integers indicating the orders of the generators and the values
268    /// are vectors of one-dimensional arrays, each of which gives the axis of a symmetry generator.
269    /// If the order value is `-1`, then the associated generators have infinite order.
270    ///
271    /// Python type: `dict[int, list[numpy.1darray[float]]]`.
272    pub fn get_generators_of_kind(
273        &self,
274        kind: &PySymmetryElementKind,
275    ) -> PyResult<HashMap<i32, Vec<Py<PyArray1<f64>>>>> {
276        self.generators
277            .get(kind)
278            .cloned()
279            .and_then(|generators| {
280                generators
281                    .iter()
282                    .map(|(order, axes_arc)| {
283                        let axes_opt = axes_arc
284                            .iter()
285                            .map(|axis_arc| Arc::into_inner(axis_arc.clone()))
286                            .collect::<Option<Vec<_>>>();
287                        axes_opt.map(|axes| (order.clone(), axes))
288                    })
289                    .collect::<Option<HashMap<_, _>>>()
290            })
291            .ok_or(PyRuntimeError::new_err(format!(
292                "Elements of kind `{kind}` not found."
293            )))
294    }
295}
296
297impl TryFrom<&Symmetry> for PySymmetry {
298    type Error = anyhow::Error;
299
300    fn try_from(sym: &Symmetry) -> Result<Self, Self::Error> {
301        let group_name = sym
302            .group_name
303            .clone()
304            .ok_or(format_err!("Symmetry group name not found."))?;
305        let elements = sym
306            .elements
307            .iter()
308            .map(|(symkind, kind_elements)| {
309                let pysymkind = PySymmetryElementKind::try_from(symkind)?;
310                let pykind_elements = kind_elements
311                    .iter()
312                    .map(|(order, order_elements)| {
313                        let order_i32 = match order {
314                            ElementOrder::Int(ord) => i32::try_from(*ord)?,
315                            ElementOrder::Inf => -1,
316                        };
317                        let pyorder_elements = order_elements
318                            .iter()
319                            .map(|ele| {
320                                Arc::new(Python::with_gil(|py| {
321                                    ele.raw_axis()
322                                        .iter()
323                                        .cloned()
324                                        .collect::<Vec<_>>()
325                                        .to_pyarray(py)
326                                        .unbind()
327                                }))
328                            })
329                            .collect::<Vec<_>>();
330                        Ok::<_, Self::Error>((order_i32, pyorder_elements))
331                    })
332                    .collect::<Result<HashMap<i32, Vec<_>>, _>>()?;
333                Ok::<_, Self::Error>((pysymkind, pykind_elements))
334            })
335            .collect::<Result<HashMap<_, _>, _>>()?;
336
337        let generators = sym
338            .generators
339            .iter()
340            .map(|(symkind, kind_generators)| {
341                let pysymkind = PySymmetryElementKind::try_from(symkind)?;
342                let pykind_generators = kind_generators
343                    .iter()
344                    .map(|(order, order_generators)| {
345                        let order_i32 = match order {
346                            ElementOrder::Int(ord) => i32::try_from(*ord)?,
347                            ElementOrder::Inf => -1,
348                        };
349                        let pyorder_generators = order_generators
350                            .iter()
351                            .map(|ele| {
352                                Arc::new(Python::with_gil(|py| {
353                                    ele.raw_axis()
354                                        .iter()
355                                        .cloned()
356                                        .collect::<Vec<_>>()
357                                        .to_pyarray(py)
358                                        .unbind()
359                                }))
360                            })
361                            .collect::<Vec<_>>();
362                        Ok::<_, Self::Error>((order_i32, pyorder_generators))
363                    })
364                    .collect::<Result<HashMap<i32, Vec<_>>, _>>()?;
365                Ok::<_, Self::Error>((pysymkind, pykind_generators))
366            })
367            .collect::<Result<HashMap<_, _>, _>>()?;
368
369        PySymmetry::builder()
370            .group_name(group_name)
371            .elements(elements)
372            .generators(generators)
373            .build()
374            .map_err(|err| format_err!(err))
375    }
376}
377
378// =========
379// Functions
380// =========
381
382/// Python-exposed function to perform symmetry-group detection and log the result via the
383/// `qsym2-output` logger at the `INFO` level.
384///
385/// See [`crate::drivers::symmetry_group_detection`] for more information.
386///
387/// # Arguments
388///
389/// * `inp_xyz` - An optional string providing the path to an XYZ file containing the molecule to
390/// be analysed. Only one of `inp_xyz` or `inp_mol` can be specified. Python type: `Optional[str]`.
391/// * `inp_mol` - An optional `PyMolecule` structure containing the molecule to be analysed. Only
392/// one of `inp_xyz` or `inp_mol` can be specified. Python type: `PyMolecule`.
393/// * `out_sym` - An optional name for the [`QSym2FileType::Sym`] file to be saved that contains
394/// the serialised results of the symmetry-group detection. Python type: `Optional[str]`.
395/// * `moi_thresholds` - Thresholds for comparing moments of inertia. Python type: `list[float]`.
396/// * `distance_thresholds` - Thresholds for comparing distances. Python type: `list[float]`.
397/// * `time_reversal` - A boolean indicating whether elements involving time reversal should also
398/// be considered. Python type: `bool`.
399/// * `write_symmetry_elements` - A boolean indicating if detected symmetry elements should be
400/// printed in the output. Python type: `bool`.
401/// * `fictitious_magnetic_field` - An optional fictitious uniform external magnetic field. Python
402/// type: `Optional[tuple[float, float, float]]`.
403/// * `fictitious_electric_field` - An optional fictitious uniform external electric field. Python
404/// type: `Optional[tuple[float, float, float]]`.
405///
406/// # Errors
407///
408/// Returns an error if any intermediate step in the symmetry-group detection procedure fails.
409#[pyfunction]
410#[pyo3(signature = (
411    inp_xyz,
412    inp_mol,
413    out_sym,
414    moi_thresholds,
415    distance_thresholds,
416    time_reversal,
417    write_symmetry_elements=true,
418    fictitious_magnetic_field=None,
419    fictitious_electric_field=None,
420))]
421pub fn detect_symmetry_group(
422    py: Python<'_>,
423    inp_xyz: Option<PathBuf>,
424    inp_mol: Option<PyMolecule>,
425    out_sym: Option<PathBuf>,
426    moi_thresholds: Vec<f64>,
427    distance_thresholds: Vec<f64>,
428    time_reversal: bool,
429    write_symmetry_elements: bool,
430    fictitious_magnetic_field: Option<[f64; 3]>,
431    fictitious_electric_field: Option<[f64; 3]>,
432) -> PyResult<(PySymmetry, Option<PySymmetry>)> {
433    py.allow_threads(|| {
434        let params = SymmetryGroupDetectionParams::builder()
435            .distance_thresholds(&distance_thresholds)
436            .moi_thresholds(&moi_thresholds)
437            .time_reversal(time_reversal)
438            .fictitious_magnetic_fields(
439                fictitious_magnetic_field
440                    .map(|bs| vec![(Point3::<f64>::origin(), Vector3::new(bs[0], bs[1], bs[2]))]),
441            )
442            .fictitious_electric_fields(
443                fictitious_electric_field
444                    .map(|es| vec![(Point3::<f64>::origin(), Vector3::new(es[0], es[1], es[2]))]),
445            )
446            .field_origin_com(true)
447            .write_symmetry_elements(write_symmetry_elements)
448            .result_save_name(out_sym)
449            .build()
450            .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
451        let inp_mol = inp_mol.map(Molecule::from);
452        let mut pd_driver = SymmetryGroupDetectionDriver::builder()
453            .parameters(&params)
454            .xyz(inp_xyz)
455            .molecule(inp_mol.as_ref())
456            .build()
457            .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
458        pd_driver
459            .run()
460            .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
461        let pyunitary_symmetry: PySymmetry = (&pd_driver
462            .result()
463            .map_err(|err| PyRuntimeError::new_err(err.to_string()))?
464            .unitary_symmetry)
465            .try_into()
466            .map_err(|err: anyhow::Error| PyRuntimeError::new_err(err.to_string()))?;
467        let pymagnetic_symmetry: Option<PySymmetry> = pd_driver
468            .result()
469            .map_err(|err| PyRuntimeError::new_err(err.to_string()))?
470            .magnetic_symmetry
471            .as_ref()
472            .map(|magsym| {
473                magsym
474                    .try_into()
475                    .map_err(|err: anyhow::Error| PyRuntimeError::new_err(err.to_string()))
476            })
477            .transpose()?;
478        Ok((pyunitary_symmetry, pymagnetic_symmetry))
479    })
480}