CLI Reference¶

AIMNet2 provides command-line tools for training, model export, conversion, and data preprocessing.

Model cache recovery¶

Official registry downloads are verified against the SHA-256 digest committed with each registry entry. Cache hits are re-hashed, and a failed download is never installed as the final cache file. A corrupt cache hit triggers one verified atomic replacement attempt; if failures persist, remove the cache and retry:

aimnet clear_model_cache

Do not replace a registry digest to work around a mismatch; investigate the artifact provenance and publish a new immutable filename when bytes change.

Official wheels and source distributions currently contain no bundled model artifacts; registry models are downloaded on demand. A downstream bundled candidate is still digest-checked and fails closed when stale. Weekly strict fleet CI verifies every official registry digest, loads every artifact with the fixed registry import policy, and checks the exact role-specific YAML defaults.

Installation¶

The aimnet entry point is installed with the core package. Training, export, and self-atomic-energy commands require the train extra:

pip install "aimnet[train]"

Commands Overview¶

Command	Purpose	Typical Use
`aimnet train`	Train a model	Training from scratch
`aimnet export`	Export trained weights to inference format	After training
`aimnet convert`	Convert legacy .jpt to new .pt format	Migrating old models
`aimnet calc_sae`	Calculate self-atomic energies	Before training
`aimnet download`	Prefetch registry model weights into the local cache	Preparing for offline/HPC use
`aimnet info`	Print environment and kernel-path diagnostics	Troubleshooting
`aimnet clear_model_cache`	Remove cached model artifacts	Recovering from a corrupt cache

aimnet train¶

Train an AIMNet2 model from a configuration file.

Basic Usage¶

aimnet train --config config.yaml --model model.yaml

Options¶

aimnet train [OPTIONS] [ARGS]...

Option	Type	Description
`--config PATH`	Optional	Extra training configuration YAML (may be passed multiple times; merged over the default)
`--model PATH`	Optional	Model architecture YAML (defaults to bundled `aimnet/models/aimnet2.yaml`)
`--load PATH`	Optional	Path to existing model weights to load before training
`--save PATH`	Optional	Path to save model weights
`--no-default-config`	Flag	Skip loading `aimnet/train/default_train.yaml`
`ARGS`	Variadic	Dot-separated overrides applied last (e.g., `data.train=mydata.h5 run_name=firstrun`)

Trailing tokens after the named options (ARGS) are positional dot-form overrides into the merged training config — they are not flags.

Device selection is controlled via CUDA_VISIBLE_DEVICES; training uses all visible GPUs in DDP mode.

Example Configuration¶

Training requires two YAML files:

config.yaml (training parameters):

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

model.yaml (architecture):

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

Complete Example¶

# 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

See train.md for detailed training documentation.

aimnet export¶

Export trained model weights to inference format (.pt).

Basic Usage¶

aimnet export weights.pt output.pt --model model.yaml --sae sae.yaml

Options¶

aimnet export INPUT OUTPUT [OPTIONS]

Positional Arguments:

Argument	Description
`INPUT`	Path to trained weights (.pt checkpoint)
`OUTPUT`	Path for exported model (.pt file)

Options:

Option	Description
`--model PATH`	Model architecture YAML (required)
`--sae PATH`	Self-atomic energies YAML (required)
`--needs-coulomb` / `--no-coulomb`	Override external Coulomb detection
`--needs-dispersion` / `--no-dispersion`	Override external DFTD3 detection
`--model-import-path PATH`	Trust an exact constructor path or namespace for this local export; repeat as needed

Export Process¶

The export process:

Loads model architecture from YAML
Strips LRCoulomb/DFTD3 modules
Adds SRCoulomb if LRCoulomb was present
Loads trained weights
Bakes SAE into atomic_shift as float64
Masks unimplemented species
Validates the complete artifact with canonical distribution rules
Serializes to a sibling temporary file and atomically replaces the output

--no-coulomb is rejected when an embedded short-range Coulomb subtraction requires an external full-Coulomb correction. Enabled dispersion requires complete s8, a1, and a2 parameters. A validation or serialization failure leaves an existing output file unchanged.

Examples¶

Basic export:

aimnet export \
  checkpoints/best.pt \
  models/aimnet2_production.pt \
  --model configs/aimnet2.yaml \
  --sae configs/sae.yaml

Override auto-detection:

# Force external Coulomb and DFTD3
aimnet export weights.pt model.pt \
  --model config.yaml \
  --sae sae.yaml \
  --needs-coulomb \
  --needs-dispersion \
  --model-import-path "my_package.models.*"

--model-import-path is an explicit trust decision for local export. It extends the default constructor allowlist; it does not weaken safe YAML parsing, forbidden-key checks, metadata validation, or state-dict validation.

Export without dispersion:

# Model without DFTD3 correction
aimnet export weights.pt model_no_d3.pt \
  --model config_no_d3.yaml \
  --sae sae.yaml

aimnet convert¶

