### NPT Simulation Setup with ASE NPTIntegrator

Source: https://isayevlab.github.io/aimnetcentral/tutorials/molecular_dynamics

Shows the setup pattern for NPT simulations using ASE's NPTIntegrator. Requires a periodic system with a unit cell; for periodic systems with charges, switch to DSF Coulomb via set_lrcoulomb_method("dsf"). The example is fully commented out and includes MaxwellBoltzmannDistribution, NPTIntegrator with timestep 0.5 fs, temperature 300 K, external stress 0.0 eV/A^3, thermostat time constant 25 fs, and barostat pfactor=None. Trajectory is saved every 10 steps for 5000 steps.

```python
from ase.md.npt import NPT as NPTIntegrator
from ase.io import read
from ase import units

# For NPT, you need a periodic system with a unit cell.
# This example shows the setup pattern; replace with your periodic system.
# atoms = read("your_periodic_system.xyz")
# atoms.pbc = True
# atoms.cell = [10.0, 10.0, 10.0]  # Angstroms

# calc = AIMNet2ASE(
#     AIMNet2Calculator("aimnet2", compile_model=True)
# )

# For periodic systems with charges, switch to DSF Coulomb
# calc.base_calc.set_lrcoulomb_method("dsf")

# atoms.calc = calc

# MaxwellBoltzmannDistribution(atoms, temperature_K=300)

# npt = NPTIntegrator(
#     atoms,
#     timestep=0.5 * units.fs,
#     temperature_K=300,
#     externalstress=0.0,  # Target pressure in eV/A^3 (0 = vacuum, no external stress)
#     ttime=25 * units.fs,   # Thermostat time constant
#     pfactor=None,          # Barostat fictitious mass (None = ASE default)
# )

# traj = Trajectory("npt.traj", "w", atoms)
# npt.attach(traj.write, interval=10)
# npt.run(steps=5000)
# traj.close()

```

--------------------------------

### Install AIMNet2 from PyPI

Source: https://isayevlab.github.io/aimnetcentral/getting_started

Basic installation command for AIMNet2 from PyPI.

```bash
pip install aimnet
```

--------------------------------

### Install AIMNet2 wrapper

Source: https://isayevlab.github.io/aimnetcentral/external/orca

Run the installer script to create an isolated virtual environment with AIMNet2 dependencies.

```bash
python install.py -e aimnet2
```

--------------------------------

### Install AIMNet2 with pysisyphus support

Source: https://isayevlab.github.io/aimnetcentral/external/pysis

Install the aimnet package with the pysis extra to enable pysisyphus integration.

```bash
pip install "aimnet[pysis]"

```

--------------------------------

### Example configuration: config.yaml

Source: https://isayevlab.github.io/aimnetcentral/cli

Training parameters YAML file specifying data paths, optimizer settings, and training options.

```yaml
data:
  train: data/train.h5
  val: data/val.h5
  batch_size: 32

optimizer:
  lr: 0.001
  weight_decay: 0.0

training:
  epochs: 100
  checkpoint_dir: checkpoints/
  log_interval: 10
```

--------------------------------

### Install AIMNet2 with training extras

Source: https://isayevlab.github.io/aimnetcentral/cli

Installs the core package along with the 'train' extra required for training, export, and self-atomic-energy commands.

```bash
pip install "aimnet[train]"
```

--------------------------------

### Install AIMNet2 with Hugging Face Hub support

Source: https://isayevlab.github.io/aimnetcentral/getting_started

Installation command for AIMNet2 with Hugging Face Hub support to load models directly from Hugging Face.

```bash
pip install "aimnet[hf]"
```

--------------------------------

### Install AIMNet2 wrapper with custom paths

Source: https://isayevlab.github.io/aimnetcentral/external/orca

Specify custom directories for the virtual environment and the resulting scripts.

```bash
python install.py --venv-dir <path/to/venv> --script-dir <path/to/bin> -e aimnet2
```

--------------------------------

### Install AIMNet2 with TorchSim and ASE

Source: https://isayevlab.github.io/aimnetcentral/getting_started

