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: You can force specific long-range behavior using needs_coulomb and needs_dispersion arguments.
• 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, ensemble_member=0, revision=None, token=None)

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.

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,
    ensemble_member: int = 0,
    revision: str | None = None,
    token: str | None = None,
):
    # 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
    metadata: Mapping[str, Any] | None = None
    registry_family: str | None = None
    # Inline org/name pattern — exactly one slash, both segments alphanumeric+._-
    # This avoids importing optional HF deps for ordinary file paths containing slashes.
    _HF_ID_RE = re.compile(r"^[a-zA-Z0-9._-]+/[a-zA-Z0-9._-]+$")
    if isinstance(model, str):
        # Check for HF repo ID or local HF-style directory
        # (lazy import to keep safetensors/huggingface_hub optional)
        _is_hf_dir = os.path.isdir(model)
        _looks_like_hf = bool(_HF_ID_RE.match(model))
        if _looks_like_hf or _is_hf_dir:
            try:
                from aimnet.calculators.hf_hub import is_hf_repo_id, load_from_hf_repo
            except ImportError:
                raise ImportError(
                    f"Loading from HF repo '{model}' requires optional dependencies. "
                    "Install with: pip install aimnet[hf]"
                ) from None
            if is_hf_repo_id(model) or _is_hf_dir:
                _model, metadata = load_from_hf_repo(
                    model,
                    ensemble_member=ensemble_member,
                    device=self.device,
                    revision=revision,
                    token=token,
                )
                self.model = _model
                self.cutoff = metadata["cutoff"]
            else:
                # _looks_like_hf matched but it's a local file path — fall through
                if not os.path.isfile(model):
                    registry_family = get_registry_model_family(model)
                p = get_model_path(model)
                self.model, metadata = load_model(p, device=self.device)
                self.cutoff = metadata["cutoff"]
        else:
            if not os.path.isfile(model):
                registry_family = get_registry_model_family(model)
            p = get_model_path(model)
            self.model, metadata = load_model(p, device=self.device)
            self.cutoff = metadata["cutoff"]
    elif isinstance(model, nn.Module):
        self.model = model.to(self.device)
        self.cutoff = getattr(self.model, "cutoff", 5.0)
        metadata = cast(Mapping[str, Any] | None, getattr(self.model, "metadata", None))
        if metadata is None:
            metadata = cast(Mapping[str, Any] | None, getattr(self.model, "_metadata", None))
    else:
        raise TypeError("Invalid model type/name.")

    if metadata is not None:
        metadata = _apply_family_defaults(metadata, registry_family)
        self.model._metadata = metadata  # type: ignore[assignment]

    # 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)
    )

    # 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
        # 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=metadata.get("coulomb_sr_rc", 4.6) if metadata is not None else 4.6,
            envelope=metadata.get("coulomb_sr_envelope", "exp") if metadata is not None else "exp",
            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)
    has_embedded_lr = metadata.get("has_embedded_lr", False) if metadata is not None else False
    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:
        self.cutoff_lr = self._default_dftd3_cutoff
    elif has_embedded_lr:
        # Embedded LR modules (D3TS, SRCoulomb) need nbmat_lr
        self.cutoff_lr = self._default_dftd3_cutoff
    else:
        self.cutoff_lr = None
    self.nb_threshold = nb_threshold
    self.cache_static = bool(cache_static)
    self._static_nbmat_cache: dict[tuple[Any, ...], tuple[Any, Any, dict[str, Tensor]]] = {}
    self._static_dftd3_cache: dict[tuple[Any, ...], tuple[Any, Any, dict[str, Any]]] = {}
    # 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
    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
    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.

calculate_hessian(forces, coord) staticmethod

Dense (N, 3, N, 3) Hessian of the energy w.r.t. real-atom coordinates.

Autograd contract (IMPORTANT): The returned dense Hessian is a detached value: it carries no autograd graph back to the coordinates or model parameters. This is by design (it is materialized via torch.func.vmap over a vjp of the already-built force graph, and the periodic Ewald/PME block is a fixed-charge finite-difference term that is non-differentiable). Forces DO compose with an upstream coordinate-builder graph, but the Hessian does not, so you cannot backpropagate through eval(..., hessian=True).

If you need the Hessian to compose (e.g. H @ v that scales with / differentiates through an outer computation) or to avoid forming the dense (N, 3, N, 3) tensor on large systems, use the matrix-free :meth:hessian_vector_product instead. For a fully-differentiable Hessian, build one externally with torch.autograd.functional.hessian(energy_fn, coords) over a closure that calls the model on differentiable coordinates (note that the periodic Ewald/PME long-range block remains a fixed-charge FD term in either case).

