1use 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#[pyclass]
49#[derive(Clone)]
50pub struct PyMolecule {
51 #[pyo3(get)]
55 pub atoms: Vec<(String, [f64; 3])>,
56
57 #[pyo3(get)]
61 pub magnetic_field: Option<[f64; 3]>,
62
63 #[pyo3(get)]
67 pub electric_field: Option<[f64; 3]>,
68
69 #[pyo3(get)]
73 pub threshold: f64,
74}
75
76#[pymethods]
77impl PyMolecule {
78 #[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#[pyclass(eq, eq_int)]
132#[derive(Clone, Hash, PartialEq, Eq)]
133pub enum PySymmetryElementKind {
134 Proper,
136
137 ProperTR,
139
140 ImproperMirrorPlane,
142
143 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#[pyclass]
186#[derive(Clone, Builder)]
187pub struct PySymmetry {
188 #[pyo3(get)]
190 group_name: String,
191
192 elements: HashMap<PySymmetryElementKind, HashMap<i32, Vec<Arc<Py<PyArray1<f64>>>>>>,
196
197 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 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 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 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#[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(¶ms)
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}