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, 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. 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,
    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

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

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

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)
    self._update_lr_nblists()

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
    self._update_lr_nblists()

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.

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``.
    """
    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
    self._update_lr_nblists()

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

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_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)
    family_key, member = model_name.rsplit("_", 1)
    if not member.isdigit() or not family_key.startswith("aimnet2-"):
        raise ValueError(f"Model {model_name} does not follow the canonical registry naming convention.")
    return family_key.removeprefix("aimnet2-")