From ed6d64c54a9b0f300c2f5efcef1993d4d1f026e8 Mon Sep 17 00:00:00 2001 From: crhysc Date: Sat, 29 Aug 2026 17:08:10 -0400 Subject: [PATCH 1/4] inverse: explicit bond-angle diffusion with continuous triplet topology Makes ALIGNN's three-body geometry an optional generative channel rather than only an input feature, and lets the line-graph topology vary continuously while the coordinates denoise. Everything is off by default: ALIGNNCSPDenoiser() with no arguments is the model that was there before, and a test asserts that enabling the angle head leaves eps_frac and eps_lattice bit-identical. Angular channel (FoldingDiff, doi:10.1038/s41467-024-45051-2). One extra head on the shared backbone's line-graph feature predicts the angular displacement the forward process introduced, trained with a wrapped smooth-L1 loss (beta = 0.1 pi). The angular latent reaches the coordinate and lattice heads through ALIGNN's ordinary angles -> bonds -> atoms path, so it is coupled rather than merely supervised. One deviation, documented rather than papered over: a crystal has no persistent theta_t to diffuse, because the triplet set is a function of the coordinates. The target is therefore computed on the triplet representation that exists at the current step, with both angles evaluated on the same periodic-image identity so it measures the corruption of one triplet rather than a change of neighbour. No new SDE or schedule was invented. Continuous topology (DimeNet arXiv:2003.03123, ReaxFF doi:10.1021/jp004368u). topology="radius" replaces the hard kNN rule with a radius candidate set, each pair weighted by the DimeNet envelope and each triplet by the product of its two bonds' weights. The envelope is reused from CutoffPolynomial, already in this repository. The weight multiplies the gate before both sums of the edge-gated average, which is the only placement making a zero-weight edge exactly equal to a deleted one. The graph is rebuilt from the current geometry every forward pass, so topology follows r_ij(t) with no annealing schedule, graph-temperature term or learned bondness network. Ablations are the deliverable, not the model: alignn/inverse/ablations.py defines A0-A4 and A6, run_angle_ablation.sh runs every arm on matched splits, budget and seeds, and the table in the README says which contrast answers which question (A5 is the A1-vs-A3 comparison, not a seventh configuration). Evaluation adds two mechanism metrics fixed before any run: generated vs held-out bond-angle distributions (FoldingDiff's diagnostic) and relaxation displacement to the nearest ALIGNN-FF minimum (MatterGen's proximity evaluation, doi:10.1038/s41586-025-08628-5). Not implemented, deliberately: DimeNet's spherical Fourier-Bessel basis (A6 substitutes the Fourier basis already in the repo), and every mechanism the brief rules out. Nothing has been trained. 32 new tests cover angular periodicity, the envelope and its first two derivatives at r_c, smooth pair and triplet fade-out, absence of a jump when a triplet enters or leaves the sparse graph, periodic-coordinate invariance, baseline equivalence, and that A4 provably cannot see the angular features. Claude-Session: https://claude.ai/code/session_01ULEhMd27yfBmKf5HXqmvFE --- README.md | 5 + alignn/inverse/README.md | 172 +++++++ alignn/inverse/__init__.py | 17 +- alignn/inverse/ablations.py | 119 +++++ alignn/inverse/angles.py | 217 ++++++++ alignn/inverse/denoiser.py | 288 ++++++++++- alignn/inverse/evaluate.py | 211 ++++++++ alignn/inverse/layers.py | 127 +++++ alignn/inverse/sample.py | 36 +- alignn/inverse/train_csp.py | 131 ++++- alignn/tests/test_inverse_angle_diffusion.py | 490 +++++++++++++++++++ scripts/atombench/angle_eval.py | 149 ++++++ scripts/atombench/run_angle_ablation.sh | 63 +++ 13 files changed, 1986 insertions(+), 39 deletions(-) create mode 100644 alignn/inverse/ablations.py create mode 100644 alignn/inverse/angles.py create mode 100644 alignn/inverse/evaluate.py create mode 100644 alignn/inverse/layers.py create mode 100644 alignn/tests/test_inverse_angle_diffusion.py create mode 100755 scripts/atombench/angle_eval.py create mode 100755 scripts/atombench/run_angle_ablation.sh diff --git a/README.md b/README.md index 11b9e97..204d5dd 100644 --- a/README.md +++ b/README.md @@ -324,6 +324,11 @@ 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 documents 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, plus the ablation suite that tells the two apart. + ## Performances diff --git a/alignn/inverse/README.md b/alignn/inverse/README.md index d68a65c..30123cb 100644 --- a/alignn/inverse/README.md +++ b/alignn/inverse/README.md @@ -168,6 +168,178 @@ because it uses the spawn start method the caller must then be under an > The step count trades speed for fidelity. The models were trained at 1000 > steps; whether 200 or 50 preserves benchmark quality has not been measured. +## Explicit bond-angle diffusion (optional) + +Everything above treats bond angles as an *input feature*: ALIGNN computes +them, the line graph propagates them, and they help the network predict where +the atoms should go. They are never themselves something the model denoises. +This section adds the option of making them one, and of letting the triplet +topology change continuously while it does. All of it is off by default — +`ALIGNNCSPDenoiser()` with no arguments is byte-for-byte the model described +above, and a test asserts that switching the angle head on leaves `eps_frac` +and `eps_lattice` unchanged. + +The hypothesis under test is narrow: **does explicit three-body geometric +denoising improve crystal generation when embedded in ALIGNN's +atom–bond–line-graph hierarchy?** It is a hypothesis, not a claim — nothing +here has been trained yet. + +### The angular channel + +The model gains one output head, reading the line-graph feature `z` of the +same backbone that already produces the coordinate score and the lattice +noise. Per triplet it predicts the angular displacement the forward process +introduced, + +``` +delta_ijk = wrap( theta_ijk(f_t, L_t) - theta_ijk(f_0, L_0) ) +``` + +trained with a wrapped smooth-L1 loss (`beta = 0.1 pi`). Both the target and +the loss are FoldingDiff's, which runs DDPM-style corruption and denoising +directly on protein bond and dihedral angles and shows that a model with an +explicit angular channel reproduces natural angular distributions. + +**One honest deviation.** FoldingDiff can diffuse a genuinely persistent +`theta_t` because a protein backbone's internal-coordinate list is fixed: +residue *i* always has the same three angles. A crystal being denoised from +noise has no such list — the triplet set is a function of the coordinates and +changes as they move, so there is no independent angular state to noise. What +is implemented is the closest well-defined thing: the angular *target* is +computed on the triplet representation that exists at the current step, with +both angles evaluated on the same periodic-image identity +`(i, j, k, n_ji, n_jk)` so the difference measures the corruption of one +triplet rather than a change of neighbour. Angles are an explicit denoising +objective with their own head and their own loss; they are not an +independently-noised variable. No new SDE, angle manifold or schedule was +invented to paper over the difference. + +The angular representation itself is untouched — ALIGNN's Gaussian RBF over +`cos(theta)` — so that the experiment measures the *objective*, not a change +of basis. + +### Continuous topology + +At large `t` the coordinates are close to uniform, and a hard neighbour-rank +rule for "does this triplet exist" is unjustified: ranks swap constantly and +the line graph jumps. `topology="radius"` replaces the kNN rule with a radius +candidate set in which every pair carries a smooth relevance + +``` +s_ij = u(r_ij ; r_c) DimeNet's polynomial envelope: u, u' and u'' + all vanish at r_c +s_ijk = s_ji * s_jk ReaxFF's treatment of valence angles: an angle + switches off smoothly as either bond dissociates +``` + +`u` is not re-derived — it is `CutoffPolynomial`, already in this repository +for the smooth property model. The weight multiplies the gate *before both* +sums of the edge-gated convolution's normalised average, which is the only +placement that makes a zero-weight edge exactly equivalent to a deleted one; +scaling the numerator alone would renormalise the survivors and would not be +continuous. Because `s` is exactly zero at and beyond `r_c`, restricting the +sparse graph to pairs inside `r_c` drops only terms already contributing +nothing, so a triplet can enter or leave without any jump — which is what the +tests check by sweeping an atom through the cutoff and watching the output. + +The graph is rebuilt from the current coordinates and lattice on every +forward pass, so the topology follows the geometry through reverse diffusion +on its own. There is no connectivity annealing schedule, no graph-temperature +term and no learned bondness network; `r_ij(t)` evolving is the whole +mechanism. + +`radius_cutoff` defaults to 5 Å — between this repository's own three-body +cutoff (3.5 Å) and DimeNet's molecular cutoff (5 Å), and close to the radius +the baseline's 12 nearest neighbours actually span, which keeps the kNN +comparison fair. + +### Ablations + +Ablations are the point, not an afterthought: a single "it improved" number +cannot separate the three claims. `alignn/inverse/ablations.py` holds the +configurations, `scripts/atombench/run_angle_ablation.sh` runs them with the +same split, optimiser, epoch budget and seed list across every arm. + +| | angular objective | topology | angle → bond coupling | +|---|---|---|---| +| **A0** baseline | – | kNN | – | +| **A1** | yes | kNN | yes | +| **A2** | – | smooth radius | – | +| **A3** proposed | yes | smooth radius | yes | +| **A4** control | yes | smooth radius | **cut** | +| **A6** | yes | smooth radius | yes, Fourier angle basis | + +Which contrast answers what: + +- **A0 vs A1** — does explicit angular denoising help on its own? +- **A0 vs A2** — or is the gain just from a better-behaved noisy graph? +- **A4 vs A3** — is the benefit genuine coupling, or would auxiliary + supervision on a shared trunk do as well? A4 keeps the angular features + evolving and supervised but zeroes their contribution to the bond + representation, so its coordinate/lattice pathway provably cannot see them + (a test perturbs the angle embedding and asserts A4's output does not move + while A3's does). +- **A1 vs A3** — the design brief's A5: hard kNN against the smooth radius + graph. It is a comparison between existing arms, not a seventh + configuration. +- **A3 vs A6** — does the angular basis matter? Run this *last*; mixing a + basis change into the primary experiment would make attribution impossible. + +```bash +bash scripts/atombench/run_angle_ablation.sh runs/data_jarvis runs/ablation 0 1 2 +``` + +Given the spread already recorded above — match rate across fifteen +independently trained models spanned 0.437–0.524 — one seed per arm cannot +settle any of these. The runner defaults to three. + +### Evaluation + +Every existing AtomBench metric is preserved. Two mechanism metrics are added, +defined before the runs so that a favourable one cannot be picked afterwards: + +- **bond-angle distributions**, generated against held-out real structures + (KL, Jensen-Shannon, and a 1-D Wasserstein distance reported in degrees) — + FoldingDiff's own diagnostic; +- **relaxation displacement**, how far a sample moves to reach its nearest + ALIGNN-FF local minimum, with the volume change and energy drop — the + proximity-to-local-minimum evaluation MatterGen uses. If explicit angular + denoising produces locally coherent geometry, its samples should need less + repair. + +```bash +python scripts/atombench/angle_eval.py runs/ablation/A3_s0/bench.csv --relax +``` + +### What is deliberately not here + +DimeNet's joint spherical Fourier–Bessel distance–angle basis is **not** +implemented; A6 substitutes the Fourier basis on theta this repository +already ships. The full SBF needs spherical Bessel zeros and a new +dependency, and the design brief sequences that ablation last, so it is +deferred rather than half-done. Also absent, deliberately: cross-attention, +separate networks per variable, a learned bondness classifier, a +time-dependent graph-temperature schedule, Jacobian angle forces, and any +loss whose role cannot be traced to one of the papers below. + +Note that atom types are *not* diffused — the generator is conditioned on +composition — so the state is really `(F, L)` and, with this extension, +`(F, L, Theta)`. The `A` in the ablation names follows the design brief's +notation. + +### References + +| | | +|---|---| +| ALIGNN — atom graph + line graph, angles update bonds update atoms | Choudhary & DeCost, *npj Comput. Mater.* **7**, 185 (2021), [10.1038/s41524-021-00650-1](https://doi.org/10.1038/s41524-021-00650-1) | +| FoldingDiff — diffusion directly on bond angles; wrapped noise, wrapped smooth-L1, angle-distribution evaluation | Wu *et al.*, *Nat. Commun.* **15**, 1059 (2024), [10.1038/s41467-024-45051-2](https://doi.org/10.1038/s41467-024-45051-2) | +| Torsional Diffusion — diffusion on angular configuration spaces | Jing *et al.*, NeurIPS 2022 | +| ReaxFF — continuous distance-dependent bond order; angle terms vanish as either bond dissociates | van Duin *et al.*, *J. Phys. Chem. A* **105**, 9396 (2001), [10.1021/jp004368u](https://doi.org/10.1021/jp004368u) | +| DimeNet — cutoff envelope whose value and first two derivatives vanish at `r_c` | Gasteiger, Groß & Günnemann, ICLR 2020, [arXiv:2003.03123](https://arxiv.org/abs/2003.03123) | +| MatterGen — one score network denoising crystal variables jointly; proximity-to-relaxed evaluation | Zeni *et al.*, *Nature* **639**, 624 (2025), [10.1038/s41586-025-08628-5](https://doi.org/10.1038/s41586-025-08628-5) | +| DiffCSP — joint equivariant diffusion of lattice and fractional coordinates | Jiao *et al.*, NeurIPS 2023 | +| CrystalDiT — the warning against unnecessary multi-stream architectural complexity | Yi *et al.*, AAAI 2026, [10.1609/aaai.v40i2.37121](https://doi.org/10.1609/aaai.v40i2.37121) | + ## Relax and rank `relax_rank.py` refines candidates with the pretrained ALIGNN force field. diff --git a/alignn/inverse/__init__.py b/alignn/inverse/__init__.py index cf9f09b..9d18fc3 100644 --- a/alignn/inverse/__init__.py +++ b/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", +] diff --git a/alignn/inverse/ablations.py b/alignn/inverse/ablations.py new file mode 100644 index 0000000..ded8eec --- /dev/null +++ b/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()] diff --git a/alignn/inverse/angles.py b/alignn/inverse/angles.py new file mode 100644 index 0000000..87f6942 --- /dev/null +++ b/alignn/inverse/angles.py @@ -0,0 +1,217 @@ +"""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, Tuple + +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) + + +def angle_histogram( + angles_deg, bins: int = 180, lo: float = 0.0, hi: float = 180.0 +) -> Tuple: + """Normalised histogram of bond angles in degrees. + + Kept here so the evaluation code and the tests share one definition. + """ + import numpy as np + + counts, edges = np.histogram( + np.asarray(angles_deg, dtype=float), bins=bins, range=(lo, hi) + ) + total = counts.sum() + density = counts / total if total else counts.astype(float) + return density, edges diff --git a/alignn/inverse/denoiser.py b/alignn/inverse/denoiser.py index bb8d759..2ef379e 100644 --- a/alignn/inverse/denoiser.py +++ b/alignn/inverse/denoiser.py @@ -18,6 +18,28 @@ The distinguishing ingredient relative to CSPNet / CDVAE / FlowMM denoisers is the ALIGNN line graph: bond *angles* are propagated alongside bond lengths, which is the three-body information that pins down coordination geometry. + +Two optional extensions turn that three-body information from a *feature* into +a *generative channel*. Both default to off, so the original model is +recovered exactly by the default configuration. + +``angle_diffusion`` + Adds a per-triplet head predicting the angular displacement the forward + process introduced, trained with FoldingDiff's wrapped smooth-L1 + objective. The head reads the line-graph feature ``z`` of the shared + backbone, and ``z`` reaches the coordinate and lattice heads through + ALIGNN's ordinary ``angles -> bonds -> atoms`` path, so the angular + channel is coupled rather than merely supervised. ``angle_feedback=False`` + cuts that coupling for the control ablation. + +``topology="radius"`` + Replaces the hard k-nearest-neighbour rule that decides which bonds may + form triplets with a radius candidate set plus a DimeNet cutoff envelope + and a ReaxFF-style product gate, so the line graph changes continuously as + the coordinates denoise instead of jumping when two neighbours swap rank. + +See :mod:`alignn.inverse.angles` for the literature these follow and for the +one place where this deviates from FoldingDiff. """ from __future__ import annotations @@ -28,16 +50,29 @@ import torch from torch import nn -from alignn.models.alignn_atomwise_pure import ( - ALIGNNConvPure, - EdgeGatedGraphConvPure, - scatter_mean, - scatter_sum, +from alignn.models.alignn_atomwise_pure import scatter_mean, scatter_sum +from alignn.models.alignn_atomwise_pure_smooth import ( + CutoffPolynomial, + FourierAngular, ) from alignn.models.utils import MLPLayer, RBFExpansion from alignn.torch_graph_builder import _line_graph_edges, torch_bond_cosines +from alignn.inverse.angles import bond_angle, triplet_relevance from alignn.inverse.diffusion import wrap_diff +from alignn.inverse.layers import ( + WeightedALIGNNConv, + WeightedEdgeGatedGraphConv, +) + +#: Line-graph topologies. ``knn`` is the original hard neighbour-rank rule; +#: ``radius`` is the smooth construction of section 4 of the design brief. +TOPOLOGIES = ("knn", "radius") + +#: Angular input bases. ``cosine_rbf`` is ALIGNN's own; ``fourier`` is the +#: learnable Fourier basis on theta already shipped in this repository, and is +#: reserved for the A6 basis ablation. +ANGLE_BASES = ("cosine_rbf", "fourier") def sinusoidal_embedding(x: torch.Tensor, dim: int, max_period: float = 1e4): @@ -112,6 +147,29 @@ def _image_offsets(device, dtype): return torch.cartesian_prod(r, r, r) # (27, 3) +def _angle_basis_layers(kind, triplet_bins, embedding_features, hidden): + """Layers expanding a bond-angle cosine to a hidden-size feature. + + ``cosine_rbf`` is ALIGNN's own representation and is left byte-for-byte + as it was; ``fourier`` swaps in the learnable Fourier basis on theta that + this repository already carries, and exists only for the A6 ablation. + DimeNet's joint spherical Fourier-Bessel distance-angle basis is *not* + implemented here — see the ablation notes. + """ + if kind == "fourier": + order = max(1, (triplet_bins - 1) // 2) + basis = FourierAngular(order=order) + n_in = basis.out_features + else: + basis = RBFExpansion(vmin=-1.0, vmax=1.0, bins=triplet_bins) + n_in = triplet_bins + return [ + basis, + MLPLayer(n_in, embedding_features), + MLPLayer(embedding_features, hidden), + ] + + class ALIGNNCSPDenoiser(nn.Module): """Predict (coordinate score, lattice noise) for a noised crystal.""" @@ -129,12 +187,84 @@ def __init__( num_species: int = 120, num_steps: int = 1000, score_channels: int = 32, + angle_diffusion: bool = False, + angle_feedback: bool = True, + topology: str = "knn", + radius_cutoff: float = 5.0, + envelope_exponent: int = 5, + gate_pair_messages: bool = False, + angle_basis: str = "cosine_rbf", ): + """Build the denoiser. + + Parameters beyond the original set, all defaulting to the original + behaviour: + + angle_diffusion + Emit an angular denoising prediction per triplet. Requires + ``alignn_layers > 0``, since the line graph is what carries + angles. + angle_feedback + Whether the angular features are allowed to reach the bond (and + hence atom, coordinate and lattice) representations. ``False`` is + ablation A4: angular supervision on a shared trunk with the + architectural coupling removed. + topology + ``"knn"`` keeps the original rule — a bond may join a triplet if + it is among the ``knn`` shortest bonds at its destination atom. + ``"radius"`` replaces it with every bond shorter than + ``radius_cutoff``, each weighted by the DimeNet envelope, with + triplets weighted by the product of their two bonds' weights. + radius_cutoff, envelope_exponent + Cutoff radius and polynomial order of that envelope. The default + 5 A sits between this repository's own three-body cutoff (3.5 A) + and DimeNet's molecular cutoff (5 A), and is close to the radius + the baseline's 12 nearest neighbours actually span in a crystal, + which keeps the A1-vs-A3 comparison fair. + gate_pair_messages + Also weight the *pair* channel — the atom-graph messages and the + per-edge terms of the coordinate score — by ``s_ij``. The pair + graph is dense rather than neighbour-ranked, so nothing is ever + inserted or deleted there and this is not needed for continuity; + it is the fuller reading of "smoothly vanishing pair + interactions" and is switched on by the smooth-topology + ablations. + angle_basis + ``"cosine_rbf"`` is ALIGNN's own angular representation and is + what every primary experiment uses. ``"fourier"`` is reserved + for the A6 basis ablation. + """ super().__init__() + if topology not in TOPOLOGIES: + raise ValueError( + f"topology must be one of {TOPOLOGIES}, got {topology!r}" + ) + if angle_basis not in ANGLE_BASES: + raise ValueError( + f"angle_basis must be one of {ANGLE_BASES}, " + f"got {angle_basis!r}" + ) + if angle_diffusion and alignn_layers <= 0: + raise ValueError( + "angle_diffusion needs alignn_layers > 0: the angular " + "channel lives on the line graph, which is not built when " + "there are no ALIGNN layers" + ) + if gate_pair_messages and topology != "radius": + raise ValueError( + "gate_pair_messages requires topology='radius'; the gate is " + "the radius envelope" + ) self.hidden_features = hidden_features self.fourier_k = fourier_k self.knn = knn self.num_steps = num_steps + self.angle_diffusion = angle_diffusion + self.angle_feedback = angle_feedback + self.topology = topology + self.radius_cutoff = radius_cutoff + self.gate_pair_messages = gate_pair_messages + self.angle_basis = angle_basis self.species_embedding = nn.Embedding(num_species, hidden_features) @@ -166,23 +296,36 @@ def __init__( self.use_line_graph = alignn_layers > 0 self.angle_embedding = ( nn.Sequential( - RBFExpansion(vmin=-1.0, vmax=1.0, bins=triplet_bins), - MLPLayer(triplet_bins, embedding_features), - MLPLayer(embedding_features, hidden_features), + *_angle_basis_layers( + angle_basis, + triplet_bins, + embedding_features, + hidden_features, + ) ) if self.use_line_graph else None ) + # DimeNet's polynomial cutoff envelope, already implemented in this + # repository for the smooth property model; u, u' and u'' all vanish + # at the cutoff. + self.envelope = ( + CutoffPolynomial( + cutoff=radius_cutoff, coeff=float(envelope_exponent) + ) + if topology == "radius" + else None + ) self.alignn_layers = nn.ModuleList( [ - ALIGNNConvPure(hidden_features, hidden_features) + WeightedALIGNNConv(hidden_features, hidden_features) for _ in range(alignn_layers) ] ) self.gcn_layers = nn.ModuleList( [ - EdgeGatedGraphConvPure(hidden_features, hidden_features) + WeightedEdgeGatedGraphConv(hidden_features, hidden_features) for _ in range(gcn_layers) ] ) @@ -210,11 +353,28 @@ def __init__( nn.SiLU(), nn.Linear(hidden_features, 6), ) + # Angular denoising head. Reads the line-graph feature of the shared + # backbone, so nothing about it is a second network: it is one more + # output head on the representation that already denoises coordinates + # and lattice. + self.angle_head = ( + nn.Sequential( + nn.Linear(hidden_features, hidden_features), + nn.SiLU(), + nn.Linear(hidden_features, 1), + ) + if angle_diffusion + else None + ) + # Start from a near-zero prediction: diffusion training is much better # behaved when the model does not begin by shouting. nn.init.zeros_(self.score_combine.weight) nn.init.zeros_(self.lattice_head[-1].weight) nn.init.zeros_(self.lattice_head[-1].bias) + if self.angle_head is not None: + nn.init.zeros_(self.angle_head[-1].weight) + nn.init.zeros_(self.angle_head[-1].bias) # ── geometry ───────────────────────────────────────────────────────── def _edge_geometry( @@ -224,9 +384,23 @@ def _edge_geometry( src: torch.Tensor, dst: torch.Tensor, edge_graph_id: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Return wrapped Δf, min-image Δf, min-image Cartesian vec, dist.""" - df = wrap_diff(frac[dst] - frac[src]) + ) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + """Return wrapped Δf, min-image Δf, Cartesian vec, dist, image. + + ``image`` is the integer cell offset ``n`` this resolution settled on, + i.e. the one for which ``Δf = f[dst] - f[src] + n``. The angular + target needs it: re-applying the same ``n`` to the *clean* structure + is what makes the target the corruption of one fixed triplet identity + instead of a comparison between two different neighbours. + """ + raw = frac[dst] - frac[src] + df = wrap_diff(raw) offsets = _image_offsets(frac.device, frac.dtype) # (27, 3) cand = df.unsqueeze(1) + offsets.unsqueeze(0) # (E, 27, 3) lat_e = lattice[edge_graph_id] # (E, 3, 3) @@ -244,7 +418,15 @@ def _edge_geometry( # The fractional difference for the *same* image, which is what the # coordinate score head combines. df_min = cand.gather(1, idx).squeeze(1) # (E, 3) - return df, df_min, r, d_all.gather(1, best.view(-1, 1)).squeeze(1) + # cand = (raw - round(raw)) + offset, so the total integer shift is: + image = offsets[best] - torch.round(raw) # (E, 3) + return ( + df, + df_min, + r, + d_all.gather(1, best.view(-1, 1)).squeeze(1), + image, + ) def _fourier(self, df: torch.Tensor) -> torch.Tensor: """Fourier features of a fractional difference (periodic, signed).""" @@ -274,16 +456,29 @@ def forward( ``cond_embedding`` is a ``(B, hidden_features)`` vector produced by a :class:`~alignn.inverse.conditioners.MultiModalConditioner` — the denoiser is deliberately agnostic to which modalities went into it. + + With ``angle_diffusion`` on, the returned dict carries an extra + ``"angle"`` entry holding the per-triplet prediction and everything + the loss needs to build its target: the angles at the *noised* + geometry, the triplet relevance weights, and the edge/triplet indices + together with the periodic images the geometry was resolved against. """ num_nodes = frac.shape[0] if pair_index is None: pair_index = dense_pair_index(natoms) src, dst, edge_graph_id = pair_index - df, df_min, r, dist = self._edge_geometry( + df, df_min, r, dist, image = self._edge_geometry( frac, lattice, src, dst, edge_graph_id ) + # Smooth pair relevance s_ij = u(r_ij; r_c). Recomputed from the + # current coordinates and lattice on every call, which is what makes + # the topology follow the geometry through reverse diffusion rather + # than being fixed up front. + s_edge = None if self.envelope is None else self.envelope(dist) + edge_w = s_edge if self.gate_pair_messages else None + # Node features: species + timestep + conditioning + lattice state. h = self.species_embedding(atomic_numbers) t_emb = self.time_mlp( @@ -305,24 +500,74 @@ def forward( y = self.edge_embedding( torch.cat([self.rbf(dist), self._fourier(df)], dim=-1) ) + angle_out: Optional[Dict[str, torch.Tensor]] = None if self.use_line_graph: - allowed = _knn_mask(dist, dst, num_nodes, self.knn) + if s_edge is None: + # Original topology: a bond may join a triplet if it is among + # the k shortest at its destination atom. + allowed = _knn_mask(dist, dst, num_nodes, self.knn) + else: + # Radius candidate set. s_ij is exactly zero at and beyond + # r_c, so dropping those bonds removes only terms that + # already contributed nothing. + allowed = s_edge > 0.0 lg_src, lg_dst = _line_graph_edges( src, dst, num_nodes, allowed=allowed ) - z = self.angle_embedding(torch_bond_cosines(r[lg_src], r[lg_dst])) + # ReaxFF-style product gate: an angle fades out when either of + # its two bonds does. + tri_w = ( + None + if s_edge is None + else triplet_relevance(s_edge, lg_src, lg_dst) + ) + cos_theta = torch_bond_cosines(r[lg_src], r[lg_dst]) + z = self.angle_embedding(cos_theta) + # A4: keep the angular features evolving and supervised, but stop + # them from reaching the bond representation. + conv_tri_w = tri_w + if not self.angle_feedback: + conv_tri_w = torch.zeros_like(cos_theta) for layer in self.alignn_layers: h, y, z = layer.forward_tensors( - src, dst, num_nodes, lg_src, lg_dst, y.shape[0], h, y, z + src, + dst, + num_nodes, + lg_src, + lg_dst, + y.shape[0], + h, + y, + z, + edge_w, + conv_tri_w, ) + if self.angle_head is not None: + angle_out = { + "eps": self.angle_head(z).squeeze(-1), + "theta_t": bond_angle(r[lg_src], r[lg_dst]), + "weight": ( + torch.ones_like(cos_theta) if tri_w is None else tri_w + ), + "lg_src": lg_src, + "lg_dst": lg_dst, + "src": src, + "dst": dst, + "image": image, + "edge_graph_id": edge_graph_id, + } for layer in self.gcn_layers: - h, y = layer.forward_tensors(src, dst, num_nodes, h, y) + h, y = layer.forward_tensors(src, dst, num_nodes, h, y, edge_w) # Coordinate score: sum the fractional edge offsets into their # destination atom, each weighted by a learned per-edge scalar. w = self.edge_weight_mlp( torch.cat([h[src], h[dst], y], dim=-1) ) # (E, C) + if edge_w is not None: + # Same continuity requirement as the messages: a pair leaving the + # cutoff must stop contributing smoothly, not abruptly. + w = w * edge_w.view(-1, 1) contrib = w.unsqueeze(-1) * df_min.unsqueeze(1) # (E, C, 3) per_node = scatter_sum(contrib, dst, num_nodes) # (N, C, 3) eps_frac = self.score_combine(per_node.transpose(1, 2)).squeeze( @@ -331,4 +576,7 @@ def forward( pooled = scatter_mean(h, node_graph_id, int(natoms.shape[0])) eps_lattice = self.lattice_head(pooled) - return {"eps_frac": eps_frac, "eps_lattice": eps_lattice} + out: Dict = {"eps_frac": eps_frac, "eps_lattice": eps_lattice} + if angle_out is not None: + out["angle"] = angle_out + return out diff --git a/alignn/inverse/evaluate.py b/alignn/inverse/evaluate.py new file mode 100644 index 0000000..8f0b337 --- /dev/null +++ b/alignn/inverse/evaluate.py @@ -0,0 +1,211 @@ +"""Mechanism-level evaluation for generated crystals. + +The AtomBench pipeline already scores match rate, RMSD, ccRMSD, lattice MAE +and KLD, and none of that changes here. What it cannot say is *why* one model +is better, and that is what an angular-diffusion experiment has to answer. +Two extra measurements are provided: + +**Bond-angle distributions.** Compare the angles a model actually generates +against the angles of held-out real structures. This is FoldingDiff's own +diagnostic (Wu et al., Nat. Commun. 15, 1059, 2024): a generative model with +an explicit angular channel should reproduce the natural angular distribution, +and a model that merely places atoms plausibly on average need not. Reported +as KL, Jensen-Shannon and 1-D Wasserstein distance between normalised +histograms, all on the same binning. + +**Relaxation displacement.** How far a generated structure has to move to +reach the nearest local minimum of the force field. MatterGen (Zeni et al., +Nature 639, 624, 2025) evaluates generated structures by how close they sit to +their relaxed counterparts; if explicit angular denoising produces locally +coherent geometry, its samples should need less repair. Reported as the +translation-corrected Cartesian RMSD between the sample and its relaxed self, +plus the fractional volume change and the energy drop. + +Nothing here is used to select a model — the suite is fixed before the +ablations are run, exactly so that a favourable metric cannot be chosen after +seeing the results. +""" + +from __future__ import annotations + +import math +from typing import Dict, List, Optional, Sequence + +import numpy as np +import torch + +__all__ = [ + "bond_angles_deg", + "collect_bond_angles", + "compare_angle_distributions", + "relaxation_displacement", + "structures_from_benchmark_csv", +] + +#: Neighbour cutoff and count used for every angle measurement. These match +#: the pure-torch graph builder's own three-body defaults so that the angles +#: being scored are the ones ALIGNN would see. +DEFAULT_ANGLE_CUTOFF = 3.5 +DEFAULT_MAX_NEIGHBORS = 12 +DEFAULT_BINS = 180 + + +def bond_angles_deg( + atoms, + cutoff: float = DEFAULT_ANGLE_CUTOFF, + max_neighbors: Optional[int] = DEFAULT_MAX_NEIGHBORS, +) -> np.ndarray: + """Every bond angle in one structure, in degrees. + + An angle is formed by each unordered pair of neighbours of a central atom, + where "neighbour" is a periodic image within ``cutoff`` (truncated to the + ``max_neighbors`` closest, as the graph builder does). + """ + from alignn.torch_graph_builder import torch_neighbor_list + + positions = torch.tensor( + np.asarray(atoms.cart_coords, dtype=float), dtype=torch.float64 + ) + lattice = torch.tensor( + np.asarray(atoms.lattice_mat, dtype=float), dtype=torch.float64 + ) + src, _dst, _shift, r = torch_neighbor_list( + positions, + lattice, + cutoff, + max_neighbors=max_neighbors, + use_matscipy_topology=False, + ) + src_np = src.numpy() + # r points from src outward, so the angle at the central atom is taken + # between two of its outgoing vectors directly. + vec = r.numpy() + out: List[float] = [] + order = np.argsort(src_np, kind="stable") + src_sorted, vec_sorted = src_np[order], vec[order] + bounds = np.searchsorted( + src_sorted, np.arange(int(positions.shape[0]) + 1) + ) + for j in range(int(positions.shape[0])): + v = vec_sorted[bounds[j] : bounds[j + 1]] + if len(v) < 2: + continue + unit = v / np.linalg.norm(v, axis=1, keepdims=True).clip(1e-12) + cos = unit @ unit.T + iu = np.triu_indices(len(v), k=1) + out.append(np.degrees(np.arccos(np.clip(cos[iu], -1.0, 1.0)))) + if not out: + return np.zeros(0) + return np.concatenate(out) + + +def collect_bond_angles( + structures: Sequence, + cutoff: float = DEFAULT_ANGLE_CUTOFF, + max_neighbors: Optional[int] = DEFAULT_MAX_NEIGHBORS, +) -> np.ndarray: + """Bond angles pooled over a list of structures, in degrees.""" + parts = [ + bond_angles_deg(a, cutoff=cutoff, max_neighbors=max_neighbors) + for a in structures + ] + parts = [p for p in parts if p.size] + return np.concatenate(parts) if parts else np.zeros(0) + + +def _density(angles: np.ndarray, bins: int) -> np.ndarray: + counts, _ = np.histogram(angles, bins=bins, range=(0.0, 180.0)) + total = counts.sum() + return counts / total if total else counts.astype(float) + + +def compare_angle_distributions( + generated: np.ndarray, + reference: np.ndarray, + bins: int = DEFAULT_BINS, + eps: float = 1e-9, +) -> Dict[str, float]: + """Distances between two pooled bond-angle distributions. + + ``wasserstein`` is exact for a 1-D histogram (the integral of the absolute + CDF difference) and is reported in degrees, so it reads directly as "the + generated angles are off by this much on average". + """ + p = _density(np.asarray(generated), bins) + q = _density(np.asarray(reference), bins) + ps, qs = p + eps, q + eps + ps, qs = ps / ps.sum(), qs / qs.sum() + kl = float((ps * np.log(ps / qs)).sum()) + m = 0.5 * (ps + qs) + js = float( + 0.5 * (ps * np.log(ps / m)).sum() + 0.5 * (qs * np.log(qs / m)).sum() + ) + width = 180.0 / bins + emd = float(np.abs(np.cumsum(p) - np.cumsum(q)).sum() * width) + return { + "kl": kl, + "js": js, + "wasserstein_deg": emd, + "n_generated": int(np.asarray(generated).size), + "n_reference": int(np.asarray(reference).size), + "bins": bins, + } + + +def _min_image_displacement(before, after) -> np.ndarray: + """Cartesian displacement per atom, minimum-image and drift-corrected. + + A relaxation is free to translate the whole cell, and the benchmark's own + metrics quotient that out, so the mean displacement is removed before the + RMSD is taken. + """ + f0 = np.asarray(before.frac_coords, dtype=float) + f1 = np.asarray(after.frac_coords, dtype=float) + df = f1 - f0 + df -= np.round(df) + df -= df.mean(axis=0, keepdims=True) + df -= np.round(df) + return df @ np.asarray(after.lattice_mat, dtype=float) + + +def relaxation_displacement( + before, + after, + energy_before: Optional[float] = None, + energy_after: Optional[float] = None, +) -> Dict[str, float]: + """How far one generated structure moved to reach its local minimum.""" + d = _min_image_displacement(before, after) + norms = np.linalg.norm(d, axis=1) + v0 = float(abs(np.linalg.det(np.asarray(before.lattice_mat, dtype=float)))) + v1 = float(abs(np.linalg.det(np.asarray(after.lattice_mat, dtype=float)))) + out = { + "rmsd_angstrom": float(np.sqrt((norms**2).mean())), + "max_displacement_angstrom": float(norms.max()) if norms.size else 0.0, + "volume_change_frac": (v1 - v0) / v0 if v0 else float("nan"), + } + if energy_before is not None and energy_after is not None: + drop = float(energy_before - energy_after) + out["energy_drop_ev_per_atom"] = ( + drop if math.isfinite(drop) else float("nan") + ) + return out + + +def structures_from_benchmark_csv(path, column: str = "prediction") -> List: + """Read one POSCAR column of an AtomBench CSV into jarvis ``Atoms``. + + The CSV written by ``scripts/atombench/generate_benchmark.py`` holds both + the generated structure (``prediction``) and the held-out reference + (``target``), so one file supplies both sides of the angle comparison. + """ + import csv + + from jarvis.core.atoms import Atoms + + out = [] + with open(path, newline="") as fh: + for row in csv.DictReader(fh): + text = row[column].replace("\\n", "\n") + out.append(Atoms.from_poscar(text)) + return out diff --git a/alignn/inverse/layers.py b/alignn/inverse/layers.py new file mode 100644 index 0000000..b2943aa --- /dev/null +++ b/alignn/inverse/layers.py @@ -0,0 +1,127 @@ +"""ALIGNN convolutions with an optional per-edge / per-triplet weight. + +These are thin subclasses of the shared pure-torch ALIGNN layers in +:mod:`alignn.models.alignn_atomwise_pure`. They exist so that the diffusion +denoiser can attenuate a message continuously instead of a graph edge simply +being present or absent, without touching the property-prediction models. + +Parameter names and shapes are *identical* to the classes they subclass +(``node_update.*`` / ``edge_update.*``), so a checkpoint trained with the +stock layers loads into these and vice versa. + +How the weight enters +--------------------- +The edge-gated convolution aggregates a normalised, gated average + + h_i = sum_j sigma_ij * Bh_j / sum_j sigma_ij . + +A weight ``w_ij`` is applied to ``sigma_ij`` *before both* sums. That is the +only placement with the property we need: an edge with ``w = 0`` leaves ``h`` +exactly as if the edge had never been in the list, so inserting or deleting it +at the cutoff produces no jump. Scaling only the numerator would instead +renormalise the surviving messages and would not be continuous. The +normalisation ``bn_nodes`` / ``bn_edges`` is ``LayerNorm``, computed per +element, so no cross-edge statistic can smuggle a discontinuity back in. + +The same class is used in both roles ALIGNN gives it — over the atom graph the +weight is a per-pair relevance ``s_ij``; over the line graph the same code +receives a per-triplet relevance ``s_ijk`` — which is why no separate triplet +machinery is needed. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from alignn.models.alignn_atomwise_pure import ( + EdgeGatedGraphConvPure, + scatter_sum, +) + +__all__ = ["WeightedEdgeGatedGraphConv", "WeightedALIGNNConv"] + + +class WeightedEdgeGatedGraphConv(EdgeGatedGraphConvPure): + """:class:`EdgeGatedGraphConvPure` with an optional per-edge weight.""" + + def forward_tensors( + self, + src: torch.Tensor, + dst: torch.Tensor, + num_nodes: int, + x: torch.Tensor, + y: torch.Tensor, + edge_weight: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Identical to the base layer when ``edge_weight`` is ``None``.""" + e_src = self.src_gate(x) + e_dst = self.dst_gate(x) + m = e_src[src] + e_dst[dst] + self.edge_gate(y) + sigma = torch.sigmoid(m) + if edge_weight is not None: + sigma = sigma * edge_weight.view(-1, 1) + + Bh = self.dst_update(x) + msg_h = Bh[src] * sigma + sum_sigma_h = scatter_sum(msg_h, dst, num_nodes) + sum_sigma = scatter_sum(sigma, dst, num_nodes) + h = sum_sigma_h / (sum_sigma + 1e-6) + x_new = self.src_update(x) + h + + x_new = F.silu(self.bn_nodes(x_new)) + y_new = F.silu(self.bn_edges(m)) + + if self.residual: + x_new = x + x_new + y_new = y + y_new + return x_new, y_new + + +class WeightedALIGNNConv(nn.Module): + """ALIGNN layer whose pair and triplet messages can both be weighted. + + Mirrors :class:`alignn.models.alignn_atomwise_pure.ALIGNNConvPure` exactly + — the line-graph convolution updates bond features, which the atom-graph + convolution then uses — with two extra optional arguments. + + ``triplet_weight = 0`` for every triplet is what ablation A4 uses: the + angular features ``z`` still evolve and still feed the angle head, but they + no longer reach the bond features, so the coordinate/lattice pathway sees + no angular information. The angular loss then acts as a pure auxiliary + task on a shared trunk, which is the control the design brief asks for. + """ + + def __init__(self, in_features: int, out_features: int): + super().__init__() + self.node_update = WeightedEdgeGatedGraphConv( + in_features, out_features + ) + self.edge_update = WeightedEdgeGatedGraphConv( + out_features, out_features + ) + + def forward_tensors( + self, + g_src: torch.Tensor, + g_dst: torch.Tensor, + g_num_nodes: int, + lg_src: torch.Tensor, + lg_dst: torch.Tensor, + lg_num_nodes: int, + x: torch.Tensor, + y: torch.Tensor, + z: torch.Tensor, + edge_weight: Optional[torch.Tensor] = None, + triplet_weight: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x, m = self.node_update.forward_tensors( + g_src, g_dst, g_num_nodes, x, y, edge_weight + ) + y, z = self.edge_update.forward_tensors( + lg_src, lg_dst, lg_num_nodes, m, z, triplet_weight + ) + return x, y, z diff --git a/alignn/inverse/sample.py b/alignn/inverse/sample.py index 0df9a0f..584c291 100644 --- a/alignn/inverse/sample.py +++ b/alignn/inverse/sample.py @@ -224,18 +224,40 @@ def to_jarvis_atoms( return out +def denoiser_config_from_run(cfg: Dict) -> Dict: + """Denoiser keyword arguments recorded in a training run's config. + + Checkpoints written before the angular channel existed simply lack the + new keys, and the denoiser defaults reproduce their original behaviour, + so old models keep loading unchanged. + """ + out = { + "hidden_features": cfg["hidden_features"], + "alignn_layers": cfg["alignn_layers"], + "gcn_layers": cfg["gcn_layers"], + "knn": cfg["knn"], + "num_steps": cfg["num_steps"], + } + for key in ( + "angle_diffusion", + "angle_feedback", + "topology", + "radius_cutoff", + "envelope_exponent", + "gate_pair_messages", + "angle_basis", + ): + if cfg.get(key) is not None: + out[key] = cfg[key] + return out + + def load_model(checkpoint_path, device, use_ema: bool = True): """Load an ALIGNN-CSP checkpoint into a ready-to-sample model.""" ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False) cfg = ckpt["config"] model = ALIGNNCSP( - denoiser_config={ - "hidden_features": cfg["hidden_features"], - "alignn_layers": cfg["alignn_layers"], - "gcn_layers": cfg["gcn_layers"], - "knn": cfg["knn"], - "num_steps": cfg["num_steps"], - }, + denoiser_config=denoiser_config_from_run(cfg), conditioner_spec=ckpt["conditioner_spec"], ).to(device) # Released checkpoints carry only the EMA weights, which are the ones diff --git a/alignn/inverse/train_csp.py b/alignn/inverse/train_csp.py index 6ecf2af..5060abe 100644 --- a/alignn/inverse/train_csp.py +++ b/alignn/inverse/train_csp.py @@ -7,11 +7,16 @@ import time from copy import deepcopy from pathlib import Path -from typing import Dict +from typing import Dict, Optional import torch from torch.utils.data import DataLoader +from alignn.inverse.ablations import ABLATIONS, ablation_config, describe +from alignn.inverse.angles import ( + angle_denoising_loss, + angular_denoising_target, +) from alignn.inverse.data import ( CrystalDataset, Normalizer, @@ -52,6 +57,31 @@ def update(self, model: torch.nn.Module): s.copy_(p) +def angle_loss_from_output(out: Dict, batch: Dict) -> Optional[torch.Tensor]: + """Angular denoising loss for one forward pass, or ``None`` if disabled. + + The target is the angular displacement the forward process applied to + each triplet, evaluated on the same periodic images the noised geometry + was resolved against, and the loss is FoldingDiff's relevance-weighted + wrapped smooth-L1. See :mod:`alignn.inverse.angles`. + """ + aux = out.get("angle") + if aux is None: + return None + target = angular_denoising_target( + aux["theta_t"], + batch["frac"], + batch["lattice"], + aux["src"], + aux["dst"], + aux["edge_graph_id"], + aux["image"], + aux["lg_src"], + aux["lg_dst"], + ) + return angle_denoising_loss(aux["eps"], target, aux["weight"]) + + def diffusion_loss( model: ALIGNNCSP, schedule: DiffusionSchedule, @@ -60,6 +90,7 @@ def diffusion_loss( cond_dropout: Dict[str, float], lattice_weight: float, frac_weight: float, + angle_weight: float = 0.0, ) -> Dict[str, torch.Tensor]: device = batch["frac"].device natoms = batch["natoms"] @@ -97,10 +128,17 @@ def diffusion_loss( ) loss_lat = torch.nn.functional.mse_loss(out["eps_lattice"], eps_lat) loss_frac = torch.nn.functional.mse_loss(out["eps_frac"], target_frac) + total = lattice_weight * loss_lat + frac_weight * loss_frac + loss_ang = angle_loss_from_output(out, batch) + if loss_ang is None: + loss_ang = loss_frac.new_zeros(()) + else: + total = total + angle_weight * loss_ang return { - "loss": lattice_weight * loss_lat + frac_weight * loss_frac, + "loss": total, "loss_lattice": loss_lat.detach(), "loss_frac": loss_frac.detach(), + "loss_angle": loss_ang.detach(), } @@ -176,6 +214,34 @@ def main(): ap.add_argument("--composition-dropout", type=float, default=0.1) ap.add_argument("--lattice-weight", type=float, default=1.0) ap.add_argument("--frac-weight", type=float, default=10.0) + + # ── angular channel and graph topology ────────────────────────────── + # --ablation picks a named configuration; the individual switches below + # override it, and on their own default to the original model. + ap.add_argument( + "--ablation", + default="A0", + choices=sorted(ABLATIONS), + help="named ablation (see alignn.inverse.ablations); " + "A0 is the unmodified baseline", + ) + ap.add_argument("--angle-diffusion", type=int, default=None) + ap.add_argument("--angle-feedback", type=int, default=None) + ap.add_argument("--topology", default=None, choices=["knn", "radius"]) + ap.add_argument("--radius-cutoff", type=float, default=None) + ap.add_argument("--envelope-exponent", type=int, default=None) + ap.add_argument("--gate-pair-messages", type=int, default=None) + ap.add_argument( + "--angle-basis", default=None, choices=["cosine_rbf", "fourier"] + ) + ap.add_argument( + "--angle-weight", + type=float, + default=1.0, + help="weight on the angular denoising loss; ignored when the " + "angular channel is off, and must be held fixed across an " + "ablation comparison", + ) ap.add_argument("--ema-decay", type=float, default=0.999) ap.add_argument("--grad-clip", type=float, default=1.0) ap.add_argument( @@ -235,17 +301,42 @@ def main(): "composition": args.composition_dropout, } + denoiser_config = { + "hidden_features": args.hidden_features, + "alignn_layers": args.alignn_layers, + "gcn_layers": args.gcn_layers, + "knn": args.knn, + "num_steps": args.num_steps, + **ablation_config(args.ablation), + } + # Explicit switches win over the named ablation, so a one-off variant does + # not need a new entry in the table. + overrides = { + "angle_diffusion": args.angle_diffusion, + "angle_feedback": args.angle_feedback, + "topology": args.topology, + "radius_cutoff": args.radius_cutoff, + "envelope_exponent": args.envelope_exponent, + "gate_pair_messages": args.gate_pair_messages, + "angle_basis": args.angle_basis, + } + for key, value in overrides.items(): + if value is None: + continue + if key in ("angle_diffusion", "angle_feedback", "gate_pair_messages"): + value = bool(value) + denoiser_config[key] = value + # Record what was actually built, so a checkpoint reconstructs itself. + for key, value in denoiser_config.items(): + setattr(args, key, value) + model = ALIGNNCSP( - denoiser_config={ - "hidden_features": args.hidden_features, - "alignn_layers": args.alignn_layers, - "gcn_layers": args.gcn_layers, - "knn": args.knn, - "num_steps": args.num_steps, - }, + denoiser_config=denoiser_config, conditioner_spec=conditioner_spec, ).to(device) n_par = sum(p.numel() for p in model.parameters()) + print(f"ablation {args.ablation}: {describe(args.ablation)}") + print(f"denoiser: {json.dumps(denoiser_config)}") print(f"parameters: {n_par / 1e6:.2f}M modalities: {model.modalities}") if args.init_from: @@ -306,7 +397,12 @@ def main(): t_start = time.time() for epoch in range(1, args.epochs + 1): model.train() - agg = {"loss": 0.0, "loss_lattice": 0.0, "loss_frac": 0.0} + agg = { + "loss": 0.0, + "loss_lattice": 0.0, + "loss_frac": 0.0, + "loss_angle": 0.0, + } nb = 0 for batch in train_dl: batch = batch_to(batch, device) @@ -318,6 +414,7 @@ def main(): cond_dropout, args.lattice_weight, args.frac_weight, + args.angle_weight, ) opt.zero_grad(set_to_none=True) losses["loss"].backward() @@ -333,7 +430,12 @@ def main(): # Validation uses the EMA weights, with the timestep draw fixed so the # curve is comparable epoch to epoch rather than dominated by noise. - val_agg = {"loss": 0.0, "loss_lattice": 0.0, "loss_frac": 0.0} + val_agg = { + "loss": 0.0, + "loss_lattice": 0.0, + "loss_frac": 0.0, + "loss_angle": 0.0, + } nv = 0 gen_state = torch.random.get_rng_state() torch.manual_seed(1234) @@ -348,6 +450,7 @@ def main(): {k: 0.0 for k in cond_dropout}, args.lattice_weight, args.frac_weight, + args.angle_weight, ) for k in val_agg: val_agg[k] += float(losses[k]) @@ -372,12 +475,18 @@ def main(): out_dir / "best_model.pt", ) if epoch % args.log_every == 0 or epoch == 1: + ang = ( + f" ang {agg['loss_angle']:.4f}/{val_agg['loss_angle']:.4f}" + if denoiser_config.get("angle_diffusion") + else "" + ) print( f"epoch {epoch:5d} train {agg['loss']:.4f} " f"(lat {agg['loss_lattice']:.4f} frac {agg['loss_frac']:.4f})" f" val {val_agg['loss']:.4f} " f"(lat {val_agg['loss_lattice']:.4f} " f"frac {val_agg['loss_frac']:.4f})" + f"{ang}" f" best {best_val:.4f} {time.time() - t_start:.0f}s", flush=True, ) diff --git a/alignn/tests/test_inverse_angle_diffusion.py b/alignn/tests/test_inverse_angle_diffusion.py new file mode 100644 index 0000000..3688eb4 --- /dev/null +++ b/alignn/tests/test_inverse_angle_diffusion.py @@ -0,0 +1,490 @@ +"""Tests for the angular diffusion channel and the smooth line-graph topology. + +These cover section 10 of the design brief: + +* the angular corruption respects angular periodicity, +* the cutoff envelope and its first two derivatives vanish at ``r_c``, +* a pair contribution goes to zero continuously as ``r -> r_c``, +* a triplet contribution goes to zero when *either* of its bonds does, +* moving one atom across the cutoff does not make the model output jump + merely because a sparse edge was inserted or deleted, +* periodic coordinate handling is still correct, +* the default configuration reproduces the original model exactly, +* the angle-enabled configuration trains and samples. +""" + +import math + +import pytest +import torch + +from alignn.inverse.ablations import ABLATIONS, ablation_config +from alignn.inverse.angles import ( + CutoffPolynomial, + angle_denoising_loss, + angular_denoising_target, + bond_angle, + triplet_relevance, + wrap_angle, +) +from alignn.inverse.data import Normalizer +from alignn.inverse.denoiser import ALIGNNCSPDenoiser, dense_pair_index +from alignn.inverse.diffusion import DiffusionSchedule, wrap_frac +from alignn.inverse.model import ALIGNNCSP +from alignn.inverse.sample import sample +from alignn.inverse.train_csp import diffusion_loss + +CUTOFF = 5.0 + + +@pytest.fixture(autouse=True) +def _float32_default(): + """Pin the default dtype for this module. + + ``test_force_reduction`` sets the global default to float64 at import + time and does not restore it, and the diffusion denoiser's timestep + embedding is float32 by construction, so without this the results of + these tests would depend on collection order. + """ + previous = torch.get_default_dtype() + torch.set_default_dtype(torch.float32) + yield + torch.set_default_dtype(previous) + + +SMALL = dict( + hidden_features=32, + embedding_features=16, + alignn_layers=2, + gcn_layers=1, + rbf_bins=16, + triplet_bins=8, + score_channels=4, + num_steps=50, +) + + +def _batch(seed=0, natoms=(3, 2), cell=4.0): + """A tiny two-crystal batch.""" + torch.manual_seed(seed) + n = torch.tensor(natoms) + total = int(n.sum()) + return { + "frac": torch.rand(total, 3), + "lattice": torch.eye(3).repeat(len(natoms), 1, 1) * cell, + "atomic_numbers": torch.randint(1, 60, (total,)), + "natoms": n, + "node_graph_id": torch.repeat_interleave(torch.arange(len(natoms)), n), + "prop": torch.zeros(len(natoms)), + } + + +def _forward(model, batch, t=25): + return model( + frac=batch["frac"], + lattice=batch["lattice"], + lattice_vec6=torch.zeros(len(batch["natoms"]), 6), + atomic_numbers=batch["atomic_numbers"], + natoms=batch["natoms"], + node_graph_id=batch["node_graph_id"], + t=torch.full((len(batch["natoms"]),), t, dtype=torch.long), + ) + + +# ── angular periodicity ────────────────────────────────────────────────── +def test_wrap_angle_is_periodic_and_in_range(): + x = torch.linspace(-20.0, 20.0, 401) + w = wrap_angle(x) + assert torch.all(w >= -math.pi) and torch.all(w < math.pi) + # Adding a full turn changes nothing. + for k in (-2, -1, 1, 2): + assert torch.allclose(wrap_angle(x + k * 2 * math.pi), w, atol=1e-5) + # And it is the identity where it should be. + inner = torch.linspace(-3.0, 3.0, 61) + assert torch.allclose(wrap_angle(inner), inner, atol=1e-6) + + +def test_angle_loss_is_wrapped(): + """A residual of 2*pi is no error at all.""" + pred = torch.tensor([1.0, 2.0]) + zero = angle_denoising_loss(pred, pred) + wrapped = angle_denoising_loss(pred + 2 * math.pi, pred) + assert float(zero) == pytest.approx(0.0, abs=1e-6) + assert float(wrapped) == pytest.approx(0.0, abs=1e-6) + + +def test_angular_target_vanishes_without_corruption(): + """theta_t == theta_0 when the noised structure *is* the clean one.""" + b = _batch(seed=3) + model = ALIGNNCSPDenoiser(**SMALL, **ablation_config("A3")) + out = _forward(model, b) + aux = out["angle"] + target = angular_denoising_target( + aux["theta_t"], + b["frac"], + b["lattice"], + aux["src"], + aux["dst"], + aux["edge_graph_id"], + aux["image"], + aux["lg_src"], + aux["lg_dst"], + ) + assert target.numel() > 0 + assert float(target.abs().max()) < 1e-4 + + +def test_angular_target_matches_a_hand_computed_rotation(): + """Bending one bond by a known angle shows up in the target.""" + # Two atoms placed so the triplet at atom 0 is a right angle, then the + # third atom is swung to 60 degrees. + lattice = torch.eye(3).unsqueeze(0) * 12.0 + clean = torch.tensor([[0.0, 0.0, 0.0], [0.25, 0.0, 0.0], [0.0, 0.25, 0.0]]) + moved = clean.clone() + moved[2] = torch.tensor([0.125, 0.125 * math.sqrt(3.0), 0.0]) + src, dst, egid = dense_pair_index(torch.tensor([3])) + model = ALIGNNCSPDenoiser(**SMALL, **ablation_config("A3")) + b = { + "frac": moved, + "lattice": lattice, + "atomic_numbers": torch.tensor([6, 6, 6]), + "natoms": torch.tensor([3]), + "node_graph_id": torch.zeros(3, dtype=torch.long), + } + aux = _forward(model, b)["angle"] + target = angular_denoising_target( + aux["theta_t"], + clean, + lattice, + aux["src"], + aux["dst"], + aux["edge_graph_id"], + aux["image"], + aux["lg_src"], + aux["lg_dst"], + ) + # The 1-0-2 triplet went from 90 to 60 degrees, i.e. -30 degrees. + is_triplet_at_0 = ( + (aux["dst"][aux["lg_src"]] == 0) + & (aux["src"][aux["lg_src"]] == 1) + & (aux["dst"][aux["lg_dst"]] == 2) + ) + assert bool(is_triplet_at_0.any()) + got = math.degrees(float(target[is_triplet_at_0][0])) + assert got == pytest.approx(-30.0, abs=0.5) + + +# ── smooth cutoff ──────────────────────────────────────────────────────── +def test_envelope_and_two_derivatives_vanish_at_cutoff(): + env = CutoffPolynomial(cutoff=CUTOFF, coeff=5.0) + r = torch.tensor([CUTOFF], dtype=torch.float64, requires_grad=True) + u = env(r) + (du,) = torch.autograd.grad(u.sum(), r, create_graph=True) + (d2u,) = torch.autograd.grad(du.sum(), r, create_graph=True) + assert float(u) == pytest.approx(0.0, abs=1e-12) + assert float(du) == pytest.approx(0.0, abs=1e-10) + assert float(d2u) == pytest.approx(0.0, abs=1e-8) + # Unit at zero separation, monotone decreasing, never negative. + grid = torch.linspace(0.0, CUTOFF, 501, dtype=torch.float64) + vals = env(grid) + assert float(env(torch.zeros(1, dtype=torch.float64))) == pytest.approx( + 1.0 + ) + assert torch.all(vals >= 0.0) + assert torch.all(vals[1:] <= vals[:-1] + 1e-12) + + +def test_pair_contribution_goes_to_zero_continuously(): + env = CutoffPolynomial(cutoff=CUTOFF, coeff=5.0) + eps = 1e-6 + inside = float(env(torch.tensor([CUTOFF - eps], dtype=torch.float64))) + outside = float(env(torch.tensor([CUTOFF + eps], dtype=torch.float64))) + assert inside == pytest.approx(0.0, abs=1e-14) + assert outside == 0.0 + # No step anywhere across the boundary. Evaluated in double precision: + # the polynomial is written as a sum of terms of order 20 that cancel to + # ~1e-5 near r_c, so in float32 the *value* carries ~1e-6 of rounding + # noise. That is harmless — it multiplies messages that are already being + # driven to zero — but it would swamp a test of the exact property. + grid = torch.linspace( + CUTOFF - 0.05, CUTOFF + 0.05, 2001, dtype=torch.float64 + ) + vals = env(grid) + step = float(vals.diff().abs().max()) + assert step < 5e-7 + # And the whole window sits within rounding distance of zero: there is + # no cliff for a message to fall off. + assert float(vals.max()) < 1e-4 + + +def test_triplet_weight_vanishes_when_either_bond_reaches_the_cutoff(): + env = CutoffPolynomial(cutoff=CUTOFF, coeff=5.0) + dist = torch.tensor([1.0, CUTOFF - 1e-7, 2.0]) + s = env(dist) + lg_src = torch.tensor([0, 0, 1, 2]) + lg_dst = torch.tensor([2, 1, 2, 0]) + w = triplet_relevance(s, lg_src, lg_dst) + # Any triplet touching edge 1 (at the cutoff) is off. + assert float(w[1]) == pytest.approx(0.0, abs=1e-12) + assert float(w[2]) == pytest.approx(0.0, abs=1e-12) + # The one made of two short bonds is not. + assert float(w[0]) > 0.1 + + +# ── no jump when an edge enters or leaves the sparse graph ─────────────── +def test_no_finite_jump_when_a_triplet_crosses_the_cutoff(): + """Sweep one atom through the radius and check the output is smooth. + + The line graph is rebuilt from scratch at every position, so triplets are + genuinely inserted and deleted during this sweep; the envelope is what + makes that invisible. + """ + torch.manual_seed(0) + model = ALIGNNCSPDenoiser( + **SMALL, **ablation_config("A3"), radius_cutoff=CUTOFF + ).eval() + cell = 20.0 + lattice = torch.eye(3).unsqueeze(0) * cell + base = torch.tensor([[0.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.0, 0.0, 0.0]]) + + def run(x): + frac = base.clone() + frac[2, 0] = x / cell + out = model( + frac=wrap_frac(frac), + lattice=lattice, + lattice_vec6=torch.zeros(1, 6), + atomic_numbers=torch.tensor([6, 6, 6]), + natoms=torch.tensor([3]), + node_graph_id=torch.zeros(3, dtype=torch.long), + t=torch.tensor([25]), + ) + return out + + # A dense sweep straight through the cutoff radius. + xs = torch.linspace(CUTOFF - 0.3, CUTOFF + 0.3, 121) + scores, n_triplets = [], [] + with torch.no_grad(): + for x in xs: + out = run(float(x)) + scores.append(out["eps_frac"].clone()) + n_triplets.append(int(out["angle"]["eps"].shape[0])) + # The triplet count really does change across the sweep. + assert len(set(n_triplets)) > 1 + steps = torch.stack( + [(a - b).abs().max() for a, b in zip(scores[1:], scores[:-1])] + ) + # Every consecutive step is small; a topology jump would show up as one + # step far larger than its neighbours. + assert float(steps.max()) < 20.0 * float(steps.median()) + 1e-6 + + +def test_gated_message_ignores_a_zero_weight_edge(): + """An edge weighted to zero leaves the aggregation exactly unchanged.""" + from alignn.inverse.layers import WeightedEdgeGatedGraphConv + + torch.manual_seed(0) + conv = WeightedEdgeGatedGraphConv(8, 8).eval() + src = torch.tensor([0, 1, 2, 0]) + dst = torch.tensor([1, 2, 0, 2]) + x = torch.randn(3, 8) + y = torch.randn(4, 8) + w = torch.tensor([1.0, 1.0, 1.0, 0.0]) + with torch.no_grad(): + gated, _ = conv.forward_tensors(src, dst, 3, x, y, w) + # Same graph with that edge physically removed. + keep = torch.tensor([0, 1, 2]) + dropped, _ = conv.forward_tensors(src[keep], dst[keep], 3, x, y[keep]) + assert torch.allclose(gated, dropped, atol=1e-6) + + +# ── periodicity of the model itself ────────────────────────────────────── +@pytest.mark.parametrize("name", sorted(ABLATIONS)) +def test_output_is_invariant_to_lattice_translations(name): + """Adding whole cells to the coordinates must change nothing.""" + b = _batch(seed=1) + model = ALIGNNCSPDenoiser(**SMALL, **ablation_config(name)).eval() + with torch.no_grad(): + a = _forward(model, b) + shifted = dict(b) + shifted["frac"] = b["frac"] + torch.tensor([1.0, -2.0, 3.0]) + c = _forward(model, shifted) + assert torch.allclose(a["eps_frac"], c["eps_frac"], atol=1e-5) + assert torch.allclose(a["eps_lattice"], c["eps_lattice"], atol=1e-5) + + +@pytest.mark.parametrize("name", sorted(ABLATIONS)) +def test_coordinate_score_is_invariant_to_a_global_shift(name): + """A rigid translation of the crystal is not a change of structure.""" + b = _batch(seed=2) + model = ALIGNNCSPDenoiser(**SMALL, **ablation_config(name)).eval() + with torch.no_grad(): + a = _forward(model, b) + shifted = dict(b) + shifted["frac"] = wrap_frac(b["frac"] + 0.137) + c = _forward(model, shifted) + assert torch.allclose(a["eps_frac"], c["eps_frac"], atol=1e-5) + + +# ── the baseline is untouched ──────────────────────────────────────────── +def test_default_config_builds_the_original_model(): + """No new parameters, and no new outputs, unless asked for.""" + model = ALIGNNCSPDenoiser(**SMALL) + keys = set(model.state_dict()) + assert not any(k.startswith("angle_head") for k in keys) + assert not any("envelope" in k for k in keys) + assert model.topology == "knn" + assert model.angle_diffusion is False + out = _forward(model, _batch()) + assert set(out) == {"eps_frac", "eps_lattice"} + + +def test_angle_head_does_not_perturb_the_structural_pathway(): + """A1 must equal A0 on eps_frac / eps_lattice, given the same weights. + + This is what makes the angular objective an addition rather than a + change: with the topology held at the baseline's, switching the angle + head on adds an output without moving the existing ones. + """ + torch.manual_seed(7) + base = ALIGNNCSPDenoiser(**SMALL, **ablation_config("A0")).eval() + angled = ALIGNNCSPDenoiser(**SMALL, **ablation_config("A1")).eval() + missing, unexpected = angled.load_state_dict( + base.state_dict(), strict=False + ) + assert not unexpected + assert all(k.startswith("angle_head") for k in missing) + b = _batch(seed=5) + with torch.no_grad(): + a, c = _forward(base, b), _forward(angled, b) + assert torch.allclose(a["eps_frac"], c["eps_frac"], atol=1e-6) + assert torch.allclose(a["eps_lattice"], c["eps_lattice"], atol=1e-6) + assert "angle" in c + + +def test_a4_cuts_the_angular_coupling(): + """A4's structural output must not depend on the angular features. + + Perturbing only the angle-embedding weights moves A3's coordinate score + and leaves A4's alone. + """ + b = _batch(seed=6) + results = {} + for name in ("A1", "A3", "A4"): + torch.manual_seed(11) # identical weights in every arm + model = ALIGNNCSPDenoiser(**SMALL, **ablation_config(name)).eval() + # The coordinate head is zero-initialised by design, which would make + # every arm read zero; give it signal first. + torch.nn.init.normal_(model.score_combine.weight, std=0.5) + with torch.no_grad(): + before = _forward(model, b)["eps_frac"].clone() + torch.manual_seed(3) + for pname, p in model.named_parameters(): + if pname.startswith("angle_embedding"): + p.add_(torch.randn_like(p) * 0.5) + after = _forward(model, b)["eps_frac"] + results[name] = float( + (after - before).abs().max() / before.abs().max() + ) + assert results["A4"] == pytest.approx(0.0, abs=1e-9) + assert results["A3"] > 1e-3 + assert results["A1"] > 1e-3 + + +# ── end-to-end ─────────────────────────────────────────────────────────── +@pytest.mark.parametrize("name", sorted(ABLATIONS)) +def test_train_step_and_sampling_run(name): + """Forward, loss, backward and reverse sampling for every ablation.""" + torch.manual_seed(0) + model = ALIGNNCSP( + denoiser_config={**SMALL, **ablation_config(name)}, + conditioner_spec={"composition": {"type": "composition"}}, + ) + schedule = DiffusionSchedule(num_steps=SMALL["num_steps"]) + normalizer = Normalizer( + lattice_mean=torch.zeros(6), + lattice_std=torch.ones(6), + prop_mean=0.0, + prop_std=1.0, + ) + b = _batch(seed=4) + losses = diffusion_loss( + model, + schedule, + normalizer, + b, + {"composition": 0.1}, + lattice_weight=1.0, + frac_weight=10.0, + angle_weight=1.0, + ) + losses["loss"].backward() + grads = [ + p.grad.abs().sum() for p in model.parameters() if p.grad is not None + ] + assert float(sum(grads)) > 0.0 + if ABLATIONS[name]["angle_diffusion"]: + assert float(losses["loss_angle"]) > 0.0 + else: + assert float(losses["loss_angle"]) == 0.0 + + out = sample( + model, + schedule, + normalizer, + b, + guidance=1.0, + n_corrector=0, + device=torch.device("cpu"), + ) + assert out["frac"].shape == b["frac"].shape + assert torch.isfinite(out["frac"]).all() + assert torch.isfinite(out["lattice"]).all() + + +def test_angle_loss_gradient_reaches_the_shared_backbone(): + """The angular objective must train the trunk, not just its own head.""" + torch.manual_seed(0) + model = ALIGNNCSPDenoiser(**SMALL, **ablation_config("A3")) + # The angle head's last layer is zero-initialised so that training starts + # from a silent prediction; that also zeroes the gradient through it, so + # this test looks at the model one step into training. + torch.nn.init.normal_(model.angle_head[-1].weight, std=0.5) + b = _batch(seed=8) + out = _forward(model, b) + aux = out["angle"] + target = angular_denoising_target( + aux["theta_t"], + b["frac"] + 0.05, + b["lattice"], + aux["src"], + aux["dst"], + aux["edge_graph_id"], + aux["image"], + aux["lg_src"], + aux["lg_dst"], + ) + angle_denoising_loss(aux["eps"], target, aux["weight"]).backward() + touched = { + name + for name, p in model.named_parameters() + if p.grad is not None and float(p.grad.abs().sum()) > 0 + } + assert any(n.startswith("alignn_layers") for n in touched) + assert any(n.startswith("angle_embedding") for n in touched) + assert any(n.startswith("edge_embedding") for n in touched) + + +def test_bond_angle_matches_a_known_geometry(): + r_ij = torch.tensor([[1.0, 0.0, 0.0]]) + r_jk = torch.tensor([[1.0, 0.0, 0.0]]) + # i -> j -> k collinear and continuing forward is a straight 180 degrees. + # The tolerance is the deliberate clamp inside bond_angle, which keeps + # acos differentiable at the poles at the cost of ~0.03 degrees there. + assert math.degrees(float(bond_angle(r_ij, r_jk))) == pytest.approx( + 180.0, abs=0.05 + ) + r_jk = torch.tensor([[0.0, 1.0, 0.0]]) + assert math.degrees(float(bond_angle(r_ij, r_jk))) == pytest.approx( + 90.0, abs=1e-2 + ) diff --git a/scripts/atombench/angle_eval.py b/scripts/atombench/angle_eval.py new file mode 100755 index 0000000..de3c8f6 --- /dev/null +++ b/scripts/atombench/angle_eval.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +"""Mechanism metrics for a generated benchmark CSV. + +Complements ``score.sh``, which reports the AtomBench benchmark numbers. This +adds the two diagnostics that say whether an *angular* channel is doing +anything: + + bond-angle distribution generated vs. held-out real structures, the + comparison FoldingDiff uses + relaxation displacement how far a sample has to move to reach the + nearest ALIGNN-FF local minimum, the proximity + MatterGen evaluates + +Both sides of the angle comparison come from the same file: the CSV written by +``generate_benchmark.py`` carries the generated structure in ``prediction`` and +the held-out reference in ``target``. + + python scripts/atombench/angle_eval.py runs/bench/alignn_csp.csv \ + --relax --limit 50 + +Writes ``angle_metrics.json`` next to the CSV. Relaxation is off by default +because it costs about a second per structure. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + +from alignn.inverse.evaluate import ( + DEFAULT_ANGLE_CUTOFF, + DEFAULT_BINS, + DEFAULT_MAX_NEIGHBORS, + collect_bond_angles, + compare_angle_distributions, + relaxation_displacement, + structures_from_benchmark_csv, +) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("csv", help="benchmark CSV from generate_benchmark.py") + ap.add_argument( + "--reference-csv", + default=None, + help="take the reference structures from a different CSV's " + "'target' column (default: the same file)", + ) + ap.add_argument("--cutoff", type=float, default=DEFAULT_ANGLE_CUTOFF) + ap.add_argument("--max-neighbors", type=int, default=DEFAULT_MAX_NEIGHBORS) + ap.add_argument("--bins", type=int, default=DEFAULT_BINS) + ap.add_argument( + "--relax", + action="store_true", + help="also relax each generated structure with ALIGNN-FF and report " + "how far it moved", + ) + ap.add_argument("--relax-steps", type=int, default=200) + ap.add_argument("--relax-fmax", type=float, default=0.05) + ap.add_argument( + "--limit", + type=int, + default=0, + help="only use the first N rows (0 = all); relaxation is the slow " + "part, so this mostly matters with --relax", + ) + ap.add_argument("--output", default=None) + args = ap.parse_args() + + csv_path = Path(args.csv) + generated = structures_from_benchmark_csv(csv_path, "prediction") + reference = structures_from_benchmark_csv( + args.reference_csv or csv_path, "target" + ) + if args.limit: + generated = generated[: args.limit] + reference = reference[: args.limit] + print(f"{len(generated)} generated, {len(reference)} reference structures") + + kw = {"cutoff": args.cutoff, "max_neighbors": args.max_neighbors} + gen_angles = collect_bond_angles(generated, **kw) + ref_angles = collect_bond_angles(reference, **kw) + metrics = { + "angle_distribution": compare_angle_distributions( + gen_angles, ref_angles, bins=args.bins + ), + "angle_cutoff": args.cutoff, + "angle_max_neighbors": args.max_neighbors, + "generated_angle_mean_deg": ( + float(np.mean(gen_angles)) if gen_angles.size else None + ), + "reference_angle_mean_deg": ( + float(np.mean(ref_angles)) if ref_angles.size else None + ), + } + d = metrics["angle_distribution"] + print( + f" bond angles: KL {d['kl']:.4f} JS {d['js']:.4f} " + f"Wasserstein {d['wasserstein_deg']:.3f} deg " + f"({d['n_generated']} vs {d['n_reference']} angles)" + ) + + if args.relax: + from alignn.inverse.relax_rank import AlignnFFRelaxer + + relaxer = AlignnFFRelaxer( + relax_cell=True, fmax=args.relax_fmax, steps=args.relax_steps + ) + rows = [] + for i, atoms in enumerate(generated): + e0 = None + try: + e0 = relaxer.energy(atoms) + except Exception as exc: # noqa: BLE001 + print(f" [{i}] single point failed: {exc}") + res = relaxer.relax(atoms) + if res.error: + print(f" [{i}] relaxation failed: {res.error}") + continue + rows.append( + relaxation_displacement( + atoms, res.atoms, e0, res.energy_per_atom + ) + ) + if rows: + keys = sorted({k for r in rows for k in r}) + summary = { + k: float(np.nanmean([r[k] for r in rows if k in r])) + for k in keys + } + summary["n_relaxed"] = len(rows) + metrics["relaxation"] = summary + print( + f" relaxation: RMSD {summary['rmsd_angstrom']:.4f} A " + f"|dV|/V {abs(summary['volume_change_frac']):.4f} " + f"over {len(rows)} structures" + ) + + out = Path(args.output or csv_path.with_name("angle_metrics.json")) + out.write_text(json.dumps(metrics, indent=2)) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/atombench/run_angle_ablation.sh b/scripts/atombench/run_angle_ablation.sh new file mode 100755 index 0000000..c04ec9c --- /dev/null +++ b/scripts/atombench/run_angle_ablation.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +# Train, generate and score the angular-diffusion ablation suite. +# +# bash scripts/atombench/run_angle_ablation.sh [seeds...] +# +# Every arm gets the same data split, the same optimiser settings, the same +# epoch budget and the same seed list; the arms differ only in the switches +# under test. That is the whole point — a single run of the proposed model +# against a single run of the baseline cannot separate the effect from the +# spread across seeds, which on splits this size has been large enough to +# invert comparisons before (see alignn/inverse/README.md). +# +# A0 baseline, angles as features only +# A1 + explicit angular denoising +# A2 + smooth radius topology, no angular objective +# A3 both: the proposed model +# A4 control: angular objective with the angle->bond coupling cut +# A6 A3 with the Fourier angular basis +# +# A5 in the design brief is the A1-vs-A3 (and A0-vs-A2) contrast, which these +# runs already provide. +set -euo pipefail + +DATA="${1:?usage: run_angle_ablation.sh [seeds...]}" +OUT="${2:?}" +shift 2 +SEEDS=("$@") +[ ${#SEEDS[@]} -eq 0 ] && SEEDS=(0 1 2) + +ARMS="${ARMS:-A0 A1 A2 A3 A4}" +EPOCHS="${EPOCHS:-3000}" +NCAND="${NCAND:-8}" +GUIDANCE="${GUIDANCE:-2.0}" +ANGLE_WEIGHT="${ANGLE_WEIGHT:-1.0}" +GPU="${GPU:-0}" +WORKERS="${WORKERS:-20}" +export SCORE_ENV="${SCORE_ENV:-}" + +for arm in $ARMS; do + for seed in "${SEEDS[@]}"; do + run="$OUT/${arm}_s${seed}" + echo "=== training $arm seed $seed -> $run" + CUDA_VISIBLE_DEVICES="$GPU" python -u -m alignn.inverse.train_csp \ + --data-dir "$DATA" --output "$run" \ + --ablation "$arm" --angle-weight "$ANGLE_WEIGHT" \ + --epochs "$EPOCHS" --seed "$seed" + + echo "=== generating $arm seed $seed" + CUDA_VISIBLE_DEVICES="$GPU" OMP_NUM_THREADS=1 python -u \ + scripts/atombench/generate_benchmark.py \ + --checkpoint "$run/best_model.pt" --data-dir "$DATA" \ + --output-csv "$run/bench.csv" \ + --num-candidates "$NCAND" --guidance "$GUIDANCE" \ + --relax cell --rank energy --relax-workers "$WORKERS" + + echo "=== mechanism metrics $arm seed $seed" + python scripts/atombench/angle_eval.py "$run/bench.csv" --relax + done +done + +echo +echo "=== benchmark scores" +bash scripts/atombench/score.sh "$OUT"/*/bench.csv From 2ba2345e5fb1df0bb86d7b9e7933b88fb2e953e5 Mon Sep 17 00:00:00 2001 From: crhysc Date: Sat, 29 Aug 2026 18:16:41 -0400 Subject: [PATCH 2/4] inverse: keep the new evaluation code on torch too The model, training and sampling path was already pure PyTorch; this makes the evaluation helpers match rather than dropping into numpy for geometry. - evaluate.py computes bond angles with the same neighbour list and the same grouped-pair construction the line-graph builder uses, and does the histogram / KL / JS / Wasserstein arithmetic in torch. Results are identical to the numpy version: fcc still gives exactly 60/90/120/180 degrees and the Wasserstein distance is still calibrated in degrees. - angles.py drops an unused numpy histogram helper that duplicated evaluate.py's. - angle_eval.py follows. No new dependencies: setup.py and environment.yml are untouched, and the new modules import nothing beyond the standard library, torch and alignn itself (jarvis only lazily, for reading POSCARs). Adds a test pinning the evaluation geometry to fcc's known bond angles. Claude-Session: https://claude.ai/code/session_01ULEhMd27yfBmKf5HXqmvFE --- alignn/inverse/README.md | 11 ++ alignn/inverse/angles.py | 19 +-- alignn/inverse/evaluate.py | 169 ++++++++++++------- alignn/tests/test_inverse_angle_diffusion.py | 38 +++++ scripts/atombench/angle_eval.py | 19 ++- 5 files changed, 172 insertions(+), 84 deletions(-) diff --git a/alignn/inverse/README.md b/alignn/inverse/README.md index 30123cb..b96e8c2 100644 --- a/alignn/inverse/README.md +++ b/alignn/inverse/README.md @@ -378,6 +378,17 @@ bash scripts/atombench/score.sh runs/bench/alignn_csp.csv `scripts/atombench/run_ablation.sh` runs the four pipeline variants (raw / rank / relax / full) so you can see what each stage contributes. +Or skip the four commands above. `task_runners/` wraps each published result +in one resumable task, with the arguments pinned, seeds handled and an sbatch +script per task: + +```bash +python task_runners/run_task.py tasks # what is available +python task_runners/run_task.py bench-jarvis # train, generate, score, x3 seeds +python task_runners/run_task.py bench-jarvis --aggregate --latex +bash task_runners/submit.sh bench-jarvis # the same, through SLURM +``` + ## What actually moved the numbers Findings from the AtomBench runs, recorded so they are not rediscovered: diff --git a/alignn/inverse/angles.py b/alignn/inverse/angles.py index 87f6942..982de21 100644 --- a/alignn/inverse/angles.py +++ b/alignn/inverse/angles.py @@ -66,7 +66,7 @@ from __future__ import annotations import math -from typing import Optional, Tuple +from typing import Optional import torch @@ -198,20 +198,3 @@ def angle_denoising_loss( if weight is None: return per_triplet.mean() return (weight * per_triplet).sum() / weight.sum().clamp_min(1e-8) - - -def angle_histogram( - angles_deg, bins: int = 180, lo: float = 0.0, hi: float = 180.0 -) -> Tuple: - """Normalised histogram of bond angles in degrees. - - Kept here so the evaluation code and the tests share one definition. - """ - import numpy as np - - counts, edges = np.histogram( - np.asarray(angles_deg, dtype=float), bins=bins, range=(lo, hi) - ) - total = counts.sum() - density = counts / total if total else counts.astype(float) - return density, edges diff --git a/alignn/inverse/evaluate.py b/alignn/inverse/evaluate.py index 8f0b337..05c807c 100644 --- a/alignn/inverse/evaluate.py +++ b/alignn/inverse/evaluate.py @@ -24,6 +24,11 @@ Nothing here is used to select a model — the suite is fixed before the ablations are run, exactly so that a favourable metric cannot be chosen after seeing the results. + +Everything is plain PyTorch, on the same neighbour list the model itself uses, +so the angles being scored are the angles ALIGNN would see. ``jarvis`` enters +only to parse POSCARs and the ASE force field only through the existing +:mod:`alignn.inverse.relax_rank`. """ from __future__ import annotations @@ -31,7 +36,6 @@ import math from typing import Dict, List, Optional, Sequence -import numpy as np import torch __all__ = [ @@ -43,18 +47,61 @@ ] #: Neighbour cutoff and count used for every angle measurement. These match -#: the pure-torch graph builder's own three-body defaults so that the angles -#: being scored are the ones ALIGNN would see. +#: the pure-torch graph builder's own three-body defaults. DEFAULT_ANGLE_CUTOFF = 3.5 DEFAULT_MAX_NEIGHBORS = 12 DEFAULT_BINS = 180 +_EPS = 1e-12 + + +def _atoms_to_tensors(atoms, dtype=torch.float64): + """Cartesian positions and lattice of a jarvis ``Atoms`` as tensors.""" + return ( + torch.tensor(atoms.cart_coords, dtype=dtype), + torch.tensor(atoms.lattice_mat, dtype=dtype), + ) + + +def _same_source_pairs(src: torch.Tensor, num_nodes: int): + """Unordered pairs of edges sharing a source node. + + Same construction the line-graph builder uses: group the edges by their + shared node, expand each group to all ordered pairs, then keep one of + each two. Returns indices into the edge list. + """ + n_edges = int(src.shape[0]) + device = src.device + if n_edges == 0: + empty = torch.empty(0, dtype=torch.long, device=device) + return empty, empty + order = torch.argsort(src, stable=True) + src_sorted = src[order] + counts = torch.bincount(src_sorted, minlength=num_nodes) + starts = torch.cumsum(counts, 0) - counts + # Every edge pairs with each edge in its own group. + per_edge = counts[src_sorted] + total = int(per_edge.sum()) + if total == 0: + empty = torch.empty(0, dtype=torch.long, device=device) + return empty, empty + positions = torch.arange(n_edges, device=device) + left = torch.repeat_interleave(positions, per_edge) + cum = torch.cumsum(per_edge, 0) + row_start = cum - per_edge + offsets = torch.arange(total, device=device) - torch.repeat_interleave( + row_start, per_edge + ) + right = torch.repeat_interleave(starts[src_sorted], per_edge) + offsets + keep = left < right + return order[left[keep]], order[right[keep]] + def bond_angles_deg( atoms, cutoff: float = DEFAULT_ANGLE_CUTOFF, max_neighbors: Optional[int] = DEFAULT_MAX_NEIGHBORS, -) -> np.ndarray: +) -> torch.Tensor: """Every bond angle in one structure, in degrees. An angle is formed by each unordered pair of neighbours of a central atom, @@ -63,12 +110,7 @@ def bond_angles_deg( """ from alignn.torch_graph_builder import torch_neighbor_list - positions = torch.tensor( - np.asarray(atoms.cart_coords, dtype=float), dtype=torch.float64 - ) - lattice = torch.tensor( - np.asarray(atoms.lattice_mat, dtype=float), dtype=torch.float64 - ) + positions, lattice = _atoms_to_tensors(atoms) src, _dst, _shift, r = torch_neighbor_list( positions, lattice, @@ -76,96 +118,95 @@ def bond_angles_deg( max_neighbors=max_neighbors, use_matscipy_topology=False, ) - src_np = src.numpy() - # r points from src outward, so the angle at the central atom is taken - # between two of its outgoing vectors directly. - vec = r.numpy() - out: List[float] = [] - order = np.argsort(src_np, kind="stable") - src_sorted, vec_sorted = src_np[order], vec[order] - bounds = np.searchsorted( - src_sorted, np.arange(int(positions.shape[0]) + 1) + left, right = _same_source_pairs(src, int(positions.shape[0])) + if left.numel() == 0: + return torch.zeros(0, dtype=positions.dtype) + # r points away from the shared atom, so the interior angle is the plain + # angle between the two outgoing vectors. + a = r[left] + b = r[right] + cos = (a * b).sum(-1) / ( + a.norm(dim=-1).clamp_min(_EPS) * b.norm(dim=-1).clamp_min(_EPS) ) - for j in range(int(positions.shape[0])): - v = vec_sorted[bounds[j] : bounds[j + 1]] - if len(v) < 2: - continue - unit = v / np.linalg.norm(v, axis=1, keepdims=True).clip(1e-12) - cos = unit @ unit.T - iu = np.triu_indices(len(v), k=1) - out.append(np.degrees(np.arccos(np.clip(cos[iu], -1.0, 1.0)))) - if not out: - return np.zeros(0) - return np.concatenate(out) + return torch.rad2deg(torch.acos(cos.clamp(-1.0, 1.0))) def collect_bond_angles( structures: Sequence, cutoff: float = DEFAULT_ANGLE_CUTOFF, max_neighbors: Optional[int] = DEFAULT_MAX_NEIGHBORS, -) -> np.ndarray: +) -> torch.Tensor: """Bond angles pooled over a list of structures, in degrees.""" parts = [ bond_angles_deg(a, cutoff=cutoff, max_neighbors=max_neighbors) for a in structures ] - parts = [p for p in parts if p.size] - return np.concatenate(parts) if parts else np.zeros(0) + parts = [p for p in parts if p.numel()] + if not parts: + return torch.zeros(0) + return torch.cat(parts) -def _density(angles: np.ndarray, bins: int) -> np.ndarray: - counts, _ = np.histogram(angles, bins=bins, range=(0.0, 180.0)) +def _density(angles: torch.Tensor, bins: int) -> torch.Tensor: + """Normalised histogram over [0, 180] degrees.""" + x = torch.as_tensor(angles, dtype=torch.float64).flatten() + if x.numel() == 0: + return torch.zeros(bins, dtype=torch.float64) + # Bucket by index rather than torch.histc so that the closed right edge + # (a perfectly straight 180-degree angle) lands in the last bin. + idx = (x.clamp(0.0, 180.0) * (bins / 180.0)).long().clamp(0, bins - 1) + counts = torch.bincount(idx, minlength=bins).to(torch.float64) total = counts.sum() - return counts / total if total else counts.astype(float) + return counts / total if total > 0 else counts def compare_angle_distributions( - generated: np.ndarray, - reference: np.ndarray, + generated: torch.Tensor, + reference: torch.Tensor, bins: int = DEFAULT_BINS, eps: float = 1e-9, ) -> Dict[str, float]: """Distances between two pooled bond-angle distributions. - ``wasserstein`` is exact for a 1-D histogram (the integral of the absolute - CDF difference) and is reported in degrees, so it reads directly as "the + ``wasserstein_deg`` is exact for a 1-D histogram (the integral of the + absolute CDF difference) and is in degrees, so it reads directly as "the generated angles are off by this much on average". """ - p = _density(np.asarray(generated), bins) - q = _density(np.asarray(reference), bins) + p = _density(generated, bins) + q = _density(reference, bins) ps, qs = p + eps, q + eps ps, qs = ps / ps.sum(), qs / qs.sum() - kl = float((ps * np.log(ps / qs)).sum()) + kl = float((ps * (ps / qs).log()).sum()) m = 0.5 * (ps + qs) js = float( - 0.5 * (ps * np.log(ps / m)).sum() + 0.5 * (qs * np.log(qs / m)).sum() + 0.5 * (ps * (ps / m).log()).sum() + 0.5 * (qs * (qs / m).log()).sum() ) width = 180.0 / bins - emd = float(np.abs(np.cumsum(p) - np.cumsum(q)).sum() * width) + emd = float((p.cumsum(0) - q.cumsum(0)).abs().sum() * width) return { "kl": kl, "js": js, "wasserstein_deg": emd, - "n_generated": int(np.asarray(generated).size), - "n_reference": int(np.asarray(reference).size), + "n_generated": int(torch.as_tensor(generated).numel()), + "n_reference": int(torch.as_tensor(reference).numel()), "bins": bins, } -def _min_image_displacement(before, after) -> np.ndarray: +def _min_image_displacement(before, after) -> torch.Tensor: """Cartesian displacement per atom, minimum-image and drift-corrected. A relaxation is free to translate the whole cell, and the benchmark's own metrics quotient that out, so the mean displacement is removed before the RMSD is taken. """ - f0 = np.asarray(before.frac_coords, dtype=float) - f1 = np.asarray(after.frac_coords, dtype=float) + f0 = torch.tensor(before.frac_coords, dtype=torch.float64) + f1 = torch.tensor(after.frac_coords, dtype=torch.float64) df = f1 - f0 - df -= np.round(df) - df -= df.mean(axis=0, keepdims=True) - df -= np.round(df) - return df @ np.asarray(after.lattice_mat, dtype=float) + df = df - df.round() + df = df - df.mean(dim=0, keepdim=True) + df = df - df.round() + return df @ torch.tensor(after.lattice_mat, dtype=torch.float64) def relaxation_displacement( @@ -176,12 +217,22 @@ def relaxation_displacement( ) -> Dict[str, float]: """How far one generated structure moved to reach its local minimum.""" d = _min_image_displacement(before, after) - norms = np.linalg.norm(d, axis=1) - v0 = float(abs(np.linalg.det(np.asarray(before.lattice_mat, dtype=float)))) - v1 = float(abs(np.linalg.det(np.asarray(after.lattice_mat, dtype=float)))) + norms = d.norm(dim=-1) + v0 = float( + torch.linalg.det( + torch.tensor(before.lattice_mat, dtype=torch.float64) + ).abs() + ) + v1 = float( + torch.linalg.det( + torch.tensor(after.lattice_mat, dtype=torch.float64) + ).abs() + ) out = { - "rmsd_angstrom": float(np.sqrt((norms**2).mean())), - "max_displacement_angstrom": float(norms.max()) if norms.size else 0.0, + "rmsd_angstrom": float(norms.pow(2).mean().sqrt()), + "max_displacement_angstrom": ( + float(norms.max()) if norms.numel() else 0.0 + ), "volume_change_frac": (v1 - v0) / v0 if v0 else float("nan"), } if energy_before is not None and energy_after is not None: diff --git a/alignn/tests/test_inverse_angle_diffusion.py b/alignn/tests/test_inverse_angle_diffusion.py index 3688eb4..49b96f9 100644 --- a/alignn/tests/test_inverse_angle_diffusion.py +++ b/alignn/tests/test_inverse_angle_diffusion.py @@ -475,6 +475,44 @@ def test_angle_loss_gradient_reaches_the_shared_backbone(): assert any(n.startswith("edge_embedding") for n in touched) +def test_evaluation_angles_match_a_known_crystal(): + """fcc has exactly 60, 90, 120 and 180 degree bond angles.""" + from jarvis.core.atoms import Atoms + + from alignn.inverse.evaluate import ( + bond_angles_deg, + compare_angle_distributions, + ) + + fcc = Atoms( + lattice_mat=[[4.05, 0, 0], [0, 4.05, 0], [0, 0, 4.05]], + coords=[ + [0.0, 0.0, 0.0], + [0.0, 0.5, 0.5], + [0.5, 0.0, 0.5], + [0.5, 0.5, 0.0], + ], + elements=["Al"] * 4, + cartesian=False, + ) + angles = bond_angles_deg(fcc) + assert angles.numel() > 0 + assert sorted({round(float(v), 1) for v in angles}) == [ + 60.0, + 90.0, + 120.0, + 180.0, + ] + # A distribution is at zero distance from itself, and the Wasserstein + # distance is calibrated in degrees. + same = compare_angle_distributions(angles, angles) + assert same["wasserstein_deg"] == pytest.approx(0.0, abs=1e-9) + assert same["js"] == pytest.approx(0.0, abs=1e-12) + shifted = torch.linspace(80.0, 120.0, 5000) + moved = compare_angle_distributions(shifted + 5.0, shifted) + assert moved["wasserstein_deg"] == pytest.approx(5.0, abs=0.05) + + def test_bond_angle_matches_a_known_geometry(): r_ij = torch.tensor([[1.0, 0.0, 0.0]]) r_jk = torch.tensor([[1.0, 0.0, 0.0]]) diff --git a/scripts/atombench/angle_eval.py b/scripts/atombench/angle_eval.py index de3c8f6..e8fe991 100755 --- a/scripts/atombench/angle_eval.py +++ b/scripts/atombench/angle_eval.py @@ -28,7 +28,7 @@ import json from pathlib import Path -import numpy as np +import torch from alignn.inverse.evaluate import ( DEFAULT_ANGLE_CUTOFF, @@ -91,10 +91,10 @@ def main(): "angle_cutoff": args.cutoff, "angle_max_neighbors": args.max_neighbors, "generated_angle_mean_deg": ( - float(np.mean(gen_angles)) if gen_angles.size else None + float(gen_angles.mean()) if gen_angles.numel() else None ), "reference_angle_mean_deg": ( - float(np.mean(ref_angles)) if ref_angles.size else None + float(ref_angles.mean()) if ref_angles.numel() else None ), } d = metrics["angle_distribution"] @@ -128,10 +128,15 @@ def main(): ) if rows: keys = sorted({k for r in rows for k in r}) - summary = { - k: float(np.nanmean([r[k] for r in rows if k in r])) - for k in keys - } + summary = {} + for k in keys: + vals = torch.tensor( + [r[k] for r in rows if k in r], dtype=torch.float64 + ) + vals = vals[torch.isfinite(vals)] + summary[k] = ( + float(vals.mean()) if vals.numel() else float("nan") + ) summary["n_relaxed"] = len(rows) metrics["relaxation"] = summary print( From 7a4923289dd9ca5ef8c3198d2a19ac63b7dfa092 Mon Sep 17 00:00:00 2001 From: crhysc Date: Sat, 29 Aug 2026 18:33:07 -0400 Subject: [PATCH 3/4] task_runners: one runner per published inverse-design result The pieces were all in scripts/atombench already; what was missing was anything tying them to a number in the paper. No multi-seed orchestration, no aggregation with error bars, and no SLURM scripts anywhere in the repo. Twelve tasks, each pinning the arguments to the existing scripts rather than reimplementing them. A task is a list of units (one SLURM array element each) and a unit is a list of stages: train -> generate -> symmetrize -> score-nosym -> score-sym. - Resumability is command-hashed, not file-existence based. Every finished stage writes /.stages/.json with the exact argv, so bumping --epochs retrains and rescores while re-submitting after a walltime kill picks up where it stopped. A failed stage writes no marker. - Stages declare their inputs, so a missing prerequisite is reported in a second ("BLOCKED, missing .../pretrain_dft3d/best_model.pt (produced by pretrain)") instead of crashing an hour in. - Configurations shared between tasks are keyed by run directory and trained once: bench-jarvis, arm A of ablation-linegraph and A0 of angle-ablation are the same model. - Both the unsymmetrised and symmetrised predictions are scored, because the lattice columns are measured after symmetrisation and the force-field paragraph is quoted before it. Each task declares which one --aggregate reads by default. Coverage is checked, not asserted. claims.py registers all 48 numbers the inverse-design section prints -- both tables, the force-field paragraph, the leakage fractions, the split sizes, the parameter match and the 2.4x step cost -- against the task that regenerates each. `run_task.py verify` is a to-do list in dependency order before you have run anything and a measured-vs-published table afterwards, and exits non-zero if a claim ever appears with no runner behind it. Parameter counts, training wall time, the seed spread and the best individual run are derived from artefacts the runner already writes, so the caption claims are measurable too. Three cost settings, each in its own run tree so a cheap run can never overwrite an expensive checkpoint: full (3000 epochs, 32 candidates, 3 seeds), --quick (300 / 8 / 2, whole test split, arms comparable to each other but not to the published numbers), and --smoke (2 epochs, 4 targets, plumbing only). --loss-only trains and stops, which needs no scoring environment and no AtomBench clone; the denoising loss is the most reproducible arm-vs-arm signal and a filter rather than a verdict, since the line-graph loss gap reproduced across machines while match rate did not move at all. SLURM: cluster.env is the only site-specific file, and submit.sh sizes the array from the arguments actually passed (--seeds 0,1 submits two elements) and queues the aggregation with --dependency=afterok. The #SBATCH walltimes are labelled placeholders; they have not been measured anywhere. Two inconsistencies in the source material are flagged rather than papered over. Table 4's "best single run" row pairs match 0.524 with RMSD 0.023, which in the model registry belong to two different checkpoints, so --aggregate names the run it is quoting. And the leakage paragraph says the quoted results use no pretraining while the released csp_supercon_alex behind the 0.485 row was fine-tuned from csp_pretrain_dft3d, so bench-alex gained --from-scratch and pretrain-transfer runs both arms. Table 3's RMSD spreads are the ones the inverse-design README calls a small-sample artifact, so those tolerances are wide and the measured spread is printed. EPOCHS["alex"] and EPOCHS["pretrain"] are not pinned by the manuscript and say so; inspect_checkpoint.py reads the argument namespace out of a released checkpoint to replace them with the published values. data-jarvis was run end to end and reproduces the 847/105/103 split. run_task.py doctor takes a real optimiser step rather than only importing torch, which catches a mismatched install before a job is queued instead of an hour into one. Claude-Session: https://claude.ai/code/session_01RvR9sYTGsrRPBoUC41Vqir --- .gitignore | 5 + task_runners/INSTRUCTIONS.md | 251 +++++ task_runners/README.md | 244 +++++ task_runners/__init__.py | 1 + task_runners/aggregate.py | 567 +++++++++++ task_runners/claims.py | 493 +++++++++ task_runners/cluster.env | 36 + task_runners/common.sh | 43 + task_runners/inspect_checkpoint.py | 167 ++++ task_runners/logs/.gitkeep | 0 task_runners/run_task.py | 617 ++++++++++++ task_runners/sbatch/ablation-linegraph.sbatch | 38 + task_runners/sbatch/aggregate.sbatch | 23 + task_runners/sbatch/angle-ablation.sbatch | 38 + task_runners/sbatch/bench-alex.sbatch | 38 + task_runners/sbatch/bench-jarvis.sbatch | 38 + task_runners/sbatch/data-alex.sbatch | 36 + task_runners/sbatch/data-jarvis.sbatch | 36 + task_runners/sbatch/data-pretrain.sbatch | 37 + task_runners/sbatch/leakage.sbatch | 37 + task_runners/sbatch/pipeline-ablation.sbatch | 38 + task_runners/sbatch/pretrain-transfer.sbatch | 38 + task_runners/sbatch/pretrain.sbatch | 38 + task_runners/sbatch/symprec-sweep.sbatch | 38 + task_runners/submit.sh | 101 ++ task_runners/tasks.py | 945 ++++++++++++++++++ 26 files changed, 3943 insertions(+) create mode 100644 task_runners/INSTRUCTIONS.md create mode 100644 task_runners/README.md create mode 100644 task_runners/__init__.py create mode 100644 task_runners/aggregate.py create mode 100644 task_runners/claims.py create mode 100644 task_runners/cluster.env create mode 100644 task_runners/common.sh create mode 100755 task_runners/inspect_checkpoint.py create mode 100644 task_runners/logs/.gitkeep create mode 100755 task_runners/run_task.py create mode 100755 task_runners/sbatch/ablation-linegraph.sbatch create mode 100755 task_runners/sbatch/aggregate.sbatch create mode 100755 task_runners/sbatch/angle-ablation.sbatch create mode 100755 task_runners/sbatch/bench-alex.sbatch create mode 100755 task_runners/sbatch/bench-jarvis.sbatch create mode 100755 task_runners/sbatch/data-alex.sbatch create mode 100755 task_runners/sbatch/data-jarvis.sbatch create mode 100755 task_runners/sbatch/data-pretrain.sbatch create mode 100755 task_runners/sbatch/leakage.sbatch create mode 100755 task_runners/sbatch/pipeline-ablation.sbatch create mode 100755 task_runners/sbatch/pretrain-transfer.sbatch create mode 100755 task_runners/sbatch/pretrain.sbatch create mode 100755 task_runners/sbatch/symprec-sweep.sbatch create mode 100755 task_runners/submit.sh create mode 100644 task_runners/tasks.py diff --git a/.gitignore b/.gitignore index c3d5e02..f50630c 100644 --- a/.gitignore +++ b/.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 diff --git a/task_runners/INSTRUCTIONS.md b/task_runners/INSTRUCTIONS.md new file mode 100644 index 0000000..6257074 --- /dev/null +++ b/task_runners/INSTRUCTIONS.md @@ -0,0 +1,251 @@ +# INSTRUCTIONS — reproducing the paper's inverse-design results + +**Audience:** whoever (or whatever) is sitting at a terminal on the machine +that will do the work, with no memory of the conversation that produced this +directory. Everything needed is below. + +`README.md` next to this file is the reference — what each task is, how the +runner works, what the output tree looks like. This file is the procedure. + +--- + +## 0. What you are reproducing + +The "Generative inverse design" section of the manuscript. Every number it +prints — both tables, the force-field paragraph, the leakage caveat, the split +sizes — is registered in `claims.py` and mapped to a task. To see that map, +and how much of it you have measured so far: + +```bash +python task_runners/run_task.py verify +``` + +Run that first and last. First it tells you which tasks you still owe; last it +tells you whether what you measured agrees with what was published. It never +writes anything. + +**Never edit a published number to match a measurement, and never report a +measurement you did not take.** If a claim will not reproduce, say so, say by +how much, and say what you think went wrong. `verify` prints the measured +spread beside the published value precisely so that a disagreement can be +argued about rather than papered over. + +--- + +## 1. Set the machine up + +```bash +cd # the directory containing task_runners/ +pip install -e . # re-run even if you installed before: + # the editable finder does not see + # alignn/inverse if it predates it +python task_runners/run_task.py doctor +``` + +`doctor` must show `[x]` on every line before you queue anything. What each +failure means: + +| line | fix | +|---|---| +| `torch` | a broken or CPU-only install. The check takes a real optimiser step, so an import that "works" can still fail here. | +| `alignn.inverse` | `pip install -e .` from the repo root | +| `jarvis-tools`, `pymatgen` | `pip install jarvis-tools pymatgen` | +| `average-minimum-distance` | `pip install average-minimum-distance` — needed for ccRMSD only | +| `AtomBench compute_metrics.py` | `git clone https://github.com/atomgptlab/atombench` and `export ATOMBENCH_REPO=` | + +Scoring dependencies may live in a separate environment; set `CSP_SCORE_ENV` +in `cluster.env` and `score.sh` will switch into it by itself. + +On a cluster, fill in `task_runners/cluster.env` — it is the only +site-specific file. Account, partition, QoS, `--gres`, conda environments, +modules, and where the run tree lives (`CSP_RUNS`, point it at scratch). + +--- + +## 2. Prove the plumbing before spending a GPU-week + +```bash +python task_runners/run_task.py bench-jarvis --smoke --device cpu \ + --runs-root /tmp/csp_smoke +``` + +Two epochs, two candidates, four targets. It exercises every stage — training, +sampling, relaxation, symmetrisation, scoring — and so catches a missing +dependency, a broken checkpoint format or an unreachable AtomBench clone in +minutes rather than after a queue wait. The numbers it produces are +meaningless and it writes into its own `train_smoke/` tree, so it cannot +touch a real run. + +--- + +## 3. Run the tasks, in this order + +Each line is safe to re-run: finished stages are skipped, and a stage whose +command changed re-runs. Everything after `data-jarvis` can also be submitted +with `bash task_runners/submit.sh `. + +```bash +# 1. data (CPU. Minutes, plus a one-off ~40 MB JARVIS download) +python task_runners/run_task.py data-jarvis +python task_runners/run_task.py data-alex # needs the DS-A/B pickles +python task_runners/run_task.py data-pretrain + +# 2. the base model (GPU, long) +python task_runners/run_task.py pretrain + +# 3. the two benchmark tables +python task_runners/run_task.py bench-jarvis # 3 seeds -> Table 4 JARVIS +python task_runners/run_task.py bench-alex # -> Table 4 Alexandria + +# 4. the line-graph ablation (Table 3). Arm A is bench-jarvis, already done +python task_runners/run_task.py ablation-linegraph + +# 5. the paragraphs +python task_runners/run_task.py pipeline-ablation # force-field loop +python task_runners/run_task.py leakage # the 18.4% / 15.4% caveat + +# 6. optional: the tolerance used before scoring, chosen on validation +python task_runners/run_task.py symprec-sweep + +# 7. finally +python task_runners/run_task.py verify +``` + +`data-alex` needs `DS-A.pk.bz2` and `DS-B.pk.bz2` from figshare DOI +`10.6084/m9.figshare.31045597`. Put them in `/data/alexandria/` or pass +`--alex-inputs A.pk.bz2 B.pk.bz2`. Nothing else needs a manual download. + +### Reading each table + +```bash +python task_runners/run_task.py ablation-linegraph --aggregate --latex +python task_runners/run_task.py bench-jarvis --aggregate --latex +python task_runners/run_task.py pipeline-ablation --aggregate --variant nosym +``` + +`--aggregate` prints mean ± one standard deviation per arm, every individual +run, parameter counts and training wall time, the declared comparison with a +Welch p-value, and the published baselines where they apply. `--latex` adds a +tabular in the shape of the manuscript's. + +--- + +## 4. Benchmarking an ablation quickly + +This is the common case: you have changed something in the denoiser and want +to know, today, whether it is worth a full run. + +**Fastest useful signal — denoising loss only.** No sampling, no relaxation, +no scoring environment, no AtomBench clone: + +```bash +python task_runners/run_task.py angle-ablation --quick --loss-only +python task_runners/run_task.py angle-ablation --aggregate +``` + +Twelve 300-epoch trainings (six arms × two seeds). The comparison block prints +each pair from `alignn.inverse.ablations.COMPARISONS` with the change, a +p-value and a *within noise* marker, plus the training cost of each arm. + +The loss is the right first metric here, and the reason is in the +inverse-design README: the line-graph loss gap reproduced to three decimals +across two machines while the downstream match rate did not move at all. Loss +tells you how precisely atoms are placed; it does not tell you how often the +right structure is found. So: + +**Then the full pipeline, still quick**, when the loss says something moved: + +```bash +python task_runners/run_task.py angle-ablation --quick +python task_runners/run_task.py angle-ablation --aggregate --latex +``` + +300 epochs, 8 candidates, 2 seeds, but the **whole** test split and the same +scoring code, so the arms are comparable to each other. They are not +comparable to the published numbers — `verify --quick` will say so. + +**Only then the real thing**, on the two arms that survived: + +```bash +python task_runners/run_task.py angle-ablation --seeds 0,1,2 +``` + +Some knobs, when you want a variant that is not in the table: + +```bash +--only-stages train,generate # or --skip-stages score-nosym +--seeds 0,1,2,3,4 # more seeds; a few percent needs them +--epochs 800 # anything explicit overrides --quick +--num-candidates 16 +``` + +A one-off configuration does not need a new entry in `ablations.py`: pass the +denoiser switches straight through, e.g. +`--only-stages train --epochs 300` with a new `--ablation` name added to +`alignn/inverse/ablations.py` if you want it to be a named arm. + +### Cost, roughly + +Relative, since absolute times depend on the machine. One JARVIS training run +at 3000 epochs is the unit. + +| | cost | +|---|---| +| `--smoke` | negligible, minutes | +| `--quick --loss-only`, one arm one seed | ~0.1 | +| `--quick`, one arm one seed | ~0.1 + generation | +| full, one arm one seed | 1 | +| an arm with the angular channel on | ×2.4 per step | +| generation, 32 candidates on 103 targets | dominated by relaxation; use `--relax-workers` | + +`angle-ablation` at full size is 18 runs. Decide with `--quick` first. + +--- + +## 5. When something disagrees + +Work through these before concluding the model changed: + +1. **Is it the seed?** Across independently trained models the JARVIS match + rate spanned 0.437–0.524 — nine structures out of 103. `--aggregate` prints every + individual run; look at the spread before believing a mean. +2. **Is it symmetrisation?** The lattice-angle and KLD columns are measured + after Niggli reduction and move a lot with the tolerance (angle MAE + 15.9 → 8.4 at `--symprec 0.1`). Compare `--variant sym` against + `--variant nosym`, and pick the tolerance with `symprec-sweep` on + *validation*, never on test. +3. **Is it the pipeline, not the model?** `pipeline-ablation` separates what + the generator contributes from what candidate selection and the force + field contribute. One sample scores 0.22; the full pipeline scores 0.52. +4. **Is it the epochs?** `EPOCHS["alex"]` and `EPOCHS["pretrain"]` in + `tasks.py` are *not* pinned by the manuscript. Recover the published + values with + `python task_runners/inspect_checkpoint.py csp_supercon_alex` and set them + before blaming anything else. +5. **Is it leakage?** If you fine-tuned from `csp_pretrain_dft3d`, some test + targets are reachable by recall. Run `leakage` and compare the filtered + score. + +Two known inconsistencies in the source material, so you are not surprised by +them: + +- The Table 4 "best single run" row pairs match 0.524 with RMSD 0.023. In the + released model registry those belong to two different checkpoints + (`csp_supercon_jarvis`: 0.524 / 0.056, `csp_supercon_jarvis_pt`: 0.515 / + 0.023). `--aggregate` names which run it is quoting, so you will see this. +- The leakage paragraph says the quoted ALIGNN-CSP results use no + pretraining, but the released `csp_supercon_alex` behind the 0.485 + Alexandria row was fine-tuned from `csp_pretrain_dft3d`. `bench-alex` + defaults to the pretrained arm; `--from-scratch` trains the other one, and + `pretrain-transfer` runs both. + +--- + +## 6. What to hand back + +- the `verify` table, unedited; +- `--aggregate --latex` output for each table you regenerated; +- a note of any claim marked `!!`, with your reading of why; +- the run tree, or at least each run's `config.json`, `history.json`, + `.stages/*.json` (which record the exact command, host and git revision) and + `bench/*/metrics.json`. diff --git a/task_runners/README.md b/task_runners/README.md new file mode 100644 index 0000000..706ed7b --- /dev/null +++ b/task_runners/README.md @@ -0,0 +1,244 @@ +# Task runners — generative inverse design + +One command per result in the manuscript's "Generative inverse design" +section. Each task pins the arguments to the scripts that already live in +`scripts/atombench/` and `alignn/inverse/`, runs them in order, resumes where +it stopped, and prints the table with error bars over seeds. + +**Step-by-step procedure: `INSTRUCTIONS.md`.** This file is the reference. + +```bash +python task_runners/run_task.py tasks # what is available +python task_runners/run_task.py doctor # is this machine ready +python task_runners/run_task.py verify # paper number -> task -> measured +python task_runners/run_task.py data-jarvis # build the split +python task_runners/run_task.py bench-jarvis --smoke --device cpu # minutes +python task_runners/run_task.py bench-jarvis # the real thing +python task_runners/run_task.py bench-jarvis --aggregate --latex +``` + +On a cluster, the same tasks go through SLURM as a job array over their units, +plus a dependent job that prints the table when the array succeeds: + +```bash +$EDITOR task_runners/cluster.env # account, partition, conda env, scratch +bash task_runners/submit.sh bench-jarvis --dry-run +bash task_runners/submit.sh bench-jarvis +``` + +## The tasks + +| task | reproduces | units | needs | +|---|---|---|---| +| `data-jarvis` | the JARVIS Supercon-3D split, 847/105/103 | 1 | — | +| `data-alex` | the Alexandria DS-A/B split, 6603/825/825 | 1 | the DS-A/B pickles | +| `data-pretrain` | the 65k dft_3d corpus, benchmark ids held out | 1 | `data-jarvis` | +| `pretrain` | `csp_pretrain_dft3d`, the base model | 1 | `data-pretrain` | +| `bench-jarvis` | Table `tab:inverse_bench`, JARVIS block | 3 seeds | `data-jarvis` | +| `bench-alex` | Table `tab:inverse_bench`, Alexandria block | 1 | `data-alex`, `pretrain` | +| `pretrain-transfer` | Alexandria from scratch vs fine-tuned | 2 | `data-alex`, `pretrain` | +| `ablation-linegraph` | Table `tab:inverse_ablation` | 2 arms × 3 seeds | `data-jarvis` | +| `angle-ablation` | the A0–A6 angular-diffusion suite | 6 arms × 3 seeds | `data-jarvis` | +| `pipeline-ablation` | "Closing the loop with the force field" | 4 | `bench-jarvis` | +| `symprec-sweep` | the symmetrisation tolerance, chosen on val | 1 | `bench-jarvis` | +| `leakage` | the 18.4% / 15.4% recall caveat | 2 | `data-pretrain`, `bench-jarvis` | + +Configurations shared between tasks are keyed by their run directory and +therefore trained **once**. `bench-jarvis`, arm A of `ablation-linegraph` and +`A0` of `angle-ablation` are the same model; running all three costs one set +of trainings, and whichever runs second finds the first one's work and skips +straight to what is missing. + +## How a task is put together + +A task is a list of **units** — independent pieces of work, one per SLURM +array element — and a unit is a list of **stages**: + +``` +train -> generate -> symmetrize -> score-nosym -> score-sym +``` + +Every finished stage writes `/.stages/.json` recording the +exact command it ran. On a re-run a stage is skipped if that command is +unchanged, and re-runs if it is not — so bumping `--epochs` retrains and +rescores, while re-submitting after a walltime kill picks up where it left +off. `--force` ignores the markers. + +Stages also declare their inputs, so a missing prerequisite is reported +immediately: + +``` +[alex-seed0/train] BLOCKED, missing input(s): + .../runs/train/pretrain_dft3d/seed0/best_model.pt (produced by pretrain) +``` + +Everything lands under one root, `--runs-root` (env `ALIGNN_RUNS`, default +`/runs`, `CSP_RUNS` in `cluster.env`): + +``` +runs/ +├── data/{jarvis,alex,pretrain}/ train.json val.json test.json +├── train//seed/ best_model.pt history.json config.json +│ └── bench/{nosym,sym}/ pred.csv metrics.json candidates.json +├── pipeline/{raw,rank,relax,full}/ +├── symprec/ +└── leakage/{jarvis,alex}/ +``` + +Both the unsymmetrised and the symmetrised predictions are scored, because +they answer different questions. Symmetrisation snaps each predicted cell onto +its detected space group, which matters a great deal for the two metrics +measured after Niggli reduction — angle MAE 15.9 → 8.4, KLD 0.030 → 0.018 — +and not at all for match rate. The manuscript's lattice columns are the +symmetrised ones (`--variant sym`, the default); the pipeline ablation is +easier to read before symmetrisation (`--variant nosym`). + +## Reading the results + +```bash +python task_runners/run_task.py ablation-linegraph --aggregate --latex +``` + +prints mean ± one standard deviation per arm, every individual run, the +declared comparison with a Welch p-value where scipy is installed, and a LaTeX +tabular in the shape of the manuscript table. + +Individual runs are always printed, and the reason is in the manuscript: +across independently trained models the match rate on 103 JARVIS targets +spanned 0.437–0.524 — nine models in the manuscript, fifteen by the time the +inverse-design README was written. A difference of a few percent in match rate +between two single runs is not a result; a difference in coordinate RMSD of +the size reported there is. The comparison output marks a change smaller than +the arms' own spread as *within noise* rather than letting a sign carry the +argument. + +## On a cluster + +`cluster.env` is the only site-specific file. Fill in what your site needs and +leave the rest empty: + +```bash +CSP_ACCOUNT="..." # sbatch --account +CSP_PARTITION="gpu" # sbatch --partition +CSP_GPU_GRES="gpu:a100:1" # sbatch --gres, GPU tasks only +CSP_MAX_CONCURRENT="4" # array throttle +CSP_ENV="alignn2" # conda env with torch + this repo +CSP_SCORE_ENV="atombench" # conda env with pymatgen + average-minimum-distance +CSP_ATOMBENCH_REPO="$HOME/atombench" +CSP_RUNS="/scratch/$USER/alignn_csp" +CSP_MODULES="cuda/12.1" # module load ... +``` + +`submit.sh` reads it, sizes the array from the arguments you actually pass +(`--seeds 0,1,2,3,4` submits five elements, not three), submits, and queues +the aggregation job with `--dependency=afterok`. The `#SBATCH` directives +inside `sbatch/.sbatch` are defaults for the default seeds; anything on +the sbatch command line overrides them, which is how `submit.sh` applies +`cluster.env`. + +**The walltimes in the sbatch headers are placeholders.** They have not been +measured on any particular machine. Check the first array element and adjust. +The one measured cost figure is relative: the angular channel costs 2.4× per +training step, so `angle-ablation`'s A1–A6 arms need more walltime than A0. + +Submit from the repository root — the log paths (`task_runners/logs/`) and +`source task_runners/common.sh` are relative to it. `submit.sh` handles that; +`sbatch` by hand does not. + +## Coverage + +`claims.py` registers all 48 numbers the inverse-design section prints — +both tables, the force-field paragraph, the leakage fractions, the split +sizes, the parameter match and the 2.4× step cost — against the task that +regenerates each one. + +```bash +python task_runners/run_task.py verify +``` + +Before you have run anything it is a to-do list in dependency order. After +you have, it is a measured-vs-published table: `ok` inside the claim's +tolerance, `~` inside one measured standard deviation, `!!` outside both, +blank for not yet run. It reads only what is on disk and writes nothing. If a +claim ever appears with no task behind it, `verify` exits non-zero and calls +it a bug — that is the check that this directory stays complete. + +## Three cost settings + +| | epochs | candidates | seeds | targets | run tree | +|---|---|---|---|---|---| +| full (default) | 3000 | 32 | 3 | all | `runs/train/` | +| `--quick` | 300 | 8 | 2 | all | `runs/train_quick/` | +| `--smoke` | 2 | 2 | 1 | 4 | `runs/train_smoke/` | + +The separate trees matter: a `--quick` run has a different training command +from a full one, so without them it would be detected as stale work and +overwrite a checkpoint that cost days. + +`--quick` keeps everything that makes two arms comparable to each other — the +whole test split, the same pipeline, the same scoring code — and cuts only +the things that scale cost. Its numbers are not comparable to the published +ones, and `verify --quick` says so rather than letting you read across. + +Cheaper still, and the right first question for an ablation: + +```bash +python task_runners/run_task.py angle-ablation --quick --loss-only +python task_runners/run_task.py angle-ablation --aggregate +``` + +`--loss-only` runs the training stage and stops. It needs no scoring +environment and no AtomBench clone, and the denoising validation loss is the +most reproducible arm-vs-arm signal there is — the line-graph loss gap +repeated to three decimals on a second machine while the match rate did not +move at all. It also will not tell you whether the right structure is found +more often, which is why it is a filter and not a verdict. + +`--only-stages` and `--skip-stages` take the stage names for anything in +between. + +## Smoke test + +`--smoke` runs every stage of a task at a size that finishes in minutes: two +epochs, two candidates, four targets, one seed. It proves the plumbing — +data layout, checkpoint format, CSV columns, the scoring environment — and +proves nothing at all about the science. + +```bash +python task_runners/run_task.py bench-jarvis --smoke --device cpu +``` + +Smoke runs write into the same tree, so use a throwaway root: +`--runs-root /tmp/csp_smoke`. + +## Two epoch counts are not pinned + +`EPOCHS["jarvis"] = 3000` is the value the inverse-design README records for +the published JARVIS runs. `EPOCHS["alex"]` and `EPOCHS["pretrain"]` are +plausible choices, not the manuscript's. To replace them with what the +released checkpoints actually used: + +```bash +python task_runners/inspect_checkpoint.py csp_supercon_alex +python task_runners/inspect_checkpoint.py csp_pretrain_dft3d --command +``` + +That reads the argument namespace `train_csp.py` stores in every checkpoint, +so it also works on your own runs (`runs/train/jarvis_A0/seed0`). + +## Requirements + +Training and generation need torch, jarvis-tools and this repository +installed (`pip install -e .` — re-run it if your editable install predates +`alignn/inverse`). Scoring runs AtomBench's own metric code and additionally +needs pymatgen and `average-minimum-distance`, plus a clone of +[atombench](https://github.com/atomgptlab/atombench) pointed at by +`ATOMBENCH_REPO`. If those live in a separate environment, set +`CSP_SCORE_ENV`; `score.sh` switches into it by itself. + +`python task_runners/run_task.py doctor` checks all of it before you queue +anything. + +The Alexandria pickles (`DS-A.pk.bz2`, `DS-B.pk.bz2`) are not downloadable +from here — fetch them from figshare DOI `10.6084/m9.figshare.31045597` and +either drop them in `runs/data/alexandria/` or pass `--alex-inputs`. diff --git a/task_runners/__init__.py b/task_runners/__init__.py new file mode 100644 index 0000000..ab4d746 --- /dev/null +++ b/task_runners/__init__.py @@ -0,0 +1 @@ +"""Reproducible runners for the manuscript's inverse-design experiments.""" diff --git a/task_runners/aggregate.py b/task_runners/aggregate.py new file mode 100644 index 0000000..b9aef10 --- /dev/null +++ b/task_runners/aggregate.py @@ -0,0 +1,567 @@ +"""Summarise a task's runs: mean +/- sd over seeds, and the LaTeX table. + +Seeds are the whole point. Across fifteen independently trained models the +match rate on the 103 JARVIS targets spanned 0.437-0.524, so a single run is +not evidence for a difference of a few percent, and this module refuses to +present one as though it were: every group is reported with its spread and +its n, and the paired comparisons print the change with both arms' spreads +next to it. +""" + +from __future__ import annotations + +import json +import math +import statistics +import sys +from pathlib import Path +from typing import Dict, List, Optional, Sequence + +from task_runners import tasks as T + +# AtomBench's own extraction of metrics.json, and the published baselines, +# rather than a second implementation that could drift from it. +sys.path.insert(0, str(T.REPO / "scripts" / "atombench")) +from collect_results import BASELINES, extract # noqa: E402 + +#: label, key, decimals, lower-is-better +COLUMNS = [ + ("loss", "loss", 3, True), + ("match", "match", 4, False), + ("RMSD", "rmsd", 3, True), + ("ccRMSD", "ccrmsd", 3, True), + ("MAE abc", "abc", 3, True), + ("MAE ang", "ang", 2, True), + ("KLD", "kld", 4, True), +] + +LATEX_HEADERS = { + "loss": r"Denoising loss $\downarrow$", + "match": r"Match rate $\uparrow$", + "rmsd": r"Coordinate RMSD (\AA) $\downarrow$", + "ccrmsd": r"ccRMSD $\downarrow$", + "abc": r"Lattice MAE, $abc$ (\AA) $\downarrow$", + "ang": r"Lattice MAE, angles ($^{\circ}$) $\downarrow$", + "kld": r"KLD $\downarrow$", +} + + +def best_val_loss(history: Optional[Path]) -> Optional[float]: + """Lowest validation denoising loss recorded during training.""" + if history is None or not history.exists(): + return None + try: + rows = json.loads(history.read_text()) + except json.JSONDecodeError: + return None + losses = [r["val"]["loss"] for r in rows if "val" in r] + return min(losses) if losses else None + + +def unit_facts(unit: T.Unit) -> Dict: + """Facts a metrics.json does not carry, read from the run directory. + + Parameter count and wall time make two of the manuscript's claims + checkable -- "matched to within 1% on parameters" and "2.4x in time per + training step" -- from artefacts the runner already writes. + """ + facts: Dict = {} + cfg = unit.rundir / "config.json" + if cfg.exists(): + try: + facts["params"] = json.loads(cfg.read_text()).get("n_parameters") + except json.JSONDecodeError: + pass + marker = unit.rundir / ".stages" / "train.json" + if marker.exists(): + try: + facts["train_s"] = json.loads(marker.read_text()).get("elapsed_s") + except json.JSONDecodeError: + pass + for name, keys in ( + ("split_meta.json", ("n_train", "n_val", "n_test")), + ("leakage.json", ("fraction", "n_leaked", "n_test")), + ): + path = unit.rundir / name + if not path.exists(): + continue + try: + data = json.loads(path.read_text()) + except json.JSONDecodeError: + continue + for key in keys: + if key in data: + facts[key] = data[key] + if "fraction" in data: + facts["leak_fraction"] = data["fraction"] + return facts + + +def collect(units: Sequence[T.Unit], variant: str) -> Dict[str, List[Dict]]: + """Group the units' results by aggregation label. + + A unit whose metrics are not keyed by ``variant`` (the tolerance sweep, + the leakage filter) contributes one row per key it does have, so those + tasks read as several groups of one rather than needing a special case. + """ + groups: Dict[str, List[Dict]] = {} + for unit in units: + loss = best_val_loss(unit.history) + if variant in unit.metrics: + keyed = {unit.group or unit.name: unit.metrics[variant]} + elif unit.metrics: + prefix = f"{unit.group}:" if len(units) > 1 else "" + keyed = {f"{prefix}{k}": v for k, v in unit.metrics.items()} + else: + keyed = {unit.group or unit.name: None} + facts = unit_facts(unit) + for label, path in keyed.items(): + row = {"unit": unit.name, "seed": unit.seed, "loss": loss} + row.update(facts) + if path is not None and Path(path).exists(): + row.update(extract(Path(path))) + elif path is not None: + row["missing"] = str(path) + groups.setdefault(label, []).append(row) + return groups + + +def stat(rows: Sequence[Dict], key: str): + """(mean, sd, n) over the rows that actually have this metric.""" + vals = [ + r[key] for r in rows if r.get(key) is not None and not _nan(r.get(key)) + ] + if not vals: + return None, None, 0 + if len(vals) == 1: + return vals[0], None, 1 + return statistics.fmean(vals), statistics.stdev(vals), len(vals) + + +def _nan(x) -> bool: + try: + return math.isnan(float(x)) + except (TypeError, ValueError): + return True + + +def cell(mean, sd, dp: int) -> str: + if mean is None: + return "-" + if sd is None: + return f"{mean:.{dp}f}" + return f"{mean:.{dp}f}+-{sd:.{dp}f}" + + +def welch_p(a: Sequence[float], b: Sequence[float]) -> Optional[float]: + """Two-sided Welch p-value, if scipy is around to give one.""" + if len(a) < 2 or len(b) < 2: + return None + try: + from scipy import stats + except ImportError: + return None + return float(stats.ttest_ind(a, b, equal_var=False).pvalue) + + +def values(rows: Sequence[Dict], key: str) -> List[float]: + return [ + r[key] for r in rows if r.get(key) is not None and not _nan(r.get(key)) + ] + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- + + +def print_table(groups: Dict[str, List[Dict]]) -> None: + label_w = max([len(k) for k in groups] + [12]) + 2 + header = ( + "group".ljust(label_w) + + "".join(f"{lab:>17}" for lab, _, _, _ in COLUMNS) + + f"{'n':>4}" + ) + print(header) + print("-" * len(header)) + for label, rows in groups.items(): + line = label.ljust(label_w) + n_max = 0 + for _, key, dp, _ in COLUMNS: + mean, sd, n = stat(rows, key) + n_max = max(n_max, n) + line += f"{cell(mean, sd, dp):>17}" + print(line + f"{n_max:>4}") + + +def print_notes(units: Sequence[T.Unit]) -> None: + """Headline numbers a metrics.json does not carry. + + Currently just the leakage fraction: how many test targets have a + structure-matcher-identical counterpart in the pretraining corpus, and so + are reachable by recall rather than by generation. + """ + for unit in units: + report = unit.rundir / "leakage.json" + if not report.exists(): + continue + try: + data = json.loads(report.read_text()) + except json.JSONDecodeError: + continue + print( + f" {unit.name}: {data['n_leaked']}/{data['n_test']} test targets" + f" ({data['fraction']:.1%}) are reachable by recall" + ) + print() + + +def print_legend(task: T.Task, groups: Dict[str, List[Dict]]) -> None: + """Spell out the short group labels, so the table can stay narrow.""" + entries = [(k, v) for k, v in task.legend.items() if k in groups] + if not entries: + return + width = max(len(k) for k, _ in entries) + print() + for label, text in entries: + print(f" {label.ljust(width)} {text}") + + +def print_runs(groups: Dict[str, List[Dict]]) -> None: + """Per-run rows: an outlier seed should be visible, not averaged away.""" + print("\nindividual runs") + for label, rows in groups.items(): + for row in rows: + if row.get("missing"): + print(f" {label:<34} {row['unit']:<24} not scored yet") + continue + bits = [] + for lab, key, dp, _ in COLUMNS: + val = row.get(key) + bits.append( + f"{lab} {val:.{dp}f}" + if val is not None and not _nan(val) + else f"{lab} -" + ) + print(f" {label:<34} {row['unit']:<24} " + " ".join(bits)) + + +def _hms(seconds: float) -> str: + hours, rest = divmod(int(seconds), 3600) + return f"{hours}h{rest // 60:02d}m" if hours else f"{rest // 60}m" + + +def print_cost(groups: Dict[str, List[Dict]]) -> None: + """Parameters and training wall time per arm. + + The time ratio equals the per-step ratio only when the arms ran the same + number of epochs on the same hardware, which every task here arranges and + a job array does not guarantee -- so it is labelled as wall time, not as + cost per step. + """ + rows = [ + (label, stat(r, "params"), stat(r, "train_s")) + for label, r in groups.items() + ] + rows = [r for r in rows if r[1][0] is not None or r[2][0] is not None] + if not rows: + return + ref = next((r[2][0] for r in rows if r[2][0]), None) + width = max(len(r[0]) for r in rows) + 2 + print("\ntraining cost") + print( + " " + "arm".ljust(width) + f"{'params':>10}{'wall time':>14}" + f"{'x vs first':>12}" + ) + for label, (params, _, _), (secs, sd, _) in rows: + par = "-" if params is None else f"{params / 1e6:.2f} M" + if secs is None: + time_txt, ratio = "-", "-" + else: + time_txt = _hms(secs) + (f" +-{_hms(sd)}" if sd else "") + ratio = f"{secs / ref:.2f}" if ref else "-" + print(f" {label.ljust(width)}{par:>10}{time_txt:>14}{ratio:>12}") + + +def print_comparisons(task: T.Task, groups: Dict[str, List[Dict]]) -> None: + if not task.comparisons: + return + print("\ncomparisons (change of the second arm relative to the first)") + for question, (ref, test) in task.comparisons.items(): + if ref not in groups or test not in groups: + missing = [g for g in (ref, test) if g not in groups] + print(f"\n {question}: not run ({', '.join(missing)})") + continue + print(f"\n {question}") + print(f" {ref} -> {test}") + for lab, key, dp, lower in COLUMNS: + m_ref, s_ref, n_ref = stat(groups[ref], key) + m_test, s_test, n_test = stat(groups[test], key) + if m_ref is None or m_test is None: + continue + change = ( + (m_test - m_ref) / m_ref * 100.0 if m_ref else float("nan") + ) + direction = "better" if (change < 0) == lower else "worse" + if abs(change) < 0.05: + direction = "unchanged" + p = welch_p(values(groups[ref], key), values(groups[test], key)) + # A change smaller than the arms' own spread is not a result. + spread = max(s_ref or 0.0, s_test or 0.0) + noise = ( + " (within noise)" + if spread and abs(m_test - m_ref) < spread + else "" + ) + p_txt = f" p={p:.3f}" if p is not None else "" + print( + f" {lab:<9} {cell(m_ref, s_ref, dp):>17} -> " + f"{cell(m_test, s_test, dp):>17}" + f" {change:+6.1f}% {direction}{p_txt}{noise}" + ) + + +def print_baselines(task: T.Task, groups: Dict[str, List[Dict]]) -> None: + if not task.baselines: + return + base = BASELINES[task.baselines] + print(f"\npublished AtomBench baselines ({task.baselines})") + label_w = max([len(k) for k in base] + [len(k) + 7 for k in groups]) + 2 + cols = [c for c in COLUMNS if c[1] != "loss"] + print( + "model".ljust(label_w) + "".join(f"{lab:>10}" for lab, _, _, _ in cols) + ) + for name, row in base.items(): + line = name.ljust(label_w) + for _, key, dp, _ in cols: + line += f"{row.get(key, float('nan')):>10.{dp}f}" + print(line) + for label, rows in groups.items(): + line = f"{label} (ours)".ljust(label_w) + for _, key, dp, _ in cols: + mean, _, _ = stat(rows, key) + line += ("-" if mean is None else f"{mean:.{dp}f}").rjust(10) + print(line) + # The manuscript also quotes the best individual run; name it here + # rather than letting a reader assume the column-wise best is one run. + scored = [r for r in rows if r.get("match") is not None] + if len(scored) > 1: + best = max(scored, key=lambda r: r["match"]) + line = f" best run ({best['unit']})".ljust(label_w) + for _, key, dp, _ in cols: + val = best.get(key) + line += ( + "-" if val is None or _nan(val) else f"({val:.{dp}f})" + ).rjust(10) + print(line) + + +# --------------------------------------------------------------------------- +# Claim resolution, for `verify` +# --------------------------------------------------------------------------- + + +def resolve(rows: Sequence[Dict], metric: str): + """(value, sd, n) for a metric name, including the derived ones. + + ``match_min`` / ``match_max`` give the seed spread; ``bestrun_`` gives + metric *m* as measured in the run with the highest match rate, which is + what the manuscript's parenthesised "best individual run" row is. + """ + if metric.startswith("bestrun_"): + key = metric[len("bestrun_") :] + scored = [r for r in rows if r.get("match") is not None] + if not scored: + return None, None, 0 + best = max(scored, key=lambda r: r["match"]) + value = best.get(key) + return (None, None, 0) if _nan(value) else (value, None, 1) + if metric in ("match_min", "match_max"): + vals = values(rows, "match") + if not vals: + return None, None, 0 + return ( + (min(vals) if metric.endswith("min") else max(vals)), + None, + len(vals), + ) + return stat(rows, metric) + + +def check( + task: T.Task, + units: Sequence[T.Unit], + claim, +) -> Dict: + """Measure one claim against the runs on disk.""" + groups = collect(units, claim.variant) + rows = groups.get(claim.group or next(iter(groups), ""), []) + value, sd, n = resolve(rows, claim.metric) + if claim.ref_group: + ref_rows = groups.get(claim.ref_group, []) + ref, _, ref_n = resolve(ref_rows, claim.metric) + if value is None or not ref: + value, sd, n = None, None, min(n, ref_n) + else: + value, sd, n = value / ref, None, min(n, ref_n) + + out = {"measured": value, "sd": sd, "n": n, "status": "not run"} + if value is None: + return out + published = claim.published + rel = (value - published) / published if published else float("inf") + out["rel"] = rel + if abs(rel) <= claim.tol: + out["status"] = "ok" + elif sd and abs(value - published) <= sd: + # Inside one standard deviation of the published number is agreement + # at this sample size, whatever the relative gap looks like. + out["status"] = "within sd" + else: + out["status"] = "differs" + return out + + +# --------------------------------------------------------------------------- +# LaTeX +# --------------------------------------------------------------------------- + + +def latex_ablation(task: T.Task, groups: Dict[str, List[Dict]]) -> str: + """Metrics down the side, arms across: the shape of Table 3.""" + labels = list(groups) + lines = [ + r"\begin{tabular}{l" + "c" * (len(labels) + 1) + "}", + r"\hline", + "Metric & " + " & ".join(labels) + r" & Change \\", + r"\hline", + ] + # The Change column is the task's own comparison, so it reads the same + # way as the printed one: the second arm relative to the first, not + # whichever arm happens to be leftmost in the table. + ref, test = next(iter(task.comparisons.values()), (labels[0], labels[-1])) + for lab, key, dp, lower in COLUMNS: + cells = [] + means = {} + for label in labels: + mean, sd, _ = stat(groups[label], key) + means[label] = mean + cells.append("-" if mean is None else _tex_cell(mean, sd, dp)) + if means.get(ref) and means.get(test) is not None: + change = (means[test] - means[ref]) / means[ref] * 100 + change_txt = f"${change:+.0f}\\%$" + else: + change_txt = "-" + lines.append( + f"{LATEX_HEADERS[key]} & " + + " & ".join(cells) + + f" & {change_txt} " + + r"\\" + ) + lines += [r"\hline", r"\end{tabular}"] + return "\n".join(lines) + + +def latex_bench(task: T.Task, groups: Dict[str, List[Dict]]) -> str: + """Models down the side, metrics across: the shape of Table 4.""" + cols = [c for c in COLUMNS if c[1] != "loss"] + lines = [ + r"\begin{tabular}{l" + "c" * len(cols) + "}", + r"\hline", + "Model & " + + " & ".join(LATEX_HEADERS[k] for _, k, _, _ in cols) + + r" \\", + r"\hline", + ] + for name, row in BASELINES[task.baselines].items(): + cells = [f"${row[k]:.{dp}f}$" for _, k, dp, _ in cols] + lines.append(f"{name} & " + " & ".join(cells) + r" \\") + for label, rows in groups.items(): + cells = [] + for _, key, dp, _ in cols: + mean, sd, _ = stat(rows, key) + cells.append("-" if mean is None else _tex_cell(mean, sd, dp)) + lines.append(f"ALIGNN-CSP ({label}) & " + " & ".join(cells) + r" \\") + lines += [r"\hline", r"\end{tabular}"] + return "\n".join(lines) + + +def _tex_cell(mean: float, sd: Optional[float], dp: int) -> str: + if sd is None: + return f"${mean:.{dp}f}$" + return f"${mean:.{dp}f}\\pm{sd:.{dp}f}$" + + +def latex_generic(groups: Dict[str, List[Dict]]) -> str: + lines = [ + r"\begin{tabular}{l" + "c" * len(COLUMNS) + "}", + r"\hline", + "Run & " + + " & ".join(LATEX_HEADERS[k] for _, k, _, _ in COLUMNS) + + r" \\", + r"\hline", + ] + for label, rows in groups.items(): + cells = [] + for _, key, dp, _ in COLUMNS: + mean, sd, _ = stat(rows, key) + cells.append("-" if mean is None else _tex_cell(mean, sd, dp)) + lines.append(f"{label} & " + " & ".join(cells) + r" \\") + lines += [r"\hline", r"\end{tabular}"] + return "\n".join(lines) + + +# --------------------------------------------------------------------------- + + +def report( + task: T.Task, + units: Sequence[T.Unit], + *, + variant: str = "sym", + latex: bool = False, +) -> int: + groups = collect(units, variant) + scored = sum( + 1 + for rows in groups.values() + for r in rows + if r.get("match") is not None or r.get("loss") is not None + ) + print(f"\n{task.name}: {task.summary}") + print(f"reproduces: {task.reproduces}") + print( + f"variant: {variant} groups: {len(groups)} with results: " + f"{scored}\n" + ) + if not any(u.metrics or u.history for u in units): + # A data-preparation task has nothing to average; say so and exit + # clean, so the dependent aggregation job is not a red herring. + print( + "This task produces inputs, not metrics -- nothing to " + "aggregate." + ) + return 0 + if not scored: + print("Nothing scored yet. Run the task first:") + print(f" python task_runners/run_task.py {task.name}") + return 1 + + print_notes(units) + print_table(groups) + print_legend(task, groups) + print_runs(groups) + print_cost(groups) + print_comparisons(task, groups) + print_baselines(task, groups) + + if latex: + print("\n% ---- LaTeX ----") + if task.baselines: + print(latex_bench(task, groups)) + elif len(groups) == 2 and task.comparisons: + print(latex_ablation(task, groups)) + else: + print(latex_generic(groups)) + print() + return 0 diff --git a/task_runners/claims.py b/task_runners/claims.py new file mode 100644 index 0000000..c9d0373 --- /dev/null +++ b/task_runners/claims.py @@ -0,0 +1,493 @@ +"""Every quantitative claim the manuscript makes about inverse design. + +This is the coverage map. Each entry names a number printed in the paper, the +task that regenerates it, and where in that task's output it appears, so +``run_task.py verify`` can answer two different questions: + +* **before running anything** -- is every claim reachable from an executable + in this directory, and which tasks would I have to run? +* **after running** -- does what I measured agree with what was published? + +A claim with no task is a claim this directory cannot reproduce, and there +should not be any; ``verify`` fails loudly if one appears. Tolerances are +generous on purpose: with 103 test targets, a match rate differing by 0.03 is +three structures, and the manuscript itself reports a spread of nine across +seeds. ``verify`` prints the measured spread next to the published value +rather than reducing agreement to a single pass/fail. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List + + +@dataclass(frozen=True) +class Claim: + """One published number, and where to find its measured counterpart.""" + + source: str + statement: str + task: str + metric: str + published: float + group: str = "" + #: When set, ``published`` is metric[group] / metric[ref_group]. + ref_group: str = "" + variant: str = "sym" + #: Relative tolerance for calling the claim reproduced. + tol: float = 0.10 + + +A = "A: line graph" +B = "B: no line graph" +CSP = "ALIGNN-CSP" + +CLAIMS: List[Claim] = [ + # -- splits ------------------------------------------------------------- + Claim( + "text", + "JARVIS Supercon-3D split is 847/105/103", + "data-jarvis", + "n_train", + 847, + tol=0.0, + ), + Claim( + "text", + "JARVIS Supercon-3D split is 847/105/103", + "data-jarvis", + "n_val", + 105, + tol=0.0, + ), + Claim( + "text", + "JARVIS Supercon-3D split is 847/105/103", + "data-jarvis", + "n_test", + 103, + tol=0.0, + ), + Claim( + "text", + "Alexandria DS-A/B split is 6603/825/825", + "data-alex", + "n_train", + 6603, + tol=0.0, + ), + Claim( + "text", + "Alexandria DS-A/B split is 6603/825/825", + "data-alex", + "n_val", + 825, + tol=0.0, + ), + Claim( + "text", + "Alexandria DS-A/B split is 6603/825/825", + "data-alex", + "n_test", + 825, + tol=0.0, + ), + # -- Table 3, tab:inverse_ablation -------------------------------------- + # Only the means are registered. The table's RMSD spreads (0.030+-0.001 + # against 0.048+-0.013) are known not to hold: the inverse-design README + # records six models per arm giving 0.031+-0.012 against 0.044+-0.011 and + # calls the original tightness a small-sample artifact. The RMSD + # tolerances below are wide for that reason, and --aggregate prints the + # measured spread so the point is visible rather than asserted. + Claim( + "Table 3", + "denoising loss, line graph", + "ablation-linegraph", + "loss", + 1.997, + group=A, + tol=0.05, + ), + Claim( + "Table 3", + "denoising loss, no line graph", + "ablation-linegraph", + "loss", + 2.351, + group=B, + tol=0.05, + ), + Claim( + "Table 3", + "coordinate RMSD, line graph", + "ablation-linegraph", + "rmsd", + 0.030, + group=A, + tol=0.35, + ), + Claim( + "Table 3", + "coordinate RMSD, no line graph", + "ablation-linegraph", + "rmsd", + 0.048, + group=B, + tol=0.35, + ), + Claim( + "Table 3", + "ccRMSD, line graph", + "ablation-linegraph", + "ccrmsd", + 0.508, + group=A, + ), + Claim( + "Table 3", + "ccRMSD, no line graph", + "ablation-linegraph", + "ccrmsd", + 0.521, + group=B, + ), + Claim( + "Table 3", + "lattice MAE abc, line graph", + "ablation-linegraph", + "abc", + 0.535, + group=A, + tol=0.15, + ), + Claim( + "Table 3", + "lattice MAE abc, no line graph", + "ablation-linegraph", + "abc", + 0.542, + group=B, + tol=0.15, + ), + Claim( + "Table 3", + "lattice MAE angles, line graph", + "ablation-linegraph", + "ang", + 9.47, + group=A, + tol=0.15, + ), + Claim( + "Table 3", + "lattice MAE angles, no line graph", + "ablation-linegraph", + "ang", + 9.76, + group=B, + tol=0.15, + ), + Claim( + "Table 3", + "match rate, line graph", + "ablation-linegraph", + "match", + 0.4725, + group=A, + tol=0.10, + ), + Claim( + "Table 3", + "match rate, no line graph", + "ablation-linegraph", + "match", + 0.4725, + group=B, + tol=0.10, + ), + Claim( + "Table 3 caption", + "parameters matched within 1%: 3.79 M", + "ablation-linegraph", + "params", + 3.79e6, + group=A, + tol=0.01, + ), + Claim( + "Table 3 caption", + "parameters matched within 1%: 3.75 M", + "ablation-linegraph", + "params", + 3.75e6, + group=B, + tol=0.01, + ), + Claim( + "text", + "angles cost 2.4x in time per training step", + "ablation-linegraph", + "train_s", + 2.4, + group=A, + ref_group=B, + tol=0.30, + ), + # -- Table 4, JARVIS Supercon-3D block ---------------------------------- + Claim( + "Table 4", + "JARVIS match rate, mean of three seeds", + "bench-jarvis", + "match", + 0.473, + group=CSP, + tol=0.10, + ), + Claim( + "Table 4", + "JARVIS coordinate RMSD", + "bench-jarvis", + "rmsd", + 0.030, + group=CSP, + tol=0.35, + ), + Claim( + "Table 4", "JARVIS ccRMSD", "bench-jarvis", "ccrmsd", 0.508, group=CSP + ), + Claim( + "Table 4", + "JARVIS lattice MAE abc", + "bench-jarvis", + "abc", + 0.535, + group=CSP, + tol=0.15, + ), + Claim( + "Table 4", + "JARVIS lattice MAE angles", + "bench-jarvis", + "ang", + 9.47, + group=CSP, + tol=0.15, + ), + Claim( + "Table 4", + "JARVIS KLD", + "bench-jarvis", + "kld", + 0.023, + group=CSP, + tol=0.30, + ), + Claim( + "Table 4", + "JARVIS best single run, match", + "bench-jarvis", + "bestrun_match", + 0.524, + group=CSP, + tol=0.10, + ), + Claim( + "Table 4", + "JARVIS best single run, RMSD", + "bench-jarvis", + "bestrun_rmsd", + 0.023, + group=CSP, + tol=0.40, + ), + Claim( + "Table 4", + "JARVIS best single run, ccRMSD", + "bench-jarvis", + "bestrun_ccrmsd", + 0.470, + group=CSP, + ), + Claim( + "Table 4", + "JARVIS best single run, MAE abc", + "bench-jarvis", + "bestrun_abc", + 0.433, + group=CSP, + tol=0.20, + ), + Claim( + "Table 4", + "JARVIS best single run, MAE angles", + "bench-jarvis", + "bestrun_ang", + 8.37, + group=CSP, + tol=0.20, + ), + Claim( + "Table 4", + "JARVIS best single run, KLD", + "bench-jarvis", + "bestrun_kld", + 0.018, + group=CSP, + tol=0.35, + ), + Claim( + "text", + "match rate across seeds spans 0.437 at the low end", + "bench-jarvis", + "match_min", + 0.437, + group=CSP, + tol=0.10, + ), + Claim( + "text", + "match rate across seeds spans 0.524 at the high end", + "bench-jarvis", + "match_max", + 0.524, + group=CSP, + tol=0.10, + ), + # -- Table 4, Alexandria DS-A/B block ----------------------------------- + Claim( + "Table 4", + "Alexandria match rate", + "bench-alex", + "match", + 0.485, + group=CSP, + tol=0.10, + ), + Claim( + "Table 4", + "Alexandria coordinate RMSD", + "bench-alex", + "rmsd", + 0.028, + group=CSP, + tol=0.35, + ), + Claim( + "Table 4", + "Alexandria ccRMSD", + "bench-alex", + "ccrmsd", + 0.343, + group=CSP, + tol=0.15, + ), + Claim( + "Table 4", + "Alexandria lattice MAE abc", + "bench-alex", + "abc", + 0.561, + group=CSP, + tol=0.15, + ), + Claim( + "Table 4", + "Alexandria lattice MAE angles", + "bench-alex", + "ang", + 10.09, + group=CSP, + tol=0.15, + ), + Claim( + "Table 4", + "Alexandria KLD", + "bench-alex", + "kld", + 0.023, + group=CSP, + tol=0.30, + ), + # -- "Closing the loop with the force field" ---------------------------- + # Quoted before symmetrisation: these are about what sampling and the + # force field contribute, not about the lattice metrics. + Claim( + "text", + "one sample, no selection or relaxation: match 0.22", + "pipeline-ablation", + "match", + 0.22, + group="raw", + variant="nosym", + tol=0.15, + ), + Claim( + "text", + "one sample, no selection or relaxation: RMSD 0.29", + "pipeline-ablation", + "rmsd", + 0.29, + group="raw", + variant="nosym", + tol=0.25, + ), + Claim( + "text", + "relaxation without selection contributes almost nothing", + "pipeline-ablation", + "match", + 0.24, + group="relax", + variant="nosym", + tol=0.15, + ), + Claim( + "text", + "32 candidates ranked and relaxed: match 0.52", + "pipeline-ablation", + "match", + 0.52, + group="full", + variant="nosym", + tol=0.10, + ), + Claim( + "text", + "32 candidates ranked and relaxed: RMSD 0.06", + "pipeline-ablation", + "rmsd", + 0.06, + group="full", + variant="nosym", + tol=0.40, + ), + # -- the leakage caveat ------------------------------------------------- + Claim( + "text", + "18.4% of JARVIS test targets are reachable by recall", + "leakage", + "leak_fraction", + 0.184, + group="jarvis:all", + tol=0.05, + ), + Claim( + "text", + "15.4% of Alexandria test targets are reachable by recall", + "leakage", + "leak_fraction", + 0.154, + group="alex:all", + tol=0.05, + ), +] + + +def tasks_needed() -> List[str]: + """Distinct tasks that have to run before every claim can be checked.""" + seen = [] + for claim in CLAIMS: + if claim.task not in seen: + seen.append(claim.task) + return seen diff --git a/task_runners/cluster.env b/task_runners/cluster.env new file mode 100644 index 0000000..8108f62 --- /dev/null +++ b/task_runners/cluster.env @@ -0,0 +1,36 @@ +# --------------------------------------------------------------------------- +# Site configuration for the SLURM submission scripts. +# +# Every value is optional. An empty one is simply not passed to sbatch, so +# the #SBATCH defaults inside task_runners/sbatch/*.sbatch apply. This is the +# one file to edit when moving to a new cluster; nothing else here is +# site-specific. +# +# The CSP_ prefix is deliberate: SLURM_ACCOUNT, SLURM_PARTITION and friends +# are set by SLURM itself inside a running job, so reusing those names here +# would collide. +# --------------------------------------------------------------------------- + +# --- scheduler ------------------------------------------------------------- +CSP_ACCOUNT="" # sbatch --account +CSP_PARTITION="" # sbatch --partition +CSP_QOS="" # sbatch --qos +CSP_CONSTRAINT="" # sbatch --constraint, e.g. "a100" +CSP_GPU_GRES="" # sbatch --gres, e.g. "gpu:1" or "gpu:a100:1" +CSP_RESERVATION="" # sbatch --reservation +CSP_MAIL_USER="" # sbatch --mail-user (with --mail-type=END,FAIL) +CSP_MAX_CONCURRENT="4" # array throttle, the %N in --array=0-9%N +CSP_SBATCH_EXTRA="" # anything else, appended to the sbatch line + +# --- environment ----------------------------------------------------------- +CSP_MODULES="" # space-separated `module load` arguments +CSP_ENV="" # conda env with torch + this repo installed +CSP_SCORE_ENV="" # conda env with pymatgen + average-minimum-distance + # (leave empty to score in CSP_ENV) +CSP_PRE_RUN_HOOK="" # extra shell evaluated before the task runs + +# --- paths ----------------------------------------------------------------- +CSP_RUNS="" # where data/checkpoints/results go + # (default /runs; point this at scratch) +CSP_ATOMBENCH_REPO="" # clone of github.com/atomgptlab/atombench, + # needed by the scoring stages diff --git a/task_runners/common.sh b/task_runners/common.sh new file mode 100644 index 0000000..d0c81b4 --- /dev/null +++ b/task_runners/common.sh @@ -0,0 +1,43 @@ +# Shared bootstrap, sourced by every script in task_runners/sbatch. +# +# Reads cluster.env, loads modules, activates the conda environment and +# exports the variables the underlying scripts read (ALIGNN_RUNS, +# ATOMBENCH_REPO, SCORE_ENV). Sourcing this from an interactive shell is a +# perfectly good way to get the same environment by hand. + +set -euo pipefail + +CSP_HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CSP_REPO="$(cd "$CSP_HERE/.." && pwd)" + +# shellcheck disable=SC1091 +[ -f "$CSP_HERE/cluster.env" ] && source "$CSP_HERE/cluster.env" + +if [ -n "${CSP_MODULES:-}" ] && command -v module >/dev/null 2>&1; then + # shellcheck disable=SC2086 + module load ${CSP_MODULES} +fi + +if [ -n "${CSP_ENV:-}" ]; then + # shellcheck disable=SC1091 + source "$(conda info --base)/etc/profile.d/conda.sh" + conda activate "$CSP_ENV" +fi + +# score.sh switches to this environment itself, if it is set. +export SCORE_ENV="${CSP_SCORE_ENV:-}" +[ -n "${CSP_ATOMBENCH_REPO:-}" ] && export ATOMBENCH_REPO="$CSP_ATOMBENCH_REPO" +export ALIGNN_RUNS="${CSP_RUNS:-$CSP_REPO/runs}" +export PYTHONUNBUFFERED=1 + +if [ -n "${CSP_PRE_RUN_HOOK:-}" ]; then + eval "$CSP_PRE_RUN_HOOK" +fi + +cd "$CSP_REPO" + +echo "repo: $CSP_REPO" +echo "runs: $ALIGNN_RUNS" +echo "python: $(command -v python)" +echo "host: $(hostname)" +echo "job: ${SLURM_JOB_ID:-none} array element ${SLURM_ARRAY_TASK_ID:-none}" diff --git a/task_runners/inspect_checkpoint.py b/task_runners/inspect_checkpoint.py new file mode 100755 index 0000000..c7b34ab --- /dev/null +++ b/task_runners/inspect_checkpoint.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +"""Print the training flags baked into an ALIGNN-CSP checkpoint. + +``train_csp.py`` stores the full argument namespace in every checkpoint and in +``config.json`` next to it, so a released model can say exactly how it was +trained. Two epoch counts in ``task_runners/tasks.py`` (Alexandria and the +dft_3d pretraining run) are *not* pinned by the manuscript; this is how to +replace them with the published values instead of guessing. + + python task_runners/inspect_checkpoint.py csp_supercon_alex + python task_runners/inspect_checkpoint.py runs/train/jarvis_A0/seed0 + python task_runners/inspect_checkpoint.py path/to/best_model.pt --command +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +# Flags that describe the run rather than the machine it ran on. +INTERESTING = [ + "epochs", + "batch_size", + "lr", + "weight_decay", + "hidden_features", + "alignn_layers", + "gcn_layers", + "knn", + "num_steps", + "sigma_min", + "sigma_max", + "prop_dropout", + "composition_dropout", + "lattice_weight", + "frac_weight", + "angle_weight", + "ablation", + "angle_diffusion", + "angle_feedback", + "topology", + "radius_cutoff", + "envelope_exponent", + "gate_pair_messages", + "angle_basis", + "ema_decay", + "grad_clip", + "augment", + "init_from", + "seed", + "n_parameters", +] + +SKIP_IN_COMMAND = {"n_parameters", "device", "log_every", "output", "data_dir"} + + +def load_config(target: str) -> dict: + """Config for a local checkpoint, a run directory, or a released name.""" + path = Path(target) + if path.is_dir(): + for name in ("config.json",): + if (path / name).exists(): + return json.loads((path / name).read_text()) + path = path / "best_model.pt" + if path.suffix == ".pt" and path.exists(): + import torch + + ckpt = torch.load(path, map_location="cpu", weights_only=False) + return ckpt.get("config", {}) + if path.exists(): + return json.loads(path.read_text()) + + # Not a path: treat it as a name in the ALIGNN 2.0 registry. + from alignn.pretrained import get_alignn2_model + + paths = get_alignn2_model(target) + cfg_path = _find(paths, "config.json") + if cfg_path: + return json.loads(Path(cfg_path).read_text()) + ckpt_path = _find(paths, "best_model.pt") + if not ckpt_path: + raise FileNotFoundError(f"no config or checkpoint for {target!r}") + import torch + + return torch.load(ckpt_path, map_location="cpu", weights_only=False).get( + "config", {} + ) + + +def _find(paths, name: str): + """Pull one artifact out of whatever get_alignn2_model returned.""" + if isinstance(paths, dict): + for value in paths.values(): + if str(value).endswith(name): + return value + return None + if isinstance(paths, (list, tuple)): + for value in paths: + if str(value).endswith(name): + return value + return None + candidate = Path(paths) + if candidate.is_dir() and (candidate / name).exists(): + return candidate / name + return candidate if str(candidate).endswith(name) else None + + +def as_command(cfg: dict) -> str: + parts = [ + "python -m alignn.inverse.train_csp", + " --data-dir DATA", + " --output OUT", + ] + for key, value in sorted(cfg.items()): + if key in SKIP_IN_COMMAND or value is None: + continue + flag = "--" + key.replace("_", "-") + if isinstance(value, bool): + value = int(value) + parts.append(f" {flag} {value}") + return " \\\n".join(parts) + + +def main() -> int: + ap = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + ap.add_argument( + "target", + help="a run directory, a .pt/config.json, or a registered model name " + "(csp_supercon_jarvis, csp_supercon_alex, csp_pretrain_dft3d, ...)", + ) + ap.add_argument( + "--command", + action="store_true", + help="print a train_csp command line instead of a table", + ) + ap.add_argument("--all", action="store_true", help="every stored key") + args = ap.parse_args() + + try: + cfg = load_config(args.target) + except Exception as exc: # noqa: BLE001 - a CLI, not a library + print(f"could not read {args.target!r}: {exc}", file=sys.stderr) + return 1 + if not cfg: + print(f"{args.target}: no config recorded", file=sys.stderr) + return 1 + + if args.command: + print(as_command(cfg)) + return 0 + + keys = sorted(cfg) if args.all else [k for k in INTERESTING if k in cfg] + width = max(len(k) for k in keys) + print(f"\n{args.target}") + for key in keys: + print(f" {key.ljust(width)} {cfg[key]}") + print() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/task_runners/logs/.gitkeep b/task_runners/logs/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/task_runners/run_task.py b/task_runners/run_task.py new file mode 100755 index 0000000..a4a87cc --- /dev/null +++ b/task_runners/run_task.py @@ -0,0 +1,617 @@ +#!/usr/bin/env python3 +"""Run one of the manuscript's inverse-design tasks. + + python task_runners/run_task.py tasks # what is available + python task_runners/run_task.py doctor # is this box ready + python task_runners/run_task.py verify # paper -> task map + python task_runners/run_task.py bench-jarvis --list # the units + python task_runners/run_task.py bench-jarvis --unit 0 + python task_runners/run_task.py bench-jarvis # all units, in order + python task_runners/run_task.py bench-jarvis --aggregate + +Three cost settings. Full is the default. ``--quick`` trains for 300 epochs +with 8 candidates on the whole test split, which is a real arm-vs-arm +comparison at roughly a tenth of the cost; ``--smoke`` is 2 epochs on 4 +targets and proves only that the plumbing works. Both write into their own +run tree, so a cheap run can never overwrite an expensive checkpoint. + +Every stage records the exact command it ran in ``/.stages``, and is +skipped on a re-run if that command has not changed. Change a hyperparameter +and the affected stages re-run; change nothing and the task resumes where it +stopped. ``--force`` ignores the markers. + +Under SLURM, ``--unit $SLURM_ARRAY_TASK_ID`` makes each array element one +unit; ``task_runners/submit.sh`` sizes the array from ``--count``. +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import subprocess +import sys +import time +from pathlib import Path +from typing import List, Optional, Sequence + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from task_runners import tasks as T # noqa: E402 + +REPO = T.REPO + + +# --------------------------------------------------------------------------- +# Stage execution +# --------------------------------------------------------------------------- + + +def marker_path(unit: T.Unit, stage: T.Stage) -> Path: + return unit.rundir / ".stages" / f"{stage.name}.json" + + +def already_done(unit: T.Unit, stage: T.Stage) -> bool: + """True if this exact command has completed here before.""" + path = marker_path(unit, stage) + if not path.exists(): + return False + try: + rec = json.loads(path.read_text()) + except json.JSONDecodeError: + return False + return rec.get("argv") == stage.argv and rec.get("env") == stage.env + + +def check_requires(stage: T.Stage) -> List[str]: + """Inputs the stage needs that are not there yet.""" + return [ + f"{path} (produced by {producer})" + for path, producer in stage.requires + if not Path(path).exists() + ] + + +def run_stage( + unit: T.Unit, stage: T.Stage, *, dry_run: bool, force: bool +) -> str: + """Run one stage. Returns 'skip', 'ok', 'blocked' or 'fail'.""" + label = f"[{unit.name}/{stage.name}]" + if already_done(unit, stage) and not force: + print(f"{label} skip (already done)") + return "skip" + + missing = check_requires(stage) + if missing and not dry_run: + print(f"{label} BLOCKED, missing input(s):") + for m in missing: + print(f" {m}") + return "blocked" + + env = dict(os.environ, PYTHONUNBUFFERED="1", **stage.env) + printable = " ".join(stage.argv) + if dry_run: + prefix = " ".join(f"{k}={v}" for k, v in stage.env.items()) + print(f"{label} {prefix + ' ' if prefix else ''}{printable}") + return "ok" + + print(f"{label} $ {printable}", flush=True) + unit.rundir.mkdir(parents=True, exist_ok=True) + started = time.time() + proc = subprocess.run(stage.argv, cwd=REPO, env=env) + elapsed = time.time() - started + if proc.returncode != 0: + print(f"{label} FAILED (exit {proc.returncode}) after {elapsed:.0f}s") + return "fail" + + path = marker_path(unit, stage) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "argv": stage.argv, + "env": stage.env, + "finished": time.strftime("%Y-%m-%dT%H:%M:%S"), + "elapsed_s": round(elapsed, 1), + "host": platform.node(), + "git": git_rev(), + }, + indent=2, + ) + ) + print(f"{label} done in {elapsed:.0f}s", flush=True) + return "ok" + + +def git_rev() -> Optional[str]: + try: + out = subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + cwd=REPO, + capture_output=True, + text=True, + ) + return out.stdout.strip() or None + except OSError: + return None + + +def select_stages( + unit: T.Unit, only: Sequence[str], skip: Sequence[str] +) -> List[T.Stage]: + """Filter a unit's stages, keeping their order.""" + stages = unit.stages + if only: + stages = [s for s in stages if s.name in only] + if skip: + stages = [s for s in stages if s.name not in skip] + return stages + + +def run_unit( + unit: T.Unit, + *, + dry_run: bool, + force: bool, + only: Sequence[str] = (), + skip: Sequence[str] = (), +) -> bool: + """Run a unit's stages in order. Stops at the first failure.""" + print(f"\n=== unit {unit.name} -> {unit.rundir}") + stages = select_stages(unit, only, skip) + if not stages: + # Silently doing nothing and reporting success is the worst outcome + # here: --loss-only on a task with no training stage would look like + # a completed run. + have = ", ".join(s.name for s in unit.stages) + print(f" nothing to run: the filter left no stages of [{have}]") + return True + for stage in stages: + status = run_stage(unit, stage, dry_run=dry_run, force=force) + if status in ("fail", "blocked"): + return False + return True + + +# --------------------------------------------------------------------------- +# Context +# --------------------------------------------------------------------------- + + +def build_ctx(args, task: T.Task) -> T.Ctx: + seeds = ( + tuple(int(s) for s in args.seeds.split(",")) + if args.seeds + else tuple(task.default_seeds) + ) + if args.smoke: + seeds = seeds[:1] + elif args.quick and not args.seeds: + seeds = seeds[: T.QUICK["seeds"]] + return T.Ctx( + runs=Path(args.runs_root).resolve(), + seeds=seeds, + epochs=args.epochs, + device=args.device, + num_candidates=args.num_candidates, + guidance=args.guidance, + symprec=args.symprec, + relax_workers=args.relax_workers, + relax_steps=args.relax_steps, + limit=args.limit, + checkpoint=args.checkpoint, + alex_inputs=tuple(args.alex_inputs or ()), + smoke=args.smoke, + quick=args.quick, + from_scratch=args.from_scratch, + ) + + +# --------------------------------------------------------------------------- +# Sub-commands that are not tasks +# --------------------------------------------------------------------------- + + +def cmd_tasks() -> int: + width = max(len(n) for n in T.TASKS) + print("\nInverse-design tasks (task_runners/tasks.py)\n") + for name, task in T.TASKS.items(): + print(f" {name.ljust(width)} {task.summary}") + print(f" {' ' * width} reproduces: {task.reproduces}") + if task.needs: + print(f" {' ' * width} needs: {', '.join(task.needs)}") + print() + return 0 + + +def cmd_verify(args) -> int: + """Map every published number to a task, and to what was measured.""" + from task_runners import aggregate, claims + + print("\nManuscript coverage -- inverse design\n") + if args.quick or args.smoke: + mode = "smoke" if args.smoke else "quick" + print( + f" NOTE: --{mode} reads a separate run tree whose arms are " + "comparable to\n each other but not to the published " + "numbers. Expect disagreement.\n" + ) + orphans = [c for c in claims.CLAIMS if c.task not in T.TASKS] + if orphans: + for claim in orphans: + print(f" NO TASK for {claim.source}: {claim.statement}") + print(f"\n{len(orphans)} claim(s) have no runner. This is a bug.") + return 2 + + # Build each task once; a task's units are the same for every claim on it. + built = {} + for name in claims.tasks_needed(): + task = T.get(name) + built[name] = (task, task.build(build_ctx(args, task))) + + header = ( + f" {'source':<16}{'claim':<52}{'published':>10}" + f"{'measured':>20}{'':>3}" + ) + print(header) + print(" " + "-" * (len(header) - 2)) + counts = {"ok": 0, "within sd": 0, "differs": 0, "not run": 0} + todo = [] + for claim in claims.CLAIMS: + task, units = built[claim.task] + result = aggregate.check(task, units, claim) + counts[result["status"]] += 1 + if result["status"] == "not run" and claim.task not in todo: + todo.append(claim.task) + measured = result["measured"] + if measured is None: + shown = "-" + elif result["sd"]: + shown = f"{measured:.4g} +-{result['sd']:.2g}" + else: + shown = f"{measured:.4g}" + mark = { + "ok": "ok", + "within sd": "~", + "differs": "!!", + "not run": "", + }[result["status"]] + print( + f" {claim.source:<16}{claim.statement[:50]:<52}" + f"{claim.published:>10.4g}{shown:>20}{mark:>3}" + ) + + print( + "\n ok = within the claim's tolerance ~ = within one measured " + "standard deviation\n !! = outside both blank = not run yet" + ) + print(f"\n {len(claims.CLAIMS)} published numbers, all mapped to a task.") + print( + f" ok {counts['ok']} within sd {counts['within sd']} " + f"differs {counts['differs']} not run {counts['not run']}" + ) + if todo: + print( + "\n to fill the gaps, in dependency order (a prerequisite that " + "is\n already finished costs nothing to re-run -- its stages " + "are skipped):" + ) + for name in _ordered(todo): + print(f" python task_runners/run_task.py {name}") + print() + return 0 + + +def _ordered(names: Sequence[str]) -> List[str]: + """Topologically sort task names by their declared prerequisites. + + Ties are broken by registration order in ``TASKS``, which is the order a + person would run them in, so the list reads as a plan rather than as + whichever claim happened to be listed first. + """ + order = list(T.TASKS) + out: List[str] = [] + + def visit(name: str) -> None: + if name in out or name not in T.TASKS: + return + for need in T.TASKS[name].needs: + visit(need) + if name not in out: + out.append(name) + + for name in sorted(names, key=lambda n: order.index(n)): + visit(name) + return out + + +def _probe(label: str, fn) -> bool: + try: + detail = fn() + except Exception as exc: # noqa: BLE001 - a doctor reports, never raises + print(f" [ ] {label}: {type(exc).__name__}: {exc}") + return False + print(f" [x] {label}: {detail}") + return True + + +def cmd_doctor(runs_root: Path) -> int: + """Check the things that make a task fail an hour in, not a minute in.""" + print("\nEnvironment") + ok = True + + def _torch(): + import torch + + cuda = ( + f"cuda {torch.version.cuda}, " + f"{torch.cuda.device_count()} device(s)" + if torch.cuda.is_available() + else "no CUDA (use --device cpu)" + ) + # Importing torch is not the same as being able to train with it: a + # mismatched install can import cleanly and then fail inside the + # optimiser, an hour into a queued job. Take one step here instead. + net = torch.nn.Linear(4, 4) + opt = torch.optim.AdamW(net.parameters(), lr=1e-3) + net(torch.zeros(1, 4)).sum().backward() + opt.step() + return f"{torch.__version__}, {cuda}, one optimiser step ok" + + ok &= _probe("torch", _torch) + ok &= _probe( + "alignn.inverse", + lambda: __import__( + "alignn.inverse.train_csp", fromlist=["main"] + ).__name__, + ) + ok &= _probe( + "jarvis-tools", + lambda: __import__("jarvis").__version__, + ) + ok &= _probe("pymatgen", lambda: __import__("pymatgen.core").core.__name__) + ok &= _probe( + "average-minimum-distance (ccRMSD)", + lambda: __import__("amd").__version__, + ) + + print("\nScoring") + compute = find_compute_metrics() + if compute: + print(f" [x] AtomBench compute_metrics.py: {compute}") + else: + ok = False + print( + " [ ] AtomBench compute_metrics.py not found. Clone\n" + " https://github.com/atomgptlab/atombench and set\n" + " ATOMBENCH_REPO to it (score.sh also looks in ~/atombench)." + ) + + print(f"\nRun root: {runs_root}") + for name in ("jarvis", "alex", "pretrain"): + path = runs_root / "data" / name + mark = "x" if (path / "train.json").exists() else " " + print(f" [{mark}] data/{name}") + + print() + return 0 if ok else 1 + + +def _names(spec: Optional[str]) -> tuple: + if not spec: + return () + return tuple(s.strip() for s in spec.split(",") if s.strip()) + + +def find_compute_metrics() -> Optional[Path]: + """Mirror score.sh's search for AtomBench's metric script.""" + candidates = [] + if os.environ.get("ATOMBENCH_REPO"): + candidates.append(Path(os.environ["ATOMBENCH_REPO"])) + candidates += [Path.home() / "atombench", REPO.parent / "atombench"] + for repo in candidates: + path = repo / "scripts" / "scripts_consolidated" / "compute_metrics.py" + if path.exists(): + return path + return None + + +# --------------------------------------------------------------------------- + + +def main() -> int: + ap = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + ap.add_argument( + "task", + help="task name, or one of: tasks (list them), doctor (check this " + "machine), verify (map the paper's numbers to tasks and to what " + "has been measured)", + ) + ap.add_argument( + "--unit", + type=int, + default=None, + help="run only this unit index (SLURM array element)", + ) + ap.add_argument("--list", action="store_true", help="list units and exit") + ap.add_argument( + "--count", + action="store_true", + help="print the number of units and exit (sizes an sbatch array)", + ) + ap.add_argument("--dry-run", action="store_true") + ap.add_argument( + "--force", action="store_true", help="re-run completed stages" + ) + ap.add_argument( + "--aggregate", + action="store_true", + help="summarise this task's results instead of running it", + ) + ap.add_argument("--latex", action="store_true", help="with --aggregate") + ap.add_argument( + "--variant", + default=None, + help="with --aggregate: which scored CSV to read, 'sym' or 'nosym' " + "(default: whichever the task's published numbers are quoted on)", + ) + + ap.add_argument( + "--runs-root", + default=os.environ.get("ALIGNN_RUNS", str(REPO / "runs")), + help="where data, checkpoints and results live " + "(env ALIGNN_RUNS; default /runs)", + ) + ap.add_argument( + "--seeds", default=None, help="comma-separated, e.g. 0,1,2" + ) + ap.add_argument("--epochs", type=int, default=None) + ap.add_argument( + "--device", default=os.environ.get("ALIGNN_DEVICE", "cuda") + ) + ap.add_argument("--num-candidates", type=int, default=None) + ap.add_argument("--guidance", type=float, default=2.0) + ap.add_argument( + "--symprec", + type=float, + default=0.1, + help="symmetrisation tolerance; choose it with symprec-sweep", + ) + ap.add_argument("--relax-steps", type=int, default=None) + ap.add_argument("--relax-workers", type=int, default=None) + ap.add_argument( + "--limit", type=int, default=None, help="first N targets only" + ) + ap.add_argument( + "--checkpoint", + default=None, + help="checkpoint for pipeline-ablation / symprec-sweep " + "(default: this task's first seed)", + ) + ap.add_argument( + "--alex-inputs", + nargs="+", + default=None, + help="DS-A.pk.bz2 DS-B.pk.bz2, in that order", + ) + ap.add_argument( + "--smoke", + action="store_true", + help="2 epochs, 2 candidates, 4 targets, one seed: plumbing only", + ) + ap.add_argument( + "--quick", + action="store_true", + help=f"{T.QUICK['epochs']} epochs, {T.QUICK['num_candidates']} " + f"candidates, {T.QUICK['seeds']} seeds, whole test split: a real " + "arm-vs-arm comparison at a fraction of the cost", + ) + ap.add_argument( + "--loss-only", + action="store_true", + help="train and stop. The denoising loss is the cheapest and most " + "reproducible arm-vs-arm signal, and needs no scoring environment", + ) + ap.add_argument( + "--only-stages", + default=None, + help="comma-separated stage names to run (train, generate, " + "symmetrize, score-nosym, score-sym)", + ) + ap.add_argument( + "--skip-stages", default=None, help="comma-separated stages to skip" + ) + ap.add_argument( + "--from-scratch", + action="store_true", + help="bench-alex: train on Alexandria alone, no pretrained init", + ) + args = ap.parse_args() + + if args.task == "tasks": + return cmd_tasks() + if args.task == "doctor": + return cmd_doctor(Path(args.runs_root).resolve()) + if args.task == "verify": + return cmd_verify(args) + + try: + task = T.get(args.task) + except KeyError as exc: + print(exc, file=sys.stderr) + return 2 + + ctx = build_ctx(args, task) + units = task.build(ctx) + + if args.count: + print(len(units)) + return 0 + + if args.aggregate: + from task_runners import aggregate + + return aggregate.report( + task, + units, + variant=args.variant or task.variant, + latex=args.latex, + ) + + if args.list: + print(f"\n{task.name}: {task.summary}") + print(f"reproduces: {task.reproduces}") + if task.needs: + print(f"needs: {', '.join(task.needs)}") + print(f"\n{len(units)} unit(s):") + for i, unit in enumerate(units): + stages = ", ".join(s.name for s in unit.stages) + print(f" {i:3d} {unit.name:<28} [{stages}]") + print(f" {unit.rundir}") + return 0 + + if args.smoke: + print("smoke mode: 2 epochs, 2 candidates, 4 targets, one seed") + elif args.quick: + print( + f"quick mode: {ctx.epochs_for('jarvis')} epochs, " + f"{T.QUICK['num_candidates']} candidates, {len(ctx.seeds)} " + "seed(s), whole test split -- arms comparable to each other, " + "not to the published numbers" + ) + + only = _names(args.only_stages) + skip = _names(args.skip_stages) + if args.loss_only: + only = ("train",) + + selected = units if args.unit is None else [units[args.unit]] + failed = [] + for unit in selected: + if not run_unit( + unit, + dry_run=args.dry_run, + force=args.force, + only=only, + skip=skip, + ): + failed.append(unit.name) + + if failed: + print(f"\n{len(failed)} unit(s) did not finish: {', '.join(failed)}") + return 1 + if not args.dry_run: + print(f"\n{len(selected)} unit(s) complete.") + print( + f"summarise with: python task_runners/run_task.py " + f"{task.name} --aggregate" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/task_runners/sbatch/ablation-linegraph.sbatch b/task_runners/sbatch/ablation-linegraph.sbatch new file mode 100755 index 0000000..88bfc11 --- /dev/null +++ b/task_runners/sbatch/ablation-linegraph.sbatch @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-ablation-linegraph +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-5 +#SBATCH --time=24:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --gres=gpu:1 +# +# Line graph vs the same budget spent on pair-graph depth. +# reproduces: Table 3 (tab:inverse_ablation) +# +# 6 array element(s) at the default seeds. +# Needs, in order: data-jarvis. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh ablation-linegraph --seeds 0,1,2,3,4 +# +# Plain `sbatch task_runners/sbatch/ablation-linegraph.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py ablation-linegraph \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/aggregate.sbatch b/task_runners/sbatch/aggregate.sbatch new file mode 100755 index 0000000..60e0c84 --- /dev/null +++ b/task_runners/sbatch/aggregate.sbatch @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-aggregate +#SBATCH --output=task_runners/logs/%x-%j.out +#SBATCH --error=task_runners/logs/%x-%j.err +#SBATCH --time=01:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=2 +#SBATCH --mem=8G +# +# Print a task's table once its array has finished. submit.sh queues this +# with --dependency=afterok on the array job, so it only runs if every unit +# succeeded; run it by hand at any time to see partial results. +# +# CSP_TASK=bench-jarvis sbatch task_runners/sbatch/aggregate.sbatch +# +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +TASK="${CSP_TASK:?set CSP_TASK to the task to summarise}" + +# shellcheck disable=SC2086 +python task_runners/run_task.py "$TASK" --aggregate --latex ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/angle-ablation.sbatch b/task_runners/sbatch/angle-ablation.sbatch new file mode 100755 index 0000000..c51a3ec --- /dev/null +++ b/task_runners/sbatch/angle-ablation.sbatch @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-angle-ablation +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-17 +#SBATCH --time=36:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --gres=gpu:1 +# +# The A0-A6 angular-diffusion suite from alignn.inverse.ablations. +# reproduces: the explicit bond-angle denoising extension (this branch) +# +# 18 array element(s) at the default seeds. +# Needs, in order: data-jarvis. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh angle-ablation --seeds 0,1,2,3,4 +# +# Plain `sbatch task_runners/sbatch/angle-ablation.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py angle-ablation \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/bench-alex.sbatch b/task_runners/sbatch/bench-alex.sbatch new file mode 100755 index 0000000..b17ce40 --- /dev/null +++ b/task_runners/sbatch/bench-alex.sbatch @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-bench-alex +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-0 +#SBATCH --time=48:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=96G +#SBATCH --gres=gpu:1 +# +# Fine-tune from the base model and benchmark on Alexandria DS-A/B. +# reproduces: Table 4 (tab:inverse_bench), Alexandria DS-A/B block +# +# 1 array element(s) at the default seeds. +# Needs, in order: data-alex -> pretrain. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh bench-alex +# +# Plain `sbatch task_runners/sbatch/bench-alex.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py bench-alex \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/bench-jarvis.sbatch b/task_runners/sbatch/bench-jarvis.sbatch new file mode 100755 index 0000000..4e26c0f --- /dev/null +++ b/task_runners/sbatch/bench-jarvis.sbatch @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-bench-jarvis +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-2 +#SBATCH --time=24:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --gres=gpu:1 +# +# Train and benchmark ALIGNN-CSP on JARVIS Supercon-3D, three seeds. +# reproduces: Table 4 (tab:inverse_bench), JARVIS Supercon-3D block +# +# 3 array element(s) at the default seeds. +# Needs, in order: data-jarvis. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh bench-jarvis --seeds 0,1,2,3,4 +# +# Plain `sbatch task_runners/sbatch/bench-jarvis.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py bench-jarvis \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/data-alex.sbatch b/task_runners/sbatch/data-alex.sbatch new file mode 100755 index 0000000..b97b1df --- /dev/null +++ b/task_runners/sbatch/data-alex.sbatch @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-data-alex +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-0 +#SBATCH --time=04:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +# +# Build the AtomBench Alexandria DS-A/B split (6603/825/825). +# reproduces: the split the Alexandria block of Table 4 is measured on +# +# 1 array element(s) at the default seeds. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh data-alex +# +# Plain `sbatch task_runners/sbatch/data-alex.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py data-alex \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/data-jarvis.sbatch b/task_runners/sbatch/data-jarvis.sbatch new file mode 100755 index 0000000..adc049a --- /dev/null +++ b/task_runners/sbatch/data-jarvis.sbatch @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-data-jarvis +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-0 +#SBATCH --time=02:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +# +# Build the AtomBench JARVIS Supercon-3D split (847/105/103). +# reproduces: the split every JARVIS number in the paper is measured on +# +# 1 array element(s) at the default seeds. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh data-jarvis +# +# Plain `sbatch task_runners/sbatch/data-jarvis.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py data-jarvis \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/data-pretrain.sbatch b/task_runners/sbatch/data-pretrain.sbatch new file mode 100755 index 0000000..1785295 --- /dev/null +++ b/task_runners/sbatch/data-pretrain.sbatch @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-data-pretrain +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-0 +#SBATCH --time=06:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +# +# Build the 65k dft_3d pretraining corpus, benchmark ids held out. +# reproduces: the corpus behind csp_pretrain_dft3d +# +# 1 array element(s) at the default seeds. +# Needs, in order: data-jarvis. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh data-pretrain +# +# Plain `sbatch task_runners/sbatch/data-pretrain.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py data-pretrain \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/leakage.sbatch b/task_runners/sbatch/leakage.sbatch new file mode 100755 index 0000000..8671c47 --- /dev/null +++ b/task_runners/sbatch/leakage.sbatch @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-leakage +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-1 +#SBATCH --time=06:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +# +# Test targets recoverable from JARVIS-DFT by recall, and the score on the complement. +# reproduces: the 18.4% / 15.4% leakage caveat +# +# 2 array element(s) at the default seeds. +# Needs, in order: data-pretrain -> bench-jarvis. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh leakage +# +# Plain `sbatch task_runners/sbatch/leakage.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py leakage \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/pipeline-ablation.sbatch b/task_runners/sbatch/pipeline-ablation.sbatch new file mode 100755 index 0000000..3c65d93 --- /dev/null +++ b/task_runners/sbatch/pipeline-ablation.sbatch @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-pipeline-ablation +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-3 +#SBATCH --time=12:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --gres=gpu:1 +# +# raw / rank / relax / full: what sampling and the force field buy. +# reproduces: 'Closing the loop with the force field' +# +# 4 array element(s) at the default seeds. +# Needs, in order: bench-jarvis. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh pipeline-ablation +# +# Plain `sbatch task_runners/sbatch/pipeline-ablation.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py pipeline-ablation \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/pretrain-transfer.sbatch b/task_runners/sbatch/pretrain-transfer.sbatch new file mode 100755 index 0000000..ea21cc5 --- /dev/null +++ b/task_runners/sbatch/pretrain-transfer.sbatch @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-pretrain-transfer +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-1 +#SBATCH --time=48:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=96G +#SBATCH --gres=gpu:1 +# +# Alexandria from scratch vs fine-tuned from the dft_3d base model. +# reproduces: the pretraining claim behind the Table 4 Alexandria row +# +# 2 array element(s) at the default seeds. +# Needs, in order: data-alex -> pretrain. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh pretrain-transfer +# +# Plain `sbatch task_runners/sbatch/pretrain-transfer.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py pretrain-transfer \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/pretrain.sbatch b/task_runners/sbatch/pretrain.sbatch new file mode 100755 index 0000000..408605a --- /dev/null +++ b/task_runners/sbatch/pretrain.sbatch @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-pretrain +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-0 +#SBATCH --time=72:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --gres=gpu:1 +# +# Train the composition-only base model on 65k dft_3d crystals. +# reproduces: csp_pretrain_dft3d, the checkpoint bench-alex fine-tunes from +# +# 1 array element(s) at the default seeds. +# Needs, in order: data-pretrain. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh pretrain +# +# Plain `sbatch task_runners/sbatch/pretrain.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py pretrain \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/sbatch/symprec-sweep.sbatch b/task_runners/sbatch/symprec-sweep.sbatch new file mode 100755 index 0000000..a1e8cc3 --- /dev/null +++ b/task_runners/sbatch/symprec-sweep.sbatch @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +#SBATCH --job-name=csp-symprec-sweep +#SBATCH --output=task_runners/logs/%x-%A_%a.out +#SBATCH --error=task_runners/logs/%x-%A_%a.err +#SBATCH --array=0-0 +#SBATCH --time=08:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --gres=gpu:1 +# +# Pick the symmetrisation tolerance on the validation split. +# reproduces: the symmetrisation step used before scoring the test split +# +# 1 array element(s) at the default seeds. +# Needs, in order: bench-jarvis. +# +# Account, partition, QoS, constraint and the exact --gres come from +# task_runners/cluster.env, applied by submit.sh, which also resizes the array +# to match the seeds actually requested: +# +# bash task_runners/submit.sh symprec-sweep +# +# Plain `sbatch task_runners/sbatch/symprec-sweep.sbatch` works too, from the repo +# root; flags given on the sbatch command line override the directives above. +# +# The walltime above is a placeholder: it has not been measured on +# any particular machine. Check the first element and adjust. +# + +cd "${SLURM_SUBMIT_DIR:-$PWD}" +source task_runners/common.sh + +# shellcheck disable=SC2086 # CSP_RUN_ARGS is a deliberate word-split +python task_runners/run_task.py symprec-sweep \ + --unit "${SLURM_ARRAY_TASK_ID:-0}" \ + ${CSP_RUN_ARGS:-} diff --git a/task_runners/submit.sh b/task_runners/submit.sh new file mode 100755 index 0000000..743009c --- /dev/null +++ b/task_runners/submit.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash +# Submit one task to SLURM: a job array over its units, plus a dependent +# aggregation job that prints the table once every element has succeeded. +# +# bash task_runners/submit.sh bench-jarvis +# bash task_runners/submit.sh angle-ablation --seeds 0,1,2,3,4 +# bash task_runners/submit.sh bench-jarvis --dry-run # show, don't submit +# +# Arguments after the task name are forwarded to run_task.py, and are used +# both to size the array and to run each element, so `--seeds 0,1` really does +# submit two elements. Scheduler metadata comes from task_runners/cluster.env. + +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "$HERE/.." && pwd)" +# shellcheck disable=SC1091 +[ -f "$HERE/cluster.env" ] && source "$HERE/cluster.env" + +TASK="${1:?usage: submit.sh [run_task.py args...]}" +shift || true + +SBATCH_FILE="$HERE/sbatch/$TASK.sbatch" +if [ ! -f "$SBATCH_FILE" ]; then + echo "no sbatch script for task '$TASK'" >&2 + echo "available:" >&2 + ls "$HERE/sbatch" | sed 's/\.sbatch$//' | sed 's/^/ /' >&2 + exit 2 +fi + +DRY=0 +ARGS=() +for arg in "$@"; do + if [ "$arg" = "--dry-run" ]; then DRY=1; else ARGS+=("$arg"); fi +done + +cd "$REPO" +export ALIGNN_RUNS="${CSP_RUNS:-$REPO/runs}" + +# Size the array from the same arguments the elements will run with, so a +# non-default --seeds cannot silently under- or over-submit. +N_UNITS="$(python "$HERE/run_task.py" "$TASK" --count "${ARGS[@]+"${ARGS[@]}"}")" +if [ "$N_UNITS" -lt 1 ]; then + echo "task '$TASK' has no units" >&2 + exit 1 +fi + +OPTS=() +[ -n "${CSP_ACCOUNT:-}" ] && OPTS+=(--account="$CSP_ACCOUNT") +[ -n "${CSP_PARTITION:-}" ] && OPTS+=(--partition="$CSP_PARTITION") +[ -n "${CSP_QOS:-}" ] && OPTS+=(--qos="$CSP_QOS") +[ -n "${CSP_CONSTRAINT:-}" ] && OPTS+=(--constraint="$CSP_CONSTRAINT") +[ -n "${CSP_RESERVATION:-}" ] && OPTS+=(--reservation="$CSP_RESERVATION") +if [ -n "${CSP_MAIL_USER:-}" ]; then + OPTS+=(--mail-user="$CSP_MAIL_USER" --mail-type=END,FAIL) +fi +# --gres only makes sense for the GPU tasks; the CPU sbatch files declare +# none, and adding one there would queue them behind GPU availability. +if [ -n "${CSP_GPU_GRES:-}" ] && grep -q '^#SBATCH --gres' "$SBATCH_FILE"; then + OPTS+=(--gres="$CSP_GPU_GRES") +fi +# shellcheck disable=SC2206 +[ -n "${CSP_SBATCH_EXTRA:-}" ] && OPTS+=(${CSP_SBATCH_EXTRA}) + +THROTTLE="${CSP_MAX_CONCURRENT:-4}" +ARRAY="0-$((N_UNITS - 1))%${THROTTLE}" + +# Forwarded through the submitting environment rather than --export=K=V, which +# does not survive values containing spaces or commas. +CSP_RUN_ARGS="${ARGS[*]+${ARGS[*]}}" +export CSP_RUN_ARGS CSP_TASK="$TASK" + +echo "task: $TASK" +echo "units: $N_UNITS array $ARRAY" +echo "run args: ${CSP_RUN_ARGS:-}" +echo "sbatch: ${OPTS[*]+${OPTS[*]}}" + +if [ "$DRY" = "1" ]; then + echo + echo "would submit:" + echo " sbatch --array=$ARRAY ${OPTS[*]+${OPTS[*]}} $SBATCH_FILE" + echo " sbatch --dependency=afterok: ${OPTS[*]+${OPTS[*]}} \\" + echo " $HERE/sbatch/aggregate.sbatch" + exit 0 +fi + +JOB="$(sbatch --parsable --export=ALL --array="$ARRAY" \ + "${OPTS[@]+"${OPTS[@]}"}" "$SBATCH_FILE")" +echo "submitted array job $JOB" + +AGG_OPTS=() +[ -n "${CSP_ACCOUNT:-}" ] && AGG_OPTS+=(--account="$CSP_ACCOUNT") +[ -n "${CSP_PARTITION:-}" ] && AGG_OPTS+=(--partition="$CSP_PARTITION") +[ -n "${CSP_QOS:-}" ] && AGG_OPTS+=(--qos="$CSP_QOS") +AGG="$(sbatch --parsable --export=ALL --dependency="afterok:$JOB" \ + "${AGG_OPTS[@]+"${AGG_OPTS[@]}"}" "$HERE/sbatch/aggregate.sbatch")" +echo "submitted aggregation job $AGG (after $JOB)" +echo +echo "watch: squeue -j $JOB,$AGG" +echo "logs: task_runners/logs/" +echo "table: python task_runners/run_task.py $TASK --aggregate" diff --git a/task_runners/tasks.py b/task_runners/tasks.py new file mode 100644 index 0000000..09402ed --- /dev/null +++ b/task_runners/tasks.py @@ -0,0 +1,945 @@ +"""Declarative definitions of the manuscript's inverse-design tasks. + +Each entry in :data:`TASKS` reproduces one claim in the "Generative inverse +design" section. A task is a list of *units* — independent pieces of work +that can run as SLURM array elements — and each unit is a list of *stages*, +which are ordinary shell commands against the scripts already in this +repository. Nothing here reimplements the science; it pins the arguments. + + task reproduces + ------------------ -------------------------------------------------- + data-jarvis the JARVIS Supercon-3D split (847/105/103) + data-alex the Alexandria DS-A/B split (6603/825/825) + data-pretrain the 65k dft_3d pretraining corpus + pretrain the csp_pretrain_dft3d base model + bench-jarvis Table 4, JARVIS block (3 seeds) + bench-alex Table 4, Alexandria block (single run) + pretrain-transfer Alexandria from scratch vs fine-tuned + ablation-linegraph Table 3, line graph vs pair-graph depth + angle-ablation the A0-A6 angular-diffusion suite (this branch) + pipeline-ablation "Closing the loop with the force field" + symprec-sweep the symmetrisation tolerance, chosen on validation + leakage the 18.4% / 15.4% recall-not-generation caveat + +Units are keyed by their run directory, so configurations shared between +tasks are trained once. ``bench-jarvis`` and arm A of ``ablation-linegraph`` +and ``A0`` of ``angle-ablation`` are the same six-layer baseline and land in +the same ``train/jarvis_A0/seed*`` directories; running all three costs one +set of trainings. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Callable, Dict, List, Optional, Sequence + +from alignn.inverse.ablations import ABLATIONS, COMPARISONS, DESCRIPTIONS + +REPO = Path(__file__).resolve().parents[1] + +# --------------------------------------------------------------------------- +# Hyperparameters. +# +# The JARVIS numbers are the ones the inverse-design README records as having +# produced the published results; the Alexandria and pretraining epoch counts +# are *not* pinned by the manuscript, so they are marked and can be recovered +# exactly from a released checkpoint with ``inspect_checkpoint.py``. +# --------------------------------------------------------------------------- + +EPOCHS = { + "jarvis": 3000, + "alex": 1000, # not pinned by the manuscript + "pretrain": 200, # not pinned by the manuscript +} + +#: Generation settings for the benchmark tables. 32 candidates with a +#: single-point energy prescreen down to 4 relaxations is the configuration +#: the README credits with match 0.524 on JARVIS. +GEN = { + "num_candidates": 32, + "prescreen_keep": 4, + "relax": "cell", + "rank": "energy", + "relax_steps": 200, +} + +#: ``--quick``: a real measurement at a fraction of the cost, for deciding +#: whether an ablation is worth a full run. Everything that changes the +#: *comparison* is untouched -- both arms still see the whole test split, the +#: same candidate pool policy and the same scoring -- so the arms stay +#: comparable to each other. They are not comparable to the published +#: numbers, which is what ``verify`` will tell you if you try. +QUICK = { + "epochs": 300, + "num_candidates": 8, + "prescreen_keep": 2, + "sample_steps": 200, + "relax_steps": 50, + "seeds": 2, +} + +#: Tolerances swept on the validation split by ``symprec-sweep``. +SYMPREC_GRID = (0.01, 0.02, 0.05, 0.1, 0.2, 0.5) + +#: The four stages of the generate -> select -> relax pipeline, isolated. +PIPELINE_VARIANTS = { + "raw": {"num_candidates": 1, "relax": "none", "rank": "none"}, + "rank": {"num_candidates": 32, "relax": "none", "rank": "energy"}, + "relax": {"num_candidates": 1, "relax": "cell", "rank": "energy"}, + "full": { + "num_candidates": 32, + "relax": "cell", + "rank": "energy", + "prescreen_keep": 4, + }, +} + + +@dataclass(frozen=True) +class Stage: + """One command, plus the marker that says it already succeeded.""" + + name: str + argv: List[str] + env: Dict[str, str] = field(default_factory=dict) + #: A path that must exist before the stage can run, with the task that + #: makes it. Checked up front so a missing prerequisite is reported + #: rather than crashed into an hour later. + requires: Sequence[tuple] = () + + +@dataclass(frozen=True) +class Unit: + """One array element: a run directory and the stages that fill it.""" + + name: str + rundir: Path + stages: List[Stage] + #: Aggregation label. Units sharing a group are averaged over seeds. + group: str = "" + seed: Optional[int] = None + #: Where this unit's scored metrics.json files live, by variant. + metrics: Dict[str, Path] = field(default_factory=dict) + #: history.json, for the denoising validation loss. + history: Optional[Path] = None + + +@dataclass(frozen=True) +class Task: + """A named group of units, plus how to summarise them.""" + + name: str + summary: str + reproduces: str + build: Callable[["Ctx"], List[Unit]] + needs: Sequence[str] = () + #: AtomBench baseline block to print alongside, if any. + baselines: Optional[str] = None + #: Pairs of groups whose difference is the point of the task. + comparisons: Dict[str, tuple] = field(default_factory=dict) + #: Short group label -> what it means, printed under the table. + legend: Dict[str, str] = field(default_factory=dict) + #: Which scored CSV --aggregate reads by default. The lattice columns + #: are measured after symmetrisation; the pipeline paragraph is quoted + #: before it. + variant: str = "sym" + #: Whether units want a GPU (drives the sbatch header we ship). + gpu: bool = True + default_seeds: Sequence[int] = (0, 1, 2) + + +@dataclass +class Ctx: + """Runtime knobs shared by every task.""" + + runs: Path + seeds: Sequence[int] = (0, 1, 2) + epochs: Optional[int] = None + device: str = "cuda" + num_candidates: Optional[int] = None + guidance: float = 2.0 + symprec: float = 0.1 + relax_workers: Optional[int] = None + relax_steps: Optional[int] = None + limit: Optional[int] = None + checkpoint: Optional[str] = None + alex_inputs: Sequence[str] = () + smoke: bool = False + quick: bool = False + from_scratch: bool = False + + # -- derived paths ------------------------------------------------------ + @property + def data(self) -> Path: + # Splits are identical in every mode, so they are shared: a quick run + # should not rebuild the data, only the models. + return self.runs / "data" + + @property + def suffix(self) -> str: + """Keeps reduced-cost runs out of the full runs' directories. + + Without this a ``--quick`` run would land on the same checkpoint path + as the real one, and because the training command differs it would + overwrite it -- days of GPU time destroyed by a 20-minute sanity + check. + """ + if self.smoke: + return "_smoke" + if self.quick: + return "_quick" + return "" + + def out(self, *parts: str) -> Path: + """A top-level output directory for this mode.""" + head, *rest = parts + return self.runs.joinpath(f"{head}{self.suffix}", *rest) + + def train_dir(self, config: str, seed: int) -> Path: + return self.out("train", config, f"seed{seed}") + + def epochs_for(self, key: str) -> int: + """--epochs wins, then --smoke, then --quick, then the table.""" + if self.epochs is not None: + return self.epochs + if self.smoke: + return 2 + if self.quick: + return QUICK["epochs"] + return EPOCHS[key] + + def candidates(self, default: int) -> int: + if self.num_candidates: + return self.num_candidates + if self.smoke: + return 2 + if self.quick: + # Never sample *more* than the variant asks for: the pipeline + # ablation's one-sample arms must stay at one sample. + return min(default, QUICK["num_candidates"]) + return default + + +# --------------------------------------------------------------------------- +# Stage builders +# --------------------------------------------------------------------------- + + +def _py(script: str, *args: str) -> List[str]: + return [ + "python", + "-u", + str(REPO / "scripts" / "atombench" / script), + *args, + ] + + +def train_stage( + ctx: Ctx, + rundir: Path, + data_dir: Path, + *, + seed: int, + epochs_key: str, + alignn_layers: int = 3, + gcn_layers: int = 3, + ablation: str = "A0", + augment: int = 0, + init_from: Optional[Path] = None, +) -> Stage: + """A single ``alignn.inverse.train_csp`` run. + + Everything not under test is fixed here: hidden size, kNN, the diffusion + schedule, batch size and learning rate. A comparison is only meaningful + if the arms differ in the switch being studied and nothing else. + """ + argv = [ + "python", + "-u", + "-m", + "alignn.inverse.train_csp", + "--data-dir", + str(data_dir), + "--output", + str(rundir), + "--epochs", + str(ctx.epochs_for(epochs_key)), + "--seed", + str(seed), + "--ablation", + ablation, + "--alignn-layers", + str(alignn_layers), + "--gcn-layers", + str(gcn_layers), + "--hidden-features", + "256", + "--knn", + "12", + "--num-steps", + "1000", + "--batch-size", + "64", + "--lr", + "1e-3", + "--augment", + str(augment), + "--device", + ctx.device, + "--log-every", + "25", + ] + if init_from is not None: + argv += ["--init-from", str(init_from)] + return Stage( + "train", + argv, + requires=(((data_dir / "train.json"), f"data ({data_dir.name})"),) + + (((init_from, "pretrain"),) if init_from is not None else ()), + ) + + +def generate_stage( + ctx: Ctx, + checkpoint: Path, + data_dir: Path, + out_csv: Path, + *, + seed: int = 0, + split: str = "test", + overrides: Optional[Dict] = None, +) -> Stage: + """Sample, optionally rank and relax, and write an AtomBench CSV.""" + cfg = dict(GEN) + cfg.update(overrides or {}) + cfg["num_candidates"] = ctx.candidates(cfg["num_candidates"]) + if ctx.relax_steps is not None: + cfg["relax_steps"] = ctx.relax_steps + if ctx.smoke: + # The models are trained at 1000 denoising steps and sampled at 1000; + # a smoke run only needs the code path, so it takes the cheap end of + # that trade and says so rather than passing off the result. + cfg["relax_steps"] = 20 + elif ctx.quick and ctx.relax_steps is None: + cfg["relax_steps"] = QUICK["relax_steps"] + if cfg.get("prescreen_keep"): + cfg["prescreen_keep"] = QUICK["prescreen_keep"] + argv = _py( + "generate_benchmark.py", + "--checkpoint", + str(checkpoint), + "--data-dir", + str(data_dir), + "--split", + split, + "--output-csv", + str(out_csv), + "--num-candidates", + str(cfg["num_candidates"]), + "--guidance", + str(ctx.guidance), + "--relax", + cfg["relax"], + "--rank", + cfg["rank"], + "--relax-steps", + str(cfg["relax_steps"]), + "--seed", + str(seed), + "--device", + ctx.device, + "--save-candidates", + str(out_csv.with_name("candidates.json")), + ) + keep = cfg.get("prescreen_keep") + if keep and cfg["relax"] != "none": + # Never prescreen to more candidates than were sampled. + argv += ["--prescreen-keep", str(min(keep, cfg["num_candidates"]))] + if ctx.smoke: + argv += ["--steps", "50"] + elif ctx.quick: + argv += ["--steps", str(QUICK["sample_steps"])] + if ctx.relax_workers is not None: + argv += ["--relax-workers", str(ctx.relax_workers)] + limit = 4 if ctx.smoke else ctx.limit + if limit is not None: + argv += ["--limit", str(limit)] + return Stage( + "generate", + argv, + # Relaxation forks CPU workers; leaving BLAS threaded oversubscribes + # the node badly, which is what run_ablation.sh guards against too. + env={"OMP_NUM_THREADS": "1"}, + requires=((checkpoint, "the training stage"),), + ) + + +def symmetrize_stage(csv_in: Path, csv_out: Path, symprec: float) -> Stage: + """Idealise each predicted cell to its detected space group.""" + return Stage( + "symmetrize", + _py( + "symmetrize_predictions.py", + "--csv", + str(csv_in), + "--out", + str(csv_out), + "--symprec", + str(symprec), + ), + requires=((csv_in, "the generate stage"),), + ) + + +def score_stage(csv: Path, name: str = "score") -> Stage: + """Run AtomBench's own metric code; writes metrics.json beside the CSV.""" + return Stage( + name, + ["bash", str(REPO / "scripts" / "atombench" / "score.sh"), str(csv)], + requires=((csv, "the stage that writes this CSV"),), + ) + + +def eval_stages( + ctx: Ctx, + rundir: Path, + checkpoint: Path, + data_dir: Path, + *, + seed: int = 0, + split: str = "test", + overrides: Optional[Dict] = None, +) -> tuple: + """Generate, symmetrise and score, unsymmetrised and symmetrised. + + Both are kept because they answer different questions: the pipeline + ablation is about what sampling and the force field contribute, which is + visible before symmetrisation, while the manuscript's lattice-angle and + KLD columns are measured after it. + """ + nosym = rundir / "bench" / "nosym" / "pred.csv" + sym = rundir / "bench" / "sym" / "pred.csv" + stages = [ + generate_stage( + ctx, + checkpoint, + data_dir, + nosym, + seed=seed, + split=split, + overrides=overrides, + ), + symmetrize_stage(nosym, sym, ctx.symprec), + score_stage(nosym, "score-nosym"), + score_stage(sym, "score-sym"), + ] + metrics = { + "nosym": nosym.with_name("metrics.json"), + "sym": sym.with_name("metrics.json"), + } + return stages, metrics + + +# --------------------------------------------------------------------------- +# Tasks +# --------------------------------------------------------------------------- + + +def _data_jarvis(ctx: Ctx) -> List[Unit]: + out = ctx.data / "jarvis" + return [ + Unit( + "jarvis", + out, + [Stage("prepare", _py("prepare_data.py", "--output", str(out)))], + group="data", + ) + ] + + +def _data_alex(ctx: Ctx) -> List[Unit]: + out = ctx.data / "alex" + inputs = list(ctx.alex_inputs) or [ + str(ctx.data / "alexandria" / "DS-A.pk.bz2"), + str(ctx.data / "alexandria" / "DS-B.pk.bz2"), + ] + return [ + Unit( + "alex", + out, + [ + Stage( + "prepare", + _py( + "prepare_alex_data.py", + "--inputs", + *inputs, + "--output", + str(out), + ), + requires=tuple( + ( + Path(p), + "the Alexandria DS-A/DS-B pickles " + "(figshare 10.6084/m9.figshare.31045597); " + "pass --alex-inputs to point elsewhere", + ) + for p in inputs + ), + ) + ], + group="data", + ) + ] + + +def _data_pretrain(ctx: Ctx) -> List[Unit]: + out = ctx.data / "pretrain" + return [ + Unit( + "pretrain", + out, + [ + Stage( + "prepare", + _py( + "prepare_pretrain_data.py", + "--output", + str(out), + "--exclude-splits", + str(ctx.data / "jarvis"), + ), + requires=( + (ctx.data / "jarvis" / "test.json", "data-jarvis"), + ), + ) + ], + group="data", + ) + ] + + +def _pretrain(ctx: Ctx) -> List[Unit]: + rundir = ctx.train_dir("pretrain_dft3d", 0) + return [ + Unit( + "pretrain_dft3d", + rundir, + [ + train_stage( + ctx, + rundir, + ctx.data / "pretrain", + seed=0, + epochs_key="pretrain", + augment=1, + ) + ], + group="pretrain", + seed=0, + history=rundir / "history.json", + ) + ] + + +def _jarvis_units( + ctx: Ctx, + config: str, + group: str, + *, + alignn_layers: int = 3, + gcn_layers: int = 3, + ablation: str = "A0", +) -> List[Unit]: + """Train-and-evaluate units on the JARVIS split, one per seed.""" + data = ctx.data / "jarvis" + units = [] + for seed in ctx.seeds: + rundir = ctx.train_dir(config, seed) + stages = [ + train_stage( + ctx, + rundir, + data, + seed=seed, + epochs_key="jarvis", + alignn_layers=alignn_layers, + gcn_layers=gcn_layers, + ablation=ablation, + # 48 basis relabellings of 847 crystals cost accuracy on a + # split this small; the README records augmentation off as + # the best small-data setting. + augment=0, + ) + ] + ev, metrics = eval_stages( + ctx, rundir, rundir / "best_model.pt", data, seed=seed + ) + units.append( + Unit( + f"{config}-seed{seed}", + rundir, + stages + ev, + group=group, + seed=seed, + metrics=metrics, + history=rundir / "history.json", + ) + ) + return units + + +def _bench_jarvis(ctx: Ctx) -> List[Unit]: + return _jarvis_units(ctx, "jarvis_A0", "ALIGNN-CSP") + + +def _alex_units( + ctx: Ctx, config: str, group: str, *, pretrained: bool +) -> List[Unit]: + """Train-and-evaluate units on Alexandria DS-A/B, one per seed.""" + data = ctx.data / "alex" + init = ( + ctx.train_dir("pretrain_dft3d", 0) / "best_model.pt" + if pretrained + else None + ) + units = [] + for seed in ctx.seeds: + rundir = ctx.train_dir(config, seed) + stages = [ + train_stage( + ctx, + rundir, + data, + seed=seed, + epochs_key="alex", + # 6603 training crystals, so the 48 basis relabellings are + # affordable here in a way they are not on the 847-crystal + # JARVIS split. + augment=1, + init_from=init, + ) + ] + ev, metrics = eval_stages( + ctx, rundir, rundir / "best_model.pt", data, seed=seed + ) + units.append( + Unit( + f"{config}-seed{seed}", + rundir, + stages + ev, + group=group, + seed=seed, + metrics=metrics, + history=rundir / "history.json", + ) + ) + return units + + +def _bench_alex(ctx: Ctx) -> List[Unit]: + # The released csp_supercon_alex, whose 0.485 match rate is the Table 4 + # Alexandria row, was fine-tuned from csp_pretrain_dft3d -- so that is the + # default here. The leakage paragraph says the quoted ALIGNN-CSP results + # use no pretraining, which does not fit that row; --from-scratch trains + # the arm that sentence describes, and pretrain-transfer runs both. + if ctx.from_scratch: + return _alex_units(ctx, "alex_scratch", "ALIGNN-CSP", pretrained=False) + return _alex_units(ctx, "alex", "ALIGNN-CSP", pretrained=True) + + +def _pretrain_transfer(ctx: Ctx) -> List[Unit]: + """Both Alexandria arms, so the pretraining claim is a measurement.""" + return _alex_units( + ctx, "alex_scratch", "from scratch", pretrained=False + ) + _alex_units(ctx, "alex", "pretrained", pretrained=True) + + +def _ablation_linegraph(ctx: Ctx) -> List[Unit]: + # Arm B spends the deleted angular budget on pair-graph depth: nine + # convolution blocks in both arms, parameters matched to within 1%. + return _jarvis_units(ctx, "jarvis_A0", "A: line graph") + _jarvis_units( + ctx, + "jarvis_nolg", + "B: no line graph", + alignn_layers=0, + gcn_layers=9, + ) + + +def _angle_ablation(ctx: Ctx) -> List[Unit]: + units: List[Unit] = [] + for name in sorted(ABLATIONS): + units += _jarvis_units(ctx, f"jarvis_{name}", name, ablation=name) + return units + + +def _default_checkpoint(ctx: Ctx) -> Path: + if ctx.checkpoint: + return Path(ctx.checkpoint) + return ctx.train_dir("jarvis_A0", ctx.seeds[0]) / "best_model.pt" + + +def _pipeline_ablation(ctx: Ctx) -> List[Unit]: + ckpt = _default_checkpoint(ctx) + data = ctx.data / "jarvis" + units = [] + for variant, overrides in PIPELINE_VARIANTS.items(): + rundir = ctx.out("pipeline", variant) + ev, metrics = eval_stages( + ctx, rundir, ckpt, data, seed=ctx.seeds[0], overrides=overrides + ) + units.append( + Unit( + variant, + rundir, + ev, + group=variant, + metrics=metrics, + ) + ) + return units + + +def _symprec_sweep(ctx: Ctx) -> List[Unit]: + """Choose the symmetrisation tolerance on validation, never on test.""" + ckpt = _default_checkpoint(ctx) + data = ctx.data / "jarvis" + root = ctx.out("symprec") + val_csv = root / "val" / "pred.csv" + stages = [ + generate_stage( + ctx, ckpt, data, val_csv, seed=ctx.seeds[0], split="val" + ), + Stage( + "sweep", + _py( + "symmetrize_predictions.py", + "--csv", + str(val_csv), + "--out", + str(root), + "--sweep", + ",".join(f"{s:g}" for s in SYMPREC_GRID), + ), + requires=((val_csv, "the generate stage"),), + ), + score_stage(val_csv, "score-nosym"), + ] + metrics = {"none": val_csv.with_name("metrics.json")} + for sp in SYMPREC_GRID: + tag = f"symprec{sp:g}".replace(".", "p") + csv = root / tag / f"pred_{tag}.csv" + stages.append(score_stage(csv, f"score-{tag}")) + metrics[f"{sp:g}"] = csv.with_name("metrics.json") + return [Unit("sweep", root, stages, group="symprec", metrics=metrics)] + + +def _leakage(ctx: Ctx) -> List[Unit]: + """Quantify how many test targets a pretrained model could recall.""" + units = [] + sources = { + "jarvis": ( + ctx.data / "jarvis", + ctx.train_dir("jarvis_A0", ctx.seeds[0]), + ), + "alex": (ctx.data / "alex", ctx.train_dir("alex", ctx.seeds[0])), + } + for name, (data, train) in sources.items(): + root = ctx.out("leakage", name) + report = root / "leakage.json" + pred = train / "bench" / "sym" / "pred.csv" + filtered = root / "filtered" / "pred.csv" + stages = [ + Stage( + "check", + _py( + "check_pretrain_leakage.py", + "--pretrain-dir", + str(ctx.data / "pretrain"), + "--test-json", + str(data / "test.json"), + "--output", + str(report), + ), + requires=( + (ctx.data / "pretrain" / "train.json", "data-pretrain"), + (data / "test.json", f"data-{name}"), + ), + ), + Stage( + "filter", + _py( + "filter_leaked.py", + "--csv", + str(pred), + "--leakage-json", + str(report), + "--out", + str(filtered), + ), + requires=((pred, f"bench-{name}"),), + ), + score_stage(filtered, "score-filtered"), + ] + units.append( + Unit( + name, + root, + stages, + group=name, + metrics={ + "all": pred.with_name("metrics.json"), + "not-leaked": filtered.with_name("metrics.json"), + }, + ) + ) + return units + + +TASKS: Dict[str, Task] = { + "data-jarvis": Task( + "data-jarvis", + "Build the AtomBench JARVIS Supercon-3D split (847/105/103).", + "the split every JARVIS number in the paper is measured on", + _data_jarvis, + gpu=False, + default_seeds=(0,), + ), + "data-alex": Task( + "data-alex", + "Build the AtomBench Alexandria DS-A/B split (6603/825/825).", + "the split the Alexandria block of Table 4 is measured on", + _data_alex, + gpu=False, + default_seeds=(0,), + ), + "data-pretrain": Task( + "data-pretrain", + "Build the 65k dft_3d pretraining corpus, benchmark ids held out.", + "the corpus behind csp_pretrain_dft3d", + _data_pretrain, + needs=("data-jarvis",), + gpu=False, + default_seeds=(0,), + ), + "pretrain": Task( + "pretrain", + "Train the composition-only base model on 65k dft_3d crystals.", + "csp_pretrain_dft3d, the checkpoint bench-alex fine-tunes from", + _pretrain, + needs=("data-pretrain",), + default_seeds=(0,), + ), + "bench-jarvis": Task( + "bench-jarvis", + "Train and benchmark ALIGNN-CSP on JARVIS Supercon-3D, three seeds.", + "Table 4 (tab:inverse_bench), JARVIS Supercon-3D block", + _bench_jarvis, + needs=("data-jarvis",), + baselines="jarvis", + ), + "bench-alex": Task( + "bench-alex", + "Fine-tune from the base model and benchmark on Alexandria DS-A/B.", + "Table 4 (tab:inverse_bench), Alexandria DS-A/B block", + _bench_alex, + needs=("data-alex", "pretrain"), + baselines="alex", + default_seeds=(0,), + ), + "pretrain-transfer": Task( + "pretrain-transfer", + "Alexandria from scratch vs fine-tuned from the dft_3d base model.", + "the pretraining claim behind the Table 4 Alexandria row", + _pretrain_transfer, + needs=("data-alex", "pretrain"), + comparisons={ + "does pretraining help on the larger split": ( + "from scratch", + "pretrained", + ) + }, + legend={ + "from scratch": "trained on Alexandria alone", + "pretrained": "fine-tuned from csp_pretrain_dft3d (65k dft_3d)", + }, + default_seeds=(0,), + ), + "ablation-linegraph": Task( + "ablation-linegraph", + "Line graph vs the same budget spent on pair-graph depth.", + "Table 3 (tab:inverse_ablation)", + _ablation_linegraph, + needs=("data-jarvis",), + comparisons={ + "does the line graph transfer to generation": ( + "B: no line graph", + "A: line graph", + ) + }, + legend={ + "A: line graph": "three ALIGNN layers, three pair-graph " + "convolutions (3.79 M parameters)", + "B: no line graph": "no angular channel, nine pair-graph " + "convolutions (3.75 M parameters)", + }, + ), + "angle-ablation": Task( + "angle-ablation", + "The A0-A6 angular-diffusion suite from alignn.inverse.ablations.", + "the explicit bond-angle denoising extension (this branch)", + _angle_ablation, + needs=("data-jarvis",), + comparisons=dict(COMPARISONS), + legend=dict(DESCRIPTIONS), + ), + "pipeline-ablation": Task( + "pipeline-ablation", + "raw / rank / relax / full: what sampling and the force field buy.", + "'Closing the loop with the force field'", + _pipeline_ablation, + needs=("bench-jarvis",), + comparisons={ + "what one sample plus the full pipeline is worth": ( + "raw", + "full", + ), + "selection or refinement": ("relax", "rank"), + }, + variant="nosym", + legend={ + "raw": "one sample per target, straight from the diffusion model", + "rank": "32 samples, lowest ALIGNN-FF energy, no relaxation", + "relax": "one sample, relaxed with ALIGNN-FF", + "full": "32 samples, energy prescreen, top 4 relaxed", + }, + default_seeds=(0,), + ), + "symprec-sweep": Task( + "symprec-sweep", + "Pick the symmetrisation tolerance on the validation split.", + "the symmetrisation step used before scoring the test split", + _symprec_sweep, + needs=("bench-jarvis",), + default_seeds=(0,), + ), + "leakage": Task( + "leakage", + "Test targets recoverable from JARVIS-DFT by recall, and the score " + "on the complement.", + "the 18.4% / 15.4% leakage caveat", + _leakage, + needs=("data-pretrain", "bench-jarvis"), + gpu=False, + default_seeds=(0,), + ), +} + + +def get(name: str) -> Task: + """Look a task up, with the available names in the error.""" + if name not in TASKS: + raise KeyError(f"unknown task {name!r}; available: {', '.join(TASKS)}") + return TASKS[name] From 262ce52564187eb943b4692cf9c9ff32ee2b3b55 Mon Sep 17 00:00:00 2001 From: crhysc Date: Sat, 29 Aug 2026 18:35:00 -0400 Subject: [PATCH 4/4] inverse: document the angular-diffusion ablations in full MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the short overview of the angular channel with a reference section written for the group: an advisor and coauthors should be able to read it without having been briefed and know exactly what is in each arm. Adds, in order: a high-level overview of what angular denoising is, why it might matter, and the three falsifiable questions the suite asks — including what a negative result looks like, agreed in advance; a precise statement of what A0 already contains, since every arm is a delta against it; the mathematics of both additions with in-text citations separating what is taken from each source from what is ours; a paragraph on each arm; a confounds section; the protocol and metrics; a provenance table; and a bibliography. The three deviations from the literature are named and argued rather than buried: the angular state is induced rather than persistent (there is no fixed triplet list in a crystal), the target fixes the periodic-image identity, and the gate placement inside the edge-gated average is forced by the exact-deletion requirement rather than chosen. Also records a confound the earlier text did not: with gate_pair_messages on, the smooth arms truncate the pair interaction range as well as smoothing the topology. `--ablation A3 --gate-pair-messages 0` separates the two with no code change, and the section recommends running it. Numbers quoted from the graph construction were measured, not estimated (back-tracking triplet fraction, kNN vs radius triplet counts, envelope values at r_c/2). Three citation details are flagged inline as needing a check against the primary sources before they reach a manuscript. Claude-Session: https://claude.ai/code/session_01ULEhMd27yfBmKf5HXqmvFE --- README.md | 10 +- alignn/inverse/README.md | 798 +++++++++++++++++++++++++++++++-------- 2 files changed, 656 insertions(+), 152 deletions(-) diff --git a/README.md b/README.md index 204d5dd..9ea0a53 100644 --- a/README.md +++ b/README.md @@ -324,10 +324,12 @@ 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 documents 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, plus the ablation suite that tells the two apart. +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. ## Performances diff --git a/alignn/inverse/README.md b/alignn/inverse/README.md index b96e8c2..7e6fe50 100644 --- a/alignn/inverse/README.md +++ b/alignn/inverse/README.md @@ -168,177 +168,679 @@ because it uses the spawn start method the caller must then be under an > The step count trades speed for fidelity. The models were trained at 1000 > steps; whether 200 or 50 preserves benchmark quality has not been measured. -## Explicit bond-angle diffusion (optional) +## Explicit bond-angle diffusion — reference + +> **Written by Claude (Anthropic)** on the `angle-diffusion` branch, at the +> request of the repository owner. This section is meant as the working reference for +> the group: it states exactly what each ablation contains, which part of each +> design decision comes from published work and which part is ours. Every +> equation below was checked against the code in `alignn/inverse/`, and the +> numbers quoted from the graph construction were measured, not estimated. +> Nothing here has been trained. Three bibliographic details are flagged +> inline as needing a check against the primary sources before they go into a +> manuscript. + +### Overview: what this is, why it might matter, what we are testing + +**The one-sentence version.** ALIGNN's distinguishing feature is that it looks +at bond *angles*, not just bond lengths; our generative model already uses +angles to *read* a noisy crystal, but never asks the network what the angles +*should be*. This branch makes bond angles a thing the model denoises, and +asks whether that helps. + +**Where angles currently sit.** ALIGNN-CSP is a diffusion model over +fractional coordinates and the lattice, conditioned on composition +[DiffCSP, MatterGen]. Two processes run on a shared timestep: the lattice +follows DDPM in a rotation-invariant log-symmetric representation +[DDPM, iDDPM], and the coordinates follow a wrapped-normal score process on +the torus [NCSN, ScoreSDE]. The denoiser is ALIGNN [ALIGNN], so it builds a +line graph and propagates bond angles alongside bond lengths. But angles enter +only as an *input feature*. The network is asked "given this noisy structure, +which way should each atom move?" and it consults the angles to answer. It is +never asked "given this noisy structure, what should the angles be?" Angles +are an input, never an output. + +**Why anyone should care.** What separates a real crystal from a plausible +cloud of atoms is mostly *local coordination geometry* — tetrahedra, +octahedra, specific angular motifs — and that is a three-body property. A +purely coordinate-space objective supervises each atom's displacement more or +less independently and only reaches angular structure indirectly, through +whatever correlations the network happens to learn. There is a direct +precedent for doing better: FoldingDiff [FoldingDiff] generates protein +backbones by representing them in internal coordinates and running the +diffusion process *on the angles themselves*, and it reproduces natural +bond-angle distributions in a way coordinate-space models do not. Torsional +diffusion [TorsDiff] makes the same general point for molecular conformers: +when the interesting degrees of freedom are angular, define the process on the +angular space. + +**The specific puzzle in our own data.** The ablation already recorded in this +README is the reason this is worth compute rather than idle curiosity. Deleting +the line graph and spending the same parameter budget on pair-graph depth costs +a large, unambiguous amount of denoising validation loss — 2.011 ± 0.018 +against 2.351 ± 0.007 across six models per arm, p = 0.002, reproducing on two +machines — and yet leaves match rate *exactly* unchanged at 0.4709. Angular +information demonstrably helps the network fit the score, and demonstrably does +not help it find the right structure more often. One reading of that gap is +that angles are being used to interpolate rather than to constrain: the network +reads them, gets a better local fit, and still lands in the wrong basin. If +that reading is right, *supervising* angles rather than merely observing them +is the natural intervention, and this branch is the test of it. + +**The second idea: connectivity should not click.** There is a separate, +smaller problem visible in the current code. To build the line graph, we choose +each atom's twelve nearest bonds and call those the ones that form angles. That +is a sensible rule for a real crystal. It is a poor rule in the middle of +reverse diffusion, when coordinates are near-uniform: neighbour ranks swap +constantly, so the set of triplets — and therefore the function the network is +computing — jumps discontinuously from one denoising step to the next, for +reasons that have nothing to do with chemistry. Physical chemistry solved this +problem long ago. ReaxFF [ReaxFF] gives bonds a continuously varying order and +lets an angular energy term fade smoothly to zero as either of its bonds +dissociates; DimeNet [DimeNet] gives graph neural networks a cutoff envelope +whose value and first two derivatives all vanish at the cutoff radius, so an +edge can enter or leave the neighbour list without any discontinuity. We adopt +both, so that the effective line graph *crystallises* as the geometry does +rather than flickering. + +**What we are testing, stated as falsifiable questions.** + +1. Does an explicit three-body denoising *objective* improve crystal + generation, over and above angles being an input feature? +2. Does continuously varying interaction topology matter in the high-noise + regime, independently of any angular objective? +3. If something improves, is it because the angular representation is + *coupled back* into the coordinate/lattice pathway, or would any auxiliary + task on a shared trunk have done the same? + +**What a negative result looks like, agreed in advance.** If A2 (smooth +topology, no angular objective) captures the whole gain, then the angular +objective is not the story and we should say so. If A4 (angular objective with +the coupling severed) matches A3, then the benefit is generic auxiliary +supervision and the three-body claim fails even if the benchmark number goes +up. Both outcomes are publishable and both are cheap to reach; the ablation +suite is designed so that we cannot avoid learning which one we are in. + +**Scope discipline.** This stays an ALIGNN. There is no transformer, no second +network, no cross-attention, no learned bond classifier — CrystalDiT +[CrystalDiT] is a recent reminder that multi-stream architectural complexity +is usually not what buys crystal-generation performance, and MatterGen +[MatterGen] denoises every crystal variable through one shared score network. +The angular channel is one extra output head on the representation that +already exists. + +### Notation -Everything above treats bond angles as an *input feature*: ALIGNN computes -them, the line graph propagates them, and they help the network predict where -the atoms should go. They are never themselves something the model denoises. -This section adds the option of making them one, and of letting the triplet -topology change continuously while it does. All of it is off by default — -`ALIGNNCSPDenoiser()` with no arguments is byte-for-byte the model described -above, and a test asserts that switching the angle head on leaves `eps_frac` -and `eps_lattice` unchanged. +``` +N atoms in a structure (batched: all crystals concatenated) +f ∈ [0,1)^3 fractional coordinates, per atom +L ∈ R^3x3 lattice matrix, rows are lattice vectors, cart = f L +t diffusion timestep, shared by the lattice and coordinate processes +i, j, k atom indices; a triplet is the angle at the shared atom +(i→j) a directed pair ("bond"), a node of the line graph +r_ij minimum-image Cartesian vector from i to j +θ_ijk interior bond angle at j, in [0, π] +s_ij smooth pair relevance in [0,1] +s_ijk smooth triplet relevance in [0,1] +h, y, z atom, bond and triplet features inside the network +``` + +Atom types are **not** diffused: the generator is conditioned on composition, +so the state is `(f, L)` and, with this branch, `(f, L, Θ)`. The `A` in the +ablation names follows the design brief's `(A, F, L, Θ)` notation and should +not be read as a claim that species are generated. + +### What A0 already contains + +Stating this precisely matters, because every arm is a delta against it. + +**The two inherited processes.** The lattice is diffused not as `L` but as +`log S` where `S = (L Lᵀ)^{1/2}`, scaled by `N^{-1/3}` and flattened to a +Frobenius-preserving 6-vector; this is rotation-invariant, and `expm` maps any +point in `R^6` back to a valid cell, so noise cannot produce a degenerate +lattice. Standard DDPM [DDPM] on that 6-vector with the cosine ᾱ schedule +[iDDPM]; the network predicts ε. Fractional coordinates live on the torus, so +they use a wrapped-normal variance-exploding score process [NCSN] with a +geometric σ ladder from 0.005 to 0.5, `f_t = wrap(f_0 + σ_t z)`; the network +predicts the σ-scaled score, computed as a softmax-weighted mean over 11 +periodic images, and sampling uses a predictor–corrector scheme [ScoreSDE]. +Conditioning uses classifier-free guidance [CFG] with each modality dropped +independently. **None of this is touched by any arm below.** + +**The graph.** The pair graph is *dense* within each cell — every ordered pair +including self-pairs, resolved to the minimum image over the 27 offsets in +{−1,0,1}³, with self-pairs forced off the zero image so a one-atom cell still +carries information about its own translates. Cells are small, and at high +noise a radius graph on the pair channel would be arbitrary, so density is a +deliberate choice rather than an oversight. + +**The line graph.** A bond may participate in triplets if it is among the `knn` += 12 shortest bonds incident on its destination atom (`_knn_mask`, ranked by +destination). For each surviving bond `A = (u→v)` and each surviving bond +`B = (v→w)`, one line-graph edge is emitted; the shared atom is `v`. Triplets +are **not** deduplicated and back-tracking triplets `i→j→i` are **not** +excluded — measured on a 3-atom dense graph, 9 of 27 triplets are +back-tracking, each contributing exactly `cos θ = −1`. This is inherited +behaviour and is identical in every arm, but it matters when reading generated +angle histograms, which will therefore carry a spike at 180°. + +**The angular feature.** `cos θ` from `torch_bond_cosines`, expanded in 40 +fixed Gaussian RBFs over [−1, 1], then two MLP layers to hidden width — ALIGNN's +own representation [ALIGNN], unchanged in every primary arm. + +**The convolution.** Edge-gated graph convolution [GatedGCN] as ALIGNN uses it: + +``` +m_ij = W_s h_i + W_d h_j + W_e y_ij +σ_ij = sigmoid(m_ij) +h'_i = W_1 h_i + ( Σ_j σ_ij ⊙ W_2 h_j ) / ( Σ_j σ_ij + 1e-6 ) +y'_ij = y_ij + SiLU(LayerNorm(m_ij)) +``` + +The same class is instantiated twice per ALIGNN layer: once on the atom graph +(atoms ← bonds) and once on the line graph (bonds ← triplets), which is how +`Θ → L(G) → G` propagation happens. Note `LayerNorm`, not `BatchNorm` — the +normalisation is per element, so no cross-edge statistic exists. That fact is +load-bearing for the continuity argument in §"Continuous topology" below. + +**The heads.** The coordinate score is *not* an MLP on the node feature. Because +the edge-gated convolution lets an edge feature only gate a source-node feature, +an aggregated message carries magnitude far more readily than direction, and +such a head does not learn. Instead the score is assembled from the edge vectors +themselves, `score_i = Σ_j w_ij(h_i, h_j, y_ij) · Δf_ij`, which is +direction-correct by construction and invariant to a global shift. The lattice +head is an MLP on the mean-pooled atom feature. + +### Addition 1 — the angular denoising channel + +#### What we take from FoldingDiff + +FoldingDiff [FoldingDiff] establishes that bond angles can be the variables of +a diffusion model rather than a derived quantity: it represents a protein +backbone in internal coordinates, corrupts those angles with wrapped noise, +trains a network to predict the angular noise, and shows the resulting samples +reproduce natural angular distributions. We take three specific things: -The hypothesis under test is narrow: **does explicit three-body geometric -denoising improve crystal generation when embedded in ALIGNN's -atom–bond–line-graph hierarchy?** It is a hypothesis, not a claim — nothing -here has been trained yet. +1. **that angles can be a denoising target at all** — the conceptual move; +2. **the wrapped residual** — an angular error is only ever defined modulo + 2π, so the loss must wrap before it penalises; +3. **the loss functional form** — smooth L1 on the wrapped residual, with + `β = 0.1π`. -### The angular channel +#### What we implement -The model gains one output head, reading the line-graph feature `z` of the -same backbone that already produces the coordinate score and the lattice -noise. Per triplet it predicts the angular displacement the forward process -introduced, +The network gains exactly one head: a two-layer MLP on the line-graph feature +`z` of the shared backbone, zero-initialised on the output layer so training +starts from a silent prediction (matching how the existing coordinate and +lattice heads are initialised). Per triplet it predicts a scalar `δ̂_ijk`. + +The target is the angular displacement the forward process actually produced: ``` -delta_ijk = wrap( theta_ijk(f_t, L_t) - theta_ijk(f_0, L_0) ) +θ_t = θ_ijk( f_t , L_t ) angle at the noised geometry +θ_0 = θ_ijk( f_0 , L_0 ) angle at the clean geometry, same triplet +δ = wrap( θ_t − θ_0 ) ∈ [−π, π) ``` -trained with a wrapped smooth-L1 loss (`beta = 0.1 pi`). Both the target and -the loss are FoldingDiff's, which runs DDPM-style corruption and denoising -directly on protein bond and dihedral angles and shows that a model with an -explicit angular channel reproduces natural angular distributions. - -**One honest deviation.** FoldingDiff can diffuse a genuinely persistent -`theta_t` because a protein backbone's internal-coordinate list is fixed: -residue *i* always has the same three angles. A crystal being denoised from -noise has no such list — the triplet set is a function of the coordinates and -changes as they move, so there is no independent angular state to noise. What -is implemented is the closest well-defined thing: the angular *target* is -computed on the triplet representation that exists at the current step, with -both angles evaluated on the same periodic-image identity -`(i, j, k, n_ji, n_jk)` so the difference measures the corruption of one -triplet rather than a change of neighbour. Angles are an explicit denoising -objective with their own head and their own loss; they are not an -independently-noised variable. No new SDE, angle manifold or schedule was -invented to paper over the difference. - -The angular representation itself is untouched — ALIGNN's Gaussian RBF over -`cos(theta)` — so that the experiment measures the *objective*, not a change -of basis. - -### Continuous topology - -At large `t` the coordinates are close to uniform, and a hard neighbour-rank -rule for "does this triplet exist" is unjustified: ranks swap constantly and -the line graph jumps. `topology="radius"` replaces the kNN rule with a radius -candidate set in which every pair carries a smooth relevance +and the loss is the relevance-weighted wrapped smooth L1 ``` -s_ij = u(r_ij ; r_c) DimeNet's polynomial envelope: u, u' and u'' - all vanish at r_c -s_ijk = s_ji * s_jk ReaxFF's treatment of valence angles: an angle - switches off smoothly as either bond dissociates +L_ang = Σ_T s_ijk · SmoothL1_β( wrap( δ̂ − δ ) ) / Σ_T s_ijk , β = 0.1π ``` -`u` is not re-derived — it is `CutoffPolynomial`, already in this repository -for the smooth property model. The weight multiplies the gate *before both* -sums of the edge-gated convolution's normalised average, which is the only -placement that makes a zero-weight edge exactly equivalent to a deleted one; -scaling the numerator alone would renormalise the survivors and would not be -continuous. Because `s` is exactly zero at and beyond `r_c`, restricting the -sparse graph to pairs inside `r_c` drops only terms already contributing -nothing, so a triplet can enter or leave without any jump — which is what the -tests check by sweeping an atom through the cutoff and watching the output. - -The graph is rebuilt from the current coordinates and lattice on every -forward pass, so the topology follows the geometry through reverse diffusion -on its own. There is no connectivity annealing schedule, no graph-temperature -term and no learned bondness network; `r_ij(t)` evolving is the whole -mechanism. - -`radius_cutoff` defaults to 5 Å — between this repository's own three-body -cutoff (3.5 Å) and DimeNet's molecular cutoff (5 Å), and close to the radius -the baseline's 12 nearest neighbours actually span, which keeps the kNN -comparison fair. - -### Ablations - -Ablations are the point, not an afterthought: a single "it improved" number -cannot separate the three claims. `alignn/inverse/ablations.py` holds the -configurations, `scripts/atombench/run_angle_ablation.sh` runs them with the -same split, optimiser, epoch budget and seed list across every arm. - -| | angular objective | topology | angle → bond coupling | -|---|---|---|---| -| **A0** baseline | – | kNN | – | -| **A1** | yes | kNN | yes | -| **A2** | – | smooth radius | – | -| **A3** proposed | yes | smooth radius | yes | -| **A4** control | yes | smooth radius | **cut** | -| **A6** | yes | smooth radius | yes, Fourier angle basis | - -Which contrast answers what: - -- **A0 vs A1** — does explicit angular denoising help on its own? -- **A0 vs A2** — or is the gain just from a better-behaved noisy graph? -- **A4 vs A3** — is the benefit genuine coupling, or would auxiliary - supervision on a shared trunk do as well? A4 keeps the angular features - evolving and supervised but zeroes their contribution to the bond - representation, so its coordinate/lattice pathway provably cannot see them - (a test perturbs the angle embedding and asserts A4's output does not move - while A3's does). -- **A1 vs A3** — the design brief's A5: hard kNN against the smooth radius - graph. It is a comparison between existing arms, not a seventh - configuration. -- **A3 vs A6** — does the angular basis matter? Run this *last*; mixing a - basis change into the primary experiment would make attribution impossible. +added to the existing objective with a fixed weight `angle_weight` (default +1.0, alongside `frac_weight` 10.0 and `lattice_weight` 1.0). + +#### Deviation 1 (important, and unavoidable) — the angular state is induced, not persistent + +**This is the main methodological caveat of the branch and should be stated in +any write-up.** FoldingDiff can diffuse a genuinely persistent `Θ_t` because a +protein backbone's internal-coordinate list is *fixed*: residue *i* always has +the same three angles, so the angles are legitimate independent state +variables. A crystal being denoised from noise has no such list — the set of +triplets is a *function of the coordinates* and changes as they move. There is +therefore no independent angular variable to noise, and no forward process +`q(Θ_t | Θ_0)` to write down. + +**What we do instead, and why.** The angular *target* is computed on the +triplet representation that exists at the current step. Angles remain an +explicit denoising objective with their own head and their own loss; they are +not an independently-noised variable. Concretely, the process is +`(f, L)`-driven and `Θ` is read off it, so the model is trained to report the +angular component of a corruption it is simultaneously being asked to undo in +Cartesian terms. This is the fallback the design brief specifies for exactly +this situation, and it is chosen over the alternative of inventing an angular +SDE with no literature standing. + +**Consequence to be honest about.** Because `Θ` is induced, `L_ang` is not an +independent diffusion loss with its own ELBO interpretation; it is a +geometrically-structured auxiliary objective on the same forward process. Any +claim in a paper should be phrased as "explicit angular supervision derived +from the joint process", not "we diffuse angles". + +#### Deviation 2 (ours) — fixed periodic-image identity for the target + +**Not from any cited source.** A triplet's identity in a periodic crystal is +`(i, j, k, n_ji, n_jk)`, where the `n` are integer cell offsets. When we +resolve the noised geometry to minimum image we record the integer offset +actually chosen, -```bash -bash scripts/atombench/run_angle_ablation.sh runs/data_jarvis runs/ablation 0 1 2 +``` +n = offset_argmin − round( f_t[dst] − f_t[src] ) +``` + +and compute `θ_0` by applying *that same* `n` to the clean coordinates, +`Δf_0 = f_0[dst] − f_0[src] + n`, rather than re-running minimum-image +resolution on the clean structure. The reason is that the two give different +answers at large `t`: re-resolving would compare the noised angle at one +neighbour against the clean angle at a *different* neighbour, so the target +would mix "this triplet bent" with "this is now a different triplet". Fixing +`n` makes `δ` measure the corruption of one specific triplet. As σ → 0 the two +constructions coincide, and a test asserts `|δ| < 1e−4` when `f_t = f_0`. + +#### Deviation 3 (ours) — relevance weighting of the loss + +**Not from FoldingDiff**, which has no per-angle weights because its angle set +is fixed. Ours is `Σ s·ℓ / Σ s`. This is what makes the objective continuous +when a triplet enters or leaves the sparse graph: a triplet at the cutoff has +weight zero and contributes nothing on either side of the boundary. In kNN +arms all weights are 1 and the expression reduces to a plain mean. + +#### A note on wrapping + +Bond angles are in `[0, π]`, so `θ_t − θ_0 ∈ [−π, π]` already, and `wrap` is +the identity except exactly at the boundary. It is applied anyway so the +objective is the wrapped one by construction rather than by an argument about +ranges, and so the same helper serves a genuinely circular variable if a +dihedral channel is ever added. `acos` is clamped at `±(1 − 1e−7)` to keep the +gradient finite for collinear triplets, which — see the back-tracking note +above — are common; the cost is ≈ 0.03° of accuracy at the poles. + +### Addition 2 — continuously weighted topology + +#### What we take from DimeNet + +DimeNet [DimeNet] introduces the polynomial envelope + +``` +u(x) = 1 + a x^p + b x^(p+1) + c x^(p+2) , x = r / r_c +a = −(p+1)(p+2)/2 , b = p(p+2) , c = −p(p+1)/2 +u(x) = 0 for x ≥ 1 ``` -Given the spread already recorded above — match rate across fifteen -independently trained models spanned 0.437–0.524 — one seed per arm cannot -settle any of these. The runner defaults to three. +for which `u(0) = 1` and `u(1) = u'(1) = u''(1) = 0`. We take the envelope +exactly, and we do not re-derive it: this repository already ships it as +`CutoffPolynomial` for the smooth property model, and we import that class. + +**Convention note, worth stating precisely because the two differ.** This +repository parameterises by `coeff`, which *is* the paper's `p`; the widely +used reference implementation parameterises by `envelope_exponent` and sets +`p = exponent + 1`. We use `coeff = 5`, i.e. `p = 5`. Both satisfy the +vanishing conditions (verified by autograd in the test suite); they differ only +in how fast the envelope decays — at `r = r_c/2`, `u = 0.7734` for `p = 5` +against `0.8555` for `p = 6`. Also note DimeNet's *code* folds a factor `1/x` +into its envelope so that it can multiply a Bessel basis; the polynomial above +is the paper's `u(d)`, which is the form we want for a weight in `[0,1]`. + +#### What we take from ReaxFF + +ReaxFF [ReaxFF] makes bond order a continuous function of interatomic +distance and multiplies every valence-angle energy term by switching factors +for *both* of its constituent bonds, so an angular interaction disappears +smoothly when either bond dissociates and appears smoothly when one forms. We +take the product structure of that gate and nothing else — no ReaxFF +parameters, no bond-order function, no element-specific tables: -### Evaluation +``` +s_ij = u( r_ij ; r_c ) pair relevance [DimeNet] +s_ijk = s_ji · s_jk triplet relevance [ReaxFF] +``` + +The behaviour we want during reverse diffusion follows directly: as `r_ji` +grows, `s_ji → 0` and every triplet through that bond fades out; as some other +`r_jl` shrinks, new triplets fade in. The effective line graph changes +continuously because `r(t)` changes, with no annealing schedule, no +graph-temperature term and no learned bondness network. + +#### Deviation 4 (ours) — where the gate is applied, and why the placement is forced -Every existing AtomBench metric is preserved. Two mechanism metrics are added, -defined before the runs so that a favourable one cannot be picked afterwards: +**Not specified by either source.** ReaxFF gates a physical energy term; +DimeNet's envelope multiplies a radial basis. Neither says what to do with a +*learned, normalised, gated message*. In the edge-gated convolution the +aggregation is a normalised average, and the weight must therefore multiply +`σ_ij` **before both sums**: -- **bond-angle distributions**, generated against held-out real structures - (KL, Jensen-Shannon, and a 1-D Wasserstein distance reported in degrees) — - FoldingDiff's own diagnostic; -- **relaxation displacement**, how far a sample moves to reach its nearest - ALIGNN-FF local minimum, with the volume change and energy drop — the - proximity-to-local-minimum evaluation MatterGen uses. If explicit angular - denoising produces locally coherent geometry, its samples should need less - repair. +``` +h'_i = W_1 h_i + ( Σ_j w_ij σ_ij ⊙ W_2 h_j ) / ( Σ_j w_ij σ_ij + 1e-6 ) +``` + +This is the only placement with the property we need. Setting `w = 0` makes the +term contribute nothing to numerator *and* denominator, so the result is +*exactly* what the same layer computes on a graph with that edge physically +removed — verified by a test that compares the two. Weighting the numerator +alone would renormalise the surviving messages and would be discontinuous in +exactly the situation the whole construction exists to avoid. Because +`LayerNorm` is per element, no batch statistic can smuggle the deleted edge +back in. + +Three places carry the weight, all with the same scalar: + +| where | weight | why | +|---|---|---| +| atom-graph messages | `s_ij` | pair channel fades with distance | +| line-graph messages | `s_ijk` | triplet channel fades with either bond | +| coordinate-score head, `w_ij ← w_ij · s_ij` | `s_ij` | a pair leaving the cutoff must stop pushing the atom smoothly, not abruptly | + +#### Deviation 5 (ours) — the sparse graph is an exact truncation, not an approximation + +**Not from either source.** Because `s` is *exactly* zero at and beyond `r_c` +— not merely small — restricting the sparse line graph to pairs with +`s_ij > 0` removes only terms already contributing exactly nothing. The radius +graph is therefore an exact sparsification of a gated dense graph, not an +approximation of it, and triplet insertion/deletion is invisible by +construction rather than by tolerance. The test suite sweeps an atom straight +through `r_c` in 121 steps, confirms the triplet count really does change +during the sweep, and asserts no step in the output exceeds 20× the median +step. + +The pair graph itself stays dense in the edge list, with `s_ij` doing the +truncation numerically. That is a deliberate choice: it keeps the sparse +structure identical to A0's, so nothing about batching or indexing differs +between arms, and it removes any risk of an atom being left with no pair +edges at all in a large cell. + +#### Deviation 6 (ours) — the graph is rebuilt every forward pass + +Every distance, every envelope value, the `allowed` mask and the entire line +graph are recomputed inside `forward()` from the current `(f, L)`. During +reverse diffusion that means the topology is re-derived at each of the T +denoising steps at no extra bookkeeping cost. The pair index is composition- +determined and is built once. This satisfies the "rebuild the radius graph +during reverse diffusion" requirement without any sampler changes. + +**Compute cost, measured.** On a 16-atom 7 Å cell the kNN construction yields +2304 triplets and the radius construction at `r_c = 5 Å` yields 3262 — about +1.4×. Triplet count scales as `Σ_j deg(j)²`, so this ratio grows with cell +size and with `r_c`; budget for it when matching wall-clock across arms. + +#### The choice of `r_c` + +`radius_cutoff` defaults to 5.0 Å. Justification, in order of weight: it is +close to the radius that the baseline's twelve nearest neighbours actually +span in a dense crystal, which is what makes A1 vs A3 a comparison of +*smoothness* rather than of *interaction range*; it sits just above this +repository's own three-body cutoff of 3.5 Å; and it matches DimeNet's +molecular cutoff. It is a free parameter and should be held fixed across every +arm of a comparison. + +### The arms + +Each arm is a keyword dict in `alignn/inverse/ablations.py`. They differ only +in the switches named; width, depth, schedule, optimiser, splits and seeds come +from the training script and must be identical across a comparison. + +#### A0 — baseline + +`angle_diffusion=False, topology="knn", gate_pair_messages=False, +angle_feedback=True` + +The model that exists on `develop`, unchanged, as described in "What A0 already +contains". Angles are an input feature; the outputs are the coordinate score +and the lattice noise and nothing else. Its first job is to reproduce the +published numbers — match 0.485 on Alexandria, 0.524 on JARVIS Supercon-3D — +before any comparison is trusted. Its second job is as the correctness anchor +for the branch: a test loads A0's state dict into A1 and asserts `eps_frac` +and `eps_lattice` come out bit-identical, which is what establishes that the +angular channel is an *addition* to the model rather than a perturbation of +it. A0 is also what a user gets from `ALIGNNCSPDenoiser()` with no arguments, +so every released checkpoint keeps loading and behaving exactly as before. + +#### A1 — angular denoising only + +`angle_diffusion=True`, graph left at baseline. + +Adds the angle head and `L_ang`; changes nothing about the graph. Same kNN +membership rule, same triplets, same 40-bin cosine RBF, same convolutions, +same everything, plus one two-layer MLP. Because the architectural delta is a +single head that feeds no other computation in the forward direction, a +difference between A0 and A1 is attributable to the *objective* and to nothing +else — this is the arm that answers question 1 in isolation, and it is +deliberately built on the baseline's topology so that question 1 and question +2 cannot contaminate each other. The angular information does reach the +structural heads, but only through ALIGNN's ordinary path, and only because +the loss reshapes the shared trunk. One caveat to report alongside the result: +adding a loss term changes the gradient balance, so `angle_weight` must be +fixed across arms and stated, or the comparison silently becomes a +hyperparameter search. + +#### A2 — smooth topology only + +`topology="radius", gate_pair_messages=True`, no angle head. + +Replaces the hard kNN membership rule with a radius candidate set, gates +triplet messages by `s_ijk`, and gates atom-graph messages and the per-edge +terms of the coordinate score by `s_ij`. No angle is ever denoised; the +outputs are still just the coordinate score and the lattice noise. This arm +exists to kill the boring explanation. Making the high-noise graph continuous +is a substantial change to the model's inductive bias on its own — at large +`t` the coordinates are near-uniform and neighbour ranks swap constantly — so +if A3 beats A0 and A2 beats A0 by the same margin, the angular objective +contributed nothing and the honest conclusion is that the topology did the +work. A2 answers question 2 with the angular channel held off. + +#### A3 — the proposed model + +`angle_diffusion=True, topology="radius", gate_pair_messages=True, +angle_feedback=True` + +Both mechanisms together, and the arm the hypothesis is about. The angular +channel is *coupled*, not bolted on: `z` reaches the coordinate and lattice +heads through ALIGNN's own hierarchy, because the line-graph convolution's +triplet feature gates the bond→bond message and the updated bond features then +gate the bond→atom message on the next layer — `Θ → L(G) → G → (f, L)` heads, +exactly the mechanism ALIGNN already provides [ALIGNN]. No Jacobian +correction, no geometric force term, no hand-designed angle-to-coordinate +update. A3 beating A0 is the headline number, but on its own it says only that +the combination helps; which half is responsible is settled by A1, A2 and A4, +not by A3. + +#### A4 — the coupling control + +`angle_diffusion=True, topology="radius", gate_pair_messages=True, +angle_feedback=False` + +Identical to A3, including parameter count, except that the triplet weight +entering the *bond aggregation* is set to zero. Mechanically, inside the +line-graph convolution: + +``` +m_T = W_s m[lg_src] + W_d m[lg_dst] + W_e z still contains z +σ_T = sigmoid(m_T) · 0 gate forced to zero +y' = SiLU(LayerNorm(W_1 m + 0/(0+1e-6))) + m bonds see no angles +z' = SiLU(LayerNorm(m_T)) + z angles still evolve +``` + +So the angular features are still computed, still updated by every ALIGNN +layer, still feed the angle head, and the angular loss still back-propagates +into the shared trunk through `m_T` — that last part is deliberate, because +gradient flow into a shared trunk is precisely what auxiliary multitask +learning *is*. What is removed is the forward path by which the angular latent +reaches the bond representation, and hence the atom, coordinate and lattice +pathway. A4 is therefore "angles supervised alongside the model" and A3 is +"angles wired into it", which is exactly the contrast question 3 asks about. +If A3 ≈ A4, the mechanism claim fails even if the benchmark number improved. + +Two things to keep straight. First, the cut is verified rather than asserted: a +test perturbs the angle-embedding weights and confirms A3's coordinate score +moves while A4's stays at *exactly* zero. Second, **A4's structural trunk is +not identical to A0's** — it is A3's with one aggregation zeroed, which is a +different function from A0's full triplet aggregation. A4 is a control for A3 +and must not be read as a second baseline. + +#### A5 — hard kNN versus smooth radius + +**Not a configuration.** This is a comparison between arms that already exist, +which is why there is no `"A5"` key. It is answered twice: A1 against A3 with +the angular objective on, and A0 against A2 with it off. Running both is worth +the compute because they can disagree informatively — if continuous topology +matters *only* when something is being denoised on the triplets, then A0↔A2 +moves little while A1↔A3 moves a lot, and that pattern is itself evidence for +the coupling story rather than for topology as a standalone improvement. + +#### A6 — angular basis + +`angle_basis="fourier"` on top of A3. + +Swaps ALIGNN's 40-bin Gaussian RBF over `cos θ` for the learnable Fourier +basis on `θ` this repository already ships (`FourierAngular`, order +`(triplet_bins − 1)//2`, giving `1 + 2·order` features). Sequenced last on +purpose: changing the angular representation at the same time as introducing +the angular objective would make attribution impossible, which is why every +primary arm keeps ALIGNN's basis untouched [ALIGNN]. + +**State its limit plainly.** DimeNet's joint spherical Fourier–Bessel `(d, θ)` +basis is **not** implemented. DimeNet reports that a joint distance–angle basis +is a stronger inductive bias than the raw angle [DimeNet], and that remains the +interesting version of this question, but it requires spherical Bessel zeros +and a new dependency, and the design brief defers the whole basis question to +last. A6 therefore answers "does the angular representation matter at all", +not "is DimeNet's SBF better". The latter is open work. + +### Confounds and threats to validity + +**Interaction range is confounded with topology smoothness in A2/A3/A4.** With +`gate_pair_messages=True` and `r_c = 5 Å`, pairs beyond 5 Å contribute exactly +zero to the atom-graph messages and to the coordinate score, whereas A0's dense +pair graph lets every atom see every other atom in the cell. So the smooth arms +differ from A0 in two ways at once: the topology is continuous *and* the +interaction range is truncated. This is what the design brief asks for +("smoothly vanishing pair interactions"), and it is defensible, but it is a +confound and should be reported as one. + +It is separable without any code change. `--ablation A3 --gate-pair-messages 0` +gives smooth *triplet* topology with A0's ungated dense pair channel — the +line graph is still radius-based and `s_ijk`-weighted, but the pair channel is +untouched. If range truncation is doing the work, that arm will sit with A0; +if smoothness is, it will sit with A3. We recommend running it for A3 at +minimum before writing anything up. + +**Back-tracking triplets inflate the 180° bin.** Inherited from the graph +builder and identical across arms, so it does not bias a comparison, but a +generated-versus-real angle histogram will show a large spike at 180° in both +distributions. Do not present that spike as physics. + +**`angle_weight` is unswept.** It is a free hyperparameter at 1.0. A negative +result for A1 or A3 at one loss weight is weak evidence; if the primary arms +come out flat, sweep it before concluding, and report the sweep. + +**Validation loss is a poor proxy here.** Already documented above in this +README: a 14.5% denoising-loss gap produced a non-significant RMSD gap and zero +change in match rate. Expect `L_ang` to fall in the angular arms *by +construction* — it is a new term being optimised — and do not report that as +evidence of anything. + +**Seed spread.** Match rate across fifteen independently trained models spanned +0.437–0.524, nine structures on a 103-target split. One seed per arm can invert +any conclusion here. The runner defaults to three; more is better. + +### Protocol and metrics + +`scripts/atombench/run_angle_ablation.sh [seeds...]` runs +every arm with the same split, optimiser settings, epoch budget and seed list, +then generates, scores and computes the mechanism metrics. + +All existing AtomBench metrics are preserved unchanged: match rate, Cartesian +RMSD, ccRMSD, lattice-parameter and lattice-angle MAE, KLD. Two mechanism +metrics are added, and **the suite was fixed before any run precisely so that a +favourable metric cannot be selected afterwards**: + +**Bond-angle distributions** — generated against held-out real structures, +pooled over the split, on a common 180-bin histogram over [0°, 180°], reported +as KL, Jensen–Shannon and 1-D Wasserstein distance. The Wasserstein figure is +exact for a 1-D histogram (the integral of the absolute CDF difference) and is +in degrees, so it reads directly as "the generated angles are off by this +much"; it is calibrated to within 0.001° on a synthetic shift test. This is +FoldingDiff's own diagnostic [FoldingDiff] and is the most direct check that +the angular channel does what it claims. + +**Relaxation displacement** — how far a sample must move to reach the nearest +ALIGNN-FF local minimum, as translation-corrected Cartesian RMSD, plus +fractional volume change and energy drop. MatterGen [MatterGen] evaluates +generated structures by their proximity to their relaxed counterparts; if +explicit angular denoising produces locally coherent geometry, its samples +should need less geometric repair, and this is the metric that would show it. ```bash python scripts/atombench/angle_eval.py runs/ablation/A3_s0/bench.csv --relax ``` -### What is deliberately not here - -DimeNet's joint spherical Fourier–Bessel distance–angle basis is **not** -implemented; A6 substitutes the Fourier basis on theta this repository -already ships. The full SBF needs spherical Bessel zeros and a new -dependency, and the design brief sequences that ablation last, so it is -deferred rather than half-done. Also absent, deliberately: cross-attention, -separate networks per variable, a learned bondness classifier, a -time-dependent graph-temperature schedule, Jacobian angle forces, and any -loss whose role cannot be traced to one of the papers below. - -Note that atom types are *not* diffused — the generator is conditioned on -composition — so the state is really `(F, L)` and, with this extension, -`(F, L, Theta)`. The `A` in the ablation names follows the design brief's -notation. - -### References - -| | | -|---|---| -| ALIGNN — atom graph + line graph, angles update bonds update atoms | Choudhary & DeCost, *npj Comput. Mater.* **7**, 185 (2021), [10.1038/s41524-021-00650-1](https://doi.org/10.1038/s41524-021-00650-1) | -| FoldingDiff — diffusion directly on bond angles; wrapped noise, wrapped smooth-L1, angle-distribution evaluation | Wu *et al.*, *Nat. Commun.* **15**, 1059 (2024), [10.1038/s41467-024-45051-2](https://doi.org/10.1038/s41467-024-45051-2) | -| Torsional Diffusion — diffusion on angular configuration spaces | Jing *et al.*, NeurIPS 2022 | -| ReaxFF — continuous distance-dependent bond order; angle terms vanish as either bond dissociates | van Duin *et al.*, *J. Phys. Chem. A* **105**, 9396 (2001), [10.1021/jp004368u](https://doi.org/10.1021/jp004368u) | -| DimeNet — cutoff envelope whose value and first two derivatives vanish at `r_c` | Gasteiger, Groß & Günnemann, ICLR 2020, [arXiv:2003.03123](https://arxiv.org/abs/2003.03123) | -| MatterGen — one score network denoising crystal variables jointly; proximity-to-relaxed evaluation | Zeni *et al.*, *Nature* **639**, 624 (2025), [10.1038/s41586-025-08628-5](https://doi.org/10.1038/s41586-025-08628-5) | -| DiffCSP — joint equivariant diffusion of lattice and fractional coordinates | Jiao *et al.*, NeurIPS 2023 | -| CrystalDiT — the warning against unnecessary multi-stream architectural complexity | Yi *et al.*, AAAI 2026, [10.1609/aaai.v40i2.37121](https://doi.org/10.1609/aaai.v40i2.37121) | +### Provenance summary + +| design choice | source | taken from the source | ours | +|---|---|---|---| +| atom + line graph, angles→bonds→atoms | [ALIGNN] | the whole hierarchy, the cosine-RBF angular feature, the edge-gated conv | nothing — used as-is | +| angles as a denoising target | [FoldingDiff] | the conceptual move, the wrapped residual, smooth-L1 with β = 0.1π | applying it where the angle set is not persistent | +| angular target definition | — | — | **ours**: `δ = wrap(θ_t − θ_0)` on the current triplet set, with a fixed periodic-image identity | +| loss weighting | — | — | **ours**: `Σ s·ℓ / Σ s`, needed for continuity at the cutoff | +| cutoff envelope | [DimeNet] | the polynomial `u`, imported from this repo's existing `CutoffPolynomial` | using it as a *topology* weight rather than a radial-basis multiplier | +| triplet gate | [ReaxFF] | the product structure `s_ijk = s_ji·s_jk` and the fade-in/fade-out behaviour | no ReaxFF parameters or bond-order function are used | +| gate placement | — | — | **ours**: multiply `σ` before *both* sums; forced by the exact-deletion requirement | +| radius truncation | — | — | **ours**: exact sparsification (`s = 0` beyond `r_c`), not an approximation | +| one shared backbone, extra head | [MatterGen], [CrystalDiT] | joint denoising through one network; the warning against multi-stream complexity | — | +| crystal diffusion formulation | [DiffCSP], [DDPM], [iDDPM], [NCSN], [ScoreSDE], [CFG] | the inherited `(f, L)` processes, schedules and guidance | untouched by this branch | +| angle-distribution metric | [FoldingDiff] | the diagnostic | 1-D Wasserstein in degrees as the headline figure | +| relaxation-proximity metric | [MatterGen] | the evaluation | translation-corrected RMSD + volume change + energy drop | +| A4 control design | — | — | **ours**: sever coupling by zeroing the triplet aggregation weight, keeping parameters and gradient flow identical | +| A6 angular basis | [DimeNet] | the motivation for a richer basis | uses this repo's Fourier basis; DimeNet's SBF is **not** implemented | + +### Bibliography + +- **[ALIGNN]** K. Choudhary and B. DeCost, "Atomistic Line Graph Neural + Network for improved materials property predictions," *npj Computational + Materials* **7**, 185 (2021). doi:10.1038/s41524-021-00650-1 +- **[FoldingDiff]** K. E. Wu, K. K. Yang, R. van den Berg, S. Alamdari, + J. Y. Zou, A. X. Lu and A. P. Amini, "Protein structure generation via + folding diffusion," *Nature Communications* **15** (2024). + doi:10.1038/s41467-024-45051-2 +- **[TorsDiff]** B. Jing, G. Corso, J. Chang, R. Barzilay and T. Jaakkola, + "Torsional Diffusion for Molecular Conformer Generation," *NeurIPS* (2022). +- **[ReaxFF]** A. C. T. van Duin, S. Dasgupta, F. Lorant and W. A. Goddard III, + "ReaxFF: A Reactive Force Field for Hydrocarbons," *J. Phys. Chem. A* **105**, + 9396–9409 (2001). doi:10.1021/jp004368u +- **[DimeNet]** J. Gasteiger (Klicpera), J. Groß and S. Günnemann, + "Directional Message Passing for Molecular Graphs," *ICLR* (2020). + arXiv:2003.03123 +- **[MatterGen]** C. Zeni *et al.*, "A generative model for inorganic materials + design," *Nature* (2025). doi:10.1038/s41586-025-08628-5 +- **[DiffCSP]** R. Jiao, W. Huang, P. Lin, J. Han, P. Chen, Y. Lu and Y. Liu, + "Crystal Structure Prediction by Joint Equivariant Diffusion," *NeurIPS* + (2023). +- **[CrystalDiT]** Yi *et al.*, "CrystalDiT: Simple Diffusion Transformers for + Crystal Generation," *AAAI* (2026). doi:10.1609/aaai.v40i2.37121 +- **[DDPM]** J. Ho, A. Jain and P. Abbeel, "Denoising Diffusion Probabilistic + Models," *NeurIPS* (2020). +- **[iDDPM]** A. Nichol and P. Dhariwal, "Improved Denoising Diffusion + Probabilistic Models," *ICML* (2021). — source of the cosine ᾱ schedule +- **[NCSN]** Y. Song and S. Ermon, "Generative Modeling by Estimating Gradients + of the Data Distribution," *NeurIPS* (2019). +- **[ScoreSDE]** Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon + and B. Poole, "Score-Based Generative Modeling through Stochastic + Differential Equations," *ICLR* (2021). — source of the predictor–corrector + sampler +- **[CFG]** J. Ho and T. Salimans, "Classifier-Free Diffusion Guidance," + (2022). arXiv:2207.12598 +- **[GatedGCN]** X. Bresson and T. Laurent, "Residual Gated Graph ConvNets," + (2017). arXiv:1711.07553 — the edge-gated convolution ALIGNN builds on +- **[SOAP]** A. P. Bartók, R. Kondor and G. Csányi, "On representing chemical + environments," *Phys. Rev. B* **87**, 184115 (2013). + doi:10.1103/PhysRevB.87.184115 — background only; not used here + +**Three citation details to verify against the primary sources before this +reaches a manuscript**, because they were carried over from the design brief or +recalled rather than checked: FoldingDiff's exact smooth-L1 `β` (we use +`0.1π` and attribute the *form* to FoldingDiff with confidence, the constant +with less); the CrystalDiT DOI and venue; and the author lists for +[MatterGen] and [DiffCSP], which are given here in abbreviated form. ## Relax and rank