torchani.transforms#

Composable functions that modify to properties when training or batching a dataset.

The API for the transforms in this module is modelled after the torchvision transforms, which you may be familiar with already. An example of their usage:

from torchani.transforms import (
    Compose,
    SubtractSAE,
    SubtractRepulsionXTB,
    SubtractTwoBodyDispersionD3,
)
from torchani.datasets import ANIBatchedDataset

symbols = ("H", "N", "O")
sae = SubtractSAE(symbols=symbols, (-0.5, -54.6, -75))
rep = SubtractRepulsionXTB(symbols=symbols)
disp = SubtractTwoBodyDispersionD3(symbols=symbols, functional='wB97X')
t = Compose([sae, rep, disp])

# Transforms will be applied automatically when iterating over the datasets
train = ANIBatchedDataset('/dataset/path/', transform=t, split='training')
valid = ANIBatchedDataset('/dataset/path/', transform=t, split='validation')

Classes

`AtomicNumbersToIndices`	Converts atomic numbers to indices
`Compose`	Compose several `torchani.transforms.Transform` into a pipeline
`Identity`	Pass-through transform
`SubtractEnergyAndForce`	Subtract the energies (and optionally forces) of a potential
`SubtractRepulsionXTB`	Subtract xTB repulsion energies (and optionally forces)
`SubtractSAE`	Subtract self atomic energies.
`SubtractTwoBodyDispersionD3`	Subtract two-body DFT-D3 energies (and optionally forces)
`Transform`	Base class for callables that modify mappings of molecule properties

class torchani.transforms.Transform(*args, **kwargs)[source]#

Base class for callables that modify mappings of molecule properties

If the callable supports only a limited number of atomic numbers (in a given order) then the atomic_numbers tensor should be defined, otherwise it should be None

forward(properties)[source]#

Transform a batch of properties

Parameters:: properties (Dict[str, Tensor]) – Input properties
Returns:: Transformed properties
Return type:: Dict[str, Tensor]

class torchani.transforms.Identity[source]#

Pass-through transform

forward(properties)[source]#

Transform a batch of properties

Parameters:: properties (Dict[str, Tensor]) – Input properties
Returns:: Transformed properties
Return type:: Dict[str, Tensor]

class torchani.transforms.SubtractEnergyAndForce(potential, subtract_force=True)[source]#

Subtract the energies (and optionally forces) of a potential

Parameters:

potential (Potential) – The potential to use for calculating energies and forces
subtract_force (bool) – Whether to subtract forces

forward(properties)[source]#

Transform a batch of properties

Parameters:: properties (Dict[str, Tensor]) – Input properties
Returns:: Transformed properties
Return type:: Dict[str, Tensor]

class torchani.transforms.SubtractRepulsionXTB(*args, subtract_force=True, **kwargs)[source]#

Subtract xTB repulsion energies (and optionally forces)

Takes same arguments as torchani.potentials.RepulsionXTB

forward(properties)[source]#

Transform a batch of properties

Parameters:: properties (Dict[str, Tensor]) – Input properties
Returns:: Transformed properties
Return type:: Dict[str, Tensor]

class torchani.transforms.SubtractTwoBodyDispersionD3(*args, subtract_force=True, **kwargs)[source]#

Subtract two-body DFT-D3 energies (and optionally forces)

Takes same arguments as torchani.potentials.TwoBodyDispersionD3

forward(properties)[source]#

Transform a batch of properties

Parameters:: properties (Dict[str, Tensor]) – Input properties
Returns:: Transformed properties
Return type:: Dict[str, Tensor]

class torchani.transforms.SubtractSAE(symbols, self_energies)[source]#

Subtract self atomic energies.

Takes same arguments as torchani.sae.SelfEnergy

forward(properties)[source]#

Transform a batch of properties

Parameters:: properties (Dict[str, Tensor]) – Input properties
Returns:: Transformed properties
Return type:: Dict[str, Tensor]

class torchani.transforms.AtomicNumbersToIndices(symbols)[source]#

Converts atomic numbers to indices

Note

Provided for backwards compatibility, if added to a transform pipeline, it should in general be the last transform.

Parameters:: symbols (Sequence[str]) – A tuple or list of strings that are valid chemical symbols. (case sensitive).

forward(properties)[source]#

Transform a batch of properties

Parameters:: properties (Dict[str, Tensor]) – Input properties
Returns:: Transformed properties
Return type:: Dict[str, Tensor]

class torchani.transforms.Compose(transforms)[source]#

Compose several torchani.transforms.Transform into a pipeline

Parameters:: transforms (Sequence[Transform]) – Transforms to compose, in order.

forward(properties)[source]#

Transform a batch of properties

Parameters:: properties (Dict[str, Tensor]) – Input properties
Returns:: Transformed properties
Return type:: Dict[str, Tensor]