API Reference

This section provides detailed documentation and source code links for the IANN package modules, classes, and functions.

Data

class iann.data.AtomsData(num_atoms: torch.Tensor, atomic_numbers: torch.Tensor, positions: torch.Tensor, cell: torch.Tensor, edge_indices: torch.Tensor, edge_vectors: torch.Tensor, num_edges: torch.Tensor, global_attr: torch.Tensor, energy: torch.Tensor | None = None, forces: torch.Tensor | None = None, stress: torch.Tensor | None = None, image_indices: torch.Tensor | None = None, atomic_energy: torch.Tensor | None = None, node_attr: torch.Tensor | None = None, node_feat: torch.Tensor | None = None, edge_dist_embedding: torch.Tensor | None = None, edge_diff_embedding: torch.Tensor | None = None, energy_variance: torch.Tensor | None = None, forces_variance: torch.Tensor | None = None, global_embedding: torch.Tensor | None = None)[source]

A NamedTuple of model inputs.

Parameters:

  • num_atoms (torch.Tensor) – The number of atoms in the system.

  • atomic_numbers (torch.Tensor) – The atomic numbers of the atoms in the system.

  • positions (torch.Tensor) – The positions of the atoms in the system.

  • cell (torch.Tensor) – The cell of the system.

  • edge_indices (torch.Tensor) – The indices of the edges in the system.

  • edge_vectors (torch.Tensor) – The vectors of the edges in the system.

  • num_edges (torch.Tensor) – The number of edges in the system.

  • energy (Optional[torch.Tensor]) – The energy of the system.

  • forces (Optional[torch.Tensor]) – The forces of the system.

  • stress (Optional[torch.Tensor]) – The stress tensor of the system (3x3 matrix).

  • image_indices (Optional[torch.Tensor]) – The image indices of the atoms in the system.

  • atomic_energy (Optional[torch.Tensor]) – The atomic energy of the system.

  • node_attr (Optional[torch.Tensor]) – The node attributes of the atoms in the system.

  • node_feat (Optional[torch.Tensor]) – The node features of the atoms in the system.

  • edge_dist_embedding (Optional[torch.Tensor]) – The edge distance embedding of the atoms in the system.

  • edge_diff_embedding (Optional[torch.Tensor]) – The edge difference embedding of the atoms in the system.

  • energy_variance (Optional[torch.Tensor]) – The energy variance of the system.

  • forces_variance (Optional[torch.Tensor]) – The forces variance of the system.

  • global_attr (Optional[torch.Tensor]) – The global attributes of the system.

  • global_embedding (Optional[torch.Tensor]) – The global embedding of the system.

class iann.data.AseDataset(*args: Any, **kwargs: Any)[source]

A class to read the data from the ASE Atoms object.

Parameters:

  • db (str or ase.Atoms (or list of ase.Atoms)) – The ASE Atoms object or the path to the ASE Atoms object.

  • cutoff (float) – The cutoff radius for the neighbor list.

  • compute_forces (bool) – Whether to compute the forces.

Returns:

atoms_data – The AtomsData object.

Return type:

AtomsData

Models

class iann.models.fastpot.FastPot(num_layers=3, num_channels=128, norm_data=True, data_mean=[0.0], data_stddev=[1.0], norm_per_atom=True, **kwargs)[source]

FastPot: Fast Potential with high-order tensor features and equivariant message passing for fast and accurate potential energy surface prediction.

__init__(num_layers=3, num_channels=128, norm_data=True, data_mean=[0.0], data_stddev=[1.0], norm_per_atom=True, **kwargs)[source]

Initialize the FastPot model.

Parameters:

  • num_layers (int) – Number of message passing layers

  • num_channels (int) – Number of feature channels

  • cutoff (float) – Interaction cutoff distance

  • lmax (int) – Maximum spherical harmonic degree

forward(data: AtomsData)[source]
Parameters:

data (AtomsData) – Input data containing atomic information

Returns:

Output data with predicted energies and optionally forces

Return type:

AtomsData

class iann.models.painn.PaiNN(...)[source]

A class to set up the PaiNN model.

__init__(num_layers=3, num_channels=128, norm_data=True, data_mean=[0.0], data_stddev=[1.0], norm_per_atom=True, **kwargs)[source]

Initialize the PaiNN model.

forward(data: AtomsData)[source]
Parameters:

data (AtomsData) – Input data for the model.

Returns:

Output data after applying the model.

Return type:

AtomsData

class iann.models.nequip.NequIP(...)[source]

A class to set up the NequIP model.

__init__(num_layers: int, num_channels: int = 64, norm_data: bool = False, norm_per_atom: bool = False, data_stddev: float = 1.0, data_mean: float = 0.0, **kwargs) → None[source]

Initialize the NequIP model.

forward(data: AtomsData)[source]
Parameters:

data (AtomsData) – Input data for the model.

Returns:

