qsym2/target/orbital/
mod.rs

1//! Orbitals.
2
3use std::fmt;
4
5use anyhow::{self, format_err};
6use approx;
7use derive_builder::Builder;
8use ndarray::{s, Array1, Array2, Ix2};
9use ndarray_einsum_beta::*;
10use ndarray_linalg::types::Lapack;
11use num_complex::{Complex, ComplexFloat};
12use num_traits::float::{Float, FloatConst};
13
14use crate::angmom::spinor_rotation_3d::{SpinConstraint, SpinOrbitCoupled, StructureConstraint};
15use crate::auxiliary::molecule::Molecule;
16use crate::basis::ao::BasisAngularOrder;
17use crate::target::density::Density;
18
19#[cfg(test)]
20mod orbital_tests;
21
22pub mod orbital_analysis;
23mod orbital_transformation;
24
25// ==================
26// Struct definitions
27// ==================
28
29/// Structure to manage molecular orbitals. Each molecular orbital is essentially a one-electron
30/// Slater determinant.
31#[derive(Builder, Clone)]
32#[builder(build_fn(validate = "Self::validate"))]
33pub struct MolecularOrbital<'a, T, SC>
34where
35    T: ComplexFloat + Lapack,
36    SC: StructureConstraint,
37{
38    /// The structure constraint associated with the coefficients describing this molecular orbital.
39    structure_constraint: SC,
40
41    /// If the structure constraint allows for multiple components, this gives the component index of
42    /// this molecular orbital.
43    component_index: usize,
44
45    /// A boolean indicating if the orbital has been acted on by an antiunitary operation. This is
46    /// so that the correct metric can be used during overlap evaluation.
47    #[builder(default = "false")]
48    complex_conjugated: bool,
49
50    /// The angular order of the basis functions with respect to which the coefficients are
51    /// expressed.
52    bao: &'a BasisAngularOrder<'a>,
53
54    /// A boolean indicating if inner products involving this molecular orbital should be the
55    /// complex-symmetric bilinear form, rather than the conventional Hermitian sesquilinear form.
56    complex_symmetric: bool,
57
58    /// The associated molecule.
59    mol: &'a Molecule,
60
61    /// The coefficients describing this molecular orbital.
62    coefficients: Array1<T>,
63
64    /// The energy of this molecular orbital.
65    #[builder(default = "None")]
66    energy: Option<T>,
67
68    /// The threshold for comparing determinants.
69    threshold: <T as ComplexFloat>::Real,
70}
71
72impl<'a, T, SC> MolecularOrbitalBuilder<'a, T, SC>
73where
74    T: ComplexFloat + Lapack,
75    SC: StructureConstraint,
76{
77    fn validate(&self) -> Result<(), String> {
78        let bao = self
79            .bao
80            .ok_or("No `BasisAngularOrder` found.".to_string())?;
81        let nbas = bao.n_funcs();
82        let coefficients = self
83            .coefficients
84            .as_ref()
85            .ok_or("No coefficients found.".to_string())?;
86        let component_index = self
87            .component_index
88            .ok_or("No `component_index` found.".to_string())?;
89
90        let structcons = self
91            .structure_constraint
92            .as_ref()
93            .ok_or("No structure constraints found.".to_string())?;
94        let coefficients_shape_check = {
95            let nrows = nbas * structcons.n_explicit_comps_per_coefficient_matrix();
96            if !coefficients.shape()[0] == nrows {
97                log::error!(
98                    "Unexpected shapes of coefficient vector: {} {} expected, but {} found.",
99                    nrows,
100                    if nrows == 1 { "row" } else { "rows" },
101                    coefficients.shape()[0],
102                );
103                false
104            } else {
105                true
106            }
107        };
108
109        let mol = self.mol.ok_or("No molecule found.".to_string())?;
110        let natoms_check = mol.atoms.len() == bao.n_atoms();
111        if !natoms_check {
112            log::error!("The number of atoms in the molecule does not match the number of local sites in the basis.");
113        }
114
115        if coefficients_shape_check && natoms_check {
116            Ok(())
117        } else {
118            Err("Molecular orbital validation failed.".to_string())
119        }
120    }
121}
122
123impl<'a, T, SC> MolecularOrbital<'a, T, SC>
124where
125    T: ComplexFloat + Clone + Lapack,
126    SC: StructureConstraint + Clone,
127{
128    /// Returns a builder to construct a new [`MolecularOrbital`].
129    pub fn builder() -> MolecularOrbitalBuilder<'a, T, SC> {
130        MolecularOrbitalBuilder::default()
131    }
132}
133
134impl<'a, T, SC> MolecularOrbital<'a, T, SC>
135where
136    T: ComplexFloat + Clone + Lapack,
137    SC: StructureConstraint,
138{
139    /// Returns a shared reference to the coefficient array.
140    pub fn coefficients(&self) -> &Array1<T> {
141        &self.coefficients
142    }
143
144    /// Returns a shared reference to the structure constraint.
145    pub fn structure_constraint(&self) -> &SC {
146        &self.structure_constraint
147    }
148
149    /// Returns a shared reference to the [`BasisAngularOrder`] description of the basis in which
150    /// the orbital coefficients are written.
151    pub fn bao(&self) -> &BasisAngularOrder {
152        self.bao
153    }
154
155    /// Returns the molecule associated with this molecular orbital.
156    pub fn mol(&self) -> &Molecule {
157        self.mol
158    }
159
160    /// Returns the complex-symmetric flag of the molecular orbital.
161    pub fn complex_symmetric(&self) -> bool {
162        self.complex_symmetric
163    }
164
165    /// Returns the threshold with which molecular orbitals are compared.
166    pub fn threshold(&self) -> <T as ComplexFloat>::Real {
167        self.threshold
168    }
169}
170
171impl<'a, T> MolecularOrbital<'a, T, SpinConstraint>
172where
173    T: ComplexFloat + Clone + Lapack,
174{
175    /// Augments the encoding of coefficients in this molecular orbital to that in the
176    /// corresponding generalised spin constraint.
177    ///
178    /// # Returns
179    ///
180    /// The equivalent molecular orbital with the coefficients encoded in the generalised spin
181    /// constraint.
182    pub fn to_generalised(&self) -> Self {
183        match self.structure_constraint {
184            SpinConstraint::Restricted(n) => {
185                let nbas = self.bao.n_funcs();
186
187                let cr = &self.coefficients;
188                let mut cg = Array1::<T>::zeros(nbas * usize::from(n));
189                let start = nbas * self.component_index;
190                let end = nbas * (self.component_index + 1);
191                cg.slice_mut(s![start..end]).assign(cr);
192                Self::builder()
193                    .coefficients(cg)
194                    .bao(self.bao)
195                    .mol(self.mol)
196                    .structure_constraint(SpinConstraint::Generalised(n, false))
197                    .component_index(0)
198                    .complex_symmetric(self.complex_symmetric)
199                    .threshold(self.threshold)
200                    .build()
201                    .expect("Unable to construct a generalised molecular orbital.")
202            }
203            SpinConstraint::Unrestricted(n, increasingm) => {
204                let nbas = self.bao.n_funcs();
205
206                let cr = &self.coefficients;
207                let mut cg = Array1::<T>::zeros(nbas * usize::from(n));
208                let start = nbas * self.component_index;
209                let end = nbas * (self.component_index + 1);
210                cg.slice_mut(s![start..end]).assign(cr);
211                Self::builder()
212                    .coefficients(cg)
213                    .bao(self.bao)
214                    .mol(self.mol)
215                    .structure_constraint(SpinConstraint::Generalised(n, increasingm))
216                    .component_index(0)
217                    .complex_symmetric(self.complex_symmetric)
218                    .threshold(self.threshold)
219                    .build()
220                    .expect("Unable to construct a generalised molecular orbital.")
221            }
222            SpinConstraint::Generalised(_, _) => self.clone(),
223        }
224    }
225}
226
227impl<'a> MolecularOrbital<'a, f64, SpinConstraint> {
228    /// Constructs the total density of the molecular orbital.
229    pub fn to_total_density(&'a self) -> Result<Density<'a, f64>, anyhow::Error> {
230        match self.structure_constraint {
231            SpinConstraint::Restricted(nspins) => {
232                let denmat = f64::from(nspins)
233                    * einsum(
234                        "m,n->mn",
235                        &[&self.coefficients.view(), &self.coefficients.view()],
236                    )
237                    .expect("Unable to construct a density matrix from the coefficient matrix.")
238                    .into_dimensionality::<Ix2>()
239                    .expect("Unable to convert the resultant density matrix to two dimensions.");
240                Density::<f64>::builder()
241                    .density_matrix(denmat)
242                    .bao(self.bao())
243                    .mol(self.mol())
244                    .complex_symmetric(self.complex_symmetric())
245                    .threshold(self.threshold())
246                    .build()
247                    .map_err(|err| format_err!(err))
248            }
249            SpinConstraint::Unrestricted(_, _) => {
250                let denmat = einsum(
251                    "m,n->mn",
252                    &[&self.coefficients.view(), &self.coefficients.view()],
253                )
254                .expect("Unable to construct a density matrix from the coefficient matrix.")
255                .into_dimensionality::<Ix2>()
256                .expect("Unable to convert the resultant density matrix to two dimensions.");
257                Density::<f64>::builder()
258                    .density_matrix(denmat)
259                    .bao(self.bao())
260                    .mol(self.mol())
261                    .complex_symmetric(self.complex_symmetric())
262                    .threshold(self.threshold())
263                    .build()
264                    .map_err(|err| format_err!(err))
265            }
266            SpinConstraint::Generalised(nspins, _) => {
267                let full_denmat = einsum(
268                    "m,n->mn",
269                    &[&self.coefficients.view(), &self.coefficients.view()],
270                )
271                .expect("Unable to construct a density matrix from the coefficient matrix.")
272                .into_dimensionality::<Ix2>()
273                .expect("Unable to convert the resultant density matrix to two dimensions.");
274                let nspatial = self.bao().n_funcs();
275                let denmat = (0..usize::from(nspins)).fold(
276                    Array2::<f64>::zeros((nspatial, nspatial)),
277                    |acc, ispin| {
278                        acc + full_denmat.slice(s![
279                            ispin * nspatial..(ispin + 1) * nspatial,
280                            ispin * nspatial..(ispin + 1) * nspatial
281                        ])
282                    },
283                );
284                Density::<f64>::builder()
285                    .density_matrix(denmat)
286                    .bao(self.bao())
287                    .mol(self.mol())
288                    .complex_symmetric(self.complex_symmetric())
289                    .threshold(self.threshold())
290                    .build()
291                    .map_err(|err| format_err!(err))
292            }
293        }
294    }
295}
296
297impl<'a, T> MolecularOrbital<'a, Complex<T>, SpinConstraint>
298where
299    T: Float + FloatConst + Lapack + From<u16>,
300    Complex<T>: Lapack,
301{
302    /// Constructs the total density of the molecular orbital.
303    pub fn to_total_density(&'a self) -> Result<Density<'a, Complex<T>>, anyhow::Error> {
304        match self.structure_constraint {
305            SpinConstraint::Restricted(nspins) => {
306                let nspins_t = Complex::<T>::from(<T as From<u16>>::from(nspins));
307                let denmat = einsum(
308                    "m,n->mn",
309                    &[
310                        &self.coefficients.view(),
311                        &self.coefficients.map(Complex::conj).view(),
312                    ],
313                )
314                .expect("Unable to construct a density matrix from the coefficient matrix.")
315                .into_dimensionality::<Ix2>()
316                .expect("Unable to convert the resultant density matrix to two dimensions.")
317                .map(|x| x * nspins_t);
318                Density::<Complex<T>>::builder()
319                    .density_matrix(denmat)
320                    .bao(self.bao())
321                    .mol(self.mol())
322                    .complex_symmetric(self.complex_symmetric())
323                    .threshold(self.threshold())
324                    .build()
325                    .map_err(|err| format_err!(err))
326            }
327            SpinConstraint::Unrestricted(_, _) => {
328                let denmat = einsum(
329                    "m,n->mn",
330                    &[
331                        &self.coefficients.view(),
332                        &self.coefficients.map(Complex::conj).view(),
333                    ],
334                )
335                .expect("Unable to construct a density matrix from the coefficient matrix.")
336                .into_dimensionality::<Ix2>()
337                .expect("Unable to convert the resultant density matrix to two dimensions.");
338                Density::<Complex<T>>::builder()
339                    .density_matrix(denmat)
340                    .bao(self.bao())
341                    .mol(self.mol())
342                    .complex_symmetric(self.complex_symmetric())
343                    .threshold(self.threshold())
344                    .build()
345                    .map_err(|err| format_err!(err))
346            }
347            SpinConstraint::Generalised(nspins, _) => {
348                let full_denmat = einsum(
349                    "m,n->mn",
350                    &[
351                        &self.coefficients.view(),
352                        &self.coefficients.map(Complex::conj).view(),
353                    ],
354                )
355                .expect("Unable to construct a density matrix from the coefficient matrix.")
356                .into_dimensionality::<Ix2>()
357                .expect("Unable to convert the resultant density matrix to two dimensions.");
358                let nspatial = self.bao().n_funcs();
359                let denmat = (0..usize::from(nspins)).fold(
360                    Array2::<Complex<T>>::zeros((nspatial, nspatial)),
361                    |acc, ispin| {
362                        acc + full_denmat.slice(s![
363                            ispin * nspatial..(ispin + 1) * nspatial,
364                            ispin * nspatial..(ispin + 1) * nspatial
365                        ])
366                    },
367                );
368                Density::<Complex<T>>::builder()
369                    .density_matrix(denmat)
370                    .bao(self.bao())
371                    .mol(self.mol())
372                    .complex_symmetric(self.complex_symmetric())
373                    .threshold(self.threshold())
374                    .build()
375                    .map_err(|err| format_err!(err))
376            }
377        }
378    }
379}
380
381impl<'a, T> MolecularOrbital<'a, Complex<T>, SpinOrbitCoupled>
382where
383    T: Float + FloatConst + Lapack + From<u16>,
384    Complex<T>: Lapack,
385{
386    /// Constructs the total density of the molecular orbital.
387    pub fn to_total_density(&'a self) -> Result<Density<'a, Complex<T>>, anyhow::Error> {
388        match self.structure_constraint {
389            SpinOrbitCoupled::JAdapted(ncomps) => {
390                let full_denmat = einsum(
391                    "m,n->mn",
392                    &[
393                        &self.coefficients.view(),
394                        &self.coefficients.map(Complex::conj).view(),
395                    ],
396                )
397                .expect("Unable to construct a density matrix from the coefficient matrix.")
398                .into_dimensionality::<Ix2>()
399                .expect("Unable to convert the resultant density matrix to two dimensions.");
400                let nspatial = self.bao().n_funcs();
401                let denmat = (0..usize::from(ncomps)).fold(
402                    Array2::<Complex<T>>::zeros((nspatial, nspatial)),
403                    |acc, icomp| {
404                        acc + full_denmat.slice(s![
405                            icomp * nspatial..(icomp + 1) * nspatial,
406                            icomp * nspatial..(icomp + 1) * nspatial
407                        ])
408                    },
409                );
410                Density::<Complex<T>>::builder()
411                    .density_matrix(denmat)
412                    .bao(self.bao())
413                    .mol(self.mol())
414                    .complex_symmetric(self.complex_symmetric())
415                    .threshold(self.threshold())
416                    .build()
417                    .map_err(|err| format_err!(err))
418            }
419        }
420    }
421}
422
423// =====================
424// Trait implementations
425// =====================
426
427// ----
428// From
429// ----
430impl<'a, T, SC> From<MolecularOrbital<'a, T, SC>> for MolecularOrbital<'a, Complex<T>, SC>
431where
432    T: Float + FloatConst + Lapack,
433    Complex<T>: Lapack,
434    SC: StructureConstraint + Clone,
435{
436    fn from(value: MolecularOrbital<'a, T, SC>) -> Self {
437        MolecularOrbital::<'a, Complex<T>, SC>::builder()
438            .coefficients(value.coefficients.map(Complex::from))
439            .bao(value.bao)
440            .mol(value.mol)
441            .structure_constraint(value.structure_constraint)
442            .component_index(value.component_index)
443            .complex_symmetric(value.complex_symmetric)
444            .threshold(value.threshold)
445            .build()
446            .expect("Unable to construct a complex molecular orbital.")
447    }
448}
449
450// ---------
451// PartialEq
452// ---------
453impl<'a, T, SC> PartialEq for MolecularOrbital<'a, T, SC>
454where
455    T: ComplexFloat<Real = f64> + Lapack,
456    SC: StructureConstraint + PartialEq,
457{
458    fn eq(&self, other: &Self) -> bool {
459        let thresh = (self.threshold * other.threshold).sqrt();
460        let coefficients_eq = approx::relative_eq!(
461            (&self.coefficients - &other.coefficients)
462                .map(|x| ComplexFloat::abs(*x).powi(2))
463                .sum()
464                .sqrt(),
465            0.0,
466            epsilon = thresh,
467            max_relative = thresh,
468        );
469        self.structure_constraint == other.structure_constraint
470            && self.component_index == other.component_index
471            && self.bao == other.bao
472            && self.mol == other.mol
473            && coefficients_eq
474    }
475}
476
477// --
478// Eq
479// --
480impl<'a, T, SC> Eq for MolecularOrbital<'a, T, SC>
481where
482    T: ComplexFloat<Real = f64> + Lapack,
483    SC: StructureConstraint + Eq,
484{
485}
486
487// -----
488// Debug
489// -----
490impl<'a, T, SC> fmt::Debug for MolecularOrbital<'a, T, SC>
491where
492    T: fmt::Debug + ComplexFloat + Lapack,
493    SC: StructureConstraint + fmt::Debug,
494{
495    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
496        write!(
497            f,
498            "MolecularOrbital[{:?} (spin index {}): coefficient array of length {}]",
499            self.structure_constraint,
500            self.component_index,
501            self.coefficients.len()
502        )?;
503        Ok(())
504    }
505}
506
507// -------
508// Display
509// -------
510impl<'a, T, SC> fmt::Display for MolecularOrbital<'a, T, SC>
511where
512    T: fmt::Display + ComplexFloat + Lapack,
513    SC: StructureConstraint + fmt::Display,
514{
515    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
516        write!(
517            f,
518            "MolecularOrbital[{} (spin index {}): coefficient array of length {}]",
519            self.structure_constraint,
520            self.component_index,
521            self.coefficients.len()
522        )?;
523        Ok(())
524    }
525}