qsym2/basis/
ao_integrals.rs

1//! Items to assist with integral evaluations on atomic-orbital basis functions.
2
3use std::collections::HashMap;
4use std::ops::Index;
5
6use anyhow::{self, format_err};
7use derive_builder::Builder;
8use nalgebra::{Point3, Vector3};
9use rayon::prelude::*;
10use reqwest;
11use serde::{Deserialize, Serialize};
12
13use crate::auxiliary::atom::{ElementMap, ANGSTROM_TO_BOHR};
14use crate::auxiliary::molecule::Molecule;
15use crate::basis::ao::{BasisShell, CartOrder, PureOrder, ShellOrder};
16
17#[cfg(test)]
18#[path = "ao_integrals_tests.rs"]
19mod ao_integrals_tests;
20
21// -------------------
22// GaussianContraction
23// -------------------
24
25/// Structure to handle primitives in a Gaussian contraction.
26#[derive(Clone, Builder, Debug)]
27pub struct GaussianContraction<E, C> {
28    /// Constituent primitives in the contraction. Each primitive has the form
29    /// $`c\exp\left[-\alpha\lvert \mathbf{r} - \mathbf{R} \rvert^2\right]`$ is characterised by a
30    /// tuple of its exponent $`\alpha`$ and coefficient $`c`$, respectively.
31    pub(crate) primitives: Vec<(E, C)>,
32}
33
34impl<E, C> GaussianContraction<E, C> {
35    /// The number of primitive Gaussians in this contraction.
36    pub fn contraction_length(&self) -> usize {
37        self.primitives.len()
38    }
39
40    /// Constituent primitives in the contraction. Each primitive has the form
41    /// $`c\exp\left[-\alpha\lvert \mathbf{r} - \mathbf{R} \rvert^2\right]`$ is characterised by a
42    /// tuple of its exponent $`\alpha`$ and coefficient $`c`$, respectively.
43    pub fn primitives(&self) -> &Vec<(E, C)> {
44        &self.primitives
45    }
46}
47
48impl<E, C> GaussianContraction<E, C>
49where
50    E: Clone,
51    C: Clone,
52{
53    /// Returns a builder for [`GaussianContraction`].
54    pub fn builder() -> GaussianContractionBuilder<E, C> {
55        GaussianContractionBuilder::<E, C>::default()
56    }
57}
58
59// ---------------------
60// BasisShellContraction
61// ---------------------
62
63// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
64// Deserialisable structs for BSE data retrieval
65// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
66
67/// Base API address of Basis Set Exchange.
68const BSE_BASE_API: &str = "https://www.basissetexchange.org/api";
69
70/// Threshold for filtering out primitives in a basis function contraction.
71const CONTRACTION_COEFF_THRESH: f64 = 1e-16;
72
73/// Structure to represent the REST API result fro, BasisSetExchange.
74#[derive(Serialize, Deserialize, Debug)]
75struct BSEResponse {
76    /// Name of the basis set.
77    name: String,
78
79    /// Version of the basis set.
80    version: String,
81
82    /// A hashmap between atomic numbers and element basis information.
83    elements: HashMap<u32, BSEElement>,
84}
85
86/// Structure to handle basis set information for an element.
87#[derive(Serialize, Deserialize, Debug)]
88struct BSEElement {
89    /// A vector of basis set information for the shells in this element.
90    electron_shells: Vec<BSEElectronShell>,
91}
92
93/// Structure to handle basis set information for a shell.
94#[derive(Serialize, Deserialize, Debug)]
95#[serde(try_from = "BSEElectronShellRaw")]
96struct BSEElectronShell {
97    /// The type of basis functions in this shell.
98    function_type: String,
99
100    /// The chemical region described by this shell.
101    region: String,
102
103    /// the angular momentum of this shell.
104    angular_momentum: Vec<u32>,
105
106    /// A vector of primitive exponents.
107    exponents: Vec<f64>,
108
109    /// A vector of vectors of primitive coefficients. Each inner vector is to be interpreted as a
110    /// separate shell with the same primitive exponents and angular momentum, but different
111    /// contraction coefficients.
112    coefficients: Vec<Vec<f64>>,
113}
114
115/// Structure to handle basis set information for a shell, as obtained raw from BasisSetExchange.
116#[derive(Deserialize)]
117struct BSEElectronShellRaw {
118    /// The type of basis functions in this shell.
119    function_type: String,
120
121    /// The chemical region described by this shell.
122    region: String,
123
124    /// the angular momentum of this shell.
125    angular_momentum: Vec<u32>,
126
127    /// A vector of primitive exponents.
128    exponents: Vec<String>,
129
130    /// A vector of vectors of primitive coefficients. Each inner vector is to be interpreted as a
131    /// separate shell with the same primitive exponents and angular momentum, but different
132    /// contraction coefficients.
133    coefficients: Vec<Vec<String>>,
134}
135
136impl TryFrom<BSEElectronShellRaw> for BSEElectronShell {
137    type Error = std::num::ParseFloatError;
138
139    fn try_from(other: BSEElectronShellRaw) -> Result<Self, Self::Error> {
140        let converted = Self {
141            function_type: other.function_type,
142            region: other.region,
143            angular_momentum: other.angular_momentum,
144            exponents: other
145                .exponents
146                .iter()
147                .map(|s| s.parse::<f64>())
148                .collect::<Result<Vec<_>, _>>()?,
149            coefficients: other
150                .coefficients
151                .iter()
152                .map(|d| {
153                    d.iter()
154                        .map(|s| s.parse::<f64>())
155                        .collect::<Result<Vec<_>, _>>()
156                })
157                .collect::<Result<Vec<_>, _>>()?,
158        };
159        Ok(converted)
160    }
161}
162
163// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
164// BasisShellContraction definition
165// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
166
167/// Structure to handle all shell information for integrals.
168#[derive(Clone, Builder, Debug)]
169pub struct BasisShellContraction<E, C> {
170    /// Basis function ordering information.
171    pub(crate) basis_shell: BasisShell,
172
173    /// The Gaussian primitives in the contraction of this shell.
174    pub(crate) contraction: GaussianContraction<E, C>,
175
176    /// The Cartesian origin $`\mathbf{R}`$ of this shell.
177    pub(crate) cart_origin: Point3<f64>,
178
179    /// The optional plane-wave $`\mathbf{k}`$ vector in the exponent
180    /// $`\exp\left[i\mathbf{k}\cdot(\mathbf{r} - \mathbf{R})\right]`$ associated with this shell.
181    /// If this is `None`, then this exponent is set to unity.
182    #[builder(default = "None")]
183    pub(crate) k: Option<Vector3<f64>>,
184}
185
186impl<E, C> BasisShellContraction<E, C> {
187    /// The basis function ordering information of this shell.
188    pub fn basis_shell(&self) -> &BasisShell {
189        &self.basis_shell
190    }
191
192    /// The Gaussian primitives in the contraction of this shell.
193    pub fn contraction(&self) -> &GaussianContraction<E, C> {
194        &self.contraction
195    }
196
197    /// The plane-wave $`\mathbf{k}`$ vector in the exponent.
198    pub fn k(&self) -> Option<&Vector3<f64>> {
199        self.k.as_ref()
200    }
201
202    /// The Cartesian origin $`\mathbf{R}`$ of this shell.
203    pub fn cart_origin(&self) -> &Point3<f64> {
204        &self.cart_origin
205    }
206
207    /// The number of primitive Gaussians in this shell.
208    pub fn contraction_length(&self) -> usize {
209        self.contraction.contraction_length()
210    }
211
212    /// Applies a uniform magnetic field to the shell and sets its plane-wave $`k`$ vector
213    /// according to
214    ///
215    /// ```math
216    ///     \mathbf{k} = \frac{1}{2} \mathbf{B} \times (\mathbf{R} - \mathbf{G}),
217    /// ```
218    ///
219    /// where $`\mathbf{B}`$ is the uniform magnetic field vector, $`\mathbf{R}`$ is the Cartesian
220    /// origin of this shell, and $`\mathbf{G}`$ the gauge origin with respect to which the
221    /// magnetic field is defined. Both $`\mathbf{R}`$ and $`\mathbf{G}`$ are points in a
222    /// space-fixed coordinate system.
223    ///
224    /// # Arguments
225    ///
226    /// * `b` - The magnetic field vector $`\mathbf{B}`$.
227    /// * `g` - The gauge origin $`\mathbf{G}`$.
228    pub fn apply_magnetic_field(&mut self, b: &Vector3<f64>, g: &Point3<f64>) -> &mut Self {
229        let k = 0.5 * b.cross(&(self.cart_origin.coords - g.coords));
230        self.k = Some(k);
231        self
232    }
233}
234
235impl<E, C> BasisShellContraction<E, C>
236where
237    E: Clone,
238    C: Clone,
239{
240    /// Returns a builder for [`BasisShellContraction`].
241    pub fn builder() -> BasisShellContractionBuilder<E, C> {
242        BasisShellContractionBuilder::<E, C>::default()
243    }
244}
245
246impl BasisShellContraction<f64, f64> {
247    /// Computes the self-overlap ($`\mathcal{l}_2`$-norm) of this shell and divides in-place the
248    /// contraction coefficients by ther square root of this, so that the functions in the shell
249    /// are always normalised.
250    pub fn renormalise(&mut self) -> &mut Self {
251        let c_self = self.clone();
252        let st = crate::integrals::shell_tuple::build_shell_tuple![
253            (&c_self, true), (&c_self, false); f64
254        ];
255        let ovs = st.overlap([0, 0]);
256        let norm = ovs[0].iter().next().unwrap();
257        let scale = 1.0 / norm.sqrt();
258        self.contraction.primitives.iter_mut().for_each(|(_, d)| {
259            *d *= scale;
260        });
261        self
262    }
263
264    /// Retrieves basis information from BasisSetExchange and constructs a vector of vectors of
265    /// [`Self`] for a specified molecule. Each inner vector is for one atom in the molecule.
266    ///
267    /// This method produces basis name and function ordering that are uniform across all atoms and
268    /// shells. The result from this method can be mutated for finer control of this.
269    ///
270    /// # Arguments
271    ///
272    /// * `mol` - A molecule.
273    /// * `basis_name` - The name of the basis set to be retrieved.
274    /// * `cart` - A boolean indicating if the shell functions should have lexicographic Cartesian
275    /// ordering. If `false`, the shell functions shall have increasing-$`m`$ pure ordering
276    /// instead.
277    /// * `optimised_contraction` - A boolean indicating if the optimised contraction version of
278    /// shells should be requested.
279    /// * `version` - The requested version of the basis set information.
280    /// * `mol_bohr` - A boolean indicating of the coordinates of the atoms in `mol` are to be
281    /// interpreted in units of Bohr. If `false`, they are assumed to be in units of Ångström and
282    /// will be converted to Bohr.
283    /// * `force_renormalisation` - A boolean indicating if each shell is renormalised by scaling
284    /// its primitive contraction coefficients by the inverse square root of its
285    /// $\mathcal{l}_2$-norm.
286    ///
287    /// # Returns
288    ///
289    /// A vector of vectors of [`Self`].
290    pub fn from_bse(
291        mol: &Molecule,
292        basis_name: &str,
293        cart: bool,
294        optimised_contraction: bool,
295        version: usize,
296        mol_bohr: bool,
297        force_renormalisation: bool,
298    ) -> Result<Vec<Vec<Self>>, anyhow::Error> {
299        let emap = ElementMap::new();
300        let bscs = mol
301            .atoms
302            .par_iter()
303            .map(|atom| {
304                let element = &atom.atomic_symbol;
305                let api_url = format!(
306                    "{BSE_BASE_API}/basis/\
307                {basis_name}/format/json/\
308                ?elements={element}\
309                &optimize_general={optimised_contraction}\
310                &version={version}"
311                );
312                let rjson: BSEResponse = reqwest::blocking::get(&api_url)?.json()?;
313                let atomic_number = emap
314                    .get(element)
315                    .ok_or(format_err!("Element {element} not found."))?
316                    .0;
317                rjson
318                    .elements
319                    .get(&atomic_number)
320                    .ok_or(format_err!(
321                        "Basis information for element {element} not found."
322                    ))
323                    .map(|element| {
324                        element
325                            .electron_shells
326                            .iter()
327                            .flat_map(|shell| {
328                                shell
329                                    .angular_momentum
330                                    .iter()
331                                    .cycle()
332                                    .zip(shell.coefficients.iter())
333                                    .map(|(&l, d)| {
334                                        let shell_order = if cart {
335                                            ShellOrder::Cart(CartOrder::lex(l))
336                                        } else {
337                                            ShellOrder::Pure(PureOrder::increasingm(l))
338                                        };
339                                        let basis_shell = BasisShell::new(l, shell_order);
340
341                                        let contraction = GaussianContraction::<f64, f64> {
342                                            primitives: shell
343                                                .exponents
344                                                .iter()
345                                                .copied()
346                                                .zip(d.iter().copied())
347                                                .filter(|(_, d)| d.abs() > CONTRACTION_COEFF_THRESH)
348                                                .collect::<Vec<(f64, f64)>>(),
349                                        };
350
351                                        let cart_origin = if mol_bohr {
352                                            atom.coordinates.clone()
353                                        } else {
354                                            atom.coordinates.clone() * ANGSTROM_TO_BOHR
355                                        };
356
357                                        if force_renormalisation {
358                                            let mut bsc = BasisShellContraction {
359                                                basis_shell,
360                                                contraction,
361                                                cart_origin,
362                                                k: None,
363                                            };
364                                            bsc.renormalise();
365                                            bsc
366                                        } else {
367                                            BasisShellContraction {
368                                                basis_shell,
369                                                contraction,
370                                                cart_origin,
371                                                k: None,
372                                            }
373                                        }
374                                    })
375                            })
376                            .collect::<Vec<_>>()
377                    })
378            })
379            .collect::<Result<Vec<_>, _>>()?;
380        Ok(bscs)
381    }
382}
383
384// --------
385// BasisSet
386// --------
387/// Structure to manage basis information for a molecule.
388#[derive(Clone, Debug)]
389pub struct BasisSet<E, C> {
390    /// A vector of vectors containing basis information for the atoms in this molecule. Each inner
391    /// vector is for one atom.
392    basis_atoms: Vec<Vec<BasisShellContraction<E, C>>>,
393
394    /// The function boundaries for the atoms in the molecule.
395    atom_boundaries: Vec<(usize, usize)>,
396
397    /// The function boundaries for the shells in the molecule.
398    shell_boundaries: Vec<(usize, usize)>,
399}
400
401impl<E, C> BasisSet<E, C> {
402    /// Returns a reference to the vector of vectors containing basis information for the atoms in
403    /// this molecule. Each inner vector is for one atom.
404    pub fn basis_atoms(&self) -> &Vec<Vec<BasisShellContraction<E, C>>> {
405        &self.basis_atoms
406    }
407
408    /// Creates a new [`BasisSet`] structure from a vector of vectors of basis shells.
409    ///
410    /// # Arguments
411    ///
412    /// * `batms` - A vector of vectors of basis shells. Each inner vector is for one atom.
413    ///
414    /// # Returns
415    ///
416    /// A new [`BasisSet`] structure.
417    pub fn new(batms: Vec<Vec<BasisShellContraction<E, C>>>) -> Self {
418        let atom_boundaries = batms
419            .iter()
420            .scan(0, |acc, batm| {
421                let atom_length = batm
422                    .iter()
423                    .map(|bs| bs.basis_shell.n_funcs())
424                    .sum::<usize>();
425                let boundary = (*acc, *acc + atom_length);
426                *acc += atom_length;
427                Some(boundary)
428            })
429            .collect::<Vec<_>>();
430        let shell_boundaries = batms
431            .iter()
432            .flatten()
433            .scan(0, |acc, bsc| {
434                let shell_length = bsc.basis_shell.n_funcs();
435                let boundary = (*acc, *acc + shell_length);
436                *acc += shell_length;
437                Some(boundary)
438            })
439            .collect::<Vec<_>>();
440        Self {
441            basis_atoms: batms,
442            atom_boundaries,
443            shell_boundaries,
444        }
445    }
446
447    /// Updates the cached shell boundaries. This is required when the shells or atoms have been
448    /// reordered.
449    fn update_shell_boundaries(&mut self) -> &mut Self {
450        self.shell_boundaries = self
451            .basis_atoms
452            .iter()
453            .flatten()
454            .scan(0, |acc, bsc| {
455                let shell_length = bsc.basis_shell.n_funcs();
456                let boundary = (*acc, *acc + shell_length);
457                *acc += shell_length;
458                Some(boundary)
459            })
460            .collect::<Vec<_>>();
461        self
462    }
463
464    /// Applies a uniform magnetic field to all shells and sets their plane-wave $`k`$ vectors.
465    /// See the documentation of [`BasisShellContraction::apply_magnetic_field`] for more
466    /// information.
467    ///
468    /// # Arguments
469    ///
470    /// * `b` - The magnetic field vector $`\mathbf{B}`$.
471    /// * `g` - The gauge origin $`\mathbf{G}`$.
472    pub fn apply_magnetic_field(&mut self, b: &Vector3<f64>, g: &Point3<f64>) -> &mut Self {
473        self.all_shells_mut().for_each(|shell| {
474            shell.apply_magnetic_field(b, g);
475        });
476        self
477    }
478
479    /// The number of shells in the basis set.
480    pub fn n_shells(&self) -> usize {
481        self.basis_atoms
482            .iter()
483            .map(|batm| batm.len())
484            .sum::<usize>()
485    }
486
487    /// The number of basis functions in the basis set.
488    pub fn n_funcs(&self) -> usize {
489        self.all_shells()
490            .map(|shell| shell.basis_shell.n_funcs())
491            .sum::<usize>()
492    }
493
494    /// Sorts the shells in each atom by their angular momenta.
495    pub fn sort_by_angular_momentum(&mut self) -> &mut Self {
496        self.basis_atoms
497            .iter_mut()
498            .for_each(|batm| batm.sort_by_key(|bsc| bsc.basis_shell.l));
499        self.update_shell_boundaries()
500    }
501
502    /// Returns the function shell boundaries.
503    pub fn shell_boundaries(&self) -> &Vec<(usize, usize)> {
504        &self.shell_boundaries
505    }
506
507    /// Returns an iterator over all shells in the basis set.
508    pub fn all_shells(&self) -> impl Iterator<Item = &BasisShellContraction<E, C>> {
509        self.basis_atoms.iter().flatten()
510    }
511
512    /// Returns a mutable iterator over all shells in the basis set.
513    pub fn all_shells_mut(&mut self) -> impl Iterator<Item = &mut BasisShellContraction<E, C>> {
514        self.basis_atoms.iter_mut().flatten()
515    }
516}
517
518impl BasisSet<f64, f64> {
519    /// Retrieves basis information from BasisSetExchange and constructs [`Self`] for a specified
520    /// molecule.
521    ///
522    /// This method produces basis name and function ordering that are uniform across all atoms and
523    /// shells. The result from this method can be mutated for finer control of this.
524    ///
525    /// # Arguments
526    ///
527    /// * `mol` - A molecule.
528    /// * `basis_name` - The name of the basis set to be retrieved.
529    /// * `cart` - A boolean indicating if the shell functions should have lexicographic Cartesian
530    /// ordering. If `false`, the shell functions shall have increasing-$`m`$ pure ordering
531    /// instead.
532    /// * `optimised_contraction` - A boolean indicating if the optimised contraction version of
533    /// shells should be requested.
534    /// * `version` - The requested version of the basis set information.
535    /// * `mol_bohr` - A boolean indicating of the coordinates of the atoms in `mol` are to be
536    /// interpreted in units of Bohr. If `false`, they are assumed to be in units of Ångström and
537    /// will be converted to Bohr.
538    /// * `force_renormalisation` - A boolean indicating if each shell is renormalised by scaling
539    /// its primitive contraction coefficients by the inverse square root of its
540    /// $\mathcal{l}_2$-norm.
541    ///
542    /// # Returns
543    ///
544    /// A [`BasisSet`] structure.
545    pub fn from_bse(
546        mol: &Molecule,
547        basis_name: &str,
548        cart: bool,
549        optimised_contraction: bool,
550        version: usize,
551        mol_bohr: bool,
552        force_renormalisation: bool,
553    ) -> Result<Self, anyhow::Error> {
554        Ok(Self::new(BasisShellContraction::<f64, f64>::from_bse(
555            mol,
556            basis_name,
557            cart,
558            optimised_contraction,
559            version,
560            mol_bohr,
561            force_renormalisation,
562        )?))
563    }
564}
565
566impl<E, C> Index<usize> for BasisSet<E, C> {
567    type Output = BasisShellContraction<E, C>;
568
569    fn index(&self, i: usize) -> &Self::Output {
570        self.basis_atoms
571            .iter()
572            .flatten()
573            .nth(i)
574            .unwrap_or_else(|| panic!("Unable to obtain the basis shell with index {i}."))
575    }
576}