kups.potential.classical ¶

@dataclass
class CosineAngleParameters:
    """UFF-style cosine angle potential parameters.

    Attributes:
        labels: Species labels, shape `(n_species,)`.
        theta0: Equilibrium angles [radians], shape `(n_species, n_species, n_species)`.
        k: Force constants [energy], shape `(n_species, n_species, n_species)`.
        linear_tol: Tolerance for detecting linear angles [radians].
    """

    labels: tuple[Label, ...] = field(static=True)  # (n_species,)
    theta0: Array  # (n_species, n_species, n_species)
    k: Array  # (n_species, n_species, n_species)
    linear_tol: Array = field(default_factory=lambda: jnp.radians(5))

    @classmethod
    def from_uff(
        cls,
        labels: tuple[str, ...],
        bond_angle: Array,
        bond_radius: Array,
        electronegativity: Array,
        effective_charge: Array,
        linear_tol: Array = jnp.radians(5),
    ) -> "CosineAngleParameters":
        r"""Create angle parameters using UFF formulas.

        Computes angle parameters from per-species atomic properties:
        - $\theta_0$ from central atom's bond angle
        - $K$ from Eq. 13 using bond lengths and effective charges

        Args:
            labels: Species labels, shape `(n_species,)`
            bond_angle: Natural valence angle [radians], shape `(n_species,)`
            bond_radius: Valence bond radii [A], shape `(n_species,)`
            electronegativity: GMP electronegativity, shape `(n_species,)`
            effective_charge: Effective atomic charge Z*, shape `(n_species,)`
            linear_tol: Tolerance for detecting linear angles [radians]

        Returns:
            CosineAngleParameters with full interaction matrices
        """
        nt = bond_angle.shape[0]

        # theta0 from central atom (index j in i-j-k)
        theta0 = bond_angle[None, :, None] * jnp.ones((nt, nt, nt))

        # Compute bond lengths for all pairs
        r_bond = compute_uff_bond_length(bond_radius, electronegativity)

        # r_ij = r_bond[i, j], r_jk = r_bond[j, k]
        r_ij = r_bond[:, :, None]
        r_jk = r_bond[None, :, :]

        # r_ik using law of cosines
        cos_theta0 = jnp.cos(theta0)
        r_ik_sq = r_ij**2 + r_jk**2 - 2 * r_ij * r_jk * cos_theta0
        r_ik = jnp.sqrt(jnp.maximum(r_ik_sq, 1e-10))

        # Force constant (Eq. 13)
        Z_i = effective_charge[:, None, None]
        Z_k = effective_charge[None, None, :]
        prefactor = 664.12 / (r_ij * r_jk)
        term = 3 * r_ij * r_jk * (1 - cos_theta0**2) - r_ik_sq * cos_theta0
        K = jnp.abs(prefactor * Z_i * Z_k / (r_ik**5) * term)

        return cls(
            labels=tuple(map(Label, labels)), theta0=theta0, k=K, linear_tol=linear_tol
        )

