### Install MESS for Google Colab

Source: https://github.com/hatemhelal/mess/blob/main/docs/optim.ipynb

Installs the MESS JAX library into the current Google Colab runtime environment. This is a prerequisite for running the subsequent code examples in a Colab notebook.

```python
import sys

if "google.colab" in sys.modules:
    !pip install mess-jax
```

--------------------------------

### Install Pre-commit Hooks with uv

Source: https://github.com/hatemhelal/mess/blob/main/CONTRIBUTING.md

Installs the pre-commit hooks using the 'uv' command. These hooks will run automatically when committing changes.

```bash
uv run pre-commit install
```

--------------------------------

### Sync Development Dependencies with uv

Source: https://github.com/hatemhelal/mess/blob/main/CONTRIBUTING.md

Creates a virtual environment with all necessary development packages. Use '--extra gpu' to include JAX for GPU execution. Assumes 'uv' is installed and available.

```bash
uv sync --extra dev
```

--------------------------------

### Install pyquante2

Source: https://github.com/hatemhelal/mess/blob/main/docs/gammanu.ipynb

Installs the pyquante2 library from a GitHub repository. This library is used for its reference implementation of the Fgamma function.

```python
!pip install pyquante2@git+https://github.com/rpmuller/pyquante2@pure
```

--------------------------------

### Install MESS Python Package

Source: https://github.com/hatemhelal/mess/blob/main/README.md

This snippet shows how to install the MESS Python package directly from its GitHub repository using pip. It requires Python 3.11+ and recommends installing JAX for optimal performance.

```bash
pip install git+https://github.com/hatemhelal/mess.git
```

--------------------------------

### Verify Hartree-Fock Energy with PySCF

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This example demonstrates how to calculate the Hartree-Fock energy using PySCF and compare it with the results obtained from MESS. It initializes a Restricted Hartree-Fock (RHF) calculation for a given molecular object and runs the SCF kernel. The output is the converged SCF energy.

```python
s = scf.RHF(scf_mol)
s.kernel()
```

--------------------------------

### Generate '6-31+G' Basis Set for Oxygen (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Illustrates the creation of a '6-31+G' basis set for a single oxygen atom. This example shows how different basis sets have varying numbers of primitives per basis function and prints the number of orbitals and their primitive function definitions.

```python
basisname = "6-31+G"
oxygen = Structure(atomic_number=np.array([Z]), position=np.zeros(3))
basis = basisset(oxygen, basisname)
print(
    f"The {basisname} basis for {Z=} has "
    f"I(Z)={basis.num_orbitals} basis functions/orbitals"
)

for i, o in enumerate(basis.orbitals):
    terms = (
        f"{coef:.3f} * r**{p.lmn} * exp(-{p.alpha:.2f} r**2)"
        for coef, p in zip(o.coefficients, o.primitives)
    )
    print(f"phi_({Z},{i})=", " + ".join(terms))

# Now count unique prims
prims = set()
for o in basis.orbitals:
    for p in o.primitives:
        prims = set.union(prims, {(*p.lmn.tolist(), float(p.alpha))})
print(f"Found {len(prims)} unique primitives for {Z=} in {basisname}")
```

--------------------------------

### Interoperability with PySCF - Python

Source: https://context7.com/hatemhelal/mess/llms.txt

Shows how to convert molecular and basis set objects between MESS and PySCF formats using `to_pyscf` and `from_pyscf`. This allows for cross-validation of results and leveraging PySCF's extensive functionalities within a MESS workflow. Examples include running a PySCF calculation and then using the converted objects in MESS.

```python
from mess import molecule, basisset
from mess.interop import to_pyscf, from_pyscf
from pyscf import dft, gto

# MESS to PySCF
mol = molecule("water")
pyscf_mol = to_pyscf(mol, basis_name="6-31g", spherical=True)

# Run PySCF calculation for comparison
mf = dft.RKS(pyscf_mol, xc="lda,vwn")
energy_pyscf = mf.kernel()
print(f"PySCF energy: {energy_pyscf} Hartree")

# PySCF to MESS
pyscf_mol = gto.M(atom="O 0 0 0; H 0 1 0; H 0 0 1", basis="6-31g")
structure, basis = from_pyscf(pyscf_mol)

# Use converted objects in MESS
from mess import Hamiltonian, minimise
H = Hamiltonian(basis, xc_method="lda")
E, C, sol = minimise(H)
print(f"MESS energy: {E} Hartree")
```

--------------------------------

### Integral Calculations - Python

Source: https://context7.com/hatemhelal/mess/llms.txt

Illustrates how to compute one-electron (overlap, kinetic, nuclear attraction) and two-electron (electron repulsion) integrals over Gaussian basis functions using MESS. These integrals are fundamental for electronic structure calculations. The example also shows how to build Coulomb and exchange matrices from a density matrix.

```python
from mess import molecule, basisset
from mess.integrals import (
    overlap_basis,
    kinetic_basis,
    nuclear_basis,
    eri_basis
)
import jax.numpy as jnp

mol = molecule("h2")
basis = basisset(mol, basis_name="sto-3g")

# One-electron integrals (NxN matrices)
S = overlap_basis(basis)  # Overlap matrix
print(f"Overlap:\n{S}")

T = kinetic_basis(basis)  # Kinetic energy
print(f"Kinetic energy:\n{T}")

V = nuclear_basis(basis)  # Nuclear attraction (num_atoms x N x N)
V_total = V.sum(axis=0)  # Sum over all nuclei
print(f"Nuclear attraction:\n{V_total}")

# Two-electron integrals (NxNxNxN tensor)
ERI = eri_basis(basis)  # Electron repulsion integrals
print(f"ERI shape: {ERI.shape}")  # (2, 2, 2, 2) for H2 with sto-3g

# Build Coulomb and exchange matrices from density matrix P
P = jnp.eye(basis.num_orbitals) * 2.0  # Example density matrix
J = jnp.einsum("kl,ijkl->ij", P, ERI)  # Coulomb matrix
K = jnp.einsum("ij,ikjl->kl", P, ERI)  # Exchange matrix
```

