Skip to content

kups.application.simulations.potentials

Force-field selection and potential construction for the unified apps.

A discriminated [PotentialConfig][kups.application.simulations.potentials.PotentialConfig] selects a force field; each config's build constructs the corresponding Potential directly with its parameters (so the parameters need not live on the simulation state) and returns the per-system cutoff used to size the neighbor list.

The torch backends (MACE/UMA) import torch lazily inside build so importing this module never requires the optional torch dependency.

BuiltPotential = tuple[Potential[State, PositionsAndCell, EmptyType, Patch[Any]], Table[SystemId, Array]]

(potential, cutoff) where cutoff sizes the neighbor list.

LjPotentialConfig

Bases: BaseModel

Lennard-Jones force field.

Source code in src/kups/application/simulations/potentials.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class LjPotentialConfig(BaseModel):
    """Lennard-Jones force field."""

    backend: Literal["lj"] = "lj"
    cutoff: float
    parameters: dict[str, tuple[float | None, float | None]]
    mixing_rule: MixingRule

    def build[State](
        self,
        state_lens: Lens[State, IsLJGraphState],
        gradient: Lens[Geometry, PositionsAndCell],
    ) -> BuiltPotential[State]:
        """Build the LJ potential with its parameters bound in."""
        params = LennardJonesParameters.from_dict(
            cutoff=self.cutoff,
            parameters=self.parameters,
            mixing_rule=self.mixing_rule,
        )
        potential = make_lennard_jones_from_state(
            state_lens, parameters=params, gradient=gradient
        )
        return potential, params.cutoff

build(state_lens, gradient)

Build the LJ potential with its parameters bound in.

Source code in src/kups/application/simulations/potentials.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def build[State](
    self,
    state_lens: Lens[State, IsLJGraphState],
    gradient: Lens[Geometry, PositionsAndCell],
) -> BuiltPotential[State]:
    """Build the LJ potential with its parameters bound in."""
    params = LennardJonesParameters.from_dict(
        cutoff=self.cutoff,
        parameters=self.parameters,
        mixing_rule=self.mixing_rule,
    )
    potential = make_lennard_jones_from_state(
        state_lens, parameters=params, gradient=gradient
    )
    return potential, params.cutoff

MaceConfig

Bases: BaseModel

PyTorch MACE checkpoint.

Source code in src/kups/application/simulations/potentials.py 
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
class MaceConfig(BaseModel):
    """PyTorch MACE checkpoint."""

    backend: Literal["mace"] = "mace"
    model_path: str | Path
    device: Literal["cpu", "cuda"] = "cuda"
    dtype: Literal["float32", "float64"] = "float32"

    def build[State](
        self,
        state_lens: Lens[State, IsTorchMliapGraphState],
        gradient: Lens[Geometry, PositionsAndCell],
    ) -> BuiltPotential[State]:
        """Load the MACE checkpoint and build the potential (lazy torch import)."""
        from kups.application.potential.mliap.torch import make_torch_mliap_from_state
        from kups.potential.mliap.torch import load_mace

        model = load_mace(
            get_model_path(self.model_path),
            device=self.device,
            dtype=self.dtype,
            compute_cell_gradients=True,
        )
        potential = make_torch_mliap_from_state(
            state_lens, parameters=model, gradient=gradient
        )
        return potential, model.cutoff

build(state_lens, gradient)

Load the MACE checkpoint and build the potential (lazy torch import).

Source code in src/kups/application/simulations/potentials.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def build[State](
    self,
    state_lens: Lens[State, IsTorchMliapGraphState],
    gradient: Lens[Geometry, PositionsAndCell],
) -> BuiltPotential[State]:
    """Load the MACE checkpoint and build the potential (lazy torch import)."""
    from kups.application.potential.mliap.torch import make_torch_mliap_from_state
    from kups.potential.mliap.torch import load_mace

    model = load_mace(
        get_model_path(self.model_path),
        device=self.device,
        dtype=self.dtype,
        compute_cell_gradients=True,
    )
    potential = make_torch_mliap_from_state(
        state_lens, parameters=model, gradient=gradient
    )
    return potential, model.cutoff

TojaxPotentialConfig

Bases: BaseModel

Jaxified MLFF (a model exported to JAX via tojax).

Source code in src/kups/application/simulations/potentials.py 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class TojaxPotentialConfig(BaseModel):
    """Jaxified MLFF (a model exported to JAX via tojax)."""

    backend: Literal["tojax"] = "tojax"
    model_path: str | Path

    def build[State](
        self,
        state_lens: Lens[State, IsTojaxedGraphState],
        gradient: Lens[Geometry, PositionsAndCell],
    ) -> BuiltPotential[State]:
        """Build the jaxified potential from the exported model."""
        model = TojaxedMliap.from_zip_file(get_model_path(self.model_path))
        potential = make_tojaxed_from_state(
            state_lens, parameters=model, gradient=gradient
        )
        return potential, model.cutoff

build(state_lens, gradient)

Build the jaxified potential from the exported model.

Source code in src/kups/application/simulations/potentials.py 
89
90
91
92
93
94
95
96
97
98
99
def build[State](
    self,
    state_lens: Lens[State, IsTojaxedGraphState],
    gradient: Lens[Geometry, PositionsAndCell],
) -> BuiltPotential[State]:
    """Build the jaxified potential from the exported model."""
    model = TojaxedMliap.from_zip_file(get_model_path(self.model_path))
    potential = make_tojaxed_from_state(
        state_lens, parameters=model, gradient=gradient
    )
    return potential, model.cutoff

UmaConfig

Bases: BaseModel

Meta FAIR Chemistry UMA checkpoint.

Source code in src/kups/application/simulations/potentials.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class UmaConfig(BaseModel):
    """Meta FAIR Chemistry UMA checkpoint."""

    backend: Literal["uma"] = "uma"
    model_path: str | Path
    device: Literal["cpu", "cuda"] = "cuda"
    task_name: UMATaskName = "omat"
    inference_settings: UMAInferenceSettings = "default"

    def build[State](
        self,
        state_lens: Lens[State, IsTorchMliapGraphState],
        gradient: Lens[Geometry, PositionsAndCell],
    ) -> BuiltPotential[State]:
        """Load the UMA checkpoint and build the potential (lazy torch import)."""
        from kups.application.potential.mliap.torch import make_torch_mliap_from_state
        from kups.potential.mliap.torch import load_uma

        model = load_uma(
            get_model_path(self.model_path),
            device=self.device,
            task_name=self.task_name,
            compute_cell_gradients=True,
            inference_settings=self.inference_settings,
        )
        potential = make_torch_mliap_from_state(
            state_lens, parameters=model, gradient=gradient
        )
        return potential, model.cutoff

build(state_lens, gradient)

Load the UMA checkpoint and build the potential (lazy torch import).

Source code in src/kups/application/simulations/potentials.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def build[State](
    self,
    state_lens: Lens[State, IsTorchMliapGraphState],
    gradient: Lens[Geometry, PositionsAndCell],
) -> BuiltPotential[State]:
    """Load the UMA checkpoint and build the potential (lazy torch import)."""
    from kups.application.potential.mliap.torch import make_torch_mliap_from_state
    from kups.potential.mliap.torch import load_uma

    model = load_uma(
        get_model_path(self.model_path),
        device=self.device,
        task_name=self.task_name,
        compute_cell_gradients=True,
        inference_settings=self.inference_settings,
    )
    potential = make_torch_mliap_from_state(
        state_lens, parameters=model, gradient=gradient
    )
    return potential, model.cutoff