convert_transforms

class AddAtomCooIndices[source]

Add a field “atom_coo_indices” to the sample, containing an index tensor of shape (2, n_basis), which can be used to construct a block diagonal sparse matrix with each block corresponding to the basis functions of one atom.

__call__(sample: OFData) → OFData[source]

Apply the transform to the sample.

See add_atom_coo_indices() for details.

class AddBasisInfo(basis_info: BasisInfo)[source]

Adds the nuclear attraction vector to the sample.

__call__(sample: OFData) → OFData[source]

Parameters:: sample – the molecule in the OFData format

__init__(basis_info: BasisInfo)[source]

Parameters:: basis_info – The basis information in the BasisInfo format.

class AddCoulombMatrix(basis_info: BasisInfo)[source]

Adds the coulomb matrix to the sample.

__call__(sample: OFData) → OFData[source]

Parameters:: sample – the molecule in the OFData format

__init__(basis_info: BasisInfo)[source]

Parameters:: basis_info – The basis information in the BasisInfo format.

class AddFullEdgeIndex(*args, **kwargs)[source]

Add a full edge index to the sample.

forward(sample: OFData) → OFData[source]

Add a full edge index to the sample and returns the same sample with the additional edge_index attribute.

Parameters:: sample (OFData) – Sample data object.
Returns:: The same sample data object with the additional edge_index attribute.
Return type:: OFData

class AddMol(basis_info: BasisInfo)[source]

Adds the mol to the sample.

__call__(sample: OFData) → OFData[source]

Parameters:: sample – the molecule in the OFData format

__init__(basis_info: BasisInfo)[source]

Parameters:: basis_info – The basis information in the BasisInfo format.

class AddNuclearAttractionVector(basis_info: BasisInfo)[source]

Adds the nuclear attraction vector to the sample.

__call__(sample: OFData) → OFData[source]

Parameters:: sample – the molecule in the OFData format

__init__(basis_info: BasisInfo)[source]

Parameters:: basis_info – The basis information in the BasisInfo format.

class AddOverlapMatrix(basis_info: BasisInfo)[source]

Adds the overlap matrix to the sample.

__call__(sample: OFData) → OFData[source]

Parameters:: sample – the molecule in the OFData format

__init__(basis_info: BasisInfo)[source]

Parameters:: basis_info – The basis information in the BasisInfo format.

class AddRadiusEdgeIndex(radius: float)[source]

Add a radius edge index to the sample.

__init__(radius: float)[source]

Parameters:: radius – The radius to use for the edge index.

forward(sample: OFData) → OFData[source]

Add a radius edge index to the sample and returns the same sample with the additional edge_index attribute.

Parameters:: sample (OFData) – Sample data object.
Returns:: The same sample data object with the additional edge_index attribute.
Return type:: OFData

class ProjectGradient(*args, **kwargs)[source]

Project gradients stored in a sample onto the tangent space of the manifold of constant electron number.

forward(sample: OFData) → OFData[source]

Project the fields of the sample as configured in __init__().

Parameters:: sample (OFData) – The sample.
Returns:: The sample with the projected gradients
Return type:: OFData

class SplitByAtom[source]

Split all basis-function wise fields by atom.

__call__(sample: OFData) → OFData[source]

Apply the transform to the sample.

See split_by_atom() for details.

class ToNumpy[source]

Convert all torch tensors in the sample to numpy arrays.

__call__(sample: OFData) → OFData[source]

Apply the transform to the sample.

See to_numpy() for details.

class ToTorch(device: device = None, float_dtype: dtype | dtype | None = None)[source]

Convert all numpy arrays in the sample to torch tensors.

__call__(sample: OFData) → OFData[source]

Apply the transform to the sample.

See to_torch() for details.

__init__(device: device = None, float_dtype: dtype | dtype | None = None)[source]

Initialize the transform.

Parameters:

device – The device to put the tensors on. Defaults to None, i.e. the pytorch default device.
float_dtype – The dtype to use for float dtypes. Defaults to the pytorch default dtype. For strings, “torch.float64” and “torch.float32” are supported.

add_atom_coo_indices(sample: OFData) → OFData[source]

Add a field “atom_coo_indices” to the sample, containing an index tensor of shape (2, n_basis), which can be used to construct a block diagonal sparse matrix with each block corresponding to the basis functions of one atom.

Parameters:: sample – The sample.
Returns:: The sample with the field “atom_coo_indices” added.

apply_to_attributes(f: Callable, sample: OFData, attributes: tuple) → OFData[source]

Simplified version of Data.apply(), with the crucial difference that errors are not quietly ignored.

Applies the function f to the attributes of the object obj specified in the tuple attributes.

Warning

This function modifies the object in-place.

Parameters:

f – The function to apply.
sample – The object to apply the function to.
attributes – The names of the attributes to apply the function to.

Returns:

The object with the attributes modified.

dtype_map(dtype: dtype, float_dtype: dtype | None | str = None) → None | dtype[source]

Map a numpy dtype to a torch dtype.

Parameters:

dtype – The numpy dtype.
float_dtype – The dtype to use for float dtypes. Defaults to the pytorch default dtype. For strings, “torch.float64” and “torch.float32” are supported.

Returns:

The torch dtype, or None if no mapping exists.

split_by_atom(sample: OFData, fields=('coeffs', 'ground_state_coeffs', 'gradient_label')) → OFData[source]

Split all basis-function wise fields by atom. Currently, not planned to be used, but demonstrates how fields can be split when needed.

Parameters:

sample – The sample.
fields – The fields to split.

Returns:

The sample with the fields split by atom.

str_to_torch_float_dtype(dtype: str | dtype) → dtype[source]

Convert a string to a torch float dtype.

Useful to set the dtype in hydra configs.

to_numpy(sample: OFData) → OFData[source]: Convert all torch tensors in the sample to numpy arrays.

to_torch(sample: OFData, device=None, float_dtype: dtype | dtype | str | None = None) → OFData[source]

Convert all numpy arrays in the sample to torch tensors.

Parameters:

sample – The sample.
device – The device to put the tensors on. Defaults to None, i.e. the pytorch default device.
float_dtype – The dtype to use for float dtypes. Defaults to the pytorch default dtype. For strings, “torch.float64” and “torch.float32” are supported to enable hydra support.

Returns:

The sample with all numpy arrays converted to torch tensors.