Output data after applying the model.

Return type:

AtomsData

class iann.models.mace.MACE(...)[source]

A class to set up the MACE model.

__init__(num_layers: int = 3, num_channels: int = 64, norm_data: bool = False, norm_per_atom: bool = False, data_stddev: float = 1.0, data_mean: float = 0.0, **kwargs) → None[source]

Initialize the MACE model.

forward(data: AtomsData)[source]
Parameters:

data (AtomsData) – Input data for the model.

Returns:

Output data after applying the model.

Return type:

AtomsData

class iann.models.equiformerV2.EquiformerV2(...)[source]

A class to set up the EquiformerV2 model.

__init__(device='cpu', num_channels='64', num_layers=3, norm_data=False, norm_per_atom=False, data_stddev=1.0, data_mean=0.0, **kwargs)[source]

Initialize the EquiformerV2 model.

forward(data: AtomsData)[source]
Parameters:

data (AtomsData) – Input data for the model.

Returns:

Output data after applying the model.

Return type:

AtomsData

Trainer

class iann.trainer.Trainer(model='painn', config=None, distributed=True, rank=None, world_size=None)[source]

Trainer class for training interatomic neural network models

__init__(model='painn', config=None, distributed=True, rank=None, world_size=None)[source]

Initialize the trainer with a model type and optional config

Parameters:

  • model (str) – Model type (“painn”, “nequip”, “mace”, or “equiformer2”)

  • config (dict, optional) – Configuration overrides

  • distributed (bool, optional) – Whether to use distributed training

  • rank (int, optional) – Rank of the current process

  • world_size (int, optional) – Total number of processes

Calculators

class iann.calculators.MLCalculator(model_path=None, model=None, config=None, energy_scale=1.0, forces_scale=1.0, device=None, verbose=False, **kwargs)[source]

Machine learning calculator for a single model.

__init__(model_path=None, model=None, config=None, energy_scale=1.0, forces_scale=1.0, device=None, verbose=False, **kwargs)[source]
Parameters:

  • model_path (str) – Path to the model (only one model).

  • model (iann.models.Model) – Model to use.

  • config (dict) – Configuration for the calculator.

  • energy_scale (float) – Energy scale.

  • forces_scale (float) – Forces scale.

  • device (str) – Device to use.

  • verbose (bool) – Verbosity.

Returns:

Dictionary with keys ‘energy’, ‘forces’.

Return type:

Dict[str, array]

calculate(atoms=None, properties=['energy', 'forces'], system_changes=['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'])[source]
Parameters:

  • atoms (ase.Atoms) – ASE atoms object.

  • properties (list of str) – energy, forces, and stress are supported.

  • system_changes (list of str) – List of changes for ASE.

class iann.calculators.EnsembleCalculator(model_paths=None, models=None, config=None, energy_scale=1.0, forces_scale=1.0, device=None, verbose=False, **kwargs)[source]

Ensemble calculator for multiple models.

__init__(model_paths=None, models=None, config=None, energy_scale=1.0, forces_scale=1.0, device=None, verbose=False, **kwargs)[source]
Parameters:

  • model_paths (list of str) – Paths to the models (several models for ensemble).

  • models (list of iann.models.Model) – Models to use.

  • config (dict) – Configuration for the calculator.

  • energy_scale (float) – Energy scale.

  • forces_scale (float) – Forces scale.

  • device (str) – Device to use.

  • verbose (bool) – Verbosity.

Returns:

Dictionary with keys ‘energy’, ‘forces’, ‘ensemble’.

Return type:

Dict[str, array]

calculate(atoms=None, properties=['energy', 'forces', 'ensemble'], system_changes=['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'])[source]
Parameters:

  • atoms (ase.Atoms) – ASE atoms object.

  • properties (list of str) – energy, forces, and ensemble are supported.

  • system_changes (list of str) – List of changes for ASE.

get_ensemble()[source]

Get the calculated ensemble

class iann.calculators.AtomicEnsembleCalculator(model_paths=None, models=None, config=None, energy_scale=1.0, forces_scale=1.0, device=None, verbose=False, **kwargs)[source]

Atomic ensemble calculator for multiple models.

__init__(model_paths=None, models=None, config=None, energy_scale=1.0, forces_scale=1.0, device=None, verbose=False, **kwargs)[source]
Parameters:

  • model_paths (list of str) – Paths to the models (several models for ensemble).

  • models (list of iann.models.Model) – Models to use.

  • config (dict) – Configuration for the calculator.

  • energy_scale (float) – Energy scale.

  • forces_scale (float) – Forces scale.

  • device (str) – Device to use.

  • verbose (bool) – Verbosity.

Returns:

Dictionary with keys ‘energy’, ‘forces’, ‘ensemble’ including atomic energy variance.

Return type:

Dict[str, array]