Source code in aimnet/calculators/calculator.py

@staticmethod
def calculate_hessian(forces: Tensor, coord: Tensor) -> Tensor:
    """Dense ``(N, 3, N, 3)`` Hessian of the energy w.r.t. real-atom coordinates.

    Autograd contract (IMPORTANT):
    The returned dense Hessian is a **detached value**: it carries no
    autograd graph back to the coordinates or model parameters. This is by
    design (it is materialized via ``torch.func.vmap`` over a vjp of the
    already-built force graph, and the periodic Ewald/PME block is a
    fixed-charge finite-difference term that is non-differentiable). Forces
    DO compose with an upstream coordinate-builder graph, but the Hessian
    does not, so you cannot backpropagate through ``eval(..., hessian=True)``.

    If you need the Hessian to *compose* (e.g. ``H @ v`` that scales with /
    differentiates through an outer computation) or to avoid forming the
    dense ``(N, 3, N, 3)`` tensor on large systems, use the matrix-free
    :meth:`hessian_vector_product` instead. For a fully-differentiable
    Hessian, build one externally with
    ``torch.autograd.functional.hessian(energy_fn, coords)`` over a closure
    that calls the model on differentiable coordinates (note that the
    periodic Ewald/PME long-range block remains a fixed-charge FD term in
    either case).
    """
    # Coord includes padding atom (shape N+1), forces only for real atoms (shape N).
    # Hessian computed only for actual atoms: (N, 3, N, 3).
    #
    # vmap-over-vjp form (not is_grads_batched=True or autograd.functional.hessian):
    # torch.library.register_vmap on aimnet::conv_sv_2d_sp_{bwd,bwd_bwd} is consulted
    # ONLY by the functorch dispatch (torch.func.vmap). The legacy batching dispatch
    # would still raise "Batching rule not implemented."
    n = forces.numel()
    eye = torch.eye(n, device=forces.device, dtype=forces.dtype)

    def vjp(go: Tensor) -> Tensor:
        return torch.autograd.grad(
            forces.flatten(),
            coord,
            grad_outputs=go,
            retain_graph=True,
            allow_unused=True,
        )[0]

    hessian = -torch.func.vmap(vjp, 0)(eye)
    return hessian.view(-1, 3, coord.shape[0], 3)[:-1, :, :-1, :]

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)

