Conversation
`RigidBodyRefinementStep` rebinds the Refinement to a resolution-truncated data view at every cutoff. `_rebind_for_data` assigns `reflection_data`, builds a fresh Scaler, and calls `_init_targets` + `reset_loss_state`. Run against the caller's own Refinement, those assignments are destructive: - `_init_targets` reconstructs `adp_target` and `geometry_target` from constructor defaults, so anything configured on them post-construction is silently reset. `adp_target['simu'].simu_sigma = 0.25` reads back as 2.0. - `reset_loss_state` discards the LossState, so a weight registered on it is gone. A key present in DEFAULT_GROUP_WEIGHTS is visibly overwritten; a custom one such as `adp/simu` simply returns None afterwards and its target falls back to the group weight. Neither rebuild is wanted by the step. Only the x-ray target depends on the data and scaler that changed; the ADP and geometry targets are built from the model alone, and `_run_one_cutoff` drops every non-xray target from the state before optimizing. Measured on a 3-cutoff run they are constructed 3 times and evaluated 6 times (registration probe plus loss refresh) purely as overhead, then deleted unused, while the x-ray target takes all 42 gradient evaluations. `run()` now points the step at a shallow clone that shares the model but owns its own attribute namespace, so every one of those assignments lands on the clone. There is nothing to restore afterwards and no window in which the caller's Refinement is inconsistent. The model is deliberately shared rather than copied: `use_rigid_xyz` swaps its xyz container in place, so refined coordinates reach the caller by object identity and no copy-back is needed. That is also what makes the change exactly equivalent rather than approximately so -- on 3E98 the refined coordinates are bit-identical to the previous behaviour, max per-atom difference 0.000e+00. `nn.Module` keeps submodules in `_modules`, so the clone copies that dict (and `_parameters` / `_buffers`) as well as `__dict__`; without it a submodule assignment on the clone would write straight through to the original. Tests: tests/integration/test_rigid_body_isolation.py, five cases -- the sigma survives, a custom LossState weight survives, targets and reflection_data keep their object identity, coordinates still reach the caller, and a normal macrocycle still runs afterwards. Verified that three of them fail when the sandbox is bypassed. Full unit + functional suite passes (1798 passed, 74 skipped). This is independent of #68. That PR gives ADP restraint parameters a constructor-level home so they survive *any* rebuild, including the `create_from_state_dict` and ensemble paths this change does not touch. Either can land without the other; together they cover both the storage and the rebuild.
Run rigid body against a sandbox so it cannot disturb its caller
Symmetry was spread across four modules and re-derived per call site. The reciprocal transform R^T h alone had seven spellings -- two functions disagreeing on axis order and dtype, plus five open-coded transposes -- and its failure mode is silent: a wrong transpose corrupts centric flags and epsilon multiplicities without raising. Symmetry now owns the operations and everything derivable from them, built on three composable primitives (apply_rotations, apply_translations, phase_factors) plus a cached reciprocal stack carrying R^T. Miller-index expansion stops being a special case a caller can get wrong: it is apply_rotations on a different op stack. SpaceGroup specialises Symmetry with the crystallographic identity and the CCP4 asymmetric-unit verbs, so a group built from a raw operation list serves non-crystallographic symmetry without pretending to be a crystal. Translation keeps two conventions on purpose. phase_factors is the complex exp(+2 pi i h.t) used to combine structure factors; expand_hkl needs a signed radian offset -2 pi h.t to expand phases. Merging them flips a sign that is invisible in P21/P212121/C2. These classes hold no refinable parameters, so they are dataclasses over DeviceMixin rather than nn.Module. That also removes a hazard: nn.Module intercepted Module-valued assignment and bypassed the spacegroup property setter. Derived state lives in one cache that .to() and copy() both clear, and the traversal clears it before moving anything so a cached sampling grid is not copied to the new device only to be discarded. Map symmetrization moves behind Symmetry.symmetrize_map, caching one operator for the most recent grid shape. The MapSymmetry factory is gone: its variable return type was being read as a boolean, which the space group answers directly. cell_params is dropped from both operators -- stored and never read, because symmetry acts on fractional coordinates. Verified: centric and systematic-absence flags match gemmi with zero mismatches across ten space groups, including the trigonal, hexagonal and cubic cases; epsilon stays Friedel-inclusive at exactly twice gemmi's count. F_calc is bitwise identical to the previous commit on 1DAW, 2DQ6 and 3A5V, and so is the interpolating map path for both combine modes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Model mixed two things: the parameters being refined, and the structure it was loaded from. ModelContext now holds the second -- unit cell, space group, atom table, link records, alternative-conformation groups, provenance and the configuration flags -- so the model's own surface is parameters and behaviour, and the crystallographic context can be handed to code that needs only that. Access is deliberately hybrid. cell, spacegroup and pdb keep forwarding properties because they carry roughly 770 call sites between them, and churning all of those would bury a real regression in rename noise. The low-traffic fields move outright and are reached as model.ctx.strip_H, .initialized, .links, .altloc_pairs, .verbose, .exclude_H_from_sf and the input paths. copy() collapses from field-by-field assignment to one ModelContext.copy(), which deep-copies the atom table and clones the cell and space group. Cloning the space group is load-bearing now that it is a mutable dataclass rather than an nn.Module: a shared reference would let an edit through one model's context reach every model copied from it. device and dtype_float stay on the model rather than moving to the context. They are live DeviceMixin trackers that the traversal rewrites in place on whichever object owns the tensors, so relocating them would either duplicate the source of truth or route the hottest device path through a property. Model.symmetry is gone; it was a bare alias for spacegroup and misleads now that Symmetry is a real class. ModelFT.copy() also loses its skip_modules guard, which only existed because SpaceGroup used to register as a submodule. Two defensive lookups had to go with it, both silent. set_adp_mode guarded on getattr(self, "initialized", False), which read the default once the field moved and turned the whole method into a no-op -- mode switches reported success and changed nothing. state_dict used hasattr(self, "altloc_pairs") and would have saved an empty grouping on every write; that path had no test, so this adds a round-trip one, checked against a null control. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The class was an nn.Module holding no parameters, and none of its 29 buffers existed at construction: build_hydrogen_topology and build_h_candidate_pairs attached them afterwards with register_buffer. That is why reading it needed hasattr guards, and why n_asu_candidates was fetched through a getattr default. As a dataclass the fields are declared with None defaults, the builders assign them directly, and the guards become plain None checks. The getattr default in the non-bonded H target is gone too: the attribute now always exists, so a default that can never fire would only mislead. The four derived placement tensors move behind reset_cache, which DeviceMixin calls on every .to(). Nothing cleared them before, so a device move left the clamped neighbour indices and the bond-length column referring to tensors on the old device. HydrogenTopology also graduates from the device-conformance UNCOVERED list to a real case, since an empty topology is now constructible. It is registered tensor_free: a fresh one is a bare shell whose tracker is all there is to check. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Restraints already define partial graphs per monomer, patch them when a link forms, and map them onto atoms by name -- then keep only flat index tensors and discard the connectivity. So connectivity gets reconstructed three incompatible ways: exclusions are inverted back out of the bond/angle/torsion index lists, riding-hydrogen parents are found by interatomic distance while the CIF bond graph is loaded and ignored, and H exclusions are derived a third time. Nothing can answer "what is atom i bonded to". torchref.topology holds that connectivity: a ResidueGraph carrying the sequence and the inter-residue links, over an AtomGraph carrying the atoms, the typed edge blocks, and a CSR bond adjacency behind neighbors(i). Edge blocks are contiguous and origin-sorted with per-origin bounds, so every subset is a view into one block. Identity stays as string arrays; every indexing structure is a tensor and moves with .to(device). Nothing consumes it yet, so this cannot change any result. What it does is establish that the graph is correct: the equivalence test compares edge sets per type and per origin against the existing builders on five structures covering alternative conformations, pre-existing hydrogens, disulfides, glycans and a nucleotide analogue with a metal. Two things only those cases surface. LINK-record bonds are a distinct origin that 3A5V and 1DAW need. And disulfide detection has to pair SG atoms, not residues: a cysteine modelled in two conformations carries two SG atoms and each forms its own bond, which a residue-keyed search reduces to one. Residues are identified by (chain, resseq, icode). The builders group on (chain, resseq) alone, so an inserted residue is merged with its predecessor and loses its intra-residue restraints. No bundled structure has an insertion code, so that path stays unexercised for now. Exclusions are offered both ways on purpose. exclusions_from_restraint_edges reproduces what the non-bonded term is given today; exclusions_12_13_14 walks the bond graph and is correct, and therefore excludes more pairs and moves the VDW loss, so wiring it in belongs in its own measured commit. The walk is pinned against a breadth-first reference, since a superset assertion would also pass for an implementation that returned everything. Adds an AlphaFold-start trajectory test with a null-control arm: a deviation from the committed reference means nothing until it is measured against the deviation between two runs of the same build. Measured null spread is 0.0001 to 0.0009 over two macro-cycles, so the tolerance is 4x that with a floor, and a run is fast enough that the test needs no slow marker. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
The restraint groups now come from the topology graph and the values layered over its edges, instead of from a flat TensorDict addressed by composed string keys. Reading them is unchanged: restraints["bond"]["all"]["indices"] resolves exactly as before, so no geometry target is edited. What changed is what happens on that read. It used to build a fresh accessor object, then a fresh per-type accessor, then a fresh dict assembled from six string-formatted buffer lookups -- per loss evaluation, per restraint type. It is now three dict lookups into a mapping assembled once, holding tensors that already exist. The per-origin indices are slices of one contiguous block, so a subset is a view rather than a copy and an in-place edit to a block needs no invalidation to be seen. Nothing is cached, so nothing can go stale on the access path. Only two events can orphan a view -- a device rebind and an edge-count change -- and both rebuild the mapping where they happen. _apply drops the views before the traversal and re-slices after, because DeviceMixin's walk recurses into dicts and would otherwise move each slice on its own, silently turning every view into an independent tensor. cat_dict was not idempotent: writing restraints["bond"]["all"] registered 'all' as an origin, so a second call folded the group into itself and doubled every bond, angle and torsion -- a 2x on the geometry weight, latent only because every call site happened to guard it. 'all' is now a span of the edge block and cannot be an origin, so that is unrepresentable rather than fixed. Restraint row order no longer depends on Python's string hash seed. It used to come from set iteration over the origin names, which reordered the rows between processes. Numerically neutral, verified two ways rather than assumed. Every (edge, property) pair matches the previous storage exactly across five structures -- including phi and psi correctly carrying no reference or sigma, and omega keeping its proline flag. Against the previous commit directly, every restraint count is identical, n_vdw included, which is what shows the exclusion set did not move; the losses agree to ~1e-7 relative, which is float32 summation-order noise from the canonical layout. Exclusions still come from the bond, angle and torsion edges rather than from bond connectivity. The connectivity-derived set is correct and excludes more pairs, so it moves the non-bonded term and belongs in its own measured commit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
A template already carries its hydrogens, with coordinates and with bonds
naming each one's parent. Generating hydrogens is therefore template
instantiation, not geometry reconstruction: align the template onto the heavy
atoms present, read the hydrogen positions off it, correct each to its library
bond length. The bond graph supplies the two things a template cannot know on
its own -- how many hydrogens a parent can still carry, and which of them have
a dihedral nobody has determined.
The previous placement fitted the template over two bond shells. For CB that
is {C, CA, CB, N, SG}, which spans chi1, and chi1 is the model's, not the
library's, so the rigid fit compromises between them. Measured on 7L84 that
aligned to 0.75 A RMSD and left 12% of side-chain hydrogens more than 1.5 A
from the atom they belong to, where they were silently dropped. Fitting the
parent and its immediate neighbours only -- the unit whose bond lengths and
angles really are library constants -- places all of them: 7L84 goes from 930
to 1064 hydrogens with none left undetermined, each at its ideal bond length to
machine precision.
Three strategies, chosen by what the graph says, in place of a seven-way
placement enum. The template frame where the template knows every heavy atom
actually bonded to the parent. Construction from the bonded neighbours where it
does not, which is the peptide-linked backbone nitrogen: it is bonded to the
preceding residue's carbon, and the free-amino-acid template has never heard of
that. An axis-preserving frame for a single-neighbour centre, which fixes every
bond angle and leaves only the rotation the scan then chooses.
Free torsions come out of the connectivity rather than a list of special cases:
a parent with exactly one heavy neighbour can rotate. That is the hydroxyl, the
thiol, the amine and the methyl, and it correctly makes an N-terminal ammonium
rotatable while an in-chain amide is not -- the test asserts a backbone
nitrogen rotates if and only if no peptide bond reaches its residue.
Waters are left alone, and for a reason rather than by omission: one heavy atom
gives no frame to align against and no bond to rotate about, so their hydrogens
could only point somewhere arbitrary.
strip_H still defaults to True, so no hydrogen enters a refinement and nothing
moves numerically. Model.generate_hydrogens, the gemmi path that round-tripped
through temporary PDBs on disk and had no callers, is gone; hydrogenate is the
single route and no longer takes lbfgs_steps or max_iter.
hydrogen_topology.py keeps its riding-placement machinery for now. It is still
live while hydrogens are absent from the model, and deleting it here would
remove the H-VDW term rather than leave the stage inert.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
strip_H now defaults to False and a model tops up the hydrogens it is missing as it loads, so hydrogens are ordinary atoms with their own coordinates and displacement parameters and they contribute to F_calc. add_hydrogens=False keeps whatever a file carries without adding more; strip_H=True still removes everything. Generation is decided per parent rather than per file. A structure deposited with some of its hydrogens gets the rest: 1AK5 arrives with 675 of roughly 2500 and ends with 3060. A does-the-table-contain-any test would have left it as deposited. Measured on the five bundled AlphaFold starts over two macro-cycles, against the previous commit: atom count roughly doubles, median dR_work +0.0188 (worse, and positive on all five), median dR_free -0.0006 (flat, three of five improve), refinement 1.63x slower, and the test suite 494s to 1277s. R-work being consistently worse while R-free does not move is what library-placed hydrogens contributing scattering they have not been refined into looks like at this cycle count. n=5 over two cycles is indicative and no more -- the number that would settle it is the 767-structure panel through compare_to_archive.py, which has not been run. The default is one line in its own commit so it can be reverted without losing the machinery. Three bugs this surfaced, none of them in the hydrogen code: The valence cap subtracted heavy neighbours but not the hydrogens a parent already carried, so a nitrogen holding its H still had budget for the free amino acid's H2. Generation was therefore not idempotent and a save/reload added one hydrogen to every linked residue, 313 of them on 1DAW. Filtering candidates by name is not sufficient on its own -- a deposited hydrogen whose name differs from the template's would still have been over-added. vdw_radii, Z, the ITC92 coefficients and the heavy-atom mask are cached behind hasattr guards and returned as they are once built. That was safe only while every change of atom set produced a fresh model; extending the table in place left the radii at the heavy-atom count while the pair list indexed the full set. load() now drops them. Anything that replicates a decided atom table and then reshapes by its own atom count breaks under generate-on-load. EnsembleModel's two factories and _new_model_from_df all pass add_hydrogens=False, and EnsembleModel defaults it off, since the replicated single copy is the atom set. Riding hydrogens are no longer placed when the model carries real ones. They approximate the sterics of hydrogens that are absent, and with real ones present they were putting phantom atoms into the structure -- 343 on 1DAW -- that push real atoms around. Nor were they the hydrogens the generator declined: the riding builder counts bonded neighbours by distance while the generator reads them off the bond graph. They stay live under strip_H=True, which is the mode they exist for. The trajectory test now carries a measured tolerance per structure rather than sizing one from two runs at test time. 6JZA is bimodal at two macro-cycles -- its trajectories land in one of two basins about 0.0074 apart -- and two runs that pick the same basin report a spread 140 times too small. Tolerances come from five runs at regeneration time and are committed with the reference, so the other four stay held to 0.002 instead of being loosened to accommodate 6JZA. Also repoints AmberTarget at hydrogenate. The previous commit deleted Model.generate_hydrogens while a live call remained, and its tests skip without OpenMM so nothing said so. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
torchref.restraints.hydrogen_topology becomes torchref.topology.riding, next to the generation path it is the alternative to. Both answer the same question over the same graph, so they belong together: hydrogens.py adds hydrogens as real atoms with their own parameters, which is the default, and riding.py reconstructs absent ones from their parents at each non-bonded evaluation, which is what a model loaded heavy-only gets. The module docstrings now say which applies when, and that they must not run together. Kept rather than deleted. Riding placement is the only thing that gives a strip_H=True model any hydrogen sterics at all, and that is a mode worth having. Two pieces stay where they are because that is where their kind lives: non_bonded_h.py is a refinement target, and place_hydrogens.py is a Triton kernel. Both now import from the new location. Suite unchanged at 1918 passed, which is what a move should do. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
subset reindexes the surviving edges instead of rebuilding, so taking part of a structure no longer means re-reading the monomer CIFs and re-matching every template. copy gives an independent graph. Both are defined at each level -- EdgeBlock, ResidueGraph, AtomGraph, Topology -- so a caller can reduce whichever it holds. An edge is dropped as soon as any of its atoms is. A bond to an atom that is gone is not a bond, and an angle missing its apex is not an angle. Worth being clear about the consequence: a subset is a geometrically weaker model rather than merely a smaller one, because every restraint crossing the boundary goes with it. A residue left with no atoms is dropped too, and so is any link that reached it, since a peptide bond to a residue that is not there would leave an edge pointing outside the graph. No re-sort is needed. The remap is monotone on the atoms it keeps -- survivors are renumbered in their existing order -- and a monotone relabelling preserves lexicographic order, so each origin's rows stay sorted among themselves and the block stays canonical. That is also why subset takes an index list as a set and returns the topology's own atom order: honouring a caller's order would quietly leave the blocks unsorted, and the 'all' group would stop being a contiguous span, which is what makes it a view rather than a copy. There is a test pinning that. The tests compare surviving edges on (chain, resseq, icode, atom name) rather than on index. An index-based check passes trivially after a remap; identity is what catches a remap that points an edge at the wrong atom. This is the primitive, not yet the win the plan claimed for Model.select. That also needs the restraint value layer and the non-bonded pair list reduced in step, or select still rebuilds. Also corrects Topology.neighbors, annotated np.ndarray since before the adjacency moved to torch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
The builders group residues on (chain, resseq), so 100 and 100A become one residue with two sets of backbone atom names. The name-to-index map keeps the first of each, and the later residue ends up with no intra-residue geometry at all. Keying on (chain, resseq, icode) is what fixes that, and until now nothing exercised it: no bundled structure has an insertion code. Synthesised in the test rather than shipped as another data file. Only columns 23-27 of 3GR5 change, so it is visibly a renumbering and nothing else, and the residues stay in file order as a real insertion would. Measured on the synthetic case: the legacy grouping finds 1104 intra-residue bonds and the graph 1119, losing none. The 15 it misses are exactly the second and third inserted residues, which get zero bonds each under the merged grouping, and the test asserts that localisation rather than just the count -- a graph that produced more restraints everywhere would pass a bare superset check. The first version of this test proved nothing and is worth recording as a trap: it compared the graph against restraints.py, which since the storage swap *builds from* the topology. Identical counts on both sides were the giveaway. The independent baseline is BondRestraintBuilder, which still keys on (chain, resseq). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
torchref.restraints.restraints becomes torchref.topology.restraints, and RestraintsNew becomes Restraints -- the "New" was historical, and the package alias already meant most callers said Restraints anyway. Moving the orchestrator in turned out to be a smaller change than extracting its pieces out one at a time. The real coupling was three production import sites, not the ninety an attribute-access count suggests, and two of those disappear here: topology/build.py was importing _lookup_link_atom back out of the restraints module, and restraints/__init__.py held the alias. Sixteen test imports were mechanical. _lookup_link_atom moves to topology/build.py, which is the only thing that uses it. That was not optional -- with the orchestrator inside topology, importing it back out would have been circular. The backwards import was the clearest sign the split sat in the wrong place. Nothing on the hot path changes. Model.restraints returns the same object, so self.restraints.restraints[...] in the geometry targets is untouched, and no file under refinement/targets/ is edited. What is left in torchref.restraints is the data layer: the monomer library, the CIF readers, the chem_mod records that patch a template when a link forms, the Numba matchers, the Ramachandran surfaces and the spatial search behind the non-bonded pair list. What is *built* from that data now lives in torchref.topology. The package docstrings say so, and that also settles where cif_dict belongs -- the loader that reads it is in topology, the library that supplies it stays put. Suite 1938 passed, unchanged by the move itself. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
torchref.restraints is gone. Every one of its seven modules had exactly one kind of importer left -- something inside torchref.topology -- so the package had already become a private implementation detail of it. Moving them in recognises that rather than imposing anything new. Placed by what each one is, not under a single label, because they are not alike: topology/monomer/library.py resolve, download and cache the CCP4 library topology/monomer/cif.py read a dictionary into its sections topology/monomer/modifications.py the chem_mod records that patch a template topology/builders.py template-to-atom matching topology/builders_numba.py the njit matchers topology/nonbonded.py the spatial search behind the VDW pair list topology/ramachandran.py the NLL surfaces Only one rename beyond the relocation: restraints_helper.py becomes monomer/cif.py, since the old name described where it sat rather than what it did. Tests followed, tests/unit/restraints to tests/unit/monomer. Two modules computed a data path by counting levels up from __file__, and library.py gained a level in the move. Left alone it would have pointed at torchref/topology/data/, where nothing lives -- and it would not have raised: the bundled monomer library would simply have appeared absent and every component would have been re-downloaded. Both paths are now pinned with an explicit parents[n], and both were checked afterwards: the bundled library resolves, ALA loads, and the Ramachandran surfaces come back at (6, 360, 360). Also fixes neighbor_search annotating three locals with List without importing it. Harmless at runtime, since annotations are not evaluated there, but wrong. Suite 1938 passed, unchanged by the move. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
DisorderFieldTensor stores disorder parameters on K nodes and gives each atom a distance-weighted mean over its k nearest, so the parameter count scales with node count rather than atom count. Storage is (K, 2) holding [log B, log sigma] per node; forward() returns per-atom isotropic B, so it fits the model.adp slot and every consumer of adp() works unchanged. Node positions are derived as the centroid of each node's anchor neighbourhood rather than refined, so a node stays inside the molecule and the optimiser has no free coordinate to drift with. Weights are a softmax over each atom's candidate nodes, so B is a convex combination of positive node values and needs no clamp. Coordinates come from an accessor injected at construction, held through ModuleReference so it stays out of state_dict, .to() and deepcopy. That leaves the inherited forward cache blind to coordinate changes, since CachedForwardMixin fingerprints parameters, buffers and call arguments only, so _fingerprint_state folds the accessor's output into the key. Nothing in Model references this yet. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
set_adp_mode("field") joins "isotropic" and "anisotropic" on the existing switch,
replacing the per-atom isotropic B with a DisorderFieldTensor whose node values are
least-squares fitted to the B it replaces. Entering runs the existing partition first,
so everything keyed off the iso/aniso split is refreshed by tested code; leaving needs
no special case, because the partition reads adp(), which a field evaluates per atom
and so materialises back into a per-atom wrapper on the way out.
Nodes are anchored on density clusters rather than single atoms. A node placed exactly
on an atom can isolate that atom by narrowing its kernel, which is per-atom refinement
wearing a node's clothes; measured, cluster anchoring lowers the worst per-atom B on
every structure tried.
Node positions carry a refinable offset from their anchor centroid, on by default. With
positions fixed a node's only way to localise is to narrow, so this is what lets a
restraint move a node toward atoms instead of only widening it.
Model.copy re-points the borrowed coordinate accessor at the copy, or the two models
share coordinates and the copy is not independent. create_from_state_dict rebuilds a
field when the saved adp storage is 2-D, reusing the saved anchor rows and inferring
from the storage width whether positions were refinable.
node_load() scatters the candidate weights back into node space. Summing weights() over
atoms does not do this -- it returns (n_atoms, k) over candidates, so the sum is a
length-k vector of per-slot totals with no meaning.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
Two restraints for the node-field ADP representation, closing different halves of the same degeneracy. Both are inert unless the model is in field mode. NodeLoadTarget (adp/node_load, weight 10.0) bars a node from being abandoned. A node can otherwise narrow its kernel until it holds a single atom and then take whatever value fits it; measured, such a node ends with a load near or below one atom against a healthy median of seven, and drives that atom's B into the thousands. The penalty is one-sided, softplus(-log(load / mean load)): an over-loaded node is not charged. Maximising load entropy would have been the symmetric choice and is wrong here, because it is optimal at uniform load and fitted fields legitimately span two orders of magnitude in kernel width. It acts on the weights alone, with no gradient to the node values, so it removes the opportunity to place an extreme B rather than pricing it. That is also its limit, and the reason for the second term. Blocking the collapse does not stop a value running away, it only stops it being confined: with the barrier alone the worst per-atom B rose at moderate node counts, spreading over a neighbourhood instead of landing on one atom. NodeSmoothnessTarget (adp/node_smoothness) prices that, as a distance-weighted mean squared difference of log B between node pairs. Level-invariant, so it cannot fight the scaler over the overall B level; scale-free in log space; and it charges only departures from the local level, leaving a genuine B gradient across a structure free. Registered at weight 0 pending its own screen, as geometry/ramachandran already is. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
Seven methods with no callers anywhere in the package or the tests, left behind when the restraint layer moved into topology and the geometry targets took over the loss maths: expand_altloc, the h_excl_hash property, torsion_deviations, nll_torsions, nll_planes, nll_vdw and adp_similarity_loss. The h_excl_hash property was only a read-only wrapper over self._h_excl_hash; that attribute is still built and passed to the non-bonded pair list, and the identically named arguments in nonbonded.py and riding.py are unrelated function parameters. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
ModelFT.create_from_state_dict carried its own copy of the block that rebuilds the parameter wrappers from the atom table, and the node-field branch had only ever been added to Model's. A model in either field mode therefore failed to restore through ModelFT -- the class every refinement uses -- with a shape mismatch naming the refinable mask rather than the representation. The existing round-trip test drives the base Model, so it passed throughout. Both classes now call Model._rebuild_wrappers_from_pdb, and _restore_adp_slot covers the u slot as well as adp: mode="field_aniso" puts the field in u, which both copies rebuilt unconditionally as a CholeskyMixedTensor, so that mode had no restore path at all. The slot is identified by its saved neighbor_list rather than by the shape of its storage, since u is two-dimensional either way and (K, 10) cannot be told from (n_atoms, 6) by rank. NodeLoadTarget and NodeSmoothnessTarget read model.adp directly and so were inert whenever the field lived in u -- the mode with the most node parameters to collapse. Both now go through Model.adp_field, which already looked in both slots. The device conformance registry gains cases for those two targets and loses three entries for the payload strategies, which are not DeviceMixin subclasses and so were never found by its AST inventory. On 1DAW, field and field_aniso now round-trip through both classes with the B and U they were saved with; the isotropic control is unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
SfFFT held an (nx, ny, nz, 3) buffer of Cartesian voxel coordinates that no structure-factor path reads. Every splat derives a voxel's position arithmetically from its index and the fractionalisation matrix, so build_electron_density used the tensor only for .device and .shape[:-1], and the CUDA wrappers already passed density_map in its place through a grid_ptr that none of the four Triton kernels ever loads. Being a registered buffer, it also went into every saved state: 187.5 MB for 3K7M at 250^3, against 24 bytes for the gridsize and voxel_size that remain. Building it cost 65.2 ms there and 5.6 ms on 1DAW, which was all of setup_grid's runtime; setup_grid is now 0.05 ms. build_electron_density takes a grid shape and a device instead. Gone with the buffer: its dead voxel_size parameter, grid_ptr in the four Triton kernels and their launch helpers, the stand-in argument in the CUDA wrappers, SfFFT.compute_real_space_grid, and ModelFT.get_radius with its two forwards, which the per-atom truncation radius had made vestigial. voxel_size is now frac_matrix @ (1 / gridsize) rather than a difference of two grid points. Shape-only callers use the new grid_shape property, and ModelFT.real_space_grid is a method over get_real_grid for the callers that want the coordinates themselves. F_calc and the density map are bit-identical before and after on 1DAW (C2), 3GR5 (P6522) and 3K7M (P432), and the analytic voxel_size reproduces the differenced value exactly. Checkpoints carrying the old buffer still restore. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
create_from_state_dict took a device argument but passed it to exactly one of the four parameter wrappers. OccupancyTensor honoured it; xyz, adp and u are built from the atom table via torch.tensor and land on CPU regardless. Restoring with the default device resolved to an accelerator therefore produced a model split across two devices, which CPU-only CI cannot see and the accelerator workflow trips over. The restore now builds on CPU throughout and moves once at the end, and only when the caller names a device. None leaves the model on CPU rather than resolving to device.current: reading a file back is not a reason to claim an accelerator, and a caller who wants one can say so or move the model itself. Measured with device.current = cuda, for isotropic, field and field_aniso alike: device=None gives all four wrappers on cpu, device=cuda gives all four on cuda:0, and the moved model computes structure factors. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
All three pass on CPU and fail under the accelerator workflow, which is the only place
that runs with TORCHREF_DEVICE set.
test_anomalous builds hkl with a plain torch.tensor while the model sits on
device.current. Nothing moves it: the reciprocal extractor takes its device from hkl by
design, so the phases end up on CPU against a map on the accelerator. Left as an error
rather than moving hkl inside forward -- a caller handing a model and its reflections
different devices should hear about it -- so the tests now build hkl on model.device.
test_entries_survive_a_device_apply clones its reference before calling
.to(torch.device("cpu")), so the two are only on the same device when the default
already is CPU. Compares against before.cpu() now.
test_empty_shell_needs_no_accessor asks for float64, which MPS does not have. Every
other case in that file inherits CPU from the tensors it passes in; only the empty shell
resolves device=None to device.current, so it is pinned to CPU.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
_run_one_cutoff borrowed the refinement's LossState, dropped every non-x-ray target from it, and restored afterwards. The restore was dead code: it snapshotted state.weights and then popped the non-x-ray names from the snapshot as well, so the finally clause wrote back only "xray" -- a value nothing in the block had touched. step() reads weights through get_effective_weight and never writes them, and the sole mutation point is set_weights(self.weighting(state)) at construction. The dropped targets were never restored at all, which was harmless only because the step runs against a sandbox clone and _rebind_for_data rebuilt the state each cutoff. The step now builds a LossState of its own holding the x-ray target alone. Nothing has to be undone afterwards, no maintenance() hook of a target that is not in the sum can fire, and the caller's state keeps its targets and any weights registered on them. Weight 1.0: with a single term the weight is a scalar on the whole objective, and 1.0 is what DEFAULT_GROUP_WEIGHTS gives x-ray anyway, so the effective objective is unchanged. _rebind_for_data drops reset_loss_state, which nothing reads now, and calls _build_xray_targets + get_scales in place of _init_targets. The x-ray target genuinely changes with each cutoff's data and target mode; TotalGeometryTarget and TotalADPTarget do not, and were being constructed five times per run -- three cutoffs, the full-resolution restore, the commit rebind -- only to go unused, the NonBondedTarget pair list among them. Counted on 3E98 that goes from 5 and 5 to 0 and 0. The cutoff schedule, the per-cutoff resolution truncation and the 6 A target-mode switch are untouched. Results are unchanged within run-to-run spread: over 3E98 and 1DAW the largest old-vs-new difference is 4.52e-02 A / +5.0e-04 R-free, against a same-code null control of 4.51e-02 A / +4.7e-04. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The file opened "tests/files/mtz/<name>.mtz" -- a path relative to the working directory -- so all 13 tests in it failed with "No such file or directory" whenever pytest ran from anywhere but the repo root. The workflows run from the root, so CI never saw it. Running the suite from tests/, which is what makes --run-slow resolve, does see it, and the failures read as a standing red in the outlier screening rather than as a missing file. All four call sites now go through the mtz_dir fixture, as the rest of the suite does; test_deposited_structures_lose_almost_nothing was already requesting pdb_dir without using it. This was the only cwd-dependent data path left under tests/. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
_percentiles built its tensors as float64 regardless of the configured dtype. It is CPU-only and feeds a report, so nothing could crash on it, but the float dtype is a configuration and hardcoding around it is how a float64 path reaches a backend that has none. The float64 that remains in the package is deliberate: the dtype registry itself, the MPS guards and dtype checks, two explicit .cpu().double() calls, the CUDA-only planarity eigh (whose forward asserts is_cuda, and where the fp64 is load-bearing for near-collinear atoms), and scalar constants that become Python floats. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A disorder-field node that carries the covariance of its displacement modes needs a q x q factor, not the 3 x 3 one the per-atom ADP wrapper uses. Same contract as the unrolled pair beside it: exp(x) + epsilon on the diagonal so the reconstruction is positive-definite for any parameter value and epsilon floors the smallest eigenvalue, an eigenvalue clamp before factorising on the way back, and a CPU-forced eigh because cuSolver's batched kernels fail on the degenerate batches a near-isotropic model produces. The 3 x 3 pair stays as it is. It runs in a forward pass, where the unrolled form is worth keeping, and raw_to_cholesky(raw, 3, eps) is asserted to reproduce it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
A node stored one ADP, so the only way to express U varying through space was to add
nodes: spatial detail cost nodes, and every added node brought another kernel that
could narrow onto a single atom. Store the covariance of the node's displacement modes
instead. With modes Psi(r) and coefficient covariance Sigma = L L^T, an atom at
displacement r from the node receives U(r) = Psi(r) Sigma Psi(r)^T, so one node's U
already varies across its whole region and cannot spike.
Positive-semidefiniteness is free at every r, because U = (Psi L)(Psi L)^T. An
arbitrary polynomial in r carries no such guarantee and goes indefinite at the edge of
the region, which is exactly where the softmax weights have not yet decayed.
The mode set is the expressiveness knob. Three translations and three rotations
reproduce the textbook TLS expression T + AS + S^T A^T + A L A^T identically, with
tr S appearing as its one flat direction; releasing the antisymmetry of the gradient
adds domains that breathe and shear as well as rotate.
constant q=3 6 payload one U per node, as AnisotropicPayload
rigid q=6 21 payload TLS
rigid_dilation q=7 28 payload TLS plus uniform breathing
affine q=12 78 payload full linear displacement field
Entering the mode seeds only the translation block from the constant-U solve and leaves
the gradient modes at the Cholesky floor, so a freshly installed field starts at the
constant-U field's own R-factor and refinement can only move away from a known state.
Measured cycle-0 R-free spread across all four rungs: 0.00004.
NodePayload.fit gains the geometric context contributions already took, because an
r-dependent payload cannot build its modes without it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
UNCOVERED exists to excuse device-bearing classes, and the node payloads are plain strategy objects, so listing them there was a category error: the AST inventory never finds them, test_uncovered_entries_still_exist reported them as classes that no longer exist, and the guard was red. With it red, NodeLoadTarget and NodeSmoothnessTarget went in with neither a case nor an excuse and nothing noticed. Drop the payload entries and give both targets a real target case. They are inert outside field mode but still device-bearing, so they construct on a plain model and must track its device like any other target. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
# Conflicts: # tests/helpers/device_cases.py
LossState.register_targets keys each component off target.name. Every ADP target declares its own path -- "adp/simu", "adp/sigd", "adp/locality" -- except these two, which declared none and so inherited ModelTarget.name, the literal string "model_target". Both registered under that, the second overwrote the first, and get_effective_weight never saw an "adp/..." path, so DEFAULT_GROUP_WEIGHTS["adp/node_load"] = 10.0 was a dead entry and the coverage barrier was in no refinement's loss at all. Every test passed because they called the target directly, which works. One even asserted the weight was present and positive, which was true and meaningless. The diagnostic that shows it is comparing TotalADPTarget's own component set against the LossState's keys. The guard added here walks every concrete ModelTarget subclass and fails any still carrying a base placeholder name, plus checks end to end that each declared component reaches the LossState with an addressable weight. Group bases are excluded by having subclasses -- they are not formally abstract, so that is the only reliable marker. Also stops node_smoothness reading float() off a grad-carrying tensor. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
simu and locality restrain by penalty the spatial smoothness a node field enforces by construction, and node_load/node_smoothness have nothing to act on off it. TotalADPTarget now builds the applicable set instead of building all of them and zero-weighting half. Registering-then-zeroing costs nothing at run time -- LossState.aggregate skips a zero-weight target -- but it leaves the correctness of the setup resting on a number, so anyone adjusting the adp group weight for their own reasons silently re-enables a restraint that double-counts the parametrisation. Whether a term applies is a property of the representation, not something to tune. sigd applies either way: it is a prior on the marginal B distribution, which a field constrains no more than a per-atom model does. The node-load tests now cover every payload. They were written against the isotropic payload, whose storage is five columns wide, and poked column 1 by index; a mode payload is 25 to 82 columns with log sigma near the end, so "acts through the weights, never the values" had to be re-established rather than assumed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
state_dict holds tensors, and the payload is a plain object that never reaches it, so a restore had to infer it from the slot. That is not a clean failure: a mode payload rebuilt as a constant-U one has the wrong storage width and surfaces as a shape mismatch, or worse as silently different ADPs. The field now registers an integer code. PAYLOAD_CODES is append-only -- a saved state dict holds the number, so renumbering would restore the wrong payload. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
…tion
The restore path now reads the payload code rather than guessing from the slot, falling
back to the old inference for state dicts written before the code existed.
set_adp_mode gains init={"fit","flat"}. "fit" (default) fits the field to the model's
current per-atom ADPs, which is right when those mean something. "flat" keeps only their
level and discards the spatial structure, for when they do not -- an AlphaFold model's B
values come from a pLDDT conversion, and fitting a smooth basis to them spends the field's
parameters reproducing structure it cannot hold.
Flattening goes through the equivalent isotropic B and hands the payload a 1-D target,
which its fit lifts to U_iso * I. Taking a median over all six U components instead sets
the off-diagonals equal to the diagonals, giving eigenvalues (3L, 0, 0) -- singular, and
NaN once the Cholesky encode takes log(diag - epsilon). Found by smoking it.
Measured: flat initialisation is worth -0.0005 R-free (p=0.29, n=60) from an AlphaFold
start, so it is an option rather than a better default.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
Model.set_adp_mode changes the representation but cannot size it -- node count follows from the reflection count, and the model has no idea how much data there is -- and cannot swap the ADP restraint set. Going through the model directly is a partial setup. set_adp_representation sizes a field from the work set, rebuilds the targets, scales and LossState, and is safe to call after construction, which is what the model's own "run once at setup" caveat is about. __init__ routes through the same method, so construction and a later switch cannot drift apart. Node cost is read off the payload object rather than tabulated, so a new payload cannot desynchronise the arithmetic from what the field allocates. Default target: 7 work reflections per ADP parameter. PDB-REDO holds ~7 across its whole resolution range and changes model form to stay there; measured on 179 of their entries a node field peaks at the same value, with both directions worse. The loss is deliberately NOT rebalanced for a field. The point of the representation is that smoothness comes from the parametrisation, so it should need less regularisation, not a reweighted version of the same priors. An earlier FIELD_GROUP_WEIGHTS raising the adp group had two side effects worth recording: adp/scaler_U and adp/scaler_log_scale sit under that group, so it multiplied the scaler regularisation by the same factor, and it made every field run differ from its baseline in a way unrelated to ADPs. Measured, it was purely harmful from an AlphaFold start (+0.0067 R-free) and irrelevant on converged models. flatten_adp_field discards the field's structure mid-refinement, keeping the level. A field fits its structure once, when installed, and nothing re-derives it, so structure inferred against wrong coordinates survives -- the same shape as freezing bulk solvent at the starting model. Worth -0.0098 R-free (p=0.002, n=60) from an AlphaFold start. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
--adp-mode gains field, field_aniso and preserve, plus --adp-mode-set, --reflections-per-adp-parameter and --adp-nodes, wired through to the constructor and echoed in the run summary and saved settings. The representation was library-only before, so none of it was reachable by a user. Five flags in this codebase were once silently no-ops, so the test asserts that non-default values arrive rather than merely that they parse. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
Takes the fuller wording from dev's working tree, which describes field_aniso as well. The preserve implementation there is identical to the committed one; only the docstring differed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TN8kHX6f9MwFxN8nGUwiQU
set_adp_mode always reparametrised, so constructing a Refinement over a deposited model isotropised its ANISOU before anything else ran, and a run meant to measure that model's own ADPs measured a converted copy instead. "preserve" returns immediately, leaving the per-atom wrappers exactly as the reader built them. The docstring also picks up "field_aniso", which the mode list had gained without a description. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…loff The page still described per-bin scaling and a plain Debye-Waller solvent term, both of which the scaler rework replaced. The overall scale is now a Chebyshev polynomial in s with n_iso_coeff coefficients, and resolution bins survive only as the seed for that fit; the solvent term is a generalised falloff parametrised by s^2_half and n, which reduces to exp(-B_s s^2) at n = 1 so a solvent B from another program still transfers. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The lab's scripts stay tracked; its cache, metrics, figures and slurm logs do not. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Brings the displacement-mode node payload (rigid/TLS and affine mode sets, node-held mode covariance), the flat field initialisation, payload persistence across save and restore, named node-field ADP targets so a weight reaches them, representation-aware restraint registration, a Refinement entry point for switching the ADP representation, and the field ADP modes on the refine CLI. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1VER's final-cycle R-free has a second basin about 0.0042 from the main one. It is rare enough that the five runs sizing the reference never sampled it, so the structure took the 0.002 floor and the basin sat outside a bound meant to cover it -- the same trap the module docstring already records for 6JZA, which was caught only because its basin is common enough to show up in five runs. R-work and R-free now carry separate bounds. R-work keeps the measured tolerance: it is reproducible to a few parts in ten thousand across CPU generation, thread count and torch version, and is what a real change in the refinement moves. R-free takes a 0.01 floor, which clears the basin while staying inside a quarter of every structure's own R-work descent -- test_tolerances_are_tight_enough_to_detect_something now checks the bound actually applied rather than the stored one, so the floor cannot widen the real bound unnoticed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR to merge dev into main
changes so far:
Version 0.7.0
ModelContext. It now holds the unit cell, space group, atom table, link records, hydrogen settings, and input paths.Symmetryas a crystallography-free class with transform primitives, and madeSpaceGroupa specialised subclass.torchref.restraintswas removed, restraint dictionaries are now plain nested dicts, and residues are identified by(chain, resseq, icode)to fix insertion-code merging.Model.hydrogenatenow aligns monomer templates onto heavy atoms present, generation is the default, andAtomGraph.exclusions_12_13_14derives non-bonded exclusions from bond connectivity.Topologyas aResidueGraphover anAtomGraphwith typed edge blocks andsubset/copyoperations that reindex surviving edges.HydrogenTopologya dataclass, changedSymmetryclasses to dataclasses overDeviceMixininstead ofnn.Module, and removed unusedCellgradient plumbing and theReciprocalSymmetryGrid/ module-level expansion functions.