torchani.arch

Construction of ANI-style models and definition of their architecture

ANI-style models are all subclasses of the ANI base class. They can be either constructed directly, or their construction can be managed by a different class, the Assembler. Think of the Assembler as a helpful friend who create the class from all the necessary components, in such a way that all parts interact in the correct way and there are no compatibility issues among them.

An ANI-style model consists of:

• torchani.aev.AEVComputer (or subclass) which is an atomic featurizer with 2-body

  and 3-body terms.

• A container for atomic networks (typically torchani.nn.ANINetworks, which takes a dictionary for networks {"H": AtomicNetwork(...), "C": AtomicNetwork(...), ...} for the mixture-of-experts model).

• A torchani.sae.SelfEnergy which takes a self energies dict (in Hartree): {"H": -12.0, "C": -75.0, ...}

These pieces are combined when the Assembler.assemble method is called.

An energy-predicting model may also have one or more torchani.potentials.Potential (torchani.potentials.RepulsionXTB, torchani.potentials.TwoBodyDispersionD3, etc.).

Each torchani.potentials.Potential has its own cutoff, and the torchani.aev.AEVComputer has two cutoffs, an angular and a radial one (the radial cutoff must be larger than the angular cutoff, and it is recommended that the angular cutoff is kept small, 3.5 Ang or less).

Functions

simple_ani	Flexible builder to create ANI-style models that predict energies
simple_aniq	Flexible builder to create ANI-style models that output charges and energies

Classes

ANI	ANI-style neural network interatomic potential
ANIq	ANI-style model that can calculate both atomic charges and energies
Assembler	Assembles an ANI model (or subclass)

class torchani.arch.ANI(symbols, aev_computer, neural_networks, energy_shifter, potentials=None, periodic_table_index=True)

ANI-style neural network interatomic potential

forward(species_coordinates, cell=None, pbc=None, charge=0, atomic=False, ensemble_values=False, _molecule_idxs=None)

Obtain a species-energies tuple from an input species-coords tuple

compute_from_neighbors(elem_idxs, coords, neighbors, charge=0, atomic=False, ensemble_values=False)

This entrypoint supports input from TorchANI neighbors

Returns a tuple:

• molecular-scalars

• atomic-scalars

class torchani.arch.ANIq(symbols, aev_computer, neural_networks, energy_shifter, potentials=None, periodic_table_index=True, charge_networks=None, charge_normalizer=None)

ANI-style model that can calculate both atomic charges and energies

Charge networks share the input features with the energy networks, and may either be fully independent of them, or share weights to some extent.

The output energies of these models don’t necessarily include a coulombic term, but they may.

compute_from_neighbors(elem_idxs, coords, neighbors, charge=0, atomic=False, ensemble_values=False)

This entrypoint supports input from TorchANI neighbors

Returns a tuple:

• molecular-scalars

• atomic-scalars

forward(species_coordinates, cell=None, pbc=None, charge=0, atomic=False, ensemble_values=False, _molecule_idxs=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchani.arch.Assembler(symbols=(), cls=<class 'torchani.arch.ANI'>, neighborlist='all_pairs', periodic_table_index=True)

Assembles an ANI model (or subclass)

set_aev_computer(angular, radial, cutoff_fn='global', strategy='pyaev', aev_computer_cls=<class 'torchani.aev._computer.AEVComputer'>)

Set the AEVComputer local atomic featurizer

set_neighborlist(neighborlist)

Set the neighborlist the model will use

set_global_cutoff_fn(cutoff_fn)

Set the default cutoff function for all modules of the model

add_potential(cls, name, cutoff=inf, cutoff_fn='global', kwargs=None)

Add a potential to the constructed ANI model

assemble(ensemble_size=1)

Construct an ANI model from the passed arguments

Parameters:

ensemble_size (int) – The size of the constructed ensemble

Returns:

ANI model, ready to train.

Return type:

ANI | ANIq

torchani.arch.simple_ani(symbols, lot, ensemble_size=1, radial_start=0.9, angular_start=0.9, radial_cutoff=5.2, angular_cutoff=3.5, radial_shifts=16, angular_shifts=8, sections=4, radial_precision=19.7, angular_precision=12.5, angular_zeta=14.1, cutoff_fn='smooth', dispersion=False, repulsion=True, container_ctor='default', container='ANINetworks', activation='gelu', bias=False, strategy='auto', periodic_table_index=True, neighborlist='all_pairs', repulsion_cutoff=True, dispersion_cutoff=8.0, self_energies=None)

Flexible builder to create ANI-style models that predict energies

Defaults are similar to ANI-2x, with some improvements.

To reproduce the ANI-2x AEV exactly use the following args:

• cutoff_fn='cosine'

• radial_start=0.8

• angular_start=0.8

• radial_cutoff=5.1

torchani.arch.simple_aniq(symbols, lot, ensemble_size=1, radial_start=0.9, angular_start=0.9, radial_cutoff=5.2, angular_cutoff=3.5, radial_shifts=16, angular_shifts=8, sections=4, radial_precision=19.7, angular_precision=12.5, angular_zeta=14.1, cutoff_fn='smooth', dispersion=False, repulsion=True, container_ctor='default', charge_container_ctor='default', container='ANINetworks', charge_container='ANINetworks', activation='gelu', bias=False, strategy='auto', merge_charge_networks=False, scale_charge_normalizer_weights=True, dummy_energies=False, use_cuda_ops=False, periodic_table_index=True, neighborlist='all_pairs', normalize=True, repulsion_cutoff=True, dispersion_cutoff=8.0, self_energies=None)

Flexible builder to create ANI-style models that output charges and energies

Defaults are similar to ANI-2x, with some improvements.

To reproduce the ANI-2x AEV exactly use the following args:

• cutoff_fn='cosine'

• radial_start=0.8

• angular_start=0.8

• radial_cutoff=5.1