--------------------------------

### Generating a Table of Boys' Integrals

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

Creates a function `boys_table` that computes a table of Boys' integrals up to a maximum order `nmax` using a recursive approach starting from a given `x` value. It then verifies the table against direct calculations.

```Python
def boys_table(x, nmax=16):
    out = np.zeros(nmax)
    fn = boys_integral(nmax, x)

    for n in range(nmax, 0, -1):
        out[n - 1] = (2 * x * fn + np.exp(-x)) / (2 * n - 1)
        fn = out[n - 1]

    return out


np.allclose(boys_table(10), np.asarray([boys_integral(n, 10) for n in range(16)]))
```

--------------------------------

### Self-Consistent Field (SCF) Solver - Python

Source: https://context7.com/hatemhelal/mess/llms.txt

Presents the `scf` function, a standard self-consistent field solver that employs fixed-point iteration with Hartree-Fock exchange. It serves as an alternative to gradient-based minimization for solving electronic structure problems. The example shows how to run an SCF calculation with specified orthonormalization and convergence criteria.

```python
from mess import molecule, basisset
from mess.scf import scf
from mess.orthnorm import cholesky

mol = molecule("h2")
basis = basisset(mol, basis_name="sto-3g")

# Run SCF calculation
E_total = scf(
    basis,
    otransform=cholesky,
    max_iters=50,
    tolerance=1e-6
)

print(f"Total HF energy: {E_total} Hartree")
```

--------------------------------

### Build Documentation with uv

Source: https://github.com/hatemhelal/mess/blob/main/CONTRIBUTING.md

Builds the project documentation from the root directory using the 'uv run task doc' command. This command facilitates the documentation generation process.

```bash
uv run task doc
```

--------------------------------

### Initialize trial matrix and optimizer

Source: https://github.com/hatemhelal/mess/blob/main/docs/batching.ipynb

Sets up the initial guess for the trial matrix 'Z' and initializes an Adam optimizer from the Optax library. The optimizer state is also initialized based on the initial guess of 'Z'.

```python
Z = jnp.tile(jnp.eye(num_orbitals), (num_confs, 1, 1))
optimiser = optax.adam(learning_rate=0.1)
state = optimiser.init(Z)
```

--------------------------------

### Import necessary libraries

Source: https://github.com/hatemhelal/mess/blob/main/docs/batching.ipynb

Imports essential libraries for numerical computation, optimization, plotting, and the MESS library itself. These include JAX for automatic differentiation and vectorization, NumPy for array operations, Optax for optimizers, and Seaborn for visualizations.

```python
import jax
import jax.numpy as jnp
import numpy as np
import optax
import seaborn as sns
from tqdm.notebook import tqdm

from mess import Hamiltonian, basisset
from mess.structure import Structure, nuclear_energy

sns.set_theme(style="whitegrid")
```

--------------------------------

### Import necessary libraries for MESS simulations

Source: https://github.com/hatemhelal/mess/blob/main/docs/optim.ipynb

Imports core libraries including JAX for automatic differentiation and numerical operations, optax for optimisation algorithms, seaborn for plotting, tqdm for progress bars, and specific modules from the MESS library for constructing Hamiltonians and molecules.

```python
import jax
import jax.numpy as jnp
import optax
import seaborn as sns
from tqdm.notebook import tqdm

from mess import Hamiltonian, basisset, molecule
from mess.structure import nuclear_energy

sns.set_theme(style="whitegrid")
```

--------------------------------

### Set up simulation components: basis set, Hamiltonian, and optimiser

Source: https://github.com/hatemhelal/mess/blob/main/docs/optim.ipynb

Initialises the necessary components for the energy minimisation. This includes defining the basis set for molecular orbitals, constructing the Hamiltonian using the PBE exchange-correlation method, and setting up the Adam optimiser with a specified learning rate.

```python
basis = basisset(mol, "6-31g")
H = Hamiltonian(basis, xc_method="pbe")
optimiser = optax.adam(learning_rate=0.1)
```

--------------------------------

### Display CPU Information

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This command displays detailed information about the CPU architecture. It includes details such as the architecture type, CPU core count, cache sizes, and supported instruction sets. This information is valuable for understanding the system's processing capabilities.

```bash
!lscpu
```

--------------------------------

### Import necessary libraries

Source: https://github.com/hatemhelal/mess/blob/main/docs/gammanu.ipynb

Imports core libraries for numerical computation, symbolic mathematics, plotting, and data handling. It also imports a reference implementation of Fgamma from pyquante2.

```python
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from pyquante2.utils import Fgamma as reference_pyquante

from functools import partial
from IPython.display import display

sns.set_theme()
%matplotlib inline
```

--------------------------------

### Run All Tests with uv

Source: https://github.com/hatemhelal/mess/blob/main/CONTRIBUTING.md

Executes all project tests using the 'uv run task test' command. This command relies on the 'uv' tool for task execution.