@dataclass
class DihedralParameters:
    """UFF dihedral/torsion potential parameters.

    Attributes:
        labels: Species labels, shape `(n_species,)`.
        V: Barrier heights [energy], shape `(n_species, n_species, n_species, n_species)`.
        n: Periodicities, shape `(n_species, n_species, n_species, n_species)`.
        phi0: Equilibrium dihedral angles [radians],
            shape `(n_species, n_species, n_species, n_species)`.
    """

    labels: tuple[Label, ...] = field(static=True)  # (n_species,)
    V: Array  # (n_species, n_species, n_species, n_species)
    n: Array  # (n_species, n_species, n_species, n_species)
    phi0: Array  # (n_species, n_species, n_species, n_species)

    @classmethod
    def from_uff(
        cls,
        labels: tuple[str, ...],
        bond_angle: Array,
        torsion_sp3: Array,
        torsion_sp2: Array,
        group: Array,
        bond_order: Array | None = None,
        hybridization_tol: float = 5.0 * 3.141592653589793 / 180.0,
    ) -> "DihedralParameters":
        r"""Create dihedral parameters using UFF formulas.

        Computes torsion parameters based on central bond hybridization following
        the UFF paper (Rappe et al. 1992). For dihedral i-j-k-l, parameters depend
        on the central j-k bond and neighboring atoms.

        Args:
            labels: Species labels, shape `(n_species,)`.
            bond_angle: Natural valence angle [radians], shape `(n_species,)`.
            torsion_sp3: sp3 torsional barrier V [kcal/mol], shape `(n_species,)`.
            torsion_sp2: sp2 torsional parameter U [kcal/mol], shape `(n_species,)`.
            group: Periodic table group (1-18), shape `(n_species,)`.
            bond_order: Bond order for j-k bond, shape `(n_species, n_species)`.
            hybridization_tol: Tolerance for hybridization detection [radians].

        Returns:
            DihedralParameters with full 4D interaction matrices.
        """
        nt = bond_angle.shape[0]

        # Determine hybridization from bond angle
        is_sp2 = jnp.abs(bond_angle - jnp.radians(120.0)) < hybridization_tol
        is_sp3 = ~(jnp.abs(bond_angle - jnp.pi) < hybridization_tol) & ~is_sp2

        # Group 6 and main group detection
        is_group6 = group == 6
        is_main_group = group > 0

        # Broadcast to 4D for dihedral i-j-k-l
        is_sp2_i = is_sp2[:, None, None, None]
        is_sp3_j = is_sp3[None, :, None, None]
        is_sp3_k = is_sp3[None, None, :, None]
        is_sp2_j = is_sp2[None, :, None, None]
        is_sp2_k = is_sp2[None, None, :, None]
        is_sp2_l = is_sp2[None, None, None, :]

        is_group6_j = is_group6[None, :, None, None]
        is_group6_k = is_group6[None, None, :, None]
        is_main_j = is_main_group[None, :, None, None]
        is_main_k = is_main_group[None, None, :, None]

        # Torsional parameters broadcast to 4D
        V_j = torsion_sp3[None, :, None, None]
        V_k = torsion_sp3[None, None, :, None]
        U_j = torsion_sp2[None, :, None, None]
        U_k = torsion_sp2[None, None, :, None]

        # Bond order for central j-k bond
        if bond_order is None:
            BO_jk = jnp.ones((1, nt, nt, 1))
        else:
            BO_jk = bond_order[None, :, :, None]

        # V from Eq.17: V = 5 * sqrt(U_j * U_k) * (1 + 4.18 * ln(BO))
        V_eq17 = (
            5.0
            * jnp.sqrt(jnp.abs(U_j * U_k))
            * (1.0 + 4.18 * jnp.log(jnp.maximum(BO_jk, 1e-10)))
        )

        # Initialize arrays
        V_arr = jnp.zeros((nt, nt, nt, nt))
        n_arr = jnp.ones((nt, nt, nt, nt), dtype=jnp.int32)
        phi0_arr = jnp.zeros((nt, nt, nt, nt))

        # === Case 1: sp3-sp3 general ===
        sp3_sp3 = is_sp3_j & is_sp3_k
        V_arr = jnp.where(sp3_sp3, jnp.sqrt(jnp.abs(V_j * V_k)), V_arr)
        n_arr = jnp.where(sp3_sp3, 3, n_arr)
        phi0_arr = jnp.where(sp3_sp3, jnp.pi, phi0_arr)

        # === Case 2: sp3-sp3 Group 6 pair (override n and phi0) ===
        sp3_sp3_group6 = sp3_sp3 & is_group6_j & is_group6_k
        n_arr = jnp.where(sp3_sp3_group6, 2, n_arr)
        phi0_arr = jnp.where(sp3_sp3_group6, jnp.pi / 2, phi0_arr)

        # === Case 3: sp2-sp2 ===
        sp2_sp2 = is_sp2_j & is_sp2_k
        V_arr = jnp.where(sp2_sp2, V_eq17, V_arr)
        n_arr = jnp.where(sp2_sp2, 2, n_arr)
        phi0_arr = jnp.where(sp2_sp2, jnp.pi, phi0_arr)

        # === Case 4: sp3-sp2 general ===
        sp3_sp2 = (is_sp3_j & is_sp2_k) | (is_sp2_j & is_sp3_k)
        V_arr = jnp.where(sp3_sp2, 1.0, V_arr)
        n_arr = jnp.where(sp3_sp2, 6, n_arr)
        phi0_arr = jnp.where(sp3_sp2, 0.0, phi0_arr)

        # === Case 5: sp3-sp2 propene-like (sp2 bonded to another sp2) ===
        propene = (is_sp2_j & is_sp3_k & is_sp2_i) | (is_sp3_j & is_sp2_k & is_sp2_l)
        V_arr = jnp.where(propene, 2.0, V_arr)
        n_arr = jnp.where(propene, 3, n_arr)
        phi0_arr = jnp.where(propene, jnp.pi, phi0_arr)

        # === Case 6: sp3-sp2 with Group 6 sp3 ===
        group6_sp3_sp2 = (is_sp3_j & is_group6_j & is_sp2_k) | (
            is_sp2_j & is_sp3_k & is_group6_k
        )
        V_arr = jnp.where(group6_sp3_sp2, V_eq17, V_arr)
        n_arr = jnp.where(group6_sp3_sp2, 2, n_arr)
        phi0_arr = jnp.where(group6_sp3_sp2, jnp.pi / 2, phi0_arr)

        # === Case 7: Non-main-group (V=0) ===
        non_main = ~is_main_j | ~is_main_k
        V_arr = jnp.where(non_main, 0.0, V_arr)

        return cls(labels=tuple(map(Label, labels)), V=V_arr, n=n_arr, phi0=phi0_arr)

