off 0.1.0__py3-none-any.whl

off/functionals/core_correction.py

@@ -0,0 +1,84 @@
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Array, Float
+from jax import vmap, lax
+
+
+class KatoCondition(eqx.Module):
+    r"""
+    Kato cusp-condition functional.
+
+    Penalizes deviations from the correct electron-nucleus cusp behaviour, weighting
+    the (Weizsacker-like) local kinetic term by a 1s-like envelope
+    w_i(r) = exp(-a Z_i |r - R_i|) around each nucleus.
+
+    Parameters
+    ----------
+    a : float, optional
+        Envelope decay prefactor, by default 2/3.
+    eps : float, optional
+        Small constant for numerical stability, by default 1e-5.
+    """
+
+    a: float
+    eps: float
+
+    def __init__(self, a=2.0/3.0, eps=1e-5):
+        self.a = a
+        self.eps = eps
+
+    def __call__(self, den, score, x, Ne, mol, xp) -> Float[Array, "batch"]:
+        r"""Uses x, score, Ne, mol; den, xp unused. Returns the cusp penalty per point."""
+        def _wi(pts, molecule):
+            r = jnp.sqrt(jnp.sum((pts - molecule['coords']) ** 2, axis=1) + self.eps ** 2)
+            return jnp.exp(-self.a * molecule['z'] * r)
+
+        wi = vmap(_wi, in_axes=(None, 0), out_axes=-1)(x, mol)
+        score_sqr = jnp.einsum('ij,ij->i', score, score)
+        weizs = (1.0 / 8.0) * lax.expand_dims(score_sqr, (1,))
+        kinetic = jnp.abs(weizs - (mol['z'] ** 2 / 2))
+        return Ne * jnp.sum(kinetic * wi, axis=-1, keepdims=True)
+
+
+class HutcheonCuspCondition(eqx.Module):
+    r"""
+    Exact nuclear-cusp cost functional.
+
+    From Hutcheon & Wibowo-Teale, Phys. Rev. B 110, 195146 (2024), Eqs. (4-5, 7):
+
+    E_{\text{CC}} = \sum_i C_i, \qquad
+    C_i(\rho) = \int |\nabla \rho + 2 Z_i \rho\, \hat{r}_i|^2\, w_i(r)\, d^3 r,
+
+    with  \hat{r}_i = (r - R_i)/|r - R_i|  and  w_i(r) = (\pi / Z_i^3)^{1/2} e^{-Z_i |r - R_i|}  [Eq. (5)].
+    Since \nabla \rho = \rho\, s (s the score), the integrand is \rho^2 |s + 2 Z_i \hat{r}_i|^2,
+    and the Kato cusp is satisfied exactly when s = -2 Z_i \hat{r}_i at R_i. With samples
+    from p = \rho / N_e,
+
+    C_i \approx N_e^2\, \mathbb{E}\left[ \rho(x)\, |s(x) + 2 Z_i \hat{r}_i(x)|^2\, w_i(x) \right].
+
+    Parameters
+    ----------
+    eps : float, optional
+        Small constant for numerical stability, by default 1e-8.
+    """
+
+    eps: float
+
+    def __init__(self, eps=1e-8):
+        self.eps = eps
+
+    def __call__(self, den, score, x, Ne, mol, xp) -> Float[Array, "batch 1"]:
+        r"""Uses x, den, score, Ne, mol; xp unused. Returns the cusp cost per point."""
+        coords, z = mol['coords'], mol['z']
+
+        def per_atom(R_i, Z_i):
+            r_vec = x - R_i
+            r_norm = jnp.sqrt(jnp.sum(r_vec ** 2, axis=-1) + self.eps ** 2)
+            r_hat = r_vec / r_norm[:, None]
+            w_i = jnp.sqrt(jnp.pi / Z_i ** 3) * jnp.exp(-Z_i * r_norm)
+            cusp_sq = jnp.sum((score + 2.0 * Z_i * r_hat) ** 2, axis=-1)
+            return cusp_sq * w_i
+
+        total = jnp.sum(jax.vmap(per_atom)(coords, z), axis=0)
+        return (Ne * den.squeeze(-1) * total)[:, None]

off/functionals/exchange_correlation.py

