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/README.md b/README.md
index 11b9e97..9ea0a53 100644
--- a/README.md
+++ b/README.md
@@ -324,6 +324,13 @@ at the trained 1000 steps, ~4.6 s at 200). Training your own model and
reproducing the AtomBench benchmarks is covered in
[alignn/inverse/README.md](alignn/inverse/README.md).
+That page also carries a full reference for an optional research extension,
+off by default: making bond angles an explicit **denoising channel** rather
+than only an input feature, with a line-graph topology that varies
+continuously as the coordinates denoise. See *Explicit bond-angle diffusion —
+reference* for the mathematics, the per-ablation breakdown, what each design
+choice takes from the literature, and the bibliography.
+
## Performances
diff --git a/alignn/inverse/README.md b/alignn/inverse/README.md
index d68a65c..7e6fe50 100644
--- a/alignn/inverse/README.md
+++ b/alignn/inverse/README.md
@@ -168,6 +168,680 @@ 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 — 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
+
+```
+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:
+
+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π`.
+
+#### What we implement
+
+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:
+
+```
+θ_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 ) ∈ [−π, π)
+```
+
+and the loss is the relevance-weighted wrapped smooth L1
+
+```
+L_ang = Σ_T s_ijk · SmoothL1_β( wrap( δ̂ − δ ) ) / Σ_T s_ijk , β = 0.1π
+```
+
+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,
+
+```
+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
+```
+
+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:
+
+```
+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
+
+**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**:
+
+```
+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
+```
+
+### 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
`relax_rank.py` refines candidates with the pretrained ALIGNN force field.
@@ -206,6 +880,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/__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..982de21
--- /dev/null
+++ b/alignn/inverse/angles.py
@@ -0,0 +1,200 @@
+"""Angular geometry and continuously-weighted topology for ALIGNN-CSP.
+
+Two pieces of machinery live here, and they are deliberately independent so
+that the ablations in :mod:`alignn.inverse.ablations` can switch one on
+without the other.
+
+**Explicit angular denoising.** ALIGNN already carries bond angles on its
+line graph, but only as an *input feature*. Here they also become a
+*denoising target*: the network predicts, per triplet, the angular
+displacement that the forward process introduced. The stochastic process and
+the loss are ported from FoldingDiff (Wu et al., Nat. Commun. 15, 1059, 2024,
+doi:10.1038/s41467-024-45051-2), which runs DDPM-style corruption and
+denoising directly on protein bond and dihedral angles with wrapped angular
+noise and a wrapped smooth-L1 objective. Torsional Diffusion (Jing et al.,
+NeurIPS 2022) is the general statement that a diffusion process can be defined
+on an angular configuration space.
+
+One difference from FoldingDiff has to be stated plainly, because it is the
+main methodological caveat of this extension. In a protein backbone the
+internal-coordinate list is *fixed*: residue i always has the same three
+angles, so a genuinely persistent state ``theta_t`` can be diffused
+independently of anything else. In a crystal being denoised from noise the
+triplet set is not fixed — it is a function of the coordinates, and it changes
+as they move. There is therefore no persistent ``theta_t`` to diffuse. What
+is implemented instead is the closest well-defined thing: the angular
+denoising *target* is computed on the triplet representation that exists at
+the current step,
+
+ delta_ijk = wrap(theta_ijk(f_t, L_t) - theta_ijk(f_0, L_0)),
+
+with both angles evaluated on the *same* periodic-image identity
+``(i, j, k, n_ji, n_jk)`` so that the difference measures the corruption of
+one specific triplet rather than a change of neighbour. Angles are still an
+explicit denoising channel with their own head and their own loss; they are
+not an independently-noised variable. Section 3 of the design brief asks for
+exactly this fallback, and asks that the distinction be documented rather than
+papered over with an invented process.
+
+**Continuously-weighted topology.** During reverse diffusion a hard
+neighbour-rank criterion for "does this triplet exist" is unjustified: at
+large ``t`` the coordinates are close to uniform, so neighbour ranks swap
+constantly and the line graph jumps discontinuously. Instead every pair
+carries a smooth relevance
+
+ s_ij = u(r_ij ; r_c),
+
+where ``u`` is the polynomial cutoff envelope introduced by DimeNet
+(Gasteiger, Gross & Gunnemann, ICLR 2020, arXiv:2003.03123), whose value and
+first two derivatives vanish at ``r_c``. That envelope is already implemented
+in this repository as
+:class:`alignn.models.alignn_atomwise_pure_smooth.CutoffPolynomial`, so it is
+reused rather than re-derived. A triplet inherits the product of its two
+constituent relevances,
+
+ s_ijk = s_ji * s_jk,
+
+which is the ReaxFF treatment of valence angles (van Duin et al.,
+J. Phys. Chem. A 105, 9396, 2001, doi:10.1021/jp004368u): bond orders vary
+continuously with distance and an angular term switches off smoothly as
+either of its bonds dissociates. Because ``s`` is exactly zero at and beyond
+``r_c``, restricting the sparse line graph to pairs inside ``r_c`` removes
+only terms that were already contributing nothing — a triplet can enter or
+leave the computational graph without any finite jump in the messages.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+
+from alignn.models.alignn_atomwise_pure_smooth import CutoffPolynomial
+from alignn.torch_graph_builder import torch_bond_cosines
+
+__all__ = [
+ "CutoffPolynomial",
+ "bond_angle",
+ "wrap_angle",
+ "pair_relevance",
+ "triplet_relevance",
+ "edge_vectors",
+ "angular_denoising_target",
+ "angle_denoising_loss",
+ "TWO_PI",
+]
+
+TWO_PI = 2.0 * math.pi
+
+# acos is not differentiable at +-1; the clamp keeps the gradient finite for
+# collinear triplets, which are common in a crystal (i -> j -> i back-tracking
+# triplets are cos = -1 exactly).
+_COS_EPS = 1.0e-7
+
+
+def bond_angle(r_ij: torch.Tensor, r_jk: torch.Tensor) -> torch.Tensor:
+ """Bond angle at the shared atom ``j`` of a triplet, in radians.
+
+ Uses ALIGNN's own cosine convention (:func:`torch_bond_cosines`), so the
+ angle is the interior angle at ``j`` and lies in ``[0, pi]``.
+ """
+ cos = torch_bond_cosines(r_ij, r_jk)
+ return torch.acos(cos.clamp(-1.0 + _COS_EPS, 1.0 - _COS_EPS))
+
+
+def wrap_angle(x: torch.Tensor) -> torch.Tensor:
+ """Wrap an angular difference into ``[-pi, pi)``.
+
+ FoldingDiff's forward process and loss are both defined modulo ``2 pi``.
+ Bond angles themselves live in ``[0, pi]``, so a difference of two of them
+ is already inside ``[-pi, pi]`` and this is a no-op up to the boundary
+ case; it is applied anyway so that the objective is the wrapped one by
+ construction rather than by an argument about ranges, and so that the same
+ helper can serve a periodic angular variable if one is ever added.
+ """
+ return torch.remainder(x + math.pi, TWO_PI) - math.pi
+
+
+def pair_relevance(
+ dist: torch.Tensor, envelope: CutoffPolynomial
+) -> torch.Tensor:
+ """Smooth per-pair relevance ``s_ij = u(r_ij; r_c)`` in ``[0, 1]``.
+
+ ``u`` is the DimeNet envelope: ``u(0) = 1`` and ``u`` together with its
+ first two derivatives vanishes at ``r_c``.
+ """
+ return envelope(dist)
+
+
+def triplet_relevance(
+ s_edge: torch.Tensor, lg_src: torch.Tensor, lg_dst: torch.Tensor
+) -> torch.Tensor:
+ """ReaxFF-style product gate ``s_ijk = s_ji * s_jk`` for each triplet."""
+ return s_edge[lg_src] * s_edge[lg_dst]
+
+
+def edge_vectors(
+ frac: torch.Tensor,
+ lattice: torch.Tensor,
+ src: torch.Tensor,
+ dst: torch.Tensor,
+ edge_graph_id: torch.Tensor,
+ image: torch.Tensor,
+) -> torch.Tensor:
+ """Cartesian edge vectors for a *given* set of periodic images.
+
+ ``image`` is the integer cell offset ``n`` that the minimum-image search
+ settled on at the noised geometry, so that ``delta_f = f[dst] - f[src] +
+ n``. Re-using the same ``n`` on a different (clean) structure is what
+ makes the angular target a corruption of one fixed triplet identity rather
+ than a comparison between two different neighbours.
+ """
+ df = frac[dst] - frac[src] + image
+ return torch.einsum("ei,eij->ej", df, lattice[edge_graph_id])
+
+
+def angular_denoising_target(
+ angle_t: torch.Tensor,
+ frac0: torch.Tensor,
+ lattice0: torch.Tensor,
+ src: torch.Tensor,
+ dst: torch.Tensor,
+ edge_graph_id: torch.Tensor,
+ image: torch.Tensor,
+ lg_src: torch.Tensor,
+ lg_dst: torch.Tensor,
+) -> torch.Tensor:
+ """Angular displacement the forward process applied to each triplet.
+
+ Returns ``wrap(theta_t - theta_0)``, the angular analogue of the noise
+ ``eps`` that FoldingDiff's network predicts.
+ """
+ r0 = edge_vectors(frac0, lattice0, src, dst, edge_graph_id, image)
+ theta0 = bond_angle(r0[lg_src], r0[lg_dst])
+ return wrap_angle(angle_t - theta0)
+
+
+def angle_denoising_loss(
+ pred: torch.Tensor,
+ target: torch.Tensor,
+ weight: Optional[torch.Tensor] = None,
+ beta: float = 0.1 * math.pi,
+) -> torch.Tensor:
+ """Relevance-weighted wrapped smooth-L1 loss on the angular channel.
+
+ The functional form — smooth L1 of the *wrapped* residual, with
+ ``beta = 0.1 pi`` — is FoldingDiff's angular objective. The weighting by
+ ``s_ijk`` is what makes the loss continuous when a triplet enters or
+ leaves the sparse line graph: a triplet at the cutoff has zero weight, so
+ it contributes nothing on either side of the boundary.
+ """
+ if pred.numel() == 0:
+ return pred.new_zeros(())
+ d = wrap_angle(pred - target)
+ per_triplet = torch.nn.functional.smooth_l1_loss(
+ d, torch.zeros_like(d), beta=beta, reduction="none"
+ )
+ if weight is None:
+ return per_triplet.mean()
+ return (weight * per_triplet).sum() / weight.sum().clamp_min(1e-8)
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..05c807c
--- /dev/null
+++ b/alignn/inverse/evaluate.py
@@ -0,0 +1,262 @@
+"""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.
+
+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
+
+import math
+from typing import Dict, List, Optional, Sequence
+
+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.
+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,
+) -> torch.Tensor:
+ """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, lattice = _atoms_to_tensors(atoms)
+ src, _dst, _shift, r = torch_neighbor_list(
+ positions,
+ lattice,
+ cutoff,
+ max_neighbors=max_neighbors,
+ use_matscipy_topology=False,
+ )
+ 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)
+ )
+ 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,
+) -> 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.numel()]
+ if not parts:
+ return torch.zeros(0)
+ return torch.cat(parts)
+
+
+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 > 0 else counts
+
+
+def compare_angle_distributions(
+ 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_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(generated, bins)
+ q = _density(reference, bins)
+ ps, qs = p + eps, q + eps
+ ps, qs = ps / ps.sum(), qs / qs.sum()
+ kl = float((ps * (ps / qs).log()).sum())
+ m = 0.5 * (ps + qs)
+ js = float(
+ 0.5 * (ps * (ps / m).log()).sum() + 0.5 * (qs * (qs / m).log()).sum()
+ )
+ width = 180.0 / bins
+ emd = float((p.cumsum(0) - q.cumsum(0)).abs().sum() * width)
+ return {
+ "kl": kl,
+ "js": js,
+ "wasserstein_deg": emd,
+ "n_generated": int(torch.as_tensor(generated).numel()),
+ "n_reference": int(torch.as_tensor(reference).numel()),
+ "bins": bins,
+ }
+
+
+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 = torch.tensor(before.frac_coords, dtype=torch.float64)
+ f1 = torch.tensor(after.frac_coords, dtype=torch.float64)
+ df = f1 - f0
+ 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(
+ 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 = 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(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:
+ 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..49b96f9
--- /dev/null
+++ b/alignn/tests/test_inverse_angle_diffusion.py
@@ -0,0 +1,528 @@
+"""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_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]])
+ # 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..e8fe991
--- /dev/null
+++ b/scripts/atombench/angle_eval.py
@@ -0,0 +1,154 @@
+#!/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 torch
+
+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(gen_angles.mean()) if gen_angles.numel() else None
+ ),
+ "reference_angle_mean_deg": (
+ float(ref_angles.mean()) if ref_angles.numel() 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 = {}
+ 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(
+ 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
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]