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:
    """Ewald summation parameters: convergence settings and reciprocal lattice vectors.

    Attributes:
        alpha: Ewald screening parameter [1/Ang], shape `(n_graphs,)`.
        cutoff: Real-space cutoff radius [Ang], shape `(n_graphs,)`.
        reciprocal_lattice_shifts: Integer k-vector coefficients,
            shape `(n_graphs, n_kvecs, 3)`.
    """

    alpha: Table[SystemId, Array]  # (n_graphs,)
    cutoff: Table[SystemId, Array]  # (n_graphs,)
    reciprocal_lattice_shifts: Table[
        SystemId, Array
    ]  # (n_graphs, n_rec_shifts, 3) integers

    @classmethod
    @no_jax_tracing
    def make(
        cls,
        charges: Table[ParticleId, IsEwaldPointData],
        unitcell: Table[SystemId, HasUnitCell],
        epsilon_total: float = 1e-8,
        real_cutoff: float | None = None,
    ) -> EwaldParameters:
        """Estimate Ewald parameters from indexed particles and systems.

        Splits particles by system index, estimates per-system parameters,
        and zero-pads k-vectors to the maximum count across systems.

        Args:
            charges: Indexed particles with charges and system assignment.
            unitcell: Indexed systems with unit cells.
            epsilon_total: Target total accuracy for the Ewald sum.
            real_cutoff: Optional real-space cutoff override; estimated if not given.

        Returns:
            ``EwaldParameters`` with estimated physics parameters.
        """
        # Unpack Indexed into per-system lists
        n_systems = len(unitcell.keys)
        sys_idx = charges.data.system.indices
        charges_list = [charges.data.charges[sys_idx == i] for i in range(n_systems)]
        unitcell_list = [unitcell.data.unitcell[i] for i in range(n_systems)]

        estimates_list = [
            estimate_ewald_parameters(
                c, u, real_cutoff=real_cutoff, epsilon_total=epsilon_total
            )
            for c, u in zip(charges_list, unitcell_list)
        ]
        shifts_list = [
            kvecs_from_kmax(u, est.k_max)
            for u, est in zip(unitcell_list, estimates_list)
        ]
        max_n_kvecs = max(len(s) for s in shifts_list)
        padded_shifts = jnp.stack(
            [jnp.pad(s, [(0, max_n_kvecs - len(s)), (0, 0)]) for s in shifts_list]
        )
        return cls(
            alpha=Table(
                unitcell.keys,
                jnp.asarray([est.alpha for est in estimates_list]),
            ),
            cutoff=Table(
                unitcell.keys,
                jnp.asarray([est.real_cutoff for est in estimates_list]),
            ),
            reciprocal_lattice_shifts=Table(unitcell.keys, padded_shifts),
        )

@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

    @classmethod
    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)

@dataclass
class LennardJonesParameters:
    """Lennard-Jones potential parameters.

    Attributes:
        labels: Species labels as ``Index``.
        sigma: Length scale parameters [Å], shape ``(n_species, n_species)``.
        epsilon: Energy well depths [energy units], shape ``(n_species, n_species)``.
        cutoff: Cutoff radius [Å], shape ``(n_systems,)``.
    """

    labels: tuple[Label, ...] = field(static=True)  # (n_species,)
    sigma: Array  # (n_species, n_species) float
    epsilon: Array  # (n_species, n_species) float
    cutoff: Table[SystemId, Array]  # (n_graphs,) float

    @classmethod
    def from_dict(
        cls,
        cutoff: float | Array,
        parameters: dict[str, tuple[float | None, float | None]],
        mixing_rule: MixingRule,
    ) -> LennardJonesParameters:
        """Create 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.
        """
        labels = tuple(parameters.keys())
        raw = [(s or 1.0, e or 0.0) for s, e in parameters.values()]
        sigma, epsilon = jnp.asarray(raw).T
        cutoff_indexed = Table((SystemId(0),), jnp.array([cutoff]))
        match mixing_rule:
            case "lorentz_berthelot":
                return cls.from_lorentz_berthelot_mixing(
                    labels, sigma, epsilon, cutoff_indexed
                )
            case _ as unreachable:
                assert_never(unreachable)

    @classmethod
    def from_lorentz_berthelot_mixing(
        cls,
        labels: tuple[str, ...],
        sigma: Array,
        epsilon: Array,
        cutoff: Table[SystemId, Array],
    ) -> LennardJonesParameters:
        """Create parameters using Lorentz-Berthelot mixing rules.

        - σᵢⱼ = (σᵢ + σⱼ) / 2 (arithmetic mean)
        - εᵢⱼ = √(εᵢ × εⱼ) (geometric mean)
        """
        assert sigma.ndim == epsilon.ndim == 1
        sigma_matrix = (sigma[:, None] + sigma) / 2
        epsilon_matrix = jnp.sqrt(epsilon[:, None] * epsilon)
        return cls(tuple(map(Label, labels)), sigma_matrix, epsilon_matrix, cutoff)