dftax 0.1.0__py3-none-any.whl

dftax/__init__.py

@@ -0,0 +1,88 @@
+"""dftax: a self-contained, pure-JAX/Equinox Kohn-Sham DFT engine.
+
+The package exposes a differentiable KS-DFT toolkit with no PySCF runtime
+dependency for the core compute path:
+
+- ``dftax.integrals``: analytical one- and two-electron integral matrices
+  (overlap S, kinetic T, nuclear attraction V, nuclear repulsion, ERIs 2c/3c/4c)
+  via the Obara-Saika recurrence, all jit/vmap/grad-friendly.
+- ``dftax.energy``: GTO basis evaluation, the Boys function, density fitting,
+  Hartree and hybrid exchange, exchange-correlation functionals, real-space
+  grids, and pointwise potentials.
+- ``dftax.utils``: chunked ``vmap`` helpers and shared types.
+- ``dftax.ks``: Kohn-Sham drivers (energy functional + DIIS SCF) for both
+  closed-shell (RKS) and open-shell (UKS) systems, plus the unified ``run_ks``.
+
+A few of the most common entry points are re-exported here for convenience;
+import the submodules directly for the full surface.
+"""
+
+from dftax.energy.gto import BasisData, extract_basis_data, eval_gto
+from dftax.integrals import (
+    overlap_matrix,
+    kinetic_matrix,
+    nuclear_attraction_matrix,
+    nuclear_repulsion,
+    eri2c_matrix,
+    eri3c_matrix,
+)
+from dftax.ks import (
+    RKS, run_rks, rks_scf, rks_minimize, rks_forces, SCFResult,
+    UKS, run_uks, uks_scf, uks_minimize, uks_forces, UKSResult,
+    run_ks,
+    run_ks_batched, run_rks_batched, run_uks_batched, BatchedResult,
+    dipole, polarizability, hessian, vibrations, ir_spectrum, raman_spectrum,
+    alchemical_deriv, implicit_density,
+)
+
+from importlib.metadata import PackageNotFoundError, version as _pkg_version
+
+try:
+    __version__ = _pkg_version("dftax")
+except PackageNotFoundError:  # running from a source tree without an install
+    __version__ = "0.0.0+local"
+del _pkg_version, PackageNotFoundError
+
+__all__ = [
+    "__version__",
+    "BasisData",
+    "extract_basis_data",
+    "eval_gto",
+    "overlap_matrix",
+    "kinetic_matrix",
+    "nuclear_attraction_matrix",
+    "nuclear_repulsion",
+    "eri2c_matrix",
+    "eri3c_matrix",
+    # restricted Kohn-Sham driver
+    "RKS",
+    "run_rks",
+    "rks_scf",
+    "rks_minimize",
+    "rks_forces",
+    "SCFResult",
+    # unrestricted (open-shell) Kohn-Sham driver
+    "UKS",
+    "run_uks",
+    "uks_scf",
+    "uks_minimize",
+    "uks_forces",
+    "UKSResult",
+    # unified dispatcher
+    "run_ks",
+    # batched (vmap over geometries)
+    "run_ks_batched",
+    "run_rks_batched",
+    "run_uks_batched",
+    "BatchedResult",
+    # response properties
+    "dipole",
+    "polarizability",
+    "hessian",
+    "vibrations",
+    "ir_spectrum",
+    "raman_spectrum",
+    "alchemical_deriv",
+    # implicit-differentiation SCF response
+    "implicit_density",
+]

dftax/basis/__init__.py

@@ -0,0 +1,5 @@
+"""PySCF-free basis-set loading via Basis Set Exchange."""
+
+from dftax.basis.loader import build_basis_data, gto_norm
+
+__all__ = ["build_basis_data", "gto_norm"]

dftax/basis/loader.py

