MolecularDiffusion.modules.models.tabasco.data.transforms¶

Attributes¶

logger

Functions¶

`apply_random_rotation`(→ tensordict.TensorDict)	Augment a batch with n_augmentations additional random rotations.
`permute_atoms`(→ tensordict.TensorDict)	Randomly permute the order of non-padded atoms in a molecule tensor.
`random_rotation`(→ tensordict.TensorDict)	Apply a single random 3D rotation to data['coords'].
`sample_uniform_rotation`(→ torch.Tensor)	Sample uniform random rotation matrices.

Module Contents¶

MolecularDiffusion.modules.models.tabasco.data.transforms.apply_random_rotation(batch: tensordict.TensorDict, n_augmentations=10) → tensordict.TensorDict¶

Augment a batch with n_augmentations additional random rotations.

Parameters:

batch – TensorDict with keys coords, padding_mask, atomics.
n_augmentations – Number of extra rotated copies to generate.

Returns:

TensorDict with keys coords, padding_mask, atomics and n_augmentations extra copies.

Reference: NVIDIA-Digital-Bio/proteina implementation.

MolecularDiffusion.modules.models.tabasco.data.transforms.permute_atoms(data: tensordict.TensorDict) → tensordict.TensorDict¶: Randomly permute the order of non-padded atoms in a molecule tensor.

MolecularDiffusion.modules.models.tabasco.data.transforms.random_rotation(data: tensordict.TensorDict) → tensordict.TensorDict¶: Apply a single random 3D rotation to data[‘coords’].

MolecularDiffusion.modules.models.tabasco.data.transforms.sample_uniform_rotation(shape: torch.Size, dtype: torch.dtype, device: torch.device) → torch.Tensor¶

Sample uniform random rotation matrices.

Reference: NVIDIA-Digital-Bio/proteina implementation.

MolecularDiffusion.modules.models.tabasco.data.transforms.logger¶