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)

• A container for atomic networks (typically torchani.nn.ANINetworks or subclass)

• If using the traditional ANINetworks, a torchani.nn.AtomicNetwork mapping, for example: {"H": AtomicNetwork(...), "C": AtomicNetwork(...), ...}.

• 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
simple_aniq	Flexible builder to create ANI-style models that output charges

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)

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)

Flexible builder to create ANI-style models

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)

Flexible builder to create ANI-style models that output charges

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