kups.potential.classical.ewald ¶

@dataclass
class EwaldCache[Gradient, Hessian]:
    """Compensated structure factors, prefactors and component outputs for updates."""

    structure_factor: KahanSummand[Array]  # (n_groups, n_kvecs, 2)
    prefactor: Array  # (n_groups, n_kvecs)
    short_range: KahanSummand[PotentialOut[Gradient, Hessian]]
    long_range: KahanSummand[PotentialOut[Gradient, Hessian]]
    self_interaction: KahanSummand[PotentialOut[Gradient, Hessian]]
    net_charge: KahanSummand[PotentialOut[Gradient, Hessian]]
    exclusion: KahanSummand[PotentialOut[Gradient, Hessian]]

    @classmethod
    @no_jax_tracing
    def make[G, H](
        cls, n_sys: int, n_kvecs: int, gradient: G = EMPTY, hessian: H = EMPTY
    ) -> EwaldCache[G, H]:
        """Create zeroed caches with the supplied derivative shapes.

        Args:
            n_sys: Number of systems.
            n_kvecs: Number of reciprocal vectors per system, including padding.
            gradient: Template defining cached gradient shapes.
            hessian: Template defining cached Hessian shapes.

        Returns:
            Structure-factor, prefactor and component-output caches set to zero.
        """
        out = PotentialOut(
            Table.arange(jnp.zeros(n_sys, dtype=float), label=SystemId),
            gradient,
            hessian,
        )
        return EwaldCache(
            KahanSummand.init(jnp.zeros((n_sys, n_kvecs, 2), dtype=float)),
            jnp.zeros((n_sys, n_kvecs), dtype=float),
            KahanSummand.init(tree_zeros_like(out)),
            KahanSummand.init(tree_zeros_like(out)),
            KahanSummand.init(tree_zeros_like(out)),
            KahanSummand.init(tree_zeros_like(out)),
            KahanSummand.init(tree_zeros_like(out)),
        )

@dataclass
class EwaldCacheDeltaPatch[State, Gradient, Hessian](Patch[State]):
    """Add each accepted system's structure-factor change to the compensated cache."""

    structure_factor_delta: Array
    system_idx: Index[SystemId]
    lens: Lens[State, EwaldCache[Gradient, Hessian]] = field(static=True)

    def __call__(self, state: State, accept: Accept) -> State:
        """Apply structure-factor updates for accepted systems.

        Args:
            state: State containing the Ewald cache.
            accept: Per-system acceptance mask.

        Returns:
            State with accepted structure factors and their compensation updated.
        """
        mask = accept[self.system_idx]
        delta = self.structure_factor_delta
        return self.lens.focus(lambda cache: cache.structure_factor).apply(
            state,
            lambda old_sf: jax.tree.map(
                lambda new, old: where_broadcast_last(mask, new, old),
                old_sf + delta,
                old_sf,
            ),
        )

@dataclass
class EwaldCachePatch[State, Gradient, Hessian](Patch[State]):
    """Accept or reject each system's structure factor and prefactor."""

    new_structure_factor: KahanSummand[Array]
    new_prefactor: Array
    system_idx: Index[SystemId]
    lens: Lens[State, EwaldCache[Gradient, Hessian]] = field(static=True)

    def __call__(self, state: State, accept: Accept) -> State:
        """Apply structure-factor updates for accepted systems.

        Args:
            state: State containing the Ewald cache.
            accept: Per-system acceptance mask.

        Returns:
            State with accepted structure factors, their compensation and
            prefactors updated.
        """
        mask = accept[self.system_idx]
        update = (self.new_structure_factor, self.new_prefactor)
        return self.lens.focus(lambda c: (c.structure_factor, c.prefactor)).apply(
            state,
            lambda current: jax.tree.map(
                lambda new, old: where_broadcast_last(mask, new, old), update, current
            ),
        )

@dataclass
class EwaldLongRangeComposer[
    State,
    Ptch: Patch[Any],
]:
    """Build current particles and retain previous values for incremental updates."""

    particles: View[State, Table[ParticleId, IsEwaldPointData]] = field(static=True)
    systems: View[State, Table[SystemId, HasCell[Periodic3D]]] = field(static=True)
    probe: Probe[State, Ptch, WithIndices[ParticleId, IsEwaldPointData]] | None = field(
        static=True
    )
    parameters: View[State, EwaldParameters] = field(static=True)
    cache: Lens[State, EwaldCache[Any, Any]] | None = field(static=True)

    def __call__(
        self, state: State, patch: Ptch | None
    ) -> Sum[EwaldLongRangeInput[State]]:
        """Compose the reciprocal-energy input for a full evaluation or proposal.

        Args:
            state: State providing particles, systems, parameters and optional cache.
            patch: Optional proposed state patch to inspect with the particle probe.

        Returns:
            A single-summand input with current particles and any probed previous values.
        """
        particles = self.particles(state)
        changes = None
        if patch is not None and self.probe is not None:
            updates = self.probe(state, patch)
            changes = WithIndices(
                updates.indices, (particles[updates.indices], updates.data)
            )
            # Only derivatives read the patched particles; energy-only
            # evaluation leaves the update to dead-code elimination.
            particles = particles.update(updates.indices, updates.data)
        inp = EwaldLongRangeInput(
            PointCloud(particles, self.systems(state)),
            self.parameters(state),
            self.cache.get(state) if self.cache else None,
            self.cache,
            changes,
        )
        return Sum(Summand(inp))

