schnetpack.transform

Transforms are applied before and/or after the model. They can be used, e.g., for calculating neighbor lists, casting, unit conversion or data augmentation. Some can applied before batching, i.e. to single systems, when loading the data. This is necessary for pre-processing and includes neighbor lists, for example. On the other hand, transforms need to be able to handle batches for post-processing. The flags is_postprocessor and is_preprocessor indicate how the tranforms may be used. The attribute mode of a transform is set automatically to either “pre” or “post”.q

class transform.Transform(*args: Any, **kwargs: Any)[source]

Base class for all transforms. The base class ensures that the reference to the data and datamodule attributes are initialized. Transforms can be used as pre- or post-processing layers. They can also be used for other parts of a model, that need to be initialized based on data.

To implement a new transform, override the forward method. Preprocessors are applied to single examples, while postprocessors operate on batches. All transforms should return a modified inputs dictionary.

Atomistic

`AddOffsets`	Add offsets to property based on the mean of the training data and/or the single atom reference calculations.
`RemoveOffsets`	Remove offsets from property based on the mean of the training data and/or the single atom reference calculations.
`SubtractCenterOfMass`	Subtract center of mass from positions.
`SubtractCenterOfGeometry`	Subtract center of geometry from positions.

Casting

`CastMap`	Cast all inputs according to type map.
`CastTo32`	Cast all float64 tensors to float32
`CastTo64`	Cast all float32 tensors to float64

Neighbor lists

`MatScipyNeighborList`	Neighborlist using the efficient implementation of the Matscipy package
`ASENeighborList`	Calculate neighbor list using ASE.
`TorchNeighborList`	Environment provider making use of neighbor lists as implemented in TorchAni
`CachedNeighborList`	Dynamic caching of neighbor lists.
`CountNeighbors`	Store the number of neighbors for each atom
`FilterNeighbors`	Filter out all neighbor list indices corresponding to interactions between a set of atoms.
`WrapPositions`	Wrap atom positions into periodic cell.
`CollectAtomTriples`	Generate the index tensors for all triples between atoms within the cutoff shell.