```bash
uv run task test
```

--------------------------------

### Import Libraries for Numerical Computation and Visualization

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

Imports necessary Python libraries including NumPy for numerical operations, SymPy for symbolic mathematics, SciPy for special functions, Matplotlib for plotting, and PySCF for computational chemistry tasks. It also includes imports for animation and display functionalities.

```Python
import numpy as np
import sympy as sp
import scipy.special as ssp
import matplotlib.pyplot as plt
import functools as ft
from mess.structure import Structure
from mess.interop import to_pyscf
from pyscf import dft
from matplotlib.animation import FuncAnimation, PillowWriter
from IPython.display import HTML
import numpy.polynomial as npp
# %matplotlib widget
```

--------------------------------

### Define batched energy function with JIT and VMAP

Source: https://github.com/hatemhelal/mess/blob/main/docs/batching.ipynb

Defines an energy function that calculates the electronic energy for a given trial matrix Z and Hamiltonian H. It uses `jax.jit` for compilation and `jax.vmap` to enable the function to operate on batches of Hamiltonians and trial matrices, significantly improving computational efficiency.

```python
@jax.jit
@jax.vmap
def energy(Z, H):
    C = H.orthonormalise(Z)
    P = H.basis.density_matrix(C)
    return H(P)
```

--------------------------------

### Run Performance Benchmarks with uv

Source: https://github.com/hatemhelal/mess/blob/main/CONTRIBUTING.md

Runs all performance benchmarks using 'uv run task bench'. This process utilizes 'pytest-benchmark' for measuring performance.

```bash
uv run task bench
```

--------------------------------

### Create and Push Feature Branch with Git

Source: https://github.com/hatemhelal/mess/blob/main/CONTRIBUTING.md

Demonstrates creating a new feature branch, making changes, and pushing it to the remote repository. A link to open a Pull Request is provided after pushing.

```bash
git checkout -b feature
...
git push --set-upstream origin feature
```

--------------------------------

### Animating H2 Integration Grid Evolution

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

Generates an animation showing the evolution of the H2 integration grid as the internuclear distance changes. The animation visualizes the weights and positions of grid points on the z=0 plane, with point size and color representing weight magnitude.

```Python
def build_grid(r: float):
    pos = np.zeros((2, 3))
    pos[0, 0] = r
    pos[1, 0] = -r
    mol = Structure(atomic_number=np.asarray([1, 1]), position=pos)
    mol = to_pyscf(mol)
    grids = dft.gen_grid.Grids(mol)
    grids.build()
    weights, points = grids.weights, grids.coords

    z0 = points[:, 2] == 0
    weights = weights[z0]
    x, y = points[z0, 0], points[z0, 1]
    return x, y, weights


data = [build_grid(r) for r in np.linspace(1.0, 20.0, 30)]
num_frames = len(data)
fig, ax = plt.subplots()
h = ax.scatter([], [], [], [], vmin=-35, vmax=55)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_xlim(-40, 40)
ax.set_ylim(-25, 25)
cb = fig.colorbar(h)


def update(i):
    x, y, w = data[i]
    h.set_offsets(np.vstack((x, y)).T)
    h.set_sizes(10 * np.abs(w))
    h.set_array(w)
    fig.colorbar(h, cax=cb.ax)
    return (h,)


ani = FuncAnimation(fig, update, frames=len(data), interval=100, blit=True)
HTML(ani.to_html5_video())
```

--------------------------------

### Import Libraries for MESS Calculations

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Imports necessary libraries for numerical operations (numpy), symbolic mathematics (sympy), audio playback (IPython.display), 3D visualization (py3Dmol), and MESS-specific modules for structure, basis sets, mesh generation, molecular orbitals, and plotting. These are foundational for most subsequent operations.

```python
import numpy as np
import sympy as sp
import IPython.display as ipd
import py3Dmol

from mess.structure import Structure
from mess.basis import basisset
from mess.mesh import uniform_mesh, molecular_orbitals
from mess.plot import plot_volume, plot_isosurfaces, plot_molecule
from itertools import product
```

--------------------------------

### Minimize Hamiltonian Energy with MESS (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This snippet demonstrates how to initialize a Hamiltonian using the PBE exchange-correlation functional and then minimize its energy. It utilizes the `minimise` function from the `mess` library and returns the energy (E) and molecular orbital coefficients (C).

```python
from mess import minimise, Hamiltonian

H = Hamiltonian(basis, xc_method="pbe")
E, C, sol = minimise(H)
float(E)
```

--------------------------------

### Create a molecule object from sample data

Source: https://github.com/hatemhelal/mess/blob/main/docs/optim.ipynb

Constructs a molecule object representing methane using sample data packaged with the pyquante project. The resulting object contains atomic numbers and positions.

```python
mol = molecule("methane")
mol
```

--------------------------------

### Run gradient descent optimization

Source: https://github.com/hatemhelal/mess/blob/main/docs/batching.ipynb

Performs a gradient descent optimization to minimize the total energy across all molecular conformations in the batch. It defines a loss function that sums the energies, uses `jax.value_and_grad` to compute loss and gradients efficiently, and iteratively updates the trial matrix 'Z' using the Adam optimizer.

```python
@jax.value_and_grad
def loss_fn(z, h):
    return jnp.sum(energy(z, h))


history = []

for _ in (bar := tqdm(range(128))):
    loss, grads = loss_fn(Z, H)
    updates, state = optimiser.update(grads, state)
    Z = optax.apply_updates(Z, updates)
    history.append(loss)
    bar.set_description(f"Loss {loss:0.06f}")
```