Installation command for AIMNet2 with TorchSim and ASE integration, requiring Python 3.12+.

```bash
pip install "aimnet[torchsim,ase]"
```

--------------------------------

### Launch training with default parameters

Source: https://isayevlab.github.io/aimnetcentral/train

Starts training with default parameters, specifying the training dataset, SAE file, and run name.

```bash
$ aimnet train data.train=dataset.h5 data.sae.energy.file=dataset_sae.yaml run_name=firstrun

```

--------------------------------

### Example configuration: model.yaml

Source: https://isayevlab.github.io/aimnetcentral/cli

Model architecture YAML file specifying the AIMNet2 model class and its hyperparameters.

```yaml
class: aimnet.models.AIMNet2
kwargs:
  nfeature: 16
  hidden: [[512, 380], [512, 380]]
  aev:
    rc_s: 5.0
    nshifts_s: 16
  outputs:
    energy_mlp:
      class: aimnet.modules.Output
      # ... output configuration
```

--------------------------------

### aimnet calc_sae examples

Source: https://isayevlab.github.io/aimnetcentral/cli

Three usage examples: default fit, limiting samples with --samples, and using the SAE file in training.

```bash
# Fit SAE from the training dataset (uses up to 100k samples by default)
aimnet calc_sae \
  data/dataset.h5 \
  configs/sae.yaml

# Limit the number of samples used for the fit
aimnet calc_sae \
  data/dataset.h5 \
  configs/sae.yaml \
  --samples 50000

# Use in training
aimnet train \
  --config train_config.yaml \
  --model model.yaml \
  data.sae.energy.file=configs/sae.yaml
```

--------------------------------

### Install AIMNet2 ASE calculator

Source: https://isayevlab.github.io/aimnetcentral/external/ase

Install the aimnet package with the ASE extra to enable the AIMNet2ASE calculator.

```bash
pip install "aimnet[ase]"
```

--------------------------------

### Complete example: training with W&B and continuing from checkpoint

Source: https://isayevlab.github.io/aimnetcentral/cli

Shows two commands: one for training with W&B logging and another for continuing from existing weights using --load and --save.

```bash
# Train with W&B logging (GPU controlled via CUDA_VISIBLE_DEVICES)
aimnet train \
  --config experiments/train_config.yaml \
  --model models/aimnet2.yaml

# Continue from existing weights
aimnet train \
  --config experiments/train_config.yaml \
  --model models/aimnet2.yaml \
  --load checkpoints/last.pt \
  --save checkpoints/continued.pt
```

--------------------------------

### Usage Example: AIMNet2ASE with ASE

Source: https://isayevlab.github.io/aimnetcentral/api/calculators

Shows how to use the AIMNet2ASE calculator with an ASE Atoms object. Requires the 'ase' extra: pip install aimnet[ase]. The calculator operates in eV and Angstrom, and this example reads a molecule from a file, assigns the calculator, and prints energy and forces.

```python
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())
```

--------------------------------

### Start oet_server

Source: https://isayevlab.github.io/aimnetcentral/external/orca

Launch the server in a separate shell to keep the AIMNet2 model resident in memory for multiple ORCA jobs.

```bash
oet_server aimnet2 --nthreads 4
```

--------------------------------

### aimnet convert example with verification

Source: https://isayevlab.github.io/aimnetcentral/cli

Converts a legacy model and verifies the output with a Python snippet that loads the new model and prints its cutoff and external Coulomb flag.

```bash
# Convert legacy model
aimnet convert \
  legacy/aimnet2_old.jpt \
  configs/aimnet2.yaml \
  models/aimnet2_new.pt

# Verify conversion
python -c "
from aimnet.calculators import AIMNet2Calculator
calc = AIMNet2Calculator('models/aimnet2_new.pt')
print(f'Cutoff: {calc.cutoff}')
print(f'Has external Coulomb: {calc.has_external_coulomb}')
"
```

--------------------------------

### Prepare Required Files: config.yaml reconstruction comment

Source: https://isayevlab.github.io/aimnetcentral/model_format

Placeholder comment shown in the migration guide when config.yaml is missing. Instructs that you may need to reconstruct it from training logs or model inspection.

