dense-evolution 8.1.80__tar.gz → 8.2.0__tar.gz

{dense_evolution-8.1.80 → dense_evolution-8.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dense-evolution
-Version: 8.1.80
+Version: 8.2.0
 Summary: High-performance quantum simulation toolkit -- Statevector/MPS engines with JIT compilation, noise modeling, VQE, QEC, quantum chemistry, and agent-native tooling
 Author-email: Salvatore Pennacchio <jtatopenn@libero.it>
 License: Business Source License 1.1
@@ -60,6 +60,10 @@ Requires-Dist: httpx>=0.27; extra == "mcp"
 Requires-Dist: pydantic>=2.0.0; extra == "mcp"
 Provides-Extra: libcint
 Requires-Dist: pyscf>=2.0.0; extra == "libcint"
+Provides-Extra: qmmm
+Requires-Dist: rdkit>=2023.9.1; extra == "qmmm"
+Provides-Extra: ase
+Requires-Dist: ase>=3.22.0; extra == "ase"
 Provides-Extra: full
 Requires-Dist: cupy-cuda12x>=12.0.0; extra == "full"
 Requires-Dist: streamlit>=1.30.0; extra == "full"
@@ -85,6 +89,7 @@ Requires-Dist: ruff>=0.4.0; extra == "docs"
 Requires-Dist: pennylane>=0.35.0; extra == "docs"
 Requires-Dist: basis_set_exchange>=0.9; extra == "docs"
 Requires-Dist: plotly>=5.0.0; extra == "docs"
+Requires-Dist: rdkit>=2023.9.1; extra == "docs"
 Dynamic: license-file
 
 <p align="center">
@@ -285,7 +290,7 @@ sim.run_chunk(circuit.to_tuples(), chunk_size_gates=500)
 - **Real noise, real mitigation.** Stochastic Kraus channels, real-device noise imported from Qiskit backends, and Zero-Noise Extrapolation to correct for it. [Noise](https://tatopenn-cell.github.io/Dense-Evolution/api/noise/) · [Mitigation](https://tatopenn-cell.github.io/Dense-Evolution/api/mitigation/) · [what noise/mitigation/healing each mean](https://tatopenn-cell.github.io/Dense-Evolution/concepts/).
 - **Differentiable VQE, from scratch.** `circuit_to_energy_fn` is the same JAX-differentiable engine real molecular VQE runs on — real Hartree-Fock Hamiltonians, UCCSD/hardware-efficient ansätze, Adam optimization. [Autodiff](https://tatopenn-cell.github.io/Dense-Evolution/api/autodiff/).
 - **OpenQASM 2.0/3.0, both directions.** A real parser, plus Qiskit/PennyLane interop bridges. [QASM Parser](https://tatopenn-cell.github.io/Dense-Evolution/api/parser/) · [Interop](https://tatopenn-cell.github.io/Dense-Evolution/api/interop/).
-- **Code-agnostic QEC decoding**, Majorana/Jordan-Wigner fermion mapping, from-scratch Hartree-Fock for elements outside PennyLane's own basis set, and a traversable-wormhole-inspired teleportation protocol — see the [full API reference](https://tatopenn-cell.github.io/Dense-Evolution/api/) for all of it.
+- **Code-agnostic QEC decoding**, Majorana/Jordan-Wigner fermion mapping, from-scratch Hartree-Fock for elements outside PennyLane's own basis set, a traversable-wormhole-inspired teleportation protocol, and real QM/MM region partitioning (ring-safe BFS, Diffuse2Seg-derived relevance propagation) — see the [full API reference](https://tatopenn-cell.github.io/Dense-Evolution/api/) for all of it.
 
 ## ▍ Benchmarks
 

{dense_evolution-8.1.80 → dense_evolution-8.2.0}/README.md

@@ -196,7 +196,7 @@ sim.run_chunk(circuit.to_tuples(), chunk_size_gates=500)
 - **Real noise, real mitigation.** Stochastic Kraus channels, real-device noise imported from Qiskit backends, and Zero-Noise Extrapolation to correct for it. [Noise](https://tatopenn-cell.github.io/Dense-Evolution/api/noise/) · [Mitigation](https://tatopenn-cell.github.io/Dense-Evolution/api/mitigation/) · [what noise/mitigation/healing each mean](https://tatopenn-cell.github.io/Dense-Evolution/concepts/).
 - **Differentiable VQE, from scratch.** `circuit_to_energy_fn` is the same JAX-differentiable engine real molecular VQE runs on — real Hartree-Fock Hamiltonians, UCCSD/hardware-efficient ansätze, Adam optimization. [Autodiff](https://tatopenn-cell.github.io/Dense-Evolution/api/autodiff/).
 - **OpenQASM 2.0/3.0, both directions.** A real parser, plus Qiskit/PennyLane interop bridges. [QASM Parser](https://tatopenn-cell.github.io/Dense-Evolution/api/parser/) · [Interop](https://tatopenn-cell.github.io/Dense-Evolution/api/interop/).
-- **Code-agnostic QEC decoding**, Majorana/Jordan-Wigner fermion mapping, from-scratch Hartree-Fock for elements outside PennyLane's own basis set, and a traversable-wormhole-inspired teleportation protocol — see the [full API reference](https://tatopenn-cell.github.io/Dense-Evolution/api/) for all of it.
+- **Code-agnostic QEC decoding**, Majorana/Jordan-Wigner fermion mapping, from-scratch Hartree-Fock for elements outside PennyLane's own basis set, a traversable-wormhole-inspired teleportation protocol, and real QM/MM region partitioning (ring-safe BFS, Diffuse2Seg-derived relevance propagation) — see the [full API reference](https://tatopenn-cell.github.io/Dense-Evolution/api/) for all of it.
 
 ## ▍ Benchmarks
 

{dense_evolution-8.1.80 → dense_evolution-8.2.0}/dense_evolution/__init__.py

@@ -57,7 +57,7 @@ from .physics.qec import (pauli_commutes, compute_syndrome, erasure_aware_decode
                            blind_minimum_weight_decode, decode_with_erasure_fallback,
                            counts_in_intervals_dimension, nearest_coset_decode)
 
-__version__ = "8.1.80"
+__version__ = "8.1.81"
 
 __all__ = [
     "__version__",

{dense_evolution-8.1.80 → dense_evolution-8.2.0}/dense_evolution/backends/chunk/geometry.py

@@ -19,7 +19,17 @@ def get_dynamic_chunk(dtype_target) -> int:
         bpe = 8
     max_elements = safe_bytes / bpe
     max_bits = int(np.floor(np.log2(max(max_elements, 2.0))))
-    return max(16, min(max_bits, 27))
+    # Upper bound raised 27 -> 30 (prog.txt / user request): the real
+    # constraint is already _device_memory_budget_bytes's own VRAM/RAM
+    # reading (a 27-bit chunk is only 2.1GB regardless of dtype -- any
+    # GPU with more than ~2.5GB free was being capped well below what it
+    # could actually hold). 30 keeps a ceiling at all (avoids ever
+    # returning a chunk_size_bits so large a single chunk's own indexing
+    # arithmetic elsewhere risks int32 overflow) without arbitrarily
+    # discarding real headroom on 16GB+ GPUs (a Kaggle T4/P100 at
+    # ~0.85*16GB/16B per complex128 element already computes max_bits
+    # near 30 on its own -- this line was the thing throwing that away).
+    return max(16, min(max_bits, 30))
 
 
 def _dtype_for_qubits(n_qubits: int):

{dense_evolution-8.1.80 → dense_evolution-8.2.0}/dense_evolution/native_hf/basis.py

@@ -22,6 +22,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 
+from dense_evolution.config import ensure_x64
 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
@@ -105,6 +106,7 @@ def load_element_shells(basis_name: str, atomic_number: int, center: jax.Array,
 
 def build_molecule_shells(atomic_numbers: list[int], geometry_bohr: np.ndarray, basis_name: str) -> list[ContractedShell]:
     """geometry_bohr: shape (n_atoms, 3), atomic units."""
+    ensure_x64()
     shells = []
     for i, (z, r) in enumerate(zip(atomic_numbers, geometry_bohr)):
         shells.extend(load_element_shells(basis_name, z, jnp.asarray(r), i))

{dense_evolution-8.1.80 → dense_evolution-8.2.0}/dense_evolution/native_hf/libcint_bridge.py

@@ -1,4 +1,10 @@
-"""Optional PySCF/libcint bridge for native_hf's electron-repulsion tensor.
+"""Optional PySCF/libcint bridge for native_hf's electron-repulsion tensor,
+and (build_overlap_and_core_hamiltonian_libcint) its one-electron overlap
+and core-Hamiltonian integrals too -- found necessary, not just faster in
+principle: a larger basis (6-31G*) ran native_hf's own one-electron JAX
+assembly out of memory during XLA JIT compilation on a real Kaggle CPU
+kernel, even with the ERI tensor already libcint-backed (see that
+function's own docstring for the concrete case).
 
 native_hf's own build_repulsion_tensor (assembly.py) computes the ERI
 tensor via JAX-jitted Obara-Saika recursions -- correct, differentiable,
@@ -134,19 +140,15 @@ def _import_pyscf():
     return gto
 
 
-def build_repulsion_tensor_libcint(atomic_numbers: list, geometry_bohr: np.ndarray, basis_name: str) -> np.ndarray:
-    """Same contract as assembly.build_repulsion_tensor, but computed via
-    PySCF/libcint instead of native_hf's own JAX recursions -- takes
-    (atomic_numbers, geometry_bohr, basis_name) rather than a shells list
-    since PySCF needs to build its own `Mole` from the same spec, not
-    native_hf's ContractedShell objects.
-
-    geometry_bohr: shape (n_atoms, 3), atomic units, same convention as
-    build_molecule_shells."""
+def _libcint_mol_and_perm_rescale(atomic_numbers: list, geometry_bohr: np.ndarray, basis_name: str):
+    """Shared Mole/permutation/rescale setup for every one- and two-electron
+    libcint integral below -- S, H_core, and the ERI tensor all live in the
+    same AO basis and need the identical native_hf<->libcint AO reordering
+    and per-AO normalization (see module docstring), so this factors out
+    what build_repulsion_tensor_libcint used to compute standalone."""
     gto = _import_pyscf()
     shells = build_molecule_shells(atomic_numbers, geometry_bohr, basis_name)
     n = n_cartesian_functions(shells)
-
     atom_spec = [
         (int(z), (float(r[0]), float(r[1]), float(r[2])))
         for z, r in zip(atomic_numbers, geometry_bohr)
@@ -158,9 +160,58 @@ def build_repulsion_tensor_libcint(atomic_numbers: list, geometry_bohr: np.ndarr
             f"for basis {basis_name!r} -- the two engines disagree on this basis "
             f"set's shell composition, not just component ordering."
         )
-
     overlap_diag = np.diag(mol.intor("int1e_ovlp"))
     rescale = 1.0 / np.sqrt(overlap_diag)
+    ao_starts = _match_pyscf_shell_ao_starts(mol, shells)
+    perm = np.empty(n, dtype=np.int64)
+    offset = 0
+    for shell, pyscf_start in zip(shells, ao_starts):
+        block_size = len(cartesian_powers(shell.degree))
+        perm[offset : offset + block_size] = pyscf_start + _permutation_libcint_to_native_hf(shell.degree)
+        offset += block_size
+    return mol, perm, rescale
+
+
+def build_overlap_and_core_hamiltonian_libcint(
+    atomic_numbers: list, geometry_bohr: np.ndarray, basis_name: str
+) -> tuple[np.ndarray, np.ndarray]:
+    """Same contract as assembly.build_overlap_matrix + build_core_hamiltonian
+    combined, but computed via PySCF/libcint instead of native_hf's own
+    JAX-jitted one-electron integrals.
+
+    Found necessary, not just faster in principle: on a real Kaggle CPU
+    kernel, native_hf's own one-electron assembly ran out of memory during
+    XLA JIT compilation for a 4-heavy-atom molecule at 6-31G* -- even with
+    the (much more expensive) ERI tensor already routed through
+    build_repulsion_tensor_libcint above, the one-electron path alone was
+    enough to exhaust available memory on that larger basis. Extending the
+    same libcint bridge to S/H_core fixed it (verified: the same molecule/
+    basis that crashed before completed in this repo's own Discovery
+    isodesmic-scission-energy experiment once this function existed).
+
+    Returns (S, H_core) in native_hf's own AO ordering/normalization, ready
+    to pass directly into scf.run_scf alongside build_repulsion_tensor_libcint's
+    own output.
+    """
+    mol, perm, rescale = _libcint_mol_and_perm_rescale(atomic_numbers, geometry_bohr, basis_name)
+    S = mol.intor("int1e_ovlp") * rescale[:, None] * rescale[None, :]
+    T = mol.intor("int1e_kin") * rescale[:, None] * rescale[None, :]
+    V_nuc = mol.intor("int1e_nuc") * rescale[:, None] * rescale[None, :]
+    H_core = T + V_nuc
+    return S[perm][:, perm], H_core[perm][:, perm]
+
+
+def build_repulsion_tensor_libcint(atomic_numbers: list, geometry_bohr: np.ndarray, basis_name: str) -> np.ndarray:
+    """Same contract as assembly.build_repulsion_tensor, but computed via
+    PySCF/libcint instead of native_hf's own JAX recursions -- takes
+    (atomic_numbers, geometry_bohr, basis_name) rather than a shells list
+    since PySCF needs to build its own `Mole` from the same spec, not
+    native_hf's ContractedShell objects.
+
+    geometry_bohr: shape (n_atoms, 3), atomic units, same convention as
+    build_molecule_shells."""
+    mol, perm, rescale = _libcint_mol_and_perm_rescale(atomic_numbers, geometry_bohr, basis_name)
+    n = perm.shape[0]
 
     V = np.asarray(mol.intor("int2e", aosym="s1")).reshape(n, n, n, n)
     V = (
@@ -171,12 +222,4 @@ def build_repulsion_tensor_libcint(atomic_numbers: list, geometry_bohr: np.ndarr
         * rescale[None, None, None, :]
     )
 
-    ao_starts = _match_pyscf_shell_ao_starts(mol, shells)
-    perm = np.empty(n, dtype=np.int64)
-    offset = 0
-    for shell, pyscf_start in zip(shells, ao_starts):
-        block_size = len(cartesian_powers(shell.degree))
-        perm[offset : offset + block_size] = pyscf_start + _permutation_libcint_to_native_hf(shell.degree)
-        offset += block_size
-
     return V[perm][:, perm][:, :, perm][:, :, :, perm]

{dense_evolution-8.1.80 → dense_evolution-8.2.0}/dense_evolution/native_hf/scf.py

@@ -92,6 +92,8 @@ import functools
 import jax
 import jax.numpy as jnp
 
+from dense_evolution.config import ensure_x64
+
 _DIIS_DIM = 8
 
 
@@ -148,29 +150,54 @@ def _diis_extrapolate(fock_history: jax.Array, error_history: jax.Array, history
     of the augmented linear system an identity-like equation instead of
     a real DIIS constraint, so they can't contribute to the solution
     regardless of their (garbage/zero) content."""
+    dtype = fock_history.dtype
+    zero = jnp.zeros((), dtype=dtype)
+    one = jnp.ones((), dtype=dtype)
     valid = jnp.arange(diis_dim) >= (diis_dim - history_count)
 
     errs_flat = error_history.reshape(diis_dim, -1)
     B_full = errs_flat @ errs_flat.T
     mask2d = valid[:, None] & valid[None, :]
-    B = jnp.where(mask2d, B_full, 0.0)
-    B = jnp.where(jnp.eye(diis_dim, dtype=bool) & ~mask2d, 1.0, B)
+    B = jnp.where(mask2d, B_full, zero)
+    B = jnp.where(jnp.eye(diis_dim, dtype=bool) & ~mask2d, one, B)
 
-    A = jnp.zeros((diis_dim + 1, diis_dim + 1))
+    A = jnp.zeros((diis_dim + 1, diis_dim + 1), dtype=dtype)
     A = A.at[:diis_dim, :diis_dim].set(B)
-    col = jnp.where(valid, -1.0, 0.0)
+    col = jnp.where(valid, -one, zero)
     A = A.at[:diis_dim, diis_dim].set(col)
     A = A.at[diis_dim, :diis_dim].set(col)
-    b = jnp.zeros(diis_dim + 1).at[diis_dim].set(-1.0)
+    b = jnp.zeros(diis_dim + 1, dtype=dtype).at[diis_dim].set(-one)
 
     solution = jnp.linalg.solve(A, b)
     is_finite = jnp.all(jnp.isfinite(solution))
-    coeffs = jnp.where(valid, jnp.where(is_finite, solution[:diis_dim], 0.0), 0.0)
+    coeffs = jnp.where(valid, jnp.where(is_finite, solution[:diis_dim], zero), zero)
 
     F_diis = jnp.tensordot(coeffs, fock_history, axes=1)
     return jnp.where(is_finite, F_diis, fock_history[-1])
 
 
+def _level_shift_fock(F_ao: jax.Array, C_prev: jax.Array, S: jax.Array, n_occupied_pairs: int, level_shift: float) -> jax.Array:
+    """Saunders & Hillier level shifting (Int. J. Quantum Chem. 7, 699
+    (1973)): push the virtual orbitals of the PREVIOUS iteration's MO
+    basis up by `level_shift` before this iteration's diagonalization,
+    to open a numerical gap and stop the occupied/virtual split from
+    flip-flopping across a near-degeneracy (see this module's own Si2
+    docstring above -- damping/DIIS already fix the textbook case, but a
+    real, harder case (a 30-atom aromatic fragment from the CASMI26
+    wiring kernel) still took 1114 iterations, swinging through three
+    wildly different intermediate energies first). Exact no-op at
+    level_shift=0.0: since C_prev is a full, S-orthonormal basis
+    (C_prev.T @ S @ C_prev = I, hence C_prev @ C_prev.T = S^{-1}), the
+    round-trip S @ C_prev @ (C_prev.T @ F_ao @ C_prev) @ C_prev.T @ S
+    reduces algebraically to exactly F_ao before any shift is added."""
+    F_mo_prev = C_prev.T @ F_ao @ C_prev
+    n = F_mo_prev.shape[0]
+    level_shift = jnp.asarray(level_shift, dtype=F_ao.dtype)
+    shift_diag = jnp.where(jnp.arange(n) >= n_occupied_pairs, level_shift, jnp.zeros((), dtype=F_ao.dtype))
+    F_mo_prev_shifted = F_mo_prev + jnp.diag(shift_diag)
+    return S @ C_prev @ F_mo_prev_shifted @ C_prev.T @ S
+
+
 def run_scf(
     S: jax.Array,
     H_core: jax.Array,
@@ -183,7 +210,9 @@ def run_scf(
     energy_tol: float = 1e-10,
     damping: float = 0.5,
     diis_dim: int = _DIIS_DIM,
+    level_shift: float = 0.0,
 ) -> HFResult:
+    ensure_x64()
     if n_electrons % 2 != 0:
         raise ValueError("Only closed-shell (even electron count) systems are supported.")
     n_occupied_pairs = n_electrons // 2
@@ -206,7 +235,7 @@ def run_scf(
         return jnp.logical_and(jnp.logical_not(converged), iteration < max_iterations)
 
     def body_fun(state):
-        iteration, P, _C, _orbital_energies, energy_prev, _converged, fock_history, error_history = state
+        iteration, P, C_prev, _orbital_energies, energy_prev, _converged, fock_history, error_history = state
 
         F, energy = _fock_and_energy(P)
         error = _diis_error(F, P, S, X)
@@ -221,7 +250,8 @@ def run_scf(
             F,
         )
 
-        orbital_energies, C_ortho = jnp.linalg.eigh(X.T @ F_step @ X)
+        F_diag = _level_shift_fock(F_step, C_prev, S, n_occupied_pairs, level_shift)
+        orbital_energies, C_ortho = jnp.linalg.eigh(X.T @ F_diag @ X)
         C = X @ C_ortho
         P_new = _density_from_coefficients(C, n_occupied_pairs)
 
@@ -239,10 +269,10 @@ def run_scf(
         P0,
         C0,
         orbital_energies0,
-        jnp.array(jnp.inf),
+        jnp.array(jnp.inf, dtype=H_core.dtype),
         jnp.array(False),
-        jnp.zeros((diis_dim, n_basis, n_basis)),
-        jnp.zeros((diis_dim, n_basis, n_basis)),
+        jnp.zeros((diis_dim, n_basis, n_basis), dtype=H_core.dtype),
+        jnp.zeros((diis_dim, n_basis, n_basis), dtype=H_core.dtype),
     )
     iteration, P, C, orbital_energies, _energy_prev, converged, _fh, _eh = jax.lax.while_loop(cond_fun, body_fun, init_state)
 

{dense_evolution-8.1.80 → dense_evolution-8.2.0}/dense_evolution/physics/__init__.py

@@ -8,6 +8,7 @@ from .fermions import majorana_pauli_terms, total_parity_operator, hubbard_hamil
 from .qec import (pauli_commutes, compute_syndrome, erasure_aware_decode, pymatching_decode,
                    blind_minimum_weight_decode, decode_with_erasure_fallback,
                    counts_in_intervals_dimension, nearest_coset_decode)
+from .spectral import has_exact_degeneracy, matrix_function_eigh, spectral_evolve
 
 __all__ = [
     "ghz_state",
@@ -18,4 +19,5 @@ __all__ = [
     "majorana_pauli_terms", "total_parity_operator", "hubbard_hamiltonian_pauli_terms", "square_lattice_edges",
     "pauli_commutes", "compute_syndrome", "erasure_aware_decode", "pymatching_decode", "blind_minimum_weight_decode",
     "decode_with_erasure_fallback", "counts_in_intervals_dimension", "nearest_coset_decode",
+    "has_exact_degeneracy", "matrix_function_eigh", "spectral_evolve",
 ]

dense_evolution-8.2.0/dense_evolution/physics/spectral.py

@@ -0,0 +1,143 @@
+"""Gauge-safe gradients for spectral functions of a Hermitian matrix
+(V f(Lambda) V^dagger, e.g. time evolution exp(-iHt)) at exact eigenvalue
+degeneracy.
+
+`jnp.linalg.eigh`'s reverse-mode gradient divides by `lambda_i - lambda_j`
+for every eigenvector pair. When two eigenvalues are exactly degenerate,
+this does not raise and does not always produce NaN -- it can silently
+return a finite, WRONG gradient, because the eigenvectors spanning a
+degenerate eigenspace are not themselves uniquely defined (any orthonormal
+basis of that subspace is an equally valid `eigh` output). Measured on a
+real Kaggle CPU kernel (Dense-Evolution-Discovery, PR #173): std eigh
+gradient error 0.98 vs Kato 4e-10 on an H with four exact doubly-degenerate
+eigenvalues -- several orders of magnitude, not a rounding difference.
+
+REFERENCES (verified against the actual paper text, not trusted at face
+value from a citation string alone):
+    Kasim, M. F., "Derivatives of partial eigendecomposition of a real
+    symmetric matrix for degenerate cases", arXiv:2011.04366 (2020).
+    Kato, T., "Perturbation Theory for Linear Operators", Springer (1995),
+    Ch. II.5.6 (the classical divided-difference formula for matrix
+    function derivatives, predating Kasim by decades).
+
+`matrix_function_eigh` uses a `jax.custom_jvp` based on Kato's
+divided-difference formula for matrix functions:
+
+    d/deps [ V(eps) f(Lambda(eps)) V(eps)^dagger ] = V (F o (V^dagger dH V)) V^dagger
+
+with F the matrix of divided differences of f:
+
+    F[i,j] = (f(lambda_i) - f(lambda_j)) / (lambda_i - lambda_j)   if lambda_i != lambda_j
+    F[i,j] = f'(lambda_i)                                          if lambda_i == lambda_j (incl. i == j)
+
+This formula does not pass through eigenvectors as an intermediate OUTPUT,
+so it is gauge-invariant: the contribution from a degenerate block uses
+f'(lambda) directly, and there is no gauge choice to make -- unlike a
+`custom_vjp` built directly on top of `eigh`'s own eigenvector output,
+which needs a compatibility condition on the perturbation direction
+(Kasim's Eq. 4.72, confirmed present in the actual paper text) that this
+formula does not.
+
+WHEN TO USE THIS:
+    - H has (or might have) exactly degenerate eigenvalues, AND
+    - the function L(H) you differentiate depends on H through eigh, AND
+    - the function is not trivially constant on degenerate blocks.
+
+For a Hamiltonian with only near-degeneracy (e.g. min_gap ~1e-5, no exact
+tie), plain `jnp.linalg.eigh` is correct and faster -- use
+`has_exact_degeneracy` to check before reaching for `spectral_evolve`.
+"""
+import functools
+
+import jax
+import jax.numpy as jnp
+
+from dense_evolution.config import ensure_x64
+
+__all__ = ["has_exact_degeneracy", "matrix_function_eigh", "spectral_evolve"]
+
+_DEGENERACY_TOL = 1e-8
+
+
+def has_exact_degeneracy(H: jax.Array, tol: float = _DEGENERACY_TOL) -> bool:
+    """True if H has at least one pair of eigenvalues closer than tol.
+
+    Diagnostic only -- call this before choosing spectral_evolve (Kato)
+    over plain jnp.linalg.eigh (std). The threshold is the same one the
+    JVP rule below uses internally, so this is the exact condition under
+    which the two methods disagree."""
+    ensure_x64()
+    w = jnp.linalg.eigvalsh(H)
+    gaps = jnp.abs(jnp.diff(jnp.sort(w)))
+    return bool((gaps < tol).any())
+
+
+def _matrix_function_divided_differences(fw, f_prime_w, w, tol):
+    lam_i = w[:, None]
+    lam_j = w[None, :]
+    gap = lam_i - lam_j
+    is_deg = jnp.abs(gap) < tol
+    safe_gap = jnp.where(is_deg, 1.0, gap)
+    F_quot = (fw[:, None] - fw[None, :]) / safe_gap
+    F_limit = f_prime_w[:, None]
+    return jnp.where(is_deg, F_limit, F_quot)
+
+
+@functools.partial(jax.custom_jvp, nondiff_argnums=(1, 2))
+def matrix_function_eigh(H: jax.Array, f, f_prime) -> jax.Array:
+    """V f(Lambda) V^dagger with a gauge-safe gradient at exact degeneracy.
+
+    Parameters
+    ----------
+    H : (n, n) Hermitian matrix.
+    f : callable, lambda (array) -> array. Applied elementwise to the
+        eigenvalues.
+    f_prime : callable, lambda (array) -> array. Analytic derivative of f,
+        used only in the JVP rule for degenerate blocks. Passed via
+        nondiff_argnums since a Python closure is not a valid JAX type to
+        trace.
+
+    Returns
+    -------
+    (n, n) complex128 matrix.
+
+    Example
+    -------
+    >>> import jax.numpy as jnp
+    >>> H = jnp.diag(jnp.array([1.0, 1.0, 2.0, 2.0], dtype=jnp.complex128))
+    >>> U = matrix_function_eigh(H, lambda w: jnp.exp(-1j * w), lambda w: -1j * jnp.exp(-1j * w))
+    """
+    ensure_x64()
+    w, v = jnp.linalg.eigh(H)
+    return v @ jnp.diag(f(w)) @ v.conj().T
+
+
+@matrix_function_eigh.defjvp
+def _matrix_function_eigh_jvp(f, f_prime, primals, tangents):
+    (H,) = primals
+    (dH,) = tangents
+
+    w, v = jnp.linalg.eigh(H)
+    fw = f(w)
+    f_prime_w = f_prime(w)
+
+    F = _matrix_function_divided_differences(fw, f_prime_w, w, _DEGENERACY_TOL)
+
+    X = v.conj().T @ dH @ v
+    dU = v @ (F * X) @ v.conj().T
+    U = v @ jnp.diag(fw) @ v.conj().T
+    return U, dU
+
+
+def spectral_evolve(H: jax.Array, t: float) -> jax.Array:
+    """exp(-i H t) with a gauge-safe gradient at exact degeneracy.
+
+    Equivalent forward to `V @ diag(exp(-1j*w*t)) @ V.conj().T` where
+    (w, V) is `jnp.linalg.eigh(H)`. Backward uses Kato's divided-difference
+    rule (see module docstring) instead of `eigh`'s own reverse-mode rule.
+    """
+    return matrix_function_eigh(
+        H,
+        f=lambda w: jnp.exp(-1j * w * t),
+        f_prime=lambda w: -1j * t * jnp.exp(-1j * w * t),
+    )

dense_evolution-8.2.0/dense_evolution/qmmm/__init__.py

@@ -0,0 +1,13 @@
+from .forces import (
+    ATOMIC_MASSES_AMU, compute_hellmann_feynman_forces, md_step, run_md_trajectory,
+    MIN_NUCLEAR_DISTANCE_ANGSTROM, ACCEL_CONVERSION,
+)
+from .region import partition_qm_mm_region, sliced_geometry, ANGSTROM_TO_BOHR, CH_BOND_BOHR
+from .propagation import propagate_relevance
+
+__all__ = [
+    "ATOMIC_MASSES_AMU", "compute_hellmann_feynman_forces", "md_step", "run_md_trajectory",
+    "MIN_NUCLEAR_DISTANCE_ANGSTROM", "ACCEL_CONVERSION",
+    "partition_qm_mm_region", "sliced_geometry", "ANGSTROM_TO_BOHR", "CH_BOND_BOHR",
+    "propagate_relevance",
+]

dense_evolution-8.2.0/dense_evolution/qmmm/ase_bridge.py

@@ -0,0 +1,97 @@
+"""
+ASE Calculator bridge for dense_evolution.native_hf (issue #288).
+
+Motivation: QM/MM point 5's basis-flexibility test (Dense-Evolution-
+Discovery, docs/qmmm_bond_order_and_embedding.md) needed a richer basis
+(6-31G*) to test a real hypothesis about STO-3G's rigidity -- native_hf
+already supports arbitrary basis names via `build_energy_fn`'s own
+`basis_name` string parameter (nothing was actually hardcoded at the
+library level; the QM/MM scripts' own `BASIS = "sto-3g"` constants were
+just a script-level convenience). What native_hf does NOT give for free
+is interop with the wider Python computational-chemistry ecosystem (ASE
+Atoms objects, ASE's own optimizers/MD drivers, other engines' Atoms
+representations) -- this bridge is exactly that interop layer, not a new
+basis-set capability.
+
+Requires ASE (the `ase` extra: pip install dense-evolution[ase]).
+"""
+import numpy as np
+
+HARTREE_TO_EV = 27.211386245988  # CODATA, matches this project's own use of scipy.constants elsewhere (dense_evolution.qmmm.forces.ACCEL_CONVERSION)
+ANGSTROM_TO_BOHR = 1.8897259886
+
+
+def _import_ase_calculator():
+    try:
+        from ase.calculators.calculator import Calculator, all_changes
+    except ImportError as exc:
+        raise ImportError(
+            "dense_evolution.qmmm.ase_bridge needs ASE, an optional dependency "
+            "(pip install dense-evolution[ase]); it is not installed."
+        ) from exc
+    return Calculator, all_changes
+
+
+class DenseEvolutionCalculator:
+    """An ASE Calculator backed by dense_evolution.native_hf's own
+    differentiable RHF energy (build_energy_fn) -- real Obara-Saika
+    integrals and SCF, not a stub. `basis_name` is a plain string, passed
+    straight through to native_hf (e.g. "sto-3g", "6-31g*", any basis
+    basis_set_exchange has data for).
+
+    ASE units: eV, Angstrom (this class converts from native_hf's own
+    Hartree/Bohr at the boundary; everything inside native_hf itself
+    stays atomic units).
+
+    Only `energy` is implemented -- native_hf's forces come from
+    dense_evolution.qmmm.forces.compute_hellmann_feynman_forces (a
+    separate, already-real Hellmann-Feynman implementation with its own
+    finite-difference derivative and its own molecule-catalog-shaped
+    calling convention); this bridge does not reimplement or wrap that
+    here.
+
+    Examples
+    --------
+    >>> from ase import Atoms
+    >>> from dense_evolution.qmmm.ase_bridge import DenseEvolutionCalculator
+    >>> h2 = Atoms('H2', positions=[[0, 0, 0], [0, 0, 0.7414]])
+    >>> h2.calc = DenseEvolutionCalculator(atomic_numbers=[1, 1], nuclear_charges=[1.0, 1.0],
+    ...                                    n_electrons=2, basis_name="sto-3g")
+    >>> round(h2.get_potential_energy(), 2)
+    -30.39
+    """
+
+    implemented_properties = ["energy"]
+
+    def __new__(cls, atomic_numbers, nuclear_charges, n_electrons, basis_name):
+        Calculator, all_changes = _import_ase_calculator()
+
+        class _Impl(Calculator):
+            implemented_properties = ["energy"]
+
+            def __init__(self):
+                super().__init__()
+                self._atomic_numbers = atomic_numbers
+                self._nuclear_charges = nuclear_charges
+                self._n_electrons = n_electrons
+                self._basis_name = basis_name
+                self._energy_fn = None
+                self._reference_geometry_bohr = None
+
+            def calculate(self, atoms=None, properties=("energy",), system_changes=all_changes):
+                super().calculate(atoms, properties, system_changes)
+                from dense_evolution.config import ensure_x64
+                from dense_evolution.native_hf.differentiable import build_energy_fn
+                ensure_x64()
+
+                geometry_bohr = np.asarray(atoms.get_positions()) * ANGSTROM_TO_BOHR
+                if self._energy_fn is None or "numbers" in system_changes:
+                    self._reference_geometry_bohr = geometry_bohr
+                    self._energy_fn = build_energy_fn(
+                        self._atomic_numbers, self._nuclear_charges, self._n_electrons,
+                        self._basis_name, self._reference_geometry_bohr,
+                    )
+                energy_hartree = float(self._energy_fn(geometry_bohr))
+                self.results["energy"] = energy_hartree * HARTREE_TO_EV
+
+        return _Impl()