Calculators

Calculator interfaces for molecular simulations using AIMNet2.

AIMNet2Calculator

The core calculator for running AIMNet2 inference. It handles model loading, device management, and application of long-range interactions (Coulomb and Dispersion).

Key Features

• Format Support: Loads both legacy .jpt models and new .pt format.
• Long-Range Interactions: Automatically attaches LRCoulomb and DFTD3 modules based on model metadata.
• Overrides: needs_coulomb and needs_dispersion resolve effective long-range behavior after artifact validation. Explicit False values can disable external components for structurally valid custom artifacts.
• Batching: Automatically batches large molecules/systems based on nb_threshold.

AIMNet2Calculator(model='aimnet2', nb_threshold=120, needs_coulomb=None, needs_dispersion=None, device=None, compile_model=False, compile_kwargs=None, cache_static=False, train=False, deterministic=False, ensemble_member=0, revision=None, token=None, *, model_import_paths=None, model_import_mode='extend')

Generic AIMNet2 calculator.

A helper class to load AIMNet2 models and perform inference.

Parameters

model : str | nn.Module Model name (from registry), path to model file, or nn.Module instance. nb_threshold : int Threshold for neighbor list batching. Molecules larger than this use flattened processing. Default is 120. needs_coulomb : bool | None Whether to add external Coulomb module. If None (default), determined from model metadata. If True/False, overrides metadata. needs_dispersion : bool | None Whether to add external DFTD3 module. If None (default), determined from model metadata. If True/False, overrides metadata. device : str | None Device to run the model on ("cuda", "cpu", or specific like "cuda:0"). If None (default), auto-detects CUDA availability. compile_model : bool Whether to compile the model with torch.compile(). Default is False. compile_kwargs : dict | None Additional keyword arguments to pass to torch.compile(). Default is None. cache_static : bool Whether to cache calculator-built neighbor matrices and explicit external DFTD3 terms for repeated static CUDA inputs. Default is False. This opt-in cache is limited to exact reuse of the same non-periodic 2D input tensors and is bypassed for Hessian, stress, training, caller-supplied mol_idx, and caller-supplied nbmat. train : bool Whether to enable training mode. Default is False (inference mode). When False, all model parameters have requires_grad=False, which improves torch.compile compatibility and reduces memory usage. Set to True only when training the model. deterministic : bool Route external DFT-D3 (and DSF Coulomb) through their differentiable pure-torch paths instead of the atomics-based nvalchemiops CUDA kernels, making repeated identical evaluations bitwise reproducible (same machine, same build). Default is False. The kernel defaults are faster on large systems but accumulate in nondeterministic order, giving run-to-run noise of ~1e-7 eV that iterative optimizers can amplify. Ewald/PME Coulomb is not covered (a one-time warning fires); the static DFTD3 cache does not apply in this mode. ensemble_member : int Zero-based member selected from a Hugging Face ensemble. revision : str | None Hugging Face repository revision, branch, or tag. token : str | None Hugging Face access token for private repositories. model_import_paths : Collection[str] | None Python imports trusted for a direct local v2 artifact or complete Hugging Face repository. Each entry is an exact dotted path or a namespace ending in .*, for example {"my_package.models.CustomModel", "my_package.layers.*"}. model_import_mode : {"extend", "replace", "unsafe"} extend adds model_import_paths to the default trusted paths; replace requires a nonempty collection and uses only those paths. unsafe cannot be combined with paths and permits arbitrary imported constructors, so use it only for locally trusted artifacts. Registry names and aliases, registry HF fallback, raw modules, and .jpt files accept only the default settings. No mode relaxes artifact or metadata validation.

Attributes

model : nn.Module The loaded AIMNet2 model. device : str Device the model is running on ("cuda" or "cpu"). cutoff : float Short-range cutoff distance in Angstroms. cutoff_lr : float | None Long-range cutoff distance, or None if no LR modules. external_coulomb : LRCoulomb | None External Coulomb module if attached. external_dftd3 : DFTD3 | None External DFTD3 module if attached.

Notes

External LR module behavior:

• For file-loaded models (str): metadata is loaded from file
• For nn.Module: metadata is read from model.metadata attribute if available
• Explicit flags (needs_coulomb, needs_dispersion) override metadata
• If no metadata and no explicit flags, no external LR modules are added

Source code in aimnet/calculators/calculator.py