```text
# If you don't have config.yaml, you may need to reconstruct it
# from training logs or model inspection

```

--------------------------------

### Test Runtime Flexibility: set_lrcoulomb_method examples

Source: https://isayevlab.github.io/aimnetcentral/model_format

Demonstrates runtime method changes for v2 models using set_lrcoulomb_method with dsf, ewald, and pme options. v1 models show a warning and do not change method.

```python
# v2 models support runtime method changes
calc_v2.set_lrcoulomb_method("dsf", cutoff=15.0)
result_dsf = calc_v2(data)

calc_v2.set_lrcoulomb_method("ewald", ewald_accuracy=1e-6)
result_ewald = calc_v2(data)

calc_v2.set_lrcoulomb_method("pme", ewald_accuracy=1e-6)
result_pme = calc_v2(data)

# v1 models show warning but don't change
calc_v1.set_lrcoulomb_method("dsf", cutoff=15.0)
# Warning: Cannot change method for legacy models

```

--------------------------------

### Complete AIMNet2 Calculator Example

Source: https://isayevlab.github.io/aimnetcentral/calculator

Creates an AIMNet2Calculator, prints its configuration, sets the long-range Coulomb method to 'dsf' with a cutoff of 12.0, prepares a random input, and runs inference with forces. The output includes energy, forces, and charges shapes.

```python
import torch
from aimnet.calculators import AIMNet2Calculator

# Create calculator
calc = AIMNet2Calculator("aimnet2")

# Check configuration
print(f"Device: {calc.device}")
print(f"Cutoff: {calc.cutoff}")
print(f"Has external Coulomb: {calc.has_external_coulomb}")
print(f"Coulomb method: {calc.coulomb_method}")

# Configure for PBC
calc.set_lrcoulomb_method("dsf", cutoff=12.0)

# Prepare input
data = {
    "coord": torch.randn(20, 3),
    "numbers": torch.randint(1, 10, (20,)),
    "charge": torch.tensor([0.0]),
}

# Run inference
result = calc.eval(data, forces=True)

print(f"Energy: {result['energy']}")
print(f"Forces shape: {result['forces'].shape}")
print(f"Charges shape: {result['charges'].shape}")
```

--------------------------------

### Quick start: single-point energy and forces with AIMNet2ASE

Source: https://isayevlab.github.io/aimnetcentral/external/ase

Create an H2O molecule, attach the AIMNet2ASE calculator using the 'aimnet2' registry alias, and compute the potential energy (eV) and forces (eV/A).

```python
from ase.build import molecule
from aimnet.calculators import AIMNet2ASE

atoms = molecule("H2O")
atoms.calc = AIMNet2ASE("aimnet2")           # registry alias
energy = atoms.get_potential_energy()        # eV
forces = atoms.get_forces()                  # eV/A
```

--------------------------------

### AIMNetCalculator __init__ method

Source: https://isayevlab.github.io/aimnetcentral/api/calculators

Constructor for the AIMNet calculator. Handles device selection, model loading via resolve_model, optional torch.compile, and setup of external LRCoulomb and DFTD3 modules based on metadata flags. Raises ValueError if needs_dispersion is True but d3_params are missing.

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

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

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

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

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

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

--------------------------------

### masses_amu(atomic_numbers)

Source: https://isayevlab.github.io/aimnetcentral/api/calculators

Returns standard atomic masses in amu for the given atomic numbers. Uses a bundled table, so no ASE install is required.

```APIDOC
## `masses_amu(atomic_numbers)`

### Description
Standard atomic masses in amu for `atomic_numbers`. Uses the table bundled in `aimnet.constants.get_masses` (the values of `ase.data.atomic_masses` stored in single precision), so no ASE install is required. Pass `atoms.get_masses()` to `analyze_hessian` instead when ASE is available and full double precision matters.

### Parameters
- **atomic_numbers** (`Any`) - Required - Integer array-like of shape `(N,)`.

### Returns
- **ndarray** - Array of shape `(N,)` with dtype float64.

### Source Code
```python
def masses_amu(atomic_numbers: Any) -> np.ndarray:
    """Standard atomic masses in amu for ``atomic_numbers``.

    Uses the table bundled in :func:`aimnet.constants.get_masses` (the values of
    ``ase.data.atomic_masses`` stored in single precision), so no ASE install is
    required. Pass ``atoms.get_masses()`` to :func:`analyze_hessian` instead when
    ASE is available and full double precision matters.

    Args:
        atomic_numbers: Integer array-like of shape ``(N,)``.
    """
    # Implementation omitted for brevity
