kups.potential.classical.cosine_angle ¶

def cosine_angle_energy(
    inp: CosineAngleInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch[Any]]:
    r"""Compute UFF-style cosine angle energy for all angles.

    Calculates energy using the general cosine form:

    $$U(\theta) = K [C_0 + C_1 \cos(\theta) + C_2 \cos(2\theta)]$$

    For near-linear angles ($\theta_0$ close to 180deg), uses:

    $$U(\theta) = K (1 + \cos\theta)$$

    Args:
        inp: Graph potential input with cosine angle parameters

    Returns:
        Total angle energy per system
    """
    graph = inp.graph
    assert graph.edges.indices.indices.shape[1] == 3, (
        "Cosine angle potential only supports triplet interactions (order=3)."
    )
    edg = graph.particles[graph.edges.indices]
    edg_species = edg.labels.indices_in(inp.parameters.labels)
    s0, s1, s2 = edg_species[:, 0], edg_species[:, 1], edg_species[:, 2]
    theta0 = inp.parameters.theta0[s0, s1, s2]
    k = inp.parameters.k[s0, s1, s2]

    # Compute current angle
    v1, v2 = graph.edge_shifts[:, 0], graph.edge_shifts[:, 1]
    cos_theta = jnp.einsum("ij,ij->i", v1, v2) / (
        jnp.linalg.norm(v1, axis=-1) * jnp.linalg.norm(v2, axis=-1)
    )
    cos_theta = jnp.clip(cos_theta, -1.0, 1.0)

    # Check for near-linear case (theta0 close to pi)
    is_linear = jnp.abs(theta0 - jnp.pi) < inp.parameters.linear_tol

    # Linear case: U = K*(1 + cos(theta))
    energy_linear = k * (1.0 + cos_theta)

    # General case: U = K * [C0 + C1*cos(theta) + C2*cos(2*theta)]
    # Using cos(2t) = 2cos^2(t) - 1
    c0, c1, c2 = _compute_cosine_coefficients(theta0)
    cos_2theta = 2.0 * cos_theta**2 - 1.0
    energy_general = k * (c0 + c1 * cos_theta + c2 * cos_2theta)

    edge_energy = jnp.where(is_linear, energy_linear, energy_general)

    total_energies = graph.edge_batch_mask.sum_over(edge_energy)
    return WithPatch(total_energies, IdPatch[Any]())

def make_cosine_angle_potential[
    State,
    Ptch: Patch[Any],
    Gradients,
    Hessians,
](
    particles_view: View[State, Table[ParticleId, IsBondedParticles]],
    edge_indices_view: View[State, Index[ParticleId]],
    systems_view: View[State, Table[SystemId, HasCell[AnyPeriodicity]]],
    parameter_view: View[State, CosineAngleParameters],
    probe: Probe[State, Ptch, IsGraphProbe[IsBondedParticles, Literal[3]]] | None,
    gradient_lens: Lens[CosineAngleInput, Gradients],
    hessian_lens: Lens[Gradients, Hessians],
    hessian_idx_view: View[State, Hessians],
    patch_idx_view: View[State, PotentialOut[Gradients, Hessians]] | None = None,
    out_cache_lens: Lens[State, PotentialOut[Gradients, Hessians]] | None = None,
) -> Potential[State, Gradients, Hessians, Ptch]:
    """Create UFF-style cosine angle potential for explicitly defined angles.

    Args:
        particles_view: Extracts particle data (positions, species) with system index
        edge_indices_view: Extracts angle connectivity (triplets)
        systems_view: Extracts indexed system data (cell)
        parameter_view: Extracts [CosineAngleParameters][kups.potential.classical.cosine_angle.CosineAngleParameters]
        probe: Graph probe for incremental particle and neighbor-list updates
        gradient_lens: Specifies gradients to compute
        hessian_lens: Specifies Hessians to compute
        hessian_idx_view: Hessian index structure
        patch_idx_view: Cached output index structure
        out_cache_lens: Cache location lens

    Returns:
        Cosine angle [Potential][kups.core.potential.Potential]
    """
    graph_fn = GraphConstructor(
        particles=particles_view,
        systems=systems_view,
        neighborlist=lambda state: FixedEdgesNeighborList[Literal[3]](
            edge_indices_view(state)
        ),
        probe=probe,
    )
    composer = LocalGraphSumComposer(
        graph_constructor=graph_fn,
        parameter_view=parameter_view,
    )
    potential = PotentialFromEnergy(
        composer=composer,
        energy_fn=cosine_angle_energy,
        gradient_lens=gradient_lens,
        hessian_lens=hessian_lens,
        hessian_idx_view=hessian_idx_view,
        cache_lens=out_cache_lens,
        patch_idx_view=patch_idx_view,
    )
    return potential