@@ -0,0 +1,174 @@
+import jax.numpy as jnp
+from jax import lax
+from .functional import Functional, CompositeFunctional, unit_coefficient
+
+
+def lda_density(den, score, x, Ne, mol, xp):
+    r"""
+    Local density approximation (LDA) / Dirac exchange functional.
+
+    See eq. 2.72 in "Time-Dependent Density-Functional Theory", Carsten A. Ullrich.
+
+    E_{\text{X}}^{\text{LDA}}[\rho] = -\frac{3}{4}\left(\frac{3}{\pi}\right)^{1/3} \int \rho(\boldsymbol{x})^{4/3} d\boldsymbol{x}
+                                    = -\frac{3}{4}\left(\frac{3}{\pi}\right)^{1/3} N_e^{4/3}\, \mathbb{E}_{\rho_\phi}\left[ \rho_\phi(\boldsymbol{x})^{1/3} \right]
+
+    Parameters
+    ----------
+    den : Array
+        Density.
+    Ne : int
+        Number of electrons.
+
+    Notes
+    -----
+    score, x, mol, xp are accepted for the shared functional interface but unused here.
+
+    Returns
+    -------
+    jax.Array
+        LDA exchange energy density (up to the rho factor).
+    """
+    l = -(3 / 4) * (Ne ** (4 / 3)) * (3 / jnp.pi) ** (1 / 3)
+    return l * den ** (1 / 3)
+
+
+def b88_density(den, score, x, Ne, mol, xp, clip_cte=1e-30, beta=0.0042):
+    r"""
+    B88 exchange functional.
+
+    See eq. 8 in https://journals.aps.org/pra/abstract/10.1103/PhysRevA.38.3098
+    See also https://github.com/ElectronicStructureLibrary/libxc/blob/master/maple/gga_exc/gga_x_b88.mpl
+
+    E_{\text{X}}^{\text{B88}}[\rho] = -\beta \int \frac{X^2}{1 + 6\beta X \sinh^{-1}(X)} \rho(\boldsymbol{x})^{4/3} d\boldsymbol{x},
+    \qquad X = \frac{|\nabla \rho|}{\rho^{4/3}}.
+
+    Parameters
+    ----------
+    den : Array
+        Density.
+    score : Array
+        Gradient of the log-density, s = (nabla rho)/rho.
+    Ne : int
+        Number of electrons.
+    clip_cte : float, optional
+        Small constant for numerical stability, by default 1e-30.
+    beta : float, optional
+        B88 parameter, by default 0.0042.
+
+    Notes
+    -----
+    x, mol, xp are accepted for the shared functional interface but unused here.
+
+    Returns
+    -------
+    jax.Array
+        B88 exchange energy density (up to the rho factor).
+    """
+    den_clipped = jnp.clip(den, clip_cte)
+    log_den = jnp.log2(den_clipped)
+    score_sqr = jnp.einsum('ij,ij->i', score, score)
+    grad_den_norm_sq = lax.expand_dims(score_sqr, (1,)) * den_clipped * den_clipped
+    log_grad_den_norm = jnp.log2(jnp.clip(grad_den_norm_sq, clip_cte)) / 2
+    log_x_sigma = log_grad_den_norm - 4 / 3.0 * log_den
+    x_sigma = 2 ** log_x_sigma
+    b88_e = -(beta * 2 ** (4 * log_den / 3 + 2 * log_x_sigma
+                           - jnp.log2(1 + 6 * beta * x_sigma * jnp.arcsinh(x_sigma))))
+    return b88_e * Ne ** (2 / 3)
+
+
+def vwn_density(den, score, x, Ne, mol, xp, clip_cte=1e-30):
+    r"""
+    VWN correlation functional.
+
+    See original paper eq. 4.4 in https://cdnsciencepub.com/doi/abs/10.1139/p80-159
+    See also the text after eq. 8.9.6.1 in https://www.theoretical-physics.com/dev/quantum/dft.html
+
+    \epsilon_{\text{C}}^{\text{VWN}} = \frac{A}{2}\left\{ \ln\frac{y^2}{Y(y)}
+        + \frac{2b}{Q}\arctan\frac{Q}{2y+b}
+        - \frac{b y_0}{Y(y_0)}\left[ \ln\frac{(y-y_0)^2}{Y(y)}
+        + \frac{2(b+2y_0)}{Q}\arctan\frac{Q}{2y+b}\right] \right\}
+
+    Parameters
+    ----------
+    den : Array
+        Density.
+    Ne : int
+        Number of electrons.
+    clip_cte : float, optional
+        Small constant for numerical stability, by default 1e-30.
+
+    Notes
+    -----
+    score, x, mol, xp are accepted for the shared functional interface but unused here.
+
+    Returns
+    -------
+    jax.Array
+        VWN correlation energy density (up to the rho factor).
+    """
+    A, b, c, x0 = 0.0621814, 3.72744, 12.9352, -0.10498
+    den_clipped = jnp.where(den > clip_cte, den, 0.0)
+    log_den = jnp.log2(jnp.clip(den_clipped, clip_cte))
+    log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_den / 3.0
+    log_x = log_rs / 2
+    x_ = 2. ** log_x
+    X = 2. ** (2. * log_x) + 2. ** (log_x + jnp.log2(b)) + c
+    X0 = x0 ** 2 + b * x0 + c
+    Q = jnp.sqrt(4 * c - b ** 2)
+    e_PF = A / 2. * (
+        2. * jnp.log(x_) - jnp.log(X) + 2. * b / Q * jnp.arctan(Q / (2. * x_ + b))
+        - b * x0 / X0 * (jnp.log((x_ - x0) ** 2. / X)
+                         + 2. * (2. * x0 + b) / Q * jnp.arctan(Q / (2. * x_ + b)))
+    )
+    return Ne * e_PF
+
+
+def pw92_density(den, score, x, Ne, mol, xp, clip_cte=1e-30):
+    r"""
+    PW92 correlation functional.
+
+    See eq. 10 in https://journals.aps.org/prb/abstract/10.1103/PhysRevB.45.13244
+
+    \epsilon_{\text{C}}^{\text{PW92}} = -2A(1 + \alpha_1 r_s)
+        \ln\left[ 1 + \frac{1}{2A(\beta_1 r_s^{1/2} + \beta_2 r_s + \beta_3 r_s^{3/2} + \beta_4 r_s^2)} \right]
+
+    Parameters
+    ----------
+    den : Array
+        Density.
+    Ne : int
+        Number of electrons.
+    clip_cte : float, optional
+        Small constant for numerical stability, by default 1e-30.
+
+    Notes
+    -----
+    score, x, mol, xp are accepted for the shared functional interface but unused here.
+
+    Returns
+    -------
+    jax.Array
+        PW92 correlation energy density (up to the rho factor).
+    """
+    A_, alpha1 = 0.031091, 0.21370
+    beta1, beta2, beta3, beta4 = 7.5957, 3.5876, 1.6382, 0.49294
+    log_den = jnp.log2(jnp.clip(den, clip_cte))
+    log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_den / 3.0
+    brs_1_2 = 2 ** (log_rs / 2 + jnp.log2(beta1))
+    ars = 2 ** (log_rs + jnp.log2(alpha1))
+    brs = 2 ** (log_rs + jnp.log2(beta2))
+    brs_3_2 = 2 ** (3 * log_rs / 2 + jnp.log2(beta3))
+    brs2 = 2 ** (2 * log_rs + jnp.log2(beta4))
+    e_PF = -2 * A_ * (1 + ars) * jnp.log(1 + (1 / (2 * A_)) / (brs_1_2 + brs + brs_3_2 + brs2))
+    return Ne * e_PF
+
+
+lda  = Functional(coefficients=unit_coefficient, energy_densities=lda_density)
+b88  = Functional(coefficients=unit_coefficient, energy_densities=b88_density)
+vwn  = Functional(coefficients=unit_coefficient, energy_densities=vwn_density)
+pw92 = Functional(coefficients=unit_coefficient, energy_densities=pw92_density)
+
+
+def lda_b88():
+    r"""LDA + B88 exchange."""
+    return CompositeFunctional(lda, b88)

