Molecular graphs, equivariance and docking

Differentiate energy into forces and train through forces

HardRaw PyTorchAbout 30 min2 public examples plus private tests

SchNetNequIPMACETorchMD-Net

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

Your task

Use the supplied Gaussian pair model to compute atomic energies and one total energy per graph. Differentiate total energy to obtain atomic forces, then return all three results and the scalar mean-squared force loss, keeping the force computation differentiable.

Function signature

forces_from_energy(
    xyz: torch.Tensor,
    batch: torch.Tensor,
    edge_index: torch.Tensor,
    edge_valid: torch.Tensor,
    amplitudes: torch.Tensor,
    rates: torch.Tensor,
    atom_bias: torch.Tensor,
    target_forces: torch.Tensor,
    num_graphs: int,
) -> dict[str, torch.Tensor]

Worked example

An energy slope gives equal and opposite forces

Use xyz=[[0,0,0],[1,0,0]], batch=[0,0], edge_index=[[0],[1]], edge_valid=[True], amplitudes=[1], rates=[0.5], atom_bias=[0,0], target_forces=[[0,0,0],[0,0,0]] and num_graphs=1. Pair energy is $e^{-1/2}$ ; each atomic energy is half of it and total graph energy is $e^{-1/2}$ . Writing $\mathbf x_i$ for coordinate row $i$ ,

\mathbf F_0=2(0.5)e^{-1/2}(\mathbf x_0-\mathbf x_1).

Forces are approximately [[-0.606531,0,0],[0.606531,0,0]]. The force loss averages six squared components:

\frac{2e^{-1}}{6}=\frac{e^{-1}}{3}\approx0.122626.

Parameters

Parameter	Shape and meaning
`xyz`torch.Tensor	$[N, 3]$ Supplied differentiable coordinates; `requires_grad=True`.
`batch`torch.Tensor	$[N]$ Graph ID for every atom.
`edge_index`torch.Tensor	$[2, E]$ Each same-graph pair $i<j$ supplied once.
`edge_valid`torch.Tensor	$[E]$ `True` means active pair.
`amplitudes`torch.Tensor	$[K]$ Learnable Gaussian energy coefficients, signed allowed.
`rates`torch.Tensor	$[K]$ Strictly positive learnable inverse-length-squared `rates`.
`atom_bias`torch.Tensor	$[N]$ Learnable coordinate-independent atomic `energies`.
`target_forces`torch.Tensor	$[N, 3]$ Supplied supervision; treated as differentiable for judging.
`num_graphs`int	Required positive number of graph output slots. Graph IDs may be interleaved and some graph slots may contain no atoms.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Inputs xyz $[N,3]$ with $N\ge 1$ , batch $[N]$ , edge_index $[2,E]$ , edge_valid $[E]$ , amplitudes $[K]$ , positive rates $[K]$ , atom_bias $[N]$ , target_forces $[N,3]$ use $K\ge 1$ and $E\ge 0$ .
xyz is supplied with requires_grad=True by BioTorch, including forward-value tests; do not detach, clone into a new leaf, mutate its requires_grad flag, or run under no_grad. num_graphs is required $G\ge 1$ ; batch contains IDs in $[0,G)$ and may be interleaved, with empty graph IDs allowed.
Each pair column contains $i<j$ in the SAME graph, supplied once; endpoints are in range. True in edge_valid means the pair contributes. All floating inputs are finite.
For active pairs define $d=x_{i}-x_{j}$ , $\mathrm{r2}=\sum_{a} d_{a}^{2}$ and $\phi_{e}=\sum_{k} \mathrm{amplitudes\_k}\cdot \operatorname{exp}(-\mathrm{rates\_k}\cdot \mathrm{r2})$ . Inactive $\phi$ is zero.
Atomic energy is atom_bias_i plus one half of every incident pair energy, and graph energy is the sum of atomic energies with that batch ID.
Compute $\mathrm{forces}=-d(\sum_{g} \mathrm{energies\_g})/\mathrm{dxyz}$ using the exact permitted exception torch.autograd.grad with create_graph=True (and retain_graph=True if needed), so force supervision trains amplitudes/rates and geometry.
The empty-edge and all-masked cases still return zero forces connected to the coordinate computation; graph energies then sum only atom biases.
Return atomic_energies $[N]$ , energies $[G]$ , forces $[N,3]$ , and $\mathrm{force\_loss}=\sum((\mathrm{forces}-\mathrm{target\_forces})^{2})/(3\cdot N)$ , a scalar mean over every coordinate component. No energy target, force clipping, learned position bias, cutoff, radius-graph rebuilding, or sampling is used.
rates are directly supplied positive values, without $\varepsilon$ or a positivity transform. Preserve dtype/device and support noncontiguous coordinates. This is a tensor-only energy-model adaptation rather than an arbitrary energy callable.

Why this matters

Forces are negative energy gradients. Training from force error also differentiates through those gradients, requiring a higher-order graph.

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 forces_from_energy(
    xyz: torch.Tensor,  # (N, 3)
    batch: torch.Tensor,  # (N,)
    edge_index: torch.Tensor,  # (2, E)
    edge_valid: torch.Tensor,  # (E,)
    amplitudes: torch.Tensor,  # (K,)
    rates: torch.Tensor,  # (K,)
    atom_bias: torch.Tensor,  # (N,)
    target_forces: torch.Tensor,  # (N, 3)
    num_graphs: int,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    atomic_energies: (N,).
    energies: (G,).
    forces: (N, 3).
    force_loss: Scalar tensor, shape ()."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement forces_from_energy")

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.