File size: 3,586 Bytes
2a87ede | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """Overlap integrals between Gaussian shells via the Obara-Saika recursion.
For two 1D Gaussians centered at A, B with exponents a, b, let
p = a+b and P = (aA+bB)/p be the center of their product (Gaussian
product theorem). The overlap S[i,j] = <(x-A)^i exp(-a(x-A)^2) |
(x-B)^j exp(-b(x-B)^2)> obeys two recursions:
Vertical (build up the first index from the base case S[0,0]):
S[i,0] = (P-A) S[i-1,0] + (i-1)/(2p) S[i-2,0]
Horizontal (build up the second index by shifting angular momentum
from center A to center B, exact for any two centers -- this is why
it needs no exponent-dependent term):
S[i,j] = (A-B) S[i,j-1] + S[i+1,j-1]
Both are linear recursions in one index with a 2-term memory, so each
maps directly onto jax.lax.scan: the whole angular-momentum ladder for
a shell pair compiles to one XLA loop instead of a Python for-loop.
"""
import functools
import jax
import jax.numpy as jnp
from dense_evolution.native_hf.gaussians import (
GaussianShell1D,
GaussianShell3D,
product_center,
product_prefactor,
)
def _base_overlap_1d(g1: GaussianShell1D, g2: GaussianShell1D) -> jax.Array:
p = g1.exponent + g2.exponent
mu = (g1.exponent * g2.exponent) / p
K = jnp.exp(-mu * jnp.square(g1.center - g2.center))
return jnp.sqrt(jnp.pi / p) * K
def _vertical_step(p_minus_a, inv_two_p, carry, i):
s_prev, s_prev2 = carry
s_next = p_minus_a * s_prev + (i - 1) * inv_two_p * s_prev2
return (s_next, s_prev), s_next
def _build_first_index(s00, g1: GaussianShell1D, g2: GaussianShell1D, n_extra: int) -> jax.Array:
"""Returns S[i, 0] for i = 0..n_extra."""
if n_extra == 0:
return s00[None]
p = g1.exponent + g2.exponent
P = (g1.exponent * g1.center + g2.exponent * g2.center) / p
step = functools.partial(_vertical_step, P - g1.center, 1.0 / (2.0 * p))
_, rest = jax.lax.scan(step, (s00, jnp.zeros_like(s00)), jnp.arange(1, n_extra + 1))
return jnp.concatenate([s00[None], rest])
def _horizontal_step(a_minus_b, column, _):
shifted = jnp.roll(column, shift=-1)
new_column = a_minus_b * column + shifted
return new_column, new_column
def _build_second_index(first_col: jax.Array, g1: GaussianShell1D, g2: GaussianShell1D) -> jax.Array:
"""first_col holds S[:,0]. Returns S[:, 0..g2.degree] (rows beyond
validity are garbage and get sliced away by the caller)."""
if g2.degree == 0:
return first_col[:, None]
step = functools.partial(_horizontal_step, g1.center - g2.center)
_, rest = jax.lax.scan(step, first_col, jnp.arange(1, g2.degree + 1))
return jnp.concatenate([first_col[:, None], jnp.moveaxis(rest, 0, -1)], axis=-1)
def overlap_1d(g1: GaussianShell1D, g2: GaussianShell1D) -> jax.Array:
"""S[i,j] for 0<=i<=g1.degree, 0<=j<=g2.degree. Shape (g1.degree+1, g2.degree+1)."""
s00 = _base_overlap_1d(g1, g2)
first_col = _build_first_index(s00, g1, g2, g1.degree + g2.degree)
full = _build_second_index(first_col, g1, g2)
return full[: g1.degree + 1, :]
@jax.jit
def overlap_3d(g1: GaussianShell3D, g2: GaussianShell3D) -> jax.Array:
"""S[ix,iy,iz,jx,jy,jz] for the full Cartesian shell pair.
Shape: (L1+1,L1+1,L1+1, L2+1,L2+1,L2+1) where L1, L2 are the shell
degrees (unphysical (i,j,k) combinations with i+j+k > degree are
simply never read by the caller).
"""
axes = [overlap_1d(g1.component(d), g2.component(d)) for d in range(3)]
return jnp.einsum("ad,be,cf->abcdef", *axes)
|