@@ -0,0 +1,160 @@
+"""Build :class:`~dftax.energy.gto.BasisData` from Basis Set Exchange data.
+
+This is the PySCF-free counterpart of ``extract_basis_data``: it fetches the
+contracted-GTO definition for each element from the ``basis_set_exchange``
+library and assembles the same normalized Cartesian-AO arrays the integral
+engine consumes. Normalization matches ``extract_basis_data`` exactly
+(primitive ``gto_norm`` times, for l<=1, a contracted norm; libcint's
+per-shell convention for l>=2).
+
+Emits Cartesian GTOs by default (``cart2sph=None``); with ``spherical=True`` it
+also builds the Cartesian->spherical transform (vendored per-l blocks under
+``data/cart2sph.npz``) so spherical d/f bases (cc-pVDZ, cc-pVTZ, def2-*) match a
+spherical PySCF reference. For l<=1 bases the two spans coincide.
+"""
+
+from __future__ import annotations
+
+import math
+from functools import lru_cache
+from importlib.resources import files
+
+import numpy as np
+import jax.numpy as jnp
+import basis_set_exchange as bse
+
+from dftax.energy.gto import BasisData, _CART_COMPONENTS, _contracted_norm
+from dftax.system.molecule import symbol_to_Z
+
+
+@lru_cache(maxsize=1)
+def _cart2sph_blocks() -> dict[int, np.ndarray]:
+    """Vendored per-l Cartesian->spherical transform blocks (ncart x nsph)."""
+    path = files("dftax.basis").joinpath("data", "cart2sph.npz")
+    with path.open("rb") as fh:
+        npz = np.load(fh)
+        return {int(k[1:]): np.asarray(npz[k], dtype=np.float64) for k in npz.files}
+
+
+def _block_diag(blocks: list[np.ndarray]) -> np.ndarray:
+    """Assemble a block-diagonal matrix from a list of 2-D blocks (no SciPy)."""
+    nr = sum(b.shape[0] for b in blocks)
+    nc = sum(b.shape[1] for b in blocks)
+    out = np.zeros((nr, nc), dtype=np.float64)
+    r = c = 0
+    for b in blocks:
+        out[r : r + b.shape[0], c : c + b.shape[1]] = b
+        r += b.shape[0]
+        c += b.shape[1]
+    return out
+
+
+def gto_norm(l: int, alpha: float) -> float:
+    """Radial normalization of a primitive GTO (matches ``pyscf.gto.gto_norm``).
+
+    ``N`` such that the primitive ``N r^l e^{-alpha r^2}`` has unit self-overlap
+    in the radial/spherical sense.
+    """
+    f = (
+        2 ** (2 * l + 3)
+        * math.factorial(l + 1)
+        * (2 * alpha) ** (l + 1.5)
+        / (math.factorial(2 * l + 2) * math.sqrt(math.pi))
+    )
+    return math.sqrt(f)
+
+
+def _iter_contractions(shell):
+    """Yield ``(l, coeff_vector)`` for each contracted function in a BSE shell.
+
+    - single angular momentum: each coefficient block is an independent
+      (general) contraction of that L;
+    - multiple angular momenta (e.g. an ``sp`` shell): block ``i`` belongs to
+      ``angular_momentum[i]``.
+    """
+    ams = shell["angular_momentum"]
+    cblocks = shell["coefficients"]
+    if len(ams) == 1:
+        for c in cblocks:
+            yield ams[0], np.asarray(c, dtype=np.float64)
+    else:
+        for l, c in zip(ams, cblocks):
+            yield l, np.asarray(c, dtype=np.float64)
+
+
+def build_basis_data(
+    symbols: list[str],
+    coords_bohr: np.ndarray,
+    basis_name: str,
+    *,
+    spherical: bool = False,
+    return_atom_index: bool = False,
+):
+    """Assemble :class:`BasisData` for a molecule from a named basis set.
+
+    Args:
+        symbols: element symbols, one per atom.
+        coords_bohr: nuclear coordinates in Bohr, shape (n_atoms, 3).
+        basis_name: a Basis Set Exchange basis name (e.g. ``"sto-3g"``).
+        spherical: if True, attach the Cartesian->spherical transform so the
+            basis uses (2l+1) spherical harmonics (standard for cc-pVXZ/def2);
+            otherwise emit Cartesian GTOs (``cart2sph=None``). For l<=1 the two
+            coincide.
+        return_atom_index: if True, also return an int array mapping each
+            Cartesian AO to its owning atom (needed to rebuild differentiable
+            centers from nuclear coordinates, e.g. for forces).
+    """
+    coords = np.asarray(coords_bohr, dtype=np.float64).reshape(-1, 3)
+    Zs = [symbol_to_Z(s) for s in symbols]
+    bdata = bse.get_basis(basis_name, elements=sorted(set(Zs)), header=False)
+    elements = bdata["elements"]
+
+    # Gather every contracted shell: (l, exponents, raw coefficients, center, atom).
+    raw_shells: list[tuple[int, np.ndarray, np.ndarray, np.ndarray, int]] = []
+    for atom_idx, (Z, center) in enumerate(zip(Zs, coords)):
+        shells = elements[str(Z)]["electron_shells"]
+        for shell in shells:
+            exps = np.asarray(shell["exponents"], dtype=np.float64)
+            for l, c_raw in _iter_contractions(shell):
+                raw_shells.append((l, exps, c_raw, center, atom_idx))
+
+    max_prim = max(len(exps) for (_, exps, _, _, _) in raw_shells)
+
+    centers, all_exps, all_coeffs, angular, atom_index = [], [], [], [], []
+    for l, exps, c_raw, center, atom_idx in raw_shells:
+        prim_norms = np.array([gto_norm(l, e) for e in exps])
+        c_prim = c_raw * prim_norms
+        for ang in _CART_COMPONENTS[l]:
+            if l <= 1:
+                c_final = c_prim * _contracted_norm(exps, c_prim, ang)
+            else:
+                c_final = c_prim.copy()
+            pad_e = np.zeros(max_prim, dtype=np.float64)
+            pad_c = np.zeros(max_prim, dtype=np.float64)
+            pad_e[: len(exps)] = exps
+            pad_c[: len(c_final)] = c_final
+            centers.append(center)
+            all_exps.append(pad_e)
+            all_coeffs.append(pad_c)
+            angular.append(ang)
+            atom_index.append(atom_idx)
+
+    cart2sph = None
+    if spherical:
+        # One cart2sph(l) block per contracted shell, in build order.
+        blocks = _cart2sph_blocks()
+        cart2sph = jnp.asarray(
+            _block_diag([blocks[l] for (l, _, _, _, _) in raw_shells])
+        )
+
+    basis = BasisData(
+        centers=jnp.asarray(np.array(centers, dtype=np.float64)),
+        exponents=jnp.asarray(np.array(all_exps, dtype=np.float64)),
+        coefficients=jnp.asarray(np.array(all_coeffs, dtype=np.float64)),
+        angular=jnp.asarray(np.array(angular, dtype=np.int32)),
+        cart2sph=cart2sph,
+        max_l=int(max(l for (l, _, _, _, _) in raw_shells)),
+    )
+    if return_atom_index:
+        return basis, np.array(atom_index, dtype=np.int64)
+    return basis

