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 the inputs the LJ kernel needs from data.

Unlike the base implementation this does not enable gradients on positions: forces come analytically from the Warp kernel, not from autograd.

Parameters:

• data (Batch) – The input batch. AtomicData is rejected; wrap it first with Batch.from_data_list([data]).

• **kwargs – Unused; accepted for interface compatibility.

Returns:

Kernel inputs: the configured input fields plus batch_idx, ptr, num_graphs, fill_value, and optional cells [B, 3, 3] / neighbor_matrix_shifts [N, K, 3].

Return type:

dict[str, Any]

Raises:

• KeyError – If a required input field is missing from data.

• TypeError – If data is an AtomicData rather than a Batch.

adapt_output(model_output, data)

Map the LJ kernel output to the framework ModelOutputs format.

Parameters:

• model_output (dict) – Raw kernel output with energy / forces and, when stress is active, virial (converted here to tensile-positive Cauchy stress -W / V).

• data (AtomicData | Batch) – Original input batch; its cell provides the volume for stress.

Returns:

OrderedDict with the active output keys.

Return type:

ModelOutputs

compute_embeddings(data, **kwargs)

Not implemented — the Lennard-Jones potential produces no embeddings.

Parameters:

• data (AtomicData | Batch) – The input system.

• **kwargs – Unused; accepted for interface compatibility.

Returns:

Never returns.

Return type:

AtomicData | Batch

Raises:

NotImplementedError – Always; the LJ potential has no learned embeddings.

distribution_spec(strategy=None)

MLIPSpec for the Lennard-Jones wrapper under domain decomposition.

Halo-only; the strategy argument is accepted for the framework contract and ignored (LJ ships no graph-parallel spec).

The LJ Warp kernels are opaque to sharded tensors, so each is wrapped in an OpAdapter that unwraps to local tensors for the kernel and re-wraps the per-atom outputs.

Returns:

The halo spec plus one OpAdapter per LJ kernel.

Return type:

MLIPSpec

Parameters:

strategy (Any)

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)

Not implemented for the Lennard-Jones wrapper.

Parameters:

• path (Path) – Output path (unused).

• as_state_dict (bool, optional) – Unused. Defaults to False.

Returns:

Never returns.

Return type:

None

Raises:

NotImplementedError – Always; the LJ wrapper carries no learned weights to export.

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 – Forwarded to adapt_input().

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 output keys the model produces this run.

Returns:

{"energy"} plus "forces" and/or "stress" when they are in model_config.active_outputs.

Return type:

set[str]