MolecularDiffusion.modules.models.difflinker.utils¶

Small numerical helpers ported near-verbatim from DiffLinker’s src/utils.py (only the subset actually used by edm.py/egnn.py in this integration; visualization/logging/EMA helpers were dropped as unused).

Exceptions¶

FoundNaNException

Common base class for all non-exit exceptions.

Functions¶

`assert_correctly_masked`(→ None)
`assert_mean_zero_with_mask`(→ None)
`assert_partial_mean_zero_with_mask`(→ None)
`remove_mean_with_mask`(→ torch.Tensor)
`remove_partial_mean_with_mask`(→ torch.Tensor)	Subtract center of mass of fragments from coordinates of all atoms.
`sample_center_gravity_zero_gaussian_with_mask`(...)
`sample_gaussian_with_mask`(→ torch.Tensor)
`split_features`(z, n_dims, num_classes, include_charges)
`sum_except_batch`(→ torch.Tensor)

Module Contents¶

exception MolecularDiffusion.modules.models.difflinker.utils.FoundNaNException(x, h)¶

Bases: Exception

Common base class for all non-exit exceptions.

Initialize self. See help(type(self)) for accurate signature.

static find_nan_idx(z)¶

only_h_nan_idx¶

only_x_nan_idx¶

x_h_nan_idx¶

MolecularDiffusion.modules.models.difflinker.utils.assert_correctly_masked(variable: torch.Tensor, node_mask: torch.Tensor) → None¶

MolecularDiffusion.modules.models.difflinker.utils.assert_mean_zero_with_mask(x: torch.Tensor, node_mask: torch.Tensor, eps: float = 1e-10) → None¶

MolecularDiffusion.modules.models.difflinker.utils.assert_partial_mean_zero_with_mask(x: torch.Tensor, node_mask: torch.Tensor, center_of_mass_mask: torch.Tensor, eps: float = 1e-10) → None¶

MolecularDiffusion.modules.models.difflinker.utils.remove_mean_with_mask(x: torch.Tensor, node_mask: torch.Tensor) → torch.Tensor¶

MolecularDiffusion.modules.models.difflinker.utils.remove_partial_mean_with_mask(x: torch.Tensor, node_mask: torch.Tensor, center_of_mass_mask: torch.Tensor) → torch.Tensor¶: Subtract center of mass of fragments from coordinates of all atoms.

MolecularDiffusion.modules.models.difflinker.utils.sample_center_gravity_zero_gaussian_with_mask(size, device, node_mask) → torch.Tensor¶

MolecularDiffusion.modules.models.difflinker.utils.sample_gaussian_with_mask(size, device, node_mask) → torch.Tensor¶

MolecularDiffusion.modules.models.difflinker.utils.split_features(z, n_dims, num_classes, include_charges)¶

MolecularDiffusion.modules.models.difflinker.utils.sum_except_batch(x: torch.Tensor) → torch.Tensor¶