Convert legacy .jpt (JIT) models to new .pt format.

Basic Usage¶

aimnet convert model.jpt config.yaml output.pt

Options¶

aimnet convert INPUT CONFIG OUTPUT

Positional Arguments:

Argument	Description
`INPUT`	Legacy .jpt model file
`CONFIG`	Original model configuration YAML
`OUTPUT`	Output .pt model file

Conversion Process¶

Loads legacy JIT model
Extracts weights from state dict
Strips embedded LR modules
Adds SRCoulomb if needed
Converts atomic_shift to float64
Generates metadata
Saves as v2 format

Example¶

# 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}')
"

See model_format.md for detailed migration guide.

aimnet calc_sae¶

Calculate self-atomic energies (SAE) from a dataset by fitting a per-element energy shift to the training data.

Basic Usage¶

aimnet calc_sae dataset.h5 output_sae.yaml

Options¶

aimnet calc_sae [OPTIONS] DS OUTPUT

Positional Arguments:

Argument	Description
`DS`	HDF5 dataset (`SizeGroupedDataset` layout) containing `numbers` and `energy`
`OUTPUT`	Output YAML file for fitted SAE values

Options:

Option	Type	Default	Description
`--samples`	int	100000	Maximum number of dataset samples used for the fit

SAE Format¶

Output YAML format:

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

Example¶

# 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

How SAE Is Computed¶

aimnet calc_sae does not require a separate single-atom dataset. It operates on the same SizeGroupedDataset HDF5 file used for training, fits a per-element energy shift via a least-squares regression on (numbers, energy), trims outliers (2nd/98th percentiles of the residual), and refits. The element list is inferred automatically from the dataset.

aimnet download¶

Prefetch registry model weights into the local cache.

Basic Usage¶

aimnet download aimnet2
aimnet download aimnet2-wb97m aimnet2-b973c
aimnet download --all

Options¶

aimnet download [MODELS]... [OPTIONS]

Argument/Option	Description
`MODELS`	Registry names or aliases (e.g. `aimnet2`) to fetch. Repeatable.
`--all`	Download every model in the registry.

Files are stored in the cache directory ($AIMNET_CACHE_DIR or ~/.cache/aimnet) and verified against their registry SHA-256 digests, so a pre-seeded cache works fully offline. When multiple models are requested (including with --all), a failure on one model is reported and the remaining models are still downloaded; the command exits non-zero if any model failed.

See offline_hpc.md for pre-seeding a cache on air-gapped or HPC compute nodes.

aimnet info¶

Print environment and kernel-path diagnostics.

Basic Usage¶

aimnet info

Reports the aimnet, torch, and warp-lang versions, CUDA availability, the registered kernel ops, and the model cache directory path. Useful for confirming the accelerated kernel path is active before troubleshooting performance or numerics.

Common Workflows¶

Complete Training Workflow¶

# 1. Calculate SAE
aimnet calc_sae single_atoms.h5 sae.yaml

# 2. Train model
aimnet train \
  --config config.yaml \
  --model aimnet2.yaml

# 3. Export best checkpoint
aimnet export \
  checkpoints/best.pt \
  production_model.pt \
  --model aimnet2.yaml \
  --sae sae.yaml

# 4. Validate
python validate_model.py production_model.pt

Legacy Migration Workflow¶

# 1. Convert .jpt to .pt
aimnet convert \
  old_model.jpt \
  original_config.yaml \
  new_model.pt

# 2. Test equivalence
python test_equivalence.py old_model.jpt new_model.pt

# 3. Update inference code
# Replace: calc = AIMNet2Calculator("old_model.jpt")
# With:    calc = AIMNet2Calculator("new_model.pt")

Environment Variables¶

Variable	Description	Default
`AIMNET_CACHE_DIR`	Model cache directory	`~/.cache/aimnet/`
`CUDA_VISIBLE_DEVICES`	GPU device selection	All available

Exit Codes¶

Code	Meaning
0	Success
1	General error
2	Invalid arguments
3	File not found
4	Configuration error

Getting Help¶

# General help
aimnet --help

# Command-specific help
aimnet train --help
aimnet export --help
aimnet convert --help
aimnet calc_sae --help

CLI Reference¶

Model cache recovery¶

Installation¶

Commands Overview¶

aimnet train¶

Basic Usage¶

Options¶

Example Configuration¶

Complete Example¶

aimnet export¶

Basic Usage¶

Options¶

Export Process¶

Examples¶

aimnet convert¶

Basic Usage¶

Options¶

Conversion Process¶

Example¶

aimnet calc_sae¶

Basic Usage¶

Options¶

SAE Format¶

Example¶

How SAE Is Computed¶

aimnet download¶

Basic Usage¶

Options¶

aimnet info¶

Basic Usage¶

Common Workflows¶

Complete Training Workflow¶

Legacy Migration Workflow¶

Environment Variables¶

Exit Codes¶

Getting Help¶

See Also¶