Skip to content

kups.application.potential.mliap.tojax

State-binding constructors for jaxified MLIAP potentials.

These adapters take a :class:~kups.core.lens.Lens into a concrete simulation state and wire its particles, systems, neighbor list, and jaxified model into the state-agnostic factories in kups.potential.mliap.tojax.

The jaxified model may live on the state (state.jaxified_model) or be passed directly via parameters=; in the latter case it is bound with a constant lens and the state need not carry a model field.

IsTojaxedGraphState

Bases: IsState[IsTojaxedParticles, HasCell[AnyPeriodicity]], IsNeighborListState[IsUniversalNeighborlistParams], Protocol

Particles, systems, and neighbor list for a jaxified graph (no model).

Source code in src/kups/application/potential/mliap/tojax.py 
40
41
42
43
44
45
class IsTojaxedGraphState(
    IsState[IsTojaxedParticles, HasCell[AnyPeriodicity]],
    IsNeighborListState[IsUniversalNeighborlistParams],
    Protocol,
):
    """Particles, systems, and neighbor list for a jaxified graph (no model)."""

IsTojaxedState

Bases: IsTojaxedGraphState, Protocol

:class:IsTojaxedGraphState that also carries the jaxified model.

Source code in src/kups/application/potential/mliap/tojax.py 
48
49
50
51
52
class IsTojaxedState(IsTojaxedGraphState, Protocol):
    """:class:`IsTojaxedGraphState` that also carries the jaxified model."""

    @property
    def jaxified_model(self) -> TojaxedMliap: ...

make_tojaxed_from_state(state, *, parameters=None, gradient=None, neighborlist_factory=AdaptiveNeighborList.from_state) 

make_tojaxed_from_state(
    state: Lens[State, IsTojaxedState],
    *,
    parameters: None = None,
    gradient: None = None,
    neighborlist_factory: NeighborListFactory[
        IsTojaxedState
    ] = ...,
) -> Potential[State, EmptyType, EmptyType, Patch[Any]]

make_tojaxed_from_state(
    state: Lens[State, IsTojaxedState],
    *,
    parameters: None = None,
    gradient: Lens[Geometry, PositionsAndCell],
    neighborlist_factory: NeighborListFactory[
        IsTojaxedState
    ] = ...,
) -> Potential[
    State, PositionsAndCell, EmptyType, Patch[Any]
]

make_tojaxed_from_state(
    state: Lens[State, IsTojaxedGraphState],
    *,
    parameters: TojaxedMliap,
    gradient: None = None,
    neighborlist_factory: NeighborListFactory[
        IsTojaxedGraphState
    ] = ...,
) -> Potential[State, EmptyType, EmptyType, Patch[Any]]

make_tojaxed_from_state(
    state: Lens[State, IsTojaxedGraphState],
    *,
    parameters: TojaxedMliap,
    gradient: Lens[Geometry, PositionsAndCell],
    neighborlist_factory: NeighborListFactory[
        IsTojaxedGraphState
    ] = ...,
) -> Potential[
    State, PositionsAndCell, EmptyType, Patch[Any]
]

Create a jaxified potential from a typed state.

Parameters:

Name Type Description Default
state Any

Lens into the sub-state providing particles, cell, and neighbor list (plus jaxified_model when parameters is not given).

 required
parameters TojaxedMliap | None

Constant jaxified model. When given it is bound with a constant lens and the state need not carry jaxified_model.

 None
gradient Lens[Geometry, PositionsAndCell] | None

Relaxation filter Lens[Geometry, PositionsAndCell] selecting the optimizer DOFs. None computes no gradients (the default). Composed with GRAPH_GEOMETRY into the potential's gradient lens.

 None

Returns:

Type Description
Any

Configured jaxified Potential.
Source code in src/kups/application/potential/mliap/tojax.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def make_tojaxed_from_state(
    state: Any,
    *,
    parameters: TojaxedMliap | None = None,
    gradient: Lens[Geometry, PositionsAndCell] | None = None,
    neighborlist_factory: NeighborListFactory[Any] = AdaptiveNeighborList.from_state,
) -> Any:
    """Create a jaxified potential from a typed state.

    Args:
        state: Lens into the sub-state providing particles, cell, and neighbor
            list (plus ``jaxified_model`` when ``parameters`` is not given).
        parameters: Constant jaxified model. When given it is bound with a
            constant lens and the state need not carry ``jaxified_model``.
        gradient: Relaxation filter ``Lens[Geometry, PositionsAndCell]`` selecting the
            optimizer DOFs. ``None`` computes no gradients (the default). Composed with
            ``GRAPH_GEOMETRY`` into the potential's gradient lens.

    Returns:
        Configured jaxified [Potential][kups.core.potential.Potential].
    """
    gradient_lens = cast(
        "Lens[Any, PositionsAndCell | EmptyType]",
        EMPTY_LENS if gradient is None else GRAPH_GEOMETRY.nest(gradient),
    )
    if parameters is not None:
        model_view = const_lens(parameters)
    else:
        model_view = state.focus(lambda x: x.jaxified_model)

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

    return make_tojaxed_potential(
        state.focus(lambda x: x.particles),
        state.focus(lambda x: x.systems),
        neighborlist_view,
        model_view,
        gradient_lens,
        EMPTY_LENS,
        EMPTY_LENS,
    )