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}