kups.potential.classical.ewald ¶

@no_jax_tracing
def estimate_ewald_parameters(
    charges: Array,
    unitcell: UnitCell,
    /,
    real_cutoff: float | None = None,
    alpha: float | None = None,
    epsilon_total: float = 1e-8,
) -> EwaldParameterEstimates:
    """Estimate optimal Ewald parameters for target accuracy.

    Not JAX-compatible (uses scipy); call before JIT compilation.
    Only works on single systems, not batched.

    Args:
        charges: Particle charges [e], shape `(n_particles,)`.
        unitcell: Unit cell parameters.
        real_cutoff: Real-space cutoff [Ang]; optimized if ``None``.
        alpha: Screening parameter [1/Ang]; optimized if ``None``.
        epsilon_total: Target total error (split equally between real/reciprocal).

    Returns:
        Optimized Ewald summation parameters.
    """
    # Note: only runs on a single system, not a batch of systems.
    # Input validation
    charges_np = np.asarray(charges)
    volume = np.asarray(unitcell.volume)
    Q2 = np.vdot(charges_np, charges_np)
    N = charges.size

    # smallest side length spanned by the unit cell
    max_radius = np.min(unitcell.perpendicular_lengths, axis=0) / 2

    # Split error budget equally
    eps_target = epsilon_total / 2

    # Length of a box shaped like the unit cell
    lattice_length = volume ** (1 / 3)

    def minimize(
        f: Callable[[float], float], bounds: tuple[float, float], n: int = 200
    ) -> float:
        attempts = np.linspace(bounds[0], bounds[1], n)  # type: ignore
        return attempts[np.argmin([f(x) for x in attempts])]

    def real_space_error(rc, alpha):
        """Standard Ewald real space error estimate"""
        if rc <= 0 or alpha <= 0:
            return np.inf
        return (Q2 / np.sqrt(N)) * (erfc(alpha * rc) / rc)

    def recip_space_error(kc, alpha):
        """Standard Ewald reciprocal space error estimate"""
        if kc <= 0 or alpha <= 0:
            return np.inf
        exp_arg = -(kc**2) / (4 * alpha**2)
        # Prevent numerical overflow/underflow
        if exp_arg < -700:  # exp(-700) ≈ 0
            return 0.0
        return (Q2 * alpha / (np.sqrt(N) * np.pi)) * np.exp(exp_arg)

    def optimal_rc(alpha):
        # Solve for rc s.t. real_space_error(rc, alpha) ≈ eps_target
        def real_error_diff(rc):
            return abs(real_space_error(rc, alpha) - eps_target)

        # Use reasonable bounds for rc based on system size
        if real_cutoff is None:
            rc_max = min(20.0, max_radius)  # Increased max rc
            rc_bounds = (0.01, rc_max)
            rc_opt = minimize(real_error_diff, rc_bounds)
        else:
            rc_opt = real_cutoff
        if real_space_error(rc_opt, alpha) > eps_target * 2:
            return np.inf
        return rc_opt

    def optimal_kc(alpha):
        """Find kc that gives the target reciprocal space error"""

        def recip_error_diff(kc):
            # Also minimize kc if all things being equal
            return abs(recip_space_error(kc, alpha) - eps_target)  # + kc * eps_target

        # Reasonable bounds for kc
        kc_bounds = (0.1, 5.0)  # Allow larger kc values
        kc_opt = minimize(recip_error_diff, kc_bounds)
        if recip_space_error(kc_opt, alpha) > eps_target * 2:
            return np.inf
        return kc_opt

    def total_cost(alpha: float) -> float:
        if alpha <= 0:
            return np.inf

        rc_opt = optimal_rc(alpha)
        if rc_opt == np.inf:
            return np.inf

        kc_opt = optimal_kc(alpha)
        if kc_opt == np.inf:
            return np.inf

        # Cost function: computational effort scales with rc^3 for real space
        # and with number of k-vectors for reciprocal space
        # Simple cost model: real space scales as rc^3, reciprocal as kc^3
        n_kvecs = (2 * kc_opt * lattice_length / 2 / np.pi + 1) ** 3
        rc_cost = rc_opt**3 * N / volume * 4 / 3 * np.pi
        cost_per_particle = rc_cost + n_kvecs
        return float(cost_per_particle)

    # Fast path without optimization
    if real_cutoff is not None:
        # First we solve for alpha given the real space cutoff by solving
        # erfc(alpha * rc) = rc * eps_target
        target_erfc = real_cutoff * eps_target
        alpha_result: scipy.optimize.OptimizeResult = scipy.optimize.minimize_scalar(  # type: ignore
            lambda z: abs(scipy.special.erfc(z) - target_erfc)
        )
        if not alpha_result.success:
            raise ValueError("Failed to find alpha for given real space cutoff.")
        alpha_opt: float = alpha_result.x / real_cutoff
        # For this alpha, we solve for kmax via the reciprocal space error estimate
        # exp(-k^2/4alpha^2) <= eps_target -> k >= 2 alpha sqrt(-ln(eps_target))
        kmax = 2 * alpha_opt * np.sqrt(-np.log(eps_target))
        return EwaldParameterEstimates(
            alpha=alpha_opt,
            real_cutoff=real_cutoff,
            k_max=kmax,
            error_real=real_space_error(real_cutoff, alpha_opt),
            error_recip=recip_space_error(kmax, alpha_opt),
            kvecs=kvecs_from_kmax(unitcell, kmax),
        )

    if alpha is None:
        alpha_opt = minimize(total_cost, (0.001, 2), n=400)
    else:
        alpha_opt = alpha
    rc_opt = optimal_rc(alpha_opt)
    kc_opt = optimal_kc(alpha_opt)

    # Verify the cutoffs are reasonable
    if rc_opt <= 0 or kc_opt <= 0 or rc_opt == np.inf or kc_opt == np.inf:
        raise ValueError("Invalid cutoff values computed")

    return EwaldParameterEstimates(
        alpha=alpha_opt,
        real_cutoff=rc_opt,
        k_max=kc_opt,
        error_real=real_space_error(rc_opt, alpha_opt),
        error_recip=recip_space_error(kc_opt, alpha_opt),
        kvecs=kvecs_from_kmax(unitcell, kc_opt),
    )

