| """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)
|
|
|