dftax/energy/__init__.py

@@ -0,0 +1 @@
+"""Energy evaluation modules."""

dftax/energy/boys.py

@@ -0,0 +1,71 @@
+"""Boys function F_n(t) in pure JAX.
+
+The Boys function is defined as:
+    F_n(t) = integral_0^1 u^{2n} exp(-t u^2) du
+
+It is related to the lower incomplete gamma function by:
+    F_n(t) = Gamma(n+0.5) * P(n+0.5, t) / (2 * t^{n+0.5})
+where P(a, x) = gammainc(a, x) is the regularized lower incomplete gamma.
+
+Key properties:
+    F_n(0) = 1 / (2n + 1)
+    dF_n/dt = -F_{n+1}(t)
+    F_0(t) = sqrt(pi) / (2*sqrt(t)) * erf(sqrt(t))  for t > 0
+"""
+
+import math
+
+import jax
+import jax.numpy as jnp
+from jax.scipy.special import gammainc, gammaln
+
+
+def boys(n: int, t: jax.Array) -> jax.Array:
+    """Boys function F_n(t) = integral_0^1 u^{2n} exp(-t u^2) du.
+
+    Accurate to better than 1e-8 vs numerical integration for n=0..10, t >= 0.
+    Fully differentiable: jax.grad(boys(n, .))(t) == -boys(n+1, t).
+    Works under jit and vmap.
+
+    Args:
+        n: Order (Python int, not a traced value).
+        t: Argument, a JAX scalar or array with t >= 0.
+
+    Returns:
+        F_n(t) as a JAX array of the same shape and dtype as t.
+    """
+    t = jnp.asarray(t)
+    dtype = t.dtype if jnp.issubdtype(t.dtype, jnp.floating) else jnp.float64
+    t = t.astype(dtype)
+
+    a = float(n) + 0.5
+
+    # --- Large-t branch: via incomplete gamma ---
+    # F_n(t) = Gamma(n+0.5) * gammainc(n+0.5, t) / (2 * t^{n+0.5})
+    # In log-space to avoid overflow:
+    #   log F_n = gammaln(a) + log(gammainc(a, t)) - a*log(t) - log(2)
+    # Use t_safe >= 1.0 to guarantee finite gradients (log, pow, gammainc).
+    # When t < 1.0 the result is discarded by jnp.where, but JAX evaluates
+    # both branches under autodiff, so 0 * NaN = NaN must be avoided.
+    safe_t = jnp.where(t < 1.0, 1.0, t)
+    log_gamma_inc = jnp.log(jnp.maximum(gammainc(a, safe_t), jnp.finfo(dtype).tiny))
+    large_t = 0.5 * jnp.exp(gammaln(a) + log_gamma_inc - a * jnp.log(safe_t))
+
+    # --- Small-t branch: Taylor series ---
+    # F_n(t) = sum_{k=0}^{K} (-t)^k / (k! * (2n + 2k + 1))
+    # The k=0 term is the constant 1/(2n+1). We separate it to avoid
+    # computing (-t)^0 = 0^0 at t=0, whose JAX gradient is NaN.
+    K = 28
+    inv_denom_0 = 1.0 / (2 * n + 1)
+    inv_denom_rest = jnp.array(
+        [1.0 / (math.factorial(k) * (2 * n + 2 * k + 1)) for k in range(1, K)],
+        dtype=dtype,
+    )
+    k_vals_rest = jnp.arange(1, K, dtype=dtype)
+    small_t = inv_denom_0 + jnp.sum(
+        inv_denom_rest * ((-t[..., None]) ** k_vals_rest), axis=-1
+    )
+
+    # Use Taylor series for t < 1 (very accurate there), incomplete gamma elsewhere.
+    # The two branches agree to machine precision at t ~ 1.
+    return jnp.where(t < 1.0, small_t, large_t)