def __init__(
    self,
    model: str | nn.Module = "aimnet2",
    nb_threshold: int = 120,
    needs_coulomb: bool | None = None,
    needs_dispersion: bool | None = None,
    device: str | None = None,
    compile_model: bool = False,
    compile_kwargs: dict | None = None,
    cache_static: bool = False,
    train: bool = False,
    deterministic: bool = False,
    ensemble_member: int = 0,
    revision: str | None = None,
    token: str | None = None,
    *,
    model_import_paths: Collection[str] | None = None,
    model_import_mode: Literal["extend", "replace", "unsafe"] = "extend",
):
    # Device selection: use provided or auto-detect
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    self.device = str(torch.device(device))
    self.model: nn.Module
    self.external_coulomb: LRCoulomb | None = None
    self.external_dftd3: DFTD3 | None = None
    # Default cutoffs for LR modules
    self._default_dsf_cutoff = 15.0
    self._default_dftd3_cutoff = 15.0
    self._default_dftd3_smoothing = 0.2

    # Load model and get metadata
    self.model, metadata, self.cutoff = resolve_model(
        model,
        device=self.device,
        ensemble_member=ensemble_member,
        revision=revision,
        token=token,
        model_import_paths=model_import_paths,
        model_import_mode=model_import_mode,
    )

    # Compile model if requested
    self._was_compiled = bool(compile_model)
    if compile_model:
        kwargs = compile_kwargs or {}
        self.model = cast(nn.Module, torch.compile(self.model, **kwargs))

    # Resolve final flags (explicit overrides metadata)
    final_needs_coulomb = (
        needs_coulomb
        if needs_coulomb is not None
        else (metadata.get("needs_coulomb", False) if metadata is not None else False)
    )
    final_needs_dispersion = (
        needs_dispersion
        if needs_dispersion is not None
        else (metadata.get("needs_dispersion", False) if metadata is not None else False)
    )
    if metadata is not None:
        validate_runtime_model_metadata(
            metadata,
            needs_coulomb=final_needs_coulomb,
            needs_dispersion=final_needs_dispersion,
        )

    # Set up external Coulomb if needed
    if final_needs_coulomb:
        sr_embedded = metadata.get("coulomb_mode") == "sr_embedded" if metadata is not None else False
        coulomb_sr_rc = metadata.get("coulomb_sr_rc") if metadata is not None else None
        coulomb_sr_envelope = metadata.get("coulomb_sr_envelope") if metadata is not None else None
        # For PBC, user can switch to DSF/Ewald via set_lrcoulomb_method()
        # When sr_embedded=True: model has SRCoulomb which subtracts SR, so external
        # should compute FULL (subtract_sr=False) to give: (NN - SR) + FULL = NN + LR
        # When sr_embedded=False: model has no SR embedded, so external should compute
        # LR only (subtract_sr=True) to avoid double-counting
        self.external_coulomb = LRCoulomb(
            key_in="charges",
            key_out="energy",
            method="simple",
            rc=4.6 if coulomb_sr_rc is None else coulomb_sr_rc,
            envelope="exp" if coulomb_sr_envelope is None else coulomb_sr_envelope,
            subtract_sr=not sr_embedded,
        )
        self.external_coulomb = self.external_coulomb.to(self.device)

    # Set up external DFTD3 if needed
    if final_needs_dispersion:
        d3_params = metadata.get("d3_params") if metadata else None
        if d3_params is None:
            raise ValueError(
                "needs_dispersion=True but d3_params not found in metadata. "
                "Provide d3_params in model metadata or set needs_dispersion=False."
            )
        self.external_dftd3 = DFTD3(
            s8=d3_params["s8"],
            a1=d3_params["a1"],
            a2=d3_params["a2"],
            s6=d3_params.get("s6", 1.0),
        )
        self.external_dftd3 = self.external_dftd3.to(self.device)

    # Determine if model has long-range modules (embedded or external).
    #
    # The module tree is ground truth and metadata is a hint, NOT the other
    # way round. Treating a present metadata dict as authoritative fixed
    # only the metadata-absent half of #118: the shipped wb97m_cpcm_v2_0.pt
    # DOES carry a metadata dict, and that dict says has_embedded_lr=False
    # and has_embedded_d3ts=False while the module tree carries
    # `outputs.d3bj` and has_externalizable_dftd3() returns True. Believing
    # it left lr=False, so no LR neighbor list was built and D3BJ.forward
    # raised KeyError on the flattened path -- i.e. every system above
    # nb_threshold (measured: 119 atoms works, 125 fails).
    #
    # A stale or wrong flag can only ever cause a missing neighbor list,
    # never a spurious one, so OR-ing is safe in the direction that matters.
    has_embedded_lr = (
        bool(metadata.get("has_embedded_lr", False)) if metadata is not None else False
    ) or self._detect_embedded_lr_modules()
    self.lr = (
        hasattr(self.model, "cutoff_lr")
        or self.external_coulomb is not None
        or self.external_dftd3 is not None
        or has_embedded_lr
    )
    # Set cutoff_lr based on model attribute or external modules
    if hasattr(self.model, "cutoff_lr"):
        self.cutoff_lr = getattr(self.model, "cutoff_lr", float("inf"))
    elif self.external_coulomb is not None:
        # For "simple" method, use inf (all pairs). For DSF, use dsf_cutoff.
        if self.external_coulomb.method == "simple":
            self.cutoff_lr = float("inf")
        else:
            self.cutoff_lr = self._default_dsf_cutoff
    elif self.external_dftd3 is not None or self._has_embedded_dispersion():
        self.cutoff_lr = self._default_dftd3_cutoff
    elif has_embedded_lr:
        # Embedded Coulomb and unknown legacy LR modules need all pairs;
        # do not silently apply the finite D3 cutoff.
        self.cutoff_lr = float("inf")
    else:
        self.cutoff_lr = None
    self.nb_threshold = nb_threshold
    self.cache_static = bool(cache_static)
    self._static_cache = StaticInputCache()
    # Identity of the last species-validated `numbers` tensor, so repeated
    # evals on the same buffer (MD/optimization loops) skip the per-step
    # D2H tolist() + set-diff in _validate_species_and_charge.
    self._species_validation_cache: tuple[tuple[Any, ...], Any, Any] | None = None

    # Create adaptive neighbor list instances
    self._nblist = AdaptiveNeighborList(cutoff=self.cutoff)

    # Track separate cutoffs for LR modules
    self._coulomb_cutoff: float | None = None
    self._dftd3_cutoff: float = self._default_dftd3_cutoff
    if self.external_coulomb is not None:
        if self.external_coulomb.method == "simple":
            self._coulomb_cutoff = float("inf")
        elif self.external_coulomb.method in ("ewald", "pme"):
            self._coulomb_cutoff = None  # Ewald/PME manage their own cutoff
        else:
            self._coulomb_cutoff = self.external_coulomb.dsf_rc
    elif self._has_embedded_coulomb():
        self._coulomb_cutoff = float("inf")
    if self.external_dftd3 is not None:
        self._dftd3_cutoff = self.external_dftd3.smoothing_off

    # Create long-range neighbor list(s) if LR modules present
    self._nblist_lr: AdaptiveNeighborList | None = None
    self._nblist_dftd3: AdaptiveNeighborList | None = None
    self._nblist_coulomb: AdaptiveNeighborList | None = None
    self._update_lr_nblists()

    # indicator if input was flattened
    self._batch: int | None = None
    self._max_mol_size: int = 0
    self._static_input_cache_key: tuple[Any, ...] | None = None
    self._static_input_cache_refs: tuple[Any, Any] | None = None
    # placeholder for tensors that require grad
    self._saved_for_grad: dict[str, Tensor] = {}
    # set flag of current Coulomb method
    self._coulomb_method: str | None = None
    if self.external_coulomb is not None:
        self._coulomb_method = self.external_coulomb.method
    elif self._has_embedded_coulomb():
        # Legacy models have embedded Coulomb with "simple" method
        self._coulomb_method = "simple"
    # Bookkeeping for the per-evaluation simple->dsf PBC auto-switch:
    # prepare_input stores the pre-switch Coulomb state here and eval()
    # restores it in its finally block.
    self._pbc_coulomb_restore: dict[str, Any] | None = None
    # Memoized DSF-side state from the auto-switch (incl. warmed-up adaptive
    # neighbor lists) so repeated periodic evals don't rebuild it every step.
    self._auto_dsf_state: dict[str, Any] | None = None
    # One-shot flag for the ignored-multiplicity warning (eval fires inside MD loops).
    self._mult_ignored_checked = False

    # Set training mode (default False for inference)
    self._train = train
    # Deterministic LR mode: route external D3 (and DSF Coulomb) through
    # their differentiable pure-torch paths, whose reductions are run-to-run
    # deterministic, instead of the atomics-based nvalchemiops kernels.
    self._deterministic = bool(deterministic)
    self._warned_deterministic_lr = False
    self.model.train(train)
    if not train:
        # Disable gradients on all parameters for inference mode
        for param in self.model.parameters():
            param.requires_grad_(False)
        if self.external_coulomb is not None:
            for param in self.external_coulomb.parameters():
                param.requires_grad_(False)
        if self.external_dftd3 is not None:
            for param in self.external_dftd3.parameters():
                param.requires_grad_(False)

    self._maybe_warn_family_mix((metadata or {}).get("family") if metadata else None)

coulomb_cutoff property

Get the current Coulomb cutoff distance.

Returns

float | None The cutoff distance for Coulomb calculations, or None if not applicable. For "simple" this is inf; for "ewald" and "pme" this is None (cutoff is estimated per call from ewald_accuracy). Use set_lrcoulomb_method() to change.

coulomb_method property

Get the current Coulomb method.

Returns

str | None One of "simple", "dsf", "ewald", "pme", or None if no external Coulomb. For legacy models with embedded Coulomb, returns None.

dftd3_cutoff property

Get the current DFTD3 cutoff distance.

Returns

float The cutoff distance for DFTD3 calculations in Angstroms.

has_external_coulomb property