@dataclass
class EwaldLongRangeInput[State]:
    """Current particles and optional changed rows for a cached update.

    ``changes_from_prev`` holds the previous and current values of the changed
    rows; energy-only evaluation reads only these, derivatives ``point_cloud``.
    """

    point_cloud: PointCloud[IsEwaldPointData, HasCell[Periodic3D]]
    parameters: EwaldParameters
    cache: EwaldCache[Any, Any] | None = None
    cache_lens: Lens[State, EwaldCache[Any, Any]] | None = None
    changes_from_prev: (
        WithIndices[ParticleId, tuple[IsEwaldPointData, IsEwaldPointData]] | None
    ) = None

    @property
    def volume(self) -> Array:
        """Cell volumes in ų, shaped ``(n_systems,)``."""
        return self.point_cloud.systems.data.cell.volume

    @property
    def kvecs(self) -> Array:
        """Convert the stored integer shifts to Cartesian reciprocal vectors.

        The result is in 1/Å, shaped ``(n_systems, n_kvecs, 3)``.
        """
        sys_idx = self.point_cloud.systems.index
        return triangular_3x3_matmul(
            self.point_cloud.systems.data.cell.inverse_vectors.mT[:, None] * 2 * jnp.pi,
            self.parameters.reciprocal_lattice_shifts[sys_idx],
            lower=False,
        )

@dataclass
class EwaldParameterEstimates:
    """Estimated optimal Ewald parameters for given accuracy."""

    alpha: float
    real_cutoff: float
    k_max: float
    error_real: float
    error_recip: float