dftax/energy/grid.py

@@ -0,0 +1,37 @@
+"""Grid-based exchange-correlation energy integration."""
+
+import jax.numpy as jnp
+from jaxtyping import Array, Float, Scalar
+
+from dftax.energy.xc import XCFunctional
+from dftax.energy.potentials import xc_potential
+
+
+def xc_energy(
+    xc: XCFunctional,
+    rho: Float[Array, "n"],              # type: ignore
+    weights: Float[Array, "n"],          # type: ignore
+    chunk_size: int | None = None,
+    grad_rho: Float[Array, "n 3"] | None = None,  # type: ignore
+    rho_thresh: float = 1e-10,
+) -> Scalar:
+    """Exchange-correlation energy ``∫ ε_xc(ρ, ∇ρ) ρ dr`` on a quadrature grid.
+
+    ``rho``/``grad_rho`` are the density (and, for GGA, its gradient) sampled at
+    the grid points; ``weights`` are the quadrature weights.
+
+    Grid points with ``ρ < rho_thresh`` are masked out with a nan-safe
+    double-``where``: the functional is evaluated on a clamped density and the
+    contribution (and its gradient) is forced to zero. This keeps the GGA
+    reduced-gradient terms, whose derivatives diverge as ρ→0, from producing
+    NaNs on the far, vanishing-density tail of an unpruned grid.
+    """
+    mask = rho > rho_thresh
+    safe_rho = jnp.where(mask, rho, 1.0)
+    if xc.xc_type == "GGA" and grad_rho is not None:
+        safe_grad = jnp.where(mask[:, None], grad_rho, 0.0)
+        eps = xc_potential(xc, safe_rho, chunk_size, grad_rho=safe_grad)
+    else:
+        eps = xc_potential(xc, safe_rho, chunk_size, grad_rho=grad_rho)
+    contrib = jnp.where(mask, weights * eps * rho, 0.0)
+    return jnp.sum(contrib)