nvalchemi.models.lj.LennardJonesModelWrapper

class nvalchemi.models.lj.LennardJonesModelWrapper(epsilon, sigma, cutoff, switch_width=0.0, half_list=False, max_neighbors=128)

Warp-accelerated Lennard-Jones potential as a model wrapper.

Parameters:

• epsilon (float) – LJ well-depth parameter (energy units, e.g. eV).

• sigma (float) – LJ zero-crossing distance (length units, e.g. Å).

• cutoff (float) – Interaction cutoff radius (same length units as positions).

• switch_width (float, optional) – Width of the C2-continuous switching region; 0.0 disables switching (hard cutoff). Defaults to 0.0.

• half_list (bool, optional) – Pass True (default) if the neighbor matrix contains each pair once (half list). Must match the half_fill argument given to NeighborListHook.

• max_neighbors (int, optional) – Maximum neighbors per atom used when building the neighbor matrix. Passed through to NeighborConfig and read by NeighborListHook. Defaults to 128.

model_config

Mutable configuration controlling which outputs are computed. Set model.model_config.compute_stresses = True to enable virial computation for NPT/NPH simulations.

Type:

ModelConfig

adapt_input(data, **kwargs)

Collect required inputs from data without enabling gradients.

Unlike the base-class implementation this method deliberately does not call positions.requires_grad_(True) because forces are computed analytically by the Warp kernel rather than via autograd.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

dict[str, Any]

adapt_output(model_output, data)

Adapts the model output to the framework’s expected format.

The super() implementation will provide the initial OrderedDict with keys that are expected to be present in the model output. This method will then map the model outputs to this OrderedDict.

Technically, this is not necessary for the LennardJonesModelWrapper, but it is included to demonstrate how to override the super() implementation.

Parameters:

• model_output (Any)

• data (AtomicData | Batch)

Return type:

OrderedDict[str, Float[Tensor, ‘B 1’] | Float[Tensor, ‘V 3’] | Float[Tensor, ‘V 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3 3’] | Float[Tensor, ‘B 3’] | None]

compute_embeddings(data, **kwargs)

Compute embeddings for the LennardJonesModelWrapper.

This method is not implemented for the LennardJonesModelWrapper, but it is included to demonstrate how to override the super() implementation.

Parameters:

• data (AtomicData | Batch)

• kwargs (Any)

Return type:

AtomicData | Batch

property embedding_shapes: dict[str, tuple[int, ...]]

Retrieves the expected shapes of the node, edge, and graph embeddings.

export_model(path, as_state_dict=False)

Export model is not implemented for LennardJonesModelWrapper.

Parameters:

• path (Path)

• as_state_dict (bool)

Return type:

None

forward(data, **kwargs)

Run the LJ kernel and return a ModelOutputs dict.

Parameters:

• data (Batch) – Batch containing positions, neighbor_matrix, num_neighbors, and optionally cell / neighbor_shifts (populated by NeighborListHook).

• kwargs (Any)

Returns:

OrderedDict with keys "energies" (shape [B, 1]), "forces" (shape [N, 3]), and optionally "stress" (shape [B, 3, 3]) — the physical virial -W_LJ in units of eV, ready for NPT/NPH barostat use.

Return type:

ModelOutputs

property model_card: ModelCard

Retrieves the model card for the model.

The model card is a Pydantic model that contains information about the model’s capabilities and requirements.

output_data()

Return the set of keys that the model produces.

Return type:

set[str]