Check if calculator has external Coulomb module attached.

Returns True for new-format models that were trained with Coulomb and have it externalized. For legacy models, Coulomb is embedded in the model itself, so this returns False.

has_external_dftd3 property

Check if calculator has external DFTD3 module attached.

Returns True for new-format models that were trained with DFTD3/D3BJ dispersion and have it externalized. For legacy models or D3TS models, dispersion is embedded in the model itself, so this returns False.

is_nse property

Return True if the model supports spin-polarized charges (NSE, num_charge_channels=2).

metadata property

Read-only view of the model's metadata dict.

Returns a read-only mapping for v2 .pt models, or None for raw nn.Module inputs that don't carry metadata. Downstream consumers should prefer this accessor over reaching into the private model._metadata attribute.

eval(data, forces=False, stress=False, hessian=False, *, validate_species=True)

Run the model on data and return the output dict.

For a single structure each value is a Tensor. For batched Hessian requests the output is collected per structure: a 3D coord batch (B, N, 3) yields values with a new leading batch dim (stacked), while a multi-molecule mol_idx input yields a per-molecule list[Tensor] for each key (molecules are independent and generally ragged).

Source code in aimnet/calculators/calculator.py

def eval(
    self, data: dict[str, Any], forces=False, stress=False, hessian=False, *, validate_species: bool = True
) -> dict[str, Any]:
    """Run the model on ``data`` and return the output dict.

    For a single structure each value is a ``Tensor``. For batched Hessian
    requests the output is collected per structure: a 3D ``coord`` batch
    (B, N, 3) yields values with a new leading batch dim (stacked), while a
    multi-molecule ``mol_idx`` input yields a per-molecule ``list[Tensor]``
    for each key (molecules are independent and generally ragged).
    """
    # Species validation — opt-out via validate_species=False.
    if validate_species:
        self._validate_species_and_charge(data)
    # Warn once if the caller requests an open-shell `mult` this model ignores.
    self._maybe_warn_mult_ignored(data)
    # Hessian + torch.compile is known to hang on the double-backward
    # path through GELU activations. Fail fast instead.
    if hessian and getattr(self, "_was_compiled", False):
        raise RuntimeError(
            "Hessian computation is incompatible with compile_model=True "
            "(Dynamo + double-backward through GELU hangs). Reconstruct calculator "
            "with compile_model=False."
        )

    if hessian:
        subsystems = self._split_hessian_batch(data)
        if subsystems is not None:
            stack = torch.as_tensor(data["coord"]).ndim == 3
            return self._eval_hessian_batched(
                subsystems, forces=forces, stress=stress, validate_species=validate_species, stack=stack
            )

    # The simple->dsf PBC auto-switch in prepare_input is scoped to this
    # evaluation: any pending restore is consumed in the finally block, so
    # an exception mid-eval cannot leave the calculator on the switched method.
    self._pbc_coulomb_restore = None
    try:
        data = self.prepare_input(data, hessian=hessian)

        if hessian and "mol_idx" in data and data["mol_idx"][-1] > 0:
            raise NotImplementedError(
                "In-call Hessian for hand-flattened multi-molecule input is not supported; "
                "pass a 3D batch or a mol_idx dict to get per-structure Hessians."
            )
        data = self.set_grad_tensors(data, forces=forces, stress=stress, hessian=hessian)
        if isinstance(self.model, torch.jit.ScriptModule):
            with torch.jit.optimized_execution(False):  # type: ignore
                data = self.model(data)
        else:
            data = self.model(data)
        # Run external modules if present
        data, coulomb_terms = self._run_external_modules(
            data, forces=forces or hessian, stress=stress, hessian=hessian
        )
        data = self.get_derivatives(
            data, forces=forces, stress=stress, hessian=hessian, coulomb_terms=coulomb_terms
        )
        data = self.process_output(data)
        return data
    finally:
        restore = self._pbc_coulomb_restore
        if restore is not None:
            self._pbc_coulomb_restore = None
            # Keep the DSF-side state (incl. warmed-up adaptive neighbor
            # lists) for the next periodic call before flipping back to the
            # pre-switch method.
            self._auto_dsf_state = self._snapshot_coulomb_state()
            self._restore_coulomb_state(restore)

from_legacy_jit(path, *, device=None, **calculator_kwargs) classmethod

Construct a calculator from a trusted legacy TorchScript model.

path must point to a trusted .jpt source. Additional keyword arguments are forwarded to :class:AIMNet2Calculator; model is rejected because the model is supplied by path.

Source code in aimnet/calculators/calculator.py

@classmethod
def from_legacy_jit(
    cls,
    path: str,
    *,
    device: str | None = None,
    **calculator_kwargs: Any,
) -> Self:
    """Construct a calculator from a trusted legacy TorchScript model.

    ``path`` must point to a trusted ``.jpt`` source. Additional keyword
    arguments are forwarded to :class:`AIMNet2Calculator`; ``model`` is
    rejected because the model is supplied by ``path``.
    """
    if "model" in calculator_kwargs:
        raise TypeError("from_legacy_jit() does not accept a model keyword argument.")
    if not uses_default_model_import_settings(
        calculator_kwargs.get("model_import_paths"),
        calculator_kwargs.get("model_import_mode", "extend"),
    ):
        raise ValueError("Import settings are not supported for .jpt sources.")
    resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    resolved_device = str(torch.device(resolved_device))
    loaded_model, _ = load_legacy_jit(path, resolved_device)
    return cls(model=loaded_model, device=resolved_device, **calculator_kwargs)

hessian_vector_product(data, vectors, *, eps=None, validate_species=True, create_graph=False)

Matrix-free Hessian-vector product(s) H @ v for one structure.

Computes H @ v without forming the dense (N, 3, N, 3) Hessian, enabling Lanczos/LOBPCG negative-curvature checks and CG-Newton preconditioning on large systems.

Parameters