def ewald_long_range_energy[State](
    inp: EwaldLongRangeInput[State],
) -> WithPatch[Table[SystemId, Energy], Patch[State]]:
    """Reciprocal-space (long-range) Ewald energy.

    Math: ``E_lr = TO_STANDARD_UNITS * sum_k P(k) * |S(k)|^2``.

    Wraps ``structure_factor`` + ``long_range`` and returns a cache patch
    for structure factor updates on MC accept/reject.
    """
    structure_out, patch = structure_factor(inp)
    energy = long_range(inp, structure_out)
    assert energy.shape == (inp.point_cloud.batch_size,), (
        f"Expected energy shape {(inp.point_cloud.batch_size,)} but got {energy.shape}."
    )
    energy = energy * TO_STANDARD_UNITS
    return WithPatch(Table.arange(energy, label=SystemId), patch)

def ewald_self_interaction_energy(
    inp: EwaldSelfInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch]:
    """Self-interaction correction for Ewald summation.

    Removes the artificial interaction of each charge with its own Gaussian
    cloud introduced by the Ewald splitting.

    Math: ``E_self = -alpha / sqrt(pi) * sum_i q_i^2 * TO_STANDARD_UNITS``.

    Summed per system via segment_sum. Positions in Ang, charges in e,
    energy in eV.
    """
    sys_idx = Index.new(list(inp.graph.systems.keys))
    energies = (
        -jax.ops.segment_sum(
            inp.graph.particles.data.charges**2,
            inp.graph.particles.data.system.indices,
            inp.graph.batch_size,
            mode="drop",
        )
        * inp.parameters.alpha[sys_idx]
        / jnp.sqrt(jnp.pi)
    )
    energies *= TO_STANDARD_UNITS
    return WithPatch(Table.arange(energies, label=SystemId), IdPatch())