--------------------------------

### Benchmark PySCF Energy Minimization (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This code measures the performance of PySCF's energy minimization routine. It uses the `%timeit` magic command to execute `s.kernel()`, which performs the SCF calculation, and captures the timing results for comparison with MESS performance.

```python
pyscf_time = %timeit -o s.kernel()
```

--------------------------------

### Generate 'sto-3g' Basis Set for Oxygen (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Demonstrates how to create a 'sto-3g' basis set for a single oxygen atom (Z=8) using the Structure and basisset functions from the MESS library. It prints the number of atomic orbitals and their constituent primitive functions, illustrating the structure of a minimal basis set.

```python
Z = 8
basisname = "sto-3g"
oxygen = Structure(atomic_number=np.array([Z]), position=np.zeros(3))
basis = basisset(oxygen, basisname)
print(f"The {basisname} basis for {Z=} has I(Z)={basis.num_orbitals}")
```

--------------------------------

### Identify JAX Device Kind

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This Python snippet identifies the type of accelerator device configured for JAX. It imports the JAX library and accesses the `device_kind` attribute of the first available device. This is useful for understanding the computational hardware being used.

```python
import jax

jax.devices()[0].device_kind
```

--------------------------------

### Plot Molecular Orbitals using py3Dmol (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Demonstrates plotting atomic orbitals for a molecule using the MESS library and py3Dmol. It first evaluates the basis functions on a uniform mesh, then uses py3Dmol to visualize the molecule and the isosurface of a specific orbital (e.g., 2px).

```python
# Evaluate the molecular orbitals on a mesh -> [num_mesh_points, num_orbitals]
mesh = uniform_mesh(n=32, b=3.0)
orbitals = molecular_orbitals(basis, mesh.points)
orbitals = np.asarray(orbitals)

# Mapping between orbital label -> index in basis set
# TODO: what is this mapping for 6-31+G?
# orbital_idx = dict(zip(["1s", "2s", "2px", "2py", "2pz"], range(5)))

view = py3Dmol.view()
plot_molecule(view, oxygen)
plot_volume(view, orbitals[:, 4], mesh.axes)
view.zoomTo()
```

--------------------------------

### Time Performance of F Function

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

This code uses the `%%timeit` magic command to benchmark the execution speed of the `F` function. It repeatedly calls `F(2, x)` to calculate the mean execution time and standard deviation, providing insights into the function's performance.

```python
%%timeit
F(2, x)
```

--------------------------------

### Visualize Electron Density with MESS and py3Dmol (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This code visualizes the electron density derived from a minimized Hamiltonian solution. It uses `py3Dmol` for rendering and functions from `mess.mesh` and `mess.plot` to generate and display the electron cloud, which is described as 'mickey mouse head shaped'.

```python
import py3Dmol
from mess.mesh import density, uniform_mesh
from mess.plot import plot_volume, plot_molecule

view = py3Dmol.view()
plot_molecule(view, mol)
mesh = uniform_mesh()
rho = density(basis, mesh, basis.density_matrix(C))
plot_volume(view, rho, mesh.axes)
```

--------------------------------

### Building 3D H2 Integration Grid

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

A function to build integration grid weights and coordinates for an H2 molecule. It returns the weights and points, which can then be used for 3D plotting.

```Python
def build_grid(r: float):
    pos = np.zeros((2, 3))
    pos[0, 0] = r
    pos[1, 0] = -r
    mol = Structure(atomic_number=np.asarray([1, 1]), position=pos)
    mol = to_pyscf(mol)
    grids = dft.gen_grid.Grids(mol)
    grids.build()
    return grids.weights, grids.coords


weights, points = build_grid(15.0)
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
```

--------------------------------

### Saving Animation as GIF

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

Saves the generated animation object as a GIF file named 'scatter.gif'. It configures the animation writer with a specified frames per second (fps) and bitrate.

```Python
writer = PillowWriter(fps=5, bitrate=1800)
ani.save("scatter.gif", writer=writer)
```

--------------------------------

### Perform energy minimisation using Adam optimiser

Source: https://github.com/hatemhelal/mess/blob/main/docs/optim.ipynb

Initialises a trial matrix Z and the state for the Adam optimiser. It then iterates for a specified number of steps, calculating the total energy and its gradient, applying updates from the optimiser, and storing the energy history. A progress bar visualises the optimisation process.

```python
Z = jnp.eye(basis.num_orbitals)
state = optimiser.init(Z)
history = []

for _ in (bar := tqdm(range(200))):
    e, grads = total_energy(Z)
    updates, state = optimiser.update(grads, state)
    Z = optax.apply_updates(Z, updates)
    history.append(e)
    bar.set_description(f"Total energy: {e:0.06f} (Hartree)")
```

--------------------------------

### Plot optimization loss history

Source: https://github.com/hatemhelal/mess/blob/main/docs/batching.ipynb

Visualizes the optimization process by plotting the batched loss over iterations using Seaborn. This plot serves as a sanity check to ensure the optimization is converging.

```python
history = jnp.stack(history)
ax = sns.lineplot(history)
ax.set_xlabel("Iteration")
ax.set_ylabel("Batched Loss (Hartree)");
```

--------------------------------

### Define and Display Symbolic Integral

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Initializes SymPy printing for Unicode and defines symbolic variables (eta, k, t) and a function (H). It then defines a definite integral H_na and displays it as an equation.

