nvalchemi.models.base.ModelConfig

pydantic model nvalchemi.models.base.ModelConfig

Unified model configuration combining capability declaration and runtime control.

ModelConfig is the contract between a model wrapper and the rest of nvalchemi: dynamics engines, composition pipelines, loss functions, and the BaseModelMixin adapters all read it to decide which inputs to prepare, which gradients to enable, and which outputs to compute. Every BaseModelMixin subclass must set a self.model_config instance in its __init__ (there is deliberately no class-level default, so each wrapper owns its own config object).

A ModelConfig has two kinds of fields:

• Capability fields (frozen at construction) describe what the model checkpoint can do. These use frozenset to signal immutability. They are set once by the wrapper’s __init__ and should not be changed at runtime. Examples: outputs, autograd_outputs, autograd_inputs, required_inputs, optional_inputs, supports_pbc, needs_pbc, neighbor_config.

• Runtime fields (mutable) control what the model should compute on each forward pass. These can be changed freely by the user: active_outputs selects the subset of outputs to compute this run, and gradient_keys enables gradients on extra input tensors.

outputs and required_inputs use free-form strings so new properties can be added without modifying this class. Well-known output keys: energy, forces, stress, hessian, dipole, charges, embeddings. Declare a neighbor-list requirement by attaching a NeighborConfig (see needs_neighborlist).

Examples

An energy-and-forces model whose forces come from autograd on positions:

>>> from nvalchemi.models.base import ModelConfig
>>> cfg = ModelConfig(
...     outputs=frozenset({"energy", "forces"}),
...     autograd_outputs=frozenset({"forces"}),
...     autograd_inputs=frozenset({"positions"}),
... )
>>> cfg.active_outputs == {"energy", "forces"}
True

Restrict a run to a single output without changing the model’s capabilities:

>>> cfg.active_outputs = {"energy"}

Declare a PBC-aware model that needs a neighbor list:

>>> from nvalchemi.models.base import NeighborConfig
>>> cfg = ModelConfig(
...     outputs=frozenset({"energy", "forces", "stress"}),
...     autograd_outputs=frozenset({"forces", "stress"}),
...     supports_pbc=True,
...     needs_pbc=True,
...     neighbor_config=NeighborConfig(cutoff=5.0),
... )
>>> cfg.needs_neighborlist
True

Notes

• extra="forbid": unknown constructor keywords raise a pydantic.ValidationError, guarding against typo’d field names.

• active_outputs defaults to a mutable copy of outputs when left as None; set it to narrow the per-run output set, and reset it to None to fall back to all outputs.

• Capability fields are frozenset values: rebind the whole field to change them (in-place mutation is impossible), which keeps the declared capabilities effectively immutable after construction.

field outputs: frozenset[str] [Optional]

All properties the model can produce.

field autograd_outputs: frozenset[str] [Optional]

Subset of outputs computed via autograd.

field autograd_inputs: frozenset[str] [Optional]

Input keys needing requires_grad for autograd outputs.

field required_inputs: frozenset[str] [Optional]

Extra required inputs beyond {positions, atomic_numbers}.

field optional_inputs: frozenset[str] [Optional]

Extra inputs used if present, silently skipped if absent.

field supports_pbc: bool = False

Whether the model supports periodic boundary conditions.

field needs_pbc: bool = False

Whether the model requires PBC inputs.

field neighbor_config: NeighborConfig | None = None

Neighbor list requirements. None means no neighbor list.

field active_outputs: set[str] | None = None

Properties to compute this run. None means use all outputs (the default).

field gradient_keys: set[str] [Optional]

Extra input keys to enable gradients for.