doc string module summary

mess/autograd_integrals.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""automatic differentiation of atomic orbital integrals"""
+
 from functools import partial
 from typing import Callable
 

mess/basis.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""basis sets of Gaussian type orbitals"""
+
 from typing import Tuple
 
 import equinox as eqx

mess/hamiltonian.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Many electron Hamiltonian with Density Functional Theory or Hartree-Fock."""
+
 from typing import Literal, Optional, Tuple, get_args
 
 import equinox as eqx

mess/integrals.py

@@ -1,4 +1,25 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""
+JAX implementation for integrals over Gaussian basis functions.
+
+Based upon the closed-form expressions derived in
+
+    Taketa, H., Huzinaga, S., & O-ohata, K. (1966). Gaussian-expansion methods for
+    molecular integrals. Journal of the physical society of Japan, 21(11), 2313-2324.
+    <https://doi.org/10.1143/JPSJ.21.2313>
+
+Hereafter referred to as the "THO paper"
+
+Related work:
+
+[1] Augspurger JD, Dykstra CE. General quantum mechanical operators. An
+    open-ended approach for one-electron integrals with Gaussian bases. Journal of
+    computational chemistry. 1990 Jan;11(1):105-11.
+    <https://doi.org/10.1002/jcc.540110113>
+
+[2] PyQuante: <https://github.com/rpmuller/pyquante2/>
+"""
+
 from dataclasses import asdict
 from functools import partial
 from itertools import product as cartesian_product
@@ -23,26 +44,6 @@
 from mess.types import Float3, FloatNxN
 from mess.units import LMAX
 
-"""
-JAX implementation for integrals over Gaussian basis functions. Based upon the
-closed-form expressions derived in
-
-    Taketa, H., Huzinaga, S., & O-ohata, K. (1966). Gaussian-expansion methods for
-    molecular integrals. Journal of the physical society of Japan, 21(11), 2313-2324.
-    <https://doi.org/10.1143/JPSJ.21.2313>
-
-Hereafter referred to as the "THO paper"
-
-Related work:
-
-[1] Augspurger JD, Dykstra CE. General quantum mechanical operators. An
-    open-ended approach for one-electron integrals with Gaussian bases. Journal of
-    computational chemistry. 1990 Jan;11(1):105-11.
-    <https://doi.org/10.1002/jcc.540110113>
-
-[2] PyQuante: <https://github.com/rpmuller/pyquante2/>
-"""
-
 BinaryPrimitiveOp = Callable[[Primitive, Primitive], float]
 
 

mess/interop.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Interoperation tools for working across MESS, PySCF, and PyQuante."""
+
 from typing import Tuple
 
 import numpy as np

mess/mesh.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Discretised sampling of orbitals and charge density."""
+
 from typing import Optional, Tuple, Union
 
 import equinox as eqx

mess/numerics.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Function decorators to automate converting between numeric formats."""
+
 from functools import wraps
 from typing import Callable
 

mess/orbital.py

@@ -1,4 +1,5 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Container for a linear combination of Gaussian Primitives (aka contraction)."""
 
 from functools import partial
 from typing import Tuple

mess/plot.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Visualisations of molecular structures and volumetric data."""
+
 import numpy as np
 import py3Dmol
 from more_itertools import chunked

mess/primitive.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Primitive Gaussian type orbitals"""
+
 from typing import Optional
 
 import equinox as eqx

mess/scf.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Vanilla self-consistent field solver implementation."""
+
 import jax.numpy as jnp
 import jax.numpy.linalg as jnl
 from jax.lax import while_loop

mess/special.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Special mathematical functions not readily available in JAX."""
+
 from functools import partial
 from itertools import combinations_with_replacement
 

mess/structure.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Container for molecular structures"""
+
 from typing import List
 
 import equinox as eqx

mess/types.py

@@ -1,4 +1,10 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Types used throughout MESS
+
+Note:
+    ``N`` represents the number of atomic orbitals.
+"""
+
 from functools import partial
 from typing import Tuple, Callable
 

mess/units.py

@@ -1,4 +1,11 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Conversion between Bohr and Angstrom units
+
+Note:
+    MESS uses atomic units internally so these conversions are only necessary when
+    working with external packages.
+"""
+
 from jaxtyping import Array
 
 # Maximum value an individual component of the angular momentum lmn can take

mess/xcfunctional.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""Core functions for common approximations to the exchange-correlation functional"""
+
 import jax
 import numpy as np
 import jax.numpy as jnp

mess/zeropad_integrals.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+"""(experimental) Gaussian orbital integrals without array padding."""
+
 from functools import partial
 
 import jax.numpy as jnp