```python
sp.init_printing(use_unicode=True)

eta = sp.Symbol("eta", positive=True, real=True)
k = sp.Symbol("k", integer=True, nonnegative=True)
t = sp.Symbol("t", real=True)
H = sp.Function("H")

H_na = sp.Integral(t ** (2 * k) * sp.exp(-eta * t**2), (t, -sp.oo, sp.oo))
ipd.display(sp.Eq(H(2 * k, eta), H_na))
```

--------------------------------

### Convergence Analysis for gammaincln

Source: https://github.com/hatemhelal/mess/blob/main/docs/gammanu.ipynb

This code analyzes the convergence of the `gammaincln` function by varying the number of terms (`N`) and data types (float64, float32). It compares the results against a SciPy reference and calculates relative errors to determine how many terms are needed for convergence across different scenarios.

```python
A = np.arange(0, 16, 4) + 0.5
Z = np.linspace(0, 20, 5)
N = np.array([32, 64, 128])

vals = np.meshgrid(A, Z, N)
A, Z, N = [t.reshape(-1) for t in vals]

gammaincln_fp64 = partial(gammaincln, dtype=np.float64)
gammaincln_fp32 = partial(gammaincln, dtype=np.float32)
actual_fp64 = np.vectorize(gammaincln_fp64)(A, Z, N)
actual_fp32 = np.vectorize(gammaincln_fp32)(A, Z, N)
expect = reference_scipy(A, np.maximum(Z, 1e-10))

df = pd.DataFrame({
    "a": A,
    "z": Z,
    "num_terms": N,
    "actual_fp64": actual_fp64,
    "actual_fp32": actual_fp32,
    "expect": expect,
})
df["rel_error_fp64"] = np.abs((df.actual_fp64 - df.expect) / df.expect)
```

--------------------------------

### Compare MESS with PySCF for Energy Calculation (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This snippet integrates MESS with PySCF to verify simulation accuracy. It converts a MESS molecule object to a PySCF-compatible format using `to_pyscf`, performs a Restricted-Kernel (RKS) calculation with the PBE functional, and then checks if the total energies are numerically close using `np.allclose`.

```python
import numpy as np
from mess.interop import to_pyscf
from pyscf import dft, scf


scf_mol = to_pyscf(mol, basis_name)
s = dft.RKS(scf_mol, xc="pbe")
s.kernel()
np.allclose(s.energy_tot(), E)
```

--------------------------------

### Create Hydrogen molecule Hamiltonian

Source: https://github.com/hatemhelal/mess/blob/main/docs/batching.ipynb

Defines a function to create a Hamiltonian for a hydrogen molecule (H2) given a bond length 'r'. It constructs the molecular structure, defines the basis set, and initializes the Hamiltonian. The code then generates a batch of Hamiltonians for various bond lengths and stacks them for batched computation.

```python
def h2_hamiltonian(r: float, basis_name: str = "sto-3g", xc_method="lda"):
    mol = Structure(
        atomic_number=np.array([1, 1]),
        position=np.array([[0.0, 0.0, 0.0], [r, 0.0, 0.0]]),
    )
    basis = basisset(mol, basis_name=basis_name)
    return Hamiltonian(basis, xc_method=xc_method)


num_confs = 64
rs = np.linspace(0.6, 10, num_confs)
H = [h2_hamiltonian(r) for r in rs]
num_orbitals = H[0].basis.num_orbitals
H = jax.tree.map(lambda *xs: jnp.stack(xs), *H)
```

--------------------------------

### Detail 'sto-3g' Basis Functions and Primitives for Oxygen (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Expands on the 'sto-3g' basis set for oxygen by iterating through each atomic orbital and printing its definition in terms of coefficients and primitive Gaussian functions. It also counts and reports the number of unique primitive functions used across all orbitals.

```python
print(f"{basisname} for {Z=} has {len(basis.orbitals)} basis functions/orbitals")
for i, o in enumerate(basis.orbitals):
    terms = (
        f"{coef:.3f} * r**{p.lmn} * exp(-{p.alpha:.2f} r**2)"
        for coef, p in zip(o.coefficients, o.primitives)
    )
    print(f"phi_({Z},{i})=", " + ".join(terms))

# Now count unique prims
prims = set()
for o in basis.orbitals:
    for p in o.primitives:
        prims = set.union(prims, {(*p.lmn.tolist(), float(p.alpha))})
print(f"Found {len(prims)} unique primitives for {Z=} in {basisname}")
```

--------------------------------

### Plot the total energy variation during optimisation

Source: https://github.com/hatemhelal/mess/blob/main/docs/optim.ipynb

Visualises the convergence of the energy minimisation by plotting the total energy recorded at each iteration. This helps in assessing the stability and efficiency of the optimisation process.

```python
history = jnp.stack(history)
ax = sns.lineplot(history)
ax.set_xlabel("Iteration")
ax.set_ylabel("Total energy (Hartree)")
```

--------------------------------

### Fitting Boys' Integral with Chebyshev Polynomials

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

Calculates Boys' integrals over a range of x values and fits them using Chebyshev polynomials of varying degrees. It then determines the optimal degree by minimizing the maximum absolute error between the fitted and actual values.

