From daa1d1de39f0cf36c52e06e9abb7238996507efa Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Fri, 3 May 2024 12:06:19 +0100 Subject: [PATCH] doc string module summary --- mess/autograd_integrals.py | 2 ++ mess/basis.py | 2 ++ mess/hamiltonian.py | 2 ++ mess/integrals.py | 41 +++++++++++++++++++------------------- mess/interop.py | 2 ++ mess/mesh.py | 2 ++ mess/numerics.py | 2 ++ mess/orbital.py | 1 + mess/plot.py | 2 ++ mess/primitive.py | 2 ++ mess/scf.py | 2 ++ mess/special.py | 2 ++ mess/structure.py | 2 ++ mess/types.py | 6 ++++++ mess/units.py | 7 +++++++ mess/xcfunctional.py | 2 ++ mess/zeropad_integrals.py | 2 ++ 17 files changed, 61 insertions(+), 20 deletions(-) diff --git a/mess/autograd_integrals.py b/mess/autograd_integrals.py index 68aa145..3c7a721 100644 --- a/mess/autograd_integrals.py +++ b/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 diff --git a/mess/basis.py b/mess/basis.py index 4e9a80c..55b53c1 100644 --- a/mess/basis.py +++ b/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 diff --git a/mess/hamiltonian.py b/mess/hamiltonian.py index 7be73f5..0a6aa48 100644 --- a/mess/hamiltonian.py +++ b/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 diff --git a/mess/integrals.py b/mess/integrals.py index c60522c..0438c9b 100644 --- a/mess/integrals.py +++ b/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. + + +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. + + +[2] PyQuante: +""" + 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. - - -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. - - -[2] PyQuante: -""" - BinaryPrimitiveOp = Callable[[Primitive, Primitive], float] diff --git a/mess/interop.py b/mess/interop.py index 25da618..0f37f97 100644 --- a/mess/interop.py +++ b/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 diff --git a/mess/mesh.py b/mess/mesh.py index 2b92ca1..9b8717a 100644 --- a/mess/mesh.py +++ b/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 diff --git a/mess/numerics.py b/mess/numerics.py index 5f8bcc8..dc4ee5c 100644 --- a/mess/numerics.py +++ b/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 diff --git a/mess/orbital.py b/mess/orbital.py index a56e3ae..dbcbb16 100644 --- a/mess/orbital.py +++ b/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 diff --git a/mess/plot.py b/mess/plot.py index 97e573b..ae34d82 100644 --- a/mess/plot.py +++ b/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 diff --git a/mess/primitive.py b/mess/primitive.py index 252a1cd..9c48e79 100644 --- a/mess/primitive.py +++ b/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 diff --git a/mess/scf.py b/mess/scf.py index aeaa583..50368d2 100644 --- a/mess/scf.py +++ b/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 diff --git a/mess/special.py b/mess/special.py index 17ff15a..e6511db 100644 --- a/mess/special.py +++ b/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 diff --git a/mess/structure.py b/mess/structure.py index 5fd6a30..0d1b0db 100644 --- a/mess/structure.py +++ b/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 diff --git a/mess/types.py b/mess/types.py index 4f45e7a..9bda79c 100644 --- a/mess/types.py +++ b/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 diff --git a/mess/units.py b/mess/units.py index c9c73d1..3cf8395 100644 --- a/mess/units.py +++ b/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 diff --git a/mess/xcfunctional.py b/mess/xcfunctional.py index 552ff22..4f86b99 100644 --- a/mess/xcfunctional.py +++ b/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 diff --git a/mess/zeropad_integrals.py b/mess/zeropad_integrals.py index 03734db..dd7a71e 100644 --- a/mess/zeropad_integrals.py +++ b/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