MolecularDiffusion.modules.models.syncogen.api.atomics.coordinates

Classes

Coordinates

Coordinates class for handling molecular coordinates.

Module Contents

class MolecularDiffusion.modules.models.syncogen.api.atomics.coordinates.Coordinates(coordinates: torch.Tensor, atom_mask: torch.Tensor | None = None, is_batched: bool | None = None)

Coordinates class for handling molecular coordinates.

Use factory methods for construction:

  • Coordinates.from_tensor(coordinates, atom_mask)

  • Coordinates.random(shape, atom_mask, device)

Or use the constructor directly with a coordinates tensor.

Initialize from coordinates tensor (canonical constructor).

Parameters:

  • coordinates – Coordinate tensor, shape (N, 3) or (B, N, 3)

  • atom_mask – Optional mask tensor, shape (N,) or (B, N). Defaults to all ones.

  • is_batched – Whether this is batched. Inferred from coordinates.dim() if None.

apply_mask(mask: torch.Tensor)

Apply a mask to the coordinates tensor.

Parameters:

mask – Tensor of shape matching self.tensor up to the last dimension

attach_bonds(bonds: torch.Tensor, bonds_mask: torch.Tensor)

Attach bonds to coordinates with explicit mask.

Parameters:

  • bonds – […, n_bonds, 3] where each entry is (i_flat, j_flat, bond_type_index)

  • bonds_mask – […, n_bonds] boolean mask of valid bonds

attach_pharmacophores(pharm_coords: torch.Tensor, pharm_padding_mask: torch.Tensor)

Attach pharmacophores to coordinates with explicit mask.

  • Keeps atoms and pharm as separate tensors internally

  • Combined mask is computed dynamically via property

center(custom_mask=None, custom_center=None)

Center using atom mask only; move pharmacophores along if attached.

clone() → Coordinates

Create a deep copy of this Coordinates object.

classmethod from_tensor(coordinates: torch.Tensor, atom_mask: torch.Tensor | None = None, is_batched: bool | None = None) → Coordinates

Construct from existing coordinates tensor.

Parameters:

  • coordinates – Coordinate tensor, shape (N, 3) or (B, N, 3)

  • atom_mask – Optional mask tensor

  • is_batched – Whether batched (inferred if None)

get_center(custom_mask=None)

Get the center of mass according to atom mask.

is_centered(tol=1e-06)

Check if coordinates are centered according to atom mask.

kabsch_align_to(reference: torch.Tensor, mask: torch.Tensor = None, weights: torch.Tensor = None)

Align (rotate) coordinates to reference using Kabsch (weighted rigid align).

Parameters:

  • reference – Tensor of shape [N, 3] to align to

  • mask – Optional mask of shape [N] for valid points

  • weights – Optional weights of shape [N]

Returns:

self

classmethod random(shape: torch.Size | Tuple[int, ...], atom_mask: torch.Tensor | None = None, device: torch.device | None = None, dtype: torch.dtype = torch.float32, is_batched: bool | None = None) → Coordinates

Construct with random normal coordinates.

Parameters:

  • shape – Shape of coordinates tensor, e.g. (N, 3) or (B, N, 3)

  • atom_mask – Optional mask tensor. Defaults to all ones.

  • device – Target device

  • dtype – Data type for coordinates

  • is_batched – Whether batched (inferred from shape if None)

random_rotate()

Apply random rotation to coordinates (and pharmacophores if attached).

random_translate(scale: float = 1.0)

Apply random translation to coordinates (and pharmacophores if attached).

scale(multiplier)

Scale coordinates (and pharmacophores if attached) by a multiplier.

set_coordinates(coordinates: torch.Tensor, apply_mask: bool = True)

Replace coordinates tensor; must match existing tensor shape. Optionally re-apply mask zeroing.

set_mask(atom_mask: torch.Tensor, apply_mask: bool = True)

Set a new atom mask; must match existing mask shape. Recomputes n_atoms and re-applies zeroing.

to(device: torch.device | None = None, dtype: torch.dtype | None = None)

Move coordinates to specified device and/or dtype.

Parameters:

  • device – Target device (e.g., ‘cuda’, ‘cpu’)

  • dtype – Target dtype (e.g., torch.float32)

to_numpy()

Convert coordinates to numpy array.

to_torch()

Convert coordinates to torch tensor.

classmethod uniform(shape: torch.Size | Tuple[int, ...], atom_mask: torch.Tensor | None = None, device: torch.device | None = None, dtype: torch.dtype = torch.float32, is_batched: bool | None = None) → Coordinates

Construct with all-zero coordinates (a ‘uniform’ tensor of zeros).

Parameters:

  • shape – Shape of coordinates tensor, e.g. (N, 3) or (B, N, 3)

  • atom_mask – Optional mask tensor. Defaults to all ones.

  • device – Target device

  • dtype – Data type for coordinates

  • is_batched – Whether batched (inferred from shape if None)

classmethod zeros(shape: torch.Size | Tuple[int, ...], atom_mask: torch.Tensor | None = None, device: torch.device | None = None, dtype: torch.dtype = torch.float32, is_batched: bool | None = None) → Coordinates

Construct with zero coordinates.

Parameters:

  • shape – Shape of coordinates tensor

  • atom_mask – Optional mask tensor

  • device – Target device

  • dtype – Data type

  • is_batched – Whether batched

property atom_and_pharmacophore_mask

Dynamically compute concatenated atom + pharmacophore mask in atoms view.

property atom_coords

Get atom coordinates tensor.

atom_mask
bonds = None
bonds_mask = None
property device

Get device of coordinates.

property dtype

Get dtype of coordinates.

has_bonds = False
has_pharmacophores = False
is_batched = None
pharm_coords = None
pharm_padding_mask = None
property pharmacophores

Get pharmacophore coordinates (everything after max_atoms).

property shape

Get shape of coordinates.

tensor