calculate(atoms=None, properties=['energy', 'forces', 'ensemble'], system_changes=['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'])[source]
Parameters:

  • atoms (ase.Atoms) – ASE atoms object.

  • properties (list of str) – energy, forces, and ensemble are supported.

  • system_changes (list of str) – List of changes for ASE.

get_ensemble()[source]

Get the calculated ensemble

Plugins

class iann.plugins.converter.LAMMPSModelWrapper(model, compute_forces=True)[source]

A wrapper that adapts model inputs/outputs for LAMMPS.

__init__(model, compute_forces=True)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(num_atoms: Tensor, atomic_numbers: Tensor, positions: Tensor, cell: Tensor, edge_indices: Tensor, edge_vectors: Tensor, num_edges: Tensor) → Dict[str, Tensor][source]

Forward input tensors matching the C++ plugin call.

Parameters:

  • num_atoms (torch.Tensor) – Number of atoms.

  • atomic_numbers (torch.Tensor) – Atomic numbers.

  • positions (torch.Tensor) – Positions.

  • cell (torch.Tensor) – Cell.

  • edge_indices (torch.Tensor) – Edge indices.

  • edge_vectors (torch.Tensor) – Edge vectors.

  • num_edges (torch.Tensor)

Returns:

Dictionary with keys ‘energy’, ‘forces’, ‘atomic_energy’.

Return type:

Dict[str, torch.Tensor]

class iann.plugins.converter.EnsembleLAMMPSModelWrapper(models, compute_forces=True)[source]

A wrapper that adapts ensemble model inputs/outputs for LAMMPS.

__init__(models, compute_forces=True)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(num_atoms: Tensor, atomic_numbers: Tensor, positions: Tensor, cell: Tensor, edge_indices: Tensor, edge_vectors: Tensor, num_edges: Tensor) → Dict[str, Tensor][source]

Forward pass that computes ensemble averages and variances.

Parameters:

  • num_atoms (torch.Tensor) – Number of atoms.

  • atomic_numbers (torch.Tensor) – Atomic numbers.

  • positions (torch.Tensor) – Positions.

  • cell (torch.Tensor) – Cell.

  • edge_indices (torch.Tensor) – Edge indices.

  • edge_vectors (torch.Tensor) – Edge vectors.

  • num_edges (torch.Tensor)

Returns:

Dictionary with keys ‘energy’, ‘forces’, ‘energy_variance’, ‘forces_variance’, ‘atomic_energy_variance’.

Return type:

Dict[str, torch.Tensor]

iann.plugins.converter.convert_model_for_lammps(model_path, model_type=None, output_path=None, debug=False, atoms=None, **kwargs)[source]

Wrap a trained model in a TorchScript-compatible wrapper for LAMMPS.

Parameters:

  • model_path (str) – Path to the trained model checkpoint

  • model_type (str) – Type of model (painn, nequip, mace, equiformer2)

  • output_path (str, optional) – Path to save the exported model

Returns:

Path to the exported TorchScript model

Return type:

str

iann.plugins.converter.convert_models_for_lammps(model_paths, model_type, output_path=None, debug=False, atoms=None, **kwargs)[source]

Convert multiple models to a single TorchScript model for LAMMPS with ensemble statistics.

Parameters:

  • model_paths (list) – List of paths to trained model checkpoints

  • model_type (str) – Type of model (painn, nequip, mace, equiformer2)

  • output_path (str, optional) – Path to save the exported model

Returns:

Path to the exported TorchScript model

Return type:

str

C++ LAMMPS Plugins

class PairIANN

LAMMPS pair style (single GPU) for using trained IANN potentials in molecular dynamics simulations.

LAMMPS Command: pair_style iann model_type model_path cutoff

Parameters:

  • model_type: Type of ML model (fastpot, painn, nequip, mace, equiformer2)

  • model_path: Path to the exported TorchScript model file

  • cutoff: Interaction cutoff distance in Å

class PairIANNMultiGPU

IANN pair style (multiple GPUs) for using trained IANN potentials in molecular dynamics simulations.

LAMMPS Command: pair_style iann/multi_gpu model_type model_path cutoff

Parameters: Same as single GPU version PairIANN

class ComputeIANNVariance

Compute style for accessing ensemble variance statistics from IANN potentials.

LAMMPS Command: compute variance all iann/variance

Output Components:

  • c_variance[1]: Energy variance

  • c_variance[2]: Force variance

  • c_variance[3]: Maximum energy variance

  • c_variance[4]: Maximum force variance

Usage Example:

# Set up compute for variance
compute variance all iann/variance

# Use in thermo output
thermo_style custom step pe ke etotal press c_variance[1] c_variance[2] c_variance[3] c_variance[4]

thermo_modify colname c_variance[1] energy_variance
thermo_modify colname c_variance[2] force_variance
thermo_modify colname c_variance[3] max_energy_variance
thermo_modify colname c_variance[4] max_force_variance

For more information about specific functions and classes, see their respective module in source code.