1use 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#[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 structure_constraint: SC,
40
41 component_index: usize,
44
45 #[builder(default = "false")]
48 complex_conjugated: bool,
49
50 bao: &'a BasisAngularOrder<'a>,
53
54 complex_symmetric: bool,
57
58 mol: &'a Molecule,
60
61 coefficients: Array1<T>,
63
64 #[builder(default = "None")]
66 energy: Option<T>,
67
68 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 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 pub fn coefficients(&self) -> &Array1<T> {
141 &self.coefficients
142 }
143
144 pub fn structure_constraint(&self) -> &SC {
146 &self.structure_constraint
147 }
148
149 pub fn bao(&self) -> &BasisAngularOrder {
152 self.bao
153 }
154
155 pub fn mol(&self) -> &Molecule {
157 self.mol
158 }
159
160 pub fn complex_symmetric(&self) -> bool {
162 self.complex_symmetric
163 }
164
165 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 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 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 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 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
423impl<'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
450impl<'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
477impl<'a, T, SC> Eq for MolecularOrbital<'a, T, SC>
481where
482 T: ComplexFloat<Real = f64> + Lapack,
483 SC: StructureConstraint + Eq,
484{
485}
486
487impl<'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
507impl<'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}