basis_transforms

Transforms that act on OFData objects.

The idea is to apply a subset of transforms as part of the loader, i.e. execute them per-molecule on the cpu, before batches of multiple molecules are assembled and moved to the GPU.

class AddLocalFrames

Add the local frames to the sample.

__call__(sample: OFData, invert: bool = False) → OFData

Add the local frames to the sample.

Parameters:

sample – The sample.

__init__()

class ApplyBasisTransformation(*args, **kwargs)

Transform all fields in the sample according to their representation.

The transformation is determined by the transformation matrix and its inverse. The representation should be saved in the sample.representations.

forward(sample: OFData, transformation_matrix: Tensor, inv_transformation_matrix: Tensor, invert: bool = False) → OFData

Parameters:

• sample – The sample to transform.

• transformation_matrix – The transformation matrix.

• inv_transformation_matrix – The inverse of the transformation matrix.

• invert – Whether to invert the transformation.

Returns:

The transformed sample.

class MasterTransformation(name: str, use_cached_data: bool = True, pre_transforms: list = None, cached_transforms: DictConfig = None, basis_transforms: list[ApplyBasisTransformation] = None, post_transforms: list = None, add_transformation_matrix: bool = False, **kwargs)

__init__(name: str, use_cached_data: bool = True, pre_transforms: list = None, cached_transforms: DictConfig = None, basis_transforms: list[ApplyBasisTransformation] = None, post_transforms: list = None, add_transformation_matrix: bool = False, **kwargs) → None

Class that handles the basis transformations.

Before using the forward method, the configure_cached_transforms method should be called which configures whether there are cached basis transforms that can be used.

Parameters:

• target – The target transformation.

• pre_transforms – List of transforms to apply before the basis transforms.

• cached_transforms – List of basis transforms that are cached and can be used if the cached option is set.

• basis_transforms – List of basis transforms to apply.

• post_transforms – List of transforms to apply after the basis transforms

• add_transformation_matrix – Whether to add the transformation matrix to the sample.

basis_transform(sample: OFData) → OFData

Apply the forward basis transforms to a sample.

forward(sample: OFData) → OFData

Apply the configured transformation to a sample.

static initialize_transformation_matrices(sample: OFData)

Add the transformation matrices to the sample.

invert_basis_transform(sample: OFData) → OFData

Invert the basis transforms of a sample.

Inverse transforms are currently used for visualization during training and for obtaining the final density in ofdft. This will be done in the torch default float dtype.

class ToGlobalNatRep(orthogonalization: Literal['symmetric', 'canonical'] = 'symmetric')

Apply the global natural reparametrization to the sample.

__call__(sample: OFData, invert: bool = False) → OFData

Naturally reparametrize the sample in-place.

Parameters:

sample – The sample to be transformed.

__init__(orthogonalization: Literal['symmetric', 'canonical'] = 'symmetric')

Parameters:

orthogonalization – The type of orthogonalization to use. Either “symmetric” (default) or “canonical”.

class ToLocalFrames(sparse: bool = True)

Apply the local frames transformation to the sample.

__call__(sample: OFData, invert: bool = False) → OFData

Transform the sample to local frames in-place.

Parameters:

sample – The sample.

__init__(sparse: bool = True)

Parameters:

sparse – Whether to use a sparse matrix for the transformation.

transform_tensor(tensor: Tensor, transformation_matrix: Tensor, inv_transformation_matrix: Tensor, representation: Representation)

Transform tensor according to given representation.

The behavior of the transformation is determined by its representation. The available options for Representation are:

• Representation.NONE or Representation.SCALAR: no transformation is applied.

• Representation.VECTOR: e.g. the density coefficients \(p\),

  \[p' = T p.\]

  This also determines the convention for the transformation matrix.

• Representation.DUAL_VECTOR and Representation.GRADIENT e.g. gradients, vectors that are applied to other vectors like the nuclear attraction vector, the basis_integrals and the inverse transformation matrix \(M = \tilde T^{-1}\) after another transformation \(T\),

  \[M' = M T^{-1}.\]

  Gradients can also be projected but the other vectors can’t.

• Representation.ENDOMORPHISM: e.g. the projection matrix,

  \[P' = T P T^{-1}.\]

• Representation.BILINEAR_FORM: e.g. the coulomb matrix,

  \[C' = T^{-T} C T^{-1}.\]

• Representation.AO: the atomic orbitals, these transform

  \[A' = A T^{-1}\]

  which is the same as the inverse vector representation.

Parameters:

• tensor – The tensor to be transformed.

• transformation_matrix – The transformation matrix.

• inv_transformation_matrix – The inverse of the transformation matrix.

• representation – The representation that defines the transformation behavior of the tensor.

Returns:

The transformed tensor.

Raises:

KeyError – If the representation is not known.

transform_tensor_with_sample(sample: OFData, tensor: Tensor, representation: Representation, invert: bool = False) → Tensor

Transform the given tensor with the transformation that the sample was transformed with.

The transformation_matrix and inv_transformation_matrix have to be added to the sample before it is transformed.

Parameters:

• sample – The sample that was already transformed and contains the transformation matrices.

• tensor – The tensor to be transformed.

• representation – The representation of the tensor to transform.

• invert – Whether to invert the transformation.

Returns:

The transformed tensor.