off/functionals/external.py

@@ -0,0 +1,49 @@
+import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Array, Float
+from jax import vmap
+
+
+class NuclearPotential(eqx.Module):
+    r"""
+    External electron-nuclei attraction potential.
+
+    V_{\text{ext}}(\boldsymbol{x}) = -N_e \sum_i \frac{Z_i}{|\boldsymbol{x} - \boldsymbol{R}_i|}
+
+    Parameters
+    ----------
+    eps : float, optional
+        Small constant for numerical stability, by default 1e-5.
+    """
+
+    eps: float
+
+    def __init__(self, eps=1e-5):
+        self.eps = eps
+
+    def __call__(self, den, score, x, Ne, mol, xp) -> Float[Array, "batch"]:
+        r"""
+        Parameters
+        ----------
+        x : Array
+            Points where the potential is evaluated.
+        Ne : int
+            Number of electrons.
+        mol : dict
+            Nuclear coordinates and charges, {'coords': ..., 'z': ...}.
+
+        Notes
+        -----
+        den, score, xp are accepted for the shared functional interface but unused here.
+
+        Returns
+        -------
+        jax.Array
+            Electron-nuclei attraction at each point (up to the rho factor).
+        """
+        def _potential(pts, molecule):
+            r = jnp.sqrt(jnp.sum((pts - molecule['coords']) ** 2, axis=1)) + self.eps
+            return molecule['z'] / r
+
+        r = vmap(_potential, in_axes=(None, 0), out_axes=-1)(x, mol)
+        return -Ne * jnp.sum(r, axis=-1, keepdims=True)

