inverse: explicit bond-angle diffusion with continuous triplet topology

.gitignore

@@ -130,3 +130,8 @@ dmypy.json
 
 # MacOS finder
 .DS_Store
+
+# task_runners output: splits, checkpoints, benchmark CSVs, SLURM logs
+runs/
+task_runners/logs/*
+!task_runners/logs/.gitkeep

README.md

@@ -324,6 +324,13 @@ at the trained 1000 steps, ~4.6 s at 200). Training your own model and
 reproducing the AtomBench benchmarks is covered in
 [alignn/inverse/README.md](alignn/inverse/README.md).
 
+That page also carries a full reference for an optional research extension,
+off by default: making bond angles an explicit **denoising channel** rather
+than only an input feature, with a line-graph topology that varies
+continuously as the coordinates denoise. See *Explicit bond-angle diffusion —
+reference* for the mathematics, the per-ablation breakdown, what each design
+choice takes from the literature, and the bibliography.
+
 <a name="performances"></a>
 ## Performances
 

alignn/inverse/__init__.py

@@ -14,6 +14,21 @@
 ``data``       dataset / collation from the AtomBench split JSONs.
 ``sample``     ancestral + Langevin-corrector sampling with classifier-free
                guidance on the conditioning property.
+``angles``     bond-angle denoising target and the smooth (DimeNet envelope /
+               ReaxFF product gate) triplet topology.
+``layers``     ALIGNN convolutions taking an optional per-edge or per-triplet
+               weight, so a message can fade out instead of being deleted.
+``ablations``  named configurations for the angular-diffusion ablation suite.
+``evaluate``   bond-angle distribution and relaxation-displacement metrics.
 """
 
-__all__ = ["data", "denoiser", "diffusion", "sample"]
+__all__ = [
+    "ablations",
+    "angles",
+    "data",
+    "denoiser",
+    "diffusion",
+    "evaluate",
+    "layers",
+    "sample",
+]

alignn/inverse/ablations.py

@@ -0,0 +1,119 @@
+"""Named configurations for the angular-diffusion ablation suite.
+
+The point of this extension is not to obtain one better model, it is to
+separate three claims that a single "it improved" number cannot:
+
+1. does an **explicit three-body denoising objective** help, over and above
+   angles being an ordinary ALIGNN input feature?
+2. does **continuously varying interaction topology** help in the noisy
+   regime, independently of any angular objective?
+3. does it matter that the angular channel is **coupled back** into the
+   coordinate/lattice pathway, rather than merely supervised alongside it?
+
+Every entry below is a keyword dict for
+:class:`alignn.inverse.denoiser.ALIGNNCSPDenoiser`.  They differ *only* in the
+switches under test — hidden size, depth, schedule, optimiser, data splits and
+seed policy come from the training script and must be held fixed across a
+comparison for it to mean anything.
+
+    A0  (A, F, L)                      current model, no angular objective
+    A1  (A, F, L, Theta)               + explicit angular denoising
+    A2  (A, F, L) + smooth topology    smooth graph, no angular objective
+    A3  (A, F, L, Theta) + smooth      the proposed model
+    A4  A3 with the coupling cut       control: auxiliary supervision only
+    A6  A3 with a Fourier angle basis  angular-representation ablation
+
+A5 in the design brief — hard kNN versus the smooth radius graph — is a
+*comparison*, not a fifth configuration: it is A1 against A3 (and A0 against
+A2), which is why there is no ``"A5"`` key.  :data:`COMPARISONS` spells out
+which pair of runs answers which question.
+
+Atom types are not diffused in this implementation.  The generator is
+conditioned on composition and solves crystal structure prediction, so the
+state is really ``(F, L)`` and, with these switches, ``(F, L, Theta)``; the
+``A`` in the names is kept only to match the design brief's notation.
+"""
+
+from __future__ import annotations
+
+from typing import Dict
+
+__all__ = ["ABLATIONS", "COMPARISONS", "ablation_config", "describe"]
+
+_SMOOTH = {
+    "topology": "radius",
+    "gate_pair_messages": True,
+}
+
+ABLATIONS: Dict[str, Dict] = {
+    "A0": {
+        "angle_diffusion": False,
+        "topology": "knn",
+        "gate_pair_messages": False,
+        "angle_feedback": True,
+    },
+    "A1": {
+        "angle_diffusion": True,
+        "topology": "knn",
+        "gate_pair_messages": False,
+        "angle_feedback": True,
+    },
+    "A2": {
+        "angle_diffusion": False,
+        "angle_feedback": True,
+        **_SMOOTH,
+    },
+    "A3": {
+        "angle_diffusion": True,
+        "angle_feedback": True,
+        **_SMOOTH,
+    },
+    "A4": {
+        "angle_diffusion": True,
+        "angle_feedback": False,
+        **_SMOOTH,
+    },
+    "A6": {
+        "angle_diffusion": True,
+        "angle_feedback": True,
+        "angle_basis": "fourier",
+        **_SMOOTH,
+    },
+}
+
+#: What each ablation is for, and which contrast it belongs to.
+DESCRIPTIONS: Dict[str, str] = {
+    "A0": "baseline: current ALIGNN 2.0 diffusion, angles as features only",
+    "A1": "explicit angular denoising, baseline kNN line-graph topology",
+    "A2": "smooth radius topology, no angular denoising objective",
+    "A3": "proposed: explicit angular denoising + smooth topology",
+    "A4": "control: angular objective with the angle->bond coupling removed",
+    "A6": "A3 with the Fourier angular basis instead of ALIGNN's cosine RBF",
+}
+
+COMPARISONS = {
+    "does explicit angular denoising help": ("A0", "A1"),
+    "does smooth topology alone help": ("A0", "A2"),
+    "do the two together help": ("A0", "A3"),
+    "is the coupling doing the work (not just auxiliary loss)": ("A4", "A3"),
+    "A5: hard kNN vs smooth radius, with angles on": ("A1", "A3"),
+    "A5: hard kNN vs smooth radius, with angles off": ("A0", "A2"),
+    "A6: does the angular basis matter": ("A3", "A6"),
+}
+
+
+def ablation_config(name: str) -> Dict:
+    """Denoiser keyword arguments for one named ablation."""
+    key = name.upper()
+    if key not in ABLATIONS:
+        raise KeyError(
+            f"unknown ablation {name!r}; available: "
+            f"{', '.join(sorted(ABLATIONS))} "
+            "(A5 is the A1-vs-A3 comparison, not a configuration)"
+        )
+    return dict(ABLATIONS[key])
+
+
+def describe(name: str) -> str:
+    """One-line description of a named ablation."""
+    return DESCRIPTIONS[name.upper()]

alignn/inverse/angles.py

@@ -0,0 +1,200 @@
+"""Angular geometry and continuously-weighted topology for ALIGNN-CSP.
+
+Two pieces of machinery live here, and they are deliberately independent so
+that the ablations in :mod:`alignn.inverse.ablations` can switch one on
+without the other.
+
+**Explicit angular denoising.**  ALIGNN already carries bond angles on its
+line graph, but only as an *input feature*.  Here they also become a
+*denoising target*: the network predicts, per triplet, the angular
+displacement that the forward process introduced.  The stochastic process and
+the loss are ported from FoldingDiff (Wu et al., Nat. Commun. 15, 1059, 2024,
+doi:10.1038/s41467-024-45051-2), which runs DDPM-style corruption and
+denoising directly on protein bond and dihedral angles with wrapped angular
+noise and a wrapped smooth-L1 objective.  Torsional Diffusion (Jing et al.,
+NeurIPS 2022) is the general statement that a diffusion process can be defined
+on an angular configuration space.
+
+One difference from FoldingDiff has to be stated plainly, because it is the
+main methodological caveat of this extension.  In a protein backbone the
+internal-coordinate list is *fixed*: residue i always has the same three
+angles, so a genuinely persistent state ``theta_t`` can be diffused
+independently of anything else.  In a crystal being denoised from noise the
+triplet set is not fixed — it is a function of the coordinates, and it changes
+as they move.  There is therefore no persistent ``theta_t`` to diffuse.  What
+is implemented instead is the closest well-defined thing: the angular
+denoising *target* is computed on the triplet representation that exists at
+the current step,
+
+    delta_ijk = wrap(theta_ijk(f_t, L_t) - theta_ijk(f_0, L_0)),
+
+with both angles evaluated on the *same* periodic-image identity
+``(i, j, k, n_ji, n_jk)`` so that the difference measures the corruption of
+one specific triplet rather than a change of neighbour.  Angles are still an
+explicit denoising channel with their own head and their own loss; they are
+not an independently-noised variable.  Section 3 of the design brief asks for
+exactly this fallback, and asks that the distinction be documented rather than
+papered over with an invented process.
+
+**Continuously-weighted topology.**  During reverse diffusion a hard
+neighbour-rank criterion for "does this triplet exist" is unjustified: at
+large ``t`` the coordinates are close to uniform, so neighbour ranks swap
+constantly and the line graph jumps discontinuously.  Instead every pair
+carries a smooth relevance
+
+    s_ij = u(r_ij ; r_c),
+
+where ``u`` is the polynomial cutoff envelope introduced by DimeNet
+(Gasteiger, Gross & Gunnemann, ICLR 2020, arXiv:2003.03123), whose value and
+first two derivatives vanish at ``r_c``.  That envelope is already implemented
+in this repository as
+:class:`alignn.models.alignn_atomwise_pure_smooth.CutoffPolynomial`, so it is
+reused rather than re-derived.  A triplet inherits the product of its two
+constituent relevances,
+
+    s_ijk = s_ji * s_jk,
+
+which is the ReaxFF treatment of valence angles (van Duin et al.,
+J. Phys. Chem. A 105, 9396, 2001, doi:10.1021/jp004368u): bond orders vary
+continuously with distance and an angular term switches off smoothly as
+either of its bonds dissociates.  Because ``s`` is exactly zero at and beyond
+``r_c``, restricting the sparse line graph to pairs inside ``r_c`` removes
+only terms that were already contributing nothing — a triplet can enter or
+leave the computational graph without any finite jump in the messages.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+
+from alignn.models.alignn_atomwise_pure_smooth import CutoffPolynomial
+from alignn.torch_graph_builder import torch_bond_cosines
+
+__all__ = [
+    "CutoffPolynomial",
+    "bond_angle",
+    "wrap_angle",
+    "pair_relevance",
+    "triplet_relevance",
+    "edge_vectors",
+    "angular_denoising_target",
+    "angle_denoising_loss",
+    "TWO_PI",
+]
+
+TWO_PI = 2.0 * math.pi
+
+# acos is not differentiable at +-1; the clamp keeps the gradient finite for
+# collinear triplets, which are common in a crystal (i -> j -> i back-tracking
+# triplets are cos = -1 exactly).
+_COS_EPS = 1.0e-7
+
+
+def bond_angle(r_ij: torch.Tensor, r_jk: torch.Tensor) -> torch.Tensor:
+    """Bond angle at the shared atom ``j`` of a triplet, in radians.
+
+    Uses ALIGNN's own cosine convention (:func:`torch_bond_cosines`), so the
+    angle is the interior angle at ``j`` and lies in ``[0, pi]``.
+    """
+    cos = torch_bond_cosines(r_ij, r_jk)
+    return torch.acos(cos.clamp(-1.0 + _COS_EPS, 1.0 - _COS_EPS))
+
+
+def wrap_angle(x: torch.Tensor) -> torch.Tensor:
+    """Wrap an angular difference into ``[-pi, pi)``.
+
+    FoldingDiff's forward process and loss are both defined modulo ``2 pi``.
+    Bond angles themselves live in ``[0, pi]``, so a difference of two of them
+    is already inside ``[-pi, pi]`` and this is a no-op up to the boundary
+    case; it is applied anyway so that the objective is the wrapped one by
+    construction rather than by an argument about ranges, and so that the same
+    helper can serve a periodic angular variable if one is ever added.
+    """
+    return torch.remainder(x + math.pi, TWO_PI) - math.pi
+
+
+def pair_relevance(
+    dist: torch.Tensor, envelope: CutoffPolynomial
+) -> torch.Tensor:
+    """Smooth per-pair relevance ``s_ij = u(r_ij; r_c)`` in ``[0, 1]``.
+
+    ``u`` is the DimeNet envelope: ``u(0) = 1`` and ``u`` together with its
+    first two derivatives vanishes at ``r_c``.
+    """
+    return envelope(dist)
+
+
+def triplet_relevance(
+    s_edge: torch.Tensor, lg_src: torch.Tensor, lg_dst: torch.Tensor
+) -> torch.Tensor:
+    """ReaxFF-style product gate ``s_ijk = s_ji * s_jk`` for each triplet."""
+    return s_edge[lg_src] * s_edge[lg_dst]
+
+
+def edge_vectors(
+    frac: torch.Tensor,
+    lattice: torch.Tensor,
+    src: torch.Tensor,
+    dst: torch.Tensor,
+    edge_graph_id: torch.Tensor,
+    image: torch.Tensor,
+) -> torch.Tensor:
+    """Cartesian edge vectors for a *given* set of periodic images.
+
+    ``image`` is the integer cell offset ``n`` that the minimum-image search
+    settled on at the noised geometry, so that ``delta_f = f[dst] - f[src] +
+    n``.  Re-using the same ``n`` on a different (clean) structure is what
+    makes the angular target a corruption of one fixed triplet identity rather
+    than a comparison between two different neighbours.
+    """
+    df = frac[dst] - frac[src] + image
+    return torch.einsum("ei,eij->ej", df, lattice[edge_graph_id])
+
+
+def angular_denoising_target(
+    angle_t: torch.Tensor,
+    frac0: torch.Tensor,
+    lattice0: torch.Tensor,
+    src: torch.Tensor,
+    dst: torch.Tensor,
+    edge_graph_id: torch.Tensor,
+    image: torch.Tensor,
+    lg_src: torch.Tensor,
+    lg_dst: torch.Tensor,
+) -> torch.Tensor:
+    """Angular displacement the forward process applied to each triplet.
+
+    Returns ``wrap(theta_t - theta_0)``, the angular analogue of the noise
+    ``eps`` that FoldingDiff's network predicts.
+    """
+    r0 = edge_vectors(frac0, lattice0, src, dst, edge_graph_id, image)
+    theta0 = bond_angle(r0[lg_src], r0[lg_dst])
+    return wrap_angle(angle_t - theta0)
+
+
+def angle_denoising_loss(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
+    beta: float = 0.1 * math.pi,
+) -> torch.Tensor:
+    """Relevance-weighted wrapped smooth-L1 loss on the angular channel.
+
+    The functional form — smooth L1 of the *wrapped* residual, with
+    ``beta = 0.1 pi`` — is FoldingDiff's angular objective.  The weighting by
+    ``s_ijk`` is what makes the loss continuous when a triplet enters or
+    leaves the sparse line graph: a triplet at the cutoff has zero weight, so
+    it contributes nothing on either side of the boundary.
+    """
+    if pred.numel() == 0:
+        return pred.new_zeros(())
+    d = wrap_angle(pred - target)
+    per_triplet = torch.nn.functional.smooth_l1_loss(
+        d, torch.zeros_like(d), beta=beta, reduction="none"
+    )
+    if weight is None:
+        return per_triplet.mean()
+    return (weight * per_triplet).sum() / weight.sum().clamp_min(1e-8)