def ewald_short_range_energy(
    inp: EwaldShortRangeInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch]:
    """Real-space (short-range) screened Coulomb energy.

    Math: ``E_sr = 1/2 * TO_STANDARD_UNITS * sum_{i<j} q_i*q_j * erfc(alpha*r_ij) / r_ij``.

    The ``erfc(alpha*r)`` damping ensures convergence within the cutoff.
    Factor 1/2 corrects for double-counted pairs from the radius graph edges.
    Positions in Ang, charges in e, energy in eV.
    """
    edg = inp.graph.particles[inp.graph.edges.indices]
    qij = edg.charges[:, 0] * edg.charges[:, 1]
    dists = jnp.linalg.norm(inp.graph.edge_shifts[:, 0], axis=-1)
    edge_systems = inp.graph.edge_batch_mask
    erfc = jax.scipy.special.erfc(inp.parameters.alpha[edge_systems] * dists)
    energies = qij * erfc / dists
    mask = dists < inp.parameters.cutoff[edge_systems]
    energies *= mask
    total = inp.graph.edge_batch_mask.sum_over(energies) / 2 * TO_STANDARD_UNITS
    return WithPatch(total, IdPatch())

def exclusion_correction_energy(
    inp: EwaldShortRangeInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch]:
    """Correction for excluded pairs (e.g., bonded atoms).

    Math: ``E_excl = E_sr(excluded) - E_vacuum(excluded)``.

    Subtracts the vacuum Coulomb energy from the short-range Ewald energy
    for excluded pairs, ensuring they have zero net Coulomb interaction
    in the total Ewald sum.
    """
    return WithPatch(
        jax.tree.map(
            jnp.subtract,
            ewald_short_range_energy(inp).data,
            coulomb_vacuum_energy(inp).data,
        ),
        IdPatch(),
    )

@no_jax_tracing
def kvecs_from_kmax(unitcell: UnitCell, kmax: float) -> Array:
    """Generate integer k-vector coefficients within a sphere of radius ``kmax``.

    Args:
        unitcell: Unit cell defining the reciprocal lattice.
        kmax: Maximum k-vector magnitude cutoff.

    Returns:
        Integer k-vector coefficients, shape ``(n_kvecs, 3)``.
    """
    rvecs = unitcell.inverse_lattice_vectors.mT * 2 * jnp.pi
    min_length = jnp.min(jnp.linalg.svd(rvecs)[1])
    n = jnp.ceil(kmax / min_length).astype(int)
    lattice = (jnp.arange(0, n + 1), jnp.arange(-n, n + 1), jnp.arange(-n, n + 1))
    vecs = jnp.stack(jnp.meshgrid(*lattice), axis=-1).reshape(-1, 3)
    kvecs = einops.einsum(vecs, rvecs, "kvecs dim1, dim1 dim2 -> kvecs dim2")
    return vecs[jnp.linalg.norm(kvecs, axis=-1) <= kmax]

def long_range(inp: EwaldLongRangeInput, structure_factor: Array) -> Energy:
    """Reciprocal-space energy from structure factors.

    Math: ``E_lr = sum_k P(k) * |S(k)|^2`` where ``P(k)`` is the prefactor
    and ``S(k)`` the structure factor.
    """
    return einops.einsum(
        prefactor(inp),
        structure_factor,
        structure_factor,
        "batch_size kvecs, batch_size kvecs two, batch_size kvecs two -> batch_size",
    )