@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 GlobalTailCorrectedLennardJonesParameters(LennardJonesParameters):
    """Lennard-Jones parameters with analytical long-range correction.

    Attributes:
        tail_corrected: Enable correction per species pair, shape ``(n_species, n_species)``.
    """

    tail_corrected: Array  # (n_species, n_species) bool

    @property
    def tail_coefficients(self) -> tuple[Array, Array]:
        """Two cutoff-independent coefficient matrices over species pairs."""
        sigma6 = self.sigma**6
        base = self.tail_corrected * self.epsilon
        return base * sigma6, base * sigma6**2

    @property
    def tail_weights(self) -> Array:
        """Current pair weights, shaped ``(2, n_systems, n_species, n_species)``."""
        term1 = (self.sigma / self.cutoff.data[:, None, None]) ** 3
        base = self.tail_corrected * self.epsilon * self.sigma**3
        return jnp.stack([base * term1, base * term1**3])

    @classmethod
    @override
    def from_dict(
        cls,
        cutoff: float | Array,
        parameters: dict[str, tuple[float | None, float | None]],
        mixing_rule: MixingRule,
        tail_correction: bool = True,
    ) -> GlobalTailCorrectedLennardJonesParameters:
        """Create tail-corrected parameters from a dict of per-species values.

        Args:
            cutoff: Cutoff radius [Angstrom].
            parameters: Map from species label to ``(sigma, epsilon)`` pair.
                ``None`` values default to ``sigma=1.0``, ``epsilon=0.0``.
            mixing_rule: Combining rule for cross-species interactions.
            tail_correction: Whether to enable tail corrections for all
                non-zero epsilon pairs.
        """
        base = LennardJonesParameters.from_dict(cutoff, parameters, mixing_rule)
        mask = (
            base.epsilon > 0
            if tail_correction
            else jnp.zeros_like(base.epsilon, dtype=bool)
        )
        return cls(
            labels=base.labels,
            sigma=base.sigma,
            epsilon=base.epsilon,
            cutoff=base.cutoff,
            tail_corrected=mask,
        )

@dataclass
class HarmonicAngleParameters:
    """Harmonic angle potential parameters.

    Attributes:
        labels: Species labels, shape `(n_species,)`
        theta0: Equilibrium angles [degrees], shape `(n_species, n_species, n_species)`
        k: Force constants [energy/degree²], shape `(n_species, n_species, n_species)`
    """

    labels: tuple[Label, ...] = field(static=True)  # (n_species,)
    theta0: Array  # (n_species, n_species, n_species)
    k: Array  # (n_species, n_species, n_species)

@dataclass
class HarmonicBondParameters:
    """Harmonic bond potential parameters.

    Attributes:
        labels: Species labels, shape `(n_species,)`
        x0: Equilibrium bond lengths [Å], shape `(n_species, n_species)`
        k: Force constants [energy/Ų], shape `(n_species, n_species)`
    """

    labels: tuple[Label, ...] = field(static=True)  # (n_species,)
    x0: Array  # (n_species, n_species)
    k: Array  # (n_species, n_species)

@dataclass
class InversionParameters:
    r"""UFF-style inversion potential parameters.

    Attributes:
        labels: Species labels, shape `(n_species,)`.
        omega0: Equilibrium out-of-plane angles [radians],
            shape `(n_species, n_species, n_species, n_species)`.
        k: Force constant or barrier [energy],
            shape `(n_species, n_species, n_species, n_species)`.
    """

    labels: tuple[Label, ...] = field(static=True)  # (n_species,)
    omega0: Array  # (n_species, n_species, n_species, n_species)
    k: Array  # (n_species, n_species, n_species, n_species)

    @classmethod
    def from_uff(
        cls,
        labels: tuple[str, ...],
        inversion_barrier: Array,
        omega0: Array | None = None,
    ) -> "InversionParameters":
        r"""Create inversion parameters using UFF formulas.

        Args:
            labels: Species labels, shape `(n_species,)`.
            inversion_barrier: Barrier/force constant for each type, shape `(n_species,)`.
            omega0: Equilibrium angle [radians], shape `(n_species,)`.

        Returns:
            InversionParameters with full interaction matrices
        """
        nt = inversion_barrier.shape[0]

        # Broadcast to 4D: center is first index
        k_arr = inversion_barrier[:, None, None, None] * jnp.ones((nt, nt, nt, nt))

        if omega0 is None:
            omega0_arr = jnp.zeros((nt, nt, nt, nt))
        else:
            omega0_arr = omega0[:, None, None, None] * jnp.ones((nt, nt, nt, nt))

        return cls(labels=tuple(map(Label, labels)), omega0=omega0_arr, k=k_arr)