File size: 8,368 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """Assembling full molecular AO integral matrices from contracted shells.
Each pair (or quartet, for repulsion) of shells contributes a block to
the overall S/T/V matrices (or ERI tensor). We loop over shell
pairs/quartets in plain Python -- for a minimal basis like STO-3G there
are only a handful of shells per atom (Si2/STO-3G: 6 shells total), so
this loop is cheap. The sum over primitives *within* a shell pair/quartet
and the slicing down to physical Cartesian components both happen
inside a single jax.jit-compiled call per shell pair/quartet: doing
that slicing eagerly (one jnp.ndarray.__getitem__ per primitive
combination) turned out to cost as much per-call dispatch overhead as
PennyLane's own scalar-at-a-time autograd loop, defeating the purpose,
so everything from "loop over primitives" to "slice out the physical
components" is fused into one compiled program per shell pair/quartet
and only that program's *output* touches eager Python.
"""
import functools
import jax
import jax.numpy as jnp
import numpy as np
from dense_evolution.native_hf.basis import ContractedShell, n_cartesian_functions
from dense_evolution.native_hf.cartesian import cartesian_powers
from dense_evolution.native_hf.gaussians import GaussianShell3D
from dense_evolution.native_hf.overlap import overlap_3d
from dense_evolution.native_hf.kinetic import kinetic_3d
from dense_evolution.native_hf.coulomb import nuclear_attraction, electron_repulsion
@functools.partial(jax.jit, static_argnames=("degree_a", "degree_b", "primitive_fn"))
def _pair_block(exponents_a, coeffs_a, center_a, degree_a, exponents_b, coeffs_b, center_b, degree_b, primitive_fn, extra=()):
"""extra: additional *traced* arrays (e.g. nuclear charges/positions)
forwarded to primitive_fn unchanged -- traced (not static) so that
calling this with a different molecule geometry reuses the same
compiled program instead of retracing. primitive_fn must itself be a
stable, module-level (not freshly-created-per-call) callable, since
it IS a static jit key."""
powers_a = jnp.asarray(cartesian_powers(degree_a))
powers_b = jnp.asarray(cartesian_powers(degree_b))
def one_primitive_pair(ea, ca, eb, cb):
ga = GaussianShell3D(degree=degree_a, exponent=ea, center=center_a)
gb = GaussianShell3D(degree=degree_b, exponent=eb, center=center_b)
full = primitive_fn(ga, gb, *extra)
sliced = full[powers_a[:, 0], powers_a[:, 1], powers_a[:, 2]][:, powers_b[:, 0], powers_b[:, 1], powers_b[:, 2]]
return ca * cb * sliced
batched = jax.vmap(
jax.vmap(one_primitive_pair, in_axes=(None, None, 0, 0)),
in_axes=(0, 0, None, None),
)(exponents_a, coeffs_a, exponents_b, coeffs_b)
return jnp.sum(batched, axis=(0, 1))
@functools.partial(jax.jit, static_argnames=("degrees", "primitive_fn"))
def _quartet_block(exponents, coeffs, centers, degrees, primitive_fn):
"""exponents/coeffs: tuples of 4 arrays (one per shell), centers: tuple
of 4 (3,) arrays, degrees: tuple of 4 static ints."""
powers = [jnp.asarray(cartesian_powers(d)) for d in degrees]
def one_quartet(ea, ca, eb, cb, ec, cc, ed, cd):
ga = GaussianShell3D(degree=degrees[0], exponent=ea, center=centers[0])
gb = GaussianShell3D(degree=degrees[1], exponent=eb, center=centers[1])
gc = GaussianShell3D(degree=degrees[2], exponent=ec, center=centers[2])
gd = GaussianShell3D(degree=degrees[3], exponent=ed, center=centers[3])
full = primitive_fn(ga, gb, gc, gd)
sliced = full[powers[0][:, 0], powers[0][:, 1], powers[0][:, 2]]
sliced = sliced[:, powers[1][:, 0], powers[1][:, 1], powers[1][:, 2]]
sliced = sliced[:, :, powers[2][:, 0], powers[2][:, 1], powers[2][:, 2]]
sliced = sliced[:, :, :, powers[3][:, 0], powers[3][:, 1], powers[3][:, 2]]
return ca * cb * cc * cd * sliced
vmapped = one_quartet
# vmap outward-in: d, c, b, a -- each adds a leading batch axis.
vmapped = jax.vmap(vmapped, in_axes=(None, None, None, None, None, None, 0, 0))
vmapped = jax.vmap(vmapped, in_axes=(None, None, None, None, 0, 0, None, None))
vmapped = jax.vmap(vmapped, in_axes=(None, None, 0, 0, None, None, None, None))
vmapped = jax.vmap(vmapped, in_axes=(0, 0, None, None, None, None, None, None))
ea, ca, eb, cb, ec, cc, ed, cd = (
exponents[0], coeffs[0], exponents[1], coeffs[1],
exponents[2], coeffs[2], exponents[3], coeffs[3],
)
batched = vmapped(ea, ca, eb, cb, ec, cc, ed, cd)
return jnp.sum(batched, axis=(0, 1, 2, 3))
def _shell_offsets(shells: list[ContractedShell]) -> list[int]:
offsets, running = [], 0
for s in shells:
offsets.append(running)
running += len(cartesian_powers(s.degree))
return offsets
def _overlap_primitive(ga, gb):
return overlap_3d(ga, gb)
def build_overlap_matrix(shells: list[ContractedShell]) -> np.ndarray:
return _build_two_index(shells, _overlap_primitive)
def build_core_hamiltonian(shells: list[ContractedShell], nuclear_charges: list[float], nuclear_positions: np.ndarray) -> np.ndarray:
# NOTE: this closure is rebuilt (and re-jitted) on every call, which
# defeats cross-call jit caching for e.g. a bond-length scan -- see
# prog.txt task "fix cross-call JIT caching". Reverted to this
# simpler, independently-verified-correct form (matches slaterform
# to 10 significant figures on Si2/STO-3G) after a traced-args
# refactor (charges/positions as jax arrays via vmap over nuclei,
# threaded through a shared `extra` argument in _pair_block) passed
# every individual/aggregate check (element-wise Hc, eigenvalues,
# trace, sum, 5x determinism) yet still produced a different,
# wrong total SCF energy end-to-end for reasons not pinned down in
# the time available -- not worth shipping an unexplained
# discrepancy just for speed.
charges = tuple(float(z) for z in nuclear_charges)
positions = tuple(jnp.asarray(r) for r in nuclear_positions)
def core_primitive(ga, gb):
T = kinetic_3d(ga, gb)
V = jnp.zeros_like(T)
for z, r in zip(charges, positions):
V = V - z * nuclear_attraction(ga, gb, r)
return -0.5 * T + V
return _build_two_index(shells, core_primitive)
def _build_two_index(shells: list[ContractedShell], primitive_fn) -> np.ndarray:
n = n_cartesian_functions(shells)
offsets = _shell_offsets(shells)
M = np.zeros((n, n))
for i, sa in enumerate(shells):
for j, sb in enumerate(shells):
block = np.array(
_pair_block(
sa.exponents, sa.coefficients, sa.center, sa.degree,
sb.exponents, sb.coefficients, sb.center, sb.degree,
primitive_fn,
)
)
oi, oj = offsets[i], offsets[j]
M[oi : oi + block.shape[0], oj : oj + block.shape[1]] = block
return M
def build_repulsion_tensor(shells: list[ContractedShell]) -> np.ndarray:
n = n_cartesian_functions(shells)
offsets = _shell_offsets(shells)
V = np.zeros((n, n, n, n))
for i, sa in enumerate(shells):
for j, sb in enumerate(shells):
for k, sc in enumerate(shells):
for l, sd in enumerate(shells):
block = np.array(
_quartet_block(
(sa.exponents, sb.exponents, sc.exponents, sd.exponents),
(sa.coefficients, sb.coefficients, sc.coefficients, sd.coefficients),
(sa.center, sb.center, sc.center, sd.center),
(sa.degree, sb.degree, sc.degree, sd.degree),
electron_repulsion,
)
)
oi, oj, ok, ol = offsets[i], offsets[j], offsets[k], offsets[l]
V[
oi : oi + block.shape[0],
oj : oj + block.shape[1],
ok : ok + block.shape[2],
ol : ol + block.shape[3],
] = block
return V
|