Molecular graphs, equivariance and docking

Project coordinates and Gaussian noise to zero center of gravity

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

Molecular EDMGeoDiff

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Center valid coordinates and supplied Gaussian noise independently within each molecule. Return the centered tensors, noisy coordinates at the supplied scale, and the Gaussian log density in the translation-free subspace.

Function signature

masked_centered_noise(
    xyz: torch.Tensor,
    noise: torch.Tensor,
    valid: torch.Tensor,
    sigma: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Center coordinates and noise separately

Take one molecule with xyz=[[0,0,0],[2,0,0]], noise=[[0,0,0],[2,2,2]], valid=[True,True] and sigma=[0.5]. The coordinate mean is [1,0,0] and noise mean is [1,1,1]. Thus centered_xyz=[[-1,0,0],[1,0,0]], centered_noise=[[-1,-1,-1],[1,1,1]], and noisy_xyz=[[-1.5,-0.5,-0.5],[1.5,0.5,0.5]]. There are $d=3(2-1)=3$ independent centered coordinates; the squared centered-noise sum is 6. The returned log density is

\log p=-3-3\left(\log(0.5)+\frac{\log(2\pi)}{2}\right)\approx-3.677374.

Parameters

Parameter	Shape and meaning
`xyz`torch.Tensor	$[B, N, 3]$ Molecular atom coordinates to center.
`noise`torch.Tensor	$[B, N, 3]$ Supplied standard-normal draws; center them separately from the coordinates.
`valid`torch.Tensor	$[B, N]$ `True` means real atom.
`sigma`torch.Tensor	$[B]$ Positive noise standard deviation for each molecule.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Inputs xyz $[B,N,3]$ , noise $[B,N,3]$ , valid $[B,N]$ , sigma $[B]$ have $B\ge 1$ , $N\ge 0$ and strictly positive finite sigma. True means a valid atom, with finite coordinates/noise; padded values may contain arbitrary NaN/Inf. noise is the supplied unit-Gaussian variate and must never be resampled.
For each molecule $n=\sum_{i} \mathrm{valid\_i}$ and $P(z)_i=\mathrm{valid}_i\left(z_i-\operatorname{mean}_{\mathrm{valid}}(z)\right)$ , with mean computed after replacing invalid values by zero and denominator $\operatorname{max}(n,1)$ .
Return $\mathrm{centered\_xyz}=P(\mathrm{xyz})$ , $\mathrm{centered\_noise}=P(\mathrm{noise})$ , and $\mathrm{noisy\_xyz}=\mathrm{centered\_xyz}+\sigma\cdot \mathrm{centered\_noise}$ , broadcasting sigma only over that molecule's atoms and xyz axis. All invalid outputs are zero.
This is unweighted arithmetic center of gravity, not a mass-weighted center. The density is measured in an orthonormal basis of the valid centered-coordinate subspace: $d=3\cdot \operatorname{max}(n-1,0)$ , and:
$\mathrm{log\_density}=-0.5\cdot \sum(\mathrm{centered\_noise}^{2})-d\cdot (\operatorname{log}(\sigma)+0.5\cdot \operatorname{log}(2\cdot \pi)).$ It is the Gaussian log density of the generated displacement $\sigma\cdot \mathrm{centered\_noise}$ , not a full $3 N$ -dimensional density.
For one valid atom, centered tensors and log_density are zero. Empty molecules, including $N=0$ , also return zeros and $\mathrm{log\_density}=0$ as an explicit bookkeeping convention. No keys are used, no $\varepsilon$ is added, and input tensors must not be mutated.
Return centered_xyz, centered_noise and noisy_xyz of shape $[B,N,3]$ , plus log_density $[B]$ , preserving dtype/device and noncontiguous layouts.

Why this matters

Removing the common translation projects both coordinates and noise into a subspace with three fewer dimensions per nonempty molecule.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def masked_centered_noise(
    xyz: torch.Tensor,  # (B, N, 3)
    noise: torch.Tensor,  # (B, N, 3)
    valid: torch.Tensor,  # (B, N)
    sigma: torch.Tensor,  # (B,)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    centered_xyz: (B, N, 3).
    centered_noise: (B, N, 3).
    noisy_xyz: (B, N, 3).
    log_density: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement masked_centered_noise")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.