```Python
num_samples = 1000
x = np.logspace(-2, 3, num_samples)
# x = np.linspace(0, 200, num_samples)
n = 1
y = boys_integral(n, x)
# plt.plot(x, y)

e = []
for d in range(0, 64):
    f = npp.Chebyshev.fit(x, y, d)
    xx, yy = f.linspace(num_samples)
    e += [np.max(np.abs(boys_integral(n, xx) - yy))]

e = np.asarray(e)
plt.plot(e)
plt.yscale("log")
plt.scatter(np.argmin(e), np.min(e), color="r")
np.argmin(e)
```

--------------------------------

### Custom Initial Guess Function - Python

Source: https://context7.com/hatemhelal/mess/llms.txt

Defines a custom function to generate an initial guess for the electronic structure calculation. The function takes a basis set object and returns an NxN matrix, which will be orthonormalized by the solver. A common implementation is to return an identity matrix scaled by a small factor.

```python
import jax.numpy as jnp

def custom_guess(basis):
    # Return any NxN matrix - will be orthonormalized
    return jnp.eye(basis.num_orbitals) * 1.1

E, C, sol = minimise(H, initial_guess_fn=custom_guess)
```

--------------------------------

### Generate and Analyze Cubic Hydrogen Lattice with MESS

Source: https://context7.com/hatemhelal/mess/llms.txt

Generates cubic lattice structures of hydrogen atoms using the 'cubic_hydrogen' function from the MESS library. It demonstrates creating a lattice, printing its properties, and performing a DFT calculation (LDA exchange-correlation) on the generated structure to obtain its energy.

```python
from mess.structure import cubic_hydrogen, nuclear_energy
from mess import basisset, Hamiltonian, minimise

# Create 2x2x2 hydrogen lattice (8 atoms)
h_lattice = cubic_hydrogen(n=2)
print(f"Number of atoms: {h_lattice.num_atoms}")  # 8
print(f"Positions:\n{h_lattice.position}")

# Run calculation on lattice
basis = basisset(h_lattice, basis_name="sto-3g")
H = Hamiltonian(basis, xc_method="lda")
E, C, sol = minimise(H)
print(f"Lattice energy: {E} Hartree")

# Single atom vs small clusters
h1 = cubic_hydrogen(1)  # 1 atom
h3 = cubic_hydrogen(3)  # 27 atoms
```

--------------------------------

### Optimize Hartree-Fock Energy with MESS

Source: https://context7.com/hatemhelal/mess/llms.txt

Calculates the optimized ground state energy for a given molecular system using the Hartree-Fock method with the MESS library. It initializes a Hamiltonian with a pure Hartree-Fock exchange-correlation method and then uses the 'minimise' function to find the optimized energy and wave function.

```python
from mess import Hamiltonian, minimise

H = Hamiltonian(basis, xc_method="hfx")  # Pure Hartree-Fock
E_opt, C, sol = minimise(H)
print(f"Optimized HF energy: {E_opt} Hartree")
```

--------------------------------

### Time Performance of Boys Function

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

This snippet uses the `%%timeit` magic command to measure the execution time of the `boys` function. It calls `boys(2, x, order=64, use_jacobi=False)` multiple times to obtain an average execution time and its standard deviation. This is useful for performance profiling.

```python
%%timeit
boys(2, x, order=64, use_jacobi=False)
```

--------------------------------

### Output Monomial Terms with Ascending Orders (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Generates and formats monomial terms in ascending order of powers (n=0, 1, 2, ...). It uses SymPy for symbolic mathematics and LaTeX for output formatting. Dependencies include 'sympy' and 'IPython.display'.

```python
import sympy as sp
import IPython.display as ipd
from itertools import product

# Assume poly_terms, i, j, x, a, b are defined symbolically
# Example placeholder definitions:
i, j = sp.symbols('i j')
a, b, t, x = sp.symbols('a b t x')
poly_terms = (a + t)**i * (b + t)**j # This is a placeholder, actual poly_terms may differ

output = r"\begin{align*}"
for ival, jval in product(range(3), range(3)):
    p = sp.simplify(sp.expand(poly_terms.subs(i, ival).subs(j, jval)))
    output += f"(a + t)^{{ival}} (b + t)^{{jval}} &= "
    output += " + ".join([
        sp.latex(p.coeff(x, n) * x**n) for n in range(ival + jval + 1)
    ])
    output += r"\\"

output += r"\end{align*}"
ipd.Latex(output)
```

--------------------------------

### Configure Basis Set for a Molecule

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This Python code snippet shows how to configure a basis set for a given molecule using the MESS library. It utilizes the `basisset` function, taking a molecule object and a basis set name (e.g., '6-31g') as input to generate the basis set parameters.

```python
from mess import basisset

basis_name = "6-31g"
basis = basisset(mol, basis_name)
basis
```

--------------------------------

### Solve and Simplify Symbolic Integral

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Takes a previously defined symbolic integral (H_na) and computes its definite integral using '.doit()' and then simplifies the result using '.simplify()'. The solved equation is then displayed.

```python
ipd.display(sp.Eq(H(2 * k, eta), H_na.doit().simplify()))
```

--------------------------------

### Calculate Hartree-Fock Energy with MESS

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This snippet shows how to initialize and use the Hartree-Fock (HF) method in MESS to calculate the ground state energy of a system. It requires a basis set and specifies the 'hfx' exchange-correlation method. The output is the calculated energy.

```python
H = Hamiltonian(basis, xc_method="hfx")
E, C, sol = minimise(H)
float(E)
```

--------------------------------

### Reference Mpmath implementation for gammainc

Source: https://github.com/hatemhelal/mess/blob/main/docs/gammanu.ipynb

