MolecularDiffusion.modules.models.ligandiff.utils

LigandDiff helpers.

remove_partial_mean_with_mask and FoundNaNException are verbatim from the target repo’s src/utils.py. ligand_groups_from_geometry is new – see its docstring.

Attributes

METAL_Z

Exceptions

FoundNaNException

Raised by Dynamics.forward when the denoiser output goes NaN.

Functions

ligand_groups_from_geometry(→ torch.Tensor)

Recover LigandDiff's ligand_group one-hot from geometry alone.

remove_partial_mean_with_mask(x, center_of_mass_mask, ...)

Subtract the masked subset's centre of mass from every coordinate.

Module Contents

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

Bases: Exception

Raised by Dynamics.forward when the denoiser output goes NaN.

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.ligandiff.utils.ligand_groups_from_geometry(coords: torch.Tensor, charges: torch.Tensor, n_slots: int = 6, scale_factor: float = 1.25) → torch.Tensor

Recover LigandDiff’s ligand_group one-hot from geometry alone.

Upstream derives this column block with molSimplify’s ligand_breakdown (generate.py:81-87): each atom is tagged with the index of the ligand it belongs to, the metal getting an all-zero row. molSimplify is not a dependency here and generation input is a bare .xyz, so the same decomposition is recovered directly: delete the metal(s), then take the connected components of the covalent-radius contact graph – which is what “a ligand” means for a mononuclear complex.

Parameters:

  • coords – (N, 3) positions.

  • charges – (N,) atomic numbers (metals identified via METAL_Z).

  • n_slots – number of ligand columns (upstream: 6). Components beyond this are folded into the last slot.

  • scale_factor – covalent-radius multiplier for the contact criterion.

Returns:

(N, n_slots) float one-hot; metal rows are all zero.

MolecularDiffusion.modules.models.ligandiff.utils.remove_partial_mean_with_mask(x, center_of_mass_mask, batch_seg)

Subtract the masked subset’s centre of mass from every coordinate.

MolecularDiffusion.modules.models.ligandiff.utils.METAL_Z: Set[int]