hessian_vector_product(data, vectors, *, eps=0.0005, 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 Central-difference step (Angstrom) for the periodic Ewald/PME long-range term. Ignored for simple/dsf. create_graph : bool If True, keep the differentiable autograd block of the HVP in the graph so it can compose with an outer loss. The Ewald/PME fixed-position finite-difference block remains detached. Default False preserves the numeric/detached operator-action behavior.

Returns

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

Notes

The autograd part (NN + short-range + simple/dsf Coulomb + DFTD3) is an exact reverse-mode product. For ewald/pme the long-range block is a fixed-charge directional finite difference (2 force evals per vector); the same charge-response and step caveats as the dense Ewald/PME Hessian apply (see :meth:aimnet.modules.lr.LRCoulomb._coul_nvalchemi_fd_hessian). This mirrors the dense :meth:calculate_hessian assembly term-by-term, 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 differentiable autograd block 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 controls the dsf-vs-ewald/pme split: dsf passes hessian=True so LRCoulomb.forward routes through its differentiable closed-form torch path (_coul_dsf_torch), keeping the dsf curvature in the autograd graph (dsf has no dense FD block, so this is free); ewald/pme pass hessian=False so the dense O(2*3N) FD block is NOT computed, while the periodic energy is still added differentiable-through-charges (capturing the charge-response curvature). The autograd vjp of the differentiable forces equals the dense autograd Hessian block (NN + short-range + DFTD3 + Coulomb-charge-response), and the directional FD helper adds the remaining full-periodic block -- matching the dense assembly term-by-term.

Dtype / differentiability / eigensolver caveats:

• Return dtype is the model dtype (typically float32) for simple and dsf, and float64 for ewald/pme (the periodic finite-difference block is accumulated in double precision, matching the dense Ewald/PME Hessian).
• With create_graph=False the returned product is detached numeric operator action. With create_graph=True the autograd block remains differentiable w.r.t. graph-attached coordinates / model parameters; vectors are still treated as numeric directions, and the periodic FD block remains detached.
• For ewald/pme the operator is symmetric only to finite-difference accuracy (O(eps^2)); for Lanczos/LOBPCG smallest/most-negative-eigenvalue (transition-state) work, pass all probe vectors together as a single (K, N, 3) batch so the charge state is frozen across the iteration, and consider symmetrizing the operator or tuning eps.
• The fixed-charge periodic approximation (and the dsf relaxed-charge vs ewald/pme fixed-charge asymmetry) is inherited from the dense Ewald/PME Hessian and can shift near-zero/negative eigenvalues for strongly polar periodic systems; see :meth:aimnet.modules.lr.LRCoulomb._coul_nvalchemi_fd_hessian.

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 = 5e-4,
    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
        Central-difference step (Angstrom) for the periodic Ewald/PME
        long-range term. Ignored for ``simple``/``dsf``.
    create_graph : bool
        If ``True``, keep the differentiable autograd block of the HVP in
        the graph so it can compose with an outer loss. The Ewald/PME
        fixed-position finite-difference block remains detached. Default
        ``False`` preserves the numeric/detached operator-action behavior.

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

    Notes
    -----
    The autograd part (NN + short-range + ``simple``/``dsf`` Coulomb +
    DFTD3) is an exact reverse-mode product. For ``ewald``/``pme`` the
    long-range block is a fixed-charge directional finite difference (2
    force evals per vector); the same charge-response and step caveats as
    the dense Ewald/PME Hessian apply (see
    :meth:`aimnet.modules.lr.LRCoulomb._coul_nvalchemi_fd_hessian`). This
    mirrors the dense :meth:`calculate_hessian` assembly term-by-term, 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 differentiable autograd block 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 controls the dsf-vs-ewald/pme split: dsf passes
    ``hessian=True`` so ``LRCoulomb.forward`` routes through its
    differentiable closed-form torch path (``_coul_dsf_torch``), keeping the
    dsf curvature in the autograd graph (dsf has no dense FD block, so this is
    free); ewald/pme pass ``hessian=False`` so the dense O(2*3N) FD block is
    NOT computed, while the periodic energy is still added
    differentiable-through-charges (capturing the charge-response curvature).
    The autograd vjp of the differentiable forces equals the dense autograd
    Hessian block (NN + short-range + DFTD3 + Coulomb-charge-response), and
    the directional FD helper adds the remaining full-periodic block --
    matching the dense assembly term-by-term.

    Dtype / differentiability / eigensolver caveats:

    * Return dtype is the model dtype (typically float32) for ``simple`` and
      ``dsf``, and **float64** for ``ewald``/``pme`` (the periodic
      finite-difference block is accumulated in double precision, matching the
      dense Ewald/PME Hessian).
    * With ``create_graph=False`` the returned product is detached numeric
      operator action. With ``create_graph=True`` the autograd block remains
      differentiable w.r.t. graph-attached coordinates / model parameters;
      vectors are still treated as numeric directions, and the periodic FD
      block remains detached.
    * For ``ewald``/``pme`` the operator is symmetric only to
      finite-difference accuracy (O(eps^2)); for Lanczos/LOBPCG
      smallest/most-negative-eigenvalue (transition-state) work, pass all
      probe vectors together as a single ``(K, N, 3)`` batch so the charge
      state is frozen across the iteration, and consider symmetrizing the
      operator or tuning ``eps``.
    * The fixed-charge periodic approximation (and the ``dsf`` relaxed-charge
      vs ``ewald``/``pme`` fixed-charge asymmetry) is inherited from the dense
      Ewald/PME Hessian and can shift near-zero/negative eigenvalues for
      strongly polar periodic systems; see
      :meth:`aimnet.modules.lr.LRCoulomb._coul_nvalchemi_fd_hessian`.
    """
    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, eps=eps, 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: "ewald"/"pme" Hessians are computed at fixed charge (finite-difference of the analytic forces; the charge-response coupling d^2E/(dq.dr) through the model's predicted charges is omitted), while "dsf" Hessians are relaxed-charge (fully autograd). Vibrational frequencies / IR intensities are therefore not directly comparable across these backends.

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: ``"ewald"``/``"pme"`` Hessians are computed at fixed charge
    (finite-difference of the analytic forces; the charge-response coupling
    ``d^2E/(dq.dr)`` through the model's predicted charges is omitted), while
    ``"dsf"`` Hessians are relaxed-charge (fully autograd). Vibrational
    frequencies / IR intensities are therefore not directly comparable across
    these backends.
    """
    if method not in ("simple", "dsf", "ewald", "pme"):
        raise ValueError(f"Invalid method: {method}")

    # 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 it with `pip install "aimnet[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 "aimnet[pysis]"`.'
        ) 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. Install it on Python 3.12+ with `pip install "aimnet[torchsim]"`.'
        ) 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.

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.

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_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