diff --git a/gbasis/integrals/overlap_n.py b/gbasis/integrals/overlap_n.py new file mode 100644 index 00000000..2c453256 --- /dev/null +++ b/gbasis/integrals/overlap_n.py @@ -0,0 +1,308 @@ +import numpy as np +from scipy.sparse import coo_matrix +from itertools import product +class PrimitiveNEngine: + + # Gaussian collapse + @staticmethod + def collapse_gaussians(alphas, centers): + alphas = np.asarray(alphas, dtype=np.float64) + centers = np.asarray(centers, dtype=np.float64) + + alpha_tot = np.sum(alphas) + + if alpha_tot <= 0.0: + raise ValueError("Total Gaussian exponent must be positive.") + + P = np.sum(alphas[:, None] * centers, axis=0) / alpha_tot + + term1 = np.sum(alphas * np.sum(centers**2, axis=1)) + term2 = alpha_tot * np.dot(P, P) + + exponent = term2 - term1 + prefactor = np.exp(exponent) + + return alpha_tot, P, prefactor + # Pure binomial Hermite shift + @staticmethod + def hermite_coefficients(l, PA): + """ + Expand (x - A)^l about P: + + (x - A)^l = sum_t E_t (x - P)^t + """ + E = np.zeros(l + 1, dtype=np.float64) + E[0] = 1.0 + + for i in range(l): + E_new = np.zeros(l + 1, dtype=np.float64) + + for t in range(i + 1): + E_new[t] += PA * E[t] + E_new[t + 1] += E[t] + + E = E_new + + return E + # Gaussian moments + @staticmethod + def gaussian_moments(alpha, max_order): + """ + Compute: + ∫ (x-P)^k exp(-alpha (x-P)^2) dx + over (-∞,∞) + """ + moments = np.zeros(max_order + 1, dtype=np.float64) + + # zeroth moment + moments[0] = np.sqrt(np.pi / alpha) + + # only even moments survive + for k in range(0, max_order - 1, 2): + moments[k + 2] = (k + 1) / (2.0 * alpha) * moments[k] + + return moments + # Full primitive N-center overlap + @staticmethod + def primitive_overlap(alphas, centers, angmoms): + + alpha_tot, P, prefactor = PrimitiveNEngine.collapse_gaussians( + alphas, centers + ) + + result = prefactor + + # factorize into x, y, z + for axis in range(3): + + # build total polynomial via convolution + E_total = np.array([1.0], dtype=np.float64) + + for i in range(len(alphas)): + l = angmoms[i][axis] + PA = P[axis] - centers[i][axis] + + E = PrimitiveNEngine.hermite_coefficients(l, PA) + + E_total = np.convolve(E_total, E) + + moments = PrimitiveNEngine.gaussian_moments( + alpha_tot, + len(E_total) - 1 + ) + + axis_integral = np.dot(E_total, moments[:len(E_total)]) + + result *= axis_integral + + return result +# Screening Function +def is_n_shell_overlap_screened(shells, tol=1e-12): + """ + Conservative exponential upper-bound screening + for N-center contracted overlap. + """ + + alpha_mins = [np.min(shell.exps) for shell in shells] + centers = [shell.coord for shell in shells] + + alpha_tot = sum(alpha_mins) + + if alpha_tot <= 0.0: + return True + + # Exponential decay from Gaussian collapse + decay_sum = 0.0 + N = len(shells) + + for i in range(N): + for j in range(i + 1, N): + Rij = centers[i] - centers[j] + Rij2 = np.dot(Rij, Rij) + decay_sum += alpha_mins[i] * alpha_mins[j] * Rij2 + + D = decay_sum / alpha_tot + + # Contraction-level magnitude bound + coeff_bound = 1.0 + norm_bound = 1.0 + + for shell in shells: + coeff_bound *= np.max(np.abs(shell.coeffs)) + norm_bound *= np.max(np.abs(shell.norm_prim_cart)) + + volume_bound = (np.pi / alpha_tot) ** 1.5 + + bound = coeff_bound * norm_bound * volume_bound * np.exp(-D) + + return bound < tol +def contracted_n_overlap(shells): + """ + Compute contracted N-center overlap for a list of + GeneralizedContractionShell objects. + + Parameters + ---------- + shells : list[GeneralizedContractionShell] + + Returns + ------- + np.ndarray + N-dimensional array over segmented contractions + and Cartesian angular components. + + Shape: + (M1, L1, M2, L2, ..., MN, LN) + """ + + N = len(shells) + + # Build shape for final tensor + shape = [] + for shell in shells: + shape.append(shell.num_seg_cont) + shape.append(shell.num_cart) + + result = np.zeros(shape, dtype=np.float64) + + # Primitive exponent index ranges + prim_ranges = [range(len(shell.exps)) for shell in shells] + + # Segmented contraction indices + seg_ranges = [range(shell.num_seg_cont) for shell in shells] + + # Cartesian angular component indices + cart_ranges = [range(shell.num_cart) for shell in shells] + + for seg_indices in product(*seg_ranges): + for cart_indices in product(*cart_ranges): + + total_value = 0.0 + for prim_indices in product(*prim_ranges): + + alphas = [] + centers = [] + angmoms = [] + coeff_prod = 1.0 + norm_prod = 1.0 + + for i, shell in enumerate(shells): + + p = prim_indices[i] + m = seg_indices[i] + c = cart_indices[i] + + alpha = shell.exps[p] + coeff = shell.coeffs[p, m] + norm = shell.norm_prim_cart[c, p] + angmom = tuple(shell.angmom_components_cart[c]) + + alphas.append(alpha) + centers.append(shell.coord) + angmoms.append(angmom) + + coeff_prod *= coeff + norm_prod *= norm + + prim_val = PrimitiveNEngine.primitive_overlap( + alphas, + centers, + angmoms + ) + + total_value += coeff_prod * norm_prod * prim_val + + index = [] + for i in range(N): + index.append(seg_indices[i]) + index.append(cart_indices[i]) + + result[tuple(index)] = total_value + + return result +def build_n_overlap_tensor(shells, tol=1e-12): + """ + Build sparse N-center overlap tensor over shells. + + Parameters + ---------- + shells : list[GeneralizedContractionShell] + tol : float + Screening tolerance + + Returns + ------- + scipy.sparse.coo_matrix + Flattened sparse tensor of shape (total_ao^N, 1) + """ + + # Total AO dimension + shell_sizes = [ + shell.num_seg_cont * shell.num_cart + for shell in shells + ] + + total_ao = sum(shell_sizes) + N = len(shells) + + data = [] + rows = [] + + # compute AO offsets per shell + offsets = [] + acc = 0 + for size in shell_sizes: + offsets.append(acc) + acc += size + + for shell_indices in product(range(len(shells)), repeat=N): + + shell_tuple = [shells[i] for i in shell_indices] + + # Screening + if is_n_shell_overlap_screened(shell_tuple, tol=tol): + continue + + block = contracted_n_overlap(shell_tuple) + + block_flat = block.reshape(-1) + + local_sizes = [ + shells[i].num_seg_cont * shells[i].num_cart + for i in shell_indices + ] + + local_offsets = [ + offsets[i] + for i in shell_indices + ] + + for local_idx, value in enumerate(block_flat): + + if abs(value) < tol: + continue + + # convert local multi-index to global index + multi = [] + tmp = local_idx + + for size in reversed(local_sizes): + multi.append(tmp % size) + tmp //= size + + multi = list(reversed(multi)) + + global_index = 0 + for k in range(N): + global_index = ( + global_index * total_ao + + local_offsets[k] + + multi[k] + ) + + rows.append(global_index) + data.append(value) + + shape = (total_ao ** N, 1) + + return coo_matrix((data, (rows, np.zeros(len(rows)))), shape=shape) \ No newline at end of file diff --git a/tests/test_n.py b/tests/test_n.py new file mode 100644 index 00000000..93cac243 --- /dev/null +++ b/tests/test_n.py @@ -0,0 +1,57 @@ +import numpy as np +import pytest + +from gbasis.parsers import parse_gbs, make_contractions +from gbasis.integrals.overlap_n import contracted_n_overlap + + +# ===================================================== +# Helper: Build hydrogen basis +# ===================================================== + +def build_h2_basis(): + + atoms = ["H", "H"] + + coords = np.array([ + [0.0, 0.0, 0.0], + [0.0, 0.0, 1.0] + ]) + + basis_dict = parse_gbs("tests/data_631g.gbs") + + basis = make_contractions( + basis_dict, + atoms, + coords, + coord_types="c" + ) + + return basis + + +# ===================================================== +# Test arbitrary orders N=1 to N=6 +# ===================================================== + +@pytest.mark.parametrize("N", [1, 2, 3, 4, 5, 6]) +def test_arbitrary_n_orders(N): + + basis = build_h2_basis() + + shells = [] + + for i in range(N): + shells.append(basis[i % len(basis)]) + + result = contracted_n_overlap(shells) + + print(f"N={N}, max abs value:", np.max(np.abs(result))) + + assert result is not None + + assert not np.isnan(result).any() + + assert not np.isinf(result).any() + + assert np.all(np.isfinite(result)) \ No newline at end of file diff --git a/tests/test_n_engine_stress.py b/tests/test_n_engine_stress.py new file mode 100644 index 00000000..76ebfebf --- /dev/null +++ b/tests/test_n_engine_stress.py @@ -0,0 +1,115 @@ +import numpy as np +import pytest + +from gbasis.parsers import parse_gbs, make_contractions +from gbasis.integrals.overlap import overlap_integral +from gbasis.integrals.overlap_n import contracted_n_overlap + + +# ===================================================== +# Helper: Build small hydrogen basis +# ===================================================== +def build_h2_basis(): + atoms = ["H", "H"] + coords = np.array([[0.0, 0.0, 0.0], + [0.0, 0.0, 1.0]]) + + basis_dict = parse_gbs("tests/data_631g.gbs") + + basis = make_contractions( + basis_dict, + atoms, + coords, + coord_types="c" + ) + + return basis + + +# ===================================================== +# 1️⃣ N=2 FULL AO MATRIX VALIDATION +# ===================================================== + +def test_n2_matches_gbasis(): + + basis = build_h2_basis() + + ref = overlap_integral(basis, screen_basis=False) + + n_shells = len(basis) + total_ao = ref.shape[0] + mine = np.zeros((total_ao, total_ao)) + + ao_i = 0 + for shell_i in basis: + size_i = shell_i.num_seg_cont * shell_i.num_cart + + ao_j = 0 + for shell_j in basis: + size_j = shell_j.num_seg_cont * shell_j.num_cart + + block = contracted_n_overlap([shell_i, shell_j]) + block = block.reshape(size_i, size_j) + + mine[ao_i:ao_i+size_i, + ao_j:ao_j+size_j] = block + + ao_j += size_j + + ao_i += size_i + + diff = np.max(np.abs(ref - mine)) + + print("N=2 max difference:", diff) + + assert np.allclose(ref, mine, atol=1e-10) + + +# ===================================================== +# 2️⃣ N=3 SYMMETRY TEST +# ===================================================== + +def test_n3_symmetry(): + + basis = build_h2_basis() + + shell1 = basis[0] + shell2 = basis[1] + shell3 = basis[0] + + S123 = contracted_n_overlap([shell1, shell2, shell3]) + S321 = contracted_n_overlap([shell3, shell2, shell1]) + + diff = np.max(np.abs(S123 - S321)) + + print("N=3 symmetry difference:", diff) + + assert np.allclose(S123, S321, atol=1e-10) + + +# ===================================================== +# 3️⃣ HIGH ANGULAR MOMENTUM + DIFFUSE TEST +# ===================================================== + +def test_high_angmom_diffuse(): + + basis = build_h2_basis() + + # choose a higher angular momentum shell if available + # try to find p-shell + p_shells = [shell for shell in basis if shell.angmom == 1] + + if not p_shells: + pytest.skip("No p-shell available in basis.") + + shell = p_shells[0] + + # artificially make exponents diffuse + shell.exps = shell.exps * 0.01 + + S = contracted_n_overlap([shell, shell]) + + print("Diffuse shell max abs value:", np.max(np.abs(S))) + + assert not np.isnan(S).any() + assert not np.isinf(S).any() diff --git a/tests/test_overlap_n.py b/tests/test_overlap_n.py new file mode 100644 index 00000000..a2155a3c --- /dev/null +++ b/tests/test_overlap_n.py @@ -0,0 +1,51 @@ +import numpy as np +import pytest + +from gbasis.integrals.overlap import Overlap +from gbasis.integrals.overlap_n import contracted_n_overlap, PrimitiveNEngine +from gbasis.contractions import GeneralizedContractionShell + + +def build_test_shell(center): + ... + return GeneralizedContractionShell(...) + + +def test_two_center_matches_existing_overlap(): + ... + assert ... + + +def test_two_center_ss_analytic(): + ... + assert ... + + +def test_primitive_px_py_pz(): + A = np.array([0.0, 0.0, 0.0]) + B = np.array([0.3, 0.1, -0.2]) + + alpha = 0.5 + + print("px-px:", + PrimitiveNEngine.primitive_overlap( + [alpha, alpha], + [A, B], + [(1,0,0), (1,0,0)] + )) + + print("py-py:", + PrimitiveNEngine.primitive_overlap( + [alpha, alpha], + [A, B], + [(0,1,0), (0,1,0)] + )) + + print("pz-pz:", + PrimitiveNEngine.primitive_overlap( + [alpha, alpha], + [A, B], + [(0,0,1), (0,0,1)] + )) + + assert True # temporary so pytest runs it