nvalchemi.models.lj.LennardJonesModelWrapper

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

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.

model_config

Mutable configuration controlling which outputs are computed. Include "stress" in model_config.active_outputs 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.

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

extra_repr()

Format the model config for nn.Module.__repr__.

Parameters:

self (Any)

Return type:

str

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_matrix_shifts (populated by NeighborListHook).

• kwargs (Any)

Returns:

OrderedDict with keys "energy" (shape [B, 1]), "forces" (shape [N, 3]), and optionally "stress" (shape [B, 3, 3]) — Cauchy stress W/V in energy units.

Return type:

ModelOutputs

output_data()

Return the set of keys that the model produces.

Return type:

set[str]