off/functionals/functional.py

@@ -0,0 +1,129 @@
+import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Array, Float
+from typing import Any, Callable, NamedTuple, Optional
+
+
+class FunctionalInputs(NamedTuple):
+    r"""
+    Container holding every input a functional might need.
+
+    It is the single bundle handed to :class:`EnergyFunctional`, which then
+    *separates* it into explicit arguments for each component functional. The leaf
+    functionals therefore receive plain ``(den, score, x, Ne, mol, xp)`` and never
+    see this object.
+
+    Fields
+    ------
+    den   : Float[Array, "batch 1"]   density shape factor rho_phi.
+    score : Float[Array, "batch d"]   score, (nabla rho)/rho = nabla log rho_phi.
+    x     : Float[Array, "batch d"]   sample / grid positions.
+    Ne    : int                       number of electrons.
+    mol   : dict                      {'coords': ..., 'z': ...} nuclear geometry / charges.
+    xp    : Float[Array, "batch d"]   second set of positions for the pairwise Hartree.
+    """
+    den:   Any
+    score: Any
+    x:     Any
+    Ne:    Any
+    mol:   Any
+    xp:    Any = None
+
+
+def unit_coefficient(self, *_):
+    r"""Constant unit weight (c = 1): turns a Functional into a fixed (non-learned) functional."""
+    return jnp.array([[1.0]])
+
+
+class Functional(eqx.Module):
+    r"""
+    Local density functional,  F[\rho] = \int c_\theta[\rho] \cdot e[\rho]\, d\boldsymbol{x}.
+
+    Every functional shares the explicit signature ``(den, score, x, Ne, mol, xp)``
+    -- all inputs are passed even if a given functional ignores some -- and returns
+    its per-point energy density of shape (batch, 1). Assembled from two callables:
+
+    energy_densities(den, score, x, Ne, mol, xp) -> e[rho]
+        the energy densities, returned up to the rho factor.
+    coefficients(self, cinputs) -> c_theta[rho]
+        the weights; :func:`unit_coefficient` (c = 1) for a fixed functional, a
+        network for a learned one.
+    coefficient_inputs(den, score, x, Ne, mol, xp) -> features   (optional)
+        features fed to ``coefficients`` (learned functionals only).
+    """
+
+    coefficients:       Callable
+    energy_densities:   Callable
+    coefficient_inputs: Optional[Callable] = None
+    
+    def __call__(self, den, score, x, Ne, mol, xp) -> Float[Array, "batch 1"]:
+        e  = self.energy_densities(den, score, x, Ne, mol, xp)
+        ci = (self.coefficient_inputs(den, score, x, Ne, mol, xp)
+              if self.coefficient_inputs is not None else None)
+        c  = self.coefficients(self, ci)
+        return jnp.sum(c * e, axis=-1, keepdims=True)
+    #Add a functional, or density. 
+
+class CompositeFunctional(eqx.Module):
+    r"""Sum of several functionals, all sharing the ``(den, score, x, Ne, mol, xp)`` signature."""
+
+    functionals: list
+
+    def __init__(self, *functionals):
+        self.functionals = functionals
+
+    def __call__(self, den, score, x, Ne, mol, xp) -> Float[Array, "batch 1"]:
+        result = 0.0
+        for func in self.functionals:
+            result = result + func(den, score, x, Ne, mol, xp)
+        return result
+
+    def __add__(self, other):
+        if isinstance(other, CompositeFunctional):
+            return CompositeFunctional(self.functionals + other.functionals)
+        return CompositeFunctional(self.functionals + [other])
+
+
+class EnergyFunctional(eqx.Module):
+    r"""
+    High-level OF-DFT energy functional.
+
+    Receives the single :class:`FunctionalInputs` bundle and *separates* it into
+    explicit arguments for each component functional (kinetic, external/nuclear,
+    Hartree, exchange, and optionally correlation and core-correction): every
+    component is called with the same ``(den, score, x, Ne, mol, xp)`` and uses only
+    what it needs. ``correlation`` and ``core_correction`` may be ``None``.
+
+    ``terms`` returns the per-component energy densities (convenient for logging);
+    ``__call__`` returns their sum. Integrate with :meth:`_integrate` (Monte-Carlo
+    measure 1/N during training, or grid weights w*rho for quadrature).
+    """
+
+    kinetic:         Any
+    external:        Any
+    hartree:         Any
+    exchange:        Any
+    correlation:     Any = None
+    core_correction: Any = None
+
+    @staticmethod
+    def _integrate(energy_density, weights):
+        r"""Quadrature: integral ~ sum_i w_i e_i. `weights` is 1/N (Monte-Carlo) or w_i*rho_i (grid)."""
+        return jnp.sum(weights * energy_density)
+
+    def terms(self, inp: FunctionalInputs) -> dict:
+        r"""Separate ``inp`` and evaluate every component on the same explicit arguments."""
+        a = (inp.den, inp.score, inp.x, inp.Ne, inp.mol, inp.xp)
+        return {
+            "kin":  self.kinetic(*a),
+            "vnuc": self.external(*a),
+            "hart": self.hartree(*a),
+            "x":    self.exchange(*a),
+            "c":    self.correlation(*a) if self.correlation is not None else 0.0,
+            "cc":   self.core_correction(*a) if self.core_correction is not None else 0.0,
+        }
+
+    def __call__(self, inp: FunctionalInputs):
+        r"""Per-point total energy density (sum of all components)."""
+        t = self.terms(inp)
+        return t["kin"] + t["vnuc"] + t["hart"] + t["x"] + t["c"] + t["cc"]

