kups.potential.classical.cosine_angle ¶

def cosine_angle_energy(
    inp: CosineAngleInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch]:
    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())

def make_cosine_angle_from_state(
    state: Any,
    probe: Any = None,
    *,
    compute_position_and_unitcell_gradients: bool = False,
) -> Any:
    """Create a cosine angle potential from a typed state, optionally with incremental updates.

    Args:
        state: Lens into the sub-state providing particles, unit cell, edges, and parameters.
        probe: Detects which particles and edges changed since the last step.
            If ``None``, no incremental updates are used.
        compute_position_and_unitcell_gradients: When True, computes gradients
            w.r.t. particle positions and lattice vectors.

    Returns:
        Configured cosine angle [Potential][kups.core.potential.Potential].
    """
    gradient_lens: Any = EMPTY_LENS
    patch_idx_view: Any = None
    if compute_position_and_unitcell_gradients:
        gradient_lens = SimpleLens[CosineAngleInput, PositionAndUnitCell](
            lambda x: PositionAndUnitCell(
                x.graph.particles.map_data(lambda p: p.positions),
                x.graph.systems.map_data(lambda s: s.unitcell),
            )
        )
        patch_idx_view = position_and_unitcell_idx_view
    param_view = state.focus(
        lambda x: (
            x.cosine_angle_parameters.data
            if isinstance(x.cosine_angle_parameters, HasCache)
            else x.cosine_angle_parameters
        )
    )
    cache_view = None
    if probe is not None:
        param_view = state.focus(lambda x: x.cosine_angle_parameters.data)
        cache_view = state.focus(lambda x: x.cosine_angle_parameters.cache)
        patch_idx_view = patch_idx_view or empty_patch_idx_view
    return make_cosine_angle_potential(
        state.focus(lambda x: x.particles),
        state.focus(lambda x: x.cosine_angle_edges),
        state.focus(lambda x: x.systems),
        param_view,
        probe,
        gradient_lens,
        EMPTY_LENS,
        EMPTY_LENS,
        patch_idx_view,
        cache_view,
    )

def make_cosine_angle_potential[
    State,
    Ptch: Patch,
    Gradients,
    Hessians,
](
    particles_view: View[State, Table[ParticleId, IsBondedParticles]],
    edges_view: View[State, Edges[Literal[3]]],
    systems_view: View[State, Table[SystemId, HasUnitCell]],
    parameter_view: View[State, CosineAngleParameters],
    probe: Probe[State, Ptch, IsEdgeSetGraphProbe[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
        edges_view: Extracts angle connectivity (triplets)
        systems_view: Extracts indexed system data (unit cell)
        parameter_view: Extracts [CosineAngleParameters][kups.potential.classical.cosine_angle.CosineAngleParameters]
        probe: Probes particle, edge, and capacity changes for incremental 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 = EdgeSetGraphConstructor(
        particles=particles_view, edges=edges_view, systems=systems_view, 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