Modules

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

Bases: Dataset

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

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

Bases: Calculator

Atomic ensemble calculator for multiple models.

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

implemented_properties: list[str] = ['energy', 'forces', 'ensemble']

Properties calculator can handle (energy, forces, …)

class iann.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]

Bases: NamedTuple

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.

atomic_energy: torch.Tensor | None

Alias for field number 12

atomic_numbers: torch.Tensor

Alias for field number 1

cell: torch.Tensor

Alias for field number 3

contiguous()[source]
edge_diff_embedding: torch.Tensor | None

Alias for field number 16

edge_dist_embedding: torch.Tensor | None

Alias for field number 15

edge_indices: torch.Tensor

Alias for field number 4

edge_vectors: torch.Tensor

Alias for field number 5

energy: torch.Tensor | None

Alias for field number 8

energy_variance: torch.Tensor | None

Alias for field number 17

forces: torch.Tensor | None

Alias for field number 9

forces_variance: torch.Tensor | None

Alias for field number 18

global_attr: torch.Tensor

Alias for field number 7

global_embedding: torch.Tensor | None

Alias for field number 19

image_indices: torch.Tensor | None

Alias for field number 11

keys()[source]
node_attr: torch.Tensor | None

Alias for field number 13

node_feat: torch.Tensor | None

Alias for field number 14

num_atoms: torch.Tensor

Alias for field number 0

num_edges: torch.Tensor

Alias for field number 6

positions: torch.Tensor

Alias for field number 2

stress: torch.Tensor | None

Alias for field number 10

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

Bases: Calculator

Ensemble calculator for multiple models.

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

implemented_properties: list[str] = ['energy', 'forces', 'ensemble']

Properties calculator can handle (energy, forces, …)

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

Bases: Calculator

Machine learning calculator for a single model.

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.

implemented_properties: list[str] = ['energy', 'forces', 'stress']

Properties calculator can handle (energy, forces, …)

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

Bases: object

Trainer class for training interatomic neural network models

eval_model()[source]

Evaluate model on validation set

train(dataset_path=None)[source]

Train the model on a dataset

Parameters:

dataset_path (str) – Path to the dataset file

iann.get_model_class(model_name)[source]

Get a model class by name with lazy loading