off/functionals/hartree.py

@@ -0,0 +1,62 @@
+import jax.numpy as jnp
+import equinox as eqx
+from jaxtyping import Array, Float
+
+
+class CoulombPotential_(eqx.Module):
+    r"""
+    Classical electron-electron repulsion (Hartree) potential, element-wise pairs.
+
+    Pairs each x_i with x'_i (the two half-batches), giving `batch` Monte-Carlo
+    pairs:
+
+    V_{\text{Hartree}}(\boldsymbol{x}, \boldsymbol{x}') = \frac{1}{2} N_e^2 \frac{1}{|\boldsymbol{x} - \boldsymbol{x}'|}
+
+    Parameters
+    ----------
+    eps : float, optional
+        Small constant for numerical stability, by default 1e-5.
+    """
+
+    eps: float
+
+    def __init__(self, eps=1e-5):
+        self.eps = eps
+
+    def __call__(self, den, score, x, Ne, mol, xp) -> Float[Array, "batch 1"]:
+        r"""Uses x, xp (paired points) and Ne; den, score, mol unused. Returns ½Ne²/|x-x'| per pair."""
+        z = jnp.sum((x - xp) ** 2 + self.eps, axis=-1, keepdims=True)
+        return 0.5 * (Ne ** 2) / jnp.sqrt(z)
+
+
+class CoulombPotential(eqx.Module):
+    r"""
+    Classical electron-electron repulsion (Hartree) potential, all-pairs estimator.
+
+    Same physics as :class:`CoulombPotential_` (true 1/|x-x'| Coulomb), but averages
+    each x_i against *every* x'_j (the full batch x batch double sum), i.e. batch^2
+    pairs instead of batch. This is the same estimator the grid quadrature uses and
+    has substantially lower Monte-Carlo variance for the same samples, at the cost
+    of a (batch, batch) distance matrix.
+
+    V_{\text{Hartree}} = \frac{1}{2} N_e^2 \left\langle \frac{1}{|\boldsymbol{x}_i - \boldsymbol{x}'_j|} \right\rangle_j
+
+    Parameters
+    ----------
+    eps : float, optional
+        Small constant for numerical stability, by default 1e-5.
+    """
+
+    eps: float
+
+    def __init__(self, eps=1e-5):
+        self.eps = eps
+
+    def __call__(self, den, score, x, Ne, mol, xp) -> Float[Array, "batch 1"]:
+        r"""Uses x, xp and Ne; den, score, mol unused. Returns the per-point Hartree potential."""
+        x2 = jnp.sum(x * x, axis=-1)
+        xp2 = jnp.sum(xp * xp, axis=-1)
+        r2 = x2[:, None] + xp2[None, :] - 2.0 * (x @ xp.T)
+        r2 = jnp.maximum(r2, 0.0) + self.eps
+        v = jnp.mean(1.0 / jnp.sqrt(r2), axis=-1, keepdims=True)
+        return 0.5 * (Ne ** 2) * v