def make_ewald_from_state(
    state: Any,
    probe: Any = None,
    *,
    compute_position_and_unitcell_gradients: bool = False,
    include_exclusion_mask: bool = False,
) -> Any:
    """Create an Ewald potential from a typed state, optionally with incremental updates.

    When ``probe`` is ``None``, builds a static potential by extracting
    components directly from the state.  When a probe is provided, the
    potential supports incremental (cached) evaluation via the probe's
    patch mechanism.

    Args:
        state: Lens focusing on the Ewald state (particles, systems,
            neighborlist, and ewald_parameters).
        probe: Probe for incremental updates. ``None`` for a static
            potential.
        compute_position_and_unitcell_gradients: When ``True``, the
            returned potential computes gradients w.r.t. particle
            positions and lattice vectors (for forces / stress).
            Gradient type becomes ``PositionAndUnitCell``.
        include_exclusion_mask: Whether to include the exclusion
            correction term in the returned potential.

    Returns:
        An ``EwaldPotential`` combining short-range, long-range,
        self-energy, and (optionally) exclusion-correction terms.
        Gradient type is ``PositionAndUnitCell`` when gradients are
        requested, ``EmptyType`` otherwise.
    """
    gradient_lens: Any = EMPTY_LENS
    patch_idx_view = empty_patch_idx_view
    if compute_position_and_unitcell_gradients:
        gradient_lens = SimpleLens[PointCloud, PositionAndUnitCell](
            lambda pc: PositionAndUnitCell(
                pc.particles.map_data(lambda p: p.positions),
                pc.systems.map_data(lambda s: s.unitcell),
            )
        )
        patch_idx_view = position_and_unitcell_idx_view
    param_view = state.focus(
        lambda x: (
            x.ewald_parameters.data
            if isinstance(x.ewald_parameters, HasCache)
            else x.ewald_parameters
        )
    )
    cache_view = None
    if probe is not None:
        param_view = state.focus(lambda x: x.ewald_parameters.data)
        cache_view = state.focus(lambda x: x.ewald_parameters.cache)
    return make_ewald_potential(
        state.focus(lambda x: x.particles),
        state.focus(lambda x: x.systems),
        state.focus(lambda x: x.neighborlist),
        param_view,
        cache_view,
        probe,
        gradient_lens,
        EMPTY_LENS,
        EMPTY_LENS,
        patch_idx_view=patch_idx_view,
        include_exclusion_mask=include_exclusion_mask,
    )

def make_ewald_long_range_potential[
    State,
    Ptch: Patch,
    Gradients,
    Hessians,
](
    particles_view: View[State, Table[ParticleId, IsEwaldPointData]],
    systems_view: View[State, Table[SystemId, HasUnitCell]],
    parameter_lens: Lens[State, EwaldParameters],
    cache_lens: Lens[State, EwaldCache] | None,
    probe: Probe[State, Ptch, WithIndices[ParticleId, IsEwaldPointData]] | None = None,
    gradient_lens: Lens[
        PointCloud[IsEwaldPointData, HasUnitCell], Gradients
    ] = EMPTY_LENS,
    hessian_lens: Lens[Gradients, Hessians] = EMPTY_LENS,
    hessian_idx_view: View[State, Hessians] = EMPTY_LENS,
    patch_idx_view: View[State, PotentialOut[Gradients, Hessians]] | None = None,
) -> Potential[State, Gradients, Hessians, Any]:
    """Create the Ewald reciprocal-space (long-range) potential."""
    return PotentialFromEnergy(
        energy_fn=ewald_long_range_energy,
        composer=EwaldLongRangeComposer(
            particles=particles_view,
            systems=systems_view,
            probe=probe,
            parameters=parameter_lens,
            cache=cache_lens,
        ),
        gradient_lens=NestedLens(
            SimpleLens[EwaldLongRangeInput, PointCloud](
                lambda state: state.point_cloud
            ),
            gradient_lens,
        ),
        hessian_lens=hessian_lens,
        cache_lens=cache_lens.focus(lambda x: x.long_range) if cache_lens else None,
        hessian_idx_view=hessian_idx_view,
        patch_idx_view=patch_idx_view,
    )