kups.potential.classical.ewald.parameters ¶

@no_jax_tracing
def estimate_ewald_parameters(
    charges: Array,
    cell: Cell[AnyPeriodicity],
    /,
    real_cutoff: float | None = None,
    alpha: float | None = None,
    epsilon_total: float = 1e-8,
) -> EwaldParameterEstimates:
    """Estimate cutoffs and screening for one system before JIT compilation.

    Split ``epsilon_total`` equally between real and reciprocal error targets.

    Args:
        charges: Charges in e for one system, shaped ``(n_particles,)``.
        cell: Single simulation cell.
        real_cutoff: Optional fixed real-space cutoff in Å; otherwise optimized.
        alpha: Optional screening parameter in 1/Å; ignored with a fixed real cutoff.
        epsilon_total: Target combined real-space and reciprocal-space error.

    Returns:
        Screening, cutoffs and estimated errors.
    """
    charges_np = np.asarray(charges)
    net_charge = float(charges_np.astype(np.float64).sum())
    if abs(net_charge) > 1e-6:
        warnings.warn(
            f"System is not charge neutral (net charge {net_charge:.4g} e); "
            "the neutralizing-background correction will be applied.",
            stacklevel=2,
        )
    volume = np.asarray(cell.volume)
    squared_charge = np.vdot(charges_np, charges_np)
    n_particles = charges.size
    max_radius = float(np.min(cell.perpendicular_lengths, axis=0)) / 2
    target = epsilon_total / 2
    lattice_length = volume ** (1 / 3)

    def real_error(cutoff: float, alpha: float):
        if cutoff <= 0 or alpha <= 0:
            return np.inf
        return (squared_charge / np.sqrt(n_particles)) * (erfc(alpha * cutoff) / cutoff)

    def reciprocal_error(cutoff: float, alpha: float):
        if cutoff <= 0 or alpha <= 0:
            return np.inf
        exponent = -(cutoff**2) / (4 * alpha**2)
        if exponent < -700:
            return 0.0
        return (squared_charge * alpha / (np.sqrt(n_particles) * np.pi)) * np.exp(
            exponent
        )

    def cutoffs(alpha: float) -> tuple[float, float]:
        return (
            _cutoff_for_error(
                lambda r: real_error(r, alpha), (0.01, min(20.0, max_radius)), target
            ),
            _cutoff_for_error(lambda k: reciprocal_error(k, alpha), (0.1, 5.0), target),
        )

    def total_cost(alpha: float) -> float:
        if alpha <= 0:
            return np.inf
        rc, kc = cutoffs(alpha)
        if rc == np.inf or kc == np.inf:
            return np.inf
        n_kvecs = (2 * kc * lattice_length / 2 / np.pi + 1) ** 3
        real_cost = rc**3 * n_particles / volume * 4 / 3 * np.pi
        return float(real_cost + n_kvecs)

    if real_cutoff is not None:
        # erfc(alpha * cutoff) = cutoff * target fixes alpha for this cutoff.
        result = scipy.optimize.minimize_scalar(
            lambda z: abs(erfc(z) - real_cutoff * target)
        )
        if not result.success:
            raise ValueError("Failed to find alpha for given real space cutoff.")
        alpha_opt = result.x / real_cutoff
        rc_opt = real_cutoff
        kc_opt = 2 * alpha_opt * np.sqrt(-np.log(target))
    else:
        alpha_opt = (
            _minimize_on_grid(total_cost, (0.001, 2), samples=400)
            if alpha is None
            else alpha
        )
        rc_opt, kc_opt = cutoffs(alpha_opt)
        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_error(rc_opt, alpha_opt),
        error_recip=reciprocal_error(kc_opt, alpha_opt),
    )

def kvecs_from_kmax(cell: Cell[AnyPeriodicity], kmax: float) -> Array:
    """Integer half-space lattice coefficients within the reciprocal kmax sphere.

    Args:
        cell: Cell defining the reciprocal lattice.
        kmax: Maximum reciprocal-vector magnitude in 1/Å.

    Returns:
        Integer half-space lattice coefficients shaped ``(n_kvecs, 3)``.
    """
    rvecs = cell.inverse_vectors.mT * 2 * jnp.pi
    vecs = reciprocal_grid_shifts(int(reciprocal_grid_bound(cell, kmax)))
    kvecs = triangular_3x3_matmul(rvecs, vecs, lower=False)
    return vecs[jnp.sum(kvecs**2, axis=-1) <= kmax**2]

@no_jax_tracing
def prepare_reciprocal_shifts(
    cells: Cell[Periodic3D], k_max: Array, *, compact: bool
) -> tuple[Array, ReciprocalGridBound]:
    """Batch sphere-filtered shifts, padding them or enclosing them in one grid.

    Compact shifts use ``(0, 0, 0)`` as their bound. Grid shifts use the tightest
    common axis bounds for all systems, including batches with different cells.
    """
    cutoffs = np.asarray(k_max)
    if np.any(~np.isfinite(cutoffs)) or np.any(cutoffs < 0):
        raise ValueError("Reciprocal cutoffs must be finite and nonnegative")
    shifts = [kvecs_from_kmax(cells[i], k) for i, k in enumerate(cutoffs)]
    if compact:
        if not shifts:
            return jnp.zeros((0, 0, 3), dtype=int), (0, 0, 0)
        capacity = max(map(len, shifts))
        padded = jnp.stack(
            [jnp.pad(s, ((0, capacity - len(s)), (0, 0))) for s in shifts]
        )
        return padded, (0, 0, 0)

    maxima = [np.max(np.abs(s), axis=0) for s in shifts]
    bx, by, bz = np.max(maxima or [(0, 0, 0)], axis=0)
    bound = (int(bx), int(by), int(bz))
    grid = reciprocal_grid_shifts(bound)
    return jnp.broadcast_to(grid, (len(cutoffs), *grid.shape)), bound

@no_jax_tracing
def reciprocal_grid_bound(cell: Cell[AnyPeriodicity], kmax: float) -> Array:
    """Uniform coefficient bound enclosing the reciprocal kmax sphere.

    Args:
        cell: Cell defining the reciprocal lattice.
        kmax: Maximum reciprocal-vector magnitude in 1/Å.

    Returns:
        Integer scalar bound enclosing the reciprocal cutoff sphere.
    """
    rvecs = cell.inverse_vectors.mT * 2 * jnp.pi
    min_length = jnp.min(jnp.linalg.svd(rvecs)[1])
    return jnp.ceil(kmax / min_length).astype(int)

def reciprocal_grid_shifts(bound: int | ReciprocalGridBound) -> Array:
    """Integer half-space grid: nonnegative x, signed y/z, x-major.

    Args:
        bound: Nonnegative maximum absolute lattice coefficients, or one bound
            for all three axes.

    Returns:
        Integer coefficients shaped ``((bx + 1) * (2*by + 1) * (2*bz + 1), 3)``.
    """
    bx, by, bz = (bound, bound, bound) if isinstance(bound, int) else bound
    axes = (
        jnp.arange(0, bx + 1),
        jnp.arange(-by, by + 1),
        jnp.arange(-bz, bz + 1),
    )
    return jnp.stack(jnp.meshgrid(*axes, indexing="ij"), axis=-1).reshape(-1, 3)