off/functionals/kinetic.py

@@ -0,0 +1,87 @@
+import jax.numpy as jnp
+from jax import lax
+from .functional import Functional, CompositeFunctional, unit_coefficient
+
+C_TF = (3. / 10.) * (3. * jnp.pi ** 2) ** (2 / 3)
+
+
+def tf_density(den, score, x, Ne, mol, xp, c=C_TF):
+    r"""
+    Thomas-Fermi kinetic functional.
+
+    See paper eq. 2 in https://pubs.aip.org/aip/jcp/article/114/2/631/184186/Thomas-Fermi-Dirac-von-Weizsacker-models-in-finite
+
+    T_{\text{TF}}[\rho] = c \int \rho(\boldsymbol{x})^{5/3} d\boldsymbol{x}
+                        = c \int \rho(\boldsymbol{x})^{2/3} \rho(\boldsymbol{x}) d\boldsymbol{x}
+    T_{\text{TF}}[\rho] = c\, N_e^{5/3}\, \mathbb{E}_{\rho_\phi}\left[ \rho_\phi(\boldsymbol{x})^{2/3} \right]
+
+    with c = \frac{3}{10}(3\pi^2)^{2/3}.
+
+    Parameters
+    ----------
+    den : Array
+        Density.
+    Ne : int
+        Number of electrons.
+    c : float, optional
+        Prefactor, by default (3/10)(3 pi^2)^(2/3).
+
+    Notes
+    -----
+    score, x, mol, xp are accepted for the shared functional interface but unused here.
+
+    Returns
+    -------
+    jax.Array
+        Thomas-Fermi kinetic energy density (up to the rho factor).
+    """
+    return c * (Ne ** (5 / 3)) * den ** (2 / 3)
+
+
+def weizsacker_density(den, score, x, Ne, mol, xp, lam=0.2):
+    r"""
+    von Weizsacker gradient correction.
+
+    See paper eq. 3 in https://pubs.aip.org/aip/jcp/article/114/2/631/184186/Thomas-Fermi-Dirac-von-Weizsacker-models-in-finite
+
+    T_{\text{Weizsacker}}[\rho] = \frac{\lambda}{8} \int \frac{(\nabla \rho)^2}{\rho} d\boldsymbol{x}
+                                = \frac{\lambda}{8} \int \rho \left(\frac{\nabla \rho}{\rho}\right)^2 d\boldsymbol{x}
+    T_{\text{Weizsacker}}[\rho] = \frac{\lambda N_e}{8}\, \mathbb{E}_{\rho_\phi}\left[ \left(\frac{\nabla \rho}{\rho}\right)^2 \right]
+
+    Parameters
+    ----------
+    score : Array
+        Gradient of the log-density, s = (nabla rho)/rho.
+    Ne : int
+        Number of electrons.
+    lam : float, optional (W. Stich, E.K.U. Gross, Z. Physik A 309(1):511, 1982)
+        Phenomenological parameter lambda, by default 0.2.
+
+    Notes
+    -----
+    den, x, mol, xp are accepted for the shared functional interface but unused here.
+
+    Returns
+    -------
+    jax.Array
+        von Weizsacker kinetic energy density (up to the rho factor).
+    """
+    score_sqr = jnp.einsum('ij,ij->i', score, score)
+    return (lam * Ne / 8.) * lax.expand_dims(score_sqr, (1,))
+
+
+tf = Functional(coefficients=unit_coefficient, energy_densities=tf_density)
+
+
+def weizsacker(lam=0.2):
+    r"""von Weizsacker functional with prefactor ``lam`` (see :func:`weizsacker_density`)."""
+    return Functional(
+        coefficients=unit_coefficient,
+        energy_densities=lambda den, score, x, Ne, mol, xp: weizsacker_density(
+            den, score, x, Ne, mol, xp, lam),
+    )
+
+
+def tf_weizsacker(lam=0.2):
+    r"""TF-lambda-W kinetic functional:  T = T_TF + lambda * T_W."""
+    return CompositeFunctional(tf, weizsacker(lam))