kups.potential.classical.morse ¶

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

    Args:
        state: Lens into the sub-state providing particles, unit cell, edges,
            and Morse bond 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``, the returned
            potential computes gradients w.r.t. particle positions and lattice
            vectors (for forces / stress).

    Returns:
        Configured Morse bond [Potential][kups.core.potential.Potential].
    """
    gradient_lens: Any = EMPTY_LENS
    patch_idx_view: Any = None
    if compute_position_and_unitcell_gradients:
        gradient_lens = SimpleLens[MorseBondInput, 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.morse_bond_parameters.data
            if isinstance(x.morse_bond_parameters, HasCache)
            else x.morse_bond_parameters
        )
    )
    cache_view = None
    if probe is not None:
        param_view = state.focus(lambda x: x.morse_bond_parameters.data)
        cache_view = state.focus(lambda x: x.morse_bond_parameters.cache)
        patch_idx_view = patch_idx_view or empty_patch_idx_view
    return make_morse_bond_potential(
        state.focus(lambda x: x.particles),
        state.focus(lambda x: x.bond_edges),
        state.focus(lambda x: x.systems),
        param_view,
        probe,
        gradient_lens,
        EMPTY_LENS,
        EMPTY_LENS,
        patch_idx_view,
        cache_view,
    )

def make_morse_bond_potential[
    State,
    P: Patch,
    Gradients,
    Hessians,
](
    particles_view: View[State, Table[ParticleId, IsBondedParticles]],
    edges_view: View[State, Edges[Literal[2]]],
    systems_view: View[State, Table[SystemId, HasUnitCell]],
    parameter_view: View[State, MorseBondParameters],
    probe: Probe[State, P, IsEdgeSetGraphProbe[IsBondedParticles, Literal[2]]] | None,
    gradient_lens: Lens[MorseBondInput, 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, P]:
    """Create Morse bond potential for explicitly defined bonds.

    Args:
        particles_view: Extracts particle data (positions, species) with system index
        edges_view: Extracts bond connectivity
        systems_view: Extracts indexed system data (unit cell)
        parameter_view: Extracts [MorseBondParameters][kups.potential.classical.morse.MorseBondParameters]
        probe: Grouped probe for incremental updates (particles, edges, capacity)
        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:
        Morse bond [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=morse_bond_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

def morse_bond_energy(
    inp: MorseBondInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch]:
    r"""Compute Morse bond energy for all bonds.

    Calculates energy as $D [1 - e^{-\alpha(r - r_0)}]^2$ for each bond.

    Args:
        inp: Graph potential input with Morse bond parameters

    Returns:
        Total bond energy per system
    """
    graph = inp.graph
    assert graph.edges.indices.indices.shape[1] == 2, (
        "Morse bond potential only supports pairwise interactions (order=2)."
    )
    edg_species = graph.particles[graph.edges.indices].labels.indices_in(
        inp.parameters.labels
    )
    r0 = inp.parameters.r0[edg_species[:, 0], edg_species[:, 1]]
    D = inp.parameters.D[edg_species[:, 0], edg_species[:, 1]]
    alpha = inp.parameters.alpha[edg_species[:, 0], edg_species[:, 1]]
    r = jnp.linalg.norm(graph.edge_shifts[:, 0], axis=-1)
    edge_energy = D * (1 - jnp.exp(-alpha * (r - r0))) ** 2
    total_energies = graph.edge_batch_mask.sum_over(edge_energy)
    return WithPatch(total_energies, IdPatch())