This function `reference_mpmath` provides a reference implementation for the lower incomplete gamma function using the `mpmath` library. It's used for cross-validation against SciPy and custom implementations, especially for high-precision calculations.

```python
def reference_mpmath(a, z):
    # https://mpmath.org/doc/current/functions/expintegrals.html#gammainc
    from mpmath import gammainc, log

    return float(log(gammainc(float(a), b=float(z))))
```

--------------------------------

### Optimize Electronic Structure with MESS

Source: https://context7.com/hatemhelal/mess/llms.txt

The `minimise` function finds molecular orbital coefficients that minimize the total electronic energy using gradient-based optimization. It returns the minimized total energy, the optimized coefficients, and a solution object containing convergence information. Custom optimizers from libraries like `optimistix` can be specified.

```python
from mess import molecule, basisset, Hamiltonian, minimise

# Set up system
mol = molecule("water")
basis = basisset(mol, basis_name="6-31g")
H = Hamiltonian(basis, xc_method="lda")

# Minimize energy to find ground state
E_total, C, sol = minimise(H)

print(f"Total energy: {E_total} Hartree")  # Ground state energy
print(f"Converged: {sol.result}")  # Optimization result status
print(f"Coefficients shape: {C.shape}")  # (num_orbitals, num_orbitals)

# Use different optimizer with max steps
import optimistix as optx
solver = optx.BFGS(atol=1e-7, rtol=1e-6)
E, C, sol = minimise(H, max_steps=100, solver=solver)
```

--------------------------------

### Create Molecular Structures with MESS

Source: https://context7.com/hatemhelal/mess/llms.txt

The `molecule` factory function creates predefined molecular structures. It returns a Structure object containing atomic numbers, positions in Bohr units, and the total electron count for the molecule. Supported molecules include 'water', 'h2', 'benzene', and 'methane'.

```python
from mess import molecule

# Create a water molecule with optimized geometry
mol = molecule("water")
print(f"Atoms: {mol.atomic_symbol}")  # ['O', 'H', 'H']
print(f"Total electrons: {mol.num_electrons}")  # 10
print(f"Atomic numbers: {mol.atomic_number}")  # [8 1 1]
print(f"Positions (Bohr): {mol.position}")  # 3x3 array

# Create H2 molecule
h2 = molecule("h2")

# Other available molecules: "benzene", "methane"
benzene = molecule("benzene")
methane = molecule("methane")
```

--------------------------------

### Reference Scipy implementation for gammainc

Source: https://github.com/hatemhelal/mess/blob/main/docs/gammanu.ipynb

This function `reference_scipy` serves as a reference for the lower incomplete gamma function calculation using SciPy's `gammaln` and `gammainc` functions. It's used for comparison and validation of other implementations.

```python
def reference_scipy(a, z):
    # https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.gammainc.html
    from scipy.special import gammaln, gammainc

    return gammaln(a) + np.log(gammainc(a, z))
```

--------------------------------

### Building H2 Molecule Structure for PySCF

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

Defines a function `build_h2` to create a Structure object representing an H2 molecule with adjustable internuclear distance. This structure is then converted for use with PySCF's DFT module to build integration grids.

```Python
def build_h2(r: float):
    pos = np.zeros((2, 3))
    pos[0, 0] = r
    pos[1, 0] = -r
    return Structure(atomic_number=np.asarray([1, 1]), position=pos)


mol = build_h2(5.0)
mol = to_pyscf(mol)
grids = dft.gen_grid.Grids(mol)
grids.build()
weights, points = grids.weights, grids.coords

z0 = points[:, 2] == 0
weights = weights[z0]
x, y = points[z0, 0], points[z0, 1]
plt.scatter(x, y, 10 * np.abs(weights), weights)
plt.colorbar()
plt.title("$H_2$ Integration Grid z = 0")
plt.xlabel("x")
plt.ylabel("y")
```

--------------------------------

### Benchmark MESS Energy Minimization (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This function benchmarks the performance of the energy minimization process in MESS. It calls the `minimise` function and ensures the results are ready using `.block_until_ready()`. The `%timeit` magic command is used to measure the execution time.

```python
def mess_benchmark():
    E, C, _ = minimise(H)
    E, C = E.block_until_ready(), C.block_until_ready()

mess_time = %timeit -o mess_benchmark()
```

--------------------------------

### Implementing the PQ function using PyQuante2

Source: https://github.com/hatemhelal/mess/blob/main/docs/gammanu.ipynb

This Python function `pq` calculates a value using `pyquante2.utils.gamm_inc`. This appears to be an alternative implementation or wrapper for the incomplete gamma function, potentially offering different performance characteristics or numerical stability.

```python
def pq(a, x):
    from pyquante2.utils import gamm_inc

    return gamm_inc(a, x)
```

--------------------------------

### Build Gaussian Basis Sets with MESS

Source: https://context7.com/hatemhelal/mess/llms.txt

The `bassett` function constructs a basis set object for a given molecular structure, fetching parameters from the Basis Set Exchange. It supports various `basis_name` options and can handle both Cartesian and spherical Gaussian basis functions. The function returns the number of orbitals, number of primitives, and electron occupancy per orbital.