data : dict Single-structure input (same keys as eval). 3D batched or multi-molecule mol_idx inputs are not supported. vectors : Tensor Direction(s), shape (N, 3) or (K, N, 3) over the real atoms. eps : float, optional Deprecated and unused since the PME energy-graph migration (every backend's product is exact reverse-mode autograd). Passing any value emits a DeprecationWarning. create_graph : bool If True, keep the HVP in the graph so it can compose with an outer loss. Default False preserves the numeric/detached operator-action behavior.

Returns

Tensor H @ v, shape (N, 3) or (K, N, 3), matching vectors.

Notes

The product is an exact reverse-mode autograd computation for every backend: the NN, short-range, simple/dsf Coulomb, DFTD3, and periodic ewald/pme energies are all in the autograd graph, so the vjp captures the full curvature, including the relaxed-charge response d^2E/(dq.dr). This mirrors the dense :meth:calculate_hessian assembly, so hessian_vector_product(v) equals H.reshape(3N, 3N) @ v to the backend's tolerance. The default return is detached; set create_graph=True when the HVP must compose with an outer computation. See :meth:calculate_hessian for the detached-Hessian contract and the fully-differentiable recipe.

Integration note: the external modules run via _run_external_modules(forces=False, hessian=(method == 'dsf')) -- ENERGY-GRAPH mode so every external term (DFTD3 + Coulomb) stays differentiable w.r.t. coord and its curvature is captured by the autograd vjp. forces=False (not True) is required: with forces=True the DFTD3 branch takes its detached explicit-force path and its second-derivative curvature is silently dropped from H @ v. The hessian flag routes dsf through its differentiable closed-form torch path (_coul_dsf_torch); ewald/pme energies are always in the graph.

Dtype / differentiability caveats:

• Return dtype is the model dtype (typically float32) for simple and dsf, and float64 for ewald/pme (preserving the contract established when the periodic block was finite-difference).
• With create_graph=False the returned product is detached numeric operator action. With create_graph=True the HVP remains differentiable w.r.t. graph-attached coordinates / model parameters; vectors are still treated as numeric directions.

Source code in aimnet/calculators/calculator.py

@torch.inference_mode(False)
@torch.enable_grad()
def hessian_vector_product(
    self,
    data: dict[str, Any],
    vectors: Tensor,
    *,
    eps: float | None = None,
    validate_species: bool = True,
    create_graph: bool = False,
) -> Tensor:
    """Matrix-free Hessian-vector product(s) ``H @ v`` for one structure.

    Computes ``H @ v`` without forming the dense ``(N, 3, N, 3)`` Hessian,
    enabling Lanczos/LOBPCG negative-curvature checks and CG-Newton
    preconditioning on large systems.

    Parameters
    ----------
    data : dict
        Single-structure input (same keys as ``eval``). 3D batched or
        multi-molecule ``mol_idx`` inputs are not supported.
    vectors : Tensor
        Direction(s), shape ``(N, 3)`` or ``(K, N, 3)`` over the real atoms.
    eps : float, optional
        Deprecated and unused since the PME energy-graph migration (every
        backend's product is exact reverse-mode autograd). Passing any
        value emits a ``DeprecationWarning``.
    create_graph : bool
        If ``True``, keep the HVP in the graph so it can compose with an
        outer loss. Default ``False`` preserves the numeric/detached
        operator-action behavior.

    Returns
    -------
    Tensor
        ``H @ v``, shape ``(N, 3)`` or ``(K, N, 3)``, matching ``vectors``.

    Notes
    -----
    The product is an exact reverse-mode autograd computation for every
    backend: the NN, short-range, ``simple``/``dsf`` Coulomb, DFTD3, and
    periodic ``ewald``/``pme`` energies are all in the autograd graph, so
    the vjp captures the full curvature, including the relaxed-charge
    response ``d^2E/(dq.dr)``. This mirrors the dense
    :meth:`calculate_hessian` assembly, so ``hessian_vector_product(v)``
    equals ``H.reshape(3N, 3N) @ v`` to the backend's tolerance. The
    default return is detached; set ``create_graph=True`` when the HVP
    must compose with an outer computation. See :meth:`calculate_hessian`
    for the detached-Hessian contract and the fully-differentiable recipe.

    Integration note: the external modules run via
    ``_run_external_modules(forces=False, hessian=(method == 'dsf'))`` --
    ENERGY-GRAPH mode so every external term (DFTD3 + Coulomb) stays
    differentiable w.r.t. ``coord`` and its curvature is captured by the
    autograd vjp. ``forces=False`` (not ``True``) is required: with
    ``forces=True`` the DFTD3 branch takes its detached explicit-force path
    and its second-derivative curvature is silently dropped from ``H @ v``.
    The ``hessian`` flag routes dsf through its differentiable closed-form
    torch path (``_coul_dsf_torch``); ewald/pme energies are always in the
    graph.

    Dtype / differentiability caveats:

    * Return dtype is the model dtype (typically float32) for ``simple`` and
      ``dsf``, and **float64** for ``ewald``/``pme`` (preserving the
      contract established when the periodic block was finite-difference).
    * With ``create_graph=False`` the returned product is detached numeric
      operator action. With ``create_graph=True`` the HVP remains
      differentiable w.r.t. graph-attached coordinates / model parameters;
      vectors are still treated as numeric directions.
    """
    if eps is not None:
        warnings.warn(
            "hessian_vector_product(eps=...) is unused since the PME energy-graph "
            "migration (all backends are exact reverse-mode autograd) and will be removed.",
            DeprecationWarning,
            stacklevel=2,
        )
    if getattr(self, "_was_compiled", False):
        raise RuntimeError(
            "hessian_vector_product is incompatible with compile_model=True "
            "(Dynamo + double-backward through GELU hangs). Reconstruct with compile_model=False."
        )
    # Same species/charge validation contract as `eval` (opt-out via
    # validate_species=False); otherwise unsupported elements / charged
    # systems would yield undefined output silently.
    if validate_species:
        self._validate_species_and_charge(data)
    # Warn once if the caller requests an open-shell `mult` this model ignores.
    self._maybe_warn_mult_ignored(data)
    coord_in = torch.as_tensor(data["coord"])
    if coord_in.ndim == 3 and coord_in.shape[0] > 1:
        raise NotImplementedError("hessian_vector_product supports a single structure only (got 3D batch).")
    if coord_in.ndim == 2 and data.get("mol_idx") is not None:
        mol_idx_t = torch.as_tensor(data["mol_idx"])
        if mol_idx_t.numel() and int(mol_idx_t.max()) > 0:
            raise NotImplementedError(
                "hessian_vector_product supports a single structure only (got mol_idx batch)."
            )

    # prepare_input / set_grad_tensors mutate instance scratch state and can
    # persistently flip Coulomb method/cutoff/list-builder state via
    # prepare_input's simple->dsf PBC auto-switch. Snapshot now and restore
    # in the finally so a single HVP call leaves the calculator unmodified.
    # The per-vector loop reads self._saved_for_grad and self.external_coulomb,
    # so the result must be fully computed inside the try before state is
    # restored.
    eval_state = self._snapshot_eval_state()
    try:
        result = self._hessian_vector_product_impl(data, vectors, create_graph=create_graph)
    finally:
        self._restore_eval_state(eval_state)
        # _restore_eval_state already undid any PBC auto-switch from
        # prepare_input; drop the (now-redundant) pending restore marker.
        self._pbc_coulomb_restore = None
    return result

mol_flatten(data, *, hessian=False)

Flatten the input data for multiple molecules. Will not flatten for batched input and molecule size below threshold.

Source code in aimnet/calculators/calculator.py

def mol_flatten(self, data: dict[str, Tensor], *, hessian: bool = False) -> dict[str, Tensor]:
    """Flatten the input data for multiple molecules.
    Will not flatten for batched input and molecule size below threshold.
    """
    ndim = data["coord"].ndim
    if ndim == 2:
        self._batch = None
        if "mol_idx" not in data:
            data["mol_idx"] = torch.zeros(data["coord"].shape[0], dtype=torch.long, device=self.device)
            self._max_mol_size = data["coord"].shape[0]
        elif data["mol_idx"][-1] == 0:
            self._max_mol_size = len(data["mol_idx"])
        else:
            self._max_mol_size = data["mol_idx"].unique(return_counts=True)[1].max().item()

    elif ndim == 3:
        B, N = data["coord"].shape[:2]
        if hessian and B != 1:
            raise NotImplementedError("Hessian calculation is not supported for batched inputs with B > 1")
        # Force flattening for PBC (cell present) to ensure make_nbmat computes proper neighbor lists with shifts
        if (
            hessian
            or self.nb_threshold < N
            or torch.device(self.device).type == "cpu"
            or data.get("cell") is not None
        ):
            self._batch = B
            data["mol_idx"] = torch.repeat_interleave(
                torch.arange(0, B, device=self.device), torch.full((B,), N, device=self.device)
            )
            for k, v in data.items():
                if k in self.atom_feature_keys:
                    data[k] = v.flatten(0, 1)
        else:
            self._batch = None
        self._max_mol_size = N
    return data

set_dftd3_cutoff(cutoff=None, smoothing_fraction=None)

Set DFTD3 cutoff and smoothing.

Parameters

cutoff : float | None Cutoff distance in Angstroms for DFTD3 calculation. Default is _default_dftd3_cutoff (15.0). smoothing_fraction : float | None Fraction of cutoff used as smoothing width. Default is _default_dftd3_smoothing (0.2).

Notes

This method only affects external DFTD3 modules attached to new-format models. For legacy models with embedded DFTD3, the smoothing is fixed.

Updates _dftd3_cutoff and rebuilds neighbor lists.

Source code in aimnet/calculators/calculator.py

def set_dftd3_cutoff(self, cutoff: float | None = None, smoothing_fraction: float | None = None) -> None:
    """Set DFTD3 cutoff and smoothing.

    Parameters
    ----------
    cutoff : float | None
        Cutoff distance in Angstroms for DFTD3 calculation.
        Default is _default_dftd3_cutoff (15.0).
    smoothing_fraction : float | None
        Fraction of cutoff used as smoothing width.
        Default is _default_dftd3_smoothing (0.2).

    Notes
    -----
    This method only affects external DFTD3 modules attached to
    new-format models. For legacy models with embedded DFTD3,
    the smoothing is fixed.

    Updates _dftd3_cutoff and rebuilds neighbor lists.
    """
    if cutoff is None:
        cutoff = self._default_dftd3_cutoff
    if smoothing_fraction is None:
        smoothing_fraction = self._default_dftd3_smoothing

    self._dftd3_cutoff = cutoff
    if self.external_dftd3 is not None:
        self.external_dftd3.set_smoothing(cutoff, smoothing_fraction)
    # Cutoffs changed; the memoized auto-switch DSF state is stale.
    self._auto_dsf_state = None
    self._update_lr_nblists()
    self._clear_static_nbmat_cache()

set_lr_cutoff(cutoff)

Set the unified long-range cutoff for all LR modules.

Parameters

cutoff : float Cutoff distance in Angstroms for LR neighbor lists.

Notes

This updates both _coulomb_cutoff and _dftd3_cutoff. Ewald/PME use their own per-call neighbor lists and ignore this cutoff.

Source code in aimnet/calculators/calculator.py

def set_lr_cutoff(self, cutoff: float) -> None:
    """Set the unified long-range cutoff for all LR modules.

    Parameters
    ----------
    cutoff : float
        Cutoff distance in Angstroms for LR neighbor lists.

    Notes
    -----
    This updates both _coulomb_cutoff and _dftd3_cutoff.
    Ewald/PME use their own per-call neighbor lists and ignore this cutoff.
    """
    # Update both cutoffs (but not for ewald/pme which manage their own)
    if self._coulomb_method not in ("ewald", "pme"):
        self._coulomb_cutoff = cutoff
    self._dftd3_cutoff = cutoff
    self.cutoff_lr = cutoff
    # Cutoffs changed; the memoized auto-switch DSF state is stale.
    self._auto_dsf_state = None
    self._update_lr_nblists()
    self._clear_static_nbmat_cache()

set_lrcoulomb_method(method, cutoff=15.0, dsf_alpha=0.2, ewald_accuracy=1e-06)

Set the long-range Coulomb method.

Parameters

method : str One of "simple", "dsf", "ewald", or "pme". cutoff : float Cutoff distance for DSF neighbor list. Default is 15.0. Silently ignored for "ewald" and "pme" (which estimate their own real-space cutoffs from ewald_accuracy). dsf_alpha : float Alpha parameter for DSF method. Default is 0.2. ewald_accuracy : float Target accuracy for Ewald and PME summation. Controls the real-space and reciprocal-space cutoffs (and PME mesh dimensions). Smaller values give higher accuracy at the cost of more computation. Default is 1e-6, matching the nvalchemiops default.

The Ewald cutoffs follow the Kolafa-Perram formula:
- eta = (V^2 / N)^(1/6) / sqrt(2*pi)
- cutoff_real = sqrt(-2 * ln(accuracy)) * eta
- cutoff_recip = sqrt(-2 * ln(accuracy)) / eta

Notes

For new-format models with external Coulomb, this updates the external module. For legacy models with embedded Coulomb, a warning is issued as those modules cannot be modified at runtime.

"ewald" and "pme" both require periodic systems (cell set); invoking the calculator without a cell raises ValueError at prepare_input.

Hessian note: "dsf", "ewald", and "pme" Hessians are all computed under the same relaxed-charge contract (fully autograd, including the charge-response coupling d^2E/(dq.dr) through the model's predicted charges). Ewald and PME evaluate the same lattice sum and are directly comparable; DSF remains a shifted-force truncation of the PES, so its curvature still differs from the full lattice sum for pairs near the cutoff.

"pme" requires nvalchemi-toolkit-ops >= 0.4.1 and raises RuntimeError on older versions (0.4.0 silently corrupts train-mode PME forces).

Source code in aimnet/calculators/calculator.py

def set_lrcoulomb_method(
    self,
    method: Literal["simple", "dsf", "ewald", "pme"],
    cutoff: float = 15.0,
    dsf_alpha: float = 0.2,
    ewald_accuracy: float = 1e-6,
):
    """Set the long-range Coulomb method.

    Parameters
    ----------
    method : str
        One of "simple", "dsf", "ewald", or "pme".
    cutoff : float
        Cutoff distance for DSF neighbor list. Default is 15.0.
        Silently ignored for "ewald" and "pme" (which estimate their own
        real-space cutoffs from ``ewald_accuracy``).
    dsf_alpha : float
        Alpha parameter for DSF method. Default is 0.2.
    ewald_accuracy : float
        Target accuracy for Ewald and PME summation. Controls the
        real-space and reciprocal-space cutoffs (and PME mesh dimensions).
        Smaller values give higher accuracy at the cost of more
        computation. Default is 1e-6, matching the nvalchemiops default.

        The Ewald cutoffs follow the Kolafa-Perram formula:
        - eta = (V^2 / N)^(1/6) / sqrt(2*pi)
        - cutoff_real = sqrt(-2 * ln(accuracy)) * eta
        - cutoff_recip = sqrt(-2 * ln(accuracy)) / eta

    Notes
    -----
    For new-format models with external Coulomb, this updates the external module.
    For legacy models with embedded Coulomb, a warning is issued as those modules
    cannot be modified at runtime.

    ``"ewald"`` and ``"pme"`` both require periodic systems (``cell`` set);
    invoking the calculator without a cell raises ``ValueError`` at
    ``prepare_input``.

    Hessian note: ``"dsf"``, ``"ewald"``, and ``"pme"`` Hessians are all
    computed under the same relaxed-charge contract (fully autograd,
    including the charge-response coupling ``d^2E/(dq.dr)`` through the
    model's predicted charges). Ewald and PME evaluate the same lattice
    sum and are directly comparable; DSF remains a shifted-force
    truncation of the PES, so its curvature still differs from the full
    lattice sum for pairs near the cutoff.

    ``"pme"`` requires nvalchemi-toolkit-ops >= 0.4.1 and raises
    ``RuntimeError`` on older versions (0.4.0 silently corrupts train-mode
    PME forces).
    """
    if method not in ("simple", "dsf", "ewald", "pme"):
        raise ValueError(f"Invalid method: {method}")
    if method == "pme":
        _require_pme_capable_nvalchemiops()

    # Warn if model has embedded Coulomb (legacy models)
    if self._has_embedded_coulomb() and self.external_coulomb is None:
        warnings.warn(
            "Model has embedded Coulomb module (legacy format). "
            "set_lrcoulomb_method() only affects external Coulomb modules. "
            "For legacy models, the Coulomb method cannot be changed at runtime.",
            stacklevel=2,
        )
        return

    # Update external LRCoulomb module if present
    if self.external_coulomb is not None:
        self.external_coulomb.method = method
        if method == "dsf":
            self.external_coulomb.dsf_alpha = dsf_alpha
            self.external_coulomb.dsf_rc = cutoff
        elif method in ("ewald", "pme"):
            self.external_coulomb.ewald_accuracy = ewald_accuracy

    # Update _coulomb_cutoff based on method
    if method == "simple":
        self._coulomb_cutoff = float("inf")
    elif method == "dsf":
        self._coulomb_cutoff = cutoff
    elif method in ("ewald", "pme"):
        # Ewald/PME estimate their own real-space cutoff per call.
        self._coulomb_cutoff = None

    # Update cutoff_lr for backward compatibility
    if self._coulomb_cutoff is not None:
        self.cutoff_lr = self._coulomb_cutoff
    else:
        # Ewald/PME - use DFTD3 cutoff if available, else None
        self.cutoff_lr = self._dftd3_cutoff if self.external_dftd3 is not None else None

    self._coulomb_method = method
    # Method/cutoff changed; the memoized auto-switch DSF state is stale.
    self._auto_dsf_state = None
    self._update_lr_nblists()
    self._clear_static_nbmat_cache()

AIMNet2ASE

ASE (Atomic Simulation Environment) calculator interface.

Installation

Requires the ase extra: pip install aimnet[ase]

This calculator integrates with ASE's Atoms object, supporting energy, forces, stress, and dipole moment calculations. It operates in eV and Angstrom.

Usage Example

from ase.io import read
from aimnet.calculators import AIMNet2ASE

atoms = read("molecule.xyz")
atoms.calc = AIMNet2ASE("aimnet2")

print(atoms.get_potential_energy())
print(atoms.get_forces())

AIMNet2ASE(base_calc='aimnet2', charge=0, mult=1, validate_species=True)

Bases: Calculator

Source code in aimnet/calculators/aimnet2ase.py

def __init__(
    self,
    base_calc: AIMNet2Calculator | str = "aimnet2",
    charge=0,
    mult=1,
    validate_species: bool = True,
):
    if _ASE_IMPORT_ERROR is not None:
        raise ImportError(
            'AIMNet2ASE requires ASE. Install with `pip install "aimnet[ase]"` — or on conda: conda install -c conda-forge ase'
        ) from _ASE_IMPORT_ERROR

    super().__init__()
    if isinstance(base_calc, str):
        base_calc = AIMNet2Calculator(base_calc)
    self.base_calc = base_calc
    self.validate_species = validate_species
    if self.base_calc.is_nse:
        self.__dict__["implemented_properties"] = [*self.__class__.implemented_properties, "spin_charges"]
    self.reset()
    self.charge = charge
    self.mult = mult
    self.update_tensors()
    # list of implemented species — read from model metadata
    _meta = getattr(base_calc.model, "_metadata", None)
    _species = _meta.get("implemented_species") if _meta is not None else None
    self.implemented_species = np.array(_species, dtype=np.int64) if _species else None

get_hessian(atoms=None)

Return Cartesian Hessian as a (3N, 3N) ndarray in eV/Å^2.

Designed for use as Sella(atoms, hessian_function=atoms.calc.get_hessian). Computed via double-backward through the AIMNet2 energy graph; cost scales as O(3N) backward passes per call. Not supported when compile_model=True or for batched / multi-molecule input.

This method intentionally bypasses the standard ASE Calculator.calculate(properties=['hessian']) flow and self.results cache. The Sella callback contract is (atoms) -> ndarray, so a direct method is the simplest match. "hessian" is therefore not advertised in implemented_properties; if that ever changes, the two paths must be reconciled.

When called with an explicit atoms argument that differs from self.atoms, the passed atoms.info is consulted for charge/mult precedence (and the calculator's stored self.charge/self.mult may be updated as a side effect, mirroring the calculate() behavior).

Source code in aimnet/calculators/aimnet2ase.py

def get_hessian(self, atoms=None):
    """Return Cartesian Hessian as a (3N, 3N) ndarray in eV/Å^2.

    Designed for use as ``Sella(atoms, hessian_function=atoms.calc.get_hessian)``.
    Computed via double-backward through the AIMNet2 energy graph; cost scales
    as O(3N) backward passes per call. Not supported when ``compile_model=True``
    or for batched / multi-molecule input.

    This method intentionally bypasses the standard ASE
    ``Calculator.calculate(properties=['hessian'])`` flow and ``self.results``
    cache. The Sella callback contract is ``(atoms) -> ndarray``, so a direct
    method is the simplest match. ``"hessian"`` is therefore not advertised in
    ``implemented_properties``; if that ever changes, the two paths must be
    reconciled.

    When called with an explicit ``atoms`` argument that differs from
    ``self.atoms``, the passed ``atoms.info`` is consulted for charge/mult
    precedence (and the calculator's stored ``self.charge``/``self.mult`` may
    be updated as a side effect, mirroring the ``calculate()`` behavior).
    """
    if atoms is None:
        atoms = getattr(self, "atoms", None)
        if atoms is None:
            raise PropertyNotImplementedError(
                "get_hessian() requires an attached Atoms object or an explicit argument."
            )
    if atoms.pbc.any():
        raise PropertyNotImplementedError(
            "Hessian for periodic systems is not supported by AIMNet2ASE.get_hessian(). "
            "For periodic transition states, use pysisyphus dimer or climbing-image NEB."
        )
    if len(atoms) > 100:
        warnings.warn(
            f"Computing AIMNet2 Hessian for {len(atoms)} atoms; "
            "the forces+hessian path retains a much larger autograd graph than forces alone "
            "(peak GPU memory is roughly 5-10x a forces-only call). Risk of OOM on smaller GPUs.",
            stacklevel=2,
        )

    self._update_charge_spin_from_info(atoms)
    self.update_tensors(atoms)

    # Pass coord as 2D (N, 3) — not batched — so mol_flatten takes the
    # ndim==2 path and calculate_hessian sees the expected (N+1, 3) coord
    # after padding. Batching (unsqueeze(0)) triggers the ndim==3 path
    # which may skip flattening when N < nb_threshold on GPU, causing
    # calculate_hessian to produce an incorrect (N, 3, 1, 3) shape.
    coord = torch.tensor(atoms.positions, dtype=self.base_calc.keys_in["coord"], device=self.base_calc.device)
    _in = {
        "coord": coord,
        "numbers": self._t_numbers,
        "charge": self._t_charge,
        "mult": self._t_mult,
    }

    results = self.base_calc(
        _in,
        forces=True,
        hessian=True,
        validate_species=self.validate_species,
    )
    H = results["hessian"].detach()  # (N, 3, N, 3)
    N = H.shape[0]
    return H.reshape(N * 3, N * 3).cpu().numpy()

AIMNet2Pysis

PySisyphus calculator interface.

Installation

Requires the pysis extra: pip install aimnet[pysis]

This interface adapts AIMNet2 for use with PySisyphus optimizers. It handles unit conversion automatically:

• Input: Converts Bohr coordinates (PySisyphus) to Angstrom (AIMNet2).
• Output: Converts eV and eV/Angstrom (AIMNet2) to Hartree and Hartree/Bohr (PySisyphus).
• Hessian: Converts eV/Angstrom^2 (AIMNet2) to Hartree/Bohr^2 (PySisyphus).

AIMNet2Pysis(model='aimnet2', charge=0, mult=1, validate_species=True, **kwargs)

Bases: Calculator

Source code in aimnet/calculators/aimnet2pysis.py

def __init__(
    self, model: AIMNet2Calculator | str = "aimnet2", charge=0, mult=1, validate_species: bool = True, **kwargs
):
    if _PYSIS_IMPORT_ERROR is not None:
        raise ImportError(
            "AIMNet2Pysis requires PySisyphus. Install it with `pip install pysisyphus` (pysisyphus is not available on conda-forge)."
        ) from _PYSIS_IMPORT_ERROR

    super().__init__(charge=charge, mult=mult, **kwargs)
    if isinstance(model, str):
        model = AIMNet2Calculator(model)
    self.model = model
    self.validate_species = validate_species
    # AFIR and some IRC paths call get_forces then get_energy at the same coord;
    # the cache serves the second call without a redundant model forward.
    self._cache_key: tuple[tuple[str, ...], bytes] | None = None
    self._cache_results: dict | None = None

AIMNet2TorchSim

TorchSim ModelInterface wrapper.

Installation

Requires the torchsim extra and Python 3.12+: pip install "aimnet[torchsim]". Add the ase extra for ASE-based input/output examples: pip install "aimnet[torchsim,ase]". On Python 3.11, the base AIMNet package is supported but the TorchSim extra is not installed.

AIMNet2TorchSim wraps an AIMNet2Calculator as a torch-sim-atomistic model for static evaluation, geometry optimization, molecular dynamics, and autobatched workloads.

Usage Example

import ase.io
import torch_sim as ts

from aimnet.calculators import AIMNet2Calculator, AIMNet2TorchSim

atoms = ase.io.read("molecule.xyz")

base_calc = AIMNet2Calculator("aimnet2")
calc = AIMNet2TorchSim(base_calc)

results = ts.static(system=atoms, model=calc)
print(results[0]["potential_energy"], results[0]["forces"])

Stress

By default compute_stress=False. Pass compute_stress=True when constructing AIMNet2TorchSim for NPT integrators and PBC cell relaxation.

TorchSim extras

AIMNet partial charges are returned as both charges and partial_charges output fields. Set per-system charge and NSE mult through TorchSim system extras.

AIMNet2TorchSim(base_calc, *, compute_forces=True, compute_stress=False, validate_species=True)

Bases: ModelInterface

Wrap an :class:AIMNet2Calculator as a TorchSim model.

Parameters

base_calc Underlying AIMNet2 calculator. AIMNet2 inference uses float32 internally, so the wrapper reports torch.float32 regardless of the incoming TorchSim state dtype. compute_forces Request AIMNet2 forces on each forward call. Keep this enabled for geometry optimization and molecular dynamics. Set it false only for energy-only static batches. compute_stress Request AIMNet2 stress on every forward call. This is required for NPT integrators and PBC cell relaxation. Leave it false for energy/force workflows to avoid retaining extra autograd state. validate_species Forward AIMNet2 calculator species and charge-domain validation. Leave enabled unless intentionally bypassing model metadata checks.

Source code in aimnet/calculators/aimnet2torchsim.py

def __init__(
    self,
    base_calc: AIMNet2Calculator,
    *,
    compute_forces: bool = True,
    compute_stress: bool = False,
    validate_species: bool = True,
) -> None:
    if _TORCHSIM_IMPORT_ERROR is not None:
        raise ImportError(
            "AIMNet2TorchSim requires TorchSim (Python 3.12+ only). Install it with `pip install torch-sim-atomistic` (torch-sim-atomistic is not available on conda-forge)."
        ) from _TORCHSIM_IMPORT_ERROR

    super().__init__()
    self._base_calc = base_calc
    self._device = torch.device(base_calc.device)
    self._dtype = torch.float32
    self._compute_forces = compute_forces
    self._compute_stress = compute_stress
    self._validate_species = validate_species
    self._memory_scales_with = "n_atoms_x_density"
    self._update_implemented_properties()

base_calc property

Underlying AIMNet2 calculator.

metadata property

Underlying model metadata, when available.

forward(state, **kwargs)

Compute AIMNet2 outputs for a TorchSim state.

Source code in aimnet/calculators/aimnet2torchsim.py

def forward(self, state: SimState, **kwargs: Any) -> dict[str, Tensor]:
    """Compute AIMNet2 outputs for a TorchSim state."""
    if state.device != self._device or state.dtype != self._dtype:
        state = state.to(self._device, self._dtype)

    data = self._state_to_aimnet2_data(state)
    results = self._base_calc(
        data,
        forces=self._compute_forces,
        stress=self._compute_stress,
        validate_species=self._validate_species,
    )
    if "charges" in results:
        results["partial_charges"] = results["charges"]
    return {key: value.detach() if torch.is_tensor(value) else value for key, value in results.items()}

Model Registry

Utilities for loading pre-trained models. Models are automatically downloaded from the remote repository to the local model cache (AIMNET_CACHE_DIR when set, otherwise ~/.cache/aimnet/) upon first use.

Every registry download and cache hit is verified against its SHA-256 digest. A corrupt cache hit triggers one verified atomic replacement attempt; persistent acquisition failures propagate after the final path is revalidated. See Model cache recovery.

Official wheels and source distributions currently bundle no model artifacts; registry models are downloaded on demand.

Bare registry names and aliases take precedence over same-named implicit local files or directories, preventing a local shadow from bypassing verification. Use an explicit path such as ./aimnet2 or an absolute path when a local artifact is intended.

Only direct local v2 artifacts and complete Hugging Face repositories accept custom import settings. See Model YAML import policy for supported paths, modes, and security boundaries.

CLI Command

You can clear the local model cache using the CLI:

aimnet clear_model_cache

model_registry

FamilyPolicy(family=None, supports_charged_systems=None, posthoc_d3_params=None, members=tuple()) dataclass

Calculator-side policy for a released model family.

Sourced from the families: section of model_registry.yaml. The neutral policy (all fields None/empty) applies no defaults and is returned for unknown or undeclared families, so raw nn.Module loads and third-party checkpoints are unaffected.

Attributes

family : str | None Canonical family tag, or None for the neutral policy. supports_charged_systems : bool | None Required charge-support declaration for the family. The calculator defaults undeclared model metadata to this value and rejects conflicting declarations. None means the family imposes no constraint. posthoc_d3_params : dict[str, float] | None DFT-D3(BJ) parameters applied post-hoc when the model does not embed dispersion. None means no post-hoc dispersion policy. members : tuple[str, ...] Registry model keys belonging to the family, in ensemble-member order.

download_assets(models, download_all)

Prefetch registry model weights for offline use.

MODELS are registry names or aliases (e.g. aimnet2). With --all, every registered model is fetched. Files are stored in the cache directory ($AIMNET_CACHE_DIR or ~/.cache/aimnet) and verified against their registry SHA-256 digests, so a pre-seeded cache works fully offline.

Source code in aimnet/calculators/model_registry.py

@click.command(short_help="Download model weights into the local cache.")
@click.argument("models", nargs=-1)
@click.option("--all", "download_all", is_flag=True, default=False, help="Download every model in the registry.")
def download_assets(models: tuple[str, ...], download_all: bool):
    """Prefetch registry model weights for offline use.

    MODELS are registry names or aliases (e.g. ``aimnet2``). With ``--all``,
    every registered model is fetched. Files are stored in the cache directory
    (``$AIMNET_CACHE_DIR`` or ``~/.cache/aimnet``) and verified against their
    registry SHA-256 digests, so a pre-seeded cache works fully offline.
    """
    registry = load_model_registry()
    if download_all:
        names = sorted(registry["models"])
    elif models:
        names = [resolve_registry_model_name(m) for m in models]
    else:
        known = "\n  ".join(sorted(registry["models"]))
        raise click.UsageError(f"Specify model names or --all. Known models:\n  {known}")
    failed = []
    for name in names:
        try:
            path = get_registry_model_path(name)
        except Exception as exc:
            click.echo(f"{name}: FAILED ({exc})", err=True)
            failed.append(name)
            continue
        click.echo(f"{name}: {path}")
    if failed:
        raise click.ClickException(f"{len(failed)} of {len(names)} downloads failed: {', '.join(failed)}")

get_cache_dir()

Return the model cache directory.

AIMNET_CACHE_DIR has priority. Otherwise AIMNet uses ~/.cache/aimnet. The directory is created on demand.

Source code in aimnet/calculators/model_registry.py

def get_cache_dir() -> str:
    """Return the model cache directory.

    ``AIMNET_CACHE_DIR`` has priority. Otherwise AIMNet uses
    ``~/.cache/aimnet``. The directory is created on demand.
    """
    cache_dir = os.environ.get("AIMNET_CACHE_DIR")
    if cache_dir is None:
        cache_dir = os.path.join(Path.home(), ".cache", "aimnet")
    os.makedirs(cache_dir, exist_ok=True)
    return cache_dir

get_family_policy(family)

Return the calculator-side policy for a model family.

Unknown or missing families return the neutral policy (no defaults applied) rather than raising, so calculators built from raw nn.Module inputs or third-party checkpoints keep working.

Source code in aimnet/calculators/model_registry.py

def get_family_policy(family: str | None) -> FamilyPolicy:
    """Return the calculator-side policy for a model family.

    Unknown or missing families return the neutral policy (no defaults applied)
    rather than raising, so calculators built from raw ``nn.Module`` inputs or
    third-party checkpoints keep working.
    """
    if family is None:
        return _NEUTRAL_FAMILY_POLICY
    model_registry = load_model_registry()
    block = model_registry.get("families", {}).get(family)
    if block is None:
        return _NEUTRAL_FAMILY_POLICY
    members = tuple(name for name, cfg in model_registry["models"].items() if cfg.get("family") == family)
    posthoc_d3_params = block.get("posthoc_d3_params")
    return FamilyPolicy(
        family=family,
        supports_charged_systems=block.get("supports_charged_systems"),
        posthoc_d3_params=dict(posthoc_d3_params) if posthoc_d3_params is not None else None,
        members=members,
    )

get_model_path(s)

Resolve a model name or explicit local path.

Bare registry names and aliases take precedence over same-named implicit local files. Absolute paths and paths beginning with ./ or ../ are always treated as explicit local paths.

Source code in aimnet/calculators/model_registry.py

def get_model_path(s: str) -> str:
    """Resolve a model name or explicit local path.

    Bare registry names and aliases take precedence over same-named implicit
    local files. Absolute paths and paths beginning with ``./`` or ``../`` are
    always treated as explicit local paths.
    """
    explicit_local = is_explicit_local_path(s)
    if not explicit_local and try_resolve_registry_model_name(s) is not None:
        return get_registry_model_path(s)
    if os.path.isfile(s):
        return s
    return get_registry_model_path(s)

get_registry_model_family(model_name)

Return the canonical family tag for a registered model name or alias.

Source code in aimnet/calculators/model_registry.py

def get_registry_model_family(model_name: str) -> str:
    """Return the canonical family tag for a registered model name or alias."""
    model_name = resolve_registry_model_name(model_name)
    model_registry = load_model_registry()
    family = model_registry["models"][model_name].get("family")
    if family is None:
        raise ValueError(f"Model {model_name} does not declare a family in the registry.")
    return family

get_registry_model_path(model_name)

Return a verified local path for a registered model artifact.

The registry entry must contain a lowercase 64-character SHA-256 digest. Cached, bundled, and downloaded bytes are verified against that digest before the path is returned.

Source code in aimnet/calculators/model_registry.py

def get_registry_model_path(model_name: str) -> str:
    """Return a verified local path for a registered model artifact.

    The registry entry must contain a lowercase 64-character SHA-256 digest.
    Cached, bundled, and downloaded bytes are verified against that digest
    before the path is returned.
    """
    model_registry = load_model_registry()
    model_name = resolve_registry_model_name(model_name)
    cfg = model_registry["models"][model_name]  # type: ignore
    expected_sha256 = cfg.get("sha256")
    if not isinstance(expected_sha256, str) or re.fullmatch(r"[0-9a-f]{64}", expected_sha256) is None:
        raise ValueError(f"Registry model {model_name!r} has no valid SHA-256 digest; refusing artifact access.")
    model_path = _maybe_download_asset(
        file=cfg["file"],
        url=cfg["url"],
        expected_sha256=expected_sha256,
    )
    return model_path

try_resolve_registry_model_name(model_name)

Resolve a registry name or alias without raising for unknown names.

Source code in aimnet/calculators/model_registry.py

def try_resolve_registry_model_name(model_name: str) -> str | None:
    """Resolve a registry name or alias without raising for unknown names."""
    model_registry = load_model_registry()
    if model_name in model_registry["aliases"]:
        model_name = model_registry["aliases"][model_name]  # type: ignore
    if model_name not in model_registry["models"]:
        return None
    return model_name