Skip to content

kups.application.potential.mliap.torch

State-binding constructors for torch MLFF potentials.

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

The model may live on the state (state.torch_mliap_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.

IsTorchMliapGraphState

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

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

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

IsTorchMliapState

Bases: IsTorchMliapGraphState, Protocol

:class:IsTorchMliapGraphState that also carries the model on the state.

Source code in src/kups/application/potential/mliap/torch.py 
48
49
50
51
52
class IsTorchMliapState(IsTorchMliapGraphState, Protocol):
    """:class:`IsTorchMliapGraphState` that also carries the model on the state."""

    @property
    def torch_mliap_model(self) -> TorchMliap: ...

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

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

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

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

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

Create a torch MLFF potential from a typed state.

Parameters:

Name Type Description Default
state Any

Lens into a sub-state providing particles, systems, and neighbor list (plus torch_mliap_model when parameters is not given).

 required
parameters TorchMliap | None

Constant torch MLFF model. When given it is bound with a constant lens and the state need not carry torch_mliap_model.

 None
gradient Lens[Geometry, PositionsAndCell] | None

Relaxation filter Lens[Geometry, PositionsAndCell] selecting the optimizer DOFs; the torch gradients are pulled back through gradient.set into ∂E/∂u. When None (default), returns the raw PositionsAndCell (force + stress) gradients.

 None
neighborlist_factory NeighborListFactory[Any]

Neighbor list factory.

 from_state

Returns:

Type Description
Any

Configured torch MLFF Potential.
Source code in src/kups/application/potential/mliap/torch.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
137
138
139
140
141
def make_torch_mliap_from_state(
    state: Any,
    *,
    parameters: TorchMliap | None = None,
    gradient: Lens[Geometry, PositionsAndCell] | None = None,
    neighborlist_factory: NeighborListFactory[Any] = AdaptiveNeighborList.from_state,
) -> Any:
    """Create a torch MLFF potential from a typed state.

    Args:
        state: Lens into a sub-state providing particles, systems, and neighbor
            list (plus ``torch_mliap_model`` when ``parameters`` is not given).
        parameters: Constant torch MLFF model. When given it is bound with a
            constant lens and the state need not carry ``torch_mliap_model``.
        gradient: Relaxation filter ``Lens[Geometry, PositionsAndCell]`` selecting the
            optimizer DOFs; the torch gradients are pulled back through ``gradient.set``
            into ``∂E/∂u``. When ``None`` (default), returns the raw
            ``PositionsAndCell`` (force + stress) gradients.
        neighborlist_factory: Neighbor list factory.

    Returns:
        Configured torch MLFF ``Potential``.
    """
    if parameters is not None:
        model_view = const_lens(parameters)
    else:
        model_view = state.focus(lambda x: x.torch_mliap_model)

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

    particles_view = state.focus(lambda x: x.particles)
    systems_view = state.focus(lambda x: x.systems)
    if gradient is not None:
        return make_torch_mliap_potential(
            particles_view,
            systems_view,
            neighborlist_view,
            model_view,
            gradient=gradient,
        )
    return make_torch_mliap_potential(
        particles_view,
        systems_view,
        neighborlist_view,
        model_view,
    )