mlcolvar.core.loss.SmartDerivatives

class mlcolvar.core.loss.SmartDerivatives(setup_device: str = 'cpu', force_all_atoms: bool = False)

Bases: Module

Utils to compute efficently (time and memory wise) the derivatives of the model output wrt the positions used to compute the input descriptors. Rather than computing explicitly the derivatives wrt the positions, we compute those wrt the descriptors (right input) and multiply them by the matrix of the derivatives of the descriptors wrt the positions (left input).

Overview

Preparation:

• Finds the non-zero entries of the derivatives of the descriptors wrt the positions (left)

• Stores such entries in a big 1D tensor

• Stores the indeces to find such entries in the original derivatives matrix

• Creates a big 1D tensor of indeces that allows properly taking together the contributions

Forward:

• Use the matrix indeces to retrieve the corresponding elements from the derivatives of output wrt the descriptors (right) into a big 1D tensor

• Get the single contributions via element-wise multiplication (i.e., of each atom to the output due to a single descriptor along a single space dimension)

• Scatter the single contributions to global contributions (of each atom to each output along each space dimension)

When working with batches or splits the scatter indeces are rescaled from the whole dataset to the batched entry.

__init__(setup_device: str = 'cpu', force_all_atoms: bool = False)

Initialize the smart derivatives object. To setup the class, use the setup method.

Parameters:

• der_desc_wrt_pos (torch.Tensor) – Tensor containing the derivatives of the descriptors wrt the atomic positions

• n_atoms (int) – Number of atoms in the systems, all the atoms should be used in at least one of the descriptors

• setup_device (str) – Device on which to perform the expensive calculations. Either ‘cpu’ or ‘cuda’, by default ‘cpu’

• force_all_atoms (bool) – Whether to allow the use of atoms that are non involved in the calculation of any descriptor, by default False

Methods

__init__([setup_device, force_all_atoms])	Initialize the smart derivatives object.
forward(x[, ref_idx])	Adds the derivatives of descriptors wrt atomic positions to the derivatives of output using the chain rule only for non-zero contributions.
move_to_cpu()	Moves the tensors used in the forward pass onto the cpu.
setup(dataset, descriptor_function, n_atoms)	Setup the smart derivatives object from a dataset and a descriptor function.

forward(x: Tensor, ref_idx: Tensor = None)

Adds the derivatives of descriptors wrt atomic positions to the derivatives of output using the chain rule only for non-zero contributions.

Parameters:

• x (torch.Tensor) – Derivatives of output wrt to descriptors

• ref_idx (torch.Tensor) – Reference indeces of the pristine dataset (i.e., before splitting, shuffling..)

Returns:

Derivatives of the output wrt atomic positions, shape (N, n_atoms, n_dim, (n_out)))

Return type:

torch.Tensor

move_to_cpu()

Moves the tensors used in the forward pass onto the cpu.

setup(dataset: DictDataset, descriptor_function: Transform, n_atoms: int, separate_boundary_dataset=False, descriptors_batch_size: int = None) → DictDataset

Setup the smart derivatives object from a dataset and a descriptor function. Returns a properly formatted new dataset with the descriptors as data.

Parameters:

• dataset (DictDataset) – Input dataset containing atomic positions as data and the needed entries

• descriptor_function (Transform) – Function to compute the descriptors from the atomic positions, it should be taken from the mlcolvar.core.tranform module

• n_atoms (int) – Number of atoms in the dataset

• separate_boundary_dataset (bool, optional) – Whether to separate the boundary dataset from the variational one, by default False NB: Should be used only for mlcolvar.cvs.committor.

• batch_size (int) – Size of batches to process data, useful for heavy computation to avoid memory overflows, if None a singel batch is used, by default None

Returns:

Updated dataset with the computed descriptors as ‘data’.

Return type:

DictDataset

Attributes

T_destination
call_super_init
dump_patches
training