1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

//! Implementation of symmetry transformations for vibrational coordinates.

use ndarray::{concatenate, s, Array2, Axis, LinalgScalar, ScalarOperand};
use ndarray_linalg::types::Lapack;
use num_complex::{Complex, ComplexFloat};

use crate::auxiliary::molecule::Molecule;
use crate::basis::ao::{BasisAngularOrder, BasisAtom, BasisShell, CartOrder, ShellOrder};
use crate::permutation::{IntoPermutation, PermutableCollection, Permutation};
use crate::symmetry::symmetry_element::SymmetryOperation;
use crate::symmetry::symmetry_transformation::{
    assemble_sh_rotation_3d_matrices, permute_array_by_atoms, ComplexConjugationTransformable,
    SpatialUnitaryTransformable, SpinUnitaryTransformable, SymmetryTransformable,
    TimeReversalTransformable, TransformationError,
};
use crate::target::vibration::VibrationalCoordinate;

// ---------------------------
// SpatialUnitaryTransformable
// ---------------------------
impl<'a, T> SpatialUnitaryTransformable for VibrationalCoordinate<'a, T>
where
    T: ComplexFloat + LinalgScalar + ScalarOperand + Copy + Lapack,
    f64: Into<T>,
{
    fn transform_spatial_mut(
        &mut self,
        rmat: &Array2<f64>,
        perm: Option<&Permutation<usize>>,
    ) -> Result<&mut Self, TransformationError> {
        let vib_bao = construct_vibration_bao(self.mol);
        let tmats: Vec<Array2<T>> = assemble_sh_rotation_3d_matrices(&vib_bao, rmat, perm)
            .map_err(|err| TransformationError(err.to_string()))?
            .iter()
            .map(|tmat| tmat.map(|&x| x.into()))
            .collect();
        let pbao = if let Some(p) = perm {
            vib_bao
                .permute(p)
                .map_err(|err| TransformationError(err.to_string()))?
        } else {
            vib_bao.clone()
        };
        let old_coeff = &self.coefficients;
        let p_coeff = if let Some(p) = perm {
            permute_array_by_atoms(old_coeff, p, &[Axis(0)], &vib_bao)
        } else {
            old_coeff.clone()
        };
        let t_p_blocks = pbao
            .shell_boundary_indices()
            .into_iter()
            .zip(tmats.iter())
            .map(|((shl_start, shl_end), tmat)| tmat.dot(&p_coeff.slice(s![shl_start..shl_end])))
            .collect::<Vec<_>>();
        let new_coefficients = concatenate(
            Axis(0),
            &t_p_blocks
                .iter()
                .map(|t_p_block| t_p_block.view())
                .collect::<Vec<_>>(),
        )
        .expect("Unable to concatenate the transformed rows for the various atoms.");
        self.coefficients = new_coefficients;
        Ok(self)
    }
}

// ------------------------
// SpinUnitaryTransformable
// ------------------------

impl<'a, T> SpinUnitaryTransformable for VibrationalCoordinate<'a, T>
where
    T: ComplexFloat + Lapack,
{
    /// Performs a spin transformation in-place.
    ///
    /// This has no effects on the vibrational coordinate as vibrational coordinates are entirely
    /// spatial.
    fn transform_spin_mut(
        &mut self,
        _dmat: &Array2<Complex<f64>>,
    ) -> Result<&mut Self, TransformationError> {
        Ok(self)
    }
}

// -------------------------------
// ComplexConjugationTransformable
// -------------------------------

impl<'a, T> ComplexConjugationTransformable for VibrationalCoordinate<'a, T>
where
    T: ComplexFloat + Lapack,
{
    /// Performs a complex conjugation in-place.
    fn transform_cc_mut(&mut self) -> Result<&mut Self, TransformationError> {
        self.coefficients.mapv_inplace(|x| x.conj());
        Ok(self)
    }
}

// -------------------------
// TimeReversalTransformable
// -------------------------
impl<'a, T> TimeReversalTransformable for VibrationalCoordinate<'a, T>
where
    T: ComplexFloat + Lapack,
{
    /// Provides a custom implementation of time reversal where the components of the vibrational
    /// coordinates are complex-conjugated to respect the antiunitarity of time reversal.
    fn transform_timerev_mut(&mut self) -> Result<&mut Self, TransformationError> {
        self.transform_cc_mut()
    }
}

// ---------------------
// SymmetryTransformable
// ---------------------
impl<'a, T> SymmetryTransformable for VibrationalCoordinate<'a, T>
where
    T: ComplexFloat + Lapack,
    VibrationalCoordinate<'a, T>: SpatialUnitaryTransformable + TimeReversalTransformable,
{
    fn sym_permute_sites_spatial(
        &self,
        symop: &SymmetryOperation,
    ) -> Result<Permutation<usize>, TransformationError> {
        symop
            .act_permute(&self.mol.molecule_ordinary_atoms())
            .ok_or(TransformationError(format!(
            "Unable to determine the atom permutation corresponding to the operation `{symop}`."
        )))
    }
}

// ---------
// Functions
// ---------

fn construct_vibration_bao(mol: &Molecule) -> BasisAngularOrder {
    let bsp_c = BasisShell::new(1, ShellOrder::Cart(CartOrder::lex(1)));
    let batms = mol
        .atoms
        .iter()
        .map(|atom| BasisAtom::new(atom, &[bsp_c.clone()]))
        .collect::<Vec<_>>();
    BasisAngularOrder::new(&batms)
}