```
```

--------------------------------

### AIMNet2ASE __init__ Source Code

Source: https://isayevlab.github.io/aimnetcentral/api/calculators

Source code of the AIMNet2ASE constructor. It initializes the calculator with a base_calc (AIMNet2Calculator or string), charge, multiplicity, and species validation. Raises ImportError if ASE is not installed. If base_calc is a string, it creates an AIMNet2Calculator. For NSE models, adds 'spin_charges' to implemented properties. Reads implemented species from model metadata.

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

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

--------------------------------

### masses_amu: Standard atomic masses in amu

Source: https://isayevlab.github.io/aimnetcentral/api/calculators

Returns standard atomic masses in amu for the given atomic numbers, using the table bundled in aimnet.constants.get_masses (values from ase.data.atomic_masses stored in single precision). No ASE install is required; pass atoms.get_masses() to analyze_hessian instead when full double precision matters.

```python
def masses_amu(atomic_numbers: Any) -> np.ndarray:
    """Standard atomic masses in amu for ``atomic_numbers``.

    Uses the table bundled in :func:`aimnet.constants.get_masses` (the values of
    ``ase.data.atomic_masses`` stored in single precision), so no ASE install is
    required. Pass ``atoms.get_masses()`` to :func:`analyze_hessian` instead when
    ASE is available and full double precision matters.

    Args:
        atomic_numbers: Integer array-like of shape ``(N,)``.
    """
```

--------------------------------

### Calculate energy of square planar PdCl4(2-)

Source: https://isayevlab.github.io/aimnetcentral/advanced/transition_metal_catalysis

Constructs a square planar PdCl4(2-) structure with approximate Pd-Cl bond length of 2.30 Angstrom and computes its energy and forces using the 'aimnet2-pd' calculator. The example includes an explicit note that starting geometries are approximate and should be optimized before comparing energies.

```python
import torch
from aimnet.calculators import AIMNet2Calculator

calc = AIMNet2Calculator("aimnet2-pd")

# PdCl4(2-): ideal square planar
# Pd at origin, 4 Cl in xy-plane
d_pd_cl = 2.30  # approximate Pd-Cl bond length in Angstrom
sq_planar = torch.tensor([
    [0.0, 0.0, 0.0],           # Pd
    [d_pd_cl, 0.0, 0.0],       # Cl
    [0.0, d_pd_cl, 0.0],       # Cl
    [-d_pd_cl, 0.0, 0.0],      # Cl
    [0.0, -d_pd_cl, 0.0],      # Cl
])
numbers_pdcl4 = torch.tensor([46, 17, 17, 17, 17])

result_sq = calc({
    "coord": sq_planar,
    "numbers": numbers_pdcl4,
    "charge": torch.tensor(-2.0),
}, forces=True)

print(f"Square planar PdCl4(2-) energy: {result_sq['energy'].item():.4f} eV")
print(f"Max force: {result_sq['forces'].abs().max().item():.4f} eV/A")
```

--------------------------------

### Install openmmml

Source: https://isayevlab.github.io/aimnetcentral/external/openmm

Install the openmmml package along with aimnet and openmm. The PyPI distribution name is openmmml (no hyphen); conda-forge uses openmm-ml.

```bash
pip install aimnet openmm openmmml

```

--------------------------------

### Configure ORCA input for server/client mode

Source: https://isayevlab.github.io/aimnetcentral/external/orca

Use the !ExtOpt keyword and specify the path to the oet_client script within the %method block.

```text
! ExtOpt Opt PAL8

%method
  ProgExt "/full/path/to/oet_client"
end

*XYZfile 0 1 mol.xyz
```

