kups.application.potential.mliap.local ¶

def make_local_mliap_from_state(
    state: Any,
    probe: Any = None,
    gradient_lens: Any = EMPTY_LENS,
    hessian_lens: Any = EMPTY_LENS,
    hessian_idx_view: Any = EMPTY_LENS,
    out_idx_view: Any = None,
    *,
    parameters: LocalMLIAPData | None = None,
    neighborlist_factory: NeighborListFactory[Any] = AdaptiveNeighborList.from_state,
) -> Any:
    """Create a local MLIAP potential from a typed state, optionally with incremental updates.

    Convenience wrapper around
    [make_local_mliap_potential][kups.potential.mliap.local.make_local_mliap_potential].
    When ``probe`` is ``None``, extracts views from a state satisfying
    [IsLocalMLIAPState][kups.application.potential.mliap.local.IsLocalMLIAPState].
    When ``probe`` is provided, additionally wires the ``PotentialOut`` cache for
    efficient incremental caching across Monte Carlo steps.

    Args:
        state: Lens into the sub-state providing particles, cell, neighbor list,
            and local MLIAP model (when ``parameters`` is not given).
        probe: Detects which particles and neighbor-list edges changed since the last
            step. ``None`` for full-only computation.
        gradient_lens: Specifies which gradients to compute (e.g., forces).
        hessian_lens: Specifies which Hessians to compute.
        hessian_idx_view: Index structure for Hessian updates.
        out_idx_view: Index into the cached output for partial updates. Only used when
            ``probe`` is not ``None``. Defaults to full re-indexing of the cached
            ``total_energies``.
        parameters: Constant MLIAP model. When given it is bound with a constant
            lens and the state need not carry ``local_mliap_model``; with a
            ``probe``, the cache is read from ``state.local_mliap_cache``.

    Returns:
        Configured local MLIAP [Potential][kups.core.potential.Potential].
    """

    def _model(x: Any) -> LocalMLIAPData:
        m = x.local_mliap_model
        return m.data if isinstance(m, HasCache) else m

    if probe is None:
        if parameters is not None:
            model_lens = const_lens(parameters)
        else:
            model_lens = state.focus(_model)

        def neighborlist_view(s: Any) -> NeighborList[Literal[2]]:
            return neighborlist_factory(state(s), model_lens(s).cutoff)

        return make_local_mliap_potential(
            state.focus(lambda x: x.particles),
            state.focus(lambda x: x.systems),
            neighborlist_view,
            model_lens,
            None,
            gradient_lens,
            hessian_lens,
            hessian_idx_view,
            None,
            None,
        )

    if parameters is not None:
        model_lens = const_lens(parameters)
        cache_view = state.focus(lambda x: x.local_mliap_cache)
    else:
        model_lens = state.focus(lambda x: x.local_mliap_model.data)
        cache_view = state.focus(lambda x: x.local_mliap_model.cache)

    if out_idx_view is None:
        out_idx_view = cache_view.focus(
            lambda c: bind(c, lambda x: x.total_energies.data).apply(
                lambda x: jnp.arange(x.size, dtype=int)
            )
        )

    def neighborlist_view(s: Any) -> NeighborList[Literal[2]]:
        return neighborlist_factory(state(s), model_lens(s).cutoff)

    return make_local_mliap_potential(
        state.focus(lambda x: x.particles),
        state.focus(lambda x: x.systems),
        neighborlist_view,
        model_lens,
        probe,
        gradient_lens,
        hessian_lens,
        hessian_idx_view,
        out_idx_view,
        cache_view,
    )