```python
from mess import molecule, basisset
import jax.numpy as jnp

# Create molecule and basis set
mol = molecule("water")
basis = basisset(mol, basis_name="6-31g", spherical=True)

print(f"Number of orbitals: {basis.num_orbitals}")  # 13
print(f"Number of primitives: {basis.num_primitives}")
print(f"Occupancy: {basis.occupancy}")  # electron occupancy per orbital

# Different basis sets
basis_sto3g = basisset(mol, basis_name="sto-3g")
basis_def2 = basisset(mol, basis_name="def2-SVP")

# Evaluate basis functions at specific positions
positions = jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
orbital_values = basis(positions)  # Shape: (num_positions, num_orbitals)
```

--------------------------------

### Define and JIT-compile the total energy function

Source: https://github.com/hatemhelal/mess/blob/main/docs/optim.ipynb

Defines a function `total_energy` that calculates the total energy of the system for a given matrix Z. It enforces orthonormalisation of orbitals, computes the density matrix, and adds the nuclear energy. JAX's `jit` and `value_and_grad` are used for just-in-time compilation and automatic differentiation, respectively, to efficiently compute both energy and its gradient.

```python
E_n = nuclear_energy(mol)


@jax.jit
@jax.value_and_grad
def total_energy(Z):
    C = H.orthonormalise(Z)
    P = basis.density_matrix(C)
    return H(P) + E_n
```

--------------------------------

### Check Integral with Odd Exponent

Source: https://github.com/hatemhelal/mess/blob/main/docs/gto_integrals.ipynb

Defines a new integral H_odd with an odd exponent for 't'. It then calculates and displays the result of this integral, demonstrating that the integral of an odd function over symmetric limits is zero.

```python
H_odd = sp.Integral(t ** (2 * k + 1) * sp.exp(-eta * t**2), (t, -sp.oo, sp.oo))
ipd.display(sp.Eq(H_odd, H_odd.doit()))
```

--------------------------------

### Visualize 3D Scatter Plot with Matplotlib

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

This snippet visualizes a 3D scatter plot using Matplotlib. It sets the view angle, plots points with sizes and colors determined by weights, and customizes the axes and grid. The `plt.show()` command displays the generated plot.

```python
ax.view_init(elev=90, azim=-90, roll=0)
ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=np.abs(weights), c=weights)
ax.set(xticklabels=[], yticklabels=[], zticklabels=[])
ax.grid(False)

plt.show()
```

--------------------------------

### Molecular Structure Container - Python

Source: https://context7.com/hatemhelal/mess/llms.txt

Demonstrates the use of the `Structure` dataclass to hold atomic numbers, Cartesian positions, and derived properties of a molecule. It supports easy access to number of atoms, electrons, and atomic symbols, and can compute nuclear repulsion energy. This class is useful for defining and manipulating molecular geometries.

```python
from mess.structure import Structure, nuclear_energy
import numpy as np

# Create custom structure
atomic_numbers = np.array([8, 1, 1])  # O, H, H
positions = np.array([
    [0.0, 0.0, 0.0],
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0]
])  # Positions in Bohr
mol = Structure(atomic_numbers, positions)

print(f"Number of atoms: {mol.num_atoms}")  # 3
print(f"Number of electrons: {mol.num_electrons}")  # 10
print(f"Atomic symbols: {mol.atomic_symbol}")  # ['O', 'H', 'H']

# Calculate nuclear repulsion energy
E_nuc = nuclear_energy(mol)
print(f"Nuclear repulsion: {E_nuc} Hartree")

# Single atom
he = Structure(np.array(2), np.array([0.0, 0.0, 0.0]))
```

--------------------------------

### Comparing gammaincln with SciPy and Mpmath

Source: https://github.com/hatemhelal/mess/blob/main/docs/gammanu.ipynb

This snippet compares the result of the `gammaincln` function with reference implementations from SciPy and Mpmath. It calculates the relative errors to assess the accuracy of the `gammaincln` implementation.

```python
prev_a = 1.5
z = 10.0

ref_sp = reference_scipy(prev_a, z)
ref_mp = reference_mpmath(prev_a, z)
ans = gammaincln(prev_a, z, 64)
abs((ref_sp - ans) / ref_sp), abs((ref_mp - ans) / ref_mp)
```

--------------------------------

### Compare Fnut to Reference Function (Python)

Source: https://github.com/hatemhelal/mess/blob/main/docs/gammanu.ipynb

This snippet compares the output of the custom `Fnut` function with a `reference_pyquante` function for specific `nu` and `t` values (0.0). It calculates the relative error between the two functions and prints the result. It also demonstrates the special case handling when both `nu` and `t` are 0.

```python
nu = 0.0
t = 0.0

abs((Fnut(nu, t) - reference_pyquante(nu, t)) / reference_pyquante(nu, t))
```

```python
Fnut(nu, t), reference_pyquante(nu, t)
```

--------------------------------

### Evaluate and Plot Legendre Polynomials

Source: https://github.com/hatemhelal/mess/blob/main/docs/quadrature.ipynb

This code snippet evaluates Legendre polynomials using `ssp.eval_sh_legendre` for orders 0 to 4 over a linear range of `x` values. It then plots each polynomial on the same axes, allowing for visual comparison of their shapes.

```python
x = np.linspace(0, 1)
for n in range(5):
    y = ssp.eval_sh_legendre(n, x)
    plt.plot(x, y)
```

--------------------------------

### Create a Water Molecule Structure

Source: https://github.com/hatemhelal/mess/blob/main/docs/tour.ipynb

This Python code snippet demonstrates how to create a water molecule structure using the `molecule` function from the MESS library. The resulting `Structure` object represents the atomic numbers and Cartesian coordinates of the atoms in the molecule.

```python
from mess import molecule

mol = molecule("water")
mol
```