@dataclass
class EwaldParameters:
    """Per-system cutoffs and reciprocal shifts prepared before evaluation.

    ``reciprocal_lattice_shifts`` stores sphere-filtered shifts padded for
    batching. ``reciprocal_shift_bound`` always stores three axis bounds.
    ``(0, 0, 0)`` selects explicit sums over the shifts; other bounds select
    factored phase grids over that box, from which the shifts are gathered.
    Rebuild parameters and the structure-factor cache when changing the cell
    or cutoff. ``reciprocal_dtype`` is the precision of structure-factor
    evaluations; ``float32`` is much faster on GPUs with slow double precision.
    """

    alpha: Table[SystemId, Array]  # (n_graphs,)
    cutoff: Table[SystemId, Array]  # (n_graphs,)
    reciprocal_lattice_shifts: Table[SystemId, Array]  # (n_graphs, n_kvecs, 3)
    k_max: Table[SystemId, Array]
    reciprocal_shift_bound: ReciprocalGridBound = field(static=True, default=(0, 0, 0))
    reciprocal_particle_chunk_size: int = field(
        static=True, default_factory=_default_reciprocal_particle_chunk
    )
    reciprocal_k_chunk_size: int = field(
        static=True, default=_DEFAULT_RECIPROCAL_K_CHUNK
    )
    reciprocal_dtype: DTypeLike = field(static=True, default=jnp.float64)

    @classmethod
    @no_jax_tracing
    def make(
        cls,
        charges: Table[ParticleId, IsEwaldPointData],
        cell: Table[SystemId, HasCell[Periodic3D]],
        epsilon_total: float = 1e-8,
        real_cutoff: float | None = None,
        *,
        reciprocal_particle_chunk_size: int | None = None,
        reciprocal_k_chunk_size: int = _DEFAULT_RECIPROCAL_K_CHUNK,
        reciprocal_dtype: DTypeLike = jnp.float64,
    ) -> EwaldParameters:
        """Estimate each system's parameters and prepare its reciprocal shifts.

        Args:
            charges: Particle table with charges and system assignments.
            cell: System table with periodic cells.
            epsilon_total: Target combined real-space and reciprocal-space error.
            real_cutoff: Optional fixed real-space cutoff in Å.
            reciprocal_particle_chunk_size: Maximum particles per response tile;
                defaults to 8192 on GPU and 1024 elsewhere.
            reciprocal_k_chunk_size: Maximum explicit k-vectors per tile; also sets the grid workspace budget.
            reciprocal_dtype: Precision of structure-factor evaluations.

        Returns:
            Ewald parameters with stored integer shifts and reciprocal cutoffs.
        """
        n_systems = len(cell.keys)
        sys_idx = np.asarray(charges.data.system.indices_in(cell.keys))
        charge_values = np.asarray(charges.data.charges)
        cell_list = [cell.data.cell[i] for i in range(n_systems)]
        estimates = [
            estimate_ewald_parameters(
                charge_values[sys_idx == i],
                u,
                real_cutoff=real_cutoff,
                epsilon_total=epsilon_total,
            )
            for i, u in enumerate(cell_list)
        ]
        return cls.from_cutoffs(
            cell,
            alpha=cell.set_data(jnp.asarray([est.alpha for est in estimates])),
            cutoff=cell.set_data(jnp.asarray([est.real_cutoff for est in estimates])),
            k_max=cell.set_data(jnp.asarray([est.k_max for est in estimates])),
            reciprocal_particle_chunk_size=reciprocal_particle_chunk_size,
            reciprocal_k_chunk_size=reciprocal_k_chunk_size,
            reciprocal_dtype=reciprocal_dtype,
        )

    @classmethod
    @no_jax_tracing
    def from_cutoffs(
        cls,
        cell: Table[SystemId, HasCell[Periodic3D]],
        alpha: Table[SystemId, Array],
        cutoff: Table[SystemId, Array],
        k_max: Table[SystemId, Array],
        *,
        reciprocal_particle_chunk_size: int | None = None,
        reciprocal_k_chunk_size: int = _DEFAULT_RECIPROCAL_K_CHUNK,
        reciprocal_dtype: DTypeLike = jnp.float64,
    ) -> EwaldParameters:
        """Prepare reciprocal shifts once from supplied physical cutoffs.

        Args:
            cell: System table with periodic cells.
            alpha: Screening parameters in 1/Å.
            cutoff: Real-space cutoffs in Å.
            k_max: Reciprocal cutoffs in 1/Å.
            reciprocal_particle_chunk_size: Maximum particles per response tile;
                defaults to 8192 on GPU and 1024 elsewhere.
            reciprocal_k_chunk_size: Maximum k-vectors per tile.
            reciprocal_dtype: Precision of structure-factor evaluations.

        Returns:
            Parameters with stored integer shifts and reciprocal cutoffs.
        """
        if reciprocal_particle_chunk_size is None:
            reciprocal_particle_chunk_size = _default_reciprocal_particle_chunk()
        shifts, bound = prepare_reciprocal_shifts(
            cell.data.cell,
            k_max[cell.index],
            grid=jax.default_backend() in RECIPROCAL_GRID_BACKENDS,
        )
        return cls(
            alpha=alpha,
            cutoff=cutoff,
            reciprocal_lattice_shifts=cell.set_data(shifts),
            k_max=k_max,
            reciprocal_shift_bound=bound,
            reciprocal_particle_chunk_size=reciprocal_particle_chunk_size,
            reciprocal_k_chunk_size=reciprocal_k_chunk_size,
            reciprocal_dtype=reciprocal_dtype,
        )

@dataclass
class EwaldPotential[State, Gradients, Hessians, P: Patch[Any]](
    SummedPotential[State, Gradients, Hessians, P]
):
    """Complete Ewald potential with named access to each component term."""

    @property
    def short_range(self) -> Potential[State, Gradients, Hessians, P]:
        """Real-space screened Coulomb potential."""
        return self.potentials[0]

    @property
    def long_range(self) -> Potential[State, Gradients, Hessians, P]:
        """Reciprocal-space potential."""
        return self.potentials[1]

    @property
    def self_interaction(self) -> Potential[State, Gradients, Hessians, P]:
        """Subtract each charge's interaction with its screening cloud."""
        return self.potentials[2]

    @property
    def net_charge(self) -> Potential[State, Gradients, Hessians, P]:
        """Neutralizing-background correction for net-charged systems."""
        return self.potentials[3]

    @property
    def exclusion_correction(self) -> Potential[State, Gradients, Hessians, P]:
        """Subtract excluded pair interactions; requires exclusions enabled."""
        return self.potentials[4]

class IsChargedTemplate(HasPositions, HasCharges, Protocol):
    """Properties needed to prepare intramolecular Coulomb energies."""

class IsEwaldPointData(HasCharges, NeighborListPoints, Protocol):
    """Particle data required by Ewald: charges, positions, system/inclusion/exclusion indices."""