--------------------------------

### Install PyTorch with CUDA 12.6 for GPU support

Source: https://isayevlab.github.io/aimnetcentral/getting_started

Command to install PyTorch with CUDA 12.6 for GPU acceleration. AIMNet2 works on CPU by default.

```bash
# For CUDA 12.6
pip install torch --index-url https://download.pytorch.org/whl/cu126
```

--------------------------------

### compile_kwargs example: use reduce-overhead mode

Source: https://isayevlab.github.io/aimnetcentral/calculator

Passes additional keyword arguments to torch.compile(). The example enables reduce-overhead mode for lower latency. See torch.compile documentation for available options.

```python
# Example: use reduce-overhead mode for lower latency
calc = AIMNet2Calculator("aimnet2", compile_model=True, compile_kwargs={"mode": "reduce-overhead"})

```

--------------------------------

### Quick Example: Water Dimer Binding Energy with AIMNet2Calculator

Source: https://isayevlab.github.io/aimnetcentral/models/aimnet2_2025

Computes the binding energy of a water dimer using the supramolecular approach (dimer minus monomers). The example uses the AIMNet2Calculator with compile_model=True and demonstrates energy calculations for the dimer and each monomer. No BSSE correction is needed because neural network potentials do not suffer from basis set superposition error.

```python
import torch
from aimnet.calculators import AIMNet2Calculator

calc = AIMNet2Calculator("aimnet2-2025", compile_model=True)

# Water dimer coordinates (Angstrom)
dimer_coords = torch.tensor([
    # Monomer A (donor)
    [ 0.0000,  0.0000,  0.1173],  # O
    [ 0.0000,  0.7572, -0.4692],  # H
    [ 0.0000, -0.7572, -0.4692],  # H
    # Monomer B (acceptor)
    [ 2.9000,  0.0000,  0.0000],  # O
    [ 3.4000,  0.7600,  0.0000],  # H
    [ 3.4000, -0.7600,  0.0000],  # H
])
dimer_numbers = torch.tensor([8, 1, 1, 8, 1, 1])

# Compute dimer energy
e_dimer = calc({
    "coord": dimer_coords,
    "numbers": dimer_numbers,
    "charge": torch.tensor(0.0),
})["energy"].item()

# Compute monomer energies separately
mono_a_coords = dimer_coords[:3]
mono_b_coords = dimer_coords[3:]
mono_numbers = torch.tensor([8, 1, 1])

e_mono_a = calc({
    "coord": mono_a_coords,
    "numbers": mono_numbers,
    "charge": torch.tensor(0.0),
})["energy"].item()

e_mono_b = calc({
    "coord": mono_b_coords,
    "numbers": mono_numbers,
    "charge": torch.tensor(0.0),
})["energy"].item()

binding_energy = e_dimer - e_mono_a - e_mono_b
print(f"Binding energy: {binding_energy * 1000:.1f} meV")
print(f"Binding energy: {binding_energy * 23.0609:.2f} kcal/mol")

```

--------------------------------

### Import aimnet to trigger Warp initialization

Source: https://isayevlab.github.io/aimnetcentral/tutorials/performance

Importing aimnet triggers Warp initialization (wp.init()) on first kernel import, costing 2-5 seconds once per Python process.

```python
import aimnet  # Triggers wp.init() on first kernel import: ~2-5 seconds

```

--------------------------------

### SAE output YAML format

Source: https://isayevlab.github.io/aimnetcentral/cli

Example of the output YAML file containing per-element self-atomic energies in eV.

```yaml
# Self-atomic energies in eV
1: -13.587 # Hydrogen
6: -1027.592 # Carbon
7: -1483.525 # Nitrogen
8: -2039.734 # Oxygen
```

--------------------------------

### Verify AIMNet2 installation

Source: https://isayevlab.github.io/aimnetcentral/getting_started

Python code to confirm AIMNet2 loads correctly. The first time you load a model, it will be downloaded automatically.

```python
from aimnet.calculators import AIMNet2Calculator

calc = AIMNet2Calculator("aimnet2")
print("AIMNet2 loaded successfully")
```