diff --git a/.github/workflows/compatibility.yml b/.github/workflows/compatibility.yml index d6437c7a..7c46cccc 100644 --- a/.github/workflows/compatibility.yml +++ b/.github/workflows/compatibility.yml @@ -21,6 +21,7 @@ on: # Also run on PRs that modify dependencies pull_request: + branches-ignore: [dev] paths: - 'pyproject.toml' - 'tox.ini' diff --git a/.github/workflows/dev-pr.yml b/.github/workflows/dev-pr.yml new file mode 100644 index 00000000..33131672 --- /dev/null +++ b/.github/workflows/dev-pr.yml @@ -0,0 +1,47 @@ +name: Dev PR Tests + +on: + pull_request: + branches: [dev] + +permissions: + contents: read + +concurrency: + group: dev-pr-tests-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + cpu: + name: CPU (Python 3.12) + runs-on: ubuntu-latest + timeout-minutes: 90 + env: + TORCHREF_DEVICE: cpu + NUMBA_CACHE_DIR: /tmp/numba_cache + + steps: + - uses: actions/checkout@v7 + + - name: Set up Python + uses: actions/setup-python@v7 + with: + python-version: '3.12' + cache: pip + cache-dependency-path: pyproject.toml + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m pip install -e ".[dev]" + + - name: Run tests on CPU + run: | + python -m pytest tests/ \ + -m "not gpu and not slow" \ + -v --tb=short -rf --durations=20 + + mps: + name: MPS + uses: ./.github/workflows/accelerator.yml diff --git a/.gitignore b/.gitignore index 9d302c4c..d534df59 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,8 @@ _temp.mtz __pycache__/ *.pyc *.pyo +# macOS finder metadata +.DS_Store # Large binary files *.ccp4 *.png @@ -84,4 +86,19 @@ graphify-out/ # Large binary scratch dirs & squashfs images *.sqsh anisotropic/ -torchref_refine_optimization/runs/ \ No newline at end of file +torchref_refine_optimization/runs/ + +# Alignment lab: keep lab/, diagnostics/, analysis/ and tests/ tracked; +# ignore run outputs (CSVs, Phaser working dirs) and scheduler logs. +alignment_lab/runs/ +alignment_lab/slurm/ +alignment_lab/**/__pycache__/ + +# Third-party Phaser source copy (from PHENIX 1.20-4459) for FRF map instrumentation +phaser_src/ + +# scaler_investigation lab outputs: the scripts are tracked, the outputs are not. +scaler_investigation/cache/ +scaler_investigation/metrics/ +scaler_investigation/figures/ +scaler_investigation/slurm/ diff --git a/AGENTS.md b/AGENTS.md index 5fd1707b..d02be91b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -186,13 +186,13 @@ Black, 88 columns, `isort` with the black profile. Ruff lint with |---|---| | `base/` | Low-level math and crystallography. `coordinates/` (Cartesian↔fractional), `reciprocal/` (basis, HKL, d-spacing, interpolation, symmetry), `direct_summation/` (F_calc by summation; eager + Triton), `electron_density/` (real-space splatting with CPU/CUDA/MPS kernels, solvent mask, radius policy), `fourier/` (FFT and grids), `scattering/` (form-factor and anomalous tables), `metrics/` (R-factors, binwise scale, loss), `targets/` (the *kernels* behind refinement targets, eager + `triton/`), `french_wilson.py`, `math_torch.py`, `alignment/` | | `io/` | `ReflectionData`, `DatasetCollection`, `FcalcDataset`; MTZ / PDB / CIF / IHM readers and writers; `read_mtz` / `read_pdb` / `read_cif` | -| `model/` | `Model` (refinable atomic parameters), `ModelFT` (adds F_calc via `SfFFT` or `SfDS`), `MixedModel`, `ModelCollection`, and the parametrizations in `parameter_wrappers.py` / `rigid_xyz.py` that decide what is refinable | +| `model/` | `Model` (refinable atomic parameters), `ModelContext` (the cell, space group, atom table, links and provenance a model is loaded with — `model.cell` / `.spacegroup` / `.pdb` forward to it, the rest is `model.ctx.*`), `ModelFT` (adds F_calc via `SfFFT` or `SfDS`), `MixedModel`, `ModelCollection`, and the parametrizations in `parameter_wrappers.py` / `rigid_xyz.py` that decide what is refinable | | `refinement/` | Drivers (`Refinement`, `LBFGSRefinement`, `RigidBodyRefinementStep`), `targets/` (`xray/`, `geometry/`, `adp/`, `collection/`, `combined.py`), `weighting/`, `optimizers/` (annealing, Langevin, preconditioned/seeded L-BFGS), `model_error_estimation/` (σ_A, σ_M), `loss_state.py`, `logger.py` | | `restraints/` | Bonds, angles, torsions, planes, chirals, VDW. Built from the CCP4 Monomer Library, resolved lazily via `get_library_manager()` — importing this package must not trigger a library download | | `scaling/` | `ScalerBase` (model-independent), `Scaler`, `CollectionScaler`, `SolventModel` (k_sol, B_sol) | -| `symmetry/` | `SpaceGroup` (buffers on `nn.Module`), `Cell`, `MapSymmetry`, `ReciprocalSymmetry`, grid utilities | +| `symmetry/` | `Symmetry` (operations plus everything derived from them), `SpaceGroup` (adds the crystallographic identity and the CCP4 ASU verbs), `Cell`. All dataclasses over `DeviceMixin`, not `nn.Module` — they hold no refinable parameters. Map and reciprocal-grid operators are private, reached through `Symmetry` | | `maps/` | `Map` (2Fo−Fc, Fcalc), `DifferenceMap` | -| `cli/` | Entry points: `torchref.refine`, `torchref.difference-refine`, `torchref.mtz2map`, `torchref.validate-ded`, `torchref.phased-difference-map`, `torchref.add-metadata`, `torchref.strip-altlocs` | +| `cli/` | Entry points: `torchref.refine`, `torchref.difference-refine`, `torchref.mtz2map`, `torchref.validate-ded`, `torchref.difference-map`, `torchref.add-metadata`, `torchref.strip-altlocs` | | `experimental/` | APIs that may change without notice: `alignment/` (Patterson MR), `kinetic/` (time-resolved), `ensemble/`, `monolithic_refinement/`, `targets/` (AMBER/GAFF2, real-space, sampled-ML phase) | | `utils/` | See §5 | | `config.py` | See §4 | diff --git a/docs/changelog.rst b/docs/changelog.rst index 30fcbb07..4c09517e 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -2,8 +2,196 @@ Changelog ========= +Unreleased +---------- +- The difference MTZ groups its columns into named datasets -- ``observed``, ``difference``, ``light_model``, ``extrapolated_light``, ``two_moment`` -- with one history line describing each, so ``FWT``/``PHWT`` reads as ``/torchref/extrapolated_light/FWT`` (the extrapolated light-state map ``2*FEXT - Fc``). Labels are unchanged and Coot still auto-opens it +- ``torchref.difference-map``, ``torchref.difference-refine`` and ``torchref.validate-ded`` gain ``--ded-weight {sigma_d,inverse_variance,none}`` and ``--sigma-d-gamma``. The difference MTZ now carries the unweighted ``DF``/``SIGDF`` on ``PHDELWT`` with one mean-one weight column per scheme, ``W_SD`` and ``W_IVW`` (MTZ type W), and the observed-to-model scale ``KSCALE``; ``DELFWT`` is no longer written, build the map with ``torchref.mtz2map -csf DF -cw W_IVW -cphi PHDELWT``. Registered in ``torchref.maps.ded_weights`` +- Added the ``sigma_D`` estimator (``torchref.refinement.model_error_estimation.sigma_d``): the expected true difference power per resolution shell, ``mean(dF_obs^2) - mean(sigma^2)`` with a fitted ``F_dark^gamma`` amplitude law and DerSimonian-Laird shrinkage of the signed shell power toward a decaying exponential in ``d*^2`` (fitted on all shells, so a dataset without a difference yields no power instead of the positive half of its noise), giving the Wiener weight ``S/(S + sigma^2)`` and, with a difference model, ``alpha``/``beta_model``. Inverse-variance weights suppress the strong reflections whose difference power is 10-70x that of weak ones; on independent half-datasets a Wiener weight with the true power raised map agreement 1.2-1.8x in effective patterns. The single-dataset estimate inherits the calibration of the reported sigmas, and on the campaign TD1 data (sigmas ~1.5x too large at high resolution) it emptied 60-90 % of the shells, so inverse variance stays the default; ``sigma_d`` reports its clamped-shell count and falls back to inverse variance with a warning when every shell is empty +- Added the ``difference_sd`` collection target (``CollectionDifferenceSigmaDTarget``): the difference Gaussian centred on ``alpha * dF_calc`` with variance ``beta_model + sigma_diff^2`` from ``sigma_D`` fitted on the free set. Selected with ``torchref.difference-refine --difference-target difference_sd``; ``difference`` stays the default +- ``torchref.mtz2map`` gains ``--column-weight``/``-cw`` (multiply the amplitudes by a weight column before the FFT) and ``--units {sigma,electrons,raw}``: ``electrons`` writes e/A^3 as ``(1/V) sum_h F(h) exp(-2 pi i h.x)`` with the amplitudes divided by the ``--column-scale``/``-ck`` factor (``KSCALE`` by default). ``-n``/``--normalize`` is a deprecated alias. ``Map`` and ``DifferenceMap`` take ``units`` too, and ``DifferenceMap`` an optional per-reflection ``scale`` +- ``ScalerBase.multiplicative_scale()`` returns the per-reflection ``K_overall * b_overall * anisotropy`` factor, every multiplicative component of ``forward`` and none of the additive solvent term +- ``torchref.validate-ded`` reports the real- and reciprocal-space correlations for unweighted, inverse-variance and ``sigma_D`` weights side by side (``by_weight`` in the JSON, a table in the summary) and records when ``sigma_D`` fell back to inverse variance +- The extrapolated-map Bayes shrinkage estimates its signal variance per resolution shell through ``sigma_D`` instead of one global ``tau^2``; ``tau_sq`` in the summary is now the count-weighted shell mean +- Inverse-variance difference weights floor ``sigma_diff`` at a tenth of its median, so a zero sigma yields a large finite weight rather than an infinite one +- Shell construction, segment sums, interpolation and the shrinkage line fit shared by ``sigma_A`` and ``sigma_D`` live in ``torchref.refinement.model_error_estimation._shells``; ``estimate_beta`` is unchanged +- Align difference-refinement tests with fixture ownership, integration placement and configured dtype/device conventions; share fresh collection setup and exercise noise statistics on deposited amplitudes. +- Read scaled observations directly in subset and collection accessors, avoiding unused sigma/amplitude corrections and removing redundant internal forwarding helpers. +- Remove one-off diagnostic scripts and consolidate difference-refinement regression tests while retaining numerical and output-format coverage. +- Restore the cell and space-group imports needed by difference-density validation setup. +- Keep ``torchref.phased-difference-map`` as an alias for ``torchref.difference-map``. +- Add a matched classic/AMBER refinement benchmark with shared prepared hydrogens, initial riding-parameter gradient calibration, and per-cycle R factors and geometry diagnostics. +- Hydrogen generation respects tetrahedral ammonium nitrogen types, retaining the third hydrogen on protonated lysine and free amino termini while peptide nitrogens retain their linked valence; terminal H1 aliases participate in dictionary bonds without renaming model atoms. +- AMBER targets consume every TorchRef-owned atom through a validated atom map, including hydrogen and riding-orientation gradients; incomplete chemistry is rejected at setup, and ensemble AMBER initialization leaves coordinates unchanged unless relaxation is explicitly requested. +- Enabling riding mode completes missing HOH hydrogens only when hydrogen generation is enabled and stripping is disabled; existing atom coordinates and refinement selections are preserved. +- Riding coordinates refine shared methyl/hydroxyl torsions and water rotation vectors, preserve orientations through selection and checkpoints, and expose their gradients to the xyz optimizer; explicit hydrogen generation includes HOH with dictionary geometry and seeded random initial orientations. +- ``hydrogens_in_xray`` (model constructor, ``Refinement``, ``torchref.refine --hydrogens-in-xray/--no-hydrogens-in-xray``, default on) decides whether hydrogens enter the structure factors; it replaces ``exclude_H_from_sf``, which remains as a deprecated inverted alias, and is carried by ``copy``, ``select``, the strip/hydrogenate helpers and state dicts. The bulk-solvent mask is built from heavy atoms only, whatever the model carries. +- ``torchref.base.coordinates.local_frame`` holds the differentiable local-frame placement (``place_local_frame`` and its inverse); the AMBER target's hydrogen placement now delegates to it. ``torchref.topology.hydrogens`` gains ``HydrogenFrames`` and ``hydrogen_frames`` (which heavy atoms each hydrogen rides on, read off the bond graph) and ``augment_atom_table_with_maps``. +- ``RidingXYZTensor`` (``torchref.model.riding_xyz``): a coordinate wrapper over the full atom table whose hydrogen rows are derived from their parent heavy atoms each forward, so only heavy atoms are refined and a force on a hydrogen lands on the atoms that carry it. ``Model.set_hydrogen_mode("riding" | "free")`` and ``Refinement.set_hydrogen_mode`` switch a loaded model; ``ModelContext.hydrogen_mode`` records the choice and it survives ``copy``, ``select``, ``shake_coords``, rigid-body passes and state dicts. ``Restraints.copy`` no longer deep-copies the coordinate, ADP and radius accessors it borrows from the model. +- Monomer templates keep the CCP4 ``type_energy`` of each atom, ``chem_mod_atom`` ``change`` rows retype linked atoms (an in-chain backbone N is ``NH1``, not ``NT3``), and the atom graph carries ``energy_type`` and ``template_h_count`` per atom. The per-type table from ``ener_lib.cif`` ships as ``torchref/data/ener_lib_atoms.csv`` (regenerate with ``python -m torchref.scripts.extract_ener_lib``). +- Hydrogen generation on model loading is off by default; use ``torchref.refine --add-hydrogens`` or ``add_hydrogens=True`` in Python to opt in. Hydrogens already present in input files are retained unless ``strip_H=True``. +- Hydrogen generation reads the user's restraint CIF: ``cif_path`` is a model constructor argument, set before loading by ``torchref.refine --cif`` and the shared CLI loader, and carried by ``hydrogenate``, ``strip_hydrogens``, ``select`` and state dicts. +- mmCIF models carry their covalent and metal ``_struct_conn`` links, as PDB LINK records already did. +- A LINK record repeated in a file, or a bond emitted once per altloc conformer, now counts once in the atom graph. +- Hydrogen count per atom is capped by the template's own hydrogen count minus extra covalent partners, so linked hetero atoms (acetyl caps, Schiff bases, glycosylated ASN, metal-bound HIS) no longer receive displaced hydrogens; shared backbone atoms of split residues keep their HA/H. +- Pull requests into ``dev`` run one Python 3.12 CPU test job and one MPS test job, without the dependency-compatibility matrix. +- Loss aggregation uses the configured floating-point dtype in eager and compiled execution, including empty and zero-weight aggregates. +- Named broad structure-compatibility cases explicitly, moved extra datasets to the slow tier, and removed eager all-model loading and swallowed reader failures. +- Consolidated CIF/MTZ loading contracts, checked configured tensor placement, and replaced ModelFT smoke checks with exercised forward-cache behavior. +- Replaced local-arithmetic target tests with configured-device production-kernel checks on deposited coordinates and explicit least-squares expectations. +- Consolidated weighting tests by API ownership and strengthened Gaussian-likelihood, gradient-norm, and cached-loss assertions. +- Organized test fixtures into focused modules and reused module-scoped loaded objects for read-only functional checks while retaining fresh objects for mutation and loading tests. +- Rigid-body refinement stores its Euler angles pre-multiplied by the chain's radius of gyration, so a unit step in an angle and a unit step in a translation displace atoms comparably. In radians against Angstroms the rotation block of the Hessian carried 190-530x the curvature of the translation block on 1DAW and 3E98 -- the geometric ``Rg**2``, 411 and 442/516 -- putting ``cond(H)`` at 1e3-5e3, which is why six parameters needed ~250 L-BFGS iterations to place. Dividing the scale out in ``forward()`` brings the ratio to 0.4-1.3 and ``cond(H)`` to 3-18. Over ten structures the step then converges rather than exhausting its iteration budget, on about half the gradient evaluations, with R-free no worse anywhere. ``RigidXYZTensor.rotation_radians`` returns the physical angle, and setting ``angle_scale`` to ones restores the unscaled parametrization. Not a fix for the one or two negative Hessian eigenvalues at the finer cutoffs -- scaling a saddle leaves it a saddle -- and those counts are unchanged +- The rigid-body step no longer co-refines the scaler in the same L-BFGS as the rigid parameters. The body target centres on ``alpha*|F_calc|`` and ``alpha`` absorbs a rescaling of ``F_calc`` exactly, so the scale had a flat direction there; ``SCALE_TARGETS`` already excludes every alpha-centred row from the scale fit for this reason, and 0.6.2 fixed the same thing in the main driver. ``refine_scaler`` (objective ``ls``) owns the scale, between cutoffs +- Fixed ``refine_rigid_body`` leaving the caller's reflection data truncated. ``cut_res`` masks in place and returns ``self``, so each cutoff stamped its resolution mask on the caller's own object and the restore had nothing to restore to -- it only looked correct because the default schedule ends at the native limit. With ``--rigid-body-cutoffs 6,4`` on a 2.05 A dataset, 20138 of 23352 reflections stayed masked out for the rest of the run, R-factors included +- Joint sigma-weighted dataset scaling now uses centered corrections owned by ``DatasetScaler`` and exposed by ``ScaledDataset`` through ``DatasetCollection.scale()``, including scaled uncertainties, partial overlap, checkpointing and anomalous reflection identities. +- Remove parameter state, scale fitting and E-value conversion from ``ReflectionData``; use ``WilsonNormaliser`` for E values, observation attributes for full arrays and subset views instead of deprecated getters or ``data()``. +- Match batched mixed-solvent accumulation to individual model evaluation to reduce single-precision cancellation error. +- ``torchref.phased-difference-map`` is now ``torchref.difference-map``, because it no longer defaults to a phased difference map. The default output is ``DELFWT``/``PHDELWT``, the inverse-variance-weighted amplitude difference on the **dark** model's phases -- the construction ``torchref.validate-ded`` correlates against, and the one the figure-4 CC was measured on. The writer computed that pair all along as ``WDF``/``PHIC_dark`` but foregrounded ``2mDFop-DFc``/``PHIC_diff`` instead, a phased *residual* that puts the light state's model phases into the observed amplitude and so biases the map toward the model under test. The headline map and the validation metric are now one object +- ``-lm``/``--light-model`` is optional on ``torchref.difference-map``. A weighted difference map needs only the dark state: the amplitude is ``|Fo_light| - |Fo_dark|``, which ``DatasetCollection.scale`` puts on one scale with no model at all, and the phase comes from the dark model. Without a light model there is no mixed model, no ``--fraction`` and no joint model-to-data fit, and ``--fraction`` is rejected rather than ignored. Still required on ``torchref.difference-refine``, which refines it +- The difference MTZ went from 33 columns (46 under ``--two-moment``) to 17, with the rest behind ``--all-columns``. The default file holds the difference map, the extrapolated map ``FWT``/``PHWT``, the observations and the flags; the gated set holds the alternative constructions -- the phased difference residuals, two further extrapolations, the intensity block. Standard CCP4 names throughout the default set, so Coot and CCP4 open both maps without being told which columns to use. The default path also runs one internal scale fit instead of three, and none at all with no light model +- Fixed the empirical-Bayes extrapolation propagating ``sigma_ext^2 = (sigma_L^2 + sigma_D^2)/f^2``. ``F_ext`` is linear in the observations with ``dF_ext/dF_dark = -(1-f)/f``, so the dark term carries a ``(1-f)^2`` weight: at f = 0.22 the term was over-weighted 1.64x, biasing ``tau^2`` low and over-shrinking every reflection. The estimator now takes the caller's already-correct propagation instead of rebuilding it, so the shrinkage weight and the written ``SIGFEXT`` cannot disagree. It also returns the shrunk amplitude it is named for; the caller used to discard the return value and recompute it. Correcting it raises tau^2 and shrinks less, which moves the default extrapolated map +- ``FWT``/``PHWT`` carries the phase of its own scale fit. The Bayes coefficients had no phase column and would have been paired with one fitted against the phase-aware amplitudes; the scaler contributes a phase through ``f_sol``, so the three extrapolation fits do not agree +- ``write_results_mtz`` is four layers, each declaring its columns' MTZ types beside the values, and it now calls ``infer_mtz_dtypes`` as the canonical writer does. Types used to come from four parallel name lists with no fallback, so a column added to the output dict and missed in the lists was written with whatever dtype numpy produced. A column with no declared type is now an error. Its column table moved onto the function from the module docstring four hundred lines away +- ``DFc_complex`` is ``DFc_phased``: the column holds a real amplitude, and the old name said otherwise. The undocumented ``Fextp``/``Fextc``/``Fextb`` suffixes are ``FEXT_PHASED``/``FEXT_SCALAR``/``FEXT``, and ``SIGFEXT_PHASED`` is written at last -- it was computed all along while the docs claimed it existed +- ``tau_sq`` and mean ``w(h)`` go in the JSON summary. With the Bayes extrapolation as the default map, they are what says whether it is over-shrunk +- Fixed ``paper/make_ded_maps.py`` pairing ``WDF`` with ``PHIC_diff`` -- the right amplitude on the model difference phase, which is not the map ``validate-ded`` reports +- Refinement output no longer inherits the input file's refinement header. It used to copy the whole thing and then append its own ``REMARK 3``, so a refined 3GR5 carried 420 header lines asserting two refinements at once -- ``PROGRAM : REFMAC 5.1.24`` with R-work 0.213 at line 5, ours at line 389 -- and a reader taking the first ``REMARK 3`` got REFMAC. The inherited block was not merely stale but contradicted the data beside it: it claimed a 5.1% / 1072-reflection test set, while the MTZ shipped with it holds 9.85% / 2063 (which torchref reads correctly). The passthrough was also inverted, keeping the statistics refinement invalidates and dropping the chemistry it does not -- SEQRES, SSBOND, DBREF, EXPDTA, COMPND, SOURCE, KEYWDS, SEQADV, HETNAM, FORMUL and SITE were all absent from the output. Now a whitelist carries the crystal, sample and chemistry records through in mandated record order (TITLE used to be emitted after REMARK 900), REMARK 2, 3 and 500 are dropped, and AUTHOR and JRNL are not inherited because they credit the deposition rather than this run. 283 header lines for the same file, 41 structural records preserved, one refinement block. ``add-metadata`` is exempt through ``supersede_refinement=False``: annotating a file is not re-refining it, so nothing there supersedes the existing REMARK 3 or AUTHOR records and both are kept +- Prior refinements are tracked through mmCIF's ``_software`` loop, which is the only place either format has room for them: ``_refine`` is singular by design, so a previous program's statistics cannot be kept without contradicting the current ones. ``pdbx_ordinal`` was hardcoded to ``1`` and the incoming loop was never read, truncating the chain to one link on every write; it now reads the input's loop and appends at ``max(ordinal) + 1``, carrying each entry's ``description`` so the chain says what every program did and not just that it ran. Added ``_pdbx_initial_refinement_model`` and ``_refine.pdbx_starting_model``, which name what the refinement started from, and ``_refine.pdbx_R_Free_selection_details``, which names the test set the reported R-free conditions on. ``from_cif_file`` no longer carries the input's ``_refine`` items through -- that was the mmCIF form of the duplicated ``REMARK 3`` +- Fixed mmCIF loop cells being written unquoted, which silently split any value containing whitespace into extra columns when the file was read back. Latent while every loop column was a single token; a multi-word ``_software.description`` exposed it. Nulls stay bare, since ``gemmi.cif.quote`` turns ``?`` into the quoted one-character string ``'?'`` +- Fixed PDB coordinate and B-factor columns being written as ``str()`` of a rounded float rather than with an explicit precision, so trailing zeros were dropped -- ``18.3`` and ``31.51`` where the format wants ``18.300`` and ``31.510`` (369 of 1329 atoms in a 3GR5 refinement), and ``95.4`` where it wants ``95.40``. Both writers were affected. Columns held and the values round-trip through any ``float()``-based reader, so nothing was numerically wrong; this is format conformance +- The header records what was refined and how. ``refinement_method`` was only ever set by the difference-refinement CLI, so ordinary runs emitted no method line at all; ``TARGET`` and ``OPTIMIZER`` now come from the same settings that reach ``refinement_history.json`` -- target family and registry key, macrocycles, refinement mode, ADP model and the scale target, the last because it changes the R-factors the same header reports. Long values wrap onto a continuation line with the colon held in column 25, the convention REFMAC uses for its own author list. ``rfree_source`` was set only when the MTZ also carried a validation column, leaving the common FreeR-only case unattributed, and generated flags did not record their seed -- both fixed, so the free-set provenance in the header is always either the file it came from or a reproducible draw +- Author-supplied header text goes in ``--output-remarks``, rendered as ``REMARK 3 OTHER REFINEMENT REMARKS`` and ``_refine.details`` and emitted only when set. Nothing in the block is editorialised: every generated line is a measured quantity or a recorded setting. Removed the deprecated ``template=`` argument to ``pdb.write`` (nothing called it) and the never-populated ``custom_remarks`` field. ``RefinementMetadata`` had no test coverage at all, which is how the duplicated ``REMARK 3`` shipped; 36 tests now cover the contract +- Fixed the rotation function contracting **unconjugated** calc coefficients on MPS. ``torch.conj`` returns a lazy view carrying a conjugate *bit*, and MPS's batched complex matmul -- what the radial ``einsum`` lowers to -- ignores it, so the contraction was silently wrong: on 1DAW by 173% of ``|xi|`` max, which reordered the entire peak list and pushed the true orientation from rank 0 out of the top 200 while the top score moved only 0.1%. Materialised with ``resolve_conj()`` at the two batched-matmul sites. Elementwise ops, ``where``, ``index_add``, ``fft`` and 2-D matmul all honour the bit and only the batched path does not, which is too narrow to guard structurally, so the guard is the contraction's own value against a host-double reference +- The alignment package runs end to end on an accelerator without float64: 123 alignment tests including the slow rotation searches pass on MPS, where every one of them failed before. Five things stood in the way, four of them the same mistake -- casting and moving in two steps. ``.to(device).to(dtype)`` puts the caller's width on the device first, which throws for a double input on a backend that has none (14 sites, now one fused ``.to`` each); ``detect_zsymm`` widened the symmetry operators to double while they sat on the device; the expansion's host-side clustering keys left a host index to meet device values; and MPS's ``linalg.inv`` trips an internal contiguity assert on a transposed 3x3 view. The fifth is a kernel gap -- MPS implements no complex cumulative op -- so the azimuthal phase ladder is built by doubling instead of by ``cumprod``, measured at 5.7e-6 against ``cumprod``'s 4.1e-6 in complex64 over 5000 angles at L=101 +- The anisotropy fit runs at the configured float dtype on the data's own device, instead of double on the host. Measured over the 16 datasets in ``tests/files/mtz``, float32 reproduces the double fit to 3.3e-5 relative in ``U`` and 4.5e-6 in the correction factor ``exp(+pi^2 s.U.s)`` it exists to produce -- the design matrix is a constant column beside ``2 pi^2 s.s`` terms of order 0.1-1, so there is no precision cliff for seven parameters to fall off. Unrotated searches on 1DAW, 3K7M, 2DQ6 and 4BX9 return the same peaks in the same order as before, with scores moved 1.8e-7 to 2.5e-4 relative. ``fit_anisotropy`` takes an optional ``device``; ``get_axis_order`` and ``hkl_symops_to_cartesian`` follow their inputs' width and place rather than forcing double +- The Wilson fit's IRLS solve is a Cholesky reuse plus two triangular solves, not ``lu_factor``/``lu_solve``. ``A`` is ``X^T W X`` plus a ridge, positive definite by construction, and MPS implements neither ``lu_solve`` nor ``cholesky_solve``, so the old path left the per-iteration solve to a CPU round trip: 1015 us against 114 us, on a 200k x 6 problem whose unavoidable ``XtW @ z`` is 966 us. Over the 16-dataset panel the two agree to 1e-5 relative in float32 and 1e-14 in float64 with identical iteration counts on every case. ``cholesky_ex`` reports rather than raises, and a non-positive-definite report falls back to a general solve, since a fully collinear basis is a thing this fit sees +- Dropped the two alignment tests that imported ``alignment_lab``: an external experiment is not the main package's to test, and they could not pass in a clean checkout. ``resolve_device`` is covered directly in ``tests/unit/utils/test_device_resolution.py``. Two test-side assertions that could not hold at the configured width were repaired -- `` = 1`` now widens after the readback rather than on the device, and the weight's mean-one bound tracks the dtype's epsilon instead of a fixed 1e-9 that is below one float32 ulp +- Nothing in the alignment package puts float64 or complex128 on the compute device any more, so it runs on backends without double (MPS). The rotation function's radial sum, Wigner contraction and FFT accumulate at the configured complex dtype rather than one step wider; the Bessel ladder runs in its argument's dtype, kept in range by its rescaling; the expansion's exact clustering keys are formed on the host in double, which the device never sees; rotation-function inputs, the dense P1 transform and the shell sums use the configured float dtype. 30/30 poses at both windows, timings unchanged; the double-precision left is host-side 3x3 rotation algebra, the Wigner-d eigendecomposition and the anisotropy fit +- Every hard-coded dtype in the alignment package either moved to the configured dtype or carries a ``# dtype-ok:`` justification, as the dtype-conformance guard requires: three casts to double dropped from the empirical sigma_A ratio and the dense P1 transform, the translation set's resolution mask and peak translations use the configured float dtype, and the rest (index tensors, host-side 3x3 rotation algebra, the rotation function's double accumulation, the anisotropy fit) are annotated +- Removed ``DirectModelEvaluator``; the translation search evaluates an ordinary P1 ``ModelFT`` directly. With the model's grid derived lazily from cell, space group and ``max_res`` there was nothing left for the wrapper to do +- The molecular-replacement pipeline carries 10 rotation candidates by default instead of 25. With symmetry mates suppressed the rotation function's first peak is the true orientation in 50 of 50 pose-gated cells, and the panel is 30/30 at either depth. Warm on one EPYC 9335 node: 1DAW 0.42 s, 2DQ6 0.67 s, 6G9X 0.67 s, 3K7M 1.02 s, 4BX9 1.11 s per alignment +- The fast translation function accumulates only the upper triangle of symmetry pairs; the lower triangle is its conjugate mirror and the diagonal a constant. Half the scatter, which is the stage's cost on high-symmetry cells: 3K7M's translation stage 1.04 s to 0.72 s. ``MRSolution.candidate_index`` records each solution's position in the rotation function's list +- The placement loop re-orients one P1 copy of the search model in place per candidate and builds the placed model for the winner only, instead of copying the model three times per candidate. ``MRSolution.model`` is ``None`` for the other candidates; ``MolecularReplacementPipeline.place`` builds it on request. Warm on one EPYC 9335 node: 1DAW 0.85 s, 2DQ6 1.35 s, 6G9X 1.30 s, 3K7M 2.35 s, 4BX9 2.55 s per alignment (from 1.1, 1.9, 2.2, 2.7, 3.8) +- The three candidate-ranking scores (likelihood, analytical R, fast translation score) pick the same candidate in all 60 cells of a six-structure x ten-seed sweep measured on poses; every earlier figure quoted for them was rotation-only. The likelihood stays the default +- The rotation function suppresses symmetry mates when it picks peaks, so its shortlist is one entry per orientation. The point group composes on the right of a peak's rotation, ``R R_g`` -- measured on real peak lists, where the left orbit finds no coincident pairs and the right finds every mate (187 of the 300 pairs among 3K7M's top 25 were mates of each other). The lab's orbit-based truth rank defaulted to the left side, which is why it disagreed with coordinate superposition. Placements unchanged, 30/30 at both windows +- Fixed ``empirical_sigma_a`` taking the level of the observed-to-calculated Wilson-curve ratio rather than its shape. The two curves carry different absolute scales, so the ratio was 0.02-0.06 on one structure and 8-12 on another and the returned ``sigma_A`` was flat at 0.15-0.35 regardless of resolution; each curve is now divided by its geometric mean first. Per-shell factors are gauge in the rotation function's correlation, so its placements are unchanged (30/30) +- The fast translation function scores the covariance of two normalised intensities -- the rotation function's LERF1 coefficient ``cw (E_obs^2 - 1) w sigma_A^2`` against the candidate's ``|E_calc(h, t)|^2``, normalised per candidate by the same Wilson fit -- instead of a raw-``|F_calc|^2`` ratio that was not a correlation and, on the four largest panel structures, was higher 40 A from the true position than at it. One FFT on a grid a third of the set's resolution apart with parabolic peak refinement replaces the 16-point coarse grid and three 100-point local refines; the Rice/Woolfson likelihood at a fixed Luzzati ``sigma_A`` picks among the top peaks and ranks the candidates. 30/30 true poses at the default window and 30/30 with the window removed, against 18/30 before +- The translation stage runs in the configured float and complex dtypes rather than hard-coded double +- Removed ``use_llg_tf``, ``n_translation_peaks`` and ``translation_grid_steps`` from the pipeline; the likelihood always picks the translation, and the grid is sized by resolution +- Removed the per-candidate ``SigmaAEstimator`` fit from the placement loop. The likelihood that ranks candidates uses one ``sigma_A`` for all of them, so no candidate is scored against a model error fitted to itself +- The translation search defaults to the rotation search's resolution window instead of all data. With all data it placed the four largest panel structures (2DQ6, 3VRJ, 4BX9, 6G9X) at the right orientation and 20-56 A from the true position on every trial, and its own score was higher at the wrong place than at the deposited pose; the benchmark had only ever checked the rotation. 30/30 true poses within 0.32 A against 18/30, at roughly a fifth of the wall clock +- The P1 copy each rotation candidate is evaluated through is gridded at two thirds of the translation window's resolution rather than at the model's default 1.0 A. Coherence with the fine grid 0.9995 or better on all four structures measured; 10-38 ms per candidate against 200-860 +- The pose-recovery harness compares against Cartesian symmetry mates and checks the translation. It compared a Cartesian rotation against fractional symmetry matrices, so in trigonal and hexagonal cells two of six and four of twelve correct mates read as 30 and 21 degrees; every recorded 2DQ6 failure was one of them +- Repaired the alignment integration tests, which imported a deleted module and passed removed arguments +- The alignment package uses the shared FFT-size and Euler-matrix helpers instead of its own copies, and drops four unused rotation utilities and two duplicated symmetry helpers. Bit-identical placements on the benchmark panel +- The translation likelihood's model error comes from the shared ``SigmaAEstimator`` instead of a local 81-point scan over every reflection. It returns ``alpha`` and ``beta`` per reflection rather than a per-shell ``sigma_A``, so the likelihood no longer assumes ```` is exactly one. Outcome-neutral over 70 seeded cells, zero flips; not measurably faster +- Fixed the translation likelihood's variance convention, which scored acentric reflections at twice the variance intended -- 90-95% of reflections. The alignment package carried its own Rice and Woolfson parameterised by the *amplitude* variance and handed both branches the same number, where the acentric branch needs half what the centric one does. It now uses ``base.targets.xray_likelihoods.rice_per_refl``, which takes the complex variance and derives the centric case from it +- Removed ``experimental/alignment/distributions.py``. Its ``stable_log_bessel_i0`` also carried a wrong asymptotic coefficient, giving -2.9e-5 at x = 50 against -5.4e-7 for the correct term; the shared implementation uses ``log(i0e(z)) + z``, which is exact +- Molecular-replacement candidates are ranked by the translation function's likelihood, not by an analytical-scale R-factor. 30/30 on the ten-structure panel against 29/30. Over a ten-seed sweep the two tie at 36/40 and the likelihood's advantage is placement rather than count -- median residual 1.43 deg against 1.62, carried by one structure. Selectable through ``rank_by``; the correlation is the third option and is the worst of the three at 32/40, despite a rank-level harness rating it best on a truth label that disagrees with coordinate superposition +- The likelihood ranking costs about 2.5x the wall clock, from a per-candidate sigma_A fit +- The molecular-replacement pipeline's ``verbose`` levels are a documented contract routed through one emitter, rather than ``if verbose > 0: print(...)`` at seventeen sites. Level 2 emits one machine-readable ``CAND`` line per rotation candidate carrying every score the selection could have used, so a wrong placement can be diagnosed from the run itself instead of from a harness that re-implements the placement loop and then disagrees with it +- The molecular-replacement pipeline places all ``n_rotation_candidates`` (now 25, was 15) and returns the best, instead of stopping once a placement beat an R-factor threshold. The old rule made the answer depend on the order the rotation function happened to produce and could accept the third candidate without scoring the tenth. Measured neutral -- identical placements on 10 structures x 3 seeds and on a 10-seed sweep of the marginal cases -- at about 1.6x the wall clock +- The Wilson normaliser converges in 8-12 IRLS iterations instead of 26-102. Its stopping rule was ``|dL|`` per reflection against 1e-10, which asks eleven significant digits of a normalisation curve; it is now relative to the improvement so far, which is scale-invariant for the same reason the absolute form was chosen +- `` = 1`` is solved in closed form for the intercept, so the identity no longer degrades as the convergence tolerance is loosened +- The Wilson fit runs in the configured float dtype rather than hardcoded double, and builds its normal-equations matrix once -- it is constant for a Gamma with a log link. Identical placements on all 30 benchmark cells, at 1.8x the speed; the unit suite went from 968s to 485s +- Molecular replacement is now a rotation search feeding a translation search and nothing else; the pipeline returns a placement and stops. End-to-end pose recovery over 10 structures x 3 seeds went 18/30 to 30/30, at about a sixth of the wall clock +- Removed the ML rescore from between the two searches. It reordered a shortlist that already contained the answer, and cost 6 of 30 placements +- Removed the post-placement dense rotation re-sampling and rigid-body polish. They refined a correct placement away from truth on 2DQ6, 3GR5 and 4BX9; refining a placement is downstream refinement's job +- The translation search weights reflections by inverse variance, which it previously did not do at all, and both searches normalise through one shared Wilson fit built once per run instead of five private ones. This is what recovered 6G9X +- The translation search's resolution window is a parameter (``tf_d_min``/``tf_d_max``) rather than a docstring; it defaults to the existing behaviour of no cut +- The alignment package takes its device from the configured default throughout, instead of reading it off whichever model or tensor was nearest +- Removed the alignment package's unreachable modules: quaternion transforms, a second Wigner implementation, clash scoring, vector sampling, the Lattman-Love interpolator, the E-value convention layer and its French-Wilson posterior, and the unused half of the spherical-harmonic expansion +- Fixed the reciprocal-space symmetry convention in the alignment package (``h.S``, not ``S.h``) +- Fixed ``hkl_symops_to_cartesian`` returning non-rotations in trigonal and hexagonal settings, which corrupted the anisotropy projection +- Fixed the overall-anisotropy fit, which regressed log intensities with no constant term and so absorbed the ``-gamma`` offset into the tensor +- Fixed molecular-replacement rotation candidates being composed onto each other instead of onto the search model +- Fixed assigning a ``SpaceGroup`` object to ``Model.spacegroup`` being a silent no-op that then made the correct name assignment raise +- The ML rescore can score orientations by weighted least squares on E-space intensities (``target='wls'``) instead of the Rice/Woolfson likelihood. Measured indistinguishable from the Rice over 10 structures x 10 seeds; both remain worse than not rescoring at all +- The rotation function estimates ``sigma_A`` from the data instead of assuming it. Total scattering per shell is rotation-invariant, so ``Sigma_obs(s)/Sigma_calc(s)`` measures the model's resolution-dependent deficiency before placement; it replaces the Luzzati falloff from an estimated coordinate error and the Babinet bulk-solvent term with its two universal constants +- Removed the relative Wilson-B match from the rotation search. It multiplied ``F_calc`` by a smooth function of ``|s|`` that the normalisation then divided straight back out +- Added ``torchref.scaling.weighting``: measurement and model error combined as one inverse-variance weight per reflection, restoring the observed-side weighting that left with the French-Wilson posterior +- ``apply_shell_variance_weights`` is off by default in the rotation function. It is a per-shell weight, and per-shell weights are absorbed by the correlation; switching it on moves nothing +- The rotation function, ML rescore and translation search now normalise through the shared Wilson normaliser by default, replacing the French-Wilson posterior on the observed side. Rank-neutral over 10 structures x 10 seeds; the rotation function is about twice as fast, since the posterior and its D-factor iteration are no longer on the default path +- The observed-side ``DFAC`` weighting went with it. Weighting is a separate concern from scaling and is being rebuilt as its own object; until then the rotation function applies no measurement-error weight +- Added ``torchref.scaling.WilsonNormaliser``: an absolute normaliser that fits ``Sigma(s)`` as a Gamma GLM with a log link and divides it out, so `` = 1`` holds as an identity of the fit rather than as a separate normalisation step +- Extracted the Chebyshev resolution basis into ``torchref.scaling.basis``, shared with the isotropic scale, and gave it an explicit range so a curve fitted on one reflection set can be evaluated on another +- Moved epsilon onto ``SpaceGroup.epsilon(hkl, friedel=)``; the alignment package's own copy disagreed with it in trigonal and hexagonal groups and dropped the centring coset. The default keeps the Friedel-folded count sigma_A is calibrated against, and the molecular-replacement likelihood asks for the conventional one +- The rotation function and the ML rescore take their E-value convention as a class, so the observed and calculated sides are normalised by one rule rather than by nine private converters +- The ML rescore now receives the observation sigmas, which the rotation function computed the French-Wilson posterior from and then discarded +- Removed the rescore's ``scat_mode``, a second knob for the decision the E convention already makes +- The Sim rescore and the three translation-search sites take their E values from the convention too, so the alignment package has one normaliser rather than nine +- Removed ``wilson_normalise`` and ``wilson_normalise_epsilon``, which the convention replaced +- Fixed the rotation search, placement pipeline and ``align`` reading ``model.initialized``, which moved to ``model.ctx``; the tests covering it are slow-marked, so the break was invisible to a default test run +- Replaced the fast rotation function's keyword surface with ``rotation_search(model, data, model_error_A)``; the caller's coordinate error is now used rather than overwritten by an estimate from the atom count +- Removed the rotation function's dead modules, engine variants, debug environment switches and unreachable knobs +- ``Model``'s iso/aniso partition is now derived on access instead of being rebuilt eagerly, so a copy cannot inherit a stale one +- The rotation function's Wigner small-d blocks are memoised, so a process running more than one search builds them once +- The rotation function now takes its working precision from ``dtypes.float`` and its device from ``resolve_device``, instead of hardcoding float64 and reading one input's device +- The rotation function's spherical-Bessel recurrence rescales by a power of two as it runs, so the ladder no longer needs float64's exponent range +- The rotation function's Wigner eigendecomposition and anisotropy fit moved to the host, so neither requires float64 on the accelerator +- Removed the rotation function's duplicate Euler, Rodrigues and reciprocal-symmetry helpers in favour of the shared primitives +- Removed the rotation function's redundant calc-side resolution mask and its second bandwidth/resolution coupling call +- ``bessel_sh_expand`` lost its unread ``chunk_size`` argument and ``french_wilson_preprocess`` its unread ``sqrt_mean_F2`` output +- The rotation function's observed-side chain (French-Wilson, LERF1, shell variance weights, relative Wilson B) now runs once per unique reflection instead of once per symmetry copy; only the geometry is unrolled +- The rotation function assigns resolution shells once and shares them, instead of the Wilson normalisation and the variance reweight deriving edges that disagreed at the shell boundaries +- The rotation function masks observations to the bandwidth-coupled resolution before the symmetry unroll rather than after +- The rotation function warns on Bijvoet-unmerged data, whose shared canonical index would weight those reflections twice +- The rotation function no longer concatenates the antipodal copy onto either reflection set: only even harmonic degrees are computed, for which it is an exact factor of two, so it scaled the rotation function by four and changed no ranking. Raw ``RotationPeak.score`` and ``RotationSolutions.scores`` are therefore a quarter of their previous values; z-scores are unchanged +- Kept the rotation function's relative Wilson-B fit: knocking it out was measured rank-neutral but worth only 2% of the runtime once the fit moved to the unique reflection set +- Added ``supports_double`` / ``widest_float_dtype`` / ``widest_complex_dtype``: where precision is load-bearing the width now comes from the device rather than a hardcoded ``float64``, so a backend without it gets the working dtype instead of an error +Version 0.7.0 +---------- +- Fixed cif reading bug discarding new mmCIF field for aniso ADPs +- Removed the stored real-space coordinate grid; ``build_electron_density`` takes a grid shape and device, and ``ModelFT.real_space_grid()`` builds one on demand +- Fixed ``ModelFT`` restore dropping a node-field ADP representation, and added the anisotropic ``field_aniso`` case; both models now share one wrapper-rebuild path +- Fixed the node load and node smoothness restraints being inert in ``field_aniso`` mode +- ``create_from_state_dict`` now restores on CPU and moves only when passed a device; it previously left three of the four parameter wrappers on CPU while claiming the default device +- Separated model configuration and provenance into ``ModelContext``. It now holds the unit cell, space group, atom table, link records, hydrogen settings, and input paths. +- Refactored ``Symmetry`` as a crystallography-free class with transform primitives, and made ``SpaceGroup`` a specialised subclass. +- Moved geometry predicates, HKL verbs, and grid-size helpers onto these classes as methods. +- Rebuilt geometry restraints from the topology instead of intra-residue builders. ``torchref.restraints`` was removed, restraint dictionaries are now plain nested dicts, and residues are identified by ``(chain, resseq, icode)`` to fix insertion-code merging. +- Reworked hydrogen generation as template instantiation over the topology. ``Model.hydrogenate`` now aligns monomer templates onto heavy atoms present, generation is the default, and ``AtomGraph.exclusions_12_13_14`` derives non-bonded exclusions from bond connectivity. +- Added ``Topology`` as a ``ResidueGraph`` over an ``AtomGraph`` with typed edge blocks and ``subset`` / ``copy`` operations that reindex surviving edges. +- Made ``HydrogenTopology`` a dataclass, changed ``Symmetry`` classes to dataclasses over ``DeviceMixin`` instead of ``nn.Module``, and removed unused ``Cell`` gradient plumbing and the ``ReciprocalSymmetryGrid`` / module-level expansion functions. + + Version 0.6.4 ---------- +- ``torchref.validate-ded`` records ``mask_source`` in its results JSON; it changes the correlation and was not recoverable from the output +- Fixed the ``--two-moment`` corrected DED coefficients using a phase-blind amplitude difference instead of the phase-aware one the uncorrected coefficients use +- Added ``paper/make_ded_maps.py``, which writes CCP4 maps from a difference-refine results MTZ +- ``CollectionScaler.refine_lbfgs_joint`` builds a row of ``XRAY_TARGETS`` instead of its own Rice likelihood, and takes ``scale_target`` (default ``ls``) +- ``CollectionScaler.refine_lbfgs_joint`` normalises its objective and registers the U penalty as its own target +- Fixed ``DatasetCollection.scale`` fitting the inter-dataset scale on the free reflections as well as the work set +- ``DatasetCollection.scale`` normalises its objective, so L-BFGS's absolute tolerances mean something +- Added ``COLLECTION_XRAY_TARGETS``, the collection target taxonomy, with an intensity difference row +- Removed ``CollectionRiceTarget``, which set ``beta = sigma_obs**2``; the ``ml`` row is the absolute channel instead +- Renamed the kinetic ``xray_weight_rice`` / ``xray/rice`` weight to ``xray_weight_ml`` / ``xray/ml`` +- ``--lambda-twin`` now requires ``--two-moment``; the activation dispersion belongs in the predicted intensity, not in a weight +- Gave the collection targets the same ``_loss_inputs``/``_per_refl`` seam as the single-dataset ones, with the observable declared per row +- Added ``--xray-mode nll_i``, a Gaussian on the observed intensities, and an ``observable`` column on the target taxonomy +- Added ``DataTarget.get_I_calc_scaled``, so the observable is a choice rather than an assumption +- Added ``gaussian_per_refl`` and ``intensity_var_from_sigma_obs``; the amplitude and intensity Gaussians are now one implementation +- The absolute variance floor in the shared Gaussian is now opt-out, since it distorts any objective whose sigmas fall below it +- Added a reader for CrystFEL ``partialator`` ``.hkl`` reflection lists, via ``ReflectionData.load_crystfel_hkl`` +- Added ``FcalcDataset.add_noise`` and the ``torchref.simulate-noisy-data`` CLI, which simulate merged intensities from a structure and report R-split and CC between two independent half-datasets +- Simulated intensities keep their negative values; only the derived amplitude is clamped, since clamping the intensity biases the weak reflections upward +- ``CollectionTwoMomentIntensityTarget`` carries a ``base_weight``, calibrated against the difference target's gradient norm so an intensity likelihood does not swamp the restraints +- Fixed non-finite observed intensities poisoning the two-moment gradient, which silently froze refinement rather than failing +- Added ``CollectionTwoMomentIntensityTarget``, fitting merged intensities as ``|F(alpha)|^2 + sigma_alpha^2 |dF|^2`` to account for crystal-to-crystal spread in activation +- Added ``--two-moment`` / ``--lambda-twin`` / ``--refine-lambda-twin`` to ``torchref.difference-refine``, and the activation moments to its JSON summary +- ``torchref.difference-refine`` writes thirteen further MTZ columns under ``--two-moment``, including decontaminated difference amplitudes and the ``DDF`` diagnostic +- Fixed ``torchref.difference-refine`` crashing at ``--verbose 0``, where the R-factors written into the deposition metadata were only computed for printing +- ``ModelCollection`` now stores populations as a shared activation fraction plus a per-timepoint branching, instead of free fractions per timepoint +- Freezing and unfreezing fractions is now collection-wide; timepoints needing independent populations use ``set_fraction_override`` +- ``add_timepoint`` raises when the requested fractions imply an activation that conflicts with one already set +- Added ``ModelCollection.sigma_alpha_sq`` and ``lambda_twin`` for the spread of activation across crystals +- Added batched ``compute_component_fcalcs`` / ``mix_component_fcalcs`` and ``DatasetCollection.component_structure_factors`` +- Added ``CollectionScaler.forward_batched`` for scaling several mixtures in one pass +- Added ``ReflectionData.get_corrected_intensities`` and scaled ``I``/``sigI`` subset views, with the unscaled values as ``I_raw``/``sigI_raw`` +- Added batched ``stack_F_obs`` / ``stack_I_obs`` / ``stack_masks`` accessors on ``DatasetCollection`` +- Fixed ``f_sol_override`` overwriting the scaler's cached ``F_sol``, so a later call without an override read the wrong solvent +- Fixed a batched ``f_sol_override`` gaining a spurious leading axis, which changed the rank of the scaled structure factors - Fixed the bulk-solvent ``F_sol`` staying at the starting model's mask for every refinement macrocycle - Fixed restraint dictionaries defining several compounds yielding restraints for only one of them - Fixed chirality restraints being dropped for the ``positiv``/``negativ`` spellings used by the CCP4 library diff --git a/docs/user_guide/cli.rst b/docs/user_guide/cli.rst index a649e62f..eadaa72d 100644 --- a/docs/user_guide/cli.rst +++ b/docs/user_guide/cli.rst @@ -29,6 +29,11 @@ and a ``refinement_history.json`` log. **Key options:** * ``-n`` / ``--n-cycles`` number of macro cycles (default 5) +* ``--add-hydrogens`` generate missing hydrogens on model loading (default off). + Hydrogens already present in the input are retained with or without this flag +* ``--hydrogens-in-xray`` / ``--no-hydrogens-in-xray`` include hydrogen atoms in the + structure-factor calculation (default on). Off keeps them in the restraints only; + the bulk-solvent mask is built from heavy atoms in either case * ``--mode`` ``separate`` (separated XYZ then ADP, default) or ``everything`` (joint XYZ+ADP) * ``--xray-mode`` one of ``ml`` (default; Read MLF at variance ε·β, conditional @@ -36,8 +41,10 @@ and a ``refinement_history.json`` log. ``ml_full`` (marginalises the measurement error rather than inflating the variance; ~4× the cost), ``nll_beta`` (the Gaussian large-signal limit of ``ml`` — diagnostic), ``nll`` (Gaussian weighted by σ_obs only, no model-error - term), ``ls`` (unit-weight least squares) or ``ls_wunit_k1`` (Phenix-style, own - global scale). ``--help`` lists them from the taxonomy table itself. + term), ``nll_i`` (as ``nll`` but on the observed *intensities*, skipping the + French–Wilson conversion), ``ls`` (unit-weight least squares) or ``ls_wunit_k1`` + (Phenix-style, own global scale). ``--help`` lists them from the taxonomy table + itself, which is authoritative. * ``--sigma-a-max`` upper bound on the per-shell Luzzati σ_A (default 0.99) * ``--no-shrink`` disable the per-shell σ_A stability shrinkage * ``--adp-mode`` ``isotropic`` (default) or ``anisotropic``, the latter refining @@ -84,7 +91,9 @@ restraints. ``-dsf``/``--dark-structure-factor``, ``-lsf``/``--light-structure-factor``, ``--fraction`` (light-state population fraction, singular), ``--weight-schedule`` annealing schedule (default ``5,3,2``), -``-n``/``--n-cycles`` macro-cycles. +``-n``/``--n-cycles`` macro-cycles, ``--difference-target {difference,difference_sd}`` +(the difference row the schedule drives; default ``difference``), ``--ded-weight`` and +``--sigma-d-gamma`` for the difference MTZ (see ``torchref.difference-map``). :API: :mod:`torchref.cli.collection_difference_refine` @@ -99,10 +108,17 @@ columns, expands to P1, and computes a real-space map via FFT. .. code-block:: bash - torchref.mtz2map -f refined.mtz -F 2FOFCWT -P PH2FOFCWT -o map.ccp4 + torchref.mtz2map -sf refined.mtz -csf 2FOFCWT -cphi PH2FOFCWT -o map.ccp4 + torchref.mtz2map -sf diff.mtz -csf DF -cw W_IVW -cphi PHDELWT -o diff.ccp4 + torchref.mtz2map -sf diff.mtz -csf DF -cw W_SD -cphi PHDELWT --units electrons -o diff_e.ccp4 -**Key options:** ``--high-res``, ``--low-res`` resolution limits, -``--gridsize`` override, ``-n`` normalize to sigma units. +**Key options:** ``--dmin``/``--dmax`` resolution limits, ``--gridsize`` override, +``-cw``/``--column-weight`` multiplies the amplitudes by a weight column before the +FFT, ``--units {sigma,electrons,raw}`` (``sigma``, the default, gives zero mean and +unit standard deviation; ``electrons`` gives e/A^3 as +``(1/V) sum_h F(h) exp(-2 pi i h.x)`` with the amplitudes divided by the +``-ck``/``--column-scale`` factor, ``KSCALE`` by default). ``-n`` is the deprecated +alias of ``--units sigma``/``raw``. :API: :mod:`torchref.cli.mtz2map` @@ -119,25 +135,64 @@ Computes real-space correlations and resolution-binned reciprocal-space CC. -dm dark.pdb -lm light.pdb **Key options:** ``--fraction``, ``--selection`` (Phenix-style atom -selection), ``--mask-radius``, ``--n-bins``. +selection), ``--mask-radius``, ``--n-bins``, ``--ded-weight`` (the headline weight +scheme; every scheme is also reported side by side, real-space in each mask and +reciprocal-space overall, as the ``by_weight`` block of the JSON and a table in the +summary, and a ``sigma_d`` fallback to inverse variance is recorded under +``weights``). :API: :mod:`torchref.cli.validate_ded` -``torchref.phased-difference-map`` -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +``torchref.difference-map`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Compute phased difference and extrapolated map coefficients without -refinement. Uses the same pipeline as ``torchref.difference-refine`` but -the input models are kept as-is. +Compute difference and extrapolated map coefficients without refinement. +Uses the same pipeline as ``torchref.difference-refine`` but the input +models are kept as-is. + +The default output is the difference map: the amplitude difference ``DF``/``SIGDF`` +on the **dark** model's phases ``PHDELWT``, with one mean-one weight column per +registered scheme beside it -- ``W_IVW``, the inverse variance ``1/sigma^2`` (the +default), and ``W_SD``, the sigma_D Wiener weight ``S/(S + sigma^2)`` built from the +expected difference power -- and ``KSCALE``, the scaler's factor from model to observed +scale. This is the construction ``torchref.validate-ded`` correlates against. Build the +map with ``torchref.mtz2map -csf DF -cw W_IVW -cphi PHDELWT``, adding +``--units electrons`` for e/A^3. It needs no light-state model, so ``-lm`` is optional: + +.. code-block:: bash + + torchref.difference-map \ + -dm dark.pdb \ + -dsf dark.mtz -lsf light.mtz -o results.mtz + +Supplying ``-lm`` (with ``--fraction``) adds the light state's amplitude and +phase and the extrapolated map ``FWT``/``PHWT``: .. code-block:: bash - torchref.phased-difference-map \ + torchref.difference-map \ -dm dark.pdb -lm light.pdb \ -dsf dark.mtz -lsf light.mtz \ --fraction 0.37 -o results.mtz -:API: :mod:`torchref.cli.phased_difference_map` +``FWT``/``PHWT`` keep the standard labels so Coot auto-opens the map, but here they +are the extrapolated light-state map ``2*FEXT - Fc``, not a ``2mFo-DFc``. The file +records this: the columns sit in named MTZ datasets -- ``observed``, ``difference``, +``light_model``, ``extrapolated_light`` and, when written, ``two_moment`` -- so Coot's +column chooser shows ``/torchref/extrapolated_light/FWT``, and ``gemmi mtz`` prints a +history line per dataset. ``torchref.difference-refine`` writes the same file. + +**Key options:** ``--ded-weight {inverse_variance,sigma_d,none}`` selects the +scheme the model-phased and two-moment difference columns carry (default +``inverse_variance``; ``sigma_d`` needs calibrated sigmas, reports how many shells +it found without difference power, and falls back to inverse variance with a warning +when that is every shell); ``--sigma-d-gamma`` fixes the +dark-amplitude exponent of the sigma_D power law instead of fitting it; +``--all-columns`` writes every alternative map coefficient and diagnostic -- the +model-phased difference, the two other extrapolations and the intensity block -- at +the cost of two further scale fits. + +:API: :mod:`torchref.cli.difference_map` Model Utilities --------------- diff --git a/docs/user_guide/scaling.rst b/docs/user_guide/scaling.rst index 14fe6d43..32f6e6cf 100644 --- a/docs/user_guide/scaling.rst +++ b/docs/user_guide/scaling.rst @@ -2,8 +2,8 @@ Scaling ======= :class:`~torchref.scaling.scaler.Scaler` puts F_calc on the observed scale and -absorbs what the atomic model does not describe: an overall (per-resolution-bin) -scale, an anisotropic correction, and the bulk solvent contribution. +absorbs what the atomic model does not describe: an overall isotropic scale, an +anisotropic correction, and the bulk solvent contribution. Basic Usage ----------- @@ -14,7 +14,7 @@ Basic Usage scaler = Scaler(model, reflection_data, verbose=1) - scaler.initialize() # initial bin scales + solvent + anisotropy + scaler.initialize() # initial scale + solvent + anisotropy scaler.refine_lbfgs() # refine the scaling parameters F_calc_scaled = scaler(F_calc) @@ -23,42 +23,70 @@ Basic Usage (``calc_initial_scale`` → ``setup_solvent`` → ``setup_anisotropy_correction``). Each factor defaults to 1 while its parameter is absent, so a freshly constructed scaler is the identity, and one on which only -``calc_initial_scale()`` has run applies the overall bin scale alone. +``calc_initial_scale()`` has run applies the overall isotropic scale alone. -Bin-wise Scaling ----------------- +Isotropic Scaling +----------------- -Reflections are binned by resolution shell and an overall scale is fitted per -bin. On by default with 20 bins: +The overall scale is a Chebyshev polynomial in :math:`s = \sin\theta/\lambda`, +evaluated per reflection: + +.. math:: + + k_{iso}(s) = \exp\left( \sum_{i} c_i\, T_i(u) \right), + \qquad u \in [-1, 1] + +with ``n_iso_coeff`` coefficients (default 6) held in ``scaler.c_iso``. Every +reflection contributes to every coefficient with a continuous weight, so there are +no bin boundaries and nothing changes discontinuously when a reflection moves +between shells. ``n_iso_coeff=1`` is a single global scale +(:math:`T_0 \equiv 1`); ``2`` spans scale-plus-overall-B. .. code-block:: python - scaler = Scaler(model, reflection_data, verbose=1, nbins=1) # single global scale + scaler = Scaler(model, reflection_data, n_iso_coeff=1) # single global scale scaler.calc_initial_scale() +Resolution bins survive only as the device that *seeds* the coefficients: the +closed-form per-bin :math:`|F_{obs}|/|F_{calc}|` ratio is projected onto the basis +by least squares, so the fit starts from the curve a binned model would have +started from. Nothing downstream is binned; ``nbins`` controls only that seed. + +Use ``scaler.iso_log_scale()`` for the per-reflection log scale, and +``scaler.get_scale()`` for a single summary number. + Bulk Solvent Model ------------------ -The solvent contribution is mask-derived, not analytic: a solvent mask is built -from the model, smoothed, and Fourier transformed to give :math:`F_{solvent}`, -which is then Debye-Waller damped and scaled. +The solvent contribution is mask-derived, not analytic: a binary solvent mask is +built from the model and Fourier transformed to give :math:`F_{solvent}`, which is +then damped and scaled. .. math:: - F_{calc}^{total} = k \cdot F_{calc}^{model} - + k_s \exp(-B_s s^2) \cdot F_{calc}^{solvent} + F_{calc}^{total} = k_{iso}(s)\, k_{aniso}(\mathbf{h}) \cdot F_{calc}^{model} + + k_s \exp\!\left( -\ln 2 \left(\frac{s^2}{s^2_{1/2}}\right)^{\!n} \right) + \cdot F_{calc}^{solvent} + +where :math:`k_s` is the solvent scale, :math:`s^2_{1/2}` the point at which the +solvent term is halved, :math:`n` how sharply it switches off, and +:math:`s = \sin\theta/\lambda` — the *half*-length of the scattering vector +(``ScalerBase._s_half_sq``), which is why the exponent carries no factor of 4. -where :math:`k` is the overall (per-bin) scale, :math:`k_s` the solvent scale, -:math:`B_s` the solvent B-factor, and :math:`s = \sin\theta/\lambda` — the -*half*-length of the scattering vector (``ScalerBase._s_half_sq``), which is why -the exponent carries no factor of 4. This is the ordinary Debye-Waller -convention written the other way round: :math:`\exp(-B_s/4d^2)`, i.e. -:math:`\exp(-B_s s^2/4)` for :math:`s = 1/d`. A :math:`B_s` from another program -therefore transfers unchanged (the default is 46 Ų). +:math:`n = 1` reduces this exactly to a Debye-Waller factor +:math:`\exp(-B_s s^2)` with :math:`B_s = \ln 2 / s^2_{1/2}`, so a solvent B from +another program transfers unchanged. Larger :math:`n` gives a plateau followed by a +sharper cutoff, which is the shape a flat bulk-solvent prior actually has: it +describes the data well at low resolution and then stops being informative. -The refined parameters are ``log_k_solvent``, ``b_solvent`` and -``phase_offset``; the last blends the mask phases toward the protein phases and -is only active when ``optimize_phase`` is set. +The refined parameters are ``log_k_solvent``, ``log_ss_half``, ``log_n_exp`` and +``phase_offset``; each falloff parameter is refined in log space so it stays +positive, and is clamped to ``SS_HALF_BOUNDS`` / ``N_EXP_BOUNDS``. The phase offset +blends the mask phases toward the protein phases and is only active when +``optimize_phase`` is set. + +PDB ``REMARK 3`` and mmCIF carry a single solvent B, which this form does not have; +``SolventModel.b_solvent_equivalent`` back-fits one from the curve for deposition. Anisotropic Scaling ------------------- @@ -77,3 +105,47 @@ by a factor of 4 — and :math:`\mathbf{U}` the 6-parameter symmetric tensor stored on the scaler. The :math:`2\pi^2` is part of the definition, not a unit choice — dropping it makes the fitted ``U`` disagree with an ADP-convention ``U`` by that factor. + +Relative scaling of observed datasets +------------------------------------ + +``DatasetCollection.scale()`` jointly fits the observed datasets with +``DatasetScaler``. Each member receives an overall log scale and six quadratic +anisotropy coefficients in a shared normalized HKL basis. Coefficients are +centered over datasets during every forward evaluation, so no reference dataset +fixes the amplitude scale. The metadata reference still supplies the collection +cell and space group. + +The target profiles a shared amplitude per reflection using inverse propagated +variance weights. Both amplitudes and their uncertainties receive the same +positive correction; intensities and their uncertainties receive its square. +The consensus, centering and uncertainty propagation all carry gradients. +The fit uses work reflections shared by at least two datasets, excludes a +reflection held out in any participating dataset, and rejects disconnected or +rank-deficient overlap. Anomalous observations retain their signed identities. + +After calling ``collection.scale()``, retrieve datasets from the collection:: + + collection = DatasetCollection(device="cpu") + collection.add_dataset("dark", dark) + collection.add_dataset("light", light) + collection.scale() + scaled_light = collection["light"] + amplitudes = scaled_light.F + uncertainties = scaled_light.F_sigma + original_amplitudes = scaled_light.F_raw + +The collection owns one scaler. Its ``ScaledDataset`` members subclass +``ReflectionData`` and retain a strong reference to that scaler. Their direct +observation attributes, subset views, stack accessors and exports expose live +corrections. Copies and selections share the scaler; independent collections +have independent scalers. Raw input datasets are copied and remain unchanged. +Adding a dataset invalidates the joint fit; call ``scale()`` again. Repeated +``scale()`` calls otherwise reuse the parameter owner. Fitted parameters are +frozen; use ``collection.scaler.requires_grad_(True)`` for custom optimization. + +``ReflectionData`` holds raw observations and no optimization parameters. +Use ``WilsonNormaliser`` for normalized E values. Work, free and validation +observations are accessed through ``data.work``, ``data.free`` and +``data.validation``; full observations are available directly as ``data.F`` and +``data.F_sigma``. diff --git a/docs/user_guide/targets.rst b/docs/user_guide/targets.rst index a0dcce54..44dbe587 100644 --- a/docs/user_guide/targets.rst +++ b/docs/user_guide/targets.rst @@ -20,10 +20,17 @@ grouped into composite targets for geometry and ADP restraints. X-ray Targets ------------- -Seven modes, selected by name. ``XRAY_TARGETS`` (in +Selected by name. ``XRAY_TARGETS`` (in ``torchref.refinement.targets.xray._specs``) is the single table behind both :func:`~torchref.refinement.targets.create_xray_target` and -``torchref.refine --help``, so the list below cannot drift from the CLI: +``torchref.refine --help``. The authoritative list is ``torchref.refine --help``, +which is generated from that table; the notes below describe the rows but are +maintained by hand, so run ``--help`` if the two disagree. + +Every row shares one forward model — the scaled complex :math:`F_{calc}` — and +declares which measured column it compares against (``spec.observable``). + +**Amplitude rows** compare :math:`F_{obs}` against :math:`|F_{calc}|`: - ``ml`` — **default**. Read MLF: variance :math:`\epsilon\beta`, conditional mean :math:`\alpha|F_{calc}|`, with a cross-validated per-shell Luzzati @@ -40,6 +47,62 @@ Seven modes, selected by name. ``XRAY_TARGETS`` (in - ``ls_wunit_k1`` — Phenix-style least squares: unit weights and a single global scale recomputed every gradient call, bypassing the scaler. +**Intensity rows** compare :math:`I_{obs}` against :math:`|F_{calc}|^2`: + +- ``nll_i`` — Gaussian NLL on the observed intensities, weighted by + :math:`\sigma(I)`. As ``nll``, but skips the French–Wilson conversion. + +Use an intensity row when the signal lives in the *quadratic* part of the data — +a population variance, an activation second moment. :math:`F_{obs}` on a merged +dataset is a French–Wilson posterior rather than a measurement: it is strictly +positive, so it reshapes the weak tail and erases negative intensities, which is +precisely the information such a signal is carried by. + +The axis is deliberately not square. There is no intensity Rice row, because Rice +and the folded normal are distributions *of an amplitude* — the intensity +analogue is the exponential / :math:`\chi^2_1` Wilson distribution, a different +primitive rather than a different variance. R-factors are reported on amplitudes +for every row regardless, so they stay comparable across the whole table. Intensity +rows are not admissible as ``--scale-target``, which fails closed on them. + +Collection X-ray Targets +------------------------ + +The multi-dataset analogues, for time-resolved and difference refinement. Same +taxonomy shape as above — ``COLLECTION_XRAY_TARGETS`` in +``torchref.refinement.targets.collection._specs``, one class per row, the +observable declared per row — and the same ``_loss_inputs`` / ``_per_refl`` seam, +batched over ``(n_datasets, n_hkl)`` on the collection's common HKL grid. + +- ``difference`` — Gaussian on each dataset's **amplitude** difference from the + collection mean, with the dataset/mean covariance propagated. The primary + optimization driver for difference refinement. +- ``difference_i`` — the same on **intensities**. The entire class is one + ``observable`` declaration: the difference-from-mean algebra does not care what + the observable is. +- ``difference_sd`` — the ``difference`` Gaussian centred on + :math:`\alpha\,\Delta F_{calc}` with variance + :math:`\beta_{model} + \sigma_{\Delta}^2`, where :math:`\alpha` and the unexplained + difference power :math:`\beta_{model}` come from a per-shell moment fit of the + observed differences on the free set (``sigma_D``, + :mod:`torchref.refinement.model_error_estimation.sigma_d`); the expected power + carries an :math:`F_{dark}^{\gamma}` dependence with one fitted :math:`\gamma`. A + poor light model inflates the variance where it fails instead of pulling the + coordinates toward noise. Select it with + ``torchref.difference-refine --difference-target difference_sd``. +- ``two_moment`` — merged **intensities** as + :math:`|F(\bar\alpha)|^2 + \sigma_\alpha^2 |\Delta F|^2`, accounting for + crystal-to-crystal spread in activation. +- ``ml`` — Read MLF per dataset at one shared Luzzati :math:`\beta`, fitted on + the pooled free reflections of every data–model pair. The **absolute** channel: + with K free base models a purely relative loss leaves the overall level + unconstrained. + +Both difference rows are offered rather than one being chosen. Amplitudes keep the +loss in the same space as the output DED map coefficients; intensities avoid the +French–Wilson posterior reshaping the weak tail a small difference lives in. Which +wins is a property of a dataset's signal-to-noise. + Geometry Targets ---------------- @@ -143,3 +206,13 @@ evaluations, and tracks what needs recomputing between line-search steps. state = refinement.complete_loss_state() optimizer = LBFGS(refinement.model.parameters(), lr=1.0, max_iter=100) state.run(optimizer, n_steps=1) # equivalent to state.step(optimizer) + +Observed-dataset scaling target +------------------------------- + +``DatasetScalingTarget`` profiles a shared amplitude consensus over the work +observations prepared by ``DatasetScaler``. Each residual is weighted by its +propagated measurement variance, with gradients through log-scale centering, +anisotropy, uncertainty propagation and the consensus. ``DatasetCollection.scale`` +uses this target through ``LossState``; it is independent of structural-model and +model-to-data scaling targets. See :doc:`scaling` for ownership and data access. diff --git a/docs/user_guide/testing.rst b/docs/user_guide/testing.rst index cbdc1c8c..66e60273 100644 --- a/docs/user_guide/testing.rst +++ b/docs/user_guide/testing.rst @@ -39,7 +39,13 @@ PDB ID d_min (Å) Space group ``tests/files/`` also holds partial sets — ``1AK5_with_H.pdb`` + ``1AK5.mtz`` (no CIF), ``7L84.pdb`` + ``7L84-sf.cif`` (no MTZ), ``test_ihm_ensemble.cif`` — so a test that globs one directory and assumes a matching file in another will -fail on those. Use ``sample_structure_pair`` / ``all_test_structures``. +fail on those. Use ``sample_structure_pair`` for the quick reference crystal, +or ``compatibility_structure_pair`` for named extended cases. The latter carries +the ``slow`` marker and selects paths without loading objects. + +``tests/helpers/structure_cases.py`` assigns bundled CIF, MTZ and SF-CIF files to +the quick or extended compatibility panel. Additional files require an explicit +assignment; directory growth does not silently expand numerical test work. Running Tests ------------- @@ -99,11 +105,11 @@ The Amber stack, if you want it: Fixtures -------- -Almost everything lives in the root ``tests/conftest.py`` and is therefore -available from every category — the ``integration/`` and ``functional/`` -conftests are docstrings only. Mock data is the exception: -``tests/unit/conftest.py``. Read those two files for the authoritative list; the -ones you will reach for most: +Reusable setup lives in ``tests/fixtures/``. The root ``tests/conftest.py`` +registers shared plugins and owns test-selection hooks. The unit conftest exposes +synthetic numerical factories; the functional conftest exposes the module-scoped, +read-only Fourier-model fixture. See ``tests/fixtures/README.md`` for ownership +and mutation rules. Common fixtures include: - Paths (session-scoped): ``tests_root``, ``project_root``, ``test_files_dir``, the per-format ``cif_dir``, ``mtz_dir``, ``pdb_dir``, ``cif_sf_dir``, and @@ -118,9 +124,11 @@ ones you will reach for most: ``mock_aniso_u``, ``mock_scattering_factors``, ``mock_weights``. - Real files: ``sample_cif_file``, ``sample_pdb_file``, ``sample_mtz_file``, ``sample_structure_factor_cif``, ``sample_structure_pair`` (matched model + - data), ``all_structure_pairs``, ``all_test_structures``. -- Loaded objects: ``loaded_model``, ``loaded_reflection_data``, - ``model_and_data``, ``initialized_scaler``. + data), ``compatibility_structure_pair`` (one named slow crystal). +- Loaded objects: ``loaded_model``, ``loaded_model_ft``, ``loaded_reflection_data``, + ``model_and_data``, ``initialized_scaler``. ``compatibility_model`` and + ``compatibility_model_and_data`` load only the current slow case and remain + function-scoped to isolate mutations. The mock-data fixtures yield a *factory* taking ``n_atoms`` / ``n_reflections`` and ``seed``; ``mock_cell`` and ``mock_cell_triclinic`` yield the tensor diff --git a/paper/figure3_performance/fcalc_benchmark/benchmark_cpu.py b/paper/figure3_performance/fcalc_benchmark/benchmark_cpu.py index f1098be3..c4ff7afa 100644 --- a/paper/figure3_performance/fcalc_benchmark/benchmark_cpu.py +++ b/paper/figure3_performance/fcalc_benchmark/benchmark_cpu.py @@ -1,14 +1,14 @@ #!/usr/bin/env python -import torch -from torchref import ReflectionData -from torchref import ModelFT +import os from time import time + +import torch from iotbx import pdb +from torchref import ModelFT, ReflectionData -import os _data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data") mtz_file = os.path.join(_data_dir, '1DAW.mtz') pdb_file = os.path.join(_data_dir, '1DAW.pdb') @@ -22,7 +22,7 @@ M = ModelFT(max_res=d_min, device=device,radius_angstrom=4.0).load_pdb(pdb_file) -hkl, _, _, _ = data() +hkl = data.hkl M(hkl, recalc=True) t_start = time() @@ -48,6 +48,4 @@ t_end = time() - print(f"Elapsed time for 10 runs of cctbx calculation: {t_end - t_start} seconds") - diff --git a/paper/figure3_performance/fcalc_benchmark/benchmark_worker.py b/paper/figure3_performance/fcalc_benchmark/benchmark_worker.py index 7e7474d8..caf9b716 100644 --- a/paper/figure3_performance/fcalc_benchmark/benchmark_worker.py +++ b/paper/figure3_performance/fcalc_benchmark/benchmark_worker.py @@ -22,6 +22,7 @@ n_threads = int(os.environ.get("TORCHREF_NUM_THREADS", 1)) import torch + from torchref import ModelFT, ReflectionData @@ -104,7 +105,7 @@ def run_benchmark(n_iterations: int, n_warmup: int, device_str: str = "cpu", data = ReflectionData(device=device).load_mtz(mtz_file) d_min = data.d_min M = ModelFT(max_res=d_min, device=device).load_pdb(pdb_file) - hkl, _, _, _ = data() + hkl = data.hkl n_atoms = M.xyz().shape[0] n_reflections = hkl.shape[0] @@ -116,7 +117,7 @@ def run_benchmark(n_iterations: int, n_warmup: int, device_str: str = "cpu", t.clone().detach().requires_grad_(True) if t is not None else None for t in aniso_ref ) - + def _forward(): sf, _ed = M.fft.compute_structure_factors(hkl, *iso, *aniso) diff --git a/paper/figure3_performance/refinement_cycle_benchmark/benchmark_worker.py b/paper/figure3_performance/refinement_cycle_benchmark/benchmark_worker.py index 1c9dedd3..d34e4f8a 100644 --- a/paper/figure3_performance/refinement_cycle_benchmark/benchmark_worker.py +++ b/paper/figure3_performance/refinement_cycle_benchmark/benchmark_worker.py @@ -25,6 +25,7 @@ n_threads = int(os.environ.get("TORCHREF_NUM_THREADS", 1)) import torch + from torchref.refinement import LBFGSRefinement @@ -118,7 +119,7 @@ def run_benchmark(n_iterations: int, n_warmup: int, device_str: str = "cpu", # Collect metadata n_atoms = len(refinement.model.pdb) - hkl, _, _, _ = refinement.reflection_data() + hkl = refinement.reflection_data.hkl n_reflections = hkl.shape[0] d_min = float(refinement.reflection_data.d_min) target_names = list(loss_state.targets.keys()) diff --git a/paper/figure4_difference_refinement/README.md b/paper/figure4_difference_refinement/README.md index 7aaf0dd4..2ba275ed 100644 --- a/paper/figure4_difference_refinement/README.md +++ b/paper/figure4_difference_refinement/README.md @@ -51,8 +51,9 @@ Key parameters: The refinement script is: `run_difference_refine.sh` -**Stage 2 — Manual real-space refinement** in Coot against the 2Fext-Fc difference -electron density map (from the `difference_data.mtz` output). This step speeds up +**Stage 2 — Manual real-space refinement** in Coot against the extrapolated map +(`FWT`/`PHWT` in the `difference_data.mtz` output; Coot opens it by name). This step +speeds up convergence for the IBL ligand, which requires large conformational changes (*trans* to *cis*). The Coot-refined coordinates were saved as `work.pdb` and used as the starting model for automated refinement (Stage 1). diff --git a/pyproject.toml b/pyproject.toml index 4f0c7ba6..4104a2e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchref" -version = "0.6.4" +version = "0.7.0" description = "Pytorch based crystallographic refinement" readme = "README.md" requires-python = ">=3.10" @@ -45,9 +45,11 @@ dependencies = [ [project.scripts] "torchref.refine" = "torchref.cli.refine:main" "torchref.difference-refine" = "torchref.cli.collection_difference_refine:main" +"torchref.simulate-noisy-data" = "torchref.cli.simulate_noisy_data:main" "torchref.mtz2map" = "torchref.cli.mtz2map:main" "torchref.validate-ded" = "torchref.cli.validate_ded:main" -"torchref.phased-difference-map" = "torchref.cli.phased_difference_map:main" +"torchref.difference-map" = "torchref.cli.difference_map:main" +"torchref.phased-difference-map" = "torchref.cli.difference_map:main" "torchref.add-metadata" = "torchref.cli.add_metadata:main" "torchref.strip-altlocs" = "torchref.cli.strip_altlocs:main" diff --git a/tests/README.md b/tests/README.md index fd283106..3be67e2f 100644 --- a/tests/README.md +++ b/tests/README.md @@ -6,7 +6,8 @@ This directory contains the complete test suite for torchref. ``` tests/ -├── conftest.py # Root fixtures (paths, devices, skip decorators) +├── conftest.py # Fixture registration and test-selection hooks +├── fixtures/ # Shared setup, grouped by responsibility (see fixtures/README.md) ├── pytest.ini # Pytest configuration ├── __init__.py ├── files/ # Test data files (CIF, PDB, MTZ) @@ -15,7 +16,7 @@ tests/ │ ├── mtz/ # Reflection MTZ files │ └── cif_sf/ # Structure factor CIF files ├── unit/ # Unit tests (fast, no I/O) -│ ├── conftest.py # Unit test fixtures (mock data) +│ ├── conftest.py # Imports scoped numerical fixtures │ ├── math_functions/ # Math module tests │ ├── model/ # Model module tests │ ├── refinement/ # Refinement module tests @@ -38,6 +39,34 @@ tests/ ## Running Tests +### Coverage ownership + +| Contract | Owner | +|---|---| +| Loss weights, aggregation, cached loss reads | `unit/refinement/test_loss_state.py` | +| Refinement's default group weights | `unit/refinement/test_loss_weighting.py` | +| Gaussian amplitude-metric values and reductions | `unit/base/test_loss.py` | +| Restraint kernel values on deposited coordinates | `unit/base/test_target_values.py` | +| Gradient RMS norm | `unit/utils/test_gradnorm.py` | +| CIF atomic fields and crystal metadata | `integration/test_io_cif.py` | +| MTZ fields, resolution bins and model/data crystal agreement | `integration/test_io_reflections.py` | +| ModelFT forward cache and grid integration | `functional/test_model_ft_functional.py` | +| Extra deposited files and input inventory | `integration/test_structure_compatibility.py`, `helpers/structure_cases.py` | +| Numerical derivatives and backend parity | `unit/test_gradient_correctness.py`, `unit/structure_factor/` | + +A production call must participate in the assertion: computing a formula only in +the test does not check its implementation. Kernel values, target registration, +device transitions, and default configuration are separate contracts even when +they exercise the same class. Keep mutation tests on fresh objects. + +The quick reader contracts use 1DAW. Extended reader compatibility runs with +`pytest tests/integration/test_structure_compatibility.py --run-slow`; each file +is a separate case and must succeed. The manifest covers the bundled CIF, MTZ +and SF-CIF inputs, including the IHM fixture and reflection-only depositions. +Adding a data file requires an explicit coverage assignment in the manifest. +Extended scaler and restraint cases use 2DQ6 (trigonal) and 3A5V (body-centred +tetragonal), with fresh objects per case and `--run-slow` required. + ### Quick Local Run (on login node, for small tests only) ```bash diff --git a/tests/RUNNING_TESTS.md b/tests/RUNNING_TESTS.md index f2c8c120..109079d7 100644 --- a/tests/RUNNING_TESTS.md +++ b/tests/RUNNING_TESTS.md @@ -92,7 +92,7 @@ pytest tests/unit/refinement/ -v pytest tests/unit/refinement/test_loss_weighting.py -v # Target/loss functions -pytest tests/unit/refinement/test_targets.py -v +pytest tests/unit/base/test_target_values.py tests/unit/base/test_loss.py -v ``` ### Scaling @@ -176,17 +176,17 @@ pytest tests/unit/model/test_parameter_wrappers.py::TestMixedTensorOperations -v #### Refinement Classes ```bash -# Fixed weighting -pytest tests/unit/refinement/test_loss_weighting.py::TestFixedWeighting -v +# Weight handling +pytest tests/unit/refinement/test_loss_state.py::TestWeightManagement -v -# Resolution-dependent weighting -pytest tests/unit/refinement/test_loss_weighting.py::TestResolutionDependentWeighting -v +# Default group weights +pytest tests/unit/refinement/test_loss_weighting.py::TestDefaultGroupWeights -v # Gaussian NLL loss -pytest tests/unit/refinement/test_targets.py::TestGaussianNLL -v +pytest tests/unit/base/test_loss.py -v # Least squares target -pytest tests/unit/refinement/test_targets.py::TestLeastSquaresTarget -v +pytest tests/unit/base/test_target_values.py -k least_squares -v ``` #### Symmetry Classes @@ -361,13 +361,14 @@ pytest tests/unit --lf -v | `math_functions/test_math_numpy.py` | `TestCoordinateTransformations`, `TestScatteringVectors`, `TestRFactorCalculations`, `TestRotation` | | `model/test_model.py` | `TestModelInitialization`, `TestModelDeviceHandling` | | `model/test_parameter_wrappers.py` | `TestMixedTensorInitialization`, `TestMixedTensorOperations`, `TestMixedTensorDeviceHandling`, `TestOccupancyTensor`, `TestPositiveMixedTensor` | -| `refinement/test_loss_weighting.py` | `TestFixedWeighting`, `TestResolutionDependentWeighting`, `TestLossWeightingModule` | -| `refinement/test_targets.py` | `TestTargetBase`, `TestGaussianNLL`, `TestLeastSquaresTarget`, `TestRiceNLL`, `TestTargetDeviceHandling`, `TestNumericStability` | +| `refinement/test_loss_weighting.py` | `TestDefaultGroupWeights` | +| `base/test_target_values.py` | Deposited-coordinate restraint values and least-squares weighting | +| `base/test_loss.py` | Gaussian NLL values and reductions | | `scaling/test_scaler.py` | `TestScalerInitialization`, `TestScalerDeviceHandling`, `TestScalingCalculations`, `TestBFactorScaling`, `TestAnisotropicScaling` | | `symmetrie/test_symmetrie.py` | `TestSymmetryInitialization`, `TestSymmetryMatrices`, `TestSymmetryApplication`, `TestSymmetryDeviceHandling`, `TestSpaceGroupMapping` | | `io/test_data.py` | `TestReflectionDataInitialization`, `TestReflectionDataDeviceMovement`, `TestReflectionDataAttributes`, `TestReflectionDataProperties`, `TestMockReflectionData` | | `restraints/test_restraints.py` | `TestRestraintsInitialization`, `TestBondRestraintCalculations`, `TestAngleRestraintCalculations`, `TestTorsionRestraintCalculations`, `TestRestraintDeviceHandling`, `TestRestraintNumericStability` | -| `utils/test_gradnorm.py` | `TestGradNorm` | +| `utils/test_gradnorm.py` | RMS norms for single/multiple parameters and zero gradients | | `utils/test_utils.py` | `TestModuleReference`, `TestCIFReader` | ### Integration Tests (`tests/integration/`) diff --git a/tests/benchmarks/compare_amber_classic.py b/tests/benchmarks/compare_amber_classic.py new file mode 100644 index 00000000..fd678d57 --- /dev/null +++ b/tests/benchmarks/compare_amber_classic.py @@ -0,0 +1,384 @@ +"""Compare classic and AMBER restraints on identical prepared benchmark starts. + +Run with the project interpreter and optional OpenMM/PDBFixer dependencies. +The experiment keeps the work/free split, X-ray target, atom set, riding wrapper, +optimizer and ADP settings fixed. AMBER's weight is calibrated once from initial +heavy-coordinate gradient RMS after the riding Jacobian, without consulting R-free. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from pathlib import Path + +import numpy as np +import torch + +from torchref import Model +from torchref.experimental.targets.amber_target import AmberTarget +from torchref.refinement.lbfgs_refinement import LBFGSRefinement + +FILES = Path(__file__).resolve().parents[1] / "files" +IDENTITY = ["chainid", "resseq", "icode", "name"] + + +def _prepare(code: str, output: Path, seed: int) -> dict: + """Complete heavy atoms once, preserve existing rows, and add H in TorchRef.""" + import openmm.app as app + from pdbfixer import PDBFixer + + source = FILES / "pdb" / f"{code}_af.pdb" + original = ( + Model(device="cpu", verbose=0, strip_H=True, add_hydrogens=False) + .load_pdb(str(source)) + .strip_altlocs() + ) + fixer = PDBFixer(filename=str(source)) + fixer.findMissingResidues() + fixer.missingResidues = {} + fixer.findMissingAtoms() + fixer.addMissingAtoms(seed=seed) + fixed_path = output / f"{code}_heavy_completed.pdb" + with fixed_path.open("w") as handle: + app.PDBFile.writeFile(fixer.topology, fixer.positions, handle, keepIds=True) + fixed = Model(device="cpu", verbose=0, add_hydrogens=False).load_pdb( + str(fixed_path) + ) + original_rows = { + tuple(row[k] for k in IDENTITY): row for _, row in original.pdb.iterrows() + } + frame = fixed.pdb.copy() + frame.attrs = original.pdb.attrs.copy() + added = [] + found = set() + for i, row in frame.iterrows(): + key = tuple(row[k] for k in IDENTITY) + if key in original_rows: + found.add(key) + for column in ["x", "y", "z", "tempfactor", "occupancy", "ATOM"]: + frame.at[i, column] = original_rows[key][column] + else: + added.append(key) + same_residue = original.pdb[ + (original.pdb.chainid == row.chainid) + & (original.pdb.resseq == row.resseq) + & (original.pdb.icode == row.icode) + ] + frame.at[i, "tempfactor"] = float(same_residue.tempfactor.mean()) + frame.at[i, "occupancy"] = 1.0 + assert found == set(original_rows), "Heavy-atom preparation dropped original atoms" + model = original._new_model_from_df(frame, strip_H=False) + torch.manual_seed(seed) + model = model.hydrogenate() + model.set_hydrogen_mode("riding") + compatibility = AmberTarget(model=model) + assert compatibility._n_omm_atoms == len(model.pdb) + path = output / f"{code}_prepared.pdb" + model.write_pdb(str(path)) + return { + "code": code, + "source": str(source), + "prepared": str(path), + "added_heavy_atoms": added, + "original_heavy_atoms": len(original.pdb), + "prepared_atoms": len(model.pdb), + "hydrogens": int(model.pdb.element.str.strip().isin(["H", "D"]).sum()), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + + +def _magnitude(values: torch.Tensor) -> dict: + """Summarize atom-vector norms, or absolute scalar parameter gradients.""" + values = values.detach().cpu() + norms = values.norm(dim=-1) if values.ndim == 2 else values.abs().flatten() + if not norms.numel(): + return {"n": 0} + return { + "n": norms.numel(), + "rms": float(norms.square().mean().sqrt()), + "median": float(norms.median()), + "p95": float(torch.quantile(norms, 0.95)), + "max": float(norms.max()), + "l2": float(values.norm()), + } + + +def _classic_loss(ref: LBFGSRefinement) -> torch.Tensor: + """Sum the active classic components, excluding the disabled Rama prior.""" + return sum( + target() + for name, target in ref.geometry_target.items() + if name != "ramachandran" + ) + + +def _gradient_snapshot(ref: LBFGSRefinement, amber: AmberTarget) -> tuple[dict, dict]: + """Record raw atomic and model-parameter gradients for both priors and X-ray.""" + # The classic prior is inactive during AMBER refinement, so its contact list + # needs maintenance before evaluating it on the final AMBER coordinates. + for target in ref.geometry_target.values(): + target.maintenance() + model = ref.model + heavy = torch.as_tensor( + ~model.pdb.element.str.strip().isin(["H", "D"]).to_numpy(), device=model.device + ) + tensors = {} + result = {} + for name, function in [ + ("classic_raw", lambda: _classic_loss(ref)), + ("amber_raw", amber.forward), + ("xray", ref.xray_target_work.forward), + ]: + model.reset_cache() + xyz = model.xyz() + leaves = model.xyz.optimization_parameters() + value = function() + gradients = torch.autograd.grad(value, [xyz] + leaves, allow_unused=True) + atomic = gradients[0] + assert atomic is not None and torch.isfinite(atomic).all(), name + tensors[name] = atomic.detach() + result[name] = { + "loss": float(value.detach()), + "heavy": _magnitude(atomic[heavy]), + "hydrogen": _magnitude(atomic[~heavy]), + "all": _magnitude(atomic), + "parameter_gradients": [ + _magnitude(g) if g is not None else {"n": 0} for g in gradients[1:] + ], + } + if name == "amber_raw": + import openmm.unit as unit + + forces = np.asarray( + amber._context.getState(getForces=True) + .getForces(asNumpy=True) + .value_in_unit(unit.kilojoules_per_mole / unit.nanometer) + ) + result[name]["clipped_atom_fraction"] = float( + (np.linalg.norm(forces, axis=1) > 10000).mean() + ) + unclipped = torch.as_tensor( + -forces[amber._model_to_omm] * 0.1 / len(xyz), + dtype=xyz.dtype, + device=xyz.device, + ) + result["amber_unclipped"] = { + "heavy": _magnitude(unclipped[heavy]), + "hydrogen": _magnitude(unclipped[~heavy]), + "all": _magnitude(unclipped), + } + for name in ["classic_raw", "amber_raw"]: + g = tensors[name][heavy].flatten() + x = tensors["xray"][heavy].flatten() + result[name]["heavy_cosine_to_xray"] = float( + torch.nn.functional.cosine_similarity(g, x, dim=0) + ) + c, a = ( + tensors["classic_raw"][heavy].flatten(), + tensors["amber_raw"][heavy].flatten(), + ) + result["classic_amber_heavy_cosine"] = float( + torch.nn.functional.cosine_similarity(c, a, dim=0) + ) + return result, tensors + + +def _metrics(ref: LBFGSRefinement) -> dict: + """Report R factors and dictionary geometry using heavy atoms alone.""" + with torch.no_grad(): + rwork, rfree = ref.get_rfactor() + model = ref.model + restraints = model.restraints + restraints.cat_dict() + heavy = torch.as_tensor( + ~model.pdb.element.str.strip().isin(["H", "D"]).to_numpy(), + device=model.device, + ) + result = {"rwork": float(rwork), "rfree": float(rfree)} + for name, method in [ + ("bond", restraints.bond_deviations), + ("angle", restraints.angle_deviations), + ]: + deviations, sigmas = method() + indices = restraints.restraints[name]["all"]["indices"] + keep = heavy[indices].all(dim=-1) + d, z = deviations[keep], deviations[keep] / sigmas[keep] + if name == "angle": + d = torch.rad2deg(d) + result[f"heavy_{name}_rms_delta"] = float(d.square().mean().sqrt()) + result[f"heavy_{name}_rms_z"] = float(z.square().mean().sqrt()) + return result + + +def _run( + code: str, + prepared: Path, + arm: str, + output: Path, + cycles: int, + seed: int, + weight: float | None, +) -> dict: + """Run classic alternating scaler/XYZ/ADP refinement with one active prior.""" + torch.manual_seed(seed) + started = time.perf_counter() + ref = LBFGSRefinement( + pdb=str(prepared), + data_file=str(FILES / "mtz" / f"{code}.mtz"), + device=torch.device("cpu"), + verbose=0, + add_hydrogens=False, + hydrogens_in_xray=True, + target_mode="ml", + ) + ref.set_hydrogen_mode("riding") + amber = AmberTarget(model=ref.model, normalize_by_atoms=True, verbose=0) + state = ref.loss_state + initial_gradients, raw = _gradient_snapshot(ref, amber) + calibration = ( + 0.2 + * initial_gradients["classic_raw"]["parameter_gradients"][0]["rms"] + / initial_gradients["amber_raw"]["parameter_gradients"][0]["rms"] + ) + if weight is None: + weight = calibration + if arm == "amber": + state.targets = { + key: target + for key, target in state.targets.items() + if not key.startswith("geometry/") + } + state.clear() + state.register_target("amber", amber) + state.set_weight("amber", weight) + state.refresh_loss_leaves() + active_names = [ + key for key in state.targets if state.get_effective_weight(key) != 0 + ] + if arm == "amber": + assert "amber" in active_names and not any( + key.startswith("geometry/") for key in active_names + ) + else: + assert "amber" not in active_names + initial = _metrics(ref) + xyz_hash = hashlib.sha256( + ref.model.xyz().detach().cpu().numpy().tobytes() + ).hexdigest() + result = { + "code": code, + "arm": arm, + "cycles": cycles, + "seed": seed, + "amber_weight": weight, + "calibration_weight": calibration, + "calibration_basis": "heavy_xyz_parameter_gradient_rms_after_riding", + "cartesian_calibration_weight": ( + 0.2 + * initial_gradients["classic_raw"]["heavy"]["rms"] + / initial_gradients["amber_raw"]["heavy"]["rms"] + ), + "classic_group_weight": 0.2, + "initial": initial, + "initial_gradients": initial_gradients, + "initial_xyz_sha256": xyz_hash, + "active_targets": active_names, + "n_atoms": len(ref.model.pdb), + "n_reflections": len(ref.reflection_data.hkl), + "reflection_split_sha256": hashlib.sha256( + ref.reflection_data.hkl.detach().cpu().numpy().tobytes() + + ref.reflection_data.rfree_flags.detach().cpu().numpy().tobytes() + ).hexdigest(), + "parameter_gradient_order": [ + "heavy_xyz_per_angstrom", + "torsion_per_radian", + "water_rotation_per_radian", + ], + "geometry_units": {"bond_rms_delta": "angstrom", "angle_rms_delta": "degree"}, + "setup_seconds": time.perf_counter() - started, + "trajectory": [], + } + target_path = output / f"{code}_{arm}.json" + target_path.write_text(json.dumps(result, indent=2)) + print( + json.dumps( + { + "event": "initial", + "code": code, + "arm": arm, + "metrics": initial, + "weight": weight, + "grad_rms_classic": initial_gradients["classic_raw"]["heavy"]["rms"], + "grad_rms_amber": initial_gradients["amber_raw"]["heavy"]["rms"], + } + ), + flush=True, + ) + for cycle in range(1, cycles + 1): + start_cycle = time.perf_counter() + ref.refine(macro_cycles=1) + entry = { + "cycle": cycle, + **_metrics(ref), + "seconds": time.perf_counter() - start_cycle, + } + result["trajectory"].append(entry) + target_path.write_text(json.dumps(result, indent=2)) + print( + json.dumps({"event": "cycle", "code": code, "arm": arm, **entry}), + flush=True, + ) + result["final"] = _metrics(ref) + result["final_gradients"], _ = _gradient_snapshot(ref, amber) + result["refinement_seconds"] = sum(x["seconds"] for x in result["trajectory"]) + result["total_seconds"] = time.perf_counter() - started + ref.model.write_pdb(str(output / f"{code}_{arm}_refined.pdb")) + target_path.write_text(json.dumps(result, indent=2)) + return result + + +def _main() -> None: + """Prepare one benchmark or run one arm in its own process.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("action", choices=["prepare", "classic", "amber"]) + parser.add_argument("--code", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--cycles", type=int, default=5) + parser.add_argument("--amber-weight", type=float, default=None) + parser.add_argument("--seed", type=int, default=20260917) + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=True) + if args.action == "prepare": + details = _prepare(args.code, args.output, args.seed) + (args.output / f"{args.code}_preparation.json").write_text( + json.dumps(details, indent=2) + ) + print(json.dumps(details), flush=True) + else: + weight = args.amber_weight + if args.action == "amber" and weight is None: + baseline = json.loads( + (args.output / f"{args.code}_classic.json").read_text() + ) + gradients = baseline["initial_gradients"] + weight = ( + 0.2 + * gradients["classic_raw"]["parameter_gradients"][0]["rms"] + / gradients["amber_raw"]["parameter_gradients"][0]["rms"] + ) + _run( + args.code, + args.output / f"{args.code}_prepared.pdb", + args.action, + args.output, + args.cycles, + args.seed, + weight, + ) + + +if __name__ == "__main__": + _main() diff --git a/tests/conftest.py b/tests/conftest.py index 84715a84..380df84f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,35 +1,27 @@ -""" -Root pytest configuration and shared fixtures for torchref tests. +"""Register shared fixture plugins and gate tests on host capabilities.""" -This module provides fixtures that are automatically available to all test files. -""" import importlib.util import shutil import warnings import pytest -import torchref -import torch -import numpy as np -from pathlib import Path +# Apply process-wide settings before the device fixtures import torch. +import torchref # noqa: F401 + +pytest_plugins = ( + "tests.fixtures.paths", + "tests.fixtures.files", + "tests.fixtures.devices", + "tests.fixtures.precision", + "tests.fixtures.objects", + "tests.fixtures.collections", +) -# Optional Amber/ensemble stack: OpenMM (pip ``[amber]`` extra) and AmberTools -# (antechamber/tleap — conda-only, detected on PATH). Tests that need them are -# tagged ``@pytest.mark.openmm`` (OpenMM only) or ``@pytest.mark.amber`` (OpenMM -# + AmberTools) and auto-skipped below when the stack is absent. _HAS_OPENMM = importlib.util.find_spec("openmm") is not None _HAS_AMBERTOOLS = bool(shutil.which("antechamber") and shutil.which("tleap")) -def _cuda_available() -> bool: - return torch.cuda.is_available() - - -def _mps_available() -> bool: - return hasattr(torch.backends, "mps") and torch.backends.mps.is_available() - - def pytest_addoption(parser): """Add custom command line options.""" parser.addoption( @@ -58,24 +50,34 @@ def pytest_addoption(parser): help="Deprecated no-op: accelerator tests now run automatically.", ) parser.addoption( - "--run-slow", - action="store_true", - default=False, - help="Run slow tests" + "--run-slow", action="store_true", default=False, help="Run slow tests" ) def pytest_configure(config): """Configure pytest markers.""" config.addinivalue_line("markers", "unit: Unit tests (fast, no I/O)") - config.addinivalue_line("markers", "integration: Integration tests (slower, real I/O)") - config.addinivalue_line("markers", "gpu: Needs any accelerator (CUDA or MPS); auto-skipped if none") - config.addinivalue_line("markers", "cuda: Needs CUDA specifically (e.g. Triton); auto-skipped if absent") - config.addinivalue_line("markers", "mps: Needs MPS specifically (Metal kernels); auto-skipped if absent") + config.addinivalue_line( + "markers", "integration: Integration tests (slower, real I/O)" + ) + config.addinivalue_line( + "markers", "gpu: Needs any accelerator (CUDA or MPS); auto-skipped if none" + ) + config.addinivalue_line( + "markers", "cuda: Needs CUDA specifically (e.g. Triton); auto-skipped if absent" + ) + config.addinivalue_line( + "markers", "mps: Needs MPS specifically (Metal kernels); auto-skipped if absent" + ) config.addinivalue_line("markers", "cuda_only: Deprecated alias for 'cuda'") config.addinivalue_line("markers", "slow: Slow tests (skipped by default)") - config.addinivalue_line("markers", "openmm: Needs OpenMM (the [amber] extra); skipped if absent") - config.addinivalue_line("markers", "amber: Needs OpenMM + AmberTools (antechamber/tleap); skipped if absent") + config.addinivalue_line( + "markers", "openmm: Needs OpenMM (the [amber] extra); skipped if absent" + ) + config.addinivalue_line( + "markers", + "amber: Needs OpenMM + AmberTools (antechamber/tleap); skipped if absent", + ) if config.getoption("--run-gpu"): # UserWarning, not DeprecationWarning: pytest.ini filters the latter, @@ -109,6 +111,8 @@ def pytest_collection_modifyitems(config, items): mask a forgotten marker, and turns "this host cannot run it" into a silent pass instead of the visible skip or the real error. """ + from tests.fixtures.devices import _cuda_available, _mps_available + has_cuda = _cuda_available() has_mps = _mps_available() @@ -138,7 +142,9 @@ def pytest_collection_modifyitems(config, items): ) skip_slow = pytest.mark.skip(reason="Need --run-slow option to run") - skip_openmm = pytest.mark.skip(reason="OpenMM not installed (pip install '.[amber]')") + skip_openmm = pytest.mark.skip( + reason="OpenMM not installed (pip install '.[amber]')" + ) skip_amber = pytest.mark.skip( reason="AmberTools (antechamber/tleap) not on PATH (conda install ambertools)" ) @@ -174,466 +180,3 @@ def pytest_collection_modifyitems(config, items): item.add_marker(skip_amber) elif "openmm" in item.keywords and not _HAS_OPENMM: item.add_marker(skip_openmm) - - -# ============================================================================= -# Path Fixtures -# ============================================================================= - -@pytest.fixture(scope="session") -def tests_root() -> Path: - """Root of the tests directory.""" - return Path(__file__).parent - - -@pytest.fixture(scope="session") -def project_root() -> Path: - """Root of the project.""" - return Path(__file__).parent.parent - - -@pytest.fixture(scope="session") -def test_files_dir(tests_root) -> Path: - """Path to test files directory.""" - return tests_root / "files" - - -@pytest.fixture(scope="session") -def cif_dir(test_files_dir) -> Path: - """Path to CIF model files.""" - return test_files_dir / "cif" - - -@pytest.fixture(scope="session") -def cif_sf_dir(test_files_dir) -> Path: - """Path to CIF structure factor files.""" - return test_files_dir / "cif_sf" - - -@pytest.fixture(scope="session") -def mtz_dir(test_files_dir) -> Path: - """Path to MTZ reflection files.""" - return test_files_dir / "mtz" - - -@pytest.fixture(scope="session") -def pdb_dir(test_files_dir) -> Path: - """Path to PDB model files.""" - return test_files_dir / "pdb" - - -@pytest.fixture(scope="session") -def external_monomer_library(project_root) -> Path: - """Path to external monomer library.""" - return project_root / "external_monomer_library" - - -# ============================================================================= -# Device Fixtures -# ============================================================================= - -@pytest.fixture(scope="session") -def cpu_device() -> torch.device: - """CPU torch device.""" - return torch.device("cpu") - - -@pytest.fixture(scope="session") -def gpu_device() -> torch.device: - """GPU torch device (only use with @pytest.mark.gpu). - - Prefers CUDA, falls back to MPS; skips if neither is available. Prefer the - backend-specific ``cuda_device`` / ``mps_device`` below when a test needs one - particular backend -- this fixture's preference order means a - ``cuda``-marked test asking for it on a dual-backend host could be handed - MPS, which is why the MPS tests used to carry a ``type != 'mps'`` skip to - undo it. - """ - accel = _accelerator() - if accel is None: - pytest.skip("No accelerator (CUDA or MPS) on this host") - return accel - - -@pytest.fixture(scope="session") -def cuda_device() -> torch.device: - """Canonical CUDA device for ``cuda``-marked tests. - - Deliberately unguarded. What runs is decided by the ``cuda`` marker in - :func:`pytest_collection_modifyitems` and nowhere else, so this fixture does - not re-check availability: on a host without CUDA the test is *meant* to - error with the real backend error rather than be quietly skipped here. - """ - return torch.device("cuda", 0) - - -@pytest.fixture(scope="session") -def mps_device() -> torch.device: - """Canonical MPS device for ``mps``-marked tests. - - Unguarded for the same reason as :func:`cuda_device` -- the ``mps`` marker - owns the decision. - """ - return torch.device("mps", 0) - - -def _accelerator() -> "torch.device | None": - """The canonical accelerator this host can actually use, or ``None``. - - Indices are filled in (``cuda:0`` / ``mps:0``) so the value compares equal - to a device read back off a real tensor -- ``torch.device('mps')`` and - ``torch.device('mps:0')`` are *not* equal even though they name the same - physical device. - """ - if _cuda_available(): - return torch.device("cuda", torch.cuda.current_device()) - if _mps_available(): - return torch.device("mps", 0) - return None - - -# Built at import time so the ``gpu`` mark is attached during *collection*. -# Adding it later (e.g. via ``request.node.add_marker`` inside the fixture) is -# too late for ``pytest_collection_modifyitems`` to gate on. -_DEVICE_PARAMS = [pytest.param(torch.device("cpu"), id="cpu")] -_ACCELERATOR = _accelerator() -if _ACCELERATOR is not None: - _DEVICE_PARAMS.append( - pytest.param( - _ACCELERATOR, - id=_ACCELERATOR.type, - # Backend-specific mark, so a CUDA-less host skips the cuda leg and - # a non-Mac skips the mps leg, each with an accurate reason. - marks=getattr(pytest.mark, _ACCELERATOR.type), - ) - ) - - -@pytest.fixture(params=_DEVICE_PARAMS) -def any_device(request) -> torch.device: - """Every device this host can actually use, one test run per device. - - The CPU leg always runs. The accelerator leg is ``gpu``-marked, so a plain - ``pytest`` run skips it and ``pytest --run-gpu`` picks up CUDA on a CUDA - box or MPS on a Mac. On a CPU-only host the accelerator parameter does not - exist at all, so there is no skip noise. - """ - return request.param - - -@pytest.fixture(scope="session") -def _device_model_cache() -> dict: - """``{device_str: ModelFT}`` built at most once per device, per session.""" - return {} - - -@pytest.fixture -def device_model_bundle(_device_model_cache, pdb_dir, any_device): - """A loaded model on ``any_device``, for target conformance tests. - - The existing ``loaded_model`` / ``model_and_data`` fixtures are - function-scoped and construct on the process default, so a - device-parametrized sweep over them would reload the structure once per - test per device. This caches one model per device instead. - - Shared mutable state: callers must treat the bundle as read-only. A test - that moves a *target* will drag the borrowed model with it, poisoning every - later test on that device -- see ``test_target_device_round_trip``, which - deliberately builds its own. - """ - key = str(any_device) - if key not in _device_model_cache: - pdb = pdb_dir / "1DAW.pdb" - if not pdb.exists(): - pytest.skip("1DAW.pdb fixture not present") - from torchref.model import ModelFT - - _device_model_cache[key] = ModelFT(device=any_device, verbose=0).load_pdb( - str(pdb) - ) - return {"model": _device_model_cache[key]} - - -@pytest.fixture -def device(request) -> torch.device: - """Default test device. - - Uses the package-wide auto-detected default (``torchref.device.current``) - so tests run on whichever device the user's machine resolved to at - import time: cuda -> mps -> cpu. Tests marked ``@pytest.mark.cuda_only`` - are skipped when CUDA is not available. - """ - from torchref.config import get_default_device - - markers = {m.name for m in request.node.iter_markers()} - if "cuda_only" in markers and not torch.cuda.is_available(): - pytest.skip("Test requires CUDA") - if "gpu" in markers and not (_cuda_available() or _mps_available()): - pytest.skip("No GPU (CUDA or MPS) available") - return get_default_device() - - -# ============================================================================= -# Numerical Fixtures -# ============================================================================= - -@pytest.fixture -def rtol() -> float: - """Relative tolerance for floating point comparisons.""" - return 1e-5 - - -@pytest.fixture -def atol() -> float: - """Absolute tolerance for floating point comparisons.""" - return 1e-8 - - -# ============================================================================= -# Sample File Fixtures -# ============================================================================= - -@pytest.fixture(scope="session") -def sample_cif_file(cif_dir): - """Return a sample CIF file for testing.""" - cif_file = cif_dir / "1DAW.cif" - if cif_file.exists(): - return cif_file - # Try any available CIF file - cif_files = list(cif_dir.glob("*.cif")) - if cif_files: - return cif_files[0] - pytest.skip("No CIF files found in test data") - - -@pytest.fixture(scope="session") -def sample_mtz_file(mtz_dir): - """Return a sample MTZ file for testing.""" - mtz_file = mtz_dir / "1DAW.mtz" - if mtz_file.exists(): - return mtz_file - # Try any available MTZ file - mtz_files = list(mtz_dir.glob("*.mtz")) - if mtz_files: - return mtz_files[0] - pytest.skip("No MTZ files found in test data") - - -@pytest.fixture(scope="session") -def sample_pdb_file(pdb_dir): - """Return a sample PDB file for testing.""" - pdb_files = sorted(pdb_dir.glob("*.pdb")) - if not pdb_files: - pytest.skip("No PDB files found in test data directory") - return pdb_files[0] - - -@pytest.fixture(scope="session") -def sample_structure_factor_cif(cif_sf_dir): - """Return a sample structure factor CIF file.""" - sf_files = sorted(cif_sf_dir.glob("*.cif")) - if not sf_files: - pytest.skip("No structure factor CIF files found") - return sf_files[0] - - -@pytest.fixture(scope="session") -def sample_structure_pair(cif_dir, mtz_dir): - """Return a matching pair of CIF model and MTZ reflections.""" - # Try to find matching files - pdb_id = "1DAW" - cif_file = cif_dir / f"{pdb_id}.cif" - mtz_file = mtz_dir / f"{pdb_id}.mtz" - - if cif_file.exists() and mtz_file.exists(): - return {"model": cif_file, "reflections": mtz_file} - - # Try to find any matching pair - cif_files = {f.stem: f for f in cif_dir.glob("*.cif")} - mtz_files = {f.stem: f for f in mtz_dir.glob("*.mtz")} - - common_ids = set(cif_files.keys()) & set(mtz_files.keys()) - if common_ids: - pdb_id = sorted(common_ids)[0] - return {"model": cif_files[pdb_id], "reflections": mtz_files[pdb_id]} - - pytest.skip("No matching CIF/MTZ pairs found in test data") - - -@pytest.fixture(scope="session") -def all_structure_pairs(cif_dir, mtz_dir): - """Return all matching pairs of CIF models and MTZ reflections.""" - cif_files = {f.stem: f for f in cif_dir.glob("*.cif")} - mtz_files = {f.stem: f for f in mtz_dir.glob("*.mtz")} - - common_ids = set(cif_files.keys()) & set(mtz_files.keys()) - - if not common_ids: - pytest.skip("No matching CIF/MTZ pairs found in test data") - - return [ - {"pdb_id": pdb_id, "model": cif_files[pdb_id], "reflections": mtz_files[pdb_id]} - for pdb_id in sorted(common_ids) - ] - - -@pytest.fixture(scope="session") -def all_cif_files(cif_dir): - """Return all available CIF test structure files.""" - cif_files = sorted(cif_dir.glob("*.cif")) - if not cif_files: - pytest.skip("No CIF files found in test data directory") - return cif_files - - -@pytest.fixture(scope="session") -def all_test_structures(all_structure_pairs): - """Return all loaded model/data pairs for comprehensive testing.""" - from torchref.model.model import Model - from torchref.io import ReflectionData - - structures = [] - for pair in all_structure_pairs: - try: - model = Model() - model.load_cif(str(pair["model"])) - - data = ReflectionData() - data.load_mtz(str(pair["reflections"])) - - structures.append({ - "pdb_id": pair["pdb_id"], - "model": model, - "data": data, - "model_path": pair["model"], - "data_path": pair["reflections"] - }) - except Exception: - # Skip structures that fail to load - continue - - if not structures: - pytest.skip("No structures could be loaded") - - return structures - - -@pytest.fixture(scope="session") -def monomer_library_path(project_root): - """Get path to the monomer library as a string. - - Returns - ------- - str - Absolute path to the external_monomer_library directory. - """ - lib_path = project_root / "external_monomer_library" - if not lib_path.exists(): - pytest.skip("Monomer library not found") - return str(lib_path) - - -# ============================================================================= -# Real Object Fixtures -# ============================================================================= - -@pytest.fixture -def loaded_model(sample_cif_file): - """Fixture providing a fully loaded Model from a real CIF file.""" - from torchref.model.model import Model - - model = Model() - model.load_cif(str(sample_cif_file)) - return model - - -@pytest.fixture -def loaded_reflection_data(sample_mtz_file): - """Fixture providing fully loaded ReflectionData from a real MTZ file.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - return data - - -@pytest.fixture -def model_and_data(sample_structure_pair): - """Fixture providing matching model and reflection data.""" - from torchref.model.model import Model - from torchref.io import ReflectionData - - model = Model() - model.load_cif(str(sample_structure_pair["model"])) - - data = ReflectionData() - data.load_mtz(str(sample_structure_pair["reflections"])) - - return {"model": model, "data": data} - - -@pytest.fixture -def model_with_symmetry(loaded_model): - """Fixture providing model with initialized symmetry.""" - from torchref.symmetry import SpaceGroup - - sg = SpaceGroup(loaded_model.spacegroup) - return {"model": loaded_model, "symmetry": sg} - - -@pytest.fixture -def initialized_scaler(model_and_data): - """Fixture providing initialized Scaler with model and data.""" - from torchref.scaling.scaler import Scaler - - model = model_and_data["model"] - data = model_and_data["data"] - - scaler = Scaler(model=model, data=data, nbins=10, verbose=0) - return scaler - - -@pytest.fixture -def model_with_restraints(loaded_model): - """Fixture providing model with built restraints.""" - from torchref.restraints import Restraints - - restraints = Restraints( - pdb=loaded_model.pdb, - xyz_fn=loaded_model.xyz, - vdw_radii_fn=loaded_model.get_vdw_radii, - verbose=0 - ) - restraints.build_restraints() - return {"model": loaded_model, "restraints": restraints} - -@pytest.fixture -def double_cpu(): - """float64/complex128 on CPU for the duration of a test; restore afterwards. - - Required rather than cosmetic for anything touching eager structure factors: - ``iso_structure_factor_torched`` casts ``hkl`` to the *global* ``dtypes.float`` - (``torchref/base/direct_summation/isotropic.py:121``), so under the default float32 - config a float64 leaf produces a dtype-mismatched matmul. - - Promoted here from three byte-similar copies in ``tests/unit/test_kernel_fixes.py``, - ``tests/unit/test_gradient_correctness.py`` and - ``tests/integration/test_dtype_config_float64.py``. This version also restores - ``sigma_cutoff_ed``, which none of those did -- so a test that changed the cutoff - leaked it into everything that ran afterwards. - """ - import torchref - from torchref.config import device as _device, dtypes as _dtypes - - f0, c0, d0 = _dtypes.float, _dtypes.complex, _device.current - s0 = torchref.sigma_cutoff_ed.value - _dtypes.float = torch.float64 - _dtypes.complex = torch.complex128 - _device.current = torch.device("cpu") - try: - yield - finally: - _dtypes.float = f0 - _dtypes.complex = c0 - _device.current = d0 - torchref.sigma_cutoff_ed.value = s0 diff --git a/tests/files/hkl/dark_half1.hkl b/tests/files/hkl/dark_half1.hkl new file mode 100644 index 00000000..7312b58c --- /dev/null +++ b/tests/files/hkl/dark_half1.hkl @@ -0,0 +1,69 @@ +CrystFEL reflection list version 2.0 +Symmetry: 1 + h k l I phase sigma(I) nmeas + -17 -13 -5 -10.07 - 7.18 2 + -17 -13 -4 14.43 - 4.44 4 + -17 -13 2 5.88 - 5.79 2 + -17 -12 1 0.00 - 2.15 3 + -17 -11 -6 8.83 - 5.40 3 + -17 -11 1 4.22 - 6.52 7 + -17 -10 1 1.33 - 1.23 21 + -17 -9 -2 4.71 - 1.77 143 + -17 -9 3 0.00 - 0.00 2 + -17 -8 -5 3.66 - 1.10 40 + -17 -8 2 0.33 - 3.93 10 + -17 -7 -1 13.54 - 7.01 102 + -17 -7 3 1.74 - 1.23 2 + -17 -6 -9 2.50 - 1.77 2 + -17 -6 -5 2.08 - 2.37 15 + -17 -5 -5 8.22 - 3.02 4 + -17 -5 2 7.73 - 3.95 3 + -17 -4 -6 5.99 - 1.83 2 + -17 -3 -9 3.62 - 0.50 2 + -17 -3 -3 -2.38 - 0.07 2 + -17 -2 -6 -3.27 - 2.31 2 + -16 -16 -3 0.04 - 3.11 2 + -16 -15 -2 5.66 - 2.46 33 + -16 -14 -5 57.87 - 31.51 75 + -16 -14 3 3.08 - 2.09 18 + -16 -13 -3 5.95 - 1.49 340 + -16 -13 5 -6.29 - 5.03 4 + -16 -12 -3 32.92 - 6.61 461 + -16 -12 5 4.64 - 1.97 28 + -16 -12 6 4.70 - 10.61 2 + -16 -11 -2 16.39 - 2.52 477 + -16 -11 5 1.65 - 0.81 80 + -16 -10 -10 -0.30 - 2.94 8 + -16 -10 -3 25.30 - 4.53 492 + -16 -10 4 22.28 - 5.35 263 + -16 -9 -6 5.07 - 0.60 481 + -16 -9 1 4.82 - 0.62 478 + -16 -8 -8 27.50 - 7.96 301 + -16 -8 -1 50.29 - 6.68 580 + -16 -7 -10 2.26 - 2.60 13 + -16 -7 -3 6.24 - 1.17 585 + -16 -7 5 2.81 - 0.56 218 + -16 -7 7 1.16 - 3.64 5 + -16 -6 -5 6.30 - 0.67 540 + -16 -6 2 8.05 - 1.50 480 + -16 -5 -8 3.85 - 0.59 257 + -16 -5 -1 10.63 - 2.22 537 + -16 -5 7 5.56 - 4.54 3 + -16 -4 -9 8.48 - 3.82 16 + -16 -4 -2 41.02 - 8.50 517 + -16 -4 6 -0.25 - 1.68 3 + -16 -3 -9 -3.72 - 3.45 4 + -16 -3 -2 5.56 - 0.66 399 + -16 -3 6 9.05 - 4.62 5 + -16 -2 -2 7.94 - 1.71 332 + -16 -1 -7 4.19 - 3.65 4 + -16 -1 1 3.88 - 0.86 79 + -16 0 -7 1.87 - 1.04 2 + -16 0 -1 2.76 - 1.76 15 + -16 1 1 10.56 - 8.25 2 + -15 -17 -1 3.23 - 1.49 21 + -15 -16 -6 17.08 - 12.72 30 + -15 -16 2 4.03 - 1.06 131 + -15 -15 -5 5.13 - 0.64 344 + -15 -15 3 24.05 - 6.08 256 +End of reflections diff --git a/tests/files/hkl/dark_half2.hkl b/tests/files/hkl/dark_half2.hkl new file mode 100644 index 00000000..4487f2e2 --- /dev/null +++ b/tests/files/hkl/dark_half2.hkl @@ -0,0 +1,69 @@ +CrystFEL reflection list version 2.0 +Symmetry: 1 + h k l I phase sigma(I) nmeas + -17 -14 -2 -4.16 - 4.19 3 + -17 -13 -5 -1.97 - 3.22 2 + -17 -13 0 11.95 - 1.25 2 + -17 -12 1 1.93 - 1.35 3 + -17 -11 1 11.89 - 4.65 10 + -17 -10 1 8.99 - 6.60 22 + -17 -10 3 -4.42 - 3.12 2 + -17 -9 -2 11.01 - 3.53 126 + -17 -9 4 -0.50 - 5.55 3 + -17 -8 -5 3.27 - 2.19 35 + -17 -8 2 3.62 - 2.17 16 + -17 -7 -7 -2.56 - 3.48 4 + -17 -7 -1 2.64 - 1.18 129 + -17 -6 -5 2.82 - 1.53 14 + -17 -5 -5 5.02 - 3.43 8 + -17 -4 -6 -0.72 - 1.87 2 + -17 -3 -6 6.03 - 4.26 2 + -17 -3 -3 3.83 - 5.60 4 + -17 -3 -2 0.00 - 0.00 2 + -17 -3 2 -0.49 - 0.35 2 + -16 -15 -2 6.92 - 3.93 44 + -16 -15 2 4.75 - 2.25 5 + -16 -14 -5 7.10 - 2.17 87 + -16 -14 3 8.33 - 1.74 10 + -16 -13 -3 6.25 - 0.66 314 + -16 -13 5 7.51 - 5.24 3 + -16 -12 -3 38.54 - 12.11 437 + -16 -12 5 13.57 - 6.30 16 + -16 -11 -2 11.11 - 1.90 548 + -16 -11 5 5.14 - 1.45 88 + -16 -11 7 -0.60 - 3.65 2 + -16 -10 -3 22.98 - 4.54 515 + -16 -10 4 21.13 - 6.31 285 + -16 -9 -6 5.19 - 0.71 441 + -16 -9 1 6.26 - 1.00 461 + -16 -9 8 -3.36 - 5.17 3 + -16 -8 -8 27.27 - 6.53 307 + -16 -8 -1 42.50 - 9.09 622 + -16 -7 -10 2.94 - 2.70 12 + -16 -7 -3 6.80 - 1.06 572 + -16 -7 5 3.34 - 0.69 225 + -16 -6 -5 5.08 - 0.69 492 + -16 -6 2 8.70 - 1.24 474 + -16 -5 -8 2.82 - 0.60 241 + -16 -5 -1 8.94 - 1.24 547 + -16 -4 -9 45.76 - 38.25 20 + -16 -4 -2 26.22 - 10.09 534 + -16 -4 6 -1.56 - 1.27 3 + -16 -3 -2 5.97 - 0.63 395 + -16 -3 6 14.18 - 6.81 3 + -16 -2 -2 8.23 - 1.19 345 + -16 -2 5 9.74 - 3.30 2 + -16 -1 -7 2.71 - 3.22 3 + -16 -1 1 5.17 - 0.98 64 + -16 0 -1 2.02 - 2.98 13 + -16 1 -4 -8.38 - 13.42 2 + -15 -18 1 2.78 - 1.97 2 + -15 -17 -1 6.44 - 1.67 27 + -15 -17 4 3.35 - 14.34 2 + -15 -16 -6 10.24 - 5.08 44 + -15 -16 2 3.49 - 1.11 117 + -15 -15 -5 6.40 - 0.66 312 + -15 -15 3 27.05 - 5.93 272 + -15 -15 6 10.83 - 7.82 3 + -15 -14 -8 5.60 - 1.55 231 +End of reflections diff --git a/tests/files/mtz/1BYW.mtz b/tests/files/mtz/1BYW.mtz new file mode 100644 index 00000000..3703e898 Binary files /dev/null and b/tests/files/mtz/1BYW.mtz differ diff --git a/tests/files/mtz/1VER.mtz b/tests/files/mtz/1VER.mtz new file mode 100644 index 00000000..66a1ad24 Binary files /dev/null and b/tests/files/mtz/1VER.mtz differ diff --git a/tests/files/mtz/6JZA.mtz b/tests/files/mtz/6JZA.mtz new file mode 100644 index 00000000..abefe62b Binary files /dev/null and b/tests/files/mtz/6JZA.mtz differ diff --git a/tests/files/mtz/6SXW.mtz b/tests/files/mtz/6SXW.mtz new file mode 100644 index 00000000..7db233a9 Binary files /dev/null and b/tests/files/mtz/6SXW.mtz differ diff --git a/tests/files/mtz/6VHI.mtz b/tests/files/mtz/6VHI.mtz new file mode 100644 index 00000000..850ba605 Binary files /dev/null and b/tests/files/mtz/6VHI.mtz differ diff --git a/tests/files/pdb/1BYW_af.pdb b/tests/files/pdb/1BYW_af.pdb new file mode 100644 index 00000000..4fc50cac --- /dev/null +++ b/tests/files/pdb/1BYW_af.pdb @@ -0,0 +1,839 @@ +REMARK TITLE 1BYW AlphaFold-start MR +REMARK Log-Likelihood Gain: 611.248 +REMARK RFZ=3.8 TFZ=9.6 PAK=2 LLG=98 TFZ==11.6 LLG=611 TFZ==27.1 PAK=3 LLG=611 TFZ==27.1 +REMARK ENSEMBLE e_Q12809 EULER 106.95 45.71 227.08 FRAC -0.656 -1.300 0.015 +CRYST1 56.100 56.100 135.500 90.00 90.00 120.00 P 65 2 2 12 +SCALE1 0.017825 0.010291 -0.000000 0.00000 +SCALE2 0.000000 0.020583 -0.000000 0.00000 +SCALE3 0.000000 0.000000 0.007380 0.00000 +ATOM 1 N SER A 26 17.596 -26.152 -4.285 1.00 35.00 N +ATOM 2 CA SER A 26 17.299 -25.042 -3.356 1.00 35.00 C +ATOM 3 C SER A 26 16.551 -23.860 -4.005 1.00 35.00 C +ATOM 4 O SER A 26 16.710 -22.706 -3.608 1.00 35.00 O +ATOM 5 CB SER A 26 18.547 -24.626 -2.556 1.00 35.00 C +ATOM 6 OG SER A 26 19.643 -24.298 -3.392 1.00 35.00 O +ATOM 7 N ARG A 27 15.737 -24.135 -5.035 1.00 25.74 N +ATOM 8 CA ARG A 27 14.947 -23.112 -5.739 1.00 25.74 C +ATOM 9 C ARG A 27 13.782 -22.621 -4.880 1.00 25.74 C +ATOM 10 O ARG A 27 13.288 -23.341 -4.015 1.00 25.74 O +ATOM 11 CB ARG A 27 14.461 -23.651 -7.091 1.00 25.74 C +ATOM 12 CG ARG A 27 15.587 -23.660 -8.133 1.00 25.74 C +ATOM 13 CD ARG A 27 15.062 -24.236 -9.450 1.00 25.74 C +ATOM 14 NE ARG A 27 16.086 -24.199 -10.512 1.00 25.74 N +ATOM 15 CZ ARG A 27 16.048 -23.513 -11.638 1.00 25.74 C +ATOM 16 NH1 ARG A 27 15.022 -22.785 -11.977 1.00 25.74 N +ATOM 17 NH2 ARG A 27 17.065 -23.548 -12.448 1.00 25.74 N +ATOM 18 N LYS A 28 13.314 -21.404 -5.156 1.00 19.82 N +ATOM 19 CA LYS A 28 12.194 -20.750 -4.466 1.00 19.82 C +ATOM 20 C LYS A 28 11.059 -20.614 -5.479 1.00 19.82 C +ATOM 21 O LYS A 28 11.216 -19.897 -6.464 1.00 19.82 O +ATOM 22 CB LYS A 28 12.679 -19.404 -3.886 1.00 19.82 C +ATOM 23 CG LYS A 28 13.865 -19.582 -2.910 1.00 19.82 C +ATOM 24 CD LYS A 28 14.401 -18.256 -2.345 1.00 19.82 C +ATOM 25 CE LYS A 28 15.348 -18.491 -1.145 1.00 19.82 C +ATOM 26 NZ LYS A 28 14.667 -18.391 0.181 1.00 19.82 N +ATOM 27 N PHE A 29 9.976 -21.374 -5.327 1.00 19.08 N +ATOM 28 CA PHE A 29 8.973 -21.528 -6.386 1.00 19.08 C +ATOM 29 C PHE A 29 7.566 -21.855 -5.870 1.00 19.08 C +ATOM 30 O PHE A 29 7.379 -22.407 -4.781 1.00 19.08 O +ATOM 31 CB PHE A 29 9.440 -22.585 -7.410 1.00 19.08 C +ATOM 32 CG PHE A 29 9.376 -24.023 -6.924 1.00 19.08 C +ATOM 33 CD1 PHE A 29 10.346 -24.511 -6.029 1.00 19.08 C +ATOM 34 CD2 PHE A 29 8.325 -24.866 -7.340 1.00 19.08 C +ATOM 35 CE1 PHE A 29 10.257 -25.824 -5.535 1.00 19.08 C +ATOM 36 CE2 PHE A 29 8.245 -26.184 -6.854 1.00 19.08 C +ATOM 37 CZ PHE A 29 9.207 -26.659 -5.946 1.00 19.08 C +ATOM 38 N ILE A 30 6.579 -21.551 -6.713 1.00 18.06 N +ATOM 39 CA ILE A 30 5.178 -21.949 -6.540 1.00 18.06 C +ATOM 40 C ILE A 30 4.677 -22.676 -7.794 1.00 18.06 C +ATOM 41 O ILE A 30 5.193 -22.465 -8.896 1.00 18.06 O +ATOM 42 CB ILE A 30 4.276 -20.753 -6.139 1.00 18.06 C +ATOM 43 CG1 ILE A 30 4.053 -19.787 -7.321 1.00 18.06 C +ATOM 44 CG2 ILE A 30 4.846 -20.039 -4.900 1.00 18.06 C +ATOM 45 CD1 ILE A 30 3.349 -18.469 -6.973 1.00 18.06 C +ATOM 46 N ILE A 31 3.662 -23.524 -7.627 1.00 19.47 N +ATOM 47 CA ILE A 31 2.945 -24.191 -8.722 1.00 19.47 C +ATOM 48 C ILE A 31 1.488 -23.730 -8.698 1.00 19.47 C +ATOM 49 O ILE A 31 0.842 -23.764 -7.650 1.00 19.47 O +ATOM 50 CB ILE A 31 3.054 -25.735 -8.658 1.00 19.47 C +ATOM 51 CG1 ILE A 31 4.522 -26.204 -8.500 1.00 19.47 C +ATOM 52 CG2 ILE A 31 2.418 -26.349 -9.924 1.00 19.47 C +ATOM 53 CD1 ILE A 31 4.695 -27.726 -8.398 1.00 19.47 C +ATOM 54 N ALA A 32 0.965 -23.339 -9.857 1.00 20.23 N +ATOM 55 CA ALA A 32 -0.410 -22.889 -10.050 1.00 20.23 C +ATOM 56 C ALA A 32 -1.144 -23.742 -11.096 1.00 20.23 C +ATOM 57 O ALA A 32 -0.541 -24.202 -12.067 1.00 20.23 O +ATOM 58 CB ALA A 32 -0.368 -21.413 -10.433 1.00 20.23 C +ATOM 59 N ASN A 33 -2.448 -23.959 -10.908 1.00 21.17 N +ATOM 60 CA ASN A 33 -3.269 -24.783 -11.800 1.00 21.17 C +ATOM 61 C ASN A 33 -3.914 -23.935 -12.911 1.00 21.17 C +ATOM 62 O ASN A 33 -4.847 -23.181 -12.645 1.00 21.17 O +ATOM 63 CB ASN A 33 -4.309 -25.528 -10.945 1.00 21.17 C +ATOM 64 CG ASN A 33 -5.199 -26.469 -11.740 1.00 21.17 C +ATOM 65 OD1 ASN A 33 -5.139 -26.594 -12.953 1.00 21.17 O +ATOM 66 ND2 ASN A 33 -6.068 -27.182 -11.067 1.00 21.17 N +ATOM 67 N ALA A 34 -3.474 -24.114 -14.159 1.00 25.11 N +ATOM 68 CA ALA A 34 -3.938 -23.341 -15.316 1.00 25.11 C +ATOM 69 C ALA A 34 -5.372 -23.682 -15.772 1.00 25.11 C +ATOM 70 O ALA A 34 -5.903 -23.033 -16.666 1.00 25.11 O +ATOM 71 CB ALA A 34 -2.938 -23.548 -16.461 1.00 25.11 C +ATOM 72 N ARG A 35 -5.999 -24.718 -15.190 1.00 32.25 N +ATOM 73 CA ARG A 35 -7.390 -25.115 -15.480 1.00 32.25 C +ATOM 74 C ARG A 35 -8.420 -24.520 -14.515 1.00 32.25 C +ATOM 75 O ARG A 35 -9.608 -24.779 -14.679 1.00 32.25 O +ATOM 76 CB ARG A 35 -7.505 -26.644 -15.527 1.00 32.25 C +ATOM 77 CG ARG A 35 -6.637 -27.268 -16.627 1.00 32.25 C +ATOM 78 CD ARG A 35 -6.897 -28.775 -16.678 1.00 32.25 C +ATOM 79 NE ARG A 35 -6.175 -29.404 -17.798 1.00 32.25 N +ATOM 80 CZ ARG A 35 -6.313 -30.647 -18.219 1.00 32.25 C +ATOM 81 NH1 ARG A 35 -7.119 -31.488 -17.629 1.00 32.25 N +ATOM 82 NH2 ARG A 35 -5.639 -31.068 -19.251 1.00 32.25 N +ATOM 83 N VAL A 36 -7.978 -23.781 -13.498 1.00 31.03 N +ATOM 84 CA VAL A 36 -8.844 -23.024 -12.580 1.00 31.03 C +ATOM 85 C VAL A 36 -8.855 -21.570 -13.038 1.00 31.03 C +ATOM 86 O VAL A 36 -7.797 -21.032 -13.343 1.00 31.03 O +ATOM 87 CB VAL A 36 -8.356 -23.163 -11.125 1.00 31.03 C +ATOM 88 CG1 VAL A 36 -9.137 -22.281 -10.143 1.00 31.03 C +ATOM 89 CG2 VAL A 36 -8.499 -24.616 -10.652 1.00 31.03 C +ATOM 90 N GLU A 37 -10.028 -20.939 -13.059 1.00 39.23 N +ATOM 91 CA GLU A 37 -10.260 -19.611 -13.655 1.00 39.23 C +ATOM 92 C GLU A 37 -9.263 -18.543 -13.166 1.00 39.23 C +ATOM 93 O GLU A 37 -8.665 -17.837 -13.973 1.00 39.23 O +ATOM 94 CB GLU A 37 -11.714 -19.191 -13.361 1.00 39.23 C +ATOM 95 CG GLU A 37 -12.729 -20.164 -13.998 1.00 39.23 C +ATOM 96 CD GLU A 37 -14.199 -19.868 -13.659 1.00 39.23 C +ATOM 97 OE1 GLU A 37 -15.059 -20.484 -14.330 1.00 39.23 O +ATOM 98 OE2 GLU A 37 -14.451 -19.092 -12.712 1.00 39.23 O +ATOM 99 N ASN A 38 -8.987 -18.509 -11.858 1.00 30.61 N +ATOM 100 CA ASN A 38 -8.056 -17.557 -11.234 1.00 30.61 C +ATOM 101 C ASN A 38 -6.583 -18.026 -11.213 1.00 30.61 C +ATOM 102 O ASN A 38 -5.766 -17.429 -10.513 1.00 30.61 O +ATOM 103 CB ASN A 38 -8.582 -17.219 -9.824 1.00 30.61 C +ATOM 104 CG ASN A 38 -9.931 -16.519 -9.841 1.00 30.61 C +ATOM 105 OD1 ASN A 38 -10.362 -15.951 -10.824 1.00 30.61 O +ATOM 106 ND2 ASN A 38 -10.645 -16.528 -8.743 1.00 30.61 N +ATOM 107 N CYS A 39 -6.243 -19.121 -11.908 1.00 24.42 N +ATOM 108 CA CYS A 39 -4.930 -19.784 -11.871 1.00 24.42 C +ATOM 109 C CYS A 39 -4.385 -19.949 -10.437 1.00 24.42 C +ATOM 110 O CYS A 39 -3.297 -19.489 -10.095 1.00 24.42 O +ATOM 111 CB CYS A 39 -3.961 -19.094 -12.841 1.00 24.42 C +ATOM 112 SG CYS A 39 -4.467 -19.431 -14.552 1.00 24.42 S +ATOM 113 N ALA A 40 -5.186 -20.577 -9.573 1.00 20.39 N +ATOM 114 CA ALA A 40 -4.914 -20.658 -8.141 1.00 20.39 C +ATOM 115 C ALA A 40 -3.633 -21.452 -7.812 1.00 20.39 C +ATOM 116 O ALA A 40 -3.345 -22.488 -8.428 1.00 20.39 O +ATOM 117 CB ALA A 40 -6.145 -21.242 -7.442 1.00 20.39 C +ATOM 118 N VAL A 41 -2.891 -20.988 -6.800 1.00 19.18 N +ATOM 119 CA VAL A 41 -1.688 -21.654 -6.277 1.00 19.18 C +ATOM 120 C VAL A 41 -2.057 -22.953 -5.549 1.00 19.18 C +ATOM 121 O VAL A 41 -2.888 -22.960 -4.638 1.00 19.18 O +ATOM 122 CB VAL A 41 -0.872 -20.707 -5.372 1.00 19.18 C +ATOM 123 CG1 VAL A 41 0.366 -21.405 -4.790 1.00 19.18 C +ATOM 124 CG2 VAL A 41 -0.376 -19.490 -6.166 1.00 19.18 C +ATOM 125 N ILE A 42 -1.411 -24.051 -5.958 1.00 20.39 N +ATOM 126 CA ILE A 42 -1.597 -25.420 -5.443 1.00 20.39 C +ATOM 127 C ILE A 42 -0.362 -25.983 -4.721 1.00 20.39 C +ATOM 128 O ILE A 42 -0.460 -27.027 -4.079 1.00 20.39 O +ATOM 129 CB ILE A 42 -2.044 -26.385 -6.568 1.00 20.39 C +ATOM 130 CG1 ILE A 42 -1.008 -26.489 -7.714 1.00 20.39 C +ATOM 131 CG2 ILE A 42 -3.429 -25.979 -7.102 1.00 20.39 C +ATOM 132 CD1 ILE A 42 -1.231 -27.688 -8.645 1.00 20.39 C +ATOM 133 N TYR A 43 0.795 -25.326 -4.831 1.00 19.32 N +ATOM 134 CA TYR A 43 1.995 -25.634 -4.049 1.00 19.32 C +ATOM 135 C TYR A 43 2.857 -24.383 -3.882 1.00 19.32 C +ATOM 136 O TYR A 43 3.029 -23.623 -4.835 1.00 19.32 O +ATOM 137 CB TYR A 43 2.827 -26.736 -4.731 1.00 19.32 C +ATOM 138 CG TYR A 43 3.981 -27.262 -3.888 1.00 19.32 C +ATOM 139 CD1 TYR A 43 5.328 -26.998 -4.228 1.00 19.32 C +ATOM 140 CD2 TYR A 43 3.690 -28.030 -2.747 1.00 19.32 C +ATOM 141 CE1 TYR A 43 6.378 -27.527 -3.439 1.00 19.32 C +ATOM 142 CE2 TYR A 43 4.735 -28.535 -1.950 1.00 19.32 C +ATOM 143 CZ TYR A 43 6.082 -28.301 -2.294 1.00 19.32 C +ATOM 144 OH TYR A 43 7.040 -28.840 -1.489 1.00 19.32 O +ATOM 145 N CYS A 44 3.468 -24.224 -2.714 1.00 19.27 N +ATOM 146 CA CYS A 44 4.596 -23.325 -2.488 1.00 19.27 C +ATOM 147 C CYS A 44 5.634 -24.056 -1.630 1.00 19.27 C +ATOM 148 O CYS A 44 5.269 -24.715 -0.649 1.00 19.27 O +ATOM 149 CB CYS A 44 4.106 -22.004 -1.870 1.00 19.27 C +ATOM 150 SG CYS A 44 3.344 -22.234 -0.236 1.00 19.27 S +ATOM 151 N ASN A 45 6.913 -23.970 -1.999 1.00 19.18 N +ATOM 152 CA ASN A 45 7.989 -24.589 -1.225 1.00 19.18 C +ATOM 153 C ASN A 45 8.447 -23.672 -0.079 1.00 19.18 C +ATOM 154 O ASN A 45 8.228 -22.460 -0.121 1.00 19.18 O +ATOM 155 CB ASN A 45 9.115 -25.041 -2.167 1.00 19.18 C +ATOM 156 CG ASN A 45 10.186 -23.997 -2.433 1.00 19.18 C +ATOM 157 OD1 ASN A 45 9.945 -22.887 -2.881 1.00 19.18 O +ATOM 158 ND2 ASN A 45 11.424 -24.345 -2.181 1.00 19.18 N +ATOM 159 N ASP A 46 9.098 -24.223 0.944 1.00 19.92 N +ATOM 160 CA ASP A 46 9.416 -23.449 2.157 1.00 19.92 C +ATOM 161 C ASP A 46 10.400 -22.300 1.868 1.00 19.92 C +ATOM 162 O ASP A 46 10.244 -21.200 2.392 1.00 19.92 O +ATOM 163 CB ASP A 46 9.896 -24.401 3.262 1.00 19.92 C +ATOM 164 CG ASP A 46 8.922 -25.576 3.396 1.00 19.92 C +ATOM 165 OD1 ASP A 46 7.744 -25.357 3.771 1.00 19.92 O +ATOM 166 OD2 ASP A 46 9.264 -26.661 2.895 1.00 19.92 O +ATOM 167 N GLY A 47 11.301 -22.481 0.894 1.00 20.45 N +ATOM 168 CA GLY A 47 12.178 -21.420 0.392 1.00 20.45 C +ATOM 169 C GLY A 47 11.446 -20.225 -0.245 1.00 20.45 C +ATOM 170 O GLY A 47 12.018 -19.133 -0.298 1.00 20.45 O +ATOM 171 N PHE A 48 10.201 -20.385 -0.708 1.00 18.85 N +ATOM 172 CA PHE A 48 9.337 -19.266 -1.106 1.00 18.85 C +ATOM 173 C PHE A 48 8.702 -18.579 0.113 1.00 18.85 C +ATOM 174 O PHE A 48 8.687 -17.349 0.173 1.00 18.85 O +ATOM 175 CB PHE A 48 8.271 -19.745 -2.102 1.00 18.85 C +ATOM 176 CG PHE A 48 7.297 -18.654 -2.507 1.00 18.85 C +ATOM 177 CD1 PHE A 48 6.076 -18.506 -1.820 1.00 18.85 C +ATOM 178 CD2 PHE A 48 7.630 -17.754 -3.537 1.00 18.85 C +ATOM 179 CE1 PHE A 48 5.190 -17.473 -2.173 1.00 18.85 C +ATOM 180 CE2 PHE A 48 6.740 -16.723 -3.889 1.00 18.85 C +ATOM 181 CZ PHE A 48 5.515 -16.590 -3.214 1.00 18.85 C +ATOM 182 N CYS A 49 8.246 -19.338 1.115 1.00 19.08 N +ATOM 183 CA CYS A 49 7.742 -18.777 2.375 1.00 19.08 C +ATOM 184 C CYS A 49 8.818 -17.945 3.099 1.00 19.08 C +ATOM 185 O CYS A 49 8.550 -16.810 3.483 1.00 19.08 O +ATOM 186 CB CYS A 49 7.189 -19.910 3.251 1.00 19.08 C +ATOM 187 SG CYS A 49 5.642 -20.529 2.523 1.00 19.08 S +ATOM 188 N GLU A 50 10.064 -18.428 3.167 1.00 20.39 N +ATOM 189 CA GLU A 50 11.226 -17.650 3.637 1.00 20.39 C +ATOM 190 C GLU A 50 11.470 -16.359 2.839 1.00 20.39 C +ATOM 191 O GLU A 50 11.959 -15.368 3.377 1.00 20.39 O +ATOM 192 CB GLU A 50 12.501 -18.463 3.415 1.00 20.39 C +ATOM 193 CG GLU A 50 12.724 -19.694 4.293 1.00 20.39 C +ATOM 194 CD GLU A 50 13.969 -20.463 3.809 1.00 20.39 C +ATOM 195 OE1 GLU A 50 14.301 -21.473 4.459 1.00 20.39 O +ATOM 196 OE2 GLU A 50 14.571 -20.049 2.772 1.00 20.39 O +ATOM 197 N LEU A 51 11.206 -16.374 1.526 1.00 20.55 N +ATOM 198 CA LEU A 51 11.493 -15.237 0.651 1.00 20.55 C +ATOM 199 C LEU A 51 10.528 -14.077 0.895 1.00 20.55 C +ATOM 200 O LEU A 51 10.959 -12.926 0.862 1.00 20.55 O +ATOM 201 CB LEU A 51 11.445 -15.691 -0.818 1.00 20.55 C +ATOM 202 CG LEU A 51 11.724 -14.583 -1.851 1.00 20.55 C +ATOM 203 CD1 LEU A 51 13.111 -13.960 -1.681 1.00 20.55 C +ATOM 204 CD2 LEU A 51 11.627 -15.156 -3.264 1.00 20.55 C +ATOM 205 N CYS A 52 9.249 -14.387 1.108 1.00 19.38 N +ATOM 206 CA CYS A 52 8.172 -13.404 1.214 1.00 19.38 C +ATOM 207 C CYS A 52 7.684 -13.152 2.653 1.00 19.38 C +ATOM 208 O CYS A 52 7.060 -12.124 2.903 1.00 19.38 O +ATOM 209 CB CYS A 52 7.040 -13.825 0.269 1.00 19.38 C +ATOM 210 SG CYS A 52 6.273 -15.389 0.787 1.00 19.38 S +ATOM 211 N GLY A 53 7.962 -14.053 3.599 1.00 19.97 N +ATOM 212 CA GLY A 53 7.482 -13.982 4.985 1.00 19.97 C +ATOM 213 C GLY A 53 6.027 -14.430 5.190 1.00 19.97 C +ATOM 214 O GLY A 53 5.541 -14.380 6.317 1.00 19.97 O +ATOM 215 N TYR A 54 5.328 -14.868 4.139 1.00 18.10 N +ATOM 216 CA TYR A 54 3.965 -15.401 4.227 1.00 18.10 C +ATOM 217 C TYR A 54 3.968 -16.922 4.421 1.00 18.10 C +ATOM 218 O TYR A 54 4.750 -17.652 3.805 1.00 18.10 O +ATOM 219 CB TYR A 54 3.144 -15.012 2.993 1.00 18.10 C +ATOM 220 CG TYR A 54 2.868 -13.526 2.849 1.00 18.10 C +ATOM 221 CD1 TYR A 54 1.772 -12.950 3.518 1.00 18.10 C +ATOM 222 CD2 TYR A 54 3.678 -12.727 2.021 1.00 18.10 C +ATOM 223 CE1 TYR A 54 1.471 -11.587 3.339 1.00 18.10 C +ATOM 224 CE2 TYR A 54 3.376 -11.363 1.830 1.00 18.10 C +ATOM 225 CZ TYR A 54 2.255 -10.797 2.474 1.00 18.10 C +ATOM 226 OH TYR A 54 1.899 -9.509 2.235 1.00 18.10 O +ATOM 227 N SER A 55 3.054 -17.412 5.256 1.00 18.36 N +ATOM 228 CA SER A 55 2.863 -18.840 5.499 1.00 18.36 C +ATOM 229 C SER A 55 2.203 -19.548 4.310 1.00 18.36 C +ATOM 230 O SER A 55 1.506 -18.943 3.491 1.00 18.36 O +ATOM 231 CB SER A 55 2.036 -19.051 6.772 1.00 18.36 C +ATOM 232 OG SER A 55 0.702 -18.648 6.557 1.00 18.36 O +ATOM 233 N ARG A 56 2.353 -20.877 4.249 1.00 18.80 N +ATOM 234 CA ARG A 56 1.694 -21.725 3.240 1.00 18.80 C +ATOM 235 C ARG A 56 0.165 -21.550 3.227 1.00 18.80 C +ATOM 236 O ARG A 56 -0.430 -21.636 2.160 1.00 18.80 O +ATOM 237 CB ARG A 56 2.152 -23.180 3.462 1.00 18.80 C +ATOM 238 CG ARG A 56 1.562 -24.201 2.473 1.00 18.80 C +ATOM 239 CD ARG A 56 2.245 -25.577 2.589 1.00 18.80 C +ATOM 240 NE ARG A 56 3.596 -25.594 1.987 1.00 18.80 N +ATOM 241 CZ ARG A 56 4.768 -25.823 2.564 1.00 18.80 C +ATOM 242 NH1 ARG A 56 4.954 -26.076 3.824 1.00 18.80 N +ATOM 243 NH2 ARG A 56 5.848 -25.773 1.860 1.00 18.80 N +ATOM 244 N ALA A 57 -0.463 -21.254 4.367 1.00 19.71 N +ATOM 245 CA ALA A 57 -1.905 -21.008 4.453 1.00 19.71 C +ATOM 246 C ALA A 57 -2.339 -19.691 3.777 1.00 19.71 C +ATOM 247 O ALA A 57 -3.406 -19.637 3.178 1.00 19.71 O +ATOM 248 CB ALA A 57 -2.308 -21.033 5.932 1.00 19.71 C +ATOM 249 N GLU A 58 -1.506 -18.648 3.828 1.00 18.85 N +ATOM 250 CA GLU A 58 -1.804 -17.335 3.231 1.00 18.85 C +ATOM 251 C GLU A 58 -1.599 -17.315 1.707 1.00 18.85 C +ATOM 252 O GLU A 58 -2.282 -16.568 1.001 1.00 18.85 O +ATOM 253 CB GLU A 58 -0.900 -16.279 3.886 1.00 18.85 C +ATOM 254 CG GLU A 58 -1.229 -16.051 5.367 1.00 18.85 C +ATOM 255 CD GLU A 58 -0.074 -15.336 6.067 1.00 18.85 C +ATOM 256 OE1 GLU A 58 -0.127 -14.102 6.237 1.00 18.85 O +ATOM 257 OE2 GLU A 58 0.915 -16.010 6.442 1.00 18.85 O +ATOM 258 N VAL A 59 -0.656 -18.133 1.219 1.00 18.76 N +ATOM 259 CA VAL A 59 -0.229 -18.230 -0.190 1.00 18.76 C +ATOM 260 C VAL A 59 -1.100 -19.192 -1.007 1.00 18.76 C +ATOM 261 O VAL A 59 -1.348 -18.949 -2.189 1.00 18.76 O +ATOM 262 CB VAL A 59 1.253 -18.669 -0.240 1.00 18.76 C +ATOM 263 CG1 VAL A 59 1.769 -18.943 -1.659 1.00 18.76 C +ATOM 264 CG2 VAL A 59 2.168 -17.604 0.381 1.00 18.76 C +ATOM 265 N MET A 60 -1.556 -20.297 -0.411 1.00 20.94 N +ATOM 266 CA MET A 60 -2.387 -21.293 -1.097 1.00 20.94 C +ATOM 267 C MET A 60 -3.745 -20.712 -1.517 1.00 20.94 C +ATOM 268 O MET A 60 -4.274 -19.807 -0.882 1.00 20.94 O +ATOM 269 CB MET A 60 -2.559 -22.529 -0.204 1.00 20.94 C +ATOM 270 CG MET A 60 -1.279 -23.379 -0.165 1.00 20.94 C +ATOM 271 SD MET A 60 -0.929 -24.371 -1.644 1.00 20.94 S +ATOM 272 CE MET A 60 -2.231 -25.622 -1.491 1.00 20.94 C +ATOM 273 N GLN A 61 -4.311 -21.232 -2.611 1.00 22.89 N +ATOM 274 CA GLN A 61 -5.579 -20.783 -3.217 1.00 22.89 C +ATOM 275 C GLN A 61 -5.620 -19.321 -3.726 1.00 22.89 C +ATOM 276 O GLN A 61 -6.571 -18.958 -4.420 1.00 22.89 O +ATOM 277 CB GLN A 61 -6.785 -21.119 -2.312 1.00 22.89 C +ATOM 278 CG GLN A 61 -6.899 -22.619 -1.983 1.00 22.89 C +ATOM 279 CD GLN A 61 -8.141 -22.965 -1.163 1.00 22.89 C +ATOM 280 OE1 GLN A 61 -8.829 -22.135 -0.603 1.00 22.89 O +ATOM 281 NE2 GLN A 61 -8.493 -24.227 -1.051 1.00 22.89 N +ATOM 282 N ARG A 62 -4.589 -18.498 -3.487 1.00 20.61 N +ATOM 283 CA ARG A 62 -4.448 -17.155 -4.081 1.00 20.61 C +ATOM 284 C ARG A 62 -4.180 -17.218 -5.599 1.00 20.61 C +ATOM 285 O ARG A 62 -3.663 -18.234 -6.077 1.00 20.61 O +ATOM 286 CB ARG A 62 -3.309 -16.384 -3.389 1.00 20.61 C +ATOM 287 CG ARG A 62 -3.503 -16.124 -1.888 1.00 20.61 C +ATOM 288 CD ARG A 62 -4.673 -15.177 -1.586 1.00 20.61 C +ATOM 289 NE ARG A 62 -4.683 -14.810 -0.158 1.00 20.61 N +ATOM 290 CZ ARG A 62 -5.267 -13.760 0.399 1.00 20.61 C +ATOM 291 NH1 ARG A 62 -5.907 -12.860 -0.298 1.00 20.61 N +ATOM 292 NH2 ARG A 62 -5.217 -13.588 1.687 1.00 20.61 N +ATOM 293 N PRO A 63 -4.465 -16.141 -6.360 1.00 19.82 N +ATOM 294 CA PRO A 63 -4.038 -16.009 -7.753 1.00 19.82 C +ATOM 295 C PRO A 63 -2.510 -16.001 -7.893 1.00 19.82 C +ATOM 296 O PRO A 63 -1.803 -15.369 -7.105 1.00 19.82 O +ATOM 297 CB PRO A 63 -4.633 -14.686 -8.259 1.00 19.82 C +ATOM 298 CG PRO A 63 -5.780 -14.403 -7.293 1.00 19.82 C +ATOM 299 CD PRO A 63 -5.258 -14.979 -5.979 1.00 19.82 C +ATOM 300 N CYS A 64 -1.984 -16.651 -8.933 1.00 21.22 N +ATOM 301 CA CYS A 64 -0.538 -16.740 -9.172 1.00 21.22 C +ATOM 302 C CYS A 64 0.122 -15.457 -9.716 1.00 21.22 C +ATOM 303 O CYS A 64 1.319 -15.460 -9.995 1.00 21.22 O +ATOM 304 CB CYS A 64 -0.258 -17.965 -10.045 1.00 21.22 C +ATOM 305 SG CYS A 64 -0.912 -17.732 -11.720 1.00 21.22 S +ATOM 306 N THR A 65 -0.627 -14.358 -9.857 1.00 21.75 N +ATOM 307 CA THR A 65 -0.081 -13.001 -10.045 1.00 21.75 C +ATOM 308 C THR A 65 0.596 -12.464 -8.780 1.00 21.75 C +ATOM 309 O THR A 65 1.449 -11.582 -8.876 1.00 21.75 O +ATOM 310 CB THR A 65 -1.194 -12.025 -10.456 1.00 21.75 C +ATOM 311 OG1 THR A 65 -2.303 -12.188 -9.604 1.00 21.75 O +ATOM 312 CG2 THR A 65 -1.672 -12.291 -11.884 1.00 21.75 C +ATOM 313 N CYS A 66 0.274 -13.034 -7.610 1.00 21.69 N +ATOM 314 CA CYS A 66 0.845 -12.680 -6.308 1.00 21.69 C +ATOM 315 C CYS A 66 0.684 -11.190 -5.956 1.00 21.69 C +ATOM 316 O CYS A 66 1.633 -10.521 -5.538 1.00 21.69 O +ATOM 317 CB CYS A 66 2.286 -13.200 -6.199 1.00 21.69 C +ATOM 318 SG CYS A 66 2.322 -15.015 -6.326 1.00 21.69 S +ATOM 319 N ASP A 67 -0.540 -10.674 -6.099 1.00 22.24 N +ATOM 320 CA ASP A 67 -0.890 -9.264 -5.858 1.00 22.24 C +ATOM 321 C ASP A 67 -0.501 -8.793 -4.438 1.00 22.24 C +ATOM 322 O ASP A 67 -0.088 -7.655 -4.242 1.00 22.24 O +ATOM 323 CB ASP A 67 -2.403 -9.061 -6.079 1.00 22.24 C +ATOM 324 CG ASP A 67 -2.972 -9.751 -7.329 1.00 22.24 C +ATOM 325 OD1 ASP A 67 -2.254 -9.860 -8.348 1.00 22.24 O +ATOM 326 OD2 ASP A 67 -4.103 -10.267 -7.233 1.00 22.24 O +ATOM 327 N PHE A 68 -0.532 -9.701 -3.456 1.00 20.50 N +ATOM 328 CA PHE A 68 -0.115 -9.484 -2.059 1.00 20.50 C +ATOM 329 C PHE A 68 1.405 -9.276 -1.858 1.00 20.50 C +ATOM 330 O PHE A 68 1.844 -8.959 -0.754 1.00 20.50 O +ATOM 331 CB PHE A 68 -0.631 -10.662 -1.216 1.00 20.50 C +ATOM 332 CG PHE A 68 -0.073 -12.019 -1.619 1.00 20.50 C +ATOM 333 CD1 PHE A 68 -0.749 -12.818 -2.563 1.00 20.50 C +ATOM 334 CD2 PHE A 68 1.132 -12.479 -1.055 1.00 20.50 C +ATOM 335 CE1 PHE A 68 -0.212 -14.059 -2.951 1.00 20.50 C +ATOM 336 CE2 PHE A 68 1.665 -13.723 -1.438 1.00 20.50 C +ATOM 337 CZ PHE A 68 0.996 -14.510 -2.392 1.00 20.50 C +ATOM 338 N LEU A 69 2.216 -9.449 -2.909 1.00 19.04 N +ATOM 339 CA LEU A 69 3.639 -9.082 -2.938 1.00 19.04 C +ATOM 340 C LEU A 69 3.886 -7.706 -3.581 1.00 19.04 C +ATOM 341 O LEU A 69 5.031 -7.252 -3.639 1.00 19.04 O +ATOM 342 CB LEU A 69 4.449 -10.164 -3.682 1.00 19.04 C +ATOM 343 CG LEU A 69 4.333 -11.598 -3.145 1.00 19.04 C +ATOM 344 CD1 LEU A 69 5.266 -12.507 -3.951 1.00 19.04 C +ATOM 345 CD2 LEU A 69 4.726 -11.709 -1.673 1.00 19.04 C +ATOM 346 N HIS A 70 2.853 -7.053 -4.124 1.00 19.42 N +ATOM 347 CA HIS A 70 2.984 -5.750 -4.773 1.00 19.42 C +ATOM 348 C HIS A 70 2.999 -4.612 -3.737 1.00 19.42 C +ATOM 349 O HIS A 70 2.330 -4.669 -2.710 1.00 19.42 O +ATOM 350 CB HIS A 70 1.861 -5.540 -5.803 1.00 19.42 C +ATOM 351 CG HIS A 70 1.856 -6.480 -6.991 1.00 19.42 C +ATOM 352 ND1 HIS A 70 2.354 -7.764 -7.066 1.00 19.42 N +ATOM 353 CD2 HIS A 70 1.309 -6.202 -8.213 1.00 19.42 C +ATOM 354 CE1 HIS A 70 2.117 -8.235 -8.302 1.00 19.42 C +ATOM 355 NE2 HIS A 70 1.510 -7.300 -9.051 1.00 19.42 N +ATOM 356 N GLY A 71 3.728 -3.537 -4.033 1.00 21.28 N +ATOM 357 CA GLY A 71 3.790 -2.337 -3.198 1.00 21.28 C +ATOM 358 C GLY A 71 4.329 -1.114 -3.954 1.00 21.28 C +ATOM 359 O GLY A 71 4.427 -1.143 -5.182 1.00 21.28 O +ATOM 360 N PRO A 72 4.693 -0.018 -3.261 1.00 22.57 N +ATOM 361 CA PRO A 72 4.994 1.273 -3.893 1.00 22.57 C +ATOM 362 C PRO A 72 6.069 1.274 -4.996 1.00 22.57 C +ATOM 363 O PRO A 72 5.952 2.039 -5.952 1.00 22.57 O +ATOM 364 CB PRO A 72 5.381 2.195 -2.733 1.00 22.57 C +ATOM 365 CG PRO A 72 4.543 1.656 -1.574 1.00 22.57 C +ATOM 366 CD PRO A 72 4.547 0.148 -1.820 1.00 22.57 C +ATOM 367 N ALA A 73 7.105 0.424 -4.907 1.00 23.71 N +ATOM 368 CA ALA A 73 8.149 0.284 -5.950 1.00 23.71 C +ATOM 369 C ALA A 73 7.781 -0.681 -7.090 1.00 23.71 C +ATOM 370 O ALA A 73 8.528 -0.778 -8.064 1.00 23.71 O +ATOM 371 CB ALA A 73 9.505 -0.082 -5.310 1.00 23.71 C +ATOM 372 N THR A 74 6.658 -1.395 -7.005 1.00 20.45 N +ATOM 373 CA THR A 74 6.246 -2.398 -7.997 1.00 20.45 C +ATOM 374 C THR A 74 5.719 -1.725 -9.268 1.00 20.45 C +ATOM 375 O THR A 74 4.653 -1.109 -9.285 1.00 20.45 O +ATOM 376 CB THR A 74 5.193 -3.349 -7.408 1.00 20.45 C +ATOM 377 OG1 THR A 74 5.648 -3.877 -6.182 1.00 20.45 O +ATOM 378 CG2 THR A 74 4.915 -4.536 -8.322 1.00 20.45 C +ATOM 379 N GLN A 75 6.455 -1.844 -10.376 1.00 20.89 N +ATOM 380 CA GLN A 75 6.039 -1.239 -11.642 1.00 20.89 C +ATOM 381 C GLN A 75 4.808 -1.952 -12.222 1.00 20.89 C +ATOM 382 O GLN A 75 4.880 -3.129 -12.579 1.00 20.89 O +ATOM 383 CB GLN A 75 7.179 -1.264 -12.672 1.00 20.89 C +ATOM 384 CG GLN A 75 8.431 -0.474 -12.252 1.00 20.89 C +ATOM 385 CD GLN A 75 9.438 -0.348 -13.396 1.00 20.89 C +ATOM 386 OE1 GLN A 75 9.094 -0.390 -14.569 1.00 20.89 O +ATOM 387 NE2 GLN A 75 10.705 -0.161 -13.105 1.00 20.89 N +ATOM 388 N ARG A 76 3.711 -1.215 -12.457 1.00 20.94 N +ATOM 389 CA ARG A 76 2.501 -1.745 -13.132 1.00 20.94 C +ATOM 390 C ARG A 76 2.803 -2.397 -14.492 1.00 20.94 C +ATOM 391 O ARG A 76 2.118 -3.337 -14.881 1.00 20.94 O +ATOM 392 CB ARG A 76 1.440 -0.642 -13.308 1.00 20.94 C +ATOM 393 CG ARG A 76 0.856 -0.125 -11.982 1.00 20.94 C +ATOM 394 CD ARG A 76 -0.265 0.894 -12.249 1.00 20.94 C +ATOM 395 NE ARG A 76 -0.803 1.461 -10.996 1.00 20.94 N +ATOM 396 CZ ARG A 76 -1.681 2.444 -10.882 1.00 20.94 C +ATOM 397 NH1 ARG A 76 -2.205 3.037 -11.920 1.00 20.94 N +ATOM 398 NH2 ARG A 76 -2.058 2.857 -9.705 1.00 20.94 N +ATOM 399 N ARG A 77 3.846 -1.936 -15.198 1.00 19.71 N +ATOM 400 CA ARG A 77 4.326 -2.541 -16.456 1.00 19.71 C +ATOM 401 C ARG A 77 4.887 -3.956 -16.260 1.00 19.71 C +ATOM 402 O ARG A 77 4.641 -4.801 -17.109 1.00 19.71 O +ATOM 403 CB ARG A 77 5.383 -1.643 -17.125 1.00 19.71 C +ATOM 404 CG ARG A 77 4.803 -0.331 -17.681 1.00 19.71 C +ATOM 405 CD ARG A 77 5.905 0.487 -18.370 1.00 19.71 C +ATOM 406 NE ARG A 77 5.386 1.742 -18.957 1.00 19.71 N +ATOM 407 CZ ARG A 77 6.110 2.713 -19.489 1.00 19.71 C +ATOM 408 NH1 ARG A 77 7.412 2.659 -19.545 1.00 19.71 N +ATOM 409 NH2 ARG A 77 5.534 3.773 -19.984 1.00 19.71 N +ATOM 410 N ALA A 78 5.596 -4.220 -15.161 1.00 21.22 N +ATOM 411 CA ALA A 78 6.129 -5.546 -14.843 1.00 21.22 C +ATOM 412 C ALA A 78 5.006 -6.505 -14.413 1.00 21.22 C +ATOM 413 O ALA A 78 4.889 -7.597 -14.962 1.00 21.22 O +ATOM 414 CB ALA A 78 7.214 -5.395 -13.769 1.00 21.22 C +ATOM 415 N ALA A 79 4.101 -6.061 -13.532 1.00 20.07 N +ATOM 416 CA ALA A 79 2.920 -6.841 -13.145 1.00 20.07 C +ATOM 417 C ALA A 79 2.064 -7.246 -14.367 1.00 20.07 C +ATOM 418 O ALA A 79 1.704 -8.413 -14.527 1.00 20.07 O +ATOM 419 CB ALA A 79 2.116 -6.016 -12.134 1.00 20.07 C +ATOM 420 N ALA A 80 1.829 -6.311 -15.297 1.00 20.18 N +ATOM 421 CA ALA A 80 1.126 -6.590 -16.548 1.00 20.18 C +ATOM 422 C ALA A 80 1.856 -7.605 -17.453 1.00 20.18 C +ATOM 423 O ALA A 80 1.197 -8.399 -18.118 1.00 20.18 O +ATOM 424 CB ALA A 80 0.887 -5.263 -17.277 1.00 20.18 C +ATOM 425 N GLN A 81 3.196 -7.631 -17.465 1.00 19.87 N +ATOM 426 CA GLN A 81 3.965 -8.636 -18.215 1.00 19.87 C +ATOM 427 C GLN A 81 3.848 -10.042 -17.610 1.00 19.87 C +ATOM 428 O GLN A 81 3.810 -11.010 -18.368 1.00 19.87 O +ATOM 429 CB GLN A 81 5.444 -8.235 -18.301 1.00 19.87 C +ATOM 430 CG GLN A 81 5.692 -7.076 -19.274 1.00 19.87 C +ATOM 431 CD GLN A 81 7.118 -6.542 -19.181 1.00 19.87 C +ATOM 432 OE1 GLN A 81 8.034 -7.176 -18.693 1.00 19.87 O +ATOM 433 NE2 GLN A 81 7.384 -5.358 -19.685 1.00 19.87 N +ATOM 434 N ILE A 82 3.752 -10.167 -16.280 1.00 20.12 N +ATOM 435 CA ILE A 82 3.503 -11.456 -15.609 1.00 20.12 C +ATOM 436 C ILE A 82 2.095 -11.955 -15.957 1.00 20.12 C +ATOM 437 O ILE A 82 1.943 -13.083 -16.425 1.00 20.12 O +ATOM 438 CB ILE A 82 3.738 -11.347 -14.080 1.00 20.12 C +ATOM 439 CG1 ILE A 82 5.220 -11.003 -13.794 1.00 20.12 C +ATOM 440 CG2 ILE A 82 3.348 -12.655 -13.363 1.00 20.12 C +ATOM 441 CD1 ILE A 82 5.530 -10.677 -12.329 1.00 20.12 C +ATOM 442 N ALA A 83 1.074 -11.098 -15.834 1.00 20.99 N +ATOM 443 CA ALA A 83 -0.294 -11.431 -16.235 1.00 20.99 C +ATOM 444 C ALA A 83 -0.392 -11.816 -17.727 1.00 20.99 C +ATOM 445 O ALA A 83 -1.056 -12.791 -18.076 1.00 20.99 O +ATOM 446 CB ALA A 83 -1.199 -10.242 -15.890 1.00 20.99 C +ATOM 447 N GLN A 84 0.324 -11.108 -18.610 1.00 21.40 N +ATOM 448 CA GLN A 84 0.380 -11.430 -20.038 1.00 21.40 C +ATOM 449 C GLN A 84 1.096 -12.762 -20.322 1.00 21.40 C +ATOM 450 O GLN A 84 0.663 -13.485 -21.214 1.00 21.40 O +ATOM 451 CB GLN A 84 1.028 -10.263 -20.803 1.00 21.40 C +ATOM 452 CG GLN A 84 0.914 -10.449 -22.326 1.00 21.40 C +ATOM 453 CD GLN A 84 1.463 -9.276 -23.132 1.00 21.40 C +ATOM 454 OE1 GLN A 84 2.246 -8.452 -22.683 1.00 21.40 O +ATOM 455 NE2 GLN A 84 1.092 -9.169 -24.389 1.00 21.40 N +ATOM 456 N ALA A 85 2.147 -13.120 -19.575 1.00 21.40 N +ATOM 457 CA ALA A 85 2.835 -14.407 -19.731 1.00 21.40 C +ATOM 458 C ALA A 85 1.938 -15.595 -19.342 1.00 21.40 C +ATOM 459 O ALA A 85 1.887 -16.597 -20.059 1.00 21.40 O +ATOM 460 CB ALA A 85 4.112 -14.385 -18.885 1.00 21.40 C +ATOM 461 N LEU A 86 1.186 -15.444 -18.244 1.00 22.18 N +ATOM 462 CA LEU A 86 0.196 -16.421 -17.783 1.00 22.18 C +ATOM 463 C LEU A 86 -0.920 -16.607 -18.827 1.00 22.18 C +ATOM 464 O LEU A 86 -1.168 -17.726 -19.267 1.00 22.18 O +ATOM 465 CB LEU A 86 -0.358 -15.973 -16.415 1.00 22.18 C +ATOM 466 CG LEU A 86 0.673 -15.982 -15.266 1.00 22.18 C +ATOM 467 CD1 LEU A 86 0.111 -15.233 -14.056 1.00 22.18 C +ATOM 468 CD2 LEU A 86 1.043 -17.401 -14.838 1.00 22.18 C +ATOM 469 N LEU A 87 -1.534 -15.516 -19.298 1.00 26.90 N +ATOM 470 CA LEU A 87 -2.592 -15.572 -20.320 1.00 26.90 C +ATOM 471 C LEU A 87 -2.087 -16.055 -21.693 1.00 26.90 C +ATOM 472 O LEU A 87 -2.812 -16.741 -22.409 1.00 26.90 O +ATOM 473 CB LEU A 87 -3.239 -14.181 -20.447 1.00 26.90 C +ATOM 474 CG LEU A 87 -4.074 -13.744 -19.227 1.00 26.90 C +ATOM 475 CD1 LEU A 87 -4.477 -12.277 -19.390 1.00 26.90 C +ATOM 476 CD2 LEU A 87 -5.349 -14.573 -19.071 1.00 26.90 C +ATOM 477 N GLY A 88 -0.855 -15.704 -22.069 1.00 31.80 N +ATOM 478 CA GLY A 88 -0.244 -16.054 -23.354 1.00 31.80 C +ATOM 479 C GLY A 88 0.363 -17.459 -23.425 1.00 31.80 C +ATOM 480 O GLY A 88 0.735 -17.898 -24.510 1.00 31.80 O +ATOM 481 N ALA A 89 0.474 -18.175 -22.297 1.00 29.00 N +ATOM 482 CA ALA A 89 1.157 -19.472 -22.209 1.00 29.00 C +ATOM 483 C ALA A 89 2.614 -19.448 -22.740 1.00 29.00 C +ATOM 484 O ALA A 89 3.088 -20.421 -23.352 1.00 29.00 O +ATOM 485 CB ALA A 89 0.271 -20.555 -22.842 1.00 29.00 C +ATOM 486 N GLU A 90 3.314 -18.337 -22.485 1.00 27.25 N +ATOM 487 CA GLU A 90 4.705 -18.081 -22.878 1.00 27.25 C +ATOM 488 C GLU A 90 5.675 -18.292 -21.704 1.00 27.25 C +ATOM 489 O GLU A 90 5.345 -17.985 -20.560 1.00 27.25 O +ATOM 490 CB GLU A 90 4.865 -16.643 -23.401 1.00 27.25 C +ATOM 491 CG GLU A 90 4.085 -16.336 -24.691 1.00 27.25 C +ATOM 492 CD GLU A 90 4.425 -14.944 -25.252 1.00 27.25 C +ATOM 493 OE1 GLU A 90 4.107 -14.676 -26.430 1.00 27.25 O +ATOM 494 OE2 GLU A 90 5.021 -14.124 -24.507 1.00 27.25 O +ATOM 495 N GLU A 91 6.898 -18.760 -21.982 1.00 22.31 N +ATOM 496 CA GLU A 91 8.000 -18.618 -21.021 1.00 22.31 C +ATOM 497 C GLU A 91 8.437 -17.150 -21.006 1.00 22.31 C +ATOM 498 O GLU A 91 8.825 -16.607 -22.045 1.00 22.31 O +ATOM 499 CB GLU A 91 9.187 -19.545 -21.361 1.00 22.31 C +ATOM 500 CG GLU A 91 10.307 -19.471 -20.295 1.00 22.31 C +ATOM 501 CD GLU A 91 11.609 -20.206 -20.672 1.00 22.31 C +ATOM 502 OE1 GLU A 91 12.638 -19.980 -19.971 1.00 22.31 O +ATOM 503 OE2 GLU A 91 11.592 -20.978 -21.657 1.00 22.31 O +ATOM 504 N ARG A 92 8.391 -16.497 -19.842 1.00 21.22 N +ATOM 505 CA ARG A 92 8.813 -15.096 -19.718 1.00 21.22 C +ATOM 506 C ARG A 92 9.473 -14.835 -18.371 1.00 21.22 C +ATOM 507 O ARG A 92 8.929 -15.175 -17.323 1.00 21.22 O +ATOM 508 CB ARG A 92 7.611 -14.173 -19.991 1.00 21.22 C +ATOM 509 CG ARG A 92 8.042 -12.726 -20.280 1.00 21.22 C +ATOM 510 CD ARG A 92 6.878 -11.847 -20.761 1.00 21.22 C +ATOM 511 NE ARG A 92 6.485 -12.127 -22.166 1.00 21.22 N +ATOM 512 CZ ARG A 92 6.309 -11.250 -23.140 1.00 21.22 C +ATOM 513 NH1 ARG A 92 6.496 -9.969 -22.968 1.00 21.22 N +ATOM 514 NH2 ARG A 92 5.932 -11.629 -24.320 1.00 21.22 N +ATOM 515 N LYS A 93 10.637 -14.185 -18.421 1.00 19.23 N +ATOM 516 CA LYS A 93 11.375 -13.713 -17.247 1.00 19.23 C +ATOM 517 C LYS A 93 11.083 -12.226 -17.069 1.00 19.23 C +ATOM 518 O LYS A 93 11.313 -11.452 -17.995 1.00 19.23 O +ATOM 519 CB LYS A 93 12.876 -14.041 -17.360 1.00 19.23 C +ATOM 520 CG LYS A 93 13.087 -15.545 -17.606 1.00 19.23 C +ATOM 521 CD LYS A 93 14.547 -16.014 -17.517 1.00 19.23 C +ATOM 522 CE LYS A 93 14.760 -17.346 -18.270 1.00 19.23 C +ATOM 523 NZ LYS A 93 13.904 -18.466 -17.793 1.00 19.23 N +ATOM 524 N VAL A 94 10.540 -11.845 -15.916 1.00 19.27 N +ATOM 525 CA VAL A 94 10.149 -10.465 -15.598 1.00 19.27 C +ATOM 526 C VAL A 94 10.836 -10.042 -14.308 1.00 19.27 C +ATOM 527 O VAL A 94 10.693 -10.698 -13.279 1.00 19.27 O +ATOM 528 CB VAL A 94 8.619 -10.306 -15.479 1.00 19.27 C +ATOM 529 CG1 VAL A 94 8.247 -8.835 -15.254 1.00 19.27 C +ATOM 530 CG2 VAL A 94 7.887 -10.781 -16.742 1.00 19.27 C +ATOM 531 N GLU A 95 11.580 -8.944 -14.360 1.00 18.80 N +ATOM 532 CA GLU A 95 12.169 -8.326 -13.176 1.00 18.80 C +ATOM 533 C GLU A 95 11.154 -7.386 -12.515 1.00 18.80 C +ATOM 534 O GLU A 95 10.538 -6.540 -13.171 1.00 18.80 O +ATOM 535 CB GLU A 95 13.473 -7.631 -13.572 1.00 18.80 C +ATOM 536 CG GLU A 95 14.206 -7.051 -12.359 1.00 18.80 C +ATOM 537 CD GLU A 95 15.613 -6.578 -12.740 1.00 18.80 C +ATOM 538 OE1 GLU A 95 16.538 -6.853 -11.947 1.00 18.80 O +ATOM 539 OE2 GLU A 95 15.749 -5.957 -13.818 1.00 18.80 O +ATOM 540 N ILE A 96 10.939 -7.564 -11.213 1.00 19.13 N +ATOM 541 CA ILE A 96 9.904 -6.868 -10.454 1.00 19.13 C +ATOM 542 C ILE A 96 10.359 -6.640 -9.007 1.00 19.13 C +ATOM 543 O ILE A 96 10.980 -7.496 -8.378 1.00 19.13 O +ATOM 544 CB ILE A 96 8.568 -7.640 -10.581 1.00 19.13 C +ATOM 545 CG1 ILE A 96 7.374 -6.802 -10.084 1.00 19.13 C +ATOM 546 CG2 ILE A 96 8.617 -9.010 -9.889 1.00 19.13 C +ATOM 547 CD1 ILE A 96 6.016 -7.367 -10.526 1.00 19.13 C +ATOM 548 N ALA A 97 10.060 -5.460 -8.469 1.00 19.27 N +ATOM 549 CA ALA A 97 10.244 -5.181 -7.050 1.00 19.27 C +ATOM 550 C ALA A 97 9.041 -5.731 -6.272 1.00 19.27 C +ATOM 551 O ALA A 97 7.909 -5.295 -6.508 1.00 19.27 O +ATOM 552 CB ALA A 97 10.442 -3.675 -6.854 1.00 19.27 C +ATOM 553 N PHE A 98 9.296 -6.664 -5.357 1.00 17.94 N +ATOM 554 CA PHE A 98 8.301 -7.235 -4.449 1.00 17.94 C +ATOM 555 C PHE A 98 8.539 -6.786 -3.009 1.00 17.94 C +ATOM 556 O PHE A 98 9.636 -6.365 -2.635 1.00 17.94 O +ATOM 557 CB PHE A 98 8.312 -8.770 -4.534 1.00 17.94 C +ATOM 558 CG PHE A 98 7.666 -9.408 -5.751 1.00 17.94 C +ATOM 559 CD1 PHE A 98 6.447 -8.934 -6.283 1.00 17.94 C +ATOM 560 CD2 PHE A 98 8.247 -10.568 -6.291 1.00 17.94 C +ATOM 561 CE1 PHE A 98 5.818 -9.619 -7.337 1.00 17.94 C +ATOM 562 CE2 PHE A 98 7.620 -11.247 -7.346 1.00 17.94 C +ATOM 563 CZ PHE A 98 6.400 -10.781 -7.867 1.00 17.94 C +ATOM 564 N TYR A 99 7.487 -6.904 -2.205 1.00 18.67 N +ATOM 565 CA TYR A 99 7.477 -6.622 -0.778 1.00 18.67 C +ATOM 566 C TYR A 99 7.245 -7.908 0.014 1.00 18.67 C +ATOM 567 O TYR A 99 6.467 -8.779 -0.383 1.00 18.67 O +ATOM 568 CB TYR A 99 6.424 -5.551 -0.472 1.00 18.67 C +ATOM 569 CG TYR A 99 6.844 -4.164 -0.926 1.00 18.67 C +ATOM 570 CD1 TYR A 99 7.308 -3.233 0.023 1.00 18.67 C +ATOM 571 CD2 TYR A 99 6.807 -3.816 -2.291 1.00 18.67 C +ATOM 572 CE1 TYR A 99 7.721 -1.954 -0.384 1.00 18.67 C +ATOM 573 CE2 TYR A 99 7.244 -2.545 -2.711 1.00 18.67 C +ATOM 574 CZ TYR A 99 7.693 -1.611 -1.753 1.00 18.67 C +ATOM 575 OH TYR A 99 8.070 -0.364 -2.137 1.00 18.67 O +ATOM 576 N ARG A 100 7.943 -8.024 1.141 1.00 18.80 N +ATOM 577 CA ARG A 100 7.717 -9.057 2.154 1.00 18.80 C +ATOM 578 C ARG A 100 6.641 -8.638 3.156 1.00 18.80 C +ATOM 579 O ARG A 100 6.328 -7.455 3.285 1.00 18.80 O +ATOM 580 CB ARG A 100 9.020 -9.288 2.915 1.00 18.80 C +ATOM 581 CG ARG A 100 10.148 -9.907 2.102 1.00 18.80 C +ATOM 582 CD ARG A 100 11.186 -10.478 3.069 1.00 18.80 C +ATOM 583 NE ARG A 100 12.375 -10.933 2.338 1.00 18.80 N +ATOM 584 CZ ARG A 100 13.589 -10.424 2.412 1.00 18.80 C +ATOM 585 NH1 ARG A 100 13.868 -9.357 3.106 1.00 18.80 N +ATOM 586 NH2 ARG A 100 14.564 -10.997 1.766 1.00 18.80 N +ATOM 587 N LYS A 101 6.175 -9.603 3.951 1.00 19.87 N +ATOM 588 CA LYS A 101 5.286 -9.394 5.107 1.00 19.87 C +ATOM 589 C LYS A 101 5.878 -8.474 6.196 1.00 19.87 C +ATOM 590 O LYS A 101 5.120 -7.850 6.925 1.00 19.87 O +ATOM 591 CB LYS A 101 4.889 -10.780 5.652 1.00 19.87 C +ATOM 592 CG LYS A 101 3.738 -10.703 6.666 1.00 19.87 C +ATOM 593 CD LYS A 101 3.222 -12.096 7.046 1.00 19.87 C +ATOM 594 CE LYS A 101 2.061 -11.963 8.036 1.00 19.87 C +ATOM 595 NZ LYS A 101 1.393 -13.265 8.265 1.00 19.87 N +ATOM 596 N ASP A 102 7.206 -8.335 6.277 1.00 20.89 N +ATOM 597 CA ASP A 102 7.896 -7.368 7.157 1.00 20.89 C +ATOM 598 C ASP A 102 8.009 -5.945 6.557 1.00 20.89 C +ATOM 599 O ASP A 102 8.635 -5.064 7.144 1.00 20.89 O +ATOM 600 CB ASP A 102 9.270 -7.929 7.595 1.00 20.89 C +ATOM 601 CG ASP A 102 10.334 -7.964 6.487 1.00 20.89 C +ATOM 602 OD1 ASP A 102 10.018 -7.603 5.334 1.00 20.89 O +ATOM 603 OD2 ASP A 102 11.478 -8.393 6.731 1.00 20.89 O +ATOM 604 N GLY A 103 7.440 -5.712 5.368 1.00 24.81 N +ATOM 605 CA GLY A 103 7.524 -4.449 4.632 1.00 24.81 C +ATOM 606 C GLY A 103 8.831 -4.236 3.856 1.00 24.81 C +ATOM 607 O GLY A 103 8.942 -3.249 3.122 1.00 24.81 O +ATOM 608 N SER A 104 9.816 -5.137 3.954 1.00 20.72 N +ATOM 609 CA SER A 104 11.075 -5.007 3.213 1.00 20.72 C +ATOM 610 C SER A 104 10.876 -5.236 1.708 1.00 20.72 C +ATOM 611 O SER A 104 10.244 -6.197 1.265 1.00 20.72 O +ATOM 612 CB SER A 104 12.200 -5.870 3.804 1.00 20.72 C +ATOM 613 OG SER A 104 11.961 -7.254 3.678 1.00 20.72 O +ATOM 614 N CYS A 105 11.425 -4.323 0.905 1.00 20.94 N +ATOM 615 CA CYS A 105 11.381 -4.377 -0.554 1.00 20.94 C +ATOM 616 C CYS A 105 12.624 -5.094 -1.094 1.00 20.94 C +ATOM 617 O CYS A 105 13.745 -4.729 -0.742 1.00 20.94 O +ATOM 618 CB CYS A 105 11.283 -2.942 -1.088 1.00 20.94 C +ATOM 619 SG CYS A 105 11.139 -2.926 -2.902 1.00 20.94 S +ATOM 620 N PHE A 106 12.439 -6.050 -2.005 1.00 19.92 N +ATOM 621 CA PHE A 106 13.524 -6.767 -2.677 1.00 19.92 C +ATOM 622 C PHE A 106 13.296 -6.832 -4.192 1.00 19.92 C +ATOM 623 O PHE A 106 12.159 -6.837 -4.669 1.00 19.92 O +ATOM 624 CB PHE A 106 13.716 -8.156 -2.043 1.00 19.92 C +ATOM 625 CG PHE A 106 12.544 -9.111 -2.182 1.00 19.92 C +ATOM 626 CD1 PHE A 106 11.428 -8.995 -1.329 1.00 19.92 C +ATOM 627 CD2 PHE A 106 12.569 -10.126 -3.160 1.00 19.92 C +ATOM 628 CE1 PHE A 106 10.343 -9.880 -1.458 1.00 19.92 C +ATOM 629 CE2 PHE A 106 11.489 -11.020 -3.278 1.00 19.92 C +ATOM 630 CZ PHE A 106 10.376 -10.897 -2.427 1.00 19.92 C +ATOM 631 N LEU A 107 14.388 -6.854 -4.959 1.00 19.97 N +ATOM 632 CA LEU A 107 14.353 -6.979 -6.415 1.00 19.97 C +ATOM 633 C LEU A 107 14.377 -8.463 -6.792 1.00 19.97 C +ATOM 634 O LEU A 107 15.264 -9.204 -6.361 1.00 19.97 O +ATOM 635 CB LEU A 107 15.523 -6.176 -7.011 1.00 19.97 C +ATOM 636 CG LEU A 107 15.481 -5.998 -8.541 1.00 19.97 C +ATOM 637 CD1 LEU A 107 14.219 -5.257 -8.999 1.00 19.97 C +ATOM 638 CD2 LEU A 107 16.698 -5.178 -8.963 1.00 19.97 C +ATOM 639 N CYS A 108 13.382 -8.907 -7.556 1.00 19.92 N +ATOM 640 CA CYS A 108 13.180 -10.314 -7.867 1.00 19.92 C +ATOM 641 C CYS A 108 12.988 -10.536 -9.368 1.00 19.92 C +ATOM 642 O CYS A 108 12.202 -9.849 -10.020 1.00 19.92 O +ATOM 643 CB CYS A 108 11.992 -10.823 -7.044 1.00 19.92 C +ATOM 644 SG CYS A 108 11.917 -12.635 -7.117 1.00 19.92 S +ATOM 645 N LEU A 109 13.683 -11.536 -9.904 1.00 18.72 N +ATOM 646 CA LEU A 109 13.433 -12.081 -11.226 1.00 18.72 C +ATOM 647 C LEU A 109 12.399 -13.206 -11.100 1.00 18.72 C +ATOM 648 O LEU A 109 12.670 -14.241 -10.487 1.00 18.72 O +ATOM 649 CB LEU A 109 14.771 -12.548 -11.820 1.00 18.72 C +ATOM 650 CG LEU A 109 14.678 -12.979 -13.294 1.00 18.72 C +ATOM 651 CD1 LEU A 109 14.375 -11.791 -14.212 1.00 18.72 C +ATOM 652 CD2 LEU A 109 16.006 -13.595 -13.735 1.00 18.72 C +ATOM 653 N VAL A 110 11.221 -12.995 -11.685 1.00 18.63 N +ATOM 654 CA VAL A 110 10.156 -13.998 -11.789 1.00 18.63 C +ATOM 655 C VAL A 110 10.303 -14.725 -13.114 1.00 18.63 C +ATOM 656 O VAL A 110 10.166 -14.128 -14.180 1.00 18.63 O +ATOM 657 CB VAL A 110 8.755 -13.372 -11.668 1.00 18.63 C +ATOM 658 CG1 VAL A 110 7.655 -14.443 -11.705 1.00 18.63 C +ATOM 659 CG2 VAL A 110 8.617 -12.636 -10.338 1.00 18.63 C +ATOM 660 N ASP A 111 10.576 -16.021 -13.038 1.00 18.90 N +ATOM 661 CA ASP A 111 10.699 -16.923 -14.174 1.00 18.90 C +ATOM 662 C ASP A 111 9.396 -17.729 -14.323 1.00 18.90 C +ATOM 663 O ASP A 111 9.165 -18.684 -13.576 1.00 18.90 O +ATOM 664 CB ASP A 111 11.954 -17.788 -13.949 1.00 18.90 C +ATOM 665 CG ASP A 111 12.380 -18.580 -15.182 1.00 18.90 C +ATOM 666 OD1 ASP A 111 11.780 -18.373 -16.257 1.00 18.90 O +ATOM 667 OD2 ASP A 111 13.419 -19.283 -15.137 1.00 18.90 O +ATOM 668 N VAL A 112 8.513 -17.293 -15.231 1.00 19.13 N +ATOM 669 CA VAL A 112 7.205 -17.920 -15.486 1.00 19.13 C +ATOM 670 C VAL A 112 7.372 -19.033 -16.519 1.00 19.13 C +ATOM 671 O VAL A 112 7.670 -18.763 -17.682 1.00 19.13 O +ATOM 672 CB VAL A 112 6.148 -16.897 -15.964 1.00 19.13 C +ATOM 673 CG1 VAL A 112 4.765 -17.553 -16.107 1.00 19.13 C +ATOM 674 CG2 VAL A 112 6.000 -15.716 -14.995 1.00 19.13 C +ATOM 675 N VAL A 113 7.146 -20.283 -16.104 1.00 20.61 N +ATOM 676 CA VAL A 113 7.329 -21.484 -16.931 1.00 20.61 C +ATOM 677 C VAL A 113 6.004 -22.258 -17.074 1.00 20.61 C +ATOM 678 O VAL A 113 5.564 -22.906 -16.117 1.00 20.61 O +ATOM 679 CB VAL A 113 8.452 -22.369 -16.352 1.00 20.61 C +ATOM 680 CG1 VAL A 113 8.599 -23.668 -17.155 1.00 20.61 C +ATOM 681 CG2 VAL A 113 9.805 -21.641 -16.378 1.00 20.61 C +ATOM 682 N PRO A 114 5.349 -22.236 -18.252 1.00 20.99 N +ATOM 683 CA PRO A 114 4.130 -23.004 -18.511 1.00 20.99 C +ATOM 684 C PRO A 114 4.428 -24.486 -18.797 1.00 20.99 C +ATOM 685 O PRO A 114 5.162 -24.826 -19.725 1.00 20.99 O +ATOM 686 CB PRO A 114 3.471 -22.312 -19.708 1.00 20.99 C +ATOM 687 CG PRO A 114 4.651 -21.723 -20.479 1.00 20.99 C +ATOM 688 CD PRO A 114 5.656 -21.370 -19.386 1.00 20.99 C +ATOM 689 N VAL A 115 3.797 -25.389 -18.044 1.00 26.22 N +ATOM 690 CA VAL A 115 3.892 -26.846 -18.221 1.00 26.22 C +ATOM 691 C VAL A 115 2.776 -27.324 -19.152 1.00 26.22 C +ATOM 692 O VAL A 115 1.600 -27.337 -18.777 1.00 26.22 O +ATOM 693 CB VAL A 115 3.856 -27.581 -16.866 1.00 26.22 C +ATOM 694 CG1 VAL A 115 4.040 -29.093 -17.052 1.00 26.22 C +ATOM 695 CG2 VAL A 115 4.967 -27.092 -15.927 1.00 26.22 C +ATOM 696 N LYS A 116 3.153 -27.714 -20.374 1.00 29.30 N +ATOM 697 CA LYS A 116 2.256 -28.202 -21.437 1.00 29.30 C +ATOM 698 C LYS A 116 2.207 -29.728 -21.457 1.00 29.30 C +ATOM 699 O LYS A 116 3.222 -30.367 -21.203 1.00 29.30 O +ATOM 700 CB LYS A 116 2.681 -27.620 -22.802 1.00 29.30 C +ATOM 701 CG LYS A 116 2.453 -26.099 -22.820 1.00 29.30 C +ATOM 702 CD LYS A 116 2.961 -25.365 -24.071 1.00 29.30 C +ATOM 703 CE LYS A 116 2.734 -23.865 -23.801 1.00 29.30 C +ATOM 704 NZ LYS A 116 3.216 -22.956 -24.872 1.00 29.30 N +ATOM 705 N ASN A 117 1.035 -30.295 -21.740 1.00 38.95 N +ATOM 706 CA ASN A 117 0.839 -31.736 -21.883 1.00 38.95 C +ATOM 707 C ASN A 117 1.174 -32.237 -23.300 1.00 38.95 C +ATOM 708 O ASN A 117 1.541 -31.450 -24.170 1.00 38.95 O +ATOM 709 CB ASN A 117 -0.581 -32.100 -21.408 1.00 38.95 C +ATOM 710 CG ASN A 117 -1.681 -31.959 -22.441 1.00 38.95 C +ATOM 711 OD1 ASN A 117 -1.671 -31.100 -23.303 1.00 38.95 O +ATOM 712 ND2 ASN A 117 -2.661 -32.824 -22.358 1.00 38.95 N +ATOM 713 N GLU A 118 1.007 -33.541 -23.539 1.00 66.30 N +ATOM 714 CA GLU A 118 1.215 -34.151 -24.863 1.00 66.30 C +ATOM 715 C GLU A 118 0.344 -33.529 -25.973 1.00 66.30 C +ATOM 716 O GLU A 118 0.789 -33.430 -27.113 1.00 66.30 O +ATOM 717 CB GLU A 118 0.891 -35.645 -24.775 1.00 66.30 C +ATOM 718 CG GLU A 118 1.871 -36.465 -23.919 1.00 66.30 C +ATOM 719 CD GLU A 118 1.388 -37.915 -23.780 1.00 66.30 C +ATOM 720 OE1 GLU A 118 1.960 -38.702 -22.991 1.00 66.30 O +ATOM 721 OE2 GLU A 118 0.455 -38.332 -24.511 1.00 66.30 O +ATOM 722 N ASP A 119 -0.864 -33.056 -25.640 1.00 49.00 N +ATOM 723 CA ASP A 119 -1.767 -32.352 -26.567 1.00 49.00 C +ATOM 724 C ASP A 119 -1.372 -30.870 -26.786 1.00 49.00 C +ATOM 725 O ASP A 119 -2.072 -30.124 -27.470 1.00 49.00 O +ATOM 726 CB ASP A 119 -3.229 -32.438 -26.082 1.00 49.00 C +ATOM 727 CG ASP A 119 -3.683 -33.822 -25.602 1.00 49.00 C +ATOM 728 OD1 ASP A 119 -3.682 -34.775 -26.408 1.00 49.00 O +ATOM 729 OD2 ASP A 119 -4.059 -33.902 -24.406 1.00 49.00 O +ATOM 730 N GLY A 120 -0.276 -30.405 -26.171 1.00 42.68 N +ATOM 731 CA GLY A 120 0.210 -29.022 -26.217 1.00 42.68 C +ATOM 732 C GLY A 120 -0.508 -28.040 -25.278 1.00 42.68 C +ATOM 733 O GLY A 120 -0.100 -26.879 -25.181 1.00 42.68 O +ATOM 734 N ALA A 121 -1.541 -28.478 -24.555 1.00 31.46 N +ATOM 735 CA ALA A 121 -2.323 -27.642 -23.647 1.00 31.46 C +ATOM 736 C ALA A 121 -1.610 -27.441 -22.300 1.00 31.46 C +ATOM 737 O ALA A 121 -1.093 -28.389 -21.705 1.00 31.46 O +ATOM 738 CB ALA A 121 -3.714 -28.266 -23.473 1.00 31.46 C +ATOM 739 N VAL A 122 -1.611 -26.212 -21.774 1.00 27.69 N +ATOM 740 CA VAL A 122 -1.051 -25.923 -20.442 1.00 27.69 C +ATOM 741 C VAL A 122 -1.907 -26.591 -19.364 1.00 27.69 C +ATOM 742 O VAL A 122 -3.127 -26.430 -19.338 1.00 27.69 O +ATOM 743 CB VAL A 122 -0.914 -24.411 -20.172 1.00 27.69 C +ATOM 744 CG1 VAL A 122 -0.179 -24.139 -18.853 1.00 27.69 C +ATOM 745 CG2 VAL A 122 -0.102 -23.706 -21.262 1.00 27.69 C +ATOM 746 N ILE A 123 -1.274 -27.337 -18.455 1.00 26.65 N +ATOM 747 CA ILE A 123 -1.929 -27.827 -17.229 1.00 26.65 C +ATOM 748 C ILE A 123 -1.638 -26.882 -16.060 1.00 26.65 C +ATOM 749 O ILE A 123 -2.513 -26.631 -15.234 1.00 26.65 O +ATOM 750 CB ILE A 123 -1.507 -29.272 -16.859 1.00 26.65 C +ATOM 751 CG1 ILE A 123 -1.402 -30.225 -18.072 1.00 26.65 C +ATOM 752 CG2 ILE A 123 -2.510 -29.836 -15.828 1.00 26.65 C +ATOM 753 CD1 ILE A 123 -1.053 -31.668 -17.667 1.00 26.65 C +ATOM 754 N MET A 124 -0.401 -26.390 -15.958 1.00 22.06 N +ATOM 755 CA MET A 124 0.113 -25.695 -14.775 1.00 22.06 C +ATOM 756 C MET A 124 1.145 -24.631 -15.145 1.00 22.06 C +ATOM 757 O MET A 124 1.802 -24.737 -16.180 1.00 22.06 O +ATOM 758 CB MET A 124 0.771 -26.713 -13.826 1.00 22.06 C +ATOM 759 CG MET A 124 -0.227 -27.689 -13.200 1.00 22.06 C +ATOM 760 SD MET A 124 0.550 -29.124 -12.439 1.00 22.06 S +ATOM 761 CE MET A 124 -0.929 -30.083 -12.028 1.00 22.06 C +ATOM 762 N PHE A 125 1.348 -23.666 -14.254 1.00 20.23 N +ATOM 763 CA PHE A 125 2.483 -22.747 -14.288 1.00 20.23 C +ATOM 764 C PHE A 125 3.406 -23.027 -13.104 1.00 20.23 C +ATOM 765 O PHE A 125 2.938 -23.192 -11.978 1.00 20.23 O +ATOM 766 CB PHE A 125 1.998 -21.293 -14.286 1.00 20.23 C +ATOM 767 CG PHE A 125 1.114 -20.943 -15.465 1.00 20.23 C +ATOM 768 CD1 PHE A 125 1.694 -20.561 -16.690 1.00 20.23 C +ATOM 769 CD2 PHE A 125 -0.287 -20.989 -15.333 1.00 20.23 C +ATOM 770 CE1 PHE A 125 0.870 -20.235 -17.782 1.00 20.23 C +ATOM 771 CE2 PHE A 125 -1.106 -20.668 -16.429 1.00 20.23 C +ATOM 772 CZ PHE A 125 -0.527 -20.311 -17.656 1.00 20.23 C +ATOM 773 N ILE A 126 4.711 -23.060 -13.359 1.00 19.42 N +ATOM 774 CA ILE A 126 5.752 -22.996 -12.333 1.00 19.42 C +ATOM 775 C ILE A 126 6.285 -21.566 -12.364 1.00 19.42 C +ATOM 776 O ILE A 126 6.754 -21.118 -13.409 1.00 19.42 O +ATOM 777 CB ILE A 126 6.871 -24.030 -12.601 1.00 19.42 C +ATOM 778 CG1 ILE A 126 6.318 -25.473 -12.651 1.00 19.42 C +ATOM 779 CG2 ILE A 126 7.969 -23.909 -11.524 1.00 19.42 C +ATOM 780 CD1 ILE A 126 7.330 -26.486 -13.203 1.00 19.42 C +ATOM 781 N LEU A 127 6.216 -20.849 -11.243 1.00 18.80 N +ATOM 782 CA LEU A 127 6.850 -19.539 -11.105 1.00 18.80 C +ATOM 783 C LEU A 127 8.037 -19.698 -10.157 1.00 18.80 C +ATOM 784 O LEU A 127 7.859 -20.092 -9.004 1.00 18.80 O +ATOM 785 CB LEU A 127 5.865 -18.454 -10.635 1.00 18.80 C +ATOM 786 CG LEU A 127 4.738 -18.082 -11.619 1.00 18.80 C +ATOM 787 CD1 LEU A 127 3.571 -19.077 -11.619 1.00 18.80 C +ATOM 788 CD2 LEU A 127 4.164 -16.720 -11.229 1.00 18.80 C +ATOM 789 N ASN A 128 9.243 -19.435 -10.657 1.00 18.80 N +ATOM 790 CA ASN A 128 10.493 -19.512 -9.900 1.00 18.80 C +ATOM 791 C ASN A 128 11.020 -18.097 -9.621 1.00 18.80 C +ATOM 792 O ASN A 128 11.064 -17.265 -10.523 1.00 18.80 O +ATOM 793 CB ASN A 128 11.460 -20.425 -10.673 1.00 18.80 C +ATOM 794 CG ASN A 128 12.876 -20.437 -10.126 1.00 18.80 C +ATOM 795 OD1 ASN A 128 13.257 -21.302 -9.355 1.00 18.80 O +ATOM 796 ND2 ASN A 128 13.722 -19.540 -10.572 1.00 18.80 N +ATOM 797 N PHE A 129 11.425 -17.841 -8.380 1.00 18.94 N +ATOM 798 CA PHE A 129 11.760 -16.520 -7.852 1.00 18.94 C +ATOM 799 C PHE A 129 13.244 -16.465 -7.474 1.00 18.94 C +ATOM 800 O PHE A 129 13.690 -17.166 -6.565 1.00 18.94 O +ATOM 801 CB PHE A 129 10.863 -16.229 -6.636 1.00 18.94 C +ATOM 802 CG PHE A 129 9.373 -16.202 -6.931 1.00 18.94 C +ATOM 803 CD1 PHE A 129 8.722 -14.989 -7.217 1.00 18.94 C +ATOM 804 CD2 PHE A 129 8.633 -17.398 -6.917 1.00 18.94 C +ATOM 805 CE1 PHE A 129 7.343 -14.975 -7.503 1.00 18.94 C +ATOM 806 CE2 PHE A 129 7.259 -17.385 -7.206 1.00 18.94 C +ATOM 807 CZ PHE A 129 6.611 -16.175 -7.505 1.00 18.94 C +ATOM 808 N GLU A 130 14.013 -15.624 -8.160 1.00 20.28 N +ATOM 809 CA GLU A 130 15.443 -15.417 -7.901 1.00 20.28 C +ATOM 810 C GLU A 130 15.685 -13.970 -7.445 1.00 20.28 C +ATOM 811 O GLU A 130 15.047 -13.047 -7.949 1.00 20.28 O +ATOM 812 CB GLU A 130 16.237 -15.845 -9.149 1.00 20.28 C +ATOM 813 CG GLU A 130 17.753 -15.961 -8.911 1.00 20.28 C +ATOM 814 CD GLU A 130 18.501 -16.717 -10.031 1.00 20.28 C +ATOM 815 OE1 GLU A 130 19.741 -16.846 -9.908 1.00 20.28 O +ATOM 816 OE2 GLU A 130 17.850 -17.217 -10.980 1.00 20.28 O +ATOM 817 N VAL A 131 16.548 -13.750 -6.448 1.00 24.28 N +ATOM 818 CA VAL A 131 16.818 -12.403 -5.907 1.00 24.28 C +ATOM 819 C VAL A 131 17.979 -11.776 -6.666 1.00 24.28 C +ATOM 820 O VAL A 131 19.063 -12.356 -6.735 1.00 24.28 O +ATOM 821 CB VAL A 131 17.085 -12.419 -4.390 1.00 24.28 C +ATOM 822 CG1 VAL A 131 17.357 -11.011 -3.840 1.00 24.28 C +ATOM 823 CG2 VAL A 131 15.871 -12.974 -3.634 1.00 24.28 C +ATOM 824 N VAL A 132 17.764 -10.582 -7.214 1.00 33.17 N +ATOM 825 CA VAL A 132 18.781 -9.868 -7.990 1.00 33.17 C +ATOM 826 C VAL A 132 19.683 -9.084 -7.037 1.00 33.17 C +ATOM 827 O VAL A 132 19.245 -8.150 -6.368 1.00 33.17 O +ATOM 828 CB VAL A 132 18.144 -8.969 -9.066 1.00 33.17 C +ATOM 829 CG1 VAL A 132 19.226 -8.299 -9.924 1.00 33.17 C +ATOM 830 CG2 VAL A 132 17.235 -9.790 -9.993 1.00 33.17 C +END diff --git a/tests/files/pdb/1VER_af.pdb b/tests/files/pdb/1VER_af.pdb new file mode 100644 index 00000000..44a7cad3 --- /dev/null +++ b/tests/files/pdb/1VER_af.pdb @@ -0,0 +1,772 @@ +REMARK TITLE 1VER AlphaFold-start MR +REMARK Log-Likelihood Gain: 535.749 +REMARK RFZ=3.2 TFZ=12.1 PAK=1 LLG=158 TFZ==12.7 LLG=536 TFZ==21.3 PAK=0 LLG=536 TFZ==21.3 +REMARK ENSEMBLE e_Q6X1E7 EULER 235.26 46.67 332.21 FRAC 0.026 0.125 -0.287 +CRYST1 97.259 97.259 65.228 90.00 90.00 90.00 I 41 2 2 16 +SCALE1 0.010282 -0.000000 -0.000000 0.00000 +SCALE2 0.000000 0.010282 -0.000000 0.00000 +SCALE3 0.000000 0.000000 0.015331 0.00000 +ATOM 1 N ALA A 1 -6.456 8.196 -34.339 1.00 60.35 N +ATOM 2 CA ALA A 1 -5.443 8.967 -33.602 1.00 60.35 C +ATOM 3 C ALA A 1 -5.050 8.199 -32.353 1.00 60.35 C +ATOM 4 O ALA A 1 -5.871 7.446 -31.853 1.00 60.35 O +ATOM 5 CB ALA A 1 -5.982 10.347 -33.228 1.00 60.35 C +ATOM 6 N TRP A 2 -3.816 8.346 -31.883 1.00 55.52 N +ATOM 7 CA TRP A 2 -3.359 7.759 -30.620 1.00 55.52 C +ATOM 8 C TRP A 2 -2.203 8.583 -30.044 1.00 55.52 C +ATOM 9 O TRP A 2 -1.570 9.360 -30.766 1.00 55.52 O +ATOM 10 CB TRP A 2 -2.938 6.297 -30.836 1.00 55.52 C +ATOM 11 CG TRP A 2 -1.742 6.093 -31.714 1.00 55.52 C +ATOM 12 CD1 TRP A 2 -1.747 6.062 -33.065 1.00 55.52 C +ATOM 13 CD2 TRP A 2 -0.358 5.854 -31.322 1.00 55.52 C +ATOM 14 NE1 TRP A 2 -0.470 5.857 -33.541 1.00 55.52 N +ATOM 15 CE2 TRP A 2 0.426 5.689 -32.505 1.00 55.52 C +ATOM 16 CE3 TRP A 2 0.295 5.692 -30.082 1.00 55.52 C +ATOM 17 CZ2 TRP A 2 1.795 5.390 -32.454 1.00 55.52 C +ATOM 18 CZ3 TRP A 2 1.661 5.357 -30.030 1.00 55.52 C +ATOM 19 CH2 TRP A 2 2.420 5.242 -31.206 1.00 55.52 C +ATOM 20 N VAL A 3 -1.925 8.434 -28.747 1.00 53.37 N +ATOM 21 CA VAL A 3 -0.815 9.141 -28.091 1.00 53.37 C +ATOM 22 C VAL A 3 0.378 8.216 -27.880 1.00 53.37 C +ATOM 23 O VAL A 3 0.337 7.265 -27.090 1.00 53.37 O +ATOM 24 CB VAL A 3 -1.253 9.834 -26.794 1.00 53.37 C +ATOM 25 CG1 VAL A 3 -0.068 10.530 -26.121 1.00 53.37 C +ATOM 26 CG2 VAL A 3 -2.290 10.919 -27.086 1.00 53.37 C +ATOM 27 N ASP A 4 1.466 8.548 -28.563 1.00 52.47 N +ATOM 28 CA ASP A 4 2.763 7.913 -28.400 1.00 52.47 C +ATOM 29 C ASP A 4 3.482 8.476 -27.186 1.00 52.47 C +ATOM 30 O ASP A 4 3.983 9.596 -27.221 1.00 52.47 O +ATOM 31 CB ASP A 4 3.585 8.074 -29.679 1.00 52.47 C +ATOM 32 CG ASP A 4 4.859 7.228 -29.660 1.00 52.47 C +ATOM 33 OD1 ASP A 4 4.966 6.321 -28.801 1.00 52.47 O +ATOM 34 OD2 ASP A 4 5.697 7.500 -30.553 1.00 52.47 O +ATOM 35 N GLN A 5 3.480 7.702 -26.101 1.00 52.12 N +ATOM 36 CA GLN A 5 4.039 8.091 -24.814 1.00 52.12 C +ATOM 37 C GLN A 5 5.332 7.326 -24.529 1.00 52.12 C +ATOM 38 O GLN A 5 5.349 6.093 -24.529 1.00 52.12 O +ATOM 39 CB GLN A 5 2.988 7.921 -23.708 1.00 52.12 C +ATOM 40 CG GLN A 5 3.567 8.283 -22.332 1.00 52.12 C +ATOM 41 CD GLN A 5 2.541 8.260 -21.209 1.00 52.12 C +ATOM 42 OE1 GLN A 5 1.334 8.358 -21.384 1.00 52.12 O +ATOM 43 NE2 GLN A 5 2.995 8.080 -19.996 1.00 52.12 N +ATOM 44 N THR A 6 6.392 8.067 -24.209 1.00 52.24 N +ATOM 45 CA THR A 6 7.719 7.535 -23.888 1.00 52.24 C +ATOM 46 C THR A 6 8.328 8.242 -22.672 1.00 52.24 C +ATOM 47 O THR A 6 8.118 9.439 -22.490 1.00 52.24 O +ATOM 48 CB THR A 6 8.680 7.646 -25.082 1.00 52.24 C +ATOM 49 OG1 THR A 6 8.770 8.974 -25.534 1.00 52.24 O +ATOM 50 CG2 THR A 6 8.240 6.788 -26.267 1.00 52.24 C +ATOM 51 N PRO A 7 9.104 7.541 -21.829 1.00 52.09 N +ATOM 52 CA PRO A 7 9.375 6.101 -21.869 1.00 52.09 C +ATOM 53 C PRO A 7 8.179 5.268 -21.372 1.00 52.09 C +ATOM 54 O PRO A 7 7.255 5.806 -20.773 1.00 52.09 O +ATOM 55 CB PRO A 7 10.600 5.935 -20.969 1.00 52.09 C +ATOM 56 CG PRO A 7 10.408 7.019 -19.906 1.00 52.09 C +ATOM 57 CD PRO A 7 9.763 8.161 -20.686 1.00 52.09 C +ATOM 58 N ARG A 8 8.179 3.946 -21.594 1.00 52.34 N +ATOM 59 CA ARG A 8 7.172 3.030 -21.005 1.00 52.34 C +ATOM 60 C ARG A 8 7.489 2.648 -19.562 1.00 52.34 C +ATOM 61 O ARG A 8 6.577 2.471 -18.762 1.00 52.34 O +ATOM 62 CB ARG A 8 7.024 1.762 -21.863 1.00 52.34 C +ATOM 63 CG ARG A 8 6.457 2.020 -23.266 1.00 52.34 C +ATOM 64 CD ARG A 8 5.070 2.664 -23.202 1.00 52.34 C +ATOM 65 NE ARG A 8 4.481 2.825 -24.542 1.00 52.34 N +ATOM 66 CZ ARG A 8 3.429 3.572 -24.809 1.00 52.34 C +ATOM 67 NH1 ARG A 8 2.802 4.203 -23.857 1.00 52.34 N +ATOM 68 NH2 ARG A 8 2.994 3.700 -26.033 1.00 52.34 N +ATOM 69 N THR A 9 8.775 2.611 -19.232 1.00 52.37 N +ATOM 70 CA THR A 9 9.287 2.323 -17.893 1.00 52.37 C +ATOM 71 C THR A 9 10.473 3.236 -17.620 1.00 52.37 C +ATOM 72 O THR A 9 11.280 3.481 -18.517 1.00 52.37 O +ATOM 73 CB THR A 9 9.711 0.853 -17.758 1.00 52.37 C +ATOM 74 OG1 THR A 9 8.664 0.013 -18.185 1.00 52.37 O +ATOM 75 CG2 THR A 9 10.033 0.463 -16.316 1.00 52.37 C +ATOM 76 N ALA A 10 10.585 3.738 -16.398 1.00 52.34 N +ATOM 77 CA ALA A 10 11.725 4.509 -15.938 1.00 52.34 C +ATOM 78 C ALA A 10 12.048 4.149 -14.487 1.00 52.34 C +ATOM 79 O ALA A 10 11.157 4.094 -13.640 1.00 52.34 O +ATOM 80 CB ALA A 10 11.418 6.002 -16.098 1.00 52.34 C +ATOM 81 N THR A 11 13.335 3.963 -14.207 1.00 52.21 N +ATOM 82 CA THR A 11 13.867 3.873 -12.846 1.00 52.21 C +ATOM 83 C THR A 11 14.741 5.096 -12.618 1.00 52.21 C +ATOM 84 O THR A 11 15.572 5.411 -13.475 1.00 52.21 O +ATOM 85 CB THR A 11 14.651 2.577 -12.621 1.00 52.21 C +ATOM 86 OG1 THR A 11 13.846 1.466 -12.950 1.00 52.21 O +ATOM 87 CG2 THR A 11 15.069 2.397 -11.162 1.00 52.21 C +ATOM 88 N LYS A 12 14.507 5.809 -11.521 1.00 52.24 N +ATOM 89 CA LYS A 12 15.185 7.063 -11.178 1.00 52.24 C +ATOM 90 C LYS A 12 15.579 7.076 -9.710 1.00 52.24 C +ATOM 91 O LYS A 12 14.946 6.396 -8.910 1.00 52.24 O +ATOM 92 CB LYS A 12 14.263 8.252 -11.518 1.00 52.24 C +ATOM 93 CG LYS A 12 14.052 8.485 -13.028 1.00 52.24 C +ATOM 94 CD LYS A 12 15.367 8.951 -13.651 1.00 52.24 C +ATOM 95 CE LYS A 12 15.367 9.158 -15.156 1.00 52.24 C +ATOM 96 NZ LYS A 12 16.764 9.466 -15.539 1.00 52.24 N +ATOM 97 N GLU A 13 16.593 7.857 -9.375 1.00 52.15 N +ATOM 98 CA GLU A 13 16.986 8.085 -7.982 1.00 52.15 C +ATOM 99 C GLU A 13 16.277 9.328 -7.423 1.00 52.15 C +ATOM 100 O GLU A 13 15.869 10.233 -8.165 1.00 52.15 O +ATOM 101 CB GLU A 13 18.516 8.186 -7.864 1.00 52.15 C +ATOM 102 CG GLU A 13 19.276 6.951 -8.390 1.00 52.15 C +ATOM 103 CD GLU A 13 18.927 5.639 -7.664 1.00 52.15 C +ATOM 104 OE1 GLU A 13 18.813 4.588 -8.346 1.00 52.15 O +ATOM 105 OE2 GLU A 13 18.782 5.677 -6.426 1.00 52.15 O +ATOM 106 N THR A 14 16.115 9.381 -6.105 1.00 52.30 N +ATOM 107 CA THR A 14 15.580 10.557 -5.412 1.00 52.30 C +ATOM 108 C THR A 14 16.392 11.821 -5.744 1.00 52.30 C +ATOM 109 O THR A 14 17.620 11.808 -5.800 1.00 52.30 O +ATOM 110 CB THR A 14 15.510 10.291 -3.904 1.00 52.30 C +ATOM 111 OG1 THR A 14 14.575 9.258 -3.665 1.00 52.30 O +ATOM 112 CG2 THR A 14 15.028 11.501 -3.097 1.00 52.30 C +ATOM 113 N GLY A 15 15.699 12.933 -5.993 1.00 52.56 N +ATOM 114 CA GLY A 15 16.276 14.211 -6.422 1.00 52.56 C +ATOM 115 C GLY A 15 16.489 14.337 -7.937 1.00 52.56 C +ATOM 116 O GLY A 15 16.622 15.457 -8.441 1.00 52.56 O +ATOM 117 N GLU A 16 16.458 13.236 -8.699 1.00 52.12 N +ATOM 118 CA GLU A 16 16.564 13.293 -10.159 1.00 52.12 C +ATOM 119 C GLU A 16 15.320 13.915 -10.813 1.00 52.12 C +ATOM 120 O GLU A 16 14.335 14.293 -10.173 1.00 52.12 O +ATOM 121 CB GLU A 16 16.846 11.909 -10.773 1.00 52.12 C +ATOM 122 CG GLU A 16 18.225 11.338 -10.430 1.00 52.12 C +ATOM 123 CD GLU A 16 18.575 10.185 -11.387 1.00 52.12 C +ATOM 124 OE1 GLU A 16 19.722 10.154 -11.876 1.00 52.12 O +ATOM 125 OE2 GLU A 16 17.657 9.419 -11.770 1.00 52.12 O +ATOM 126 N SER A 17 15.355 14.020 -12.143 1.00 52.30 N +ATOM 127 CA SER A 17 14.186 14.381 -12.933 1.00 52.30 C +ATOM 128 C SER A 17 13.852 13.332 -13.986 1.00 52.30 C +ATOM 129 O SER A 17 14.725 12.663 -14.547 1.00 52.30 O +ATOM 130 CB SER A 17 14.323 15.791 -13.514 1.00 52.30 C +ATOM 131 OG SER A 17 15.416 15.911 -14.403 1.00 52.30 O +ATOM 132 N LEU A 18 12.559 13.214 -14.274 1.00 52.12 N +ATOM 133 CA LEU A 18 12.024 12.423 -15.376 1.00 52.12 C +ATOM 134 C LEU A 18 11.352 13.364 -16.366 1.00 52.12 C +ATOM 135 O LEU A 18 10.617 14.257 -15.957 1.00 52.12 O +ATOM 136 CB LEU A 18 11.036 11.376 -14.835 1.00 52.12 C +ATOM 137 CG LEU A 18 10.328 10.557 -15.931 1.00 52.12 C +ATOM 138 CD1 LEU A 18 11.302 9.637 -16.671 1.00 52.12 C +ATOM 139 CD2 LEU A 18 9.224 9.713 -15.308 1.00 52.12 C +ATOM 140 N THR A 19 11.556 13.127 -17.660 1.00 51.97 N +ATOM 141 CA THR A 19 10.744 13.745 -18.713 1.00 51.97 C +ATOM 142 C THR A 19 9.962 12.659 -19.432 1.00 51.97 C +ATOM 143 O THR A 19 10.552 11.730 -19.981 1.00 51.97 O +ATOM 144 CB THR A 19 11.592 14.575 -19.681 1.00 51.97 C +ATOM 145 OG1 THR A 19 12.311 15.541 -18.941 1.00 51.97 O +ATOM 146 CG2 THR A 19 10.731 15.326 -20.694 1.00 51.97 C +ATOM 147 N ILE A 20 8.638 12.775 -19.403 1.00 51.97 N +ATOM 148 CA ILE A 20 7.713 11.930 -20.156 1.00 51.97 C +ATOM 149 C ILE A 20 7.302 12.715 -21.392 1.00 51.97 C +ATOM 150 O ILE A 20 6.773 13.816 -21.269 1.00 51.97 O +ATOM 151 CB ILE A 20 6.490 11.535 -19.302 1.00 51.97 C +ATOM 152 CG1 ILE A 20 6.946 10.792 -18.027 1.00 51.97 C +ATOM 153 CG2 ILE A 20 5.529 10.673 -20.145 1.00 51.97 C +ATOM 154 CD1 ILE A 20 5.808 10.484 -17.053 1.00 51.97 C +ATOM 155 N ASN A 21 7.538 12.154 -22.568 1.00 52.03 N +ATOM 156 CA ASN A 21 7.199 12.748 -23.851 1.00 52.03 C +ATOM 157 C ASN A 21 5.957 12.078 -24.426 1.00 52.03 C +ATOM 158 O ASN A 21 5.810 10.860 -24.341 1.00 52.03 O +ATOM 159 CB ASN A 21 8.388 12.617 -24.809 1.00 52.03 C +ATOM 160 CG ASN A 21 9.630 13.280 -24.257 1.00 52.03 C +ATOM 161 OD1 ASN A 21 9.708 14.487 -24.130 1.00 52.03 O +ATOM 162 ND2 ASN A 21 10.625 12.510 -23.883 1.00 52.03 N +ATOM 163 N CYS A 22 5.101 12.879 -25.044 1.00 52.30 N +ATOM 164 CA CYS A 22 3.874 12.452 -25.689 1.00 52.30 C +ATOM 165 C CYS A 22 3.735 13.100 -27.063 1.00 52.30 C +ATOM 166 O CYS A 22 3.972 14.297 -27.224 1.00 52.30 O +ATOM 167 CB CYS A 22 2.692 12.748 -24.774 1.00 52.30 C +ATOM 168 SG CYS A 22 2.789 11.705 -23.306 1.00 52.30 S +ATOM 169 N VAL A 23 3.346 12.311 -28.061 1.00 52.59 N +ATOM 170 CA VAL A 23 3.093 12.788 -29.425 1.00 52.59 C +ATOM 171 C VAL A 23 1.748 12.254 -29.891 1.00 52.59 C +ATOM 172 O VAL A 23 1.531 11.045 -29.904 1.00 52.59 O +ATOM 173 CB VAL A 23 4.223 12.381 -30.393 1.00 52.59 C +ATOM 174 CG1 VAL A 23 4.024 13.047 -31.760 1.00 52.59 C +ATOM 175 CG2 VAL A 23 5.611 12.781 -29.871 1.00 52.59 C +ATOM 176 N LEU A 24 0.836 13.143 -30.281 1.00 53.60 N +ATOM 177 CA LEU A 24 -0.433 12.754 -30.889 1.00 53.60 C +ATOM 178 C LEU A 24 -0.192 12.378 -32.357 1.00 53.60 C +ATOM 179 O LEU A 24 0.072 13.242 -33.197 1.00 53.60 O +ATOM 180 CB LEU A 24 -1.443 13.896 -30.710 1.00 53.60 C +ATOM 181 CG LEU A 24 -2.798 13.653 -31.397 1.00 53.60 C +ATOM 182 CD1 LEU A 24 -3.580 12.521 -30.731 1.00 53.60 C +ATOM 183 CD2 LEU A 24 -3.634 14.931 -31.345 1.00 53.60 C +ATOM 184 N ARG A 25 -0.295 11.083 -32.660 1.00 54.15 N +ATOM 185 CA ARG A 25 -0.084 10.508 -33.992 1.00 54.15 C +ATOM 186 C ARG A 25 -1.406 10.165 -34.669 1.00 54.15 C +ATOM 187 O ARG A 25 -2.436 9.989 -34.013 1.00 54.15 O +ATOM 188 CB ARG A 25 0.828 9.279 -33.895 1.00 54.15 C +ATOM 189 CG ARG A 25 2.229 9.626 -33.365 1.00 54.15 C +ATOM 190 CD ARG A 25 3.134 8.402 -33.523 1.00 54.15 C +ATOM 191 NE ARG A 25 4.442 8.569 -32.874 1.00 54.15 N +ATOM 192 CZ ARG A 25 5.508 9.189 -33.321 1.00 54.15 C +ATOM 193 NH1 ARG A 25 5.533 9.803 -34.474 1.00 54.15 N +ATOM 194 NH2 ARG A 25 6.571 9.172 -32.580 1.00 54.15 N +ATOM 195 N ASP A 26 -1.368 10.099 -35.998 1.00 56.95 N +ATOM 196 CA ASP A 26 -2.491 9.685 -36.850 1.00 56.95 C +ATOM 197 C ASP A 26 -3.791 10.456 -36.561 1.00 56.95 C +ATOM 198 O ASP A 26 -4.900 9.914 -36.576 1.00 56.95 O +ATOM 199 CB ASP A 26 -2.631 8.154 -36.819 1.00 56.95 C +ATOM 200 CG ASP A 26 -1.289 7.462 -37.082 1.00 56.95 C +ATOM 201 OD1 ASP A 26 -0.525 7.994 -37.918 1.00 56.95 O +ATOM 202 OD2 ASP A 26 -0.997 6.484 -36.358 1.00 56.95 O +ATOM 203 N ALA A 27 -3.641 11.740 -36.223 1.00 61.02 N +ATOM 204 CA ALA A 27 -4.730 12.676 -35.997 1.00 61.02 C +ATOM 205 C ALA A 27 -4.819 13.652 -37.171 1.00 61.02 C +ATOM 206 O ALA A 27 -3.866 14.385 -37.445 1.00 61.02 O +ATOM 207 CB ALA A 27 -4.508 13.404 -34.666 1.00 61.02 C +ATOM 208 N SER A 28 -5.987 13.729 -37.810 1.00 61.79 N +ATOM 209 CA SER A 28 -6.305 14.756 -38.816 1.00 61.79 C +ATOM 210 C SER A 28 -6.506 16.148 -38.203 1.00 61.79 C +ATOM 211 O SER A 28 -6.575 17.147 -38.909 1.00 61.79 O +ATOM 212 CB SER A 28 -7.562 14.344 -39.586 1.00 61.79 C +ATOM 213 OG SER A 28 -8.628 14.103 -38.681 1.00 61.79 O +ATOM 214 N TYR A 29 -6.568 16.224 -36.874 1.00 65.00 N +ATOM 215 CA TYR A 29 -6.814 17.434 -36.106 1.00 65.00 C +ATOM 216 C TYR A 29 -5.589 17.933 -35.333 1.00 65.00 C +ATOM 217 O TYR A 29 -4.653 17.175 -35.074 1.00 65.00 O +ATOM 218 CB TYR A 29 -7.964 17.158 -35.151 1.00 65.00 C +ATOM 219 CG TYR A 29 -7.847 15.931 -34.278 1.00 65.00 C +ATOM 220 CD1 TYR A 29 -8.427 14.715 -34.679 1.00 65.00 C +ATOM 221 CD2 TYR A 29 -7.257 16.049 -33.004 1.00 65.00 C +ATOM 222 CE1 TYR A 29 -8.385 13.610 -33.809 1.00 65.00 C +ATOM 223 CE2 TYR A 29 -7.247 14.958 -32.122 1.00 65.00 C +ATOM 224 CZ TYR A 29 -7.793 13.729 -32.533 1.00 65.00 C +ATOM 225 OH TYR A 29 -7.756 12.653 -31.714 1.00 65.00 O +ATOM 226 N GLY A 30 -5.606 19.213 -34.953 1.00 61.64 N +ATOM 227 CA GLY A 30 -4.603 19.828 -34.079 1.00 61.64 C +ATOM 228 C GLY A 30 -4.761 19.450 -32.601 1.00 61.64 C +ATOM 229 O GLY A 30 -5.850 19.087 -32.155 1.00 61.64 O +ATOM 230 N LEU A 31 -3.657 19.540 -31.855 1.00 57.69 N +ATOM 231 CA LEU A 31 -3.628 19.392 -30.400 1.00 57.69 C +ATOM 232 C LEU A 31 -4.125 20.690 -29.743 1.00 57.69 C +ATOM 233 O LEU A 31 -3.459 21.721 -29.853 1.00 57.69 O +ATOM 234 CB LEU A 31 -2.178 19.054 -29.988 1.00 57.69 C +ATOM 235 CG LEU A 31 -1.949 18.920 -28.469 1.00 57.69 C +ATOM 236 CD1 LEU A 31 -2.709 17.726 -27.894 1.00 57.69 C +ATOM 237 CD2 LEU A 31 -0.458 18.723 -28.193 1.00 57.69 C +ATOM 238 N GLU A 32 -5.269 20.648 -29.053 1.00 59.12 N +ATOM 239 CA GLU A 32 -5.853 21.850 -28.437 1.00 59.12 C +ATOM 240 C GLU A 32 -5.440 22.033 -26.985 1.00 59.12 C +ATOM 241 O GLU A 32 -4.960 23.106 -26.607 1.00 59.12 O +ATOM 242 CB GLU A 32 -7.388 21.848 -28.536 1.00 59.12 C +ATOM 243 CG GLU A 32 -7.910 22.046 -29.962 1.00 59.12 C +ATOM 244 CD GLU A 32 -7.343 23.314 -30.624 1.00 59.12 C +ATOM 245 OE1 GLU A 32 -7.181 23.269 -31.862 1.00 59.12 O +ATOM 246 OE2 GLU A 32 -7.005 24.275 -29.888 1.00 59.12 O +ATOM 247 N SER A 33 -5.614 20.981 -26.189 1.00 57.46 N +ATOM 248 CA SER A 33 -5.222 20.946 -24.784 1.00 57.46 C +ATOM 249 C SER A 33 -4.653 19.582 -24.413 1.00 57.46 C +ATOM 250 O SER A 33 -4.801 18.591 -25.135 1.00 57.46 O +ATOM 251 CB SER A 33 -6.386 21.349 -23.867 1.00 57.46 C +ATOM 252 OG SER A 33 -7.534 20.564 -24.096 1.00 57.46 O +ATOM 253 N THR A 34 -3.955 19.562 -23.289 1.00 54.48 N +ATOM 254 CA THR A 34 -3.237 18.412 -22.755 1.00 54.48 C +ATOM 255 C THR A 34 -3.698 18.149 -21.327 1.00 54.48 C +ATOM 256 O THR A 34 -4.196 19.035 -20.635 1.00 54.48 O +ATOM 257 CB THR A 34 -1.723 18.657 -22.797 1.00 54.48 C +ATOM 258 OG1 THR A 34 -1.458 19.921 -22.253 1.00 54.48 O +ATOM 259 CG2 THR A 34 -1.187 18.688 -24.228 1.00 54.48 C +ATOM 260 N GLY A 35 -3.562 16.905 -20.887 1.00 54.02 N +ATOM 261 CA GLY A 35 -3.855 16.504 -19.518 1.00 54.02 C +ATOM 262 C GLY A 35 -2.839 15.492 -19.032 1.00 54.02 C +ATOM 263 O GLY A 35 -2.244 14.757 -19.822 1.00 54.02 O +ATOM 264 N TRP A 36 -2.643 15.453 -17.722 1.00 52.70 N +ATOM 265 CA TRP A 36 -1.684 14.563 -17.092 1.00 52.70 C +ATOM 266 C TRP A 36 -2.336 13.833 -15.941 1.00 52.70 C +ATOM 267 O TRP A 36 -3.095 14.409 -15.164 1.00 52.70 O +ATOM 268 CB TRP A 36 -0.453 15.348 -16.652 1.00 52.70 C +ATOM 269 CG TRP A 36 0.336 15.867 -17.801 1.00 52.70 C +ATOM 270 CD1 TRP A 36 0.187 17.083 -18.365 1.00 52.70 C +ATOM 271 CD2 TRP A 36 1.329 15.168 -18.609 1.00 52.70 C +ATOM 272 NE1 TRP A 36 1.031 17.194 -19.445 1.00 52.70 N +ATOM 273 CE2 TRP A 36 1.751 16.043 -19.654 1.00 52.70 C +ATOM 274 CE3 TRP A 36 1.918 13.883 -18.564 1.00 52.70 C +ATOM 275 CZ2 TRP A 36 2.711 15.667 -20.605 1.00 52.70 C +ATOM 276 CZ3 TRP A 36 2.902 13.506 -19.498 1.00 52.70 C +ATOM 277 CH2 TRP A 36 3.287 14.388 -20.526 1.00 52.70 C +ATOM 278 N TYR A 37 -2.021 12.553 -15.847 1.00 52.59 N +ATOM 279 CA TYR A 37 -2.644 11.635 -14.922 1.00 52.59 C +ATOM 280 C TYR A 37 -1.592 10.808 -14.219 1.00 52.59 C +ATOM 281 O TYR A 37 -0.558 10.466 -14.802 1.00 52.59 O +ATOM 282 CB TYR A 37 -3.623 10.726 -15.659 1.00 52.59 C +ATOM 283 CG TYR A 37 -4.676 11.483 -16.428 1.00 52.59 C +ATOM 284 CD1 TYR A 37 -5.787 12.005 -15.744 1.00 52.59 C +ATOM 285 CD2 TYR A 37 -4.496 11.734 -17.802 1.00 52.59 C +ATOM 286 CE1 TYR A 37 -6.717 12.803 -16.428 1.00 52.59 C +ATOM 287 CE2 TYR A 37 -5.429 12.525 -18.494 1.00 52.59 C +ATOM 288 CZ TYR A 37 -6.531 13.066 -17.798 1.00 52.59 C +ATOM 289 OH TYR A 37 -7.421 13.851 -18.441 1.00 52.59 O +ATOM 290 N ARG A 38 -1.882 10.458 -12.974 1.00 52.70 N +ATOM 291 CA ARG A 38 -1.032 9.617 -12.156 1.00 52.70 C +ATOM 292 C ARG A 38 -1.856 8.627 -11.355 1.00 52.70 C +ATOM 293 O ARG A 38 -2.767 9.013 -10.630 1.00 52.70 O +ATOM 294 CB ARG A 38 -0.201 10.511 -11.238 1.00 52.70 C +ATOM 295 CG ARG A 38 0.836 9.672 -10.510 1.00 52.70 C +ATOM 296 CD ARG A 38 1.625 10.545 -9.554 1.00 52.70 C +ATOM 297 NE ARG A 38 2.523 9.662 -8.828 1.00 52.70 N +ATOM 298 CZ ARG A 38 3.013 9.856 -7.640 1.00 52.70 C +ATOM 299 NH1 ARG A 38 3.100 11.026 -7.066 1.00 52.70 N +ATOM 300 NH2 ARG A 38 3.396 8.798 -7.013 1.00 52.70 N +ATOM 301 N THR A 39 -1.441 7.374 -11.390 1.00 52.21 N +ATOM 302 CA THR A 39 -1.868 6.337 -10.454 1.00 52.21 C +ATOM 303 C THR A 39 -0.689 6.019 -9.546 1.00 52.21 C +ATOM 304 O THR A 39 0.365 5.599 -10.023 1.00 52.21 O +ATOM 305 CB THR A 39 -2.342 5.097 -11.214 1.00 52.21 C +ATOM 306 OG1 THR A 39 -3.352 5.448 -12.129 1.00 52.21 O +ATOM 307 CG2 THR A 39 -2.941 4.061 -10.277 1.00 52.21 C +ATOM 308 N LYS A 40 -0.827 6.267 -8.240 1.00 53.01 N +ATOM 309 CA LYS A 40 0.237 5.973 -7.264 1.00 53.01 C +ATOM 310 C LYS A 40 0.509 4.471 -7.198 1.00 53.01 C +ATOM 311 O LYS A 40 -0.401 3.669 -7.403 1.00 53.01 O +ATOM 312 CB LYS A 40 -0.129 6.510 -5.874 1.00 53.01 C +ATOM 313 CG LYS A 40 -0.181 8.044 -5.836 1.00 53.01 C +ATOM 314 CD LYS A 40 -0.454 8.540 -4.411 1.00 53.01 C +ATOM 315 CE LYS A 40 -0.497 10.071 -4.390 1.00 53.01 C +ATOM 316 NZ LYS A 40 -0.703 10.589 -3.014 1.00 53.01 N +ATOM 317 N LEU A 41 1.746 4.095 -6.874 1.00 54.35 N +ATOM 318 CA LEU A 41 2.094 2.697 -6.623 1.00 54.35 C +ATOM 319 C LEU A 41 1.163 2.098 -5.549 1.00 54.35 C +ATOM 320 O LEU A 41 0.953 2.706 -4.503 1.00 54.35 O +ATOM 321 CB LEU A 41 3.576 2.627 -6.215 1.00 54.35 C +ATOM 322 CG LEU A 41 4.111 1.195 -6.041 1.00 54.35 C +ATOM 323 CD1 LEU A 41 4.117 0.424 -7.364 1.00 54.35 C +ATOM 324 CD2 LEU A 41 5.540 1.246 -5.507 1.00 54.35 C +ATOM 325 N GLY A 42 0.576 0.932 -5.834 1.00 55.01 N +ATOM 326 CA GLY A 42 -0.383 0.262 -4.944 1.00 55.01 C +ATOM 327 C GLY A 42 -1.817 0.809 -4.986 1.00 55.01 C +ATOM 328 O GLY A 42 -2.678 0.292 -4.282 1.00 55.01 O +ATOM 329 N SER A 43 -2.103 1.823 -5.809 1.00 54.15 N +ATOM 330 CA SER A 43 -3.457 2.349 -6.013 1.00 54.15 C +ATOM 331 C SER A 43 -4.025 1.929 -7.369 1.00 54.15 C +ATOM 332 O SER A 43 -3.289 1.731 -8.330 1.00 54.15 O +ATOM 333 CB SER A 43 -3.448 3.873 -5.888 1.00 54.15 C +ATOM 334 OG SER A 43 -4.772 4.362 -5.777 1.00 54.15 O +ATOM 335 N THR A 44 -5.351 1.849 -7.465 1.00 55.37 N +ATOM 336 CA THR A 44 -6.091 1.742 -8.733 1.00 55.37 C +ATOM 337 C THR A 44 -6.693 3.080 -9.169 1.00 55.37 C +ATOM 338 O THR A 44 -7.153 3.208 -10.302 1.00 55.37 O +ATOM 339 CB THR A 44 -7.205 0.696 -8.609 1.00 55.37 C +ATOM 340 OG1 THR A 44 -7.989 0.987 -7.475 1.00 55.37 O +ATOM 341 CG2 THR A 44 -6.639 -0.714 -8.432 1.00 55.37 C +ATOM 342 N ASN A 45 -6.661 4.095 -8.298 1.00 53.71 N +ATOM 343 CA ASN A 45 -7.279 5.390 -8.551 1.00 53.71 C +ATOM 344 C ASN A 45 -6.339 6.283 -9.365 1.00 53.71 C +ATOM 345 O ASN A 45 -5.258 6.662 -8.909 1.00 53.71 O +ATOM 346 CB ASN A 45 -7.703 6.034 -7.219 1.00 53.71 C +ATOM 347 CG ASN A 45 -8.846 5.294 -6.542 1.00 53.71 C +ATOM 348 OD1 ASN A 45 -9.605 4.564 -7.147 1.00 53.71 O +ATOM 349 ND2 ASN A 45 -9.024 5.478 -5.254 1.00 53.71 N +ATOM 350 N GLU A 46 -6.772 6.633 -10.572 1.00 53.37 N +ATOM 351 CA GLU A 46 -6.091 7.600 -11.430 1.00 53.37 C +ATOM 352 C GLU A 46 -6.462 9.030 -11.007 1.00 53.37 C +ATOM 353 O GLU A 46 -7.636 9.380 -10.891 1.00 53.37 O +ATOM 354 CB GLU A 46 -6.427 7.296 -12.900 1.00 53.37 C +ATOM 355 CG GLU A 46 -5.594 8.133 -13.877 1.00 53.37 C +ATOM 356 CD GLU A 46 -5.841 7.735 -15.345 1.00 53.37 C +ATOM 357 OE1 GLU A 46 -6.739 8.331 -15.988 1.00 53.37 O +ATOM 358 OE2 GLU A 46 -5.127 6.847 -15.871 1.00 53.37 O +ATOM 359 N GLN A 47 -5.451 9.860 -10.762 1.00 53.11 N +ATOM 360 CA GLN A 47 -5.591 11.243 -10.311 1.00 53.11 C +ATOM 361 C GLN A 47 -5.096 12.200 -11.390 1.00 53.11 C +ATOM 362 O GLN A 47 -4.102 11.922 -12.058 1.00 53.11 O +ATOM 363 CB GLN A 47 -4.811 11.451 -9.003 1.00 53.11 C +ATOM 364 CG GLN A 47 -5.320 10.535 -7.880 1.00 53.11 C +ATOM 365 CD GLN A 47 -4.635 10.781 -6.538 1.00 53.11 C +ATOM 366 OE1 GLN A 47 -3.447 11.046 -6.415 1.00 53.11 O +ATOM 367 NE2 GLN A 47 -5.371 10.677 -5.452 1.00 53.11 N +ATOM 368 N THR A 48 -5.754 13.349 -11.538 1.00 53.37 N +ATOM 369 CA THR A 48 -5.270 14.411 -12.432 1.00 53.37 C +ATOM 370 C THR A 48 -4.097 15.136 -11.776 1.00 53.37 C +ATOM 371 O THR A 48 -4.146 15.465 -10.591 1.00 53.37 O +ATOM 372 CB THR A 48 -6.386 15.394 -12.809 1.00 53.37 C +ATOM 373 OG1 THR A 48 -7.453 14.685 -13.401 1.00 53.37 O +ATOM 374 CG2 THR A 48 -5.933 16.433 -13.833 1.00 53.37 C +ATOM 375 N ILE A 49 -3.046 15.400 -12.544 1.00 53.01 N +ATOM 376 CA ILE A 49 -1.899 16.197 -12.116 1.00 53.01 C +ATOM 377 C ILE A 49 -2.172 17.654 -12.478 1.00 53.01 C +ATOM 378 O ILE A 49 -2.444 17.975 -13.633 1.00 53.01 O +ATOM 379 CB ILE A 49 -0.586 15.698 -12.755 1.00 53.01 C +ATOM 380 CG1 ILE A 49 -0.385 14.182 -12.533 1.00 53.01 C +ATOM 381 CG2 ILE A 49 0.608 16.509 -12.211 1.00 53.01 C +ATOM 382 CD1 ILE A 49 0.888 13.650 -13.201 1.00 53.01 C +ATOM 383 N SER A 50 -2.065 18.543 -11.495 1.00 53.55 N +ATOM 384 CA SER A 50 -2.044 19.985 -11.744 1.00 53.55 C +ATOM 385 C SER A 50 -0.624 20.418 -12.106 1.00 53.55 C +ATOM 386 O SER A 50 0.319 20.119 -11.371 1.00 53.55 O +ATOM 387 CB SER A 50 -2.552 20.747 -10.522 1.00 53.55 C +ATOM 388 OG SER A 50 -3.897 20.383 -10.279 1.00 53.55 O +ATOM 389 N ILE A 51 -0.468 21.105 -13.237 1.00 52.97 N +ATOM 390 CA ILE A 51 0.829 21.637 -13.671 1.00 52.97 C +ATOM 391 C ILE A 51 1.240 22.799 -12.764 1.00 52.97 C +ATOM 392 O ILE A 51 0.428 23.662 -12.436 1.00 52.97 O +ATOM 393 CB ILE A 51 0.791 22.028 -15.166 1.00 52.97 C +ATOM 394 CG1 ILE A 51 0.425 20.829 -16.071 1.00 52.97 C +ATOM 395 CG2 ILE A 51 2.129 22.639 -15.625 1.00 52.97 C +ATOM 396 CD1 ILE A 51 1.311 19.589 -15.883 1.00 52.97 C +ATOM 397 N GLY A 52 2.508 22.810 -12.354 1.00 53.91 N +ATOM 398 CA GLY A 52 3.049 23.753 -11.383 1.00 53.91 C +ATOM 399 C GLY A 52 4.042 23.099 -10.419 1.00 53.91 C +ATOM 400 O GLY A 52 4.124 21.875 -10.288 1.00 53.91 O +ATOM 401 N GLY A 53 4.831 23.929 -9.732 1.00 53.30 N +ATOM 402 CA GLY A 53 5.849 23.463 -8.788 1.00 53.30 C +ATOM 403 C GLY A 53 6.890 22.553 -9.453 1.00 53.30 C +ATOM 404 O GLY A 53 7.659 22.996 -10.308 1.00 53.30 O +ATOM 405 N ARG A 54 6.924 21.275 -9.052 1.00 52.77 N +ATOM 406 CA ARG A 54 7.851 20.263 -9.594 1.00 52.77 C +ATOM 407 C ARG A 54 7.434 19.719 -10.967 1.00 52.77 C +ATOM 408 O ARG A 54 8.280 19.160 -11.662 1.00 52.77 O +ATOM 409 CB ARG A 54 8.026 19.105 -8.593 1.00 52.77 C +ATOM 410 CG ARG A 54 8.646 19.543 -7.252 1.00 52.77 C +ATOM 411 CD ARG A 54 9.078 18.362 -6.366 1.00 52.77 C +ATOM 412 NE ARG A 54 7.972 17.434 -6.040 1.00 52.77 N +ATOM 413 CZ ARG A 54 7.911 16.150 -6.370 1.00 52.77 C +ATOM 414 NH1 ARG A 54 8.808 15.546 -7.068 1.00 52.77 N +ATOM 415 NH2 ARG A 54 6.909 15.402 -6.011 1.00 52.77 N +ATOM 416 N TYR A 55 6.169 19.890 -11.354 1.00 52.34 N +ATOM 417 CA TYR A 55 5.595 19.411 -12.611 1.00 52.34 C +ATOM 418 C TYR A 55 5.617 20.525 -13.655 1.00 52.34 C +ATOM 419 O TYR A 55 4.897 21.515 -13.528 1.00 52.34 O +ATOM 420 CB TYR A 55 4.159 18.932 -12.349 1.00 52.34 C +ATOM 421 CG TYR A 55 4.070 17.769 -11.385 1.00 52.34 C +ATOM 422 CD1 TYR A 55 4.463 16.478 -11.786 1.00 52.34 C +ATOM 423 CD2 TYR A 55 3.612 17.989 -10.073 1.00 52.34 C +ATOM 424 CE1 TYR A 55 4.434 15.410 -10.866 1.00 52.34 C +ATOM 425 CE2 TYR A 55 3.566 16.921 -9.157 1.00 52.34 C +ATOM 426 CZ TYR A 55 3.990 15.634 -9.545 1.00 52.34 C +ATOM 427 OH TYR A 55 3.978 14.625 -8.632 1.00 52.34 O +ATOM 428 N VAL A 56 6.444 20.372 -14.687 1.00 52.18 N +ATOM 429 CA VAL A 56 6.623 21.381 -15.737 1.00 52.18 C +ATOM 430 C VAL A 56 6.230 20.802 -17.079 1.00 52.18 C +ATOM 431 O VAL A 56 6.845 19.845 -17.548 1.00 52.18 O +ATOM 432 CB VAL A 56 8.061 21.917 -15.752 1.00 52.18 C +ATOM 433 CG1 VAL A 56 8.246 23.013 -16.812 1.00 52.18 C +ATOM 434 CG2 VAL A 56 8.396 22.506 -14.376 1.00 52.18 C +ATOM 435 N GLU A 57 5.218 21.400 -17.694 1.00 52.27 N +ATOM 436 CA GLU A 57 4.759 21.019 -19.020 1.00 52.27 C +ATOM 437 C GLU A 57 5.422 21.865 -20.114 1.00 52.27 C +ATOM 438 O GLU A 57 5.607 23.074 -19.962 1.00 52.27 O +ATOM 439 CB GLU A 57 3.232 21.093 -19.080 1.00 52.27 C +ATOM 440 CG GLU A 57 2.716 20.540 -20.410 1.00 52.27 C +ATOM 441 CD GLU A 57 1.194 20.533 -20.470 1.00 52.27 C +ATOM 442 OE1 GLU A 57 0.658 19.478 -20.875 1.00 52.27 O +ATOM 443 OE2 GLU A 57 0.576 21.559 -20.133 1.00 52.27 O +ATOM 444 N THR A 58 5.726 21.219 -21.237 1.00 52.27 N +ATOM 445 CA THR A 58 6.052 21.874 -22.507 1.00 52.27 C +ATOM 446 C THR A 58 5.073 21.378 -23.562 1.00 52.27 C +ATOM 447 O THR A 58 4.890 20.169 -23.690 1.00 52.27 O +ATOM 448 CB THR A 58 7.487 21.555 -22.951 1.00 52.27 C +ATOM 449 OG1 THR A 58 8.425 21.874 -21.949 1.00 52.27 O +ATOM 450 CG2 THR A 58 7.910 22.343 -24.191 1.00 52.27 C +ATOM 451 N VAL A 59 4.479 22.280 -24.348 1.00 52.66 N +ATOM 452 CA VAL A 59 3.542 21.931 -25.428 1.00 52.66 C +ATOM 453 C VAL A 59 4.005 22.552 -26.741 1.00 52.66 C +ATOM 454 O VAL A 59 4.218 23.759 -26.830 1.00 52.66 O +ATOM 455 CB VAL A 59 2.094 22.362 -25.105 1.00 52.66 C +ATOM 456 CG1 VAL A 59 1.114 21.881 -26.185 1.00 52.66 C +ATOM 457 CG2 VAL A 59 1.621 21.804 -23.758 1.00 52.66 C +ATOM 458 N ASN A 60 4.107 21.732 -27.782 1.00 53.04 N +ATOM 459 CA ASN A 60 4.303 22.149 -29.162 1.00 53.04 C +ATOM 460 C ASN A 60 3.112 21.674 -30.004 1.00 53.04 C +ATOM 461 O ASN A 60 3.023 20.516 -30.420 1.00 53.04 O +ATOM 462 CB ASN A 60 5.658 21.635 -29.669 1.00 53.04 C +ATOM 463 CG ASN A 60 5.980 22.157 -31.063 1.00 53.04 C +ATOM 464 OD1 ASN A 60 5.114 22.477 -31.867 1.00 53.04 O +ATOM 465 ND2 ASN A 60 7.244 22.299 -31.381 1.00 53.04 N +ATOM 466 N LYS A 61 2.182 22.595 -30.271 1.00 56.00 N +ATOM 467 CA LYS A 61 0.980 22.309 -31.067 1.00 56.00 C +ATOM 468 C LYS A 61 1.296 21.984 -32.532 1.00 56.00 C +ATOM 469 O LYS A 61 0.570 21.196 -33.129 1.00 56.00 O +ATOM 470 CB LYS A 61 -0.017 23.476 -30.967 1.00 56.00 C +ATOM 471 CG LYS A 61 -0.608 23.638 -29.556 1.00 56.00 C +ATOM 472 CD LYS A 61 -1.783 24.633 -29.558 1.00 56.00 C +ATOM 473 CE LYS A 61 -2.524 24.578 -28.215 1.00 56.00 C +ATOM 474 NZ LYS A 61 -3.864 25.221 -28.270 1.00 56.00 N +ATOM 475 N GLY A 62 2.377 22.539 -33.091 1.00 56.00 N +ATOM 476 CA GLY A 62 2.778 22.325 -34.486 1.00 56.00 C +ATOM 477 C GLY A 62 3.212 20.886 -34.759 1.00 56.00 C +ATOM 478 O GLY A 62 2.730 20.261 -35.699 1.00 56.00 O +ATOM 479 N SER A 63 4.049 20.322 -33.886 1.00 54.35 N +ATOM 480 CA SER A 63 4.440 18.903 -33.927 1.00 54.35 C +ATOM 481 C SER A 63 3.478 17.978 -33.173 1.00 54.35 C +ATOM 482 O SER A 63 3.738 16.780 -33.075 1.00 54.35 O +ATOM 483 CB SER A 63 5.868 18.737 -33.403 1.00 54.35 C +ATOM 484 OG SER A 63 5.987 19.331 -32.127 1.00 54.35 O +ATOM 485 N LYS A 64 2.387 18.522 -32.614 1.00 53.98 N +ATOM 486 CA LYS A 64 1.427 17.819 -31.744 1.00 53.98 C +ATOM 487 C LYS A 64 2.121 17.027 -30.630 1.00 53.98 C +ATOM 488 O LYS A 64 1.682 15.938 -30.255 1.00 53.98 O +ATOM 489 CB LYS A 64 0.468 16.948 -32.571 1.00 53.98 C +ATOM 490 CG LYS A 64 -0.140 17.660 -33.782 1.00 53.98 C +ATOM 491 CD LYS A 64 -1.089 16.720 -34.531 1.00 53.98 C +ATOM 492 CE LYS A 64 -1.446 17.369 -35.869 1.00 53.98 C +ATOM 493 NZ LYS A 64 -2.462 16.583 -36.602 1.00 53.98 N +ATOM 494 N SER A 65 3.225 17.573 -30.126 1.00 52.94 N +ATOM 495 CA SER A 65 4.061 16.949 -29.113 1.00 52.94 C +ATOM 496 C SER A 65 4.038 17.761 -27.833 1.00 52.94 C +ATOM 497 O SER A 65 3.945 18.985 -27.846 1.00 52.94 O +ATOM 498 CB SER A 65 5.484 16.701 -29.617 1.00 52.94 C +ATOM 499 OG SER A 65 6.188 17.905 -29.850 1.00 52.94 O +ATOM 500 N PHE A 66 4.114 17.077 -26.709 1.00 52.37 N +ATOM 501 CA PHE A 66 4.060 17.681 -25.393 1.00 52.37 C +ATOM 502 C PHE A 66 4.794 16.790 -24.402 1.00 52.37 C +ATOM 503 O PHE A 66 4.981 15.599 -24.641 1.00 52.37 O +ATOM 504 CB PHE A 66 2.601 17.949 -24.992 1.00 52.37 C +ATOM 505 CG PHE A 66 1.678 16.742 -24.955 1.00 52.37 C +ATOM 506 CD1 PHE A 66 1.256 16.110 -26.144 1.00 52.37 C +ATOM 507 CD2 PHE A 66 1.140 16.324 -23.728 1.00 52.37 C +ATOM 508 CE1 PHE A 66 0.351 15.036 -26.099 1.00 52.37 C +ATOM 509 CE2 PHE A 66 0.193 15.292 -23.690 1.00 52.37 C +ATOM 510 CZ PHE A 66 -0.183 14.623 -24.868 1.00 52.37 C +ATOM 511 N SER A 67 5.265 17.362 -23.306 1.00 52.06 N +ATOM 512 CA SER A 67 6.060 16.625 -22.334 1.00 52.06 C +ATOM 513 C SER A 67 5.827 17.121 -20.924 1.00 52.06 C +ATOM 514 O SER A 67 5.685 18.325 -20.720 1.00 52.06 O +ATOM 515 CB SER A 67 7.555 16.712 -22.666 1.00 52.06 C +ATOM 516 OG SER A 67 8.008 18.052 -22.560 1.00 52.06 O +ATOM 517 N LEU A 68 5.896 16.210 -19.961 1.00 52.00 N +ATOM 518 CA LEU A 68 5.911 16.519 -18.542 1.00 52.00 C +ATOM 519 C LEU A 68 7.290 16.226 -17.985 1.00 52.00 C +ATOM 520 O LEU A 68 7.729 15.074 -17.971 1.00 52.00 O +ATOM 521 CB LEU A 68 4.847 15.691 -17.816 1.00 52.00 C +ATOM 522 CG LEU A 68 4.703 15.997 -16.321 1.00 52.00 C +ATOM 523 CD1 LEU A 68 4.141 17.403 -16.117 1.00 52.00 C +ATOM 524 CD2 LEU A 68 3.749 14.996 -15.670 1.00 52.00 C +ATOM 525 N ARG A 69 7.950 17.260 -17.476 1.00 52.03 N +ATOM 526 CA ARG A 69 9.151 17.111 -16.667 1.00 52.03 C +ATOM 527 C ARG A 69 8.796 17.184 -15.190 1.00 52.03 C +ATOM 528 O ARG A 69 8.312 18.210 -14.719 1.00 52.03 O +ATOM 529 CB ARG A 69 10.208 18.121 -17.101 1.00 52.03 C +ATOM 530 CG ARG A 69 11.525 17.856 -16.357 1.00 52.03 C +ATOM 531 CD ARG A 69 12.689 18.560 -17.053 1.00 52.03 C +ATOM 532 NE ARG A 69 12.544 20.028 -17.036 1.00 52.03 N +ATOM 533 CZ ARG A 69 13.131 20.858 -16.201 1.00 52.03 C +ATOM 534 NH1 ARG A 69 13.798 20.443 -15.160 1.00 52.03 N +ATOM 535 NH2 ARG A 69 13.062 22.142 -16.414 1.00 52.03 N +ATOM 536 N ILE A 70 9.085 16.110 -14.470 1.00 52.12 N +ATOM 537 CA ILE A 70 8.953 16.010 -13.017 1.00 52.12 C +ATOM 538 C ILE A 70 10.340 16.248 -12.423 1.00 52.12 C +ATOM 539 O ILE A 70 11.272 15.507 -12.732 1.00 52.12 O +ATOM 540 CB ILE A 70 8.381 14.638 -12.593 1.00 52.12 C +ATOM 541 CG1 ILE A 70 7.186 14.202 -13.475 1.00 52.12 C +ATOM 542 CG2 ILE A 70 7.995 14.727 -11.104 1.00 52.12 C +ATOM 543 CD1 ILE A 70 6.652 12.803 -13.158 1.00 52.12 C +ATOM 544 N ARG A 71 10.497 17.313 -11.639 1.00 52.18 N +ATOM 545 CA ARG A 71 11.752 17.675 -10.954 1.00 52.18 C +ATOM 546 C ARG A 71 11.779 17.122 -9.533 1.00 52.18 C +ATOM 547 O ARG A 71 10.710 16.913 -8.966 1.00 52.18 O +ATOM 548 CB ARG A 71 11.900 19.201 -10.902 1.00 52.18 C +ATOM 549 CG ARG A 71 11.897 19.869 -12.281 1.00 52.18 C +ATOM 550 CD ARG A 71 12.065 21.382 -12.096 1.00 52.18 C +ATOM 551 NE ARG A 71 11.908 22.122 -13.363 1.00 52.18 N +ATOM 552 CZ ARG A 71 11.803 23.436 -13.479 1.00 52.18 C +ATOM 553 NH1 ARG A 71 11.825 24.226 -12.447 1.00 52.18 N +ATOM 554 NH2 ARG A 71 11.669 23.998 -14.647 1.00 52.18 N +ATOM 555 N ASP A 72 12.975 17.007 -8.952 1.00 52.34 N +ATOM 556 CA ASP A 72 13.161 16.649 -7.536 1.00 52.34 C +ATOM 557 C ASP A 72 12.325 15.411 -7.178 1.00 52.34 C +ATOM 558 O ASP A 72 11.403 15.491 -6.370 1.00 52.34 O +ATOM 559 CB ASP A 72 12.898 17.893 -6.655 1.00 52.34 C +ATOM 560 CG ASP A 72 12.951 17.598 -5.148 1.00 52.34 C +ATOM 561 OD1 ASP A 72 13.940 16.956 -4.750 1.00 52.34 O +ATOM 562 OD2 ASP A 72 12.010 18.028 -4.419 1.00 52.34 O +ATOM 563 N LEU A 73 12.528 14.314 -7.918 1.00 52.18 N +ATOM 564 CA LEU A 73 11.759 13.079 -7.764 1.00 52.18 C +ATOM 565 C LEU A 73 11.878 12.531 -6.348 1.00 52.18 C +ATOM 566 O LEU A 73 12.938 12.597 -5.734 1.00 52.18 O +ATOM 567 CB LEU A 73 12.249 12.016 -8.761 1.00 52.18 C +ATOM 568 CG LEU A 73 11.742 12.217 -10.194 1.00 52.18 C +ATOM 569 CD1 LEU A 73 12.417 11.195 -11.104 1.00 52.18 C +ATOM 570 CD2 LEU A 73 10.227 12.014 -10.293 1.00 52.18 C +ATOM 571 N ARG A 74 10.792 11.951 -5.847 1.00 52.56 N +ATOM 572 CA ARG A 74 10.731 11.376 -4.502 1.00 52.56 C +ATOM 573 C ARG A 74 10.187 9.959 -4.559 1.00 52.56 C +ATOM 574 O ARG A 74 9.497 9.609 -5.516 1.00 52.56 O +ATOM 575 CB ARG A 74 9.884 12.276 -3.600 1.00 52.56 C +ATOM 576 CG ARG A 74 10.269 13.754 -3.697 1.00 52.56 C +ATOM 577 CD ARG A 74 9.405 14.599 -2.781 1.00 52.56 C +ATOM 578 NE ARG A 74 9.786 16.010 -2.905 1.00 52.56 N +ATOM 579 CZ ARG A 74 9.273 16.997 -2.220 1.00 52.56 C +ATOM 580 NH1 ARG A 74 8.263 16.810 -1.407 1.00 52.56 N +ATOM 581 NH2 ARG A 74 9.801 18.177 -2.371 1.00 52.56 N +ATOM 582 N VAL A 75 10.450 9.142 -3.545 1.00 52.53 N +ATOM 583 CA VAL A 75 9.976 7.746 -3.509 1.00 52.53 C +ATOM 584 C VAL A 75 8.452 7.679 -3.662 1.00 52.53 C +ATOM 585 O VAL A 75 7.926 6.854 -4.419 1.00 52.53 O +ATOM 586 CB VAL A 75 10.441 7.038 -2.226 1.00 52.53 C +ATOM 587 CG1 VAL A 75 9.987 5.574 -2.208 1.00 52.53 C +ATOM 588 CG2 VAL A 75 11.972 7.060 -2.131 1.00 52.53 C +ATOM 589 N GLU A 76 7.736 8.623 -3.047 1.00 53.22 N +ATOM 590 CA GLU A 76 6.289 8.769 -3.165 1.00 53.22 C +ATOM 591 C GLU A 76 5.813 9.219 -4.546 1.00 53.22 C +ATOM 592 O GLU A 76 4.601 9.288 -4.733 1.00 53.22 O +ATOM 593 CB GLU A 76 5.718 9.707 -2.084 1.00 53.22 C +ATOM 594 CG GLU A 76 6.015 11.210 -2.266 1.00 53.22 C +ATOM 595 CD GLU A 76 7.331 11.668 -1.632 1.00 53.22 C +ATOM 596 OE1 GLU A 76 7.440 12.894 -1.388 1.00 53.22 O +ATOM 597 OE2 GLU A 76 8.220 10.811 -1.433 1.00 53.22 O +ATOM 598 N ASP A 77 6.703 9.530 -5.499 1.00 52.80 N +ATOM 599 CA ASP A 77 6.388 9.798 -6.908 1.00 52.80 C +ATOM 600 C ASP A 77 6.268 8.530 -7.760 1.00 52.80 C +ATOM 601 O ASP A 77 5.757 8.600 -8.883 1.00 52.80 O +ATOM 602 CB ASP A 77 7.330 10.811 -7.564 1.00 52.80 C +ATOM 603 CG ASP A 77 7.140 12.220 -7.017 1.00 52.80 C +ATOM 604 OD1 ASP A 77 5.982 12.656 -6.809 1.00 52.80 O +ATOM 605 OD2 ASP A 77 8.153 12.938 -6.898 1.00 52.80 O +ATOM 606 N SER A 78 6.622 7.369 -7.209 1.00 52.15 N +ATOM 607 CA SER A 78 6.479 6.077 -7.881 1.00 52.15 C +ATOM 608 C SER A 78 5.024 5.770 -8.255 1.00 52.15 C +ATOM 609 O SER A 78 4.070 6.114 -7.541 1.00 52.15 O +ATOM 610 CB SER A 78 7.053 4.949 -7.024 1.00 52.15 C +ATOM 611 OG SER A 78 8.418 5.192 -6.756 1.00 52.15 O +ATOM 612 N GLY A 79 4.833 5.127 -9.402 1.00 52.40 N +ATOM 613 CA GLY A 79 3.518 4.811 -9.951 1.00 52.40 C +ATOM 614 C GLY A 79 3.468 4.943 -11.468 1.00 52.40 C +ATOM 615 O GLY A 79 4.493 5.059 -12.134 1.00 52.40 O +ATOM 616 N THR A 80 2.261 4.943 -12.020 1.00 52.09 N +ATOM 617 CA THR A 80 2.038 5.018 -13.465 1.00 52.09 C +ATOM 618 C THR A 80 1.551 6.403 -13.853 1.00 52.09 C +ATOM 619 O THR A 80 0.552 6.884 -13.324 1.00 52.09 O +ATOM 620 CB THR A 80 1.063 3.937 -13.939 1.00 52.09 C +ATOM 621 OG1 THR A 80 1.562 2.676 -13.558 1.00 52.09 O +ATOM 622 CG2 THR A 80 0.923 3.912 -15.462 1.00 52.09 C +ATOM 623 N TYR A 81 2.230 7.029 -14.806 1.00 52.12 N +ATOM 624 CA TYR A 81 1.868 8.334 -15.351 1.00 52.12 C +ATOM 625 C TYR A 81 1.284 8.173 -16.744 1.00 52.12 C +ATOM 626 O TYR A 81 1.817 7.402 -17.538 1.00 52.12 O +ATOM 627 CB TYR A 81 3.101 9.235 -15.392 1.00 52.12 C +ATOM 628 CG TYR A 81 3.573 9.670 -14.024 1.00 52.12 C +ATOM 629 CD1 TYR A 81 3.278 10.963 -13.559 1.00 52.12 C +ATOM 630 CD2 TYR A 81 4.297 8.773 -13.216 1.00 52.12 C +ATOM 631 CE1 TYR A 81 3.719 11.366 -12.285 1.00 52.12 C +ATOM 632 CE2 TYR A 81 4.721 9.167 -11.938 1.00 52.12 C +ATOM 633 CZ TYR A 81 4.451 10.472 -11.476 1.00 52.12 C +ATOM 634 OH TYR A 81 4.884 10.862 -10.252 1.00 52.12 O +ATOM 635 N LYS A 82 0.232 8.914 -17.079 1.00 52.24 N +ATOM 636 CA LYS A 82 -0.352 8.949 -18.426 1.00 52.24 C +ATOM 637 C LYS A 82 -0.552 10.383 -18.876 1.00 52.24 C +ATOM 638 O LYS A 82 -0.943 11.237 -18.086 1.00 52.24 O +ATOM 639 CB LYS A 82 -1.677 8.180 -18.483 1.00 52.24 C +ATOM 640 CG LYS A 82 -1.480 6.667 -18.346 1.00 52.24 C +ATOM 641 CD LYS A 82 -2.838 5.976 -18.407 1.00 52.24 C +ATOM 642 CE LYS A 82 -2.704 4.477 -18.150 1.00 52.24 C +ATOM 643 NZ LYS A 82 -4.049 3.893 -17.940 1.00 52.24 N +ATOM 644 N CYS A 83 -0.322 10.627 -20.154 1.00 52.94 N +ATOM 645 CA CYS A 83 -0.692 11.879 -20.792 1.00 52.94 C +ATOM 646 C CYS A 83 -2.014 11.735 -21.547 1.00 52.94 C +ATOM 647 O CYS A 83 -2.379 10.641 -21.972 1.00 52.94 O +ATOM 648 CB CYS A 83 0.455 12.302 -21.687 1.00 52.94 C +ATOM 649 SG CYS A 83 0.900 11.121 -22.957 1.00 52.94 S +ATOM 650 N GLY A 84 -2.731 12.837 -21.710 1.00 54.23 N +ATOM 651 CA GLY A 84 -3.952 12.942 -22.497 1.00 54.23 C +ATOM 652 C GLY A 84 -3.827 14.073 -23.506 1.00 54.23 C +ATOM 653 O GLY A 84 -3.327 15.146 -23.170 1.00 54.23 O +ATOM 654 N ALA A 85 -4.276 13.827 -24.729 1.00 54.96 N +ATOM 655 CA ALA A 85 -4.416 14.822 -25.780 1.00 54.96 C +ATOM 656 C ALA A 85 -5.905 15.041 -26.051 1.00 54.96 C +ATOM 657 O ALA A 85 -6.640 14.086 -26.311 1.00 54.96 O +ATOM 658 CB ALA A 85 -3.671 14.334 -27.024 1.00 54.96 C +ATOM 659 N PHE A 86 -6.337 16.298 -26.011 1.00 59.79 N +ATOM 660 CA PHE A 86 -7.739 16.672 -26.140 1.00 59.79 C +ATOM 661 C PHE A 86 -7.946 17.584 -27.343 1.00 59.79 C +ATOM 662 O PHE A 86 -7.140 18.482 -27.624 1.00 59.79 O +ATOM 663 CB PHE A 86 -8.224 17.344 -24.852 1.00 59.79 C +ATOM 664 CG PHE A 86 -8.093 16.478 -23.623 1.00 59.79 C +ATOM 665 CD1 PHE A 86 -9.159 15.652 -23.226 1.00 59.79 C +ATOM 666 CD2 PHE A 86 -6.880 16.458 -22.910 1.00 59.79 C +ATOM 667 CE1 PHE A 86 -8.991 14.771 -22.147 1.00 59.79 C +ATOM 668 CE2 PHE A 86 -6.702 15.543 -21.864 1.00 59.79 C +ATOM 669 CZ PHE A 86 -7.747 14.678 -21.512 1.00 59.79 C +ATOM 670 N ARG A 87 -9.073 17.379 -28.024 1.00 71.53 N +ATOM 671 CA ARG A 87 -9.612 18.296 -29.031 1.00 71.53 C +ATOM 672 C ARG A 87 -10.963 18.826 -28.564 1.00 71.53 C +ATOM 673 O ARG A 87 -11.807 18.054 -28.117 1.00 71.53 O +ATOM 674 CB ARG A 87 -9.716 17.570 -30.381 1.00 71.53 C +ATOM 675 CG ARG A 87 -10.288 18.457 -31.503 1.00 71.53 C +ATOM 676 CD ARG A 87 -10.497 17.614 -32.760 1.00 71.53 C +ATOM 677 NE ARG A 87 -11.175 18.312 -33.877 1.00 71.53 N +ATOM 678 CZ ARG A 87 -10.691 19.204 -34.724 1.00 71.53 C +ATOM 679 NH1 ARG A 87 -9.539 19.782 -34.564 1.00 71.53 N +ATOM 680 NH2 ARG A 87 -11.371 19.567 -35.773 1.00 71.53 N +ATOM 681 N SER A 100 -10.711 12.500 -26.368 1.00 75.07 N +ATOM 682 CA SER A 100 -9.606 12.365 -25.425 1.00 75.07 C +ATOM 683 C SER A 100 -8.804 11.120 -25.765 1.00 75.07 C +ATOM 684 O SER A 100 -9.229 10.011 -25.444 1.00 75.07 O +ATOM 685 CB SER A 100 -10.165 12.258 -24.004 1.00 75.07 C +ATOM 686 OG SER A 100 -9.156 11.838 -23.101 1.00 75.07 O +ATOM 687 N GLU A 101 -7.613 11.297 -26.322 1.00 57.63 N +ATOM 688 CA GLU A 101 -6.672 10.196 -26.527 1.00 57.63 C +ATOM 689 C GLU A 101 -5.696 10.144 -25.351 1.00 57.63 C +ATOM 690 O GLU A 101 -4.983 11.116 -25.095 1.00 57.63 O +ATOM 691 CB GLU A 101 -5.938 10.362 -27.867 1.00 57.63 C +ATOM 692 CG GLU A 101 -6.856 10.343 -29.099 1.00 57.63 C +ATOM 693 CD GLU A 101 -7.627 9.026 -29.285 1.00 57.63 C +ATOM 694 OE1 GLU A 101 -8.585 9.053 -30.089 1.00 57.63 O +ATOM 695 OE2 GLU A 101 -7.223 8.021 -28.657 1.00 57.63 O +ATOM 696 N LYS A 102 -5.653 9.026 -24.617 1.00 54.69 N +ATOM 697 CA LYS A 102 -4.690 8.818 -23.522 1.00 54.69 C +ATOM 698 C LYS A 102 -3.501 7.983 -23.990 1.00 54.69 C +ATOM 699 O LYS A 102 -3.648 7.044 -24.767 1.00 54.69 O +ATOM 700 CB LYS A 102 -5.357 8.218 -22.271 1.00 54.69 C +ATOM 701 CG LYS A 102 -6.273 9.217 -21.540 1.00 54.69 C +ATOM 702 CD LYS A 102 -6.785 8.629 -20.213 1.00 54.69 C +ATOM 703 CE LYS A 102 -7.642 9.646 -19.448 1.00 54.69 C +ATOM 704 NZ LYS A 102 -8.100 9.113 -18.138 1.00 54.69 N +ATOM 705 N GLY A 103 -2.316 8.300 -23.483 1.00 54.06 N +ATOM 706 CA GLY A 103 -1.117 7.498 -23.685 1.00 54.06 C +ATOM 707 C GLY A 103 -1.180 6.184 -22.907 1.00 54.06 C +ATOM 708 O GLY A 103 -1.830 6.081 -21.864 1.00 54.06 O +ATOM 709 N ALA A 104 -0.464 5.165 -23.392 1.00 53.55 N +ATOM 710 CA ALA A 104 -0.459 3.849 -22.734 1.00 53.55 C +ATOM 711 C ALA A 104 0.223 3.842 -21.346 1.00 53.55 C +ATOM 712 O ALA A 104 0.059 2.893 -20.584 1.00 53.55 O +ATOM 713 CB ALA A 104 0.106 2.772 -23.676 1.00 53.55 C +ATOM 714 N GLY A 105 0.938 4.914 -20.993 1.00 52.73 N +ATOM 715 CA GLY A 105 1.519 5.111 -19.671 1.00 52.73 C +ATOM 716 C GLY A 105 3.029 4.903 -19.579 1.00 52.73 C +ATOM 717 O GLY A 105 3.655 4.312 -20.463 1.00 52.73 O +ATOM 718 N THR A 106 3.582 5.426 -18.485 1.00 52.06 N +ATOM 719 CA THR A 106 4.975 5.319 -18.046 1.00 52.06 C +ATOM 720 C THR A 106 4.976 4.810 -16.614 1.00 52.06 C +ATOM 721 O THR A 106 4.476 5.502 -15.726 1.00 52.06 O +ATOM 722 CB THR A 106 5.691 6.681 -18.070 1.00 52.06 C +ATOM 723 OG1 THR A 106 5.652 7.248 -19.356 1.00 52.06 O +ATOM 724 CG2 THR A 106 7.155 6.577 -17.635 1.00 52.06 C +ATOM 725 N VAL A 107 5.539 3.630 -16.376 1.00 52.09 N +ATOM 726 CA VAL A 107 5.760 3.105 -15.025 1.00 52.09 C +ATOM 727 C VAL A 107 7.047 3.715 -14.478 1.00 52.09 C +ATOM 728 O VAL A 107 8.125 3.470 -15.017 1.00 52.09 O +ATOM 729 CB VAL A 107 5.812 1.566 -15.021 1.00 52.09 C +ATOM 730 CG1 VAL A 107 5.985 1.036 -13.593 1.00 52.09 C +ATOM 731 CG2 VAL A 107 4.516 0.976 -15.596 1.00 52.09 C +ATOM 732 N LEU A 108 6.931 4.532 -13.437 1.00 52.09 N +ATOM 733 CA LEU A 108 8.050 5.153 -12.742 1.00 52.09 C +ATOM 734 C LEU A 108 8.297 4.449 -11.406 1.00 52.09 C +ATOM 735 O LEU A 108 7.401 4.382 -10.561 1.00 52.09 O +ATOM 736 CB LEU A 108 7.773 6.655 -12.558 1.00 52.09 C +ATOM 737 CG LEU A 108 8.841 7.391 -11.724 1.00 52.09 C +ATOM 738 CD1 LEU A 108 10.249 7.278 -12.322 1.00 52.09 C +ATOM 739 CD2 LEU A 108 8.458 8.863 -11.595 1.00 52.09 C +ATOM 740 N THR A 109 9.536 4.021 -11.204 1.00 52.21 N +ATOM 741 CA THR A 109 10.065 3.575 -9.914 1.00 52.21 C +ATOM 742 C THR A 109 11.139 4.557 -9.459 1.00 52.21 C +ATOM 743 O THR A 109 12.126 4.749 -10.169 1.00 52.21 O +ATOM 744 CB THR A 109 10.647 2.160 -10.036 1.00 52.21 C +ATOM 745 OG1 THR A 109 9.661 1.273 -10.515 1.00 52.21 O +ATOM 746 CG2 THR A 109 11.122 1.608 -8.692 1.00 52.21 C +ATOM 747 N VAL A 110 10.953 5.172 -8.294 1.00 52.24 N +ATOM 748 CA VAL A 110 11.960 6.037 -7.665 1.00 52.24 C +ATOM 749 C VAL A 110 12.599 5.295 -6.494 1.00 52.24 C +ATOM 750 O VAL A 110 11.874 4.704 -5.690 1.00 52.24 O +ATOM 751 CB VAL A 110 11.392 7.402 -7.240 1.00 52.24 C +ATOM 752 CG1 VAL A 110 12.515 8.335 -6.766 1.00 52.24 C +ATOM 753 CG2 VAL A 110 10.661 8.093 -8.402 1.00 52.24 C +ATOM 754 N LYS A 111 13.930 5.294 -6.442 1.00 53.86 N +ATOM 755 CA LYS A 111 14.738 4.691 -5.378 1.00 53.86 C +ATOM 756 C LYS A 111 15.293 5.735 -4.411 1.00 53.86 C +ATOM 757 O LYS A 111 15.460 6.915 -4.814 1.00 53.86 O +ATOM 758 CB LYS A 111 15.848 3.825 -5.980 1.00 53.86 C +ATOM 759 CG LYS A 111 15.296 2.604 -6.728 1.00 53.86 C +ATOM 760 CD LYS A 111 16.476 1.780 -7.241 1.00 53.86 C +ATOM 761 CE LYS A 111 16.000 0.549 -8.010 1.00 53.86 C +ATOM 762 NZ LYS A 111 17.180 -0.193 -8.509 1.00 53.86 N +ATOM 763 OXT LYS A 111 15.442 5.319 -3.248 1.00 53.86 O +END diff --git a/tests/files/pdb/6JZA_af.pdb b/tests/files/pdb/6JZA_af.pdb new file mode 100644 index 00000000..1289d872 --- /dev/null +++ b/tests/files/pdb/6JZA_af.pdb @@ -0,0 +1,565 @@ +REMARK TITLE 6JZA AlphaFold-start MR +REMARK Log-Likelihood Gain: 290.062 +REMARK RFZ=5.0 TFZ=10.1 PAK=3 LLG=90 TFZ==11.1 LLG=290 TFZ==18.5 PAK=3 LLG=290 TFZ==18.5 +REMARK ENSEMBLE e_Q62356 EULER 10.07 141.28 335.28 FRAC 1.438 0.406 -0.004 +CRYST1 45.731 45.731 68.966 90.00 90.00 120.00 P 62 6 +SCALE1 0.021867 0.012625 -0.000000 0.00000 +SCALE2 0.000000 0.025250 -0.000000 0.00000 +SCALE3 0.000000 0.000000 0.014500 0.00000 +ATOM 1 N SER A 24 -14.449 0.658 7.069 1.00 71.34 N +ATOM 2 CA SER A 24 -13.548 1.394 6.170 1.00 71.34 C +ATOM 3 C SER A 24 -12.513 2.289 6.872 1.00 71.34 C +ATOM 4 O SER A 24 -12.870 3.179 7.657 1.00 71.34 O +ATOM 5 CB SER A 24 -14.369 2.231 5.184 1.00 71.34 C +ATOM 6 OG SER A 24 -13.494 2.913 4.305 1.00 71.34 O +ATOM 7 N LYS A 25 -11.232 2.116 6.489 1.00 45.13 N +ATOM 8 CA LYS A 25 -10.060 2.914 6.921 1.00 45.13 C +ATOM 9 C LYS A 25 -10.339 4.420 6.866 1.00 45.13 C +ATOM 10 O LYS A 25 -9.947 5.143 7.779 1.00 45.13 O +ATOM 11 CB LYS A 25 -8.827 2.524 6.054 1.00 45.13 C +ATOM 12 CG LYS A 25 -7.497 3.241 6.407 1.00 45.13 C +ATOM 13 CD LYS A 25 -6.283 2.810 5.534 1.00 45.13 C +ATOM 14 CE LYS A 25 -5.005 3.574 5.955 1.00 45.13 C +ATOM 15 NZ LYS A 25 -3.740 3.241 5.218 1.00 45.13 N +ATOM 16 N SER A 26 -11.061 4.894 5.845 1.00 42.41 N +ATOM 17 CA SER A 26 -11.353 6.324 5.670 1.00 42.41 C +ATOM 18 C SER A 26 -12.208 6.917 6.792 1.00 42.41 C +ATOM 19 O SER A 26 -11.999 8.070 7.153 1.00 42.41 O +ATOM 20 CB SER A 26 -12.029 6.587 4.320 1.00 42.41 C +ATOM 21 OG SER A 26 -13.296 5.959 4.240 1.00 42.41 O +ATOM 22 N LYS A 27 -13.132 6.144 7.380 1.00 35.39 N +ATOM 23 CA LYS A 27 -14.018 6.618 8.455 1.00 35.39 C +ATOM 24 C LYS A 27 -13.308 6.669 9.808 1.00 35.39 C +ATOM 25 O LYS A 27 -13.566 7.577 10.591 1.00 35.39 O +ATOM 26 CB LYS A 27 -15.282 5.742 8.486 1.00 35.39 C +ATOM 27 CG LYS A 27 -16.377 6.319 9.399 1.00 35.39 C +ATOM 28 CD LYS A 27 -17.669 5.497 9.292 1.00 35.39 C +ATOM 29 CE LYS A 27 -18.774 6.110 10.162 1.00 35.39 C +ATOM 30 NZ LYS A 27 -20.050 5.357 10.032 1.00 35.39 N +ATOM 31 N ILE A 28 -12.412 5.716 10.068 1.00 34.55 N +ATOM 32 CA ILE A 28 -11.653 5.621 11.326 1.00 34.55 C +ATOM 33 C ILE A 28 -10.507 6.645 11.345 1.00 34.55 C +ATOM 34 O ILE A 28 -10.301 7.325 12.346 1.00 34.55 O +ATOM 35 CB ILE A 28 -11.166 4.166 11.541 1.00 34.55 C +ATOM 36 CG1 ILE A 28 -12.378 3.204 11.629 1.00 34.55 C +ATOM 37 CG2 ILE A 28 -10.308 4.062 12.816 1.00 34.55 C +ATOM 38 CD1 ILE A 28 -11.999 1.720 11.605 1.00 34.55 C +ATOM 39 N CYS A 29 -9.807 6.813 10.221 1.00 31.90 N +ATOM 40 CA CYS A 29 -8.683 7.743 10.095 1.00 31.90 C +ATOM 41 C CYS A 29 -9.076 9.192 9.738 1.00 31.90 C +ATOM 42 O CYS A 29 -8.187 10.033 9.626 1.00 31.90 O +ATOM 43 CB CYS A 29 -7.683 7.166 9.087 1.00 31.90 C +ATOM 44 SG CYS A 29 -6.793 5.694 9.661 1.00 31.90 S +ATOM 45 N ALA A 30 -10.367 9.508 9.560 1.00 34.07 N +ATOM 46 CA ALA A 30 -10.842 10.787 9.007 1.00 34.07 C +ATOM 47 C ALA A 30 -10.264 12.042 9.691 1.00 34.07 C +ATOM 48 O ALA A 30 -9.937 13.014 9.017 1.00 34.07 O +ATOM 49 CB ALA A 30 -12.375 10.814 9.107 1.00 34.07 C +ATOM 50 N ASN A 31 -10.119 11.998 11.020 1.00 34.94 N +ATOM 51 CA ASN A 31 -9.646 13.111 11.850 1.00 34.94 C +ATOM 52 C ASN A 31 -8.299 12.803 12.538 1.00 34.94 C +ATOM 53 O ASN A 31 -7.971 13.407 13.558 1.00 34.94 O +ATOM 54 CB ASN A 31 -10.759 13.490 12.847 1.00 34.94 C +ATOM 55 CG ASN A 31 -12.036 13.978 12.183 1.00 34.94 C +ATOM 56 OD1 ASN A 31 -12.050 14.577 11.126 1.00 34.94 O +ATOM 57 ND2 ASN A 31 -13.174 13.747 12.793 1.00 34.94 N +ATOM 58 N VAL A 32 -7.532 11.832 12.027 1.00 32.98 N +ATOM 59 CA VAL A 32 -6.283 11.367 12.649 1.00 32.98 C +ATOM 60 C VAL A 32 -5.085 12.022 11.971 1.00 32.98 C +ATOM 61 O VAL A 32 -4.778 11.738 10.815 1.00 32.98 O +ATOM 62 CB VAL A 32 -6.175 9.832 12.620 1.00 32.98 C +ATOM 63 CG1 VAL A 32 -4.892 9.353 13.316 1.00 32.98 C +ATOM 64 CG2 VAL A 32 -7.371 9.193 13.336 1.00 32.98 C +ATOM 65 N PHE A 33 -4.383 12.888 12.703 1.00 32.98 N +ATOM 66 CA PHE A 33 -3.172 13.544 12.214 1.00 32.98 C +ATOM 67 C PHE A 33 -1.918 12.759 12.619 1.00 32.98 C +ATOM 68 O PHE A 33 -1.538 12.736 13.789 1.00 32.98 O +ATOM 69 CB PHE A 33 -3.137 14.997 12.699 1.00 32.98 C +ATOM 70 CG PHE A 33 -1.996 15.785 12.088 1.00 32.98 C +ATOM 71 CD1 PHE A 33 -0.781 15.934 12.785 1.00 32.98 C +ATOM 72 CD2 PHE A 33 -2.139 16.341 10.804 1.00 32.98 C +ATOM 73 CE1 PHE A 33 0.283 16.644 12.201 1.00 32.98 C +ATOM 74 CE2 PHE A 33 -1.074 17.051 10.220 1.00 32.98 C +ATOM 75 CZ PHE A 33 0.136 17.205 10.919 1.00 32.98 C +ATOM 76 N CYS A 34 -1.252 12.148 11.640 1.00 32.50 N +ATOM 77 CA CYS A 34 0.053 11.520 11.829 1.00 32.50 C +ATOM 78 C CYS A 34 1.177 12.485 11.422 1.00 32.50 C +ATOM 79 O CYS A 34 1.078 13.180 10.412 1.00 32.50 O +ATOM 80 CB CYS A 34 0.106 10.192 11.066 1.00 32.50 C +ATOM 81 SG CYS A 34 -1.128 8.968 11.592 1.00 32.50 S +ATOM 82 N GLY A 35 2.259 12.523 12.207 1.00 34.43 N +ATOM 83 CA GLY A 35 3.436 13.339 11.894 1.00 34.43 C +ATOM 84 C GLY A 35 4.155 12.892 10.613 1.00 34.43 C +ATOM 85 O GLY A 35 3.933 11.789 10.116 1.00 34.43 O +ATOM 86 N ALA A 36 5.055 13.730 10.094 1.00 34.49 N +ATOM 87 CA ALA A 36 5.772 13.453 8.848 1.00 34.49 C +ATOM 88 C ALA A 36 6.441 12.061 8.842 1.00 34.49 C +ATOM 89 O ALA A 36 6.994 11.614 9.852 1.00 34.49 O +ATOM 90 CB ALA A 36 6.784 14.575 8.592 1.00 34.49 C +ATOM 91 N GLY A 37 6.367 11.375 7.696 1.00 33.83 N +ATOM 92 CA GLY A 37 6.853 10.000 7.536 1.00 33.83 C +ATOM 93 C GLY A 37 5.910 8.912 8.060 1.00 33.83 C +ATOM 94 O GLY A 37 6.257 7.735 7.964 1.00 33.83 O +ATOM 95 N ARG A 38 4.734 9.274 8.596 1.00 30.29 N +ATOM 96 CA ARG A 38 3.727 8.337 9.117 1.00 30.29 C +ATOM 97 C ARG A 38 2.366 8.491 8.440 1.00 30.29 C +ATOM 98 O ARG A 38 1.992 9.584 8.027 1.00 30.29 O +ATOM 99 CB ARG A 38 3.588 8.481 10.642 1.00 30.29 C +ATOM 100 CG ARG A 38 4.874 8.105 11.383 1.00 30.29 C +ATOM 101 CD ARG A 38 4.719 8.353 12.886 1.00 30.29 C +ATOM 102 NE ARG A 38 5.971 8.070 13.614 1.00 30.29 N +ATOM 103 CZ ARG A 38 7.040 8.845 13.694 1.00 30.29 C +ATOM 104 NH1 ARG A 38 7.106 10.007 13.100 1.00 30.29 N +ATOM 105 NH2 ARG A 38 8.084 8.449 14.365 1.00 30.29 N +ATOM 106 N GLU A 39 1.615 7.399 8.385 1.00 30.88 N +ATOM 107 CA GLU A 39 0.220 7.356 7.941 1.00 30.88 C +ATOM 108 C GLU A 39 -0.675 6.675 8.981 1.00 30.88 C +ATOM 109 O GLU A 39 -0.208 5.886 9.806 1.00 30.88 O +ATOM 110 CB GLU A 39 0.086 6.694 6.560 1.00 30.88 C +ATOM 111 CG GLU A 39 0.422 5.194 6.533 1.00 30.88 C +ATOM 112 CD GLU A 39 0.177 4.619 5.133 1.00 30.88 C +ATOM 113 OE1 GLU A 39 1.182 4.287 4.463 1.00 30.88 O +ATOM 114 OE2 GLU A 39 -1.025 4.560 4.750 1.00 30.88 O +ATOM 115 N CYS A 40 -1.972 6.983 8.947 1.00 30.58 N +ATOM 116 CA CYS A 40 -2.944 6.321 9.808 1.00 30.58 C +ATOM 117 C CYS A 40 -3.281 4.932 9.253 1.00 30.58 C +ATOM 118 O CYS A 40 -3.612 4.796 8.076 1.00 30.58 O +ATOM 119 CB CYS A 40 -4.183 7.204 9.986 1.00 30.58 C +ATOM 120 SG CYS A 40 -5.493 6.484 11.016 1.00 30.58 S +ATOM 121 N ALA A 41 -3.244 3.925 10.117 1.00 32.14 N +ATOM 122 CA ALA A 41 -3.722 2.566 9.902 1.00 32.14 C +ATOM 123 C ALA A 41 -4.756 2.208 10.983 1.00 32.14 C +ATOM 124 O ALA A 41 -4.919 2.930 11.967 1.00 32.14 O +ATOM 125 CB ALA A 41 -2.516 1.621 9.896 1.00 32.14 C +ATOM 126 N VAL A 42 -5.465 1.097 10.795 1.00 32.40 N +ATOM 127 CA VAL A 42 -6.477 0.602 11.737 1.00 32.40 C +ATOM 128 C VAL A 42 -5.957 -0.689 12.361 1.00 32.40 C +ATOM 129 O VAL A 42 -5.517 -1.578 11.637 1.00 32.40 O +ATOM 130 CB VAL A 42 -7.833 0.401 11.033 1.00 32.40 C +ATOM 131 CG1 VAL A 42 -8.916 -0.054 12.012 1.00 32.40 C +ATOM 132 CG2 VAL A 42 -8.317 1.705 10.380 1.00 32.40 C +ATOM 133 N THR A 43 -5.972 -0.783 13.689 1.00 34.61 N +ATOM 134 CA THR A 43 -5.578 -2.000 14.417 1.00 34.61 C +ATOM 135 C THR A 43 -6.641 -3.093 14.290 1.00 34.61 C +ATOM 136 O THR A 43 -7.794 -2.817 13.958 1.00 34.61 O +ATOM 137 CB THR A 43 -5.317 -1.710 15.907 1.00 34.61 C +ATOM 138 OG1 THR A 43 -6.512 -1.350 16.567 1.00 34.61 O +ATOM 139 CG2 THR A 43 -4.292 -0.599 16.138 1.00 34.61 C +ATOM 140 N GLU A 44 -6.304 -4.332 14.658 1.00 35.06 N +ATOM 141 CA GLU A 44 -7.265 -5.448 14.758 1.00 35.06 C +ATOM 142 C GLU A 44 -8.467 -5.138 15.677 1.00 35.06 C +ATOM 143 O GLU A 44 -9.536 -5.721 15.527 1.00 35.06 O +ATOM 144 CB GLU A 44 -6.534 -6.685 15.303 1.00 35.06 C +ATOM 145 CG GLU A 44 -5.388 -7.165 14.394 1.00 35.06 C +ATOM 146 CD GLU A 44 -4.625 -8.367 14.976 1.00 35.06 C +ATOM 147 OE1 GLU A 44 -3.838 -8.959 14.208 1.00 35.06 O +ATOM 148 OE2 GLU A 44 -4.790 -8.644 16.187 1.00 35.06 O +ATOM 149 N LYS A 45 -8.322 -4.178 16.603 1.00 36.49 N +ATOM 150 CA LYS A 45 -9.389 -3.702 17.501 1.00 36.49 C +ATOM 151 C LYS A 45 -10.303 -2.639 16.875 1.00 36.49 C +ATOM 152 O LYS A 45 -11.256 -2.200 17.513 1.00 36.49 O +ATOM 153 CB LYS A 45 -8.765 -3.176 18.802 1.00 36.49 C +ATOM 154 CG LYS A 45 -7.976 -4.253 19.557 1.00 36.49 C +ATOM 155 CD LYS A 45 -7.431 -3.654 20.853 1.00 36.49 C +ATOM 156 CE LYS A 45 -6.629 -4.700 21.624 1.00 36.49 C +ATOM 157 NZ LYS A 45 -6.074 -4.094 22.854 1.00 36.49 N +ATOM 158 N GLY A 46 -10.022 -2.196 15.648 1.00 37.16 N +ATOM 159 CA GLY A 46 -10.744 -1.106 14.988 1.00 37.16 C +ATOM 160 C GLY A 46 -10.334 0.298 15.452 1.00 37.16 C +ATOM 161 O GLY A 46 -11.093 1.244 15.242 1.00 37.16 O +ATOM 162 N GLU A 47 -9.165 0.442 16.082 1.00 35.46 N +ATOM 163 CA GLU A 47 -8.645 1.711 16.613 1.00 35.46 C +ATOM 164 C GLU A 47 -7.634 2.343 15.634 1.00 35.46 C +ATOM 165 O GLU A 47 -6.886 1.611 14.979 1.00 35.46 O +ATOM 166 CB GLU A 47 -7.983 1.490 17.988 1.00 35.46 C +ATOM 167 CG GLU A 47 -8.937 0.910 19.049 1.00 35.46 C +ATOM 168 CD GLU A 47 -8.221 0.465 20.337 1.00 35.46 C +ATOM 169 OE1 GLU A 47 -8.859 0.530 21.412 1.00 35.46 O +ATOM 170 OE2 GLU A 47 -7.067 -0.022 20.247 1.00 35.46 O +ATOM 171 N PRO A 48 -7.564 3.681 15.510 1.00 32.55 N +ATOM 172 CA PRO A 48 -6.562 4.338 14.676 1.00 32.55 C +ATOM 173 C PRO A 48 -5.164 4.297 15.310 1.00 32.55 C +ATOM 174 O PRO A 48 -4.997 4.567 16.497 1.00 32.55 O +ATOM 175 CB PRO A 48 -7.050 5.774 14.516 1.00 32.55 C +ATOM 176 CG PRO A 48 -7.835 6.033 15.803 1.00 32.55 C +ATOM 177 CD PRO A 48 -8.429 4.667 16.143 1.00 32.55 C +ATOM 178 N THR A 49 -4.139 4.046 14.497 1.00 31.28 N +ATOM 179 CA THR A 49 -2.721 4.107 14.890 1.00 31.28 C +ATOM 180 C THR A 49 -1.874 4.752 13.791 1.00 31.28 C +ATOM 181 O THR A 49 -2.222 4.665 12.618 1.00 31.28 O +ATOM 182 CB THR A 49 -2.201 2.712 15.269 1.00 31.28 C +ATOM 183 OG1 THR A 49 -0.880 2.822 15.743 1.00 31.28 O +ATOM 184 CG2 THR A 49 -2.198 1.698 14.123 1.00 31.28 C +ATOM 185 N CYS A 50 -0.764 5.406 14.145 1.00 31.37 N +ATOM 186 CA CYS A 50 0.148 6.032 13.182 1.00 31.37 C +ATOM 187 C CYS A 50 1.393 5.162 12.965 1.00 31.37 C +ATOM 188 O CYS A 50 2.279 5.122 13.820 1.00 31.37 O +ATOM 189 CB CYS A 50 0.530 7.441 13.653 1.00 31.37 C +ATOM 190 SG CYS A 50 -0.793 8.678 13.584 1.00 31.37 S +ATOM 191 N LEU A 51 1.482 4.513 11.805 1.00 30.16 N +ATOM 192 CA LEU A 51 2.615 3.674 11.395 1.00 30.16 C +ATOM 193 C LEU A 51 3.510 4.430 10.405 1.00 30.16 C +ATOM 194 O LEU A 51 3.098 5.450 9.859 1.00 30.16 O +ATOM 195 CB LEU A 51 2.090 2.353 10.802 1.00 30.16 C +ATOM 196 CG LEU A 51 1.162 1.552 11.735 1.00 30.16 C +ATOM 197 CD1 LEU A 51 0.739 0.252 11.055 1.00 30.16 C +ATOM 198 CD2 LEU A 51 1.822 1.198 13.072 1.00 30.16 C +ATOM 199 N CYS A 52 4.737 3.962 10.168 1.00 29.93 N +ATOM 200 CA CYS A 52 5.580 4.527 9.108 1.00 29.93 C +ATOM 201 C CYS A 52 4.925 4.313 7.734 1.00 29.93 C +ATOM 202 O CYS A 52 4.349 3.255 7.500 1.00 29.93 O +ATOM 203 CB CYS A 52 6.984 3.910 9.157 1.00 29.93 C +ATOM 204 SG CYS A 52 7.830 4.040 10.759 1.00 29.93 S +ATOM 205 N ILE A 53 5.023 5.303 6.839 1.00 30.29 N +ATOM 206 CA ILE A 53 4.440 5.230 5.487 1.00 30.29 C +ATOM 207 C ILE A 53 4.979 4.001 4.750 1.00 30.29 C +ATOM 208 O ILE A 53 6.196 3.837 4.638 1.00 30.29 O +ATOM 209 CB ILE A 53 4.722 6.536 4.700 1.00 30.29 C +ATOM 210 CG1 ILE A 53 3.884 7.686 5.297 1.00 30.29 C +ATOM 211 CG2 ILE A 53 4.415 6.394 3.196 1.00 30.29 C +ATOM 212 CD1 ILE A 53 4.285 9.087 4.821 1.00 30.29 C +ATOM 213 N GLU A 54 4.091 3.174 4.197 1.00 31.75 N +ATOM 214 CA GLU A 54 4.504 1.959 3.484 1.00 31.75 C +ATOM 215 C GLU A 54 5.297 2.311 2.220 1.00 31.75 C +ATOM 216 O GLU A 54 6.438 1.869 2.059 1.00 31.75 O +ATOM 217 CB GLU A 54 3.292 1.088 3.129 1.00 31.75 C +ATOM 218 CG GLU A 54 2.658 0.440 4.370 1.00 31.75 C +ATOM 219 CD GLU A 54 1.479 -0.488 4.029 1.00 31.75 C +ATOM 220 OE1 GLU A 54 0.945 -1.098 4.983 1.00 31.75 O +ATOM 221 OE2 GLU A 54 1.110 -0.584 2.835 1.00 31.75 O +ATOM 222 N GLN A 55 4.732 3.180 1.372 1.00 30.66 N +ATOM 223 CA GLN A 55 5.321 3.607 0.103 1.00 30.66 C +ATOM 224 C GLN A 55 5.095 5.105 -0.151 1.00 30.66 C +ATOM 225 O GLN A 55 3.965 5.583 -0.253 1.00 30.66 O +ATOM 226 CB GLN A 55 4.731 2.741 -1.026 1.00 30.66 C +ATOM 227 CG GLN A 55 5.326 3.026 -2.415 1.00 30.66 C +ATOM 228 CD GLN A 55 6.843 2.872 -2.443 1.00 30.66 C +ATOM 229 OE1 GLN A 55 7.391 1.813 -2.207 1.00 30.66 O +ATOM 230 NE2 GLN A 55 7.578 3.937 -2.685 1.00 30.66 N +ATOM 231 N CYS A 56 6.182 5.859 -0.323 1.00 30.75 N +ATOM 232 CA CYS A 56 6.113 7.262 -0.728 1.00 30.75 C +ATOM 233 C CYS A 56 5.858 7.409 -2.237 1.00 30.75 C +ATOM 234 O CYS A 56 6.242 6.561 -3.047 1.00 30.75 O +ATOM 235 CB CYS A 56 7.400 7.966 -0.299 1.00 30.75 C +ATOM 236 SG CYS A 56 7.523 8.269 1.480 1.00 30.75 S +ATOM 237 N LYS A 57 5.240 8.531 -2.636 1.00 30.92 N +ATOM 238 CA LYS A 57 5.068 8.881 -4.055 1.00 30.92 C +ATOM 239 C LYS A 57 6.439 9.180 -4.694 1.00 30.92 C +ATOM 240 O LYS A 57 7.204 9.951 -4.109 1.00 30.92 O +ATOM 241 CB LYS A 57 4.113 10.074 -4.228 1.00 30.92 C +ATOM 242 CG LYS A 57 2.670 9.740 -3.823 1.00 30.92 C +ATOM 243 CD LYS A 57 1.738 10.939 -4.049 1.00 30.92 C +ATOM 244 CE LYS A 57 0.315 10.568 -3.613 1.00 30.92 C +ATOM 245 NZ LYS A 57 -0.626 11.707 -3.759 1.00 30.92 N +ATOM 246 N PRO A 58 6.744 8.669 -5.904 1.00 32.25 N +ATOM 247 CA PRO A 58 8.070 8.766 -6.527 1.00 32.25 C +ATOM 248 C PRO A 58 8.342 10.136 -7.185 1.00 32.25 C +ATOM 249 O PRO A 58 8.824 10.231 -8.313 1.00 32.25 O +ATOM 250 CB PRO A 58 8.129 7.578 -7.498 1.00 32.25 C +ATOM 251 CG PRO A 58 6.682 7.471 -7.971 1.00 32.25 C +ATOM 252 CD PRO A 58 5.899 7.782 -6.695 1.00 32.25 C +ATOM 253 N HIS A 59 8.017 11.233 -6.499 1.00 32.25 N +ATOM 254 CA HIS A 59 8.328 12.580 -6.973 1.00 32.25 C +ATOM 255 C HIS A 59 9.828 12.855 -6.825 1.00 32.25 C +ATOM 256 O HIS A 59 10.356 12.794 -5.717 1.00 32.25 O +ATOM 257 CB HIS A 59 7.489 13.618 -6.213 1.00 32.25 C +ATOM 258 CG HIS A 59 6.001 13.472 -6.422 1.00 32.25 C +ATOM 259 ND1 HIS A 59 5.338 13.469 -7.629 1.00 32.25 N +ATOM 260 CD2 HIS A 59 5.055 13.309 -5.446 1.00 32.25 C +ATOM 261 CE1 HIS A 59 4.027 13.305 -7.383 1.00 32.25 C +ATOM 262 NE2 HIS A 59 3.804 13.199 -6.063 1.00 32.25 N +ATOM 263 N LYS A 60 10.507 13.198 -7.927 1.00 33.09 N +ATOM 264 CA LYS A 60 11.918 13.609 -7.930 1.00 33.09 C +ATOM 265 C LYS A 60 12.045 15.058 -7.445 1.00 33.09 C +ATOM 266 O LYS A 60 11.944 15.987 -8.242 1.00 33.09 O +ATOM 267 CB LYS A 60 12.525 13.423 -9.335 1.00 33.09 C +ATOM 268 CG LYS A 60 12.644 11.945 -9.740 1.00 33.09 C +ATOM 269 CD LYS A 60 13.327 11.775 -11.103 1.00 33.09 C +ATOM 270 CE LYS A 60 13.480 10.282 -11.424 1.00 33.09 C +ATOM 271 NZ LYS A 60 14.259 10.054 -12.669 1.00 33.09 N +ATOM 272 N ARG A 61 12.227 15.244 -6.138 1.00 30.75 N +ATOM 273 CA ARG A 61 12.522 16.535 -5.492 1.00 30.75 C +ATOM 274 C ARG A 61 13.591 16.284 -4.426 1.00 30.75 C +ATOM 275 O ARG A 61 13.229 16.181 -3.258 1.00 30.75 O +ATOM 276 CB ARG A 61 11.257 17.160 -4.881 1.00 30.75 C +ATOM 277 CG ARG A 61 10.282 17.688 -5.939 1.00 30.75 C +ATOM 278 CD ARG A 61 9.179 18.494 -5.247 1.00 30.75 C +ATOM 279 NE ARG A 61 8.215 19.040 -6.220 1.00 30.75 N +ATOM 280 CZ ARG A 61 7.236 19.886 -5.945 1.00 30.75 C +ATOM 281 NH1 ARG A 61 7.023 20.329 -4.737 1.00 30.75 N +ATOM 282 NH2 ARG A 61 6.447 20.310 -6.891 1.00 30.75 N +ATOM 283 N PRO A 62 14.853 16.061 -4.833 1.00 31.10 N +ATOM 284 CA PRO A 62 15.883 15.557 -3.935 1.00 31.10 C +ATOM 285 C PRO A 62 16.106 16.509 -2.760 1.00 31.10 C +ATOM 286 O PRO A 62 16.022 17.724 -2.932 1.00 31.10 O +ATOM 287 CB PRO A 62 17.141 15.374 -4.792 1.00 31.10 C +ATOM 288 CG PRO A 62 16.929 16.379 -5.924 1.00 31.10 C +ATOM 289 CD PRO A 62 15.419 16.332 -6.144 1.00 31.10 C +ATOM 290 N VAL A 63 16.389 15.946 -1.588 1.00 29.66 N +ATOM 291 CA VAL A 63 16.729 16.689 -0.366 1.00 29.66 C +ATOM 292 C VAL A 63 17.931 16.047 0.317 1.00 29.66 C +ATOM 293 O VAL A 63 18.053 14.817 0.337 1.00 29.66 O +ATOM 294 CB VAL A 63 15.539 16.803 0.610 1.00 29.66 C +ATOM 295 CG1 VAL A 63 14.389 17.606 0.001 1.00 29.66 C +ATOM 296 CG2 VAL A 63 14.971 15.450 1.064 1.00 29.66 C +ATOM 297 N CYS A 64 18.800 16.863 0.908 1.00 31.19 N +ATOM 298 CA CYS A 64 19.884 16.382 1.752 1.00 31.19 C +ATOM 299 C CYS A 64 19.392 16.239 3.195 1.00 31.19 C +ATOM 300 O CYS A 64 18.864 17.188 3.779 1.00 31.19 O +ATOM 301 CB CYS A 64 21.097 17.307 1.632 1.00 31.19 C +ATOM 302 SG CYS A 64 22.580 16.630 2.432 1.00 31.19 S +ATOM 303 N GLY A 65 19.546 15.048 3.771 1.00 32.71 N +ATOM 304 CA GLY A 65 19.258 14.812 5.180 1.00 32.71 C +ATOM 305 C GLY A 65 20.388 15.287 6.090 1.00 32.71 C +ATOM 306 O GLY A 65 21.559 15.262 5.717 1.00 32.71 O +ATOM 307 N SER A 66 20.055 15.612 7.339 1.00 36.07 N +ATOM 308 CA SER A 66 21.001 15.966 8.411 1.00 36.07 C +ATOM 309 C SER A 66 22.017 14.862 8.752 1.00 36.07 C +ATOM 310 O SER A 66 22.950 15.086 9.516 1.00 36.07 O +ATOM 311 CB SER A 66 20.205 16.354 9.664 1.00 36.07 C +ATOM 312 OG SER A 66 19.422 15.264 10.139 1.00 36.07 O +ATOM 313 N ASN A 67 21.842 13.663 8.189 1.00 38.83 N +ATOM 314 CA ASN A 67 22.766 12.530 8.241 1.00 38.83 C +ATOM 315 C ASN A 67 23.707 12.439 7.017 1.00 38.83 C +ATOM 316 O ASN A 67 24.416 11.446 6.871 1.00 38.83 O +ATOM 317 CB ASN A 67 21.933 11.249 8.421 1.00 38.83 C +ATOM 318 CG ASN A 67 21.043 10.928 7.230 1.00 38.83 C +ATOM 319 OD1 ASN A 67 20.897 11.685 6.281 1.00 38.83 O +ATOM 320 ND2 ASN A 67 20.387 9.796 7.269 1.00 38.83 N +ATOM 321 N GLY A 68 23.688 13.425 6.114 1.00 37.01 N +ATOM 322 CA GLY A 68 24.505 13.458 4.899 1.00 37.01 C +ATOM 323 C GLY A 68 24.048 12.521 3.775 1.00 37.01 C +ATOM 324 O GLY A 68 24.778 12.363 2.797 1.00 37.01 O +ATOM 325 N LYS A 69 22.865 11.896 3.880 1.00 34.67 N +ATOM 326 CA LYS A 69 22.283 11.075 2.805 1.00 34.67 C +ATOM 327 C LYS A 69 21.334 11.900 1.938 1.00 34.67 C +ATOM 328 O LYS A 69 20.519 12.665 2.451 1.00 34.67 O +ATOM 329 CB LYS A 69 21.554 9.849 3.374 1.00 34.67 C +ATOM 330 CG LYS A 69 22.478 8.878 4.122 1.00 34.67 C +ATOM 331 CD LYS A 69 21.677 7.637 4.534 1.00 34.67 C +ATOM 332 CE LYS A 69 22.519 6.669 5.368 1.00 34.67 C +ATOM 333 NZ LYS A 69 21.696 5.515 5.814 1.00 34.67 N +ATOM 334 N THR A 70 21.395 11.697 0.623 1.00 32.30 N +ATOM 335 CA THR A 70 20.398 12.252 -0.304 1.00 32.30 C +ATOM 336 C THR A 70 19.152 11.374 -0.324 1.00 32.30 C +ATOM 337 O THR A 70 19.250 10.172 -0.568 1.00 32.30 O +ATOM 338 CB THR A 70 20.927 12.376 -1.738 1.00 32.30 C +ATOM 339 OG1 THR A 70 22.088 13.169 -1.783 1.00 32.30 O +ATOM 340 CG2 THR A 70 19.896 13.057 -2.640 1.00 32.30 C +ATOM 341 N TYR A 71 17.983 11.984 -0.159 1.00 30.33 N +ATOM 342 CA TYR A 71 16.684 11.330 -0.302 1.00 30.33 C +ATOM 343 C TYR A 71 15.975 11.814 -1.564 1.00 30.33 C +ATOM 344 O TYR A 71 16.109 12.972 -1.954 1.00 30.33 O +ATOM 345 CB TYR A 71 15.839 11.580 0.950 1.00 30.33 C +ATOM 346 CG TYR A 71 16.443 10.959 2.189 1.00 30.33 C +ATOM 347 CD1 TYR A 71 16.097 9.644 2.548 1.00 30.33 C +ATOM 348 CD2 TYR A 71 17.400 11.670 2.940 1.00 30.33 C +ATOM 349 CE1 TYR A 71 16.704 9.039 3.661 1.00 30.33 C +ATOM 350 CE2 TYR A 71 18.013 11.062 4.052 1.00 30.33 C +ATOM 351 CZ TYR A 71 17.662 9.746 4.414 1.00 30.33 C +ATOM 352 OH TYR A 71 18.272 9.156 5.471 1.00 30.33 O +ATOM 353 N LEU A 72 15.178 10.940 -2.187 1.00 30.09 N +ATOM 354 CA LEU A 72 14.424 11.243 -3.412 1.00 30.09 C +ATOM 355 C LEU A 72 13.457 12.430 -3.241 1.00 30.09 C +ATOM 356 O LEU A 72 13.212 13.168 -4.200 1.00 30.09 O +ATOM 357 CB LEU A 72 13.654 9.964 -3.796 1.00 30.09 C +ATOM 358 CG LEU A 72 12.828 10.048 -5.091 1.00 30.09 C +ATOM 359 CD1 LEU A 72 13.715 10.233 -6.322 1.00 30.09 C +ATOM 360 CD2 LEU A 72 12.015 8.769 -5.279 1.00 30.09 C +ATOM 361 N ASN A 73 12.890 12.566 -2.036 1.00 30.00 N +ATOM 362 CA ASN A 73 12.028 13.656 -1.583 1.00 30.00 C +ATOM 363 C ASN A 73 11.841 13.633 -0.050 1.00 30.00 C +ATOM 364 O ASN A 73 12.241 12.681 0.621 1.00 30.00 O +ATOM 365 CB ASN A 73 10.678 13.575 -2.326 1.00 30.00 C +ATOM 366 CG ASN A 73 9.973 12.245 -2.142 1.00 30.00 C +ATOM 367 OD1 ASN A 73 9.608 11.875 -1.042 1.00 30.00 O +ATOM 368 ND2 ASN A 73 9.724 11.524 -3.209 1.00 30.00 N +ATOM 369 N HIS A 74 11.177 14.663 0.494 1.00 29.88 N +ATOM 370 CA HIS A 74 10.857 14.784 1.926 1.00 29.88 C +ATOM 371 C HIS A 74 10.117 13.572 2.521 1.00 29.88 C +ATOM 372 O HIS A 74 10.296 13.275 3.702 1.00 29.88 O +ATOM 373 CB HIS A 74 9.967 16.014 2.166 1.00 29.88 C +ATOM 374 CG HIS A 74 10.550 17.336 1.744 1.00 29.88 C +ATOM 375 ND1 HIS A 74 11.376 18.151 2.486 1.00 29.88 N +ATOM 376 CD2 HIS A 74 10.278 17.998 0.581 1.00 29.88 C +ATOM 377 CE1 HIS A 74 11.615 19.269 1.770 1.00 29.88 C +ATOM 378 NE2 HIS A 74 10.947 19.214 0.607 1.00 29.88 N +ATOM 379 N CYS A 75 9.273 12.879 1.743 1.00 29.69 N +ATOM 380 CA CYS A 75 8.543 11.713 2.246 1.00 29.69 C +ATOM 381 C CYS A 75 9.516 10.572 2.551 1.00 29.69 C +ATOM 382 O CYS A 75 9.475 10.040 3.655 1.00 29.69 O +ATOM 383 CB CYS A 75 7.436 11.307 1.261 1.00 29.69 C +ATOM 384 SG CYS A 75 6.372 9.927 1.770 1.00 29.69 S +ATOM 385 N GLU A 76 10.438 10.260 1.631 1.00 29.08 N +ATOM 386 CA GLU A 76 11.446 9.213 1.852 1.00 29.08 C +ATOM 387 C GLU A 76 12.381 9.547 3.024 1.00 29.08 C +ATOM 388 O GLU A 76 12.689 8.664 3.819 1.00 29.08 O +ATOM 389 CB GLU A 76 12.266 8.961 0.575 1.00 29.08 C +ATOM 390 CG GLU A 76 11.460 8.366 -0.592 1.00 29.08 C +ATOM 391 CD GLU A 76 10.921 6.936 -0.378 1.00 29.08 C +ATOM 392 OE1 GLU A 76 10.264 6.454 -1.329 1.00 29.08 O +ATOM 393 OE2 GLU A 76 11.113 6.341 0.710 1.00 29.08 O +ATOM 394 N LEU A 77 12.765 10.820 3.198 1.00 29.01 N +ATOM 395 CA LEU A 77 13.571 11.260 4.347 1.00 29.01 C +ATOM 396 C LEU A 77 12.863 10.984 5.678 1.00 29.01 C +ATOM 397 O LEU A 77 13.439 10.386 6.585 1.00 29.01 O +ATOM 398 CB LEU A 77 13.934 12.745 4.161 1.00 29.01 C +ATOM 399 CG LEU A 77 14.865 13.328 5.248 1.00 29.01 C +ATOM 400 CD1 LEU A 77 15.594 14.539 4.669 1.00 29.01 C +ATOM 401 CD2 LEU A 77 14.119 13.827 6.494 1.00 29.01 C +ATOM 402 N HIS A 78 11.599 11.397 5.810 1.00 29.29 N +ATOM 403 CA HIS A 78 10.855 11.186 7.053 1.00 29.29 C +ATOM 404 C HIS A 78 10.401 9.731 7.243 1.00 29.29 C +ATOM 405 O HIS A 78 10.247 9.284 8.380 1.00 29.29 O +ATOM 406 CB HIS A 78 9.681 12.161 7.127 1.00 29.29 C +ATOM 407 CG HIS A 78 10.103 13.598 7.284 1.00 29.29 C +ATOM 408 ND1 HIS A 78 10.880 14.111 8.300 1.00 29.29 N +ATOM 409 CD2 HIS A 78 9.750 14.648 6.480 1.00 29.29 C +ATOM 410 CE1 HIS A 78 10.986 15.436 8.111 1.00 29.29 C +ATOM 411 NE2 HIS A 78 10.306 15.811 7.018 1.00 29.29 N +ATOM 412 N ARG A 79 10.218 8.974 6.157 1.00 29.85 N +ATOM 413 CA ARG A 79 9.968 7.530 6.187 1.00 29.85 C +ATOM 414 C ARG A 79 11.190 6.781 6.715 1.00 29.85 C +ATOM 415 O ARG A 79 11.035 5.984 7.634 1.00 29.85 O +ATOM 416 CB ARG A 79 9.548 7.079 4.783 1.00 29.85 C +ATOM 417 CG ARG A 79 9.267 5.579 4.681 1.00 29.85 C +ATOM 418 CD ARG A 79 8.772 5.256 3.269 1.00 29.85 C +ATOM 419 NE ARG A 79 8.588 3.812 3.095 1.00 29.85 N +ATOM 420 CZ ARG A 79 9.493 2.940 2.703 1.00 29.85 C +ATOM 421 NH1 ARG A 79 10.697 3.299 2.348 1.00 29.85 N +ATOM 422 NH2 ARG A 79 9.172 1.681 2.644 1.00 29.85 N +ATOM 423 N ASP A 80 12.393 7.093 6.234 1.00 30.45 N +ATOM 424 CA ASP A 80 13.647 6.536 6.763 1.00 30.45 C +ATOM 425 C ASP A 80 13.857 6.926 8.236 1.00 30.45 C +ATOM 426 O ASP A 80 14.147 6.070 9.068 1.00 30.45 O +ATOM 427 CB ASP A 80 14.818 6.999 5.885 1.00 30.45 C +ATOM 428 CG ASP A 80 16.141 6.297 6.215 1.00 30.45 C +ATOM 429 OD1 ASP A 80 16.128 5.061 6.374 1.00 30.45 O +ATOM 430 OD2 ASP A 80 17.185 6.990 6.245 1.00 30.45 O +ATOM 431 N ALA A 81 13.589 8.186 8.604 1.00 30.49 N +ATOM 432 CA ALA A 81 13.603 8.639 10.000 1.00 30.49 C +ATOM 433 C ALA A 81 12.647 7.829 10.898 1.00 30.49 C +ATOM 434 O ALA A 81 12.968 7.514 12.043 1.00 30.49 O +ATOM 435 CB ALA A 81 13.220 10.126 10.029 1.00 30.49 C +ATOM 436 N CYS A 82 11.470 7.473 10.373 1.00 29.93 N +ATOM 437 CA CYS A 82 10.490 6.651 11.073 1.00 29.93 C +ATOM 438 C CYS A 82 10.958 5.194 11.216 1.00 29.93 C +ATOM 439 O CYS A 82 10.958 4.670 12.328 1.00 29.93 O +ATOM 440 CB CYS A 82 9.142 6.771 10.349 1.00 29.93 C +ATOM 441 SG CYS A 82 7.747 6.023 11.227 1.00 29.93 S +ATOM 442 N LEU A 83 11.396 4.567 10.119 1.00 30.29 N +ATOM 443 CA LEU A 83 11.817 3.160 10.078 1.00 30.29 C +ATOM 444 C LEU A 83 13.090 2.891 10.897 1.00 30.29 C +ATOM 445 O LEU A 83 13.215 1.837 11.510 1.00 30.29 O +ATOM 446 CB LEU A 83 12.038 2.750 8.608 1.00 30.29 C +ATOM 447 CG LEU A 83 10.759 2.666 7.753 1.00 30.29 C +ATOM 448 CD1 LEU A 83 11.135 2.455 6.285 1.00 30.29 C +ATOM 449 CD2 LEU A 83 9.845 1.517 8.180 1.00 30.29 C +ATOM 450 N THR A 84 14.023 3.844 10.928 1.00 32.88 N +ATOM 451 CA THR A 84 15.299 3.729 11.661 1.00 32.88 C +ATOM 452 C THR A 84 15.241 4.258 13.097 1.00 32.88 C +ATOM 453 O THR A 84 16.249 4.213 13.800 1.00 32.88 O +ATOM 454 CB THR A 84 16.432 4.459 10.924 1.00 32.88 C +ATOM 455 OG1 THR A 84 16.127 5.831 10.822 1.00 32.88 O +ATOM 456 CG2 THR A 84 16.690 3.914 9.520 1.00 32.88 C +ATOM 457 N GLY A 85 14.115 4.843 13.526 1.00 33.49 N +ATOM 458 CA GLY A 85 14.005 5.551 14.810 1.00 33.49 C +ATOM 459 C GLY A 85 14.919 6.784 14.943 1.00 33.49 C +ATOM 460 O GLY A 85 15.079 7.314 16.044 1.00 33.49 O +ATOM 461 N SER A 86 15.532 7.244 13.849 1.00 39.08 N +ATOM 462 CA SER A 86 16.556 8.294 13.857 1.00 39.08 C +ATOM 463 C SER A 86 15.974 9.698 13.690 1.00 39.08 C +ATOM 464 O SER A 86 15.013 9.926 12.956 1.00 39.08 O +ATOM 465 CB SER A 86 17.605 8.024 12.776 1.00 39.08 C +ATOM 466 OG SER A 86 18.312 6.842 13.086 1.00 39.08 O +ATOM 467 N LYS A 87 16.605 10.695 14.323 1.00 35.06 N +ATOM 468 CA LYS A 87 16.247 12.116 14.169 1.00 35.06 C +ATOM 469 C LYS A 87 16.859 12.698 12.887 1.00 35.06 C +ATOM 470 O LYS A 87 17.856 13.412 12.945 1.00 35.06 O +ATOM 471 CB LYS A 87 16.643 12.916 15.422 1.00 35.06 C +ATOM 472 CG LYS A 87 15.813 12.529 16.655 1.00 35.06 C +ATOM 473 CD LYS A 87 16.219 13.382 17.864 1.00 35.06 C +ATOM 474 CE LYS A 87 15.402 12.974 19.096 1.00 35.06 C +ATOM 475 NZ LYS A 87 15.821 13.725 20.307 1.00 35.06 N +ATOM 476 N ILE A 88 16.261 12.384 11.740 1.00 33.60 N +ATOM 477 CA ILE A 88 16.664 12.924 10.434 1.00 33.60 C +ATOM 478 C ILE A 88 15.792 14.142 10.101 1.00 33.60 C +ATOM 479 O ILE A 88 14.561 14.064 10.123 1.00 33.60 O +ATOM 480 CB ILE A 88 16.638 11.862 9.308 1.00 33.60 C +ATOM 481 CG1 ILE A 88 17.138 10.478 9.797 1.00 33.60 C +ATOM 482 CG2 ILE A 88 17.469 12.368 8.112 1.00 33.60 C +ATOM 483 CD1 ILE A 88 17.068 9.371 8.743 1.00 33.60 C +ATOM 484 N GLN A 89 16.434 15.267 9.803 1.00 33.60 N +ATOM 485 CA GLN A 89 15.808 16.507 9.336 1.00 33.60 C +ATOM 486 C GLN A 89 16.334 16.852 7.938 1.00 33.60 C +ATOM 487 O GLN A 89 17.340 16.293 7.505 1.00 33.60 O +ATOM 488 CB GLN A 89 16.066 17.638 10.350 1.00 33.60 C +ATOM 489 CG GLN A 89 15.365 17.372 11.694 1.00 33.60 C +ATOM 490 CD GLN A 89 15.540 18.497 12.713 1.00 33.60 C +ATOM 491 OE1 GLN A 89 16.344 19.399 12.587 1.00 33.60 O +ATOM 492 NE2 GLN A 89 14.790 18.488 13.792 1.00 33.60 N +ATOM 493 N VAL A 90 15.659 17.752 7.221 1.00 31.75 N +ATOM 494 CA VAL A 90 16.197 18.330 5.978 1.00 31.75 C +ATOM 495 C VAL A 90 17.284 19.335 6.358 1.00 31.75 C +ATOM 496 O VAL A 90 17.021 20.221 7.166 1.00 31.75 O +ATOM 497 CB VAL A 90 15.092 19.030 5.160 1.00 31.75 C +ATOM 498 CG1 VAL A 90 15.636 19.631 3.858 1.00 31.75 C +ATOM 499 CG2 VAL A 90 13.950 18.078 4.779 1.00 31.75 C +ATOM 500 N ASP A 91 18.478 19.192 5.786 1.00 35.39 N +ATOM 501 CA ASP A 91 19.572 20.168 5.904 1.00 35.39 C +ATOM 502 C ASP A 91 19.423 21.248 4.816 1.00 35.39 C +ATOM 503 O ASP A 91 19.393 22.444 5.099 1.00 35.39 O +ATOM 504 CB ASP A 91 20.911 19.403 5.822 1.00 35.39 C +ATOM 505 CG ASP A 91 22.153 20.234 6.175 1.00 35.39 C +ATOM 506 OD1 ASP A 91 22.048 21.122 7.046 1.00 35.39 O +ATOM 507 OD2 ASP A 91 23.233 19.931 5.616 1.00 35.39 O +ATOM 508 N TYR A 92 19.231 20.821 3.561 1.00 32.82 N +ATOM 509 CA TYR A 92 18.906 21.690 2.424 1.00 32.82 C +ATOM 510 C TYR A 92 18.239 20.917 1.278 1.00 32.82 C +ATOM 511 O TYR A 92 18.379 19.697 1.155 1.00 32.82 O +ATOM 512 CB TYR A 92 20.167 22.417 1.918 1.00 32.82 C +ATOM 513 CG TYR A 92 21.357 21.527 1.607 1.00 32.82 C +ATOM 514 CD1 TYR A 92 22.245 21.190 2.644 1.00 32.82 C +ATOM 515 CD2 TYR A 92 21.587 21.044 0.304 1.00 32.82 C +ATOM 516 CE1 TYR A 92 23.348 20.361 2.393 1.00 32.82 C +ATOM 517 CE2 TYR A 92 22.710 20.234 0.038 1.00 32.82 C +ATOM 518 CZ TYR A 92 23.587 19.889 1.090 1.00 32.82 C +ATOM 519 OH TYR A 92 24.681 19.118 0.867 1.00 32.82 O +ATOM 520 N ASP A 93 17.535 21.648 0.412 1.00 31.28 N +ATOM 521 CA ASP A 93 16.960 21.113 -0.821 1.00 31.28 C +ATOM 522 C ASP A 93 18.042 20.871 -1.887 1.00 31.28 C +ATOM 523 O ASP A 93 18.984 21.652 -2.040 1.00 31.28 O +ATOM 524 CB ASP A 93 15.875 22.057 -1.364 1.00 31.28 C +ATOM 525 CG ASP A 93 14.681 22.229 -0.419 1.00 31.28 C +ATOM 526 OD1 ASP A 93 14.275 21.227 0.212 1.00 31.28 O +ATOM 527 OD2 ASP A 93 14.163 23.366 -0.367 1.00 31.28 O +ATOM 528 N GLY A 94 17.888 19.799 -2.662 1.00 31.90 N +ATOM 529 CA GLY A 94 18.866 19.319 -3.636 1.00 31.90 C +ATOM 530 C GLY A 94 19.494 17.981 -3.241 1.00 31.90 C +ATOM 531 O GLY A 94 19.171 17.384 -2.217 1.00 31.90 O +ATOM 532 N HIS A 95 20.406 17.485 -4.075 1.00 31.99 N +ATOM 533 CA HIS A 95 21.263 16.371 -3.668 1.00 31.99 C +ATOM 534 C HIS A 95 22.215 16.862 -2.576 1.00 31.99 C +ATOM 535 O HIS A 95 22.663 18.011 -2.641 1.00 31.99 O +ATOM 536 CB HIS A 95 22.040 15.803 -4.866 1.00 31.99 C +ATOM 537 CG HIS A 95 21.154 15.231 -5.947 1.00 31.99 C +ATOM 538 ND1 HIS A 95 20.505 15.945 -6.928 1.00 31.99 N +ATOM 539 CD2 HIS A 95 20.859 13.909 -6.156 1.00 31.99 C +ATOM 540 CE1 HIS A 95 19.827 15.078 -7.695 1.00 31.99 C +ATOM 541 NE2 HIS A 95 19.988 13.825 -7.249 1.00 31.99 N +ATOM 542 N CYS A 96 22.574 16.005 -1.618 1.00 35.86 N +ATOM 543 CA CYS A 96 23.725 16.299 -0.781 1.00 35.86 C +ATOM 544 C CYS A 96 24.915 16.573 -1.698 1.00 35.86 C +ATOM 545 O CYS A 96 25.245 15.748 -2.554 1.00 35.86 O +ATOM 546 CB CYS A 96 24.059 15.156 0.182 1.00 35.86 C +ATOM 547 SG CYS A 96 22.856 14.842 1.493 1.00 35.86 S +ATOM 548 N LYS A 97 25.547 17.733 -1.512 1.00 40.03 N +ATOM 549 CA LYS A 97 26.876 17.993 -2.052 1.00 40.03 C +ATOM 550 C LYS A 97 27.724 16.808 -1.637 1.00 40.03 C +ATOM 551 O LYS A 97 27.757 16.481 -0.446 1.00 40.03 O +ATOM 552 CB LYS A 97 27.457 19.289 -1.468 1.00 40.03 C +ATOM 553 CG LYS A 97 26.636 20.519 -1.873 1.00 40.03 C +ATOM 554 CD LYS A 97 27.161 21.781 -1.184 1.00 40.03 C +ATOM 555 CE LYS A 97 26.246 22.958 -1.538 1.00 40.03 C +ATOM 556 NZ LYS A 97 26.659 24.199 -0.841 1.00 40.03 N +END diff --git a/tests/files/pdb/6SXW_af.pdb b/tests/files/pdb/6SXW_af.pdb new file mode 100644 index 00000000..b9b2b5f1 --- /dev/null +++ b/tests/files/pdb/6SXW_af.pdb @@ -0,0 +1,582 @@ +REMARK TITLE 6SXW AlphaFold-start MR +REMARK Log-Likelihood Gain: 115.305 +REMARK RFZ=3.9 TFZ=6.0 PAK=4 LLG=62 TFZ==8.9 LLG=115 TFZ==11.7 PAK=3 LLG=115 TFZ==11.7 +REMARK ENSEMBLE e_Q14966 EULER 152.07 35.69 68.73 FRAC -0.063 0.329 -0.463 +CRYST1 57.435 57.435 46.619 90.00 90.00 120.00 P 32 2 1 6 +SCALE1 0.017411 0.010052 -0.000000 0.00000 +SCALE2 0.000000 0.020104 -0.000000 0.00000 +SCALE3 0.000000 0.000000 0.021450 0.00000 +ATOM 1 N SER A 676 34.529 -7.571 -16.411 1.00 94.52 N +ATOM 2 CA SER A 676 34.729 -7.576 -14.942 1.00 94.52 C +ATOM 3 C SER A 676 33.552 -6.849 -14.270 1.00 94.52 C +ATOM 4 O SER A 676 33.584 -6.606 -13.062 1.00 94.52 O +ATOM 5 CB SER A 676 36.054 -6.852 -14.569 1.00 94.52 C +ATOM 6 OG SER A 676 36.419 -5.954 -15.634 1.00 94.52 O +ATOM 7 N VAL A 677 32.385 -6.533 -15.015 1.00 82.63 N +ATOM 8 CA VAL A 677 31.339 -5.697 -14.389 1.00 82.63 C +ATOM 9 C VAL A 677 30.093 -6.562 -14.177 1.00 82.63 C +ATOM 10 O VAL A 677 29.701 -7.342 -15.047 1.00 82.63 O +ATOM 11 CB VAL A 677 30.979 -4.503 -15.333 1.00 82.63 C +ATOM 12 CG1 VAL A 677 29.844 -3.646 -14.685 1.00 82.63 C +ATOM 13 CG2 VAL A 677 32.244 -3.615 -15.558 1.00 82.63 C +ATOM 14 N LEU A 678 29.467 -6.599 -12.962 1.00 79.72 N +ATOM 15 CA LEU A 678 28.226 -7.303 -12.610 1.00 79.72 C +ATOM 16 C LEU A 678 27.077 -6.302 -12.454 1.00 79.72 C +ATOM 17 O LEU A 678 27.247 -5.217 -11.894 1.00 79.72 O +ATOM 18 CB LEU A 678 28.426 -8.024 -11.252 1.00 79.72 C +ATOM 19 CG LEU A 678 29.039 -9.447 -11.407 1.00 79.72 C +ATOM 20 CD1 LEU A 678 30.437 -9.475 -10.730 1.00 79.72 C +ATOM 21 CD2 LEU A 678 28.110 -10.489 -10.725 1.00 79.72 C +ATOM 22 N LEU A 679 25.853 -6.668 -13.068 1.00 78.50 N +ATOM 23 CA LEU A 679 24.638 -5.850 -12.902 1.00 78.50 C +ATOM 24 C LEU A 679 23.778 -6.463 -11.790 1.00 78.50 C +ATOM 25 O LEU A 679 23.391 -7.632 -11.851 1.00 78.50 O +ATOM 26 CB LEU A 679 23.840 -5.849 -14.236 1.00 78.50 C +ATOM 27 CG LEU A 679 22.472 -5.128 -14.141 1.00 78.50 C +ATOM 28 CD1 LEU A 679 22.680 -3.595 -13.935 1.00 78.50 C +ATOM 29 CD2 LEU A 679 21.669 -5.378 -15.465 1.00 78.50 C +ATOM 30 N ILE A 680 23.620 -5.671 -10.651 1.00 78.39 N +ATOM 31 CA ILE A 680 22.746 -6.118 -9.555 1.00 78.39 C +ATOM 32 C ILE A 680 21.386 -5.419 -9.677 1.00 78.39 C +ATOM 33 O ILE A 680 21.294 -4.189 -9.664 1.00 78.39 O +ATOM 34 CB ILE A 680 23.398 -5.733 -8.188 1.00 78.39 C +ATOM 35 CG1 ILE A 680 24.809 -6.392 -8.113 1.00 78.39 C +ATOM 36 CG2 ILE A 680 22.498 -6.240 -7.022 1.00 78.39 C +ATOM 37 CD1 ILE A 680 25.628 -5.844 -6.904 1.00 78.39 C +ATOM 38 N THR A 681 20.226 -6.268 -9.909 1.00 81.42 N +ATOM 39 CA THR A 681 18.877 -5.704 -10.117 1.00 81.42 C +ATOM 40 C THR A 681 17.986 -6.097 -8.936 1.00 81.42 C +ATOM 41 O THR A 681 18.337 -6.964 -8.134 1.00 81.42 O +ATOM 42 CB THR A 681 18.238 -6.256 -11.425 1.00 81.42 C +ATOM 43 OG1 THR A 681 18.258 -7.707 -11.378 1.00 81.42 O +ATOM 44 CG2 THR A 681 19.042 -5.769 -12.661 1.00 81.42 C +ATOM 45 N GLU A 682 16.931 -5.350 -8.516 1.00 85.81 N +ATOM 46 CA GLU A 682 15.919 -5.580 -7.458 1.00 85.81 C +ATOM 47 C GLU A 682 16.470 -5.205 -6.081 1.00 85.81 C +ATOM 48 O GLU A 682 16.314 -5.956 -5.117 1.00 85.81 O +ATOM 49 CB GLU A 682 15.479 -7.069 -7.434 1.00 85.81 C +ATOM 50 CG GLU A 682 14.725 -7.448 -8.744 1.00 85.81 C +ATOM 51 CD GLU A 682 14.272 -8.914 -8.690 1.00 85.81 C +ATOM 52 OE1 GLU A 682 13.740 -9.399 -9.715 1.00 85.81 O +ATOM 53 OE2 GLU A 682 14.437 -9.567 -7.636 1.00 85.81 O +ATOM 54 N LEU A 683 17.154 -3.932 -6.063 1.00 85.07 N +ATOM 55 CA LEU A 683 17.701 -3.438 -4.781 1.00 85.07 C +ATOM 56 C LEU A 683 16.568 -2.964 -3.865 1.00 85.07 C +ATOM 57 O LEU A 683 15.506 -2.547 -4.334 1.00 85.07 O +ATOM 58 CB LEU A 683 18.634 -2.215 -5.056 1.00 85.07 C +ATOM 59 CG LEU A 683 19.923 -2.595 -5.847 1.00 85.07 C +ATOM 60 CD1 LEU A 683 20.664 -1.281 -6.229 1.00 85.07 C +ATOM 61 CD2 LEU A 683 20.846 -3.492 -4.961 1.00 85.07 C +ATOM 62 N PRO A 684 16.545 -3.278 -2.559 1.00 92.40 N +ATOM 63 CA PRO A 684 15.508 -2.852 -1.604 1.00 92.40 C +ATOM 64 C PRO A 684 15.227 -1.348 -1.695 1.00 92.40 C +ATOM 65 O PRO A 684 16.125 -0.552 -1.986 1.00 92.40 O +ATOM 66 CB PRO A 684 16.132 -3.183 -0.216 1.00 92.40 C +ATOM 67 CG PRO A 684 17.460 -3.864 -0.555 1.00 92.40 C +ATOM 68 CD PRO A 684 17.690 -3.758 -2.068 1.00 92.40 C +ATOM 69 N GLY A 687 16.645 1.925 0.502 1.00123.00 N +ATOM 70 CA GLY A 687 17.986 2.553 0.386 1.00123.00 C +ATOM 71 C GLY A 687 19.095 1.497 0.508 1.00123.00 C +ATOM 72 O GLY A 687 19.091 0.670 1.422 1.00123.00 O +ATOM 73 N CYS A 688 19.809 1.012 -0.585 1.00 99.77 N +ATOM 74 CA CYS A 688 21.037 0.178 -0.620 1.00 99.77 C +ATOM 75 C CYS A 688 22.277 1.074 -0.743 1.00 99.77 C +ATOM 76 O CYS A 688 22.352 1.947 -1.610 1.00 99.77 O +ATOM 77 CB CYS A 688 20.949 -0.762 -1.853 1.00 99.77 C +ATOM 78 SG CYS A 688 22.336 -1.913 -1.925 1.00 99.77 S +ATOM 79 N THR A 689 23.150 1.166 0.410 1.00 84.01 N +ATOM 80 CA THR A 689 24.392 1.981 0.381 1.00 84.01 C +ATOM 81 C THR A 689 25.506 1.215 -0.349 1.00 84.01 C +ATOM 82 O THR A 689 25.415 0.004 -0.559 1.00 84.01 O +ATOM 83 CB THR A 689 24.872 2.285 1.830 1.00 84.01 C +ATOM 84 OG1 THR A 689 25.042 1.040 2.566 1.00 84.01 O +ATOM 85 CG2 THR A 689 23.794 3.126 2.585 1.00 84.01 C +ATOM 86 N GLU A 690 26.338 1.969 -1.166 1.00 81.06 N +ATOM 87 CA GLU A 690 27.506 1.371 -1.862 1.00 81.06 C +ATOM 88 C GLU A 690 28.294 0.448 -0.938 1.00 81.06 C +ATOM 89 O GLU A 690 28.810 -0.577 -1.390 1.00 81.06 O +ATOM 90 CB GLU A 690 28.418 2.550 -2.336 1.00 81.06 C +ATOM 91 CG GLU A 690 29.631 2.016 -3.125 1.00 81.06 C +ATOM 92 CD GLU A 690 30.503 3.171 -3.644 1.00 81.06 C +ATOM 93 OE1 GLU A 690 31.574 2.894 -4.230 1.00 81.06 O +ATOM 94 OE2 GLU A 690 30.114 4.345 -3.455 1.00 81.06 O +ATOM 95 N GLU A 691 28.306 0.612 0.473 1.00 82.93 N +ATOM 96 CA GLU A 691 29.049 -0.228 1.444 1.00 82.93 C +ATOM 97 C GLU A 691 28.365 -1.591 1.610 1.00 82.93 C +ATOM 98 O GLU A 691 29.049 -2.601 1.790 1.00 82.93 O +ATOM 99 CB GLU A 691 29.068 0.488 2.826 1.00 82.93 C +ATOM 100 CG GLU A 691 30.098 1.658 2.846 1.00 82.93 C +ATOM 101 CD GLU A 691 29.421 3.003 2.555 1.00 82.93 C +ATOM 102 OE1 GLU A 691 30.118 3.922 2.069 1.00 82.93 O +ATOM 103 OE2 GLU A 691 28.207 3.140 2.825 1.00 82.93 O +ATOM 104 N ASP A 692 26.998 -1.569 1.539 1.00 80.57 N +ATOM 105 CA ASP A 692 26.258 -2.847 1.652 1.00 80.57 C +ATOM 106 C ASP A 692 26.531 -3.753 0.445 1.00 80.57 C +ATOM 107 O ASP A 692 26.680 -4.965 0.612 1.00 80.57 O +ATOM 108 CB ASP A 692 24.730 -2.549 1.714 1.00 80.57 C +ATOM 109 CG ASP A 692 24.369 -1.839 3.020 1.00 80.57 C +ATOM 110 OD1 ASP A 692 25.101 -1.969 4.024 1.00 80.57 O +ATOM 111 OD2 ASP A 692 23.321 -1.158 3.036 1.00 80.57 O +ATOM 112 N VAL A 693 26.559 -3.116 -0.793 1.00 80.51 N +ATOM 113 CA VAL A 693 26.854 -3.903 -2.018 1.00 80.51 C +ATOM 114 C VAL A 693 28.316 -4.361 -2.003 1.00 80.51 C +ATOM 115 O VAL A 693 28.602 -5.509 -2.353 1.00 80.51 O +ATOM 116 CB VAL A 693 26.616 -3.003 -3.274 1.00 80.51 C +ATOM 117 CG1 VAL A 693 26.969 -3.803 -4.563 1.00 80.51 C +ATOM 118 CG2 VAL A 693 25.113 -2.568 -3.317 1.00 80.51 C +ATOM 119 N ARG A 694 29.273 -3.532 -1.450 1.00 79.10 N +ATOM 120 CA ARG A 694 30.703 -3.912 -1.391 1.00 79.10 C +ATOM 121 C ARG A 694 30.924 -5.023 -0.363 1.00 79.10 C +ATOM 122 O ARG A 694 31.700 -5.947 -0.618 1.00 79.10 O +ATOM 123 CB ARG A 694 31.533 -2.668 -0.951 1.00 79.10 C +ATOM 124 CG ARG A 694 32.118 -1.935 -2.194 1.00 79.10 C +ATOM 125 CD ARG A 694 33.042 -0.782 -1.698 1.00 79.10 C +ATOM 126 NE ARG A 694 32.685 0.455 -2.457 1.00 79.10 N +ATOM 127 CZ ARG A 694 33.351 1.607 -2.267 1.00 79.10 C +ATOM 128 NH1 ARG A 694 34.384 1.713 -1.371 1.00 79.10 N +ATOM 129 NH2 ARG A 694 33.008 2.712 -2.962 1.00 79.10 N +ATOM 130 N LYS A 695 30.163 -5.013 0.796 1.00 82.01 N +ATOM 131 CA LYS A 695 30.323 -6.041 1.850 1.00 82.01 C +ATOM 132 C LYS A 695 29.786 -7.389 1.366 1.00 82.01 C +ATOM 133 O LYS A 695 30.343 -8.431 1.720 1.00 82.01 O +ATOM 134 CB LYS A 695 29.508 -5.607 3.100 1.00 82.01 C +ATOM 135 CG LYS A 695 30.303 -4.554 3.930 1.00 82.01 C +ATOM 136 CD LYS A 695 29.464 -4.231 5.211 1.00 82.01 C +ATOM 137 CE LYS A 695 30.151 -3.033 5.947 1.00 82.01 C +ATOM 138 NZ LYS A 695 29.356 -2.714 7.243 1.00 82.01 N +ATOM 139 N LEU A 696 28.763 -7.365 0.491 1.00 82.63 N +ATOM 140 CA LEU A 696 28.166 -8.624 -0.003 1.00 82.63 C +ATOM 141 C LEU A 696 29.084 -9.308 -1.018 1.00 82.63 C +ATOM 142 O LEU A 696 29.158 -10.538 -1.043 1.00 82.63 O +ATOM 143 CB LEU A 696 26.823 -8.297 -0.720 1.00 82.63 C +ATOM 144 CG LEU A 696 25.645 -8.042 0.271 1.00 82.63 C +ATOM 145 CD1 LEU A 696 24.443 -7.477 -0.521 1.00 82.63 C +ATOM 146 CD2 LEU A 696 25.246 -9.383 0.975 1.00 82.63 C +ATOM 147 N PHE A 697 29.892 -8.468 -1.873 1.00 78.92 N +ATOM 148 CA PHE A 697 30.645 -9.108 -2.976 1.00 78.92 C +ATOM 149 C PHE A 697 32.139 -9.139 -2.638 1.00 78.92 C +ATOM 150 O PHE A 697 32.939 -9.661 -3.416 1.00 78.92 O +ATOM 151 CB PHE A 697 30.453 -8.279 -4.274 1.00 78.92 C +ATOM 152 CG PHE A 697 29.075 -8.501 -4.881 1.00 78.92 C +ATOM 153 CD1 PHE A 697 28.819 -9.644 -5.649 1.00 78.92 C +ATOM 154 CD2 PHE A 697 28.061 -7.562 -4.663 1.00 78.92 C +ATOM 155 CE1 PHE A 697 27.549 -9.838 -6.225 1.00 78.92 C +ATOM 156 CE2 PHE A 697 26.787 -7.756 -5.230 1.00 78.92 C +ATOM 157 CZ PHE A 697 26.528 -8.878 -6.019 1.00 78.92 C +ATOM 158 N GLN A 698 32.538 -8.659 -1.379 1.00 80.78 N +ATOM 159 CA GLN A 698 33.949 -8.639 -0.925 1.00 80.78 C +ATOM 160 C GLN A 698 34.498 -10.065 -0.773 1.00 80.78 C +ATOM 161 O GLN A 698 35.660 -10.311 -1.100 1.00 80.78 O +ATOM 162 CB GLN A 698 33.994 -7.919 0.458 1.00 80.78 C +ATOM 163 CG GLN A 698 35.468 -7.702 0.913 1.00 80.78 C +ATOM 164 CD GLN A 698 35.504 -6.797 2.146 1.00 80.78 C +ATOM 165 OE1 GLN A 698 34.471 -6.417 2.705 1.00 80.78 O +ATOM 166 NE2 GLN A 698 36.749 -6.453 2.598 1.00 80.78 N +ATOM 167 N PRO A 699 33.647 -11.135 -0.432 1.00 81.13 N +ATOM 168 CA PRO A 699 34.169 -12.508 -0.295 1.00 81.13 C +ATOM 169 C PRO A 699 34.519 -13.121 -1.658 1.00 81.13 C +ATOM 170 O PRO A 699 35.332 -14.047 -1.715 1.00 81.13 O +ATOM 171 CB PRO A 699 32.972 -13.294 0.328 1.00 81.13 C +ATOM 172 CG PRO A 699 32.125 -12.193 0.988 1.00 81.13 C +ATOM 173 CD PRO A 699 32.339 -10.890 0.218 1.00 81.13 C +ATOM 174 N PHE A 700 33.982 -12.551 -2.752 1.00 84.18 N +ATOM 175 CA PHE A 700 34.242 -13.189 -4.066 1.00 84.18 C +ATOM 176 C PHE A 700 35.420 -12.508 -4.766 1.00 84.18 C +ATOM 177 O PHE A 700 35.965 -13.056 -5.726 1.00 84.18 O +ATOM 178 CB PHE A 700 32.979 -13.022 -4.956 1.00 84.18 C +ATOM 179 CG PHE A 700 31.813 -13.849 -4.433 1.00 84.18 C +ATOM 180 CD1 PHE A 700 31.740 -15.224 -4.706 1.00 84.18 C +ATOM 181 CD2 PHE A 700 30.804 -13.224 -3.678 1.00 84.18 C +ATOM 182 CE1 PHE A 700 30.649 -15.978 -4.230 1.00 84.18 C +ATOM 183 CE2 PHE A 700 29.707 -13.975 -3.207 1.00 84.18 C +ATOM 184 CZ PHE A 700 29.636 -15.349 -3.476 1.00 84.18 C +ATOM 185 N GLY A 701 36.025 -11.363 -4.140 1.00 83.02 N +ATOM 186 CA GLY A 701 37.188 -10.678 -4.756 1.00 83.02 C +ATOM 187 C GLY A 701 37.115 -9.167 -4.506 1.00 83.02 C +ATOM 188 O GLY A 701 36.188 -8.663 -3.867 1.00 83.02 O +ATOM 189 N LYS A 702 38.208 -8.431 -4.841 1.00 82.09 N +ATOM 190 CA LYS A 702 38.268 -6.967 -4.624 1.00 82.09 C +ATOM 191 C LYS A 702 37.358 -6.224 -5.611 1.00 82.09 C +ATOM 192 O LYS A 702 37.376 -6.481 -6.816 1.00 82.09 O +ATOM 193 CB LYS A 702 39.727 -6.481 -4.846 1.00 82.09 C +ATOM 194 CG LYS A 702 40.643 -6.931 -3.668 1.00 82.09 C +ATOM 195 CD LYS A 702 42.071 -6.353 -3.923 1.00 82.09 C +ATOM 196 CE LYS A 702 43.022 -6.915 -2.828 1.00 82.09 C +ATOM 197 NZ LYS A 702 44.452 -6.357 -3.055 1.00 82.09 N +ATOM 198 N VAL A 703 36.384 -5.375 -4.926 1.00 79.10 N +ATOM 199 CA VAL A 703 35.438 -4.585 -5.750 1.00 79.10 C +ATOM 200 C VAL A 703 36.123 -3.282 -6.190 1.00 79.10 C +ATOM 201 O VAL A 703 36.643 -2.522 -5.369 1.00 79.10 O +ATOM 202 CB VAL A 703 34.182 -4.230 -4.901 1.00 79.10 C +ATOM 203 CG1 VAL A 703 33.191 -3.404 -5.762 1.00 79.10 C +ATOM 204 CG2 VAL A 703 33.486 -5.572 -4.403 1.00 79.10 C +ATOM 205 N ASN A 704 36.416 -3.002 -7.564 1.00 81.49 N +ATOM 206 CA ASN A 704 37.085 -1.787 -8.101 1.00 81.49 C +ATOM 207 C ASN A 704 36.156 -0.571 -8.061 1.00 81.49 C +ATOM 208 O ASN A 704 36.561 0.490 -7.581 1.00 81.49 O +ATOM 209 CB ASN A 704 37.497 -2.057 -9.579 1.00 81.49 C +ATOM 210 CG ASN A 704 38.755 -2.922 -9.632 1.00 81.49 C +ATOM 211 OD1 ASN A 704 39.471 -3.089 -8.641 1.00 81.49 O +ATOM 212 ND2 ASN A 704 39.039 -3.483 -10.848 1.00 81.49 N +ATOM 213 N ASP A 705 34.852 -0.822 -8.410 1.00 81.13 N +ATOM 214 CA ASP A 705 33.948 0.350 -8.462 1.00 81.13 C +ATOM 215 C ASP A 705 32.495 -0.130 -8.350 1.00 81.13 C +ATOM 216 O ASP A 705 32.152 -1.235 -8.776 1.00 81.13 O +ATOM 217 CB ASP A 705 34.140 1.103 -9.821 1.00 81.13 C +ATOM 218 CG ASP A 705 33.637 2.540 -9.712 1.00 81.13 C +ATOM 219 OD1 ASP A 705 33.220 2.983 -8.620 1.00 81.13 O +ATOM 220 OD2 ASP A 705 33.663 3.233 -10.752 1.00 81.13 O +ATOM 221 N VAL A 706 31.669 0.670 -7.535 1.00 79.16 N +ATOM 222 CA VAL A 706 30.222 0.370 -7.424 1.00 79.16 C +ATOM 223 C VAL A 706 29.431 1.613 -7.865 1.00 79.16 C +ATOM 224 O VAL A 706 29.708 2.734 -7.432 1.00 79.16 O +ATOM 225 CB VAL A 706 29.865 0.043 -5.950 1.00 79.16 C +ATOM 226 CG1 VAL A 706 28.335 -0.267 -5.832 1.00 79.16 C +ATOM 227 CG2 VAL A 706 30.681 -1.192 -5.472 1.00 79.16 C +ATOM 228 N LEU A 707 28.602 1.418 -8.760 1.00 78.80 N +ATOM 229 CA LEU A 707 27.719 2.519 -9.194 1.00 78.80 C +ATOM 230 C LEU A 707 26.261 2.116 -8.956 1.00 78.80 C +ATOM 231 O LEU A 707 25.772 1.128 -9.510 1.00 78.80 O +ATOM 232 CB LEU A 707 27.953 2.778 -10.715 1.00 78.80 C +ATOM 233 CG LEU A 707 27.086 3.938 -11.281 1.00 78.80 C +ATOM 234 CD1 LEU A 707 27.628 5.311 -10.756 1.00 78.80 C +ATOM 235 CD2 LEU A 707 27.151 3.902 -12.843 1.00 78.80 C +ATOM 236 N ILE A 708 25.550 2.846 -7.998 1.00 78.22 N +ATOM 237 CA ILE A 708 24.142 2.519 -7.677 1.00 78.22 C +ATOM 238 C ILE A 708 23.226 3.522 -8.390 1.00 78.22 C +ATOM 239 O ILE A 708 23.408 4.739 -8.292 1.00 78.22 O +ATOM 240 CB ILE A 708 23.923 2.631 -6.138 1.00 78.22 C +ATOM 241 CG1 ILE A 708 24.841 1.581 -5.436 1.00 78.22 C +ATOM 242 CG2 ILE A 708 22.417 2.356 -5.795 1.00 78.22 C +ATOM 243 CD1 ILE A 708 24.791 1.726 -3.882 1.00 78.22 C +ATOM 244 N VAL A 709 22.333 2.951 -9.228 1.00 81.79 N +ATOM 245 CA VAL A 709 21.320 3.808 -9.887 1.00 81.79 C +ATOM 246 C VAL A 709 19.976 3.599 -9.166 1.00 81.79 C +ATOM 247 O VAL A 709 19.249 2.639 -9.432 1.00 81.79 O +ATOM 248 CB VAL A 709 21.162 3.402 -11.381 1.00 81.79 C +ATOM 249 CG1 VAL A 709 20.179 4.383 -12.080 1.00 81.79 C +ATOM 250 CG2 VAL A 709 22.555 3.454 -12.085 1.00 81.79 C +ATOM 251 N PRO A 710 19.620 4.476 -8.112 1.00 87.69 N +ATOM 252 CA PRO A 710 18.490 4.315 -7.179 1.00 87.69 C +ATOM 253 C PRO A 710 17.142 4.326 -7.912 1.00 87.69 C +ATOM 254 O PRO A 710 16.234 3.587 -7.521 1.00 87.69 O +ATOM 255 CB PRO A 710 18.589 5.577 -6.263 1.00 87.69 C +ATOM 256 CG PRO A 710 20.067 5.992 -6.374 1.00 87.69 C +ATOM 257 CD PRO A 710 20.579 5.578 -7.759 1.00 87.69 C +ATOM 258 N TYR A 711 16.935 5.124 -9.079 1.00 94.92 N +ATOM 259 CA TYR A 711 15.594 5.222 -9.710 1.00 94.92 C +ATOM 260 C TYR A 711 15.307 3.972 -10.547 1.00 94.92 C +ATOM 261 O TYR A 711 14.140 3.649 -10.788 1.00 94.92 O +ATOM 262 CB TYR A 711 15.551 6.484 -10.630 1.00 94.92 C +ATOM 263 CG TYR A 711 16.593 6.428 -11.745 1.00 94.92 C +ATOM 264 CD1 TYR A 711 17.884 6.945 -11.533 1.00 94.92 C +ATOM 265 CD2 TYR A 711 16.247 5.895 -12.995 1.00 94.92 C +ATOM 266 CE1 TYR A 711 18.838 6.924 -12.573 1.00 94.92 C +ATOM 267 CE2 TYR A 711 17.193 5.884 -14.038 1.00 94.92 C +ATOM 268 CZ TYR A 711 18.476 6.375 -13.804 1.00 94.92 C +ATOM 269 OH TYR A 711 19.422 6.318 -14.825 1.00 94.92 O +ATOM 270 N ARG A 712 16.487 3.132 -10.918 1.00 87.79 N +ATOM 271 CA ARG A 712 16.259 1.905 -11.721 1.00 87.79 C +ATOM 272 C ARG A 712 16.395 0.673 -10.826 1.00 87.79 C +ATOM 273 O ARG A 712 16.232 -0.454 -11.298 1.00 87.79 O +ATOM 274 CB ARG A 712 17.329 1.813 -12.848 1.00 87.79 C +ATOM 275 CG ARG A 712 17.145 2.955 -13.878 1.00 87.79 C +ATOM 276 CD ARG A 712 18.072 2.681 -15.099 1.00 87.79 C +ATOM 277 NE ARG A 712 17.902 3.787 -16.074 1.00 87.79 N +ATOM 278 CZ ARG A 712 18.267 3.665 -17.357 1.00 87.79 C +ATOM 279 NH1 ARG A 712 18.771 2.491 -17.855 1.00 87.79 N +ATOM 280 NH2 ARG A 712 18.097 4.732 -18.192 1.00 87.79 N +ATOM 281 N LYS A 713 16.518 0.883 -9.350 1.00 82.39 N +ATOM 282 CA LYS A 713 16.755 -0.252 -8.429 1.00 82.39 C +ATOM 283 C LYS A 713 17.810 -1.214 -8.984 1.00 82.39 C +ATOM 284 O LYS A 713 17.653 -2.434 -8.905 1.00 82.39 O +ATOM 285 CB LYS A 713 15.430 -1.042 -8.214 1.00 82.39 C +ATOM 286 CG LYS A 713 14.358 -0.152 -7.528 1.00 82.39 C +ATOM 287 CD LYS A 713 13.089 -1.026 -7.297 1.00 82.39 C +ATOM 288 CE LYS A 713 11.982 -0.116 -6.694 1.00 82.39 C +ATOM 289 NZ LYS A 713 10.736 -0.998 -6.383 1.00 82.39 N +ATOM 290 N GLU A 714 18.873 -0.559 -9.716 1.00 78.86 N +ATOM 291 CA GLU A 714 19.970 -1.375 -10.288 1.00 78.86 C +ATOM 292 C GLU A 714 21.319 -0.850 -9.784 1.00 78.86 C +ATOM 293 O GLU A 714 21.451 0.321 -9.424 1.00 78.86 O +ATOM 294 CB GLU A 714 19.959 -1.276 -11.845 1.00 78.86 C +ATOM 295 CG GLU A 714 18.644 -1.894 -12.427 1.00 78.86 C +ATOM 296 CD GLU A 714 18.576 -1.666 -13.939 1.00 78.86 C +ATOM 297 OE1 GLU A 714 17.578 -2.108 -14.553 1.00 78.86 O +ATOM 298 OE2 GLU A 714 19.509 -1.057 -14.510 1.00 78.86 O +ATOM 299 N ALA A 715 22.326 -1.751 -9.558 1.00 79.28 N +ATOM 300 CA ALA A 715 23.696 -1.351 -9.186 1.00 79.28 C +ATOM 301 C ALA A 715 24.715 -2.054 -10.092 1.00 79.28 C +ATOM 302 O ALA A 715 24.502 -3.183 -10.537 1.00 79.28 O +ATOM 303 CB ALA A 715 23.958 -1.756 -7.701 1.00 79.28 C +ATOM 304 N TYR A 716 25.700 -1.265 -10.632 1.00 77.35 N +ATOM 305 CA TYR A 716 26.805 -1.832 -11.447 1.00 77.35 C +ATOM 306 C TYR A 716 28.043 -2.063 -10.572 1.00 77.35 C +ATOM 307 O TYR A 716 28.501 -1.167 -9.860 1.00 77.35 O +ATOM 308 CB TYR A 716 27.187 -0.828 -12.580 1.00 77.35 C +ATOM 309 CG TYR A 716 26.071 -0.711 -13.608 1.00 77.35 C +ATOM 310 CD1 TYR A 716 26.082 -1.531 -14.750 1.00 77.35 C +ATOM 311 CD2 TYR A 716 25.038 0.226 -13.431 1.00 77.35 C +ATOM 312 CE1 TYR A 716 25.069 -1.407 -15.727 1.00 77.35 C +ATOM 313 CE2 TYR A 716 24.017 0.351 -14.402 1.00 77.35 C +ATOM 314 CZ TYR A 716 24.061 -0.465 -15.529 1.00 77.35 C +ATOM 315 OH TYR A 716 23.040 -0.364 -16.470 1.00 77.35 O +ATOM 316 N LEU A 717 28.448 -3.326 -10.593 1.00 77.35 N +ATOM 317 CA LEU A 717 29.594 -3.689 -9.740 1.00 77.35 C +ATOM 318 C LEU A 717 30.770 -4.160 -10.612 1.00 77.35 C +ATOM 319 O LEU A 717 30.637 -5.083 -11.417 1.00 77.35 O +ATOM 320 CB LEU A 717 29.149 -4.837 -8.795 1.00 77.35 C +ATOM 321 CG LEU A 717 30.295 -5.388 -7.897 1.00 77.35 C +ATOM 322 CD1 LEU A 717 30.687 -4.321 -6.844 1.00 77.35 C +ATOM 323 CD2 LEU A 717 29.802 -6.683 -7.188 1.00 77.35 C +ATOM 324 N GLU A 718 31.909 -3.371 -10.518 1.00 77.15 N +ATOM 325 CA GLU A 718 33.144 -3.758 -11.250 1.00 77.15 C +ATOM 326 C GLU A 718 34.075 -4.557 -10.329 1.00 77.15 C +ATOM 327 O GLU A 718 34.492 -4.079 -9.272 1.00 77.15 O +ATOM 328 CB GLU A 718 33.860 -2.458 -11.729 1.00 77.15 C +ATOM 329 CG GLU A 718 35.087 -2.806 -12.632 1.00 77.15 C +ATOM 330 CD GLU A 718 35.764 -1.521 -13.116 1.00 77.15 C +ATOM 331 OE1 GLU A 718 36.751 -1.633 -13.877 1.00 77.15 O +ATOM 332 OE2 GLU A 718 35.316 -0.414 -12.744 1.00 77.15 O +ATOM 333 N MET A 719 34.284 -5.885 -10.745 1.00 77.15 N +ATOM 334 CA MET A 719 35.192 -6.744 -9.958 1.00 77.15 C +ATOM 335 C MET A 719 36.608 -6.713 -10.542 1.00 77.15 C +ATOM 336 O MET A 719 36.825 -6.241 -11.660 1.00 77.15 O +ATOM 337 CB MET A 719 34.682 -8.205 -10.006 1.00 77.15 C +ATOM 338 CG MET A 719 33.218 -8.320 -9.473 1.00 77.15 C +ATOM 339 SD MET A 719 33.117 -7.945 -7.678 1.00 77.15 S +ATOM 340 CE MET A 719 34.130 -9.305 -6.985 1.00 77.15 C +ATOM 341 N GLU A 720 37.704 -7.014 -9.691 1.00 81.93 N +ATOM 342 CA GLU A 720 39.123 -6.959 -10.124 1.00 81.93 C +ATOM 343 C GLU A 720 39.421 -8.073 -11.134 1.00 81.93 C +ATOM 344 O GLU A 720 40.190 -7.857 -12.074 1.00 81.93 O +ATOM 345 CB GLU A 720 40.016 -7.141 -8.866 1.00 81.93 C +ATOM 346 CG GLU A 720 41.513 -6.877 -9.201 1.00 81.93 C +ATOM 347 CD GLU A 720 42.358 -6.887 -7.923 1.00 81.93 C +ATOM 348 OE1 GLU A 720 43.552 -6.520 -8.008 1.00 81.93 O +ATOM 349 OE2 GLU A 720 41.836 -7.263 -6.849 1.00 81.93 O +ATOM 350 N PHE A 721 38.789 -9.221 -10.938 1.00 81.86 N +ATOM 351 CA PHE A 721 39.149 -10.354 -11.823 1.00 81.86 C +ATOM 352 C PHE A 721 37.871 -10.956 -12.417 1.00 81.86 C +ATOM 353 O PHE A 721 36.804 -10.927 -11.800 1.00 81.86 O +ATOM 354 CB PHE A 721 39.874 -11.457 -10.988 1.00 81.86 C +ATOM 355 CG PHE A 721 41.165 -10.935 -10.386 1.00 81.86 C +ATOM 356 CD1 PHE A 721 42.315 -10.807 -11.182 1.00 81.86 C +ATOM 357 CD2 PHE A 721 41.212 -10.621 -9.020 1.00 81.86 C +ATOM 358 CE1 PHE A 721 43.523 -10.347 -10.606 1.00 81.86 C +ATOM 359 CE2 PHE A 721 42.416 -10.164 -8.441 1.00 81.86 C +ATOM 360 CZ PHE A 721 43.571 -10.024 -9.236 1.00 81.86 C +ATOM 361 N LYS A 722 37.814 -11.266 -13.658 1.00 80.38 N +ATOM 362 CA LYS A 722 36.694 -11.957 -14.340 1.00 80.38 C +ATOM 363 C LYS A 722 36.345 -13.266 -13.626 1.00 80.38 C +ATOM 364 O LYS A 722 35.182 -13.674 -13.629 1.00 80.38 O +ATOM 365 CB LYS A 722 37.131 -12.299 -15.796 1.00 80.38 C +ATOM 366 CG LYS A 722 37.222 -11.006 -16.660 1.00 80.38 C +ATOM 367 CD LYS A 722 37.586 -11.429 -18.115 1.00 80.38 C +ATOM 368 CE LYS A 722 37.707 -10.133 -18.977 1.00 80.38 C +ATOM 369 NZ LYS A 722 37.857 -10.535 -20.476 1.00 80.38 N +ATOM 370 N GLU A 723 37.378 -13.983 -12.949 1.00 83.42 N +ATOM 371 CA GLU A 723 37.162 -15.239 -12.190 1.00 83.42 C +ATOM 372 C GLU A 723 36.237 -14.999 -10.989 1.00 83.42 C +ATOM 373 O GLU A 723 35.466 -15.889 -10.623 1.00 83.42 O +ATOM 374 CB GLU A 723 38.528 -15.761 -11.659 1.00 83.42 C +ATOM 375 CG GLU A 723 39.395 -16.315 -12.834 1.00 83.42 C +ATOM 376 CD GLU A 723 40.779 -16.727 -12.326 1.00 83.42 C +ATOM 377 OE1 GLU A 723 41.614 -17.137 -13.166 1.00 83.42 O +ATOM 378 OE2 GLU A 723 41.026 -16.653 -11.103 1.00 83.42 O +ATOM 379 N ALA A 724 36.322 -13.728 -10.486 1.00 79.22 N +ATOM 380 CA ALA A 724 35.442 -13.384 -9.348 1.00 79.22 C +ATOM 381 C ALA A 724 33.972 -13.300 -9.783 1.00 79.22 C +ATOM 382 O ALA A 724 33.087 -13.719 -9.033 1.00 79.22 O +ATOM 383 CB ALA A 724 35.892 -12.011 -8.766 1.00 79.22 C +ATOM 384 N ILE A 725 33.748 -12.865 -10.974 1.00 80.71 N +ATOM 385 CA ILE A 725 32.361 -12.775 -11.491 1.00 80.71 C +ATOM 386 C ILE A 725 31.802 -14.184 -11.709 1.00 80.71 C +ATOM 387 O ILE A 725 30.652 -14.450 -11.353 1.00 80.71 O +ATOM 388 CB ILE A 725 32.374 -12.011 -12.857 1.00 80.71 C +ATOM 389 CG1 ILE A 725 32.777 -10.520 -12.582 1.00 80.71 C +ATOM 390 CG2 ILE A 725 30.964 -12.099 -13.524 1.00 80.71 C +ATOM 391 CD1 ILE A 725 32.861 -9.713 -13.912 1.00 80.71 C +ATOM 392 N THR A 726 32.624 -15.132 -12.283 1.00 81.93 N +ATOM 393 CA THR A 726 32.187 -16.533 -12.482 1.00 81.93 C +ATOM 394 C THR A 726 31.918 -17.206 -11.131 1.00 81.93 C +ATOM 395 O THR A 726 30.960 -17.973 -11.011 1.00 81.93 O +ATOM 396 CB THR A 726 33.299 -17.338 -13.215 1.00 81.93 C +ATOM 397 OG1 THR A 726 34.586 -17.097 -12.577 1.00 81.93 O +ATOM 398 CG2 THR A 726 33.382 -16.893 -14.712 1.00 81.93 C +ATOM 399 N ALA A 727 32.780 -16.892 -10.214 1.00 80.71 N +ATOM 400 CA ALA A 727 32.549 -17.467 -8.872 1.00 80.71 C +ATOM 401 C ALA A 727 31.254 -16.923 -8.256 1.00 80.71 C +ATOM 402 O ALA A 727 30.516 -17.675 -7.616 1.00 80.71 O +ATOM 403 CB ALA A 727 33.762 -17.107 -7.951 1.00 80.71 C +ATOM 404 N ILE A 728 30.977 -15.615 -8.497 1.00 78.74 N +ATOM 405 CA ILE A 728 29.733 -15.010 -7.954 1.00 78.74 C +ATOM 406 C ILE A 728 28.513 -15.635 -8.634 1.00 78.74 C +ATOM 407 O ILE A 728 27.537 -15.968 -7.956 1.00 78.74 O +ATOM 408 CB ILE A 728 29.746 -13.474 -8.242 1.00 78.74 C +ATOM 409 CG1 ILE A 728 30.848 -12.821 -7.358 1.00 78.74 C +ATOM 410 CG2 ILE A 728 28.352 -12.871 -7.902 1.00 78.74 C +ATOM 411 CD1 ILE A 728 31.089 -11.322 -7.730 1.00 78.74 C +ATOM 412 N MET A 729 28.591 -15.898 -9.937 1.00 79.52 N +ATOM 413 CA MET A 729 27.440 -16.479 -10.669 1.00 79.52 C +ATOM 414 C MET A 729 27.235 -17.937 -10.254 1.00 79.52 C +ATOM 415 O MET A 729 26.092 -18.389 -10.159 1.00 79.52 O +ATOM 416 CB MET A 729 27.736 -16.443 -12.197 1.00 79.52 C +ATOM 417 CG MET A 729 27.669 -14.990 -12.762 1.00 79.52 C +ATOM 418 SD MET A 729 25.980 -14.272 -12.559 1.00 79.52 S +ATOM 419 CE MET A 729 25.062 -15.207 -13.850 1.00 79.52 C +ATOM 420 N LYS A 730 28.331 -18.649 -9.912 1.00 80.92 N +ATOM 421 CA LYS A 730 28.208 -20.042 -9.438 1.00 80.92 C +ATOM 422 C LYS A 730 27.614 -20.083 -8.023 1.00 80.92 C +ATOM 423 O LYS A 730 26.806 -20.964 -7.723 1.00 80.92 O +ATOM 424 CB LYS A 730 29.629 -20.670 -9.414 1.00 80.92 C +ATOM 425 CG LYS A 730 29.555 -22.214 -9.283 1.00 80.92 C +ATOM 426 CD LYS A 730 31.010 -22.772 -9.365 1.00 80.92 C +ATOM 427 CE LYS A 730 30.938 -24.318 -9.152 1.00 80.92 C +ATOM 428 NZ LYS A 730 32.375 -24.912 -9.254 1.00 80.92 N +ATOM 429 N TYR A 731 27.980 -19.027 -7.234 1.00 80.71 N +ATOM 430 CA TYR A 731 27.467 -18.972 -5.848 1.00 80.71 C +ATOM 431 C TYR A 731 25.989 -18.567 -5.841 1.00 80.71 C +ATOM 432 O TYR A 731 25.206 -19.120 -5.065 1.00 80.71 O +ATOM 433 CB TYR A 731 28.305 -17.916 -5.069 1.00 80.71 C +ATOM 434 CG TYR A 731 27.816 -17.754 -3.636 1.00 80.71 C +ATOM 435 CD1 TYR A 731 27.037 -16.640 -3.274 1.00 80.71 C +ATOM 436 CD2 TYR A 731 28.156 -18.715 -2.665 1.00 80.71 C +ATOM 437 CE1 TYR A 731 26.606 -16.481 -1.932 1.00 80.71 C +ATOM 438 CE2 TYR A 731 27.728 -18.558 -1.331 1.00 80.71 C +ATOM 439 CZ TYR A 731 26.962 -17.438 -0.998 1.00 80.71 C +ATOM 440 OH TYR A 731 26.540 -17.312 0.323 1.00 80.71 O +ATOM 441 N ILE A 732 25.522 -17.731 -6.739 1.00 81.42 N +ATOM 442 CA ILE A 732 24.125 -17.246 -6.758 1.00 81.42 C +ATOM 443 C ILE A 732 23.209 -18.381 -7.230 1.00 81.42 C +ATOM 444 O ILE A 732 22.034 -18.416 -6.861 1.00 81.42 O +ATOM 445 CB ILE A 732 24.009 -16.052 -7.757 1.00 81.42 C +ATOM 446 CG1 ILE A 732 24.783 -14.843 -7.153 1.00 81.42 C +ATOM 447 CG2 ILE A 732 22.496 -15.686 -7.971 1.00 81.42 C +ATOM 448 CD1 ILE A 732 24.953 -13.706 -8.210 1.00 81.42 C +ATOM 449 N GLU A 733 23.862 -19.428 -8.068 1.00 84.80 N +ATOM 450 CA GLU A 733 23.053 -20.594 -8.493 1.00 84.80 C +ATOM 451 C GLU A 733 22.780 -21.527 -7.311 1.00 84.80 C +ATOM 452 O GLU A 733 21.724 -22.160 -7.264 1.00 84.80 O +ATOM 453 CB GLU A 733 23.847 -21.391 -9.574 1.00 84.80 C +ATOM 454 CG GLU A 733 23.740 -20.682 -10.964 1.00 84.80 C +ATOM 455 CD GLU A 733 24.601 -21.413 -11.999 1.00 84.80 C +ATOM 456 OE1 GLU A 733 24.720 -20.894 -13.132 1.00 84.80 O +ATOM 457 OE2 GLU A 733 25.145 -22.495 -11.686 1.00 84.80 O +ATOM 458 N THR A 734 23.619 -21.471 -6.333 1.00 87.17 N +ATOM 459 CA THR A 734 23.429 -22.404 -5.196 1.00 87.17 C +ATOM 460 C THR A 734 22.825 -21.660 -4.001 1.00 87.17 C +ATOM 461 O THR A 734 22.037 -22.243 -3.253 1.00 87.17 O +ATOM 462 CB THR A 734 24.798 -23.000 -4.763 1.00 87.17 C +ATOM 463 OG1 THR A 734 25.753 -21.932 -4.507 1.00 87.17 O +ATOM 464 CG2 THR A 734 25.363 -23.921 -5.893 1.00 87.17 C +ATOM 465 N THR A 735 23.070 -20.285 -3.908 1.00 88.64 N +ATOM 466 CA THR A 735 22.624 -19.550 -2.706 1.00 88.64 C +ATOM 467 C THR A 735 22.057 -18.190 -3.143 1.00 88.64 C +ATOM 468 O THR A 735 22.741 -17.387 -3.782 1.00 88.64 O +ATOM 469 CB THR A 735 23.818 -19.292 -1.743 1.00 88.64 C +ATOM 470 OG1 THR A 735 24.448 -20.577 -1.468 1.00 88.64 O +ATOM 471 CG2 THR A 735 23.290 -18.734 -0.377 1.00 88.64 C +ATOM 472 N PRO A 736 20.660 -18.022 -3.023 1.00 89.74 N +ATOM 473 CA PRO A 736 20.066 -16.723 -3.378 1.00 89.74 C +ATOM 474 C PRO A 736 20.628 -15.585 -2.506 1.00 89.74 C +ATOM 475 O PRO A 736 20.801 -15.739 -1.294 1.00 89.74 O +ATOM 476 CB PRO A 736 18.555 -16.916 -3.077 1.00 89.74 C +ATOM 477 CG PRO A 736 18.500 -18.121 -2.127 1.00 89.74 C +ATOM 478 CD PRO A 736 19.782 -18.943 -2.327 1.00 89.74 C +ATOM 479 N LEU A 737 21.210 -14.518 -3.081 1.00 89.74 N +ATOM 480 CA LEU A 737 21.807 -13.406 -2.305 1.00 89.74 C +ATOM 481 C LEU A 737 20.716 -12.420 -1.878 1.00 89.74 C +ATOM 482 O LEU A 737 19.801 -12.108 -2.643 1.00 89.74 O +ATOM 483 CB LEU A 737 22.821 -12.645 -3.208 1.00 89.74 C +ATOM 484 CG LEU A 737 24.103 -13.462 -3.521 1.00 89.74 C +ATOM 485 CD1 LEU A 737 24.988 -12.662 -4.515 1.00 89.74 C +ATOM 486 CD2 LEU A 737 24.909 -13.716 -2.204 1.00 89.74 C +ATOM 487 N THR A 738 20.666 -12.153 -0.533 1.00 87.99 N +ATOM 488 CA THR A 738 19.670 -11.205 0.012 1.00 87.99 C +ATOM 489 C THR A 738 20.392 -9.963 0.559 1.00 87.99 C +ATOM 490 O THR A 738 21.461 -10.061 1.164 1.00 87.99 O +ATOM 491 CB THR A 738 18.871 -11.856 1.175 1.00 87.99 C +ATOM 492 OG1 THR A 738 19.778 -12.435 2.152 1.00 87.99 O +ATOM 493 CG2 THR A 738 17.976 -13.018 0.625 1.00 87.99 C +ATOM 494 N ILE A 739 19.968 -8.755 -0.010 1.00 89.42 N +ATOM 495 CA ILE A 739 20.478 -7.486 0.570 1.00 89.42 C +ATOM 496 C ILE A 739 19.360 -6.834 1.384 1.00 89.42 C +ATOM 497 O ILE A 739 18.259 -6.590 0.884 1.00 89.42 O +ATOM 498 CB ILE A 739 20.902 -6.520 -0.582 1.00 89.42 C +ATOM 499 CG1 ILE A 739 22.065 -7.188 -1.379 1.00 89.42 C +ATOM 500 CG2 ILE A 739 21.355 -5.148 0.033 1.00 89.42 C +ATOM 501 CD1 ILE A 739 22.486 -6.329 -2.626 1.00 89.42 C +ATOM 502 N LYS A 740 19.533 -6.663 2.793 1.00 99.46 N +ATOM 503 CA LYS A 740 18.545 -6.096 3.741 1.00 99.46 C +ATOM 504 C LYS A 740 17.238 -6.887 3.706 1.00 99.46 C +ATOM 505 O LYS A 740 16.150 -6.309 3.687 1.00 99.46 O +ATOM 506 CB LYS A 740 18.242 -4.614 3.390 1.00 99.46 C +ATOM 507 CG LYS A 740 19.470 -3.705 3.659 1.00 99.46 C +ATOM 508 CD LYS A 740 19.006 -2.219 3.525 1.00 99.46 C +ATOM 509 CE LYS A 740 20.185 -1.306 3.979 1.00 99.46 C +ATOM 510 NZ LYS A 740 19.741 0.184 3.894 1.00 99.46 N +ATOM 511 N GLY A 741 17.259 -8.244 3.545 1.00 98.66 N +ATOM 512 CA GLY A 741 16.050 -9.093 3.695 1.00 98.66 C +ATOM 513 C GLY A 741 15.324 -9.283 2.357 1.00 98.66 C +ATOM 514 O GLY A 741 14.354 -10.041 2.297 1.00 98.66 O +ATOM 515 N LYS A 742 15.843 -8.593 1.113 1.00 93.82 N +ATOM 516 CA LYS A 742 15.189 -8.743 -0.206 1.00 93.82 C +ATOM 517 C LYS A 742 16.146 -9.455 -1.175 1.00 93.82 C +ATOM 518 O LYS A 742 17.354 -9.213 -1.173 1.00 93.82 O +ATOM 519 CB LYS A 742 14.849 -7.327 -0.758 1.00 93.82 C +ATOM 520 CG LYS A 742 13.785 -7.421 -1.892 1.00 93.82 C +ATOM 521 CD LYS A 742 13.133 -6.017 -2.056 1.00 93.82 C +ATOM 522 CE LYS A 742 12.091 -6.094 -3.209 1.00 93.82 C +ATOM 523 NZ LYS A 742 11.286 -4.766 -3.261 1.00 93.82 N +ATOM 524 N SER A 743 15.687 -10.558 -1.941 1.00 89.53 N +ATOM 525 CA SER A 743 16.507 -11.349 -2.891 1.00 89.53 C +ATOM 526 C SER A 743 16.900 -10.506 -4.114 1.00 89.53 C +ATOM 527 O SER A 743 16.091 -9.750 -4.655 1.00 89.53 O +ATOM 528 CB SER A 743 15.698 -12.593 -3.353 1.00 89.53 C +ATOM 529 OG SER A 743 14.473 -12.178 -4.002 1.00 89.53 O +ATOM 530 N VAL A 744 18.150 -10.311 -4.307 1.00 81.63 N +ATOM 531 CA VAL A 744 18.685 -9.467 -5.401 1.00 81.63 C +ATOM 532 C VAL A 744 19.081 -10.375 -6.575 1.00 81.63 C +ATOM 533 O VAL A 744 19.483 -11.526 -6.386 1.00 81.63 O +ATOM 534 CB VAL A 744 19.952 -8.703 -4.906 1.00 81.63 C +ATOM 535 CG1 VAL A 744 20.460 -7.771 -6.037 1.00 81.63 C +ATOM 536 CG2 VAL A 744 19.570 -7.856 -3.650 1.00 81.63 C +ATOM 537 N LYS A 745 18.814 -9.920 -7.854 1.00 83.42 N +ATOM 538 CA LYS A 745 19.214 -10.657 -9.075 1.00 83.42 C +ATOM 539 C LYS A 745 20.511 -10.075 -9.655 1.00 83.42 C +ATOM 540 O LYS A 745 20.648 -8.861 -9.826 1.00 83.42 O +ATOM 541 CB LYS A 745 18.094 -10.523 -10.142 1.00 83.42 C +ATOM 542 CG LYS A 745 16.851 -11.358 -9.743 1.00 83.42 C +ATOM 543 CD LYS A 745 15.811 -11.245 -10.898 1.00 83.42 C +ATOM 544 CE LYS A 745 14.520 -12.002 -10.450 1.00 83.42 C +ATOM 545 NZ LYS A 745 13.451 -11.801 -11.552 1.00 83.42 N +ATOM 546 N ILE A 746 21.454 -10.971 -9.913 1.00 80.44 N +ATOM 547 CA ILE A 746 22.785 -10.513 -10.382 1.00 80.44 C +ATOM 548 C ILE A 746 23.002 -11.041 -11.806 1.00 80.44 C +ATOM 549 O ILE A 746 22.769 -12.217 -12.096 1.00 80.44 O +ATOM 550 CB ILE A 746 23.891 -11.084 -9.453 1.00 80.44 C +ATOM 551 CG1 ILE A 746 23.646 -10.544 -8.007 1.00 80.44 C +ATOM 552 CG2 ILE A 746 25.291 -10.632 -9.974 1.00 80.44 C +ATOM 553 CD1 ILE A 746 24.581 -11.235 -6.959 1.00 80.44 C +ATOM 554 N CYS A 747 23.333 -10.123 -12.793 1.00 85.99 N +ATOM 555 CA CYS A 747 23.587 -10.518 -14.199 1.00 85.99 C +ATOM 556 C CYS A 747 24.879 -9.859 -14.698 1.00 85.99 C +ATOM 557 O CYS A 747 25.321 -8.838 -14.168 1.00 85.99 O +ATOM 558 CB CYS A 747 22.409 -10.033 -15.099 1.00 85.99 C +ATOM 559 SG CYS A 747 20.843 -10.766 -14.588 1.00 85.99 S +ATOM 560 N VAL A 748 25.605 -10.516 -15.664 1.00 86.19 N +ATOM 561 CA VAL A 748 26.815 -9.971 -16.326 1.00 86.19 C +ATOM 562 C VAL A 748 26.386 -9.231 -17.606 1.00 86.19 C +ATOM 563 O VAL A 748 25.668 -9.776 -18.449 1.00 86.19 O +ATOM 564 CB VAL A 748 27.772 -11.135 -16.713 1.00 86.19 C +ATOM 565 CG1 VAL A 748 29.056 -10.565 -17.402 1.00 86.19 C +ATOM 566 CG2 VAL A 748 28.166 -11.926 -15.427 1.00 86.19 C +ATOM 567 N PRO A 749 26.469 -7.862 -17.658 1.00124.40 N +ATOM 568 CA PRO A 749 26.052 -7.030 -18.798 1.00124.40 C +ATOM 569 C PRO A 749 26.883 -7.322 -20.054 1.00124.40 C +ATOM 570 O PRO A 749 28.095 -7.546 -19.975 1.00124.40 O +ATOM 571 CB PRO A 749 26.326 -5.577 -18.308 1.00124.40 C +ATOM 572 CG PRO A 749 27.150 -5.755 -17.028 1.00124.40 C +ATOM 573 CD PRO A 749 27.131 -7.245 -16.649 1.00124.40 C +END diff --git a/tests/files/pdb/6VHI_af.pdb b/tests/files/pdb/6VHI_af.pdb new file mode 100644 index 00000000..768e082d --- /dev/null +++ b/tests/files/pdb/6VHI_af.pdb @@ -0,0 +1,843 @@ +REMARK TITLE 6VHI AlphaFold-start MR +REMARK Log-Likelihood Gain: 840.246 +REMARK RFZ=6.6 TFZ=8.6 PAK=1 LLG=201 TFZ==9.8 LLG=840 TFZ==15.4 PAK=0 LLG=840 TFZ==15.4 +REMARK ENSEMBLE e_Q9H6K1 EULER 225.84 50.33 132.02 FRAC 0.232 0.008 0.384 +CRYST1 135.931 40.977 23.291 90.00 91.43 90.00 C 1 2 1 4 +SCALE1 0.007357 -0.000000 0.000184 0.00000 +SCALE2 0.000000 0.024404 -0.000000 0.00000 +SCALE3 0.000000 0.000000 0.042948 0.00000 +ATOM 1 N SER A 71 45.472 4.798 1.271 1.00 67.13 N +ATOM 2 CA SER A 71 44.644 3.773 0.692 1.00 67.13 C +ATOM 3 C SER A 71 43.230 4.194 1.034 1.00 67.13 C +ATOM 4 O SER A 71 42.831 4.194 2.201 1.00 67.13 O +ATOM 5 CB SER A 71 45.023 2.445 1.345 1.00 67.13 C +ATOM 6 OG SER A 71 44.092 1.425 1.065 1.00 67.13 O +ATOM 7 N VAL A 72 42.495 4.663 0.025 1.00 39.14 N +ATOM 8 CA VAL A 72 41.076 4.947 0.202 1.00 39.14 C +ATOM 9 C VAL A 72 40.455 3.637 0.687 1.00 39.14 C +ATOM 10 O VAL A 72 40.539 2.632 -0.029 1.00 39.14 O +ATOM 11 CB VAL A 72 40.407 5.451 -1.085 1.00 39.14 C +ATOM 12 CG1 VAL A 72 38.901 5.670 -0.877 1.00 39.14 C +ATOM 13 CG2 VAL A 72 41.025 6.786 -1.519 1.00 39.14 C +ATOM 14 N PRO A 73 39.880 3.608 1.903 1.00 30.83 N +ATOM 15 CA PRO A 73 39.348 2.371 2.441 1.00 30.83 C +ATOM 16 C PRO A 73 38.279 1.829 1.493 1.00 30.83 C +ATOM 17 O PRO A 73 37.608 2.575 0.781 1.00 30.83 O +ATOM 18 CB PRO A 73 38.834 2.705 3.843 1.00 30.83 C +ATOM 19 CG PRO A 73 38.611 4.216 3.805 1.00 30.83 C +ATOM 20 CD PRO A 73 39.659 4.714 2.820 1.00 30.83 C +ATOM 21 N SER A 74 38.132 0.511 1.454 1.00 28.76 N +ATOM 22 CA SER A 74 37.152 -0.151 0.594 1.00 28.76 C +ATOM 23 C SER A 74 36.127 -0.870 1.456 1.00 28.76 C +ATOM 24 O SER A 74 36.469 -1.488 2.463 1.00 28.76 O +ATOM 25 CB SER A 74 37.850 -1.078 -0.398 1.00 28.76 C +ATOM 26 OG SER A 74 38.668 -0.313 -1.264 1.00 28.76 O +ATOM 27 N MET A 75 34.858 -0.742 1.085 1.00 27.74 N +ATOM 28 CA MET A 75 33.746 -1.432 1.727 1.00 27.74 C +ATOM 29 C MET A 75 32.887 -2.047 0.636 1.00 27.74 C +ATOM 30 O MET A 75 32.533 -1.344 -0.313 1.00 27.74 O +ATOM 31 CB MET A 75 32.949 -0.466 2.626 1.00 27.74 C +ATOM 32 CG MET A 75 31.817 -1.159 3.397 1.00 27.74 C +ATOM 33 SD MET A 75 30.284 -1.452 2.470 1.00 27.74 S +ATOM 34 CE MET A 75 29.527 0.193 2.450 1.00 27.74 C +ATOM 35 N SER A 76 32.569 -3.329 0.786 1.00 27.59 N +ATOM 36 CA SER A 76 31.589 -4.020 -0.046 1.00 27.59 C +ATOM 37 C SER A 76 30.373 -4.429 0.770 1.00 27.59 C +ATOM 38 O SER A 76 30.493 -4.898 1.909 1.00 27.59 O +ATOM 39 CB SER A 76 32.196 -5.229 -0.757 1.00 27.59 C +ATOM 40 OG SER A 76 32.797 -6.115 0.165 1.00 27.59 O +ATOM 41 N PHE A 77 29.204 -4.289 0.160 1.00 27.59 N +ATOM 42 CA PHE A 77 27.968 -4.843 0.693 1.00 27.59 C +ATOM 43 C PHE A 77 27.985 -6.377 0.616 1.00 27.59 C +ATOM 44 O PHE A 77 28.399 -6.944 -0.394 1.00 27.59 O +ATOM 45 CB PHE A 77 26.805 -4.250 -0.098 1.00 27.59 C +ATOM 46 CG PHE A 77 25.453 -4.826 0.261 1.00 27.59 C +ATOM 47 CD1 PHE A 77 24.733 -5.558 -0.702 1.00 27.59 C +ATOM 48 CD2 PHE A 77 24.920 -4.649 1.551 1.00 27.59 C +ATOM 49 CE1 PHE A 77 23.458 -6.058 -0.398 1.00 27.59 C +ATOM 50 CE2 PHE A 77 23.654 -5.172 1.857 1.00 27.59 C +ATOM 51 CZ PHE A 77 22.908 -5.841 0.873 1.00 27.59 C +ATOM 52 N VAL A 78 27.544 -7.049 1.684 1.00 27.68 N +ATOM 53 CA VAL A 78 27.426 -8.516 1.722 1.00 27.68 C +ATOM 54 C VAL A 78 25.966 -8.930 1.572 1.00 27.68 C +ATOM 55 O VAL A 78 25.635 -9.654 0.639 1.00 27.68 O +ATOM 56 CB VAL A 78 28.047 -9.100 3.005 1.00 27.68 C +ATOM 57 CG1 VAL A 78 27.951 -10.630 3.059 1.00 27.68 C +ATOM 58 CG2 VAL A 78 29.524 -8.714 3.159 1.00 27.68 C +ATOM 59 N GLU A 79 25.100 -8.482 2.485 1.00 27.97 N +ATOM 60 CA GLU A 79 23.667 -8.790 2.476 1.00 27.97 C +ATOM 61 C GLU A 79 22.875 -7.863 3.407 1.00 27.97 C +ATOM 62 O GLU A 79 23.415 -7.257 4.342 1.00 27.97 O +ATOM 63 CB GLU A 79 23.405 -10.263 2.867 1.00 27.97 C +ATOM 64 CG GLU A 79 23.835 -10.612 4.302 1.00 27.97 C +ATOM 65 CD GLU A 79 23.704 -12.100 4.662 1.00 27.97 C +ATOM 66 OE1 GLU A 79 24.290 -12.474 5.712 1.00 27.97 O +ATOM 67 OE2 GLU A 79 23.073 -12.875 3.909 1.00 27.97 O +ATOM 68 N ASP A 80 21.567 -7.795 3.177 1.00 28.13 N +ATOM 69 CA ASP A 80 20.620 -7.248 4.136 1.00 28.13 C +ATOM 70 C ASP A 80 20.364 -8.294 5.224 1.00 28.13 C +ATOM 71 O ASP A 80 19.803 -9.356 4.963 1.00 28.13 O +ATOM 72 CB ASP A 80 19.324 -6.859 3.419 1.00 28.13 C +ATOM 73 CG ASP A 80 19.513 -5.630 2.537 1.00 28.13 C +ATOM 74 OD1 ASP A 80 20.029 -4.615 3.060 1.00 28.13 O +ATOM 75 OD2 ASP A 80 19.034 -5.648 1.391 1.00 28.13 O +ATOM 76 N VAL A 81 20.781 -7.997 6.458 1.00 28.61 N +ATOM 77 CA VAL A 81 20.559 -8.888 7.611 1.00 28.61 C +ATOM 78 C VAL A 81 19.122 -8.764 8.100 1.00 28.61 C +ATOM 79 O VAL A 81 18.500 -9.752 8.488 1.00 28.61 O +ATOM 80 CB VAL A 81 21.526 -8.562 8.760 1.00 28.61 C +ATOM 81 CG1 VAL A 81 21.308 -9.441 9.998 1.00 28.61 C +ATOM 82 CG2 VAL A 81 22.980 -8.729 8.321 1.00 28.61 C +ATOM 83 N THR A 82 18.590 -7.543 8.077 1.00 30.33 N +ATOM 84 CA THR A 82 17.172 -7.282 8.308 1.00 30.33 C +ATOM 85 C THR A 82 16.602 -6.489 7.146 1.00 30.33 C +ATOM 86 O THR A 82 17.297 -5.676 6.529 1.00 30.33 O +ATOM 87 CB THR A 82 16.892 -6.560 9.637 1.00 30.33 C +ATOM 88 OG1 THR A 82 17.370 -5.240 9.621 1.00 30.33 O +ATOM 89 CG2 THR A 82 17.499 -7.256 10.853 1.00 30.33 C +ATOM 90 N ILE A 83 15.303 -6.685 6.909 1.00 35.29 N +ATOM 91 CA ILE A 83 14.522 -5.929 5.932 1.00 35.29 C +ATOM 92 C ILE A 83 15.094 -6.092 4.527 1.00 35.29 C +ATOM 93 O ILE A 83 15.875 -5.257 4.079 1.00 35.29 O +ATOM 94 CB ILE A 83 14.377 -4.455 6.364 1.00 35.29 C +ATOM 95 CG1 ILE A 83 13.967 -4.374 7.852 1.00 35.29 C +ATOM 96 CG2 ILE A 83 13.449 -3.689 5.404 1.00 35.29 C +ATOM 97 CD1 ILE A 83 13.558 -2.992 8.325 1.00 35.29 C +ATOM 98 N GLY A 84 14.722 -7.159 3.825 1.00 35.16 N +ATOM 99 CA GLY A 84 15.080 -7.304 2.416 1.00 35.16 C +ATOM 100 C GLY A 84 14.500 -6.191 1.534 1.00 35.16 C +ATOM 101 O GLY A 84 13.671 -5.380 1.960 1.00 35.16 O +ATOM 102 N GLU A 85 14.927 -6.163 0.276 1.00 37.53 N +ATOM 103 CA GLU A 85 14.412 -5.209 -0.704 1.00 37.53 C +ATOM 104 C GLU A 85 12.884 -5.330 -0.854 1.00 37.53 C +ATOM 105 O GLU A 85 12.357 -6.373 -1.244 1.00 37.53 O +ATOM 106 CB GLU A 85 15.098 -5.398 -2.065 1.00 37.53 C +ATOM 107 CG GLU A 85 16.610 -5.119 -2.027 1.00 37.53 C +ATOM 108 CD GLU A 85 17.252 -5.134 -3.427 1.00 37.53 C +ATOM 109 OE1 GLU A 85 18.382 -4.612 -3.554 1.00 37.53 O +ATOM 110 OE2 GLU A 85 16.613 -5.649 -4.376 1.00 37.53 O +ATOM 111 N GLY A 86 12.166 -4.241 -0.555 1.00 48.33 N +ATOM 112 CA GLY A 86 10.720 -4.141 -0.777 1.00 48.33 C +ATOM 113 C GLY A 86 9.847 -4.823 0.279 1.00 48.33 C +ATOM 114 O GLY A 86 8.644 -4.972 0.058 1.00 48.33 O +ATOM 115 N GLU A 87 10.402 -5.224 1.426 1.00 34.03 N +ATOM 116 CA GLU A 87 9.588 -5.766 2.515 1.00 34.03 C +ATOM 117 C GLU A 87 8.555 -4.749 3.031 1.00 34.03 C +ATOM 118 O GLU A 87 8.796 -3.540 3.105 1.00 34.03 O +ATOM 119 CB GLU A 87 10.447 -6.300 3.667 1.00 34.03 C +ATOM 120 CG GLU A 87 11.120 -7.635 3.319 1.00 34.03 C +ATOM 121 CD GLU A 87 11.831 -8.278 4.521 1.00 34.03 C +ATOM 122 OE1 GLU A 87 12.473 -9.329 4.330 1.00 34.03 O +ATOM 123 OE2 GLU A 87 11.774 -7.721 5.636 1.00 34.03 O +ATOM 124 N SER A 88 7.376 -5.263 3.392 1.00 30.92 N +ATOM 125 CA SER A 88 6.235 -4.451 3.812 1.00 30.92 C +ATOM 126 C SER A 88 6.243 -4.153 5.311 1.00 30.92 C +ATOM 127 O SER A 88 6.265 -5.079 6.133 1.00 30.92 O +ATOM 128 CB SER A 88 4.917 -5.110 3.403 1.00 30.92 C +ATOM 129 OG SER A 88 4.797 -6.338 4.099 1.00 30.92 O +ATOM 130 N ILE A 89 6.110 -2.881 5.672 1.00 29.48 N +ATOM 131 CA ILE A 89 6.045 -2.394 7.052 1.00 29.48 C +ATOM 132 C ILE A 89 4.849 -1.441 7.241 1.00 29.48 C +ATOM 133 O ILE A 89 4.552 -0.669 6.328 1.00 29.48 O +ATOM 134 CB ILE A 89 7.403 -1.760 7.427 1.00 29.48 C +ATOM 135 CG1 ILE A 89 7.462 -1.522 8.946 1.00 29.48 C +ATOM 136 CG2 ILE A 89 7.699 -0.466 6.641 1.00 29.48 C +ATOM 137 CD1 ILE A 89 8.880 -1.317 9.469 1.00 29.48 C +ATOM 138 N PRO A 90 4.140 -1.457 8.389 1.00 29.41 N +ATOM 139 CA PRO A 90 3.084 -0.480 8.640 1.00 29.41 C +ATOM 140 C PRO A 90 3.612 0.951 8.704 1.00 29.41 C +ATOM 141 O PRO A 90 4.796 1.147 8.991 1.00 29.41 O +ATOM 142 CB PRO A 90 2.396 -0.868 9.948 1.00 29.41 C +ATOM 143 CG PRO A 90 2.788 -2.323 10.126 1.00 29.41 C +ATOM 144 CD PRO A 90 4.171 -2.418 9.478 1.00 29.41 C +ATOM 145 N PRO A 91 2.752 1.959 8.491 1.00 28.80 N +ATOM 146 CA PRO A 91 3.092 3.352 8.761 1.00 28.80 C +ATOM 147 C PRO A 91 3.526 3.597 10.214 1.00 28.80 C +ATOM 148 O PRO A 91 3.135 2.847 11.111 1.00 28.80 O +ATOM 149 CB PRO A 91 1.832 4.153 8.422 1.00 28.80 C +ATOM 150 CG PRO A 91 1.060 3.256 7.458 1.00 28.80 C +ATOM 151 CD PRO A 91 1.408 1.853 7.943 1.00 28.80 C +ATOM 152 N ASP A 92 4.320 4.648 10.436 1.00 28.98 N +ATOM 153 CA ASP A 92 4.732 5.150 11.762 1.00 28.98 C +ATOM 154 C ASP A 92 5.309 4.068 12.702 1.00 28.98 C +ATOM 155 O ASP A 92 5.138 4.091 13.924 1.00 28.98 O +ATOM 156 CB ASP A 92 3.595 5.969 12.400 1.00 28.98 C +ATOM 157 CG ASP A 92 3.221 7.195 11.563 1.00 28.98 C +ATOM 158 OD1 ASP A 92 4.101 8.063 11.358 1.00 28.98 O +ATOM 159 OD2 ASP A 92 2.049 7.290 11.129 1.00 28.98 O +ATOM 160 N THR A 93 5.989 3.080 12.124 1.00 29.52 N +ATOM 161 CA THR A 93 6.493 1.906 12.832 1.00 29.52 C +ATOM 162 C THR A 93 8.002 1.979 12.971 1.00 29.52 C +ATOM 163 O THR A 93 8.728 2.125 11.990 1.00 29.52 O +ATOM 164 CB THR A 93 6.061 0.618 12.130 1.00 29.52 C +ATOM 165 OG1 THR A 93 4.660 0.549 12.126 1.00 29.52 O +ATOM 166 CG2 THR A 93 6.526 -0.632 12.874 1.00 29.52 C +ATOM 167 N GLN A 94 8.487 1.834 14.204 1.00 29.10 N +ATOM 168 CA GLN A 94 9.916 1.731 14.481 1.00 29.10 C +ATOM 169 C GLN A 94 10.461 0.362 14.063 1.00 29.10 C +ATOM 170 O GLN A 94 9.856 -0.667 14.367 1.00 29.10 O +ATOM 171 CB GLN A 94 10.195 1.991 15.965 1.00 29.10 C +ATOM 172 CG GLN A 94 9.966 3.458 16.340 1.00 29.10 C +ATOM 173 CD GLN A 94 10.372 3.716 17.785 1.00 29.10 C +ATOM 174 OE1 GLN A 94 11.540 3.755 18.129 1.00 29.10 O +ATOM 175 NE2 GLN A 94 9.436 3.881 18.692 1.00 29.10 N +ATOM 176 N PHE A 95 11.624 0.355 13.417 1.00 29.61 N +ATOM 177 CA PHE A 95 12.324 -0.850 12.980 1.00 29.61 C +ATOM 178 C PHE A 95 13.842 -0.651 13.004 1.00 29.61 C +ATOM 179 O PHE A 95 14.331 0.479 12.971 1.00 29.61 O +ATOM 180 CB PHE A 95 11.851 -1.232 11.574 1.00 29.61 C +ATOM 181 CG PHE A 95 12.211 -0.233 10.486 1.00 29.61 C +ATOM 182 CD1 PHE A 95 11.333 0.812 10.147 1.00 29.61 C +ATOM 183 CD2 PHE A 95 13.439 -0.339 9.812 1.00 29.61 C +ATOM 184 CE1 PHE A 95 11.667 1.713 9.123 1.00 29.61 C +ATOM 185 CE2 PHE A 95 13.756 0.535 8.760 1.00 29.61 C +ATOM 186 CZ PHE A 95 12.881 1.578 8.436 1.00 29.61 C +ATOM 187 N VAL A 96 14.593 -1.749 13.028 1.00 28.98 N +ATOM 188 CA VAL A 96 16.052 -1.736 12.885 1.00 28.98 C +ATOM 189 C VAL A 96 16.413 -2.277 11.512 1.00 28.98 C +ATOM 190 O VAL A 96 16.067 -3.411 11.166 1.00 28.98 O +ATOM 191 CB VAL A 96 16.748 -2.531 14.002 1.00 28.98 C +ATOM 192 CG1 VAL A 96 18.273 -2.499 13.830 1.00 28.98 C +ATOM 193 CG2 VAL A 96 16.437 -1.931 15.374 1.00 28.98 C +ATOM 194 N LYS A 97 17.129 -1.466 10.735 1.00 28.58 N +ATOM 195 CA LYS A 97 17.767 -1.904 9.496 1.00 28.58 C +ATOM 196 C LYS A 97 19.209 -2.291 9.790 1.00 28.58 C +ATOM 197 O LYS A 97 19.953 -1.477 10.332 1.00 28.58 O +ATOM 198 CB LYS A 97 17.680 -0.792 8.451 1.00 28.58 C +ATOM 199 CG LYS A 97 18.364 -1.171 7.134 1.00 28.58 C +ATOM 200 CD LYS A 97 17.681 -2.302 6.365 1.00 28.58 C +ATOM 201 CE LYS A 97 18.470 -2.534 5.079 1.00 28.58 C +ATOM 202 NZ LYS A 97 17.877 -3.582 4.235 1.00 28.58 N +ATOM 203 N THR A 98 19.594 -3.499 9.404 1.00 27.97 N +ATOM 204 CA THR A 98 20.945 -4.023 9.581 1.00 27.97 C +ATOM 205 C THR A 98 21.491 -4.496 8.248 1.00 27.97 C +ATOM 206 O THR A 98 20.879 -5.328 7.578 1.00 27.97 O +ATOM 207 CB THR A 98 20.978 -5.160 10.608 1.00 27.97 C +ATOM 208 OG1 THR A 98 20.423 -4.732 11.827 1.00 27.97 O +ATOM 209 CG2 THR A 98 22.402 -5.625 10.916 1.00 27.97 C +ATOM 210 N TRP A 99 22.667 -3.993 7.892 1.00 27.59 N +ATOM 211 CA TRP A 99 23.434 -4.462 6.745 1.00 27.59 C +ATOM 212 C TRP A 99 24.658 -5.211 7.228 1.00 27.59 C +ATOM 213 O TRP A 99 25.315 -4.784 8.181 1.00 27.59 O +ATOM 214 CB TRP A 99 23.873 -3.290 5.871 1.00 27.59 C +ATOM 215 CG TRP A 99 22.770 -2.428 5.356 1.00 27.59 C +ATOM 216 CD1 TRP A 99 22.139 -2.570 4.170 1.00 27.59 C +ATOM 217 CD2 TRP A 99 22.173 -1.259 5.987 1.00 27.59 C +ATOM 218 NE1 TRP A 99 21.259 -1.532 3.984 1.00 27.59 N +ATOM 219 CE2 TRP A 99 21.272 -0.673 5.054 1.00 27.59 C +ATOM 220 CE3 TRP A 99 22.316 -0.617 7.235 1.00 27.59 C +ATOM 221 CZ2 TRP A 99 20.622 0.533 5.303 1.00 27.59 C +ATOM 222 CZ3 TRP A 99 21.604 0.561 7.525 1.00 27.59 C +ATOM 223 CH2 TRP A 99 20.783 1.154 6.552 1.00 27.59 C +ATOM 224 N ARG A 100 25.004 -6.281 6.525 1.00 27.53 N +ATOM 225 CA ARG A 100 26.305 -6.921 6.653 1.00 27.53 C +ATOM 226 C ARG A 100 27.241 -6.311 5.627 1.00 27.53 C +ATOM 227 O ARG A 100 26.941 -6.290 4.433 1.00 27.53 O +ATOM 228 CB ARG A 100 26.128 -8.424 6.492 1.00 27.53 C +ATOM 229 CG ARG A 100 27.424 -9.189 6.771 1.00 27.53 C +ATOM 230 CD ARG A 100 27.096 -10.661 7.026 1.00 27.53 C +ATOM 231 NE ARG A 100 26.527 -10.794 8.371 1.00 27.53 N +ATOM 232 CZ ARG A 100 25.793 -11.768 8.849 1.00 27.53 C +ATOM 233 NH1 ARG A 100 25.457 -12.807 8.140 1.00 27.53 N +ATOM 234 NH2 ARG A 100 25.358 -11.676 10.074 1.00 27.53 N +ATOM 235 N ILE A 101 28.371 -5.809 6.100 1.00 27.47 N +ATOM 236 CA ILE A 101 29.395 -5.171 5.273 1.00 27.47 C +ATOM 237 C ILE A 101 30.746 -5.832 5.507 1.00 27.47 C +ATOM 238 O ILE A 101 30.992 -6.407 6.569 1.00 27.47 O +ATOM 239 CB ILE A 101 29.443 -3.645 5.501 1.00 27.47 C +ATOM 240 CG1 ILE A 101 29.886 -3.270 6.933 1.00 27.47 C +ATOM 241 CG2 ILE A 101 28.087 -3.012 5.129 1.00 27.47 C +ATOM 242 CD1 ILE A 101 30.003 -1.758 7.152 1.00 27.47 C +ATOM 243 N GLN A 102 31.627 -5.726 4.520 1.00 27.47 N +ATOM 244 CA GLN A 102 32.981 -6.260 4.587 1.00 27.47 C +ATOM 245 C GLN A 102 34.004 -5.177 4.251 1.00 27.47 C +ATOM 246 O GLN A 102 33.794 -4.381 3.332 1.00 27.47 O +ATOM 247 CB GLN A 102 33.089 -7.458 3.636 1.00 27.47 C +ATOM 248 CG GLN A 102 34.405 -8.225 3.814 1.00 27.47 C +ATOM 249 CD GLN A 102 34.567 -9.385 2.841 1.00 27.47 C +ATOM 250 OE1 GLN A 102 33.974 -9.450 1.778 1.00 27.47 O +ATOM 251 NE2 GLN A 102 35.416 -10.333 3.170 1.00 27.47 N +ATOM 252 N ASN A 103 35.141 -5.177 4.949 1.00 27.62 N +ATOM 253 CA ASN A 103 36.320 -4.463 4.472 1.00 27.62 C +ATOM 254 C ASN A 103 36.949 -5.238 3.304 1.00 27.62 C +ATOM 255 O ASN A 103 37.710 -6.184 3.496 1.00 27.62 O +ATOM 256 CB ASN A 103 37.299 -4.223 5.627 1.00 27.62 C +ATOM 257 CG ASN A 103 38.545 -3.469 5.188 1.00 27.62 C +ATOM 258 OD1 ASN A 103 38.676 -3.032 4.052 1.00 27.62 O +ATOM 259 ND2 ASN A 103 39.480 -3.251 6.075 1.00 27.62 N +ATOM 260 N SER A 104 36.622 -4.824 2.083 1.00 28.00 N +ATOM 261 CA SER A 104 37.174 -5.380 0.842 1.00 28.00 C +ATOM 262 C SER A 104 38.548 -4.803 0.473 1.00 28.00 C +ATOM 263 O SER A 104 39.088 -5.127 -0.584 1.00 28.00 O +ATOM 264 CB SER A 104 36.160 -5.196 -0.289 1.00 28.00 C +ATOM 265 OG SER A 104 35.740 -3.844 -0.391 1.00 28.00 O +ATOM 266 N GLY A 105 39.119 -3.947 1.327 1.00 29.69 N +ATOM 267 CA GLY A 105 40.433 -3.341 1.149 1.00 29.69 C +ATOM 268 C GLY A 105 41.587 -4.241 1.592 1.00 29.69 C +ATOM 269 O GLY A 105 41.400 -5.298 2.194 1.00 29.69 O +ATOM 270 N ALA A 106 42.805 -3.790 1.288 1.00 30.78 N +ATOM 271 CA ALA A 106 44.046 -4.456 1.691 1.00 30.78 C +ATOM 272 C ALA A 106 44.539 -4.039 3.089 1.00 30.78 C +ATOM 273 O ALA A 106 45.422 -4.688 3.643 1.00 30.78 O +ATOM 274 CB ALA A 106 45.104 -4.165 0.620 1.00 30.78 C +ATOM 275 N GLU A 107 43.976 -2.970 3.653 1.00 30.51 N +ATOM 276 CA GLU A 107 44.381 -2.376 4.929 1.00 30.51 C +ATOM 277 C GLU A 107 43.173 -2.227 5.860 1.00 30.51 C +ATOM 278 O GLU A 107 42.025 -2.173 5.407 1.00 30.51 O +ATOM 279 CB GLU A 107 45.047 -1.014 4.691 1.00 30.51 C +ATOM 280 CG GLU A 107 46.362 -1.132 3.901 1.00 30.51 C +ATOM 281 CD GLU A 107 47.047 0.221 3.657 1.00 30.51 C +ATOM 282 OE1 GLU A 107 48.167 0.204 3.097 1.00 30.51 O +ATOM 283 OE2 GLU A 107 46.429 1.259 3.977 1.00 30.51 O +ATOM 284 N ALA A 108 43.431 -2.168 7.167 1.00 28.98 N +ATOM 285 CA ALA A 108 42.396 -1.903 8.154 1.00 28.98 C +ATOM 286 C ALA A 108 41.783 -0.511 7.939 1.00 28.98 C +ATOM 287 O ALA A 108 42.471 0.437 7.552 1.00 28.98 O +ATOM 288 CB ALA A 108 42.984 -2.050 9.560 1.00 28.98 C +ATOM 289 N TRP A 109 40.485 -0.358 8.206 1.00 29.13 N +ATOM 290 CA TRP A 109 39.857 0.960 8.120 1.00 29.13 C +ATOM 291 C TRP A 109 40.442 1.928 9.154 1.00 29.13 C +ATOM 292 O TRP A 109 40.658 1.539 10.305 1.00 29.13 O +ATOM 293 CB TRP A 109 38.343 0.859 8.283 1.00 29.13 C +ATOM 294 CG TRP A 109 37.592 0.266 7.137 1.00 29.13 C +ATOM 295 CD1 TRP A 109 38.084 -0.040 5.913 1.00 29.13 C +ATOM 296 CD2 TRP A 109 36.201 -0.165 7.125 1.00 29.13 C +ATOM 297 NE1 TRP A 109 37.092 -0.604 5.141 1.00 29.13 N +ATOM 298 CE2 TRP A 109 35.921 -0.727 5.847 1.00 29.13 C +ATOM 299 CE3 TRP A 109 35.162 -0.181 8.081 1.00 29.13 C +ATOM 300 CZ2 TRP A 109 34.685 -1.295 5.534 1.00 29.13 C +ATOM 301 CZ3 TRP A 109 33.904 -0.727 7.764 1.00 29.13 C +ATOM 302 CH2 TRP A 109 33.667 -1.285 6.496 1.00 29.13 C +ATOM 303 N PRO A 110 40.631 3.211 8.797 1.00 32.24 N +ATOM 304 CA PRO A 110 41.151 4.192 9.737 1.00 32.24 C +ATOM 305 C PRO A 110 40.203 4.362 10.938 1.00 32.24 C +ATOM 306 O PRO A 110 39.006 4.068 10.845 1.00 32.24 O +ATOM 307 CB PRO A 110 41.325 5.481 8.930 1.00 32.24 C +ATOM 308 CG PRO A 110 40.299 5.345 7.806 1.00 32.24 C +ATOM 309 CD PRO A 110 40.290 3.843 7.529 1.00 32.24 C +ATOM 310 N PRO A 111 40.704 4.837 12.088 1.00 31.71 N +ATOM 311 CA PRO A 111 39.854 5.135 13.233 1.00 31.71 C +ATOM 312 C PRO A 111 38.844 6.244 12.904 1.00 31.71 C +ATOM 313 O PRO A 111 39.123 7.165 12.135 1.00 31.71 O +ATOM 314 CB PRO A 111 40.812 5.518 14.361 1.00 31.71 C +ATOM 315 CG PRO A 111 42.039 6.053 13.620 1.00 31.71 C +ATOM 316 CD PRO A 111 42.086 5.193 12.359 1.00 31.71 C +ATOM 317 N GLY A 112 37.665 6.173 13.524 1.00 31.76 N +ATOM 318 CA GLY A 112 36.633 7.207 13.391 1.00 31.76 C +ATOM 319 C GLY A 112 35.790 7.122 12.114 1.00 31.76 C +ATOM 320 O GLY A 112 35.154 8.107 11.743 1.00 31.76 O +ATOM 321 N VAL A 113 35.771 5.970 11.440 1.00 29.24 N +ATOM 322 CA VAL A 113 34.854 5.725 10.319 1.00 29.24 C +ATOM 323 C VAL A 113 33.401 5.618 10.787 1.00 29.24 C +ATOM 324 O VAL A 113 33.093 5.194 11.903 1.00 29.24 O +ATOM 325 CB VAL A 113 35.256 4.500 9.477 1.00 29.24 C +ATOM 326 CG1 VAL A 113 36.559 4.783 8.726 1.00 29.24 C +ATOM 327 CG2 VAL A 113 35.393 3.215 10.298 1.00 29.24 C +ATOM 328 N CYS A 114 32.483 6.017 9.915 1.00 27.97 N +ATOM 329 CA CYS A 114 31.052 6.006 10.183 1.00 27.97 C +ATOM 330 C CYS A 114 30.259 5.683 8.916 1.00 27.97 C +ATOM 331 O CYS A 114 30.725 5.919 7.804 1.00 27.97 O +ATOM 332 CB CYS A 114 30.647 7.353 10.807 1.00 27.97 C +ATOM 333 SG CYS A 114 31.055 8.747 9.715 1.00 27.97 S +ATOM 334 N LEU A 115 29.046 5.169 9.085 1.00 27.56 N +ATOM 335 CA LEU A 115 28.061 5.053 8.016 1.00 27.56 C +ATOM 336 C LEU A 115 27.163 6.293 8.063 1.00 27.56 C +ATOM 337 O LEU A 115 26.636 6.628 9.125 1.00 27.56 O +ATOM 338 CB LEU A 115 27.284 3.741 8.205 1.00 27.56 C +ATOM 339 CG LEU A 115 26.287 3.441 7.075 1.00 27.56 C +ATOM 340 CD1 LEU A 115 26.996 2.949 5.817 1.00 27.56 C +ATOM 341 CD2 LEU A 115 25.265 2.400 7.527 1.00 27.56 C +ATOM 342 N LYS A 116 26.992 6.978 6.930 1.00 28.00 N +ATOM 343 CA LYS A 116 26.211 8.216 6.818 1.00 28.00 C +ATOM 344 C LYS A 116 25.102 8.102 5.791 1.00 28.00 C +ATOM 345 O LYS A 116 25.282 7.492 4.741 1.00 28.00 O +ATOM 346 CB LYS A 116 27.125 9.394 6.464 1.00 28.00 C +ATOM 347 CG LYS A 116 27.885 9.912 7.685 1.00 28.00 C +ATOM 348 CD LYS A 116 28.687 11.161 7.312 1.00 28.00 C +ATOM 349 CE LYS A 116 29.409 11.706 8.546 1.00 28.00 C +ATOM 350 NZ LYS A 116 30.404 12.737 8.170 1.00 28.00 N +ATOM 351 N TYR A 117 23.980 8.741 6.085 1.00 27.74 N +ATOM 352 CA TYR A 117 22.911 8.966 5.122 1.00 27.74 C +ATOM 353 C TYR A 117 23.386 9.952 4.047 1.00 27.74 C +ATOM 354 O TYR A 117 23.942 11.002 4.381 1.00 27.74 O +ATOM 355 CB TYR A 117 21.689 9.499 5.872 1.00 27.74 C +ATOM 356 CG TYR A 117 20.620 10.027 4.948 1.00 27.74 C +ATOM 357 CD1 TYR A 117 20.561 11.408 4.673 1.00 27.74 C +ATOM 358 CD2 TYR A 117 19.743 9.132 4.309 1.00 27.74 C +ATOM 359 CE1 TYR A 117 19.612 11.901 3.764 1.00 27.74 C +ATOM 360 CE2 TYR A 117 18.779 9.626 3.415 1.00 27.74 C +ATOM 361 CZ TYR A 117 18.722 11.005 3.141 1.00 27.74 C +ATOM 362 OH TYR A 117 17.800 11.476 2.279 1.00 27.74 O +ATOM 363 N VAL A 118 23.161 9.628 2.771 1.00 28.20 N +ATOM 364 CA VAL A 118 23.642 10.435 1.631 1.00 28.20 C +ATOM 365 C VAL A 118 22.538 10.927 0.693 1.00 28.20 C +ATOM 366 O VAL A 118 22.778 11.872 -0.056 1.00 28.20 O +ATOM 367 CB VAL A 118 24.752 9.713 0.841 1.00 28.20 C +ATOM 368 CG1 VAL A 118 25.987 9.467 1.716 1.00 28.20 C +ATOM 369 CG2 VAL A 118 24.309 8.377 0.236 1.00 28.20 C +ATOM 370 N GLY A 119 21.331 10.355 0.734 1.00 28.58 N +ATOM 371 CA GLY A 119 20.226 10.796 -0.122 1.00 28.58 C +ATOM 372 C GLY A 119 18.998 9.885 -0.097 1.00 28.58 C +ATOM 373 O GLY A 119 19.038 8.794 0.466 1.00 28.58 O +ATOM 374 N GLY A 120 17.912 10.332 -0.732 1.00 28.09 N +ATOM 375 CA GLY A 120 16.630 9.621 -0.789 1.00 28.09 C +ATOM 376 C GLY A 120 15.628 10.052 0.287 1.00 28.09 C +ATOM 377 O GLY A 120 15.576 11.215 0.684 1.00 28.09 O +ATOM 378 N ASP A 121 14.804 9.106 0.720 1.00 28.06 N +ATOM 379 CA ASP A 121 13.859 9.246 1.816 1.00 28.06 C +ATOM 380 C ASP A 121 14.557 8.986 3.157 1.00 28.06 C +ATOM 381 O ASP A 121 15.352 8.054 3.306 1.00 28.06 O +ATOM 382 CB ASP A 121 12.683 8.283 1.617 1.00 28.06 C +ATOM 383 CG ASP A 121 11.942 8.547 0.308 1.00 28.06 C +ATOM 384 OD1 ASP A 121 11.344 9.642 0.197 1.00 28.06 O +ATOM 385 OD2 ASP A 121 11.927 7.624 -0.537 1.00 28.06 O +ATOM 386 N GLN A 122 14.246 9.811 4.157 1.00 29.65 N +ATOM 387 CA GLN A 122 14.744 9.634 5.519 1.00 29.65 C +ATOM 388 C GLN A 122 13.735 8.879 6.377 1.00 29.65 C +ATOM 389 O GLN A 122 12.553 9.215 6.421 1.00 29.65 O +ATOM 390 CB GLN A 122 15.103 10.979 6.162 1.00 29.65 C +ATOM 391 CG GLN A 122 16.380 11.561 5.550 1.00 29.65 C +ATOM 392 CD GLN A 122 16.849 12.828 6.254 1.00 29.65 C +ATOM 393 OE1 GLN A 122 16.097 13.597 6.828 1.00 29.65 O +ATOM 394 NE2 GLN A 122 18.132 13.109 6.227 1.00 29.65 N +ATOM 395 N PHE A 123 14.239 7.909 7.134 1.00 28.83 N +ATOM 396 CA PHE A 123 13.453 7.049 8.015 1.00 28.83 C +ATOM 397 C PHE A 123 13.706 7.372 9.492 1.00 28.83 C +ATOM 398 O PHE A 123 13.829 6.477 10.316 1.00 28.83 O +ATOM 399 CB PHE A 123 13.734 5.585 7.651 1.00 28.83 C +ATOM 400 CG PHE A 123 13.483 5.263 6.196 1.00 28.83 C +ATOM 401 CD1 PHE A 123 12.167 5.305 5.719 1.00 28.83 C +ATOM 402 CD2 PHE A 123 14.538 4.983 5.313 1.00 28.83 C +ATOM 403 CE1 PHE A 123 11.887 5.052 4.368 1.00 28.83 C +ATOM 404 CE2 PHE A 123 14.262 4.743 3.957 1.00 28.83 C +ATOM 405 CZ PHE A 123 12.941 4.757 3.486 1.00 28.83 C +ATOM 406 N VAL A 126 18.376 9.830 12.765 1.00 28.83 N +ATOM 407 CA VAL A 126 19.736 9.291 12.705 1.00 28.83 C +ATOM 408 C VAL A 126 20.239 9.419 11.272 1.00 28.83 C +ATOM 409 O VAL A 126 19.663 8.841 10.358 1.00 28.83 O +ATOM 410 CB VAL A 126 19.802 7.829 13.196 1.00 28.83 C +ATOM 411 CG1 VAL A 126 21.247 7.313 13.179 1.00 28.83 C +ATOM 412 CG2 VAL A 126 19.276 7.706 14.633 1.00 28.83 C +ATOM 413 N ASN A 127 21.319 10.174 11.079 1.00 28.72 N +ATOM 414 CA ASN A 127 21.975 10.373 9.783 1.00 28.72 C +ATOM 415 C ASN A 127 23.429 9.870 9.766 1.00 28.72 C +ATOM 416 O ASN A 127 24.089 9.936 8.730 1.00 28.72 O +ATOM 417 CB ASN A 127 21.874 11.863 9.414 1.00 28.72 C +ATOM 418 CG ASN A 127 22.624 12.789 10.361 1.00 28.72 C +ATOM 419 OD1 ASN A 127 23.250 12.389 11.331 1.00 28.72 O +ATOM 420 ND2 ASN A 127 22.576 14.075 10.117 1.00 28.72 N +ATOM 421 N MET A 128 23.932 9.391 10.907 1.00 28.06 N +ATOM 422 CA MET A 128 25.293 8.908 11.087 1.00 28.06 C +ATOM 423 C MET A 128 25.335 7.859 12.201 1.00 28.06 C +ATOM 424 O MET A 128 24.785 8.084 13.279 1.00 28.06 O +ATOM 425 CB MET A 128 26.212 10.089 11.436 1.00 28.06 C +ATOM 426 CG MET A 128 27.689 9.677 11.498 1.00 28.06 C +ATOM 427 SD MET A 128 28.800 11.002 12.038 1.00 28.06 S +ATOM 428 CE MET A 128 28.587 10.850 13.835 1.00 28.06 C +ATOM 429 N VAL A 129 26.039 6.752 11.970 1.00 27.77 N +ATOM 430 CA VAL A 129 26.379 5.762 13.001 1.00 27.77 C +ATOM 431 C VAL A 129 27.867 5.437 12.950 1.00 27.77 C +ATOM 432 O VAL A 129 28.434 5.252 11.874 1.00 27.77 O +ATOM 433 CB VAL A 129 25.515 4.488 12.917 1.00 27.77 C +ATOM 434 CG1 VAL A 129 24.047 4.809 13.218 1.00 27.77 C +ATOM 435 CG2 VAL A 129 25.592 3.766 11.569 1.00 27.77 C +ATOM 436 N MET A 130 28.519 5.397 14.113 1.00 27.93 N +ATOM 437 CA MET A 130 29.929 5.015 14.198 1.00 27.93 C +ATOM 438 C MET A 130 30.089 3.532 13.876 1.00 27.93 C +ATOM 439 O MET A 130 29.303 2.705 14.335 1.00 27.93 O +ATOM 440 CB MET A 130 30.512 5.323 15.585 1.00 27.93 C +ATOM 441 CG MET A 130 30.592 6.826 15.874 1.00 27.93 C +ATOM 442 SD MET A 130 31.508 7.824 14.665 1.00 27.93 S +ATOM 443 CE MET A 130 33.162 7.095 14.785 1.00 27.93 C +ATOM 444 N VAL A 131 31.132 3.204 13.119 1.00 27.77 N +ATOM 445 CA VAL A 131 31.474 1.825 12.765 1.00 27.77 C +ATOM 446 C VAL A 131 32.889 1.554 13.262 1.00 27.77 C +ATOM 447 O VAL A 131 33.743 2.440 13.256 1.00 27.77 O +ATOM 448 CB VAL A 131 31.316 1.582 11.254 1.00 27.77 C +ATOM 449 CG1 VAL A 131 31.556 0.114 10.883 1.00 27.77 C +ATOM 450 CG2 VAL A 131 29.901 1.938 10.772 1.00 27.77 C +ATOM 451 N ARG A 132 33.142 0.337 13.745 1.00 28.30 N +ATOM 452 CA ARG A 132 34.491 -0.049 14.170 1.00 28.30 C +ATOM 453 C ARG A 132 35.443 -0.098 12.969 1.00 28.30 C +ATOM 454 O ARG A 132 35.020 -0.366 11.847 1.00 28.30 O +ATOM 455 CB ARG A 132 34.463 -1.356 14.980 1.00 28.30 C +ATOM 456 CG ARG A 132 34.150 -2.609 14.147 1.00 28.30 C +ATOM 457 CD ARG A 132 34.194 -3.863 15.031 1.00 28.30 C +ATOM 458 NE ARG A 132 34.073 -5.111 14.252 1.00 28.30 N +ATOM 459 CZ ARG A 132 33.114 -6.015 14.275 1.00 28.30 C +ATOM 460 NH1 ARG A 132 32.055 -5.908 15.032 1.00 28.30 N +ATOM 461 NH2 ARG A 132 33.236 -7.059 13.517 1.00 28.30 N +ATOM 462 N SER A 133 36.726 0.117 13.223 1.00 28.33 N +ATOM 463 CA SER A 133 37.782 -0.118 12.238 1.00 28.33 C +ATOM 464 C SER A 133 37.900 -1.615 11.967 1.00 28.33 C +ATOM 465 O SER A 133 38.291 -2.359 12.861 1.00 28.33 O +ATOM 466 CB SER A 133 39.104 0.424 12.777 1.00 28.33 C +ATOM 467 OG SER A 133 39.020 1.834 12.847 1.00 28.33 O +ATOM 468 N LEU A 134 37.508 -2.043 10.769 1.00 27.71 N +ATOM 469 CA LEU A 134 37.559 -3.442 10.353 1.00 27.71 C +ATOM 470 C LEU A 134 38.918 -3.782 9.750 1.00 27.71 C +ATOM 471 O LEU A 134 39.428 -3.019 8.926 1.00 27.71 O +ATOM 472 CB LEU A 134 36.436 -3.724 9.338 1.00 27.71 C +ATOM 473 CG LEU A 134 35.000 -3.565 9.855 1.00 27.71 C +ATOM 474 CD1 LEU A 134 34.012 -4.013 8.774 1.00 27.71 C +ATOM 475 CD2 LEU A 134 34.704 -4.384 11.109 1.00 27.71 C +ATOM 476 N GLU A 135 39.459 -4.942 10.103 1.00 27.62 N +ATOM 477 CA GLU A 135 40.640 -5.521 9.459 1.00 27.62 C +ATOM 478 C GLU A 135 40.329 -5.957 8.013 1.00 27.62 C +ATOM 479 O GLU A 135 39.160 -6.155 7.665 1.00 27.62 O +ATOM 480 CB GLU A 135 41.144 -6.715 10.288 1.00 27.62 C +ATOM 481 CG GLU A 135 41.747 -6.303 11.641 1.00 27.62 C +ATOM 482 CD GLU A 135 42.983 -5.396 11.516 1.00 27.62 C +ATOM 483 OE1 GLU A 135 43.247 -4.639 12.478 1.00 27.62 O +ATOM 484 OE2 GLU A 135 43.671 -5.474 10.472 1.00 27.62 O +ATOM 485 N PRO A 136 41.340 -6.122 7.138 1.00 27.93 N +ATOM 486 CA PRO A 136 41.131 -6.642 5.789 1.00 27.93 C +ATOM 487 C PRO A 136 40.327 -7.948 5.797 1.00 27.93 C +ATOM 488 O PRO A 136 40.638 -8.875 6.544 1.00 27.93 O +ATOM 489 CB PRO A 136 42.531 -6.857 5.207 1.00 27.93 C +ATOM 490 CG PRO A 136 43.366 -5.827 5.960 1.00 27.93 C +ATOM 491 CD PRO A 136 42.749 -5.829 7.354 1.00 27.93 C +ATOM 492 N GLN A 137 39.309 -8.029 4.938 1.00 27.74 N +ATOM 493 CA GLN A 137 38.361 -9.145 4.821 1.00 27.74 C +ATOM 494 C GLN A 137 37.431 -9.369 6.025 1.00 27.74 C +ATOM 495 O GLN A 137 36.572 -10.251 5.951 1.00 27.74 O +ATOM 496 CB GLN A 137 39.076 -10.450 4.411 1.00 27.74 C +ATOM 497 CG GLN A 137 39.935 -10.313 3.147 1.00 27.74 C +ATOM 498 CD GLN A 137 39.087 -9.943 1.939 1.00 27.74 C +ATOM 499 OE1 GLN A 137 38.170 -10.656 1.563 1.00 27.74 O +ATOM 500 NE2 GLN A 137 39.346 -8.829 1.290 1.00 27.74 N +ATOM 501 N GLU A 138 37.536 -8.583 7.100 1.00 27.53 N +ATOM 502 CA GLU A 138 36.629 -8.680 8.242 1.00 27.53 C +ATOM 503 C GLU A 138 35.209 -8.234 7.858 1.00 27.53 C +ATOM 504 O GLU A 138 35.012 -7.282 7.096 1.00 27.53 O +ATOM 505 CB GLU A 138 37.174 -7.882 9.438 1.00 27.53 C +ATOM 506 CG GLU A 138 36.445 -8.202 10.757 1.00 27.53 C +ATOM 507 CD GLU A 138 36.756 -7.214 11.892 1.00 27.53 C +ATOM 508 OE1 GLU A 138 35.890 -7.099 12.801 1.00 27.53 O +ATOM 509 OE2 GLU A 138 37.800 -6.535 11.836 1.00 27.53 O +ATOM 510 N ILE A 139 34.217 -8.933 8.413 1.00 27.62 N +ATOM 511 CA ILE A 139 32.788 -8.677 8.231 1.00 27.62 C +ATOM 512 C ILE A 139 32.198 -8.091 9.519 1.00 27.62 C +ATOM 513 O ILE A 139 32.517 -8.538 10.626 1.00 27.62 O +ATOM 514 CB ILE A 139 32.066 -9.969 7.783 1.00 27.62 C +ATOM 515 CG1 ILE A 139 32.581 -10.418 6.395 1.00 27.62 C +ATOM 516 CG2 ILE A 139 30.543 -9.774 7.741 1.00 27.62 C +ATOM 517 CD1 ILE A 139 32.001 -11.746 5.891 1.00 27.62 C +ATOM 518 N ALA A 140 31.290 -7.128 9.379 1.00 27.71 N +ATOM 519 CA ALA A 140 30.510 -6.598 10.489 1.00 27.71 C +ATOM 520 C ALA A 140 29.062 -6.303 10.101 1.00 27.71 C +ATOM 521 O ALA A 140 28.768 -5.914 8.971 1.00 27.71 O +ATOM 522 CB ALA A 140 31.197 -5.348 11.035 1.00 27.71 C +ATOM 523 N ASP A 141 28.181 -6.431 11.089 1.00 27.56 N +ATOM 524 CA ASP A 141 26.790 -6.010 10.993 1.00 27.56 C +ATOM 525 C ASP A 141 26.682 -4.567 11.505 1.00 27.56 C +ATOM 526 O ASP A 141 27.117 -4.254 12.618 1.00 27.56 O +ATOM 527 CB ASP A 141 25.887 -6.992 11.763 1.00 27.56 C +ATOM 528 CG ASP A 141 25.845 -8.403 11.149 1.00 27.56 C +ATOM 529 OD1 ASP A 141 26.363 -8.615 10.032 1.00 27.56 O +ATOM 530 OD2 ASP A 141 25.316 -9.337 11.791 1.00 27.56 O +ATOM 531 N VAL A 142 26.121 -3.680 10.684 1.00 27.71 N +ATOM 532 CA VAL A 142 25.889 -2.270 11.016 1.00 27.71 C +ATOM 533 C VAL A 142 24.393 -2.018 11.040 1.00 27.71 C +ATOM 534 O VAL A 142 23.709 -2.232 10.042 1.00 27.71 O +ATOM 535 CB VAL A 142 26.606 -1.316 10.045 1.00 27.71 C +ATOM 536 CG1 VAL A 142 26.386 0.152 10.446 1.00 27.71 C +ATOM 537 CG2 VAL A 142 28.116 -1.574 10.060 1.00 27.71 C +ATOM 538 N SER A 143 23.899 -1.555 12.185 1.00 27.93 N +ATOM 539 CA SER A 143 22.474 -1.355 12.442 1.00 27.93 C +ATOM 540 C SER A 143 22.132 0.118 12.610 1.00 27.93 C +ATOM 541 O SER A 143 22.856 0.862 13.274 1.00 27.93 O +ATOM 542 CB SER A 143 22.035 -2.128 13.685 1.00 27.93 C +ATOM 543 OG SER A 143 22.235 -3.514 13.500 1.00 27.93 O +ATOM 544 N VAL A 144 20.989 0.526 12.068 1.00 27.90 N +ATOM 545 CA VAL A 144 20.421 1.864 12.243 1.00 27.90 C +ATOM 546 C VAL A 144 18.960 1.734 12.671 1.00 27.90 C +ATOM 547 O VAL A 144 18.167 1.045 12.028 1.00 27.90 O +ATOM 548 CB VAL A 144 20.567 2.702 10.960 1.00 27.90 C +ATOM 549 CG1 VAL A 144 20.002 4.115 11.154 1.00 27.90 C +ATOM 550 CG2 VAL A 144 22.030 2.866 10.521 1.00 27.90 C +ATOM 551 N GLN A 145 18.609 2.407 13.769 1.00 28.58 N +ATOM 552 CA GLN A 145 17.225 2.538 14.222 1.00 28.58 C +ATOM 553 C GLN A 145 16.486 3.538 13.332 1.00 28.58 C +ATOM 554 O GLN A 145 16.947 4.664 13.134 1.00 28.58 O +ATOM 555 CB GLN A 145 17.189 2.982 15.692 1.00 28.58 C +ATOM 556 CG GLN A 145 15.762 3.049 16.269 1.00 28.58 C +ATOM 557 CD GLN A 145 15.109 1.677 16.371 1.00 28.58 C +ATOM 558 OE1 GLN A 145 15.652 0.757 16.952 1.00 28.58 O +ATOM 559 NE2 GLN A 145 13.943 1.475 15.800 1.00 28.58 N +ATOM 560 N MET A 146 15.329 3.127 12.830 1.00 28.20 N +ATOM 561 CA MET A 146 14.532 3.875 11.870 1.00 28.20 C +ATOM 562 C MET A 146 13.034 3.833 12.221 1.00 28.20 C +ATOM 563 O MET A 146 12.608 3.103 13.120 1.00 28.20 O +ATOM 564 CB MET A 146 14.799 3.291 10.478 1.00 28.20 C +ATOM 565 CG MET A 146 16.231 3.512 9.978 1.00 28.20 C +ATOM 566 SD MET A 146 16.476 3.058 8.240 1.00 28.20 S +ATOM 567 CE MET A 146 18.257 3.336 8.121 1.00 28.20 C +ATOM 568 N CYS A 147 12.244 4.646 11.522 1.00 28.44 N +ATOM 569 CA CYS A 147 10.795 4.759 11.592 1.00 28.44 C +ATOM 570 C CYS A 147 10.223 4.886 10.173 1.00 28.44 C +ATOM 571 O CYS A 147 10.719 5.683 9.369 1.00 28.44 O +ATOM 572 CB CYS A 147 10.423 5.980 12.442 1.00 28.44 C +ATOM 573 SG CYS A 147 8.645 5.926 12.798 1.00 28.44 S +ATOM 574 N SER A 148 9.199 4.095 9.851 1.00 28.76 N +ATOM 575 CA SER A 148 8.544 4.144 8.545 1.00 28.76 C +ATOM 576 C SER A 148 7.710 5.428 8.383 1.00 28.76 C +ATOM 577 O SER A 148 7.158 5.937 9.360 1.00 28.76 O +ATOM 578 CB SER A 148 7.675 2.905 8.321 1.00 28.76 C +ATOM 579 OG SER A 148 6.758 2.769 9.377 1.00 28.76 O +ATOM 580 N PRO A 149 7.589 5.972 7.158 1.00 28.37 N +ATOM 581 CA PRO A 149 6.727 7.114 6.874 1.00 28.37 C +ATOM 582 C PRO A 149 5.262 6.874 7.255 1.00 28.37 C +ATOM 583 O PRO A 149 4.771 5.749 7.236 1.00 28.37 O +ATOM 584 CB PRO A 149 6.875 7.371 5.370 1.00 28.37 C +ATOM 585 CG PRO A 149 8.283 6.867 5.072 1.00 28.37 C +ATOM 586 CD PRO A 149 8.409 5.667 6.001 1.00 28.37 C +ATOM 587 N SER A 150 4.530 7.954 7.528 1.00 28.58 N +ATOM 588 CA SER A 150 3.101 7.906 7.881 1.00 28.58 C +ATOM 589 C SER A 150 2.172 7.611 6.696 1.00 28.58 C +ATOM 590 O SER A 150 1.003 7.264 6.873 1.00 28.58 O +ATOM 591 CB SER A 150 2.695 9.234 8.519 1.00 28.58 C +ATOM 592 OG SER A 150 2.979 10.339 7.664 1.00 28.58 O +ATOM 593 N ARG A 151 2.665 7.771 5.463 1.00 28.44 N +ATOM 594 CA ARG A 151 1.896 7.541 4.236 1.00 28.44 C +ATOM 595 C ARG A 151 2.333 6.242 3.580 1.00 28.44 C +ATOM 596 O ARG A 151 3.518 5.937 3.531 1.00 28.44 O +ATOM 597 CB ARG A 151 2.025 8.728 3.275 1.00 28.44 C +ATOM 598 CG ARG A 151 1.420 10.005 3.871 1.00 28.44 C +ATOM 599 CD ARG A 151 1.453 11.128 2.834 1.00 28.44 C +ATOM 600 NE ARG A 151 0.927 12.382 3.400 1.00 28.44 N +ATOM 601 CZ ARG A 151 0.465 13.417 2.722 1.00 28.44 C +ATOM 602 NH1 ARG A 151 0.414 13.420 1.418 1.00 28.44 N +ATOM 603 NH2 ARG A 151 0.047 14.479 3.350 1.00 28.44 N +ATOM 604 N ALA A 152 1.359 5.508 3.054 1.00 28.83 N +ATOM 605 CA ALA A 152 1.630 4.312 2.274 1.00 28.83 C +ATOM 606 C ALA A 152 2.339 4.661 0.954 1.00 28.83 C +ATOM 607 O ALA A 152 2.025 5.680 0.331 1.00 28.83 O +ATOM 608 CB ALA A 152 0.318 3.552 2.046 1.00 28.83 C +ATOM 609 N GLY A 153 3.264 3.805 0.524 1.00 28.65 N +ATOM 610 CA GLY A 153 4.062 4.016 -0.682 1.00 28.65 C +ATOM 611 C GLY A 153 5.399 3.279 -0.663 1.00 28.65 C +ATOM 612 O GLY A 153 5.753 2.640 0.327 1.00 28.65 O +ATOM 613 N MET A 154 6.137 3.389 -1.768 1.00 28.54 N +ATOM 614 CA MET A 154 7.527 2.934 -1.859 1.00 28.54 C +ATOM 615 C MET A 154 8.466 4.060 -1.463 1.00 28.54 C +ATOM 616 O MET A 154 8.323 5.175 -1.963 1.00 28.54 O +ATOM 617 CB MET A 154 7.881 2.476 -3.280 1.00 28.54 C +ATOM 618 CG MET A 154 7.188 1.170 -3.640 1.00 28.54 C +ATOM 619 SD MET A 154 7.703 -0.238 -2.631 1.00 28.54 S +ATOM 620 CE MET A 154 6.063 -0.963 -2.478 1.00 28.54 C +ATOM 621 N TYR A 155 9.436 3.741 -0.615 1.00 28.23 N +ATOM 622 CA TYR A 155 10.436 4.686 -0.136 1.00 28.23 C +ATOM 623 C TYR A 155 11.827 4.074 -0.224 1.00 28.23 C +ATOM 624 O TYR A 155 11.995 2.877 0.024 1.00 28.23 O +ATOM 625 CB TYR A 155 10.116 5.086 1.302 1.00 28.23 C +ATOM 626 CG TYR A 155 8.798 5.808 1.462 1.00 28.23 C +ATOM 627 CD1 TYR A 155 8.762 7.209 1.353 1.00 28.23 C +ATOM 628 CD2 TYR A 155 7.611 5.091 1.699 1.00 28.23 C +ATOM 629 CE1 TYR A 155 7.551 7.901 1.493 1.00 28.23 C +ATOM 630 CE2 TYR A 155 6.391 5.779 1.843 1.00 28.23 C +ATOM 631 CZ TYR A 155 6.362 7.185 1.736 1.00 28.23 C +ATOM 632 OH TYR A 155 5.201 7.878 1.840 1.00 28.23 O +ATOM 633 N GLN A 156 12.826 4.899 -0.536 1.00 28.09 N +ATOM 634 CA GLN A 156 14.209 4.448 -0.661 1.00 28.09 C +ATOM 635 C GLN A 156 15.206 5.435 -0.051 1.00 28.09 C +ATOM 636 O GLN A 156 15.204 6.613 -0.385 1.00 28.09 O +ATOM 637 CB GLN A 156 14.540 4.107 -2.123 1.00 28.09 C +ATOM 638 CG GLN A 156 14.425 5.304 -3.080 1.00 28.09 C +ATOM 639 CD GLN A 156 14.741 4.940 -4.521 1.00 28.09 C +ATOM 640 OE1 GLN A 156 14.499 3.851 -5.006 1.00 28.09 O +ATOM 641 NE2 GLN A 156 15.301 5.853 -5.283 1.00 28.09 N +ATOM 642 N GLY A 157 16.102 4.951 0.806 1.00 28.16 N +ATOM 643 CA GLY A 157 17.142 5.752 1.454 1.00 28.16 C +ATOM 644 C GLY A 157 18.532 5.193 1.175 1.00 28.16 C +ATOM 645 O GLY A 157 18.732 3.982 1.220 1.00 28.16 O +ATOM 646 N GLN A 158 19.495 6.069 0.906 1.00 27.77 N +ATOM 647 CA GLN A 158 20.874 5.724 0.568 1.00 27.77 C +ATOM 648 C GLN A 158 21.821 6.028 1.723 1.00 27.77 C +ATOM 649 O GLN A 158 21.773 7.107 2.324 1.00 27.77 O +ATOM 650 CB GLN A 158 21.333 6.520 -0.655 1.00 27.77 C +ATOM 651 CG GLN A 158 20.577 6.170 -1.937 1.00 27.77 C +ATOM 652 CD GLN A 158 20.934 7.112 -3.083 1.00 27.77 C +ATOM 653 OE1 GLN A 158 21.740 8.021 -2.983 1.00 27.77 O +ATOM 654 NE2 GLN A 158 20.331 6.955 -4.232 1.00 27.77 N +ATOM 655 N TRP A 159 22.741 5.102 1.971 1.00 27.65 N +ATOM 656 CA TRP A 159 23.777 5.210 2.988 1.00 27.65 C +ATOM 657 C TRP A 159 25.143 4.863 2.406 1.00 27.65 C +ATOM 658 O TRP A 159 25.261 4.040 1.497 1.00 27.65 O +ATOM 659 CB TRP A 159 23.429 4.308 4.173 1.00 27.65 C +ATOM 660 CG TRP A 159 22.212 4.724 4.935 1.00 27.65 C +ATOM 661 CD1 TRP A 159 20.931 4.534 4.546 1.00 27.65 C +ATOM 662 CD2 TRP A 159 22.146 5.405 6.219 1.00 27.65 C +ATOM 663 NE1 TRP A 159 20.078 5.085 5.481 1.00 27.65 N +ATOM 664 CE2 TRP A 159 20.774 5.620 6.544 1.00 27.65 C +ATOM 665 CE3 TRP A 159 23.110 5.857 7.144 1.00 27.65 C +ATOM 666 CZ2 TRP A 159 20.384 6.246 7.736 1.00 27.65 C +ATOM 667 CZ3 TRP A 159 22.733 6.473 8.349 1.00 27.65 C +ATOM 668 CH2 TRP A 159 21.374 6.657 8.646 1.00 27.65 C +ATOM 669 N ARG A 160 26.192 5.497 2.928 1.00 27.71 N +ATOM 670 CA ARG A 160 27.567 5.264 2.486 1.00 27.71 C +ATOM 671 C ARG A 160 28.561 5.452 3.622 1.00 27.71 C +ATOM 672 O ARG A 160 28.360 6.292 4.501 1.00 27.71 O +ATOM 673 CB ARG A 160 27.866 6.192 1.303 1.00 27.71 C +ATOM 674 CG ARG A 160 29.084 5.699 0.523 1.00 27.71 C +ATOM 675 CD ARG A 160 29.239 6.488 -0.772 1.00 27.71 C +ATOM 676 NE ARG A 160 30.281 5.873 -1.599 1.00 27.71 N +ATOM 677 CZ ARG A 160 30.981 6.430 -2.561 1.00 27.71 C +ATOM 678 NH1 ARG A 160 30.847 7.686 -2.891 1.00 27.71 N +ATOM 679 NH2 ARG A 160 31.829 5.689 -3.207 1.00 27.71 N +ATOM 680 N MET A 161 29.639 4.672 3.602 1.00 27.71 N +ATOM 681 CA MET A 161 30.728 4.823 4.563 1.00 27.71 C +ATOM 682 C MET A 161 31.466 6.148 4.356 1.00 27.71 C +ATOM 683 O MET A 161 31.601 6.637 3.233 1.00 27.71 O +ATOM 684 CB MET A 161 31.700 3.643 4.481 1.00 27.71 C +ATOM 685 CG MET A 161 31.093 2.337 4.998 1.00 27.71 C +ATOM 686 SD MET A 161 30.615 2.329 6.749 1.00 27.71 S +ATOM 687 CE MET A 161 32.240 2.508 7.525 1.00 27.71 C +ATOM 688 N CYS A 162 31.943 6.725 5.452 1.00 30.02 N +ATOM 689 CA CYS A 162 32.618 8.011 5.516 1.00 30.02 C +ATOM 690 C CYS A 162 33.783 7.931 6.510 1.00 30.02 C +ATOM 691 O CYS A 162 33.624 7.444 7.633 1.00 30.02 O +ATOM 692 CB CYS A 162 31.578 9.061 5.929 1.00 30.02 C +ATOM 693 SG CYS A 162 32.312 10.719 6.057 1.00 30.02 S +ATOM 694 N THR A 163 34.952 8.432 6.117 1.00 31.56 N +ATOM 695 CA THR A 163 36.110 8.570 7.011 1.00 31.56 C +ATOM 696 C THR A 163 35.899 9.701 8.020 1.00 31.56 C +ATOM 697 O THR A 163 35.032 10.562 7.844 1.00 31.56 O +ATOM 698 CB THR A 163 37.406 8.829 6.224 1.00 31.56 C +ATOM 699 OG1 THR A 163 37.321 10.107 5.642 1.00 31.56 O +ATOM 700 CG2 THR A 163 37.675 7.786 5.140 1.00 31.56 C +ATOM 701 N ALA A 164 36.751 9.754 9.050 1.00 35.93 N +ATOM 702 CA ALA A 164 36.788 10.865 10.004 1.00 35.93 C +ATOM 703 C ALA A 164 37.072 12.229 9.335 1.00 35.93 C +ATOM 704 O ALA A 164 36.637 13.262 9.836 1.00 35.93 O +ATOM 705 CB ALA A 164 37.852 10.541 11.057 1.00 35.93 C +ATOM 706 N THR A 165 37.758 12.235 8.185 1.00 39.39 N +ATOM 707 CA THR A 165 38.047 13.438 7.385 1.00 39.39 C +ATOM 708 C THR A 165 36.894 13.863 6.471 1.00 39.39 C +ATOM 709 O THR A 165 36.987 14.899 5.822 1.00 39.39 O +ATOM 710 CB THR A 165 39.326 13.258 6.549 1.00 39.39 C +ATOM 711 OG1 THR A 165 39.245 12.120 5.713 1.00 39.39 O +ATOM 712 CG2 THR A 165 40.560 13.060 7.428 1.00 39.39 C +ATOM 713 N GLY A 166 35.801 13.092 6.410 1.00 40.61 N +ATOM 714 CA GLY A 166 34.644 13.396 5.564 1.00 40.61 C +ATOM 715 C GLY A 166 34.676 12.766 4.167 1.00 40.61 C +ATOM 716 O GLY A 166 33.775 13.025 3.373 1.00 40.61 O +ATOM 717 N LEU A 167 35.675 11.932 3.856 1.00 36.14 N +ATOM 718 CA LEU A 167 35.775 11.244 2.568 1.00 36.14 C +ATOM 719 C LEU A 167 34.804 10.058 2.520 1.00 36.14 C +ATOM 720 O LEU A 167 34.902 9.137 3.333 1.00 36.14 O +ATOM 721 CB LEU A 167 37.234 10.805 2.331 1.00 36.14 C +ATOM 722 CG LEU A 167 37.475 10.079 0.994 1.00 36.14 C +ATOM 723 CD1 LEU A 167 37.301 11.018 -0.200 1.00 36.14 C +ATOM 724 CD2 LEU A 167 38.893 9.510 0.968 1.00 36.14 C +ATOM 725 N TYR A 168 33.894 10.052 1.548 1.00 30.92 N +ATOM 726 CA TYR A 168 33.012 8.910 1.297 1.00 30.92 C +ATOM 727 C TYR A 168 33.734 7.786 0.555 1.00 30.92 C +ATOM 728 O TYR A 168 34.484 8.043 -0.386 1.00 30.92 O +ATOM 729 CB TYR A 168 31.777 9.337 0.507 1.00 30.92 C +ATOM 730 CG TYR A 168 30.809 10.189 1.295 1.00 30.92 C +ATOM 731 CD1 TYR A 168 29.997 9.593 2.278 1.00 30.92 C +ATOM 732 CD2 TYR A 168 30.727 11.574 1.053 1.00 30.92 C +ATOM 733 CE1 TYR A 168 29.105 10.385 3.023 1.00 30.92 C +ATOM 734 CE2 TYR A 168 29.833 12.366 1.796 1.00 30.92 C +ATOM 735 CZ TYR A 168 29.026 11.771 2.789 1.00 30.92 C +ATOM 736 OH TYR A 168 28.172 12.532 3.522 1.00 30.92 O +ATOM 737 N TYR A 169 33.454 6.538 0.927 1.00 29.73 N +ATOM 738 CA TYR A 169 34.078 5.366 0.315 1.00 29.73 C +ATOM 739 C TYR A 169 33.153 4.144 0.286 1.00 29.73 C +ATOM 740 O TYR A 169 32.112 4.118 0.943 1.00 29.73 O +ATOM 741 CB TYR A 169 35.387 5.057 1.045 1.00 29.73 C +ATOM 742 CG TYR A 169 35.226 4.489 2.443 1.00 29.73 C +ATOM 743 CD1 TYR A 169 35.126 5.356 3.550 1.00 29.73 C +ATOM 744 CD2 TYR A 169 35.244 3.094 2.642 1.00 29.73 C +ATOM 745 CE1 TYR A 169 35.100 4.825 4.854 1.00 29.73 C +ATOM 746 CE2 TYR A 169 35.219 2.561 3.942 1.00 29.73 C +ATOM 747 CZ TYR A 169 35.166 3.429 5.051 1.00 29.73 C +ATOM 748 OH TYR A 169 35.162 2.922 6.308 1.00 29.73 O +ATOM 749 N GLY A 170 33.542 3.139 -0.504 1.00 28.83 N +ATOM 750 CA GLY A 170 32.797 1.892 -0.673 1.00 28.83 C +ATOM 751 C GLY A 170 31.461 2.034 -1.400 1.00 28.83 C +ATOM 752 O GLY A 170 31.140 3.101 -1.950 1.00 28.83 O +ATOM 753 N ASP A 171 30.713 0.934 -1.396 1.00 27.84 N +ATOM 754 CA ASP A 171 29.398 0.813 -2.025 1.00 27.84 C +ATOM 755 C ASP A 171 28.353 1.730 -1.370 1.00 27.84 C +ATOM 756 O ASP A 171 28.417 2.046 -0.178 1.00 27.84 O +ATOM 757 CB ASP A 171 28.920 -0.652 -1.981 1.00 27.84 C +ATOM 758 CG ASP A 171 29.793 -1.623 -2.786 1.00 27.84 C +ATOM 759 OD1 ASP A 171 30.561 -1.148 -3.653 1.00 27.84 O +ATOM 760 OD2 ASP A 171 29.672 -2.840 -2.514 1.00 27.84 O +ATOM 761 N VAL A 172 27.360 2.146 -2.161 1.00 27.80 N +ATOM 762 CA VAL A 172 26.128 2.750 -1.634 1.00 27.80 C +ATOM 763 C VAL A 172 25.182 1.617 -1.256 1.00 27.80 C +ATOM 764 O VAL A 172 24.826 0.807 -2.111 1.00 27.80 O +ATOM 765 CB VAL A 172 25.461 3.698 -2.649 1.00 27.80 C +ATOM 766 CG1 VAL A 172 24.166 4.304 -2.087 1.00 27.80 C +ATOM 767 CG2 VAL A 172 26.390 4.863 -3.011 1.00 27.80 C +ATOM 768 N ILE A 173 24.766 1.582 0.006 1.00 27.71 N +ATOM 769 CA ILE A 173 23.800 0.608 0.524 1.00 27.71 C +ATOM 770 C ILE A 173 22.440 1.268 0.726 1.00 27.71 C +ATOM 771 O ILE A 173 22.344 2.480 0.939 1.00 27.71 O +ATOM 772 CB ILE A 173 24.316 -0.112 1.785 1.00 27.71 C +ATOM 773 CG1 ILE A 173 24.438 0.842 2.990 1.00 27.71 C +ATOM 774 CG2 ILE A 173 25.632 -0.836 1.455 1.00 27.71 C +ATOM 775 CD1 ILE A 173 25.070 0.201 4.227 1.00 27.71 C +ATOM 776 N TRP A 174 21.383 0.466 0.647 1.00 27.97 N +ATOM 777 CA TRP A 174 20.023 0.970 0.513 1.00 27.97 C +ATOM 778 C TRP A 174 19.114 0.476 1.632 1.00 27.97 C +ATOM 779 O TRP A 174 19.258 -0.626 2.158 1.00 27.97 O +ATOM 780 CB TRP A 174 19.473 0.608 -0.872 1.00 27.97 C +ATOM 781 CG TRP A 174 20.271 1.154 -2.017 1.00 27.97 C +ATOM 782 CD1 TRP A 174 21.413 0.615 -2.503 1.00 27.97 C +ATOM 783 CD2 TRP A 174 19.989 2.310 -2.862 1.00 27.97 C +ATOM 784 NE1 TRP A 174 21.885 1.379 -3.548 1.00 27.97 N +ATOM 785 CE2 TRP A 174 21.017 2.405 -3.849 1.00 27.97 C +ATOM 786 CE3 TRP A 174 18.939 3.253 -2.929 1.00 27.97 C +ATOM 787 CZ2 TRP A 174 20.984 3.353 -4.877 1.00 27.97 C +ATOM 788 CZ3 TRP A 174 18.872 4.182 -3.985 1.00 27.97 C +ATOM 789 CH2 TRP A 174 19.874 4.205 -4.974 1.00 27.97 C +ATOM 790 N VAL A 175 18.137 1.305 1.970 1.00 28.47 N +ATOM 791 CA VAL A 175 16.888 0.884 2.602 1.00 28.47 C +ATOM 792 C VAL A 175 15.816 1.028 1.536 1.00 28.47 C +ATOM 793 O VAL A 175 15.683 2.119 0.989 1.00 28.47 O +ATOM 794 CB VAL A 175 16.503 1.769 3.799 1.00 28.47 C +ATOM 795 CG1 VAL A 175 15.450 1.069 4.663 1.00 28.47 C +ATOM 796 CG2 VAL A 175 17.652 2.177 4.712 1.00 28.47 C +ATOM 797 N ILE A 176 15.059 -0.025 1.244 1.00 28.69 N +ATOM 798 CA ILE A 176 13.919 0.027 0.321 1.00 28.69 C +ATOM 799 C ILE A 176 12.722 -0.564 1.054 1.00 28.69 C +ATOM 800 O ILE A 176 12.764 -1.726 1.450 1.00 28.69 O +ATOM 801 CB ILE A 176 14.213 -0.721 -0.998 1.00 28.69 C +ATOM 802 CG1 ILE A 176 15.450 -0.124 -1.709 1.00 28.69 C +ATOM 803 CG2 ILE A 176 12.975 -0.663 -1.914 1.00 28.69 C +ATOM 804 CD1 ILE A 176 15.893 -0.902 -2.954 1.00 28.69 C +ATOM 805 N LEU A 177 11.676 0.237 1.250 1.00 29.85 N +ATOM 806 CA LEU A 177 10.506 -0.139 2.041 1.00 29.85 C +ATOM 807 C LEU A 177 9.227 -0.019 1.229 1.00 29.85 C +ATOM 808 O LEU A 177 8.999 1.007 0.584 1.00 29.85 O +ATOM 809 CB LEU A 177 10.380 0.759 3.283 1.00 29.85 C +ATOM 810 CG LEU A 177 11.569 0.691 4.240 1.00 29.85 C +ATOM 811 CD1 LEU A 177 11.363 1.682 5.378 1.00 29.85 C +ATOM 812 CD2 LEU A 177 11.761 -0.692 4.852 1.00 29.85 C +ATOM 813 N SER A 178 8.357 -1.018 1.383 1.00 28.94 N +ATOM 814 CA SER A 178 6.935 -0.859 1.101 1.00 28.94 C +ATOM 815 C SER A 178 6.203 -0.461 2.379 1.00 28.94 C +ATOM 816 O SER A 178 6.109 -1.250 3.321 1.00 28.94 O +ATOM 817 CB SER A 178 6.321 -2.142 0.540 1.00 28.94 C +ATOM 818 OG SER A 178 5.009 -1.773 0.163 1.00 28.94 O +ATOM 819 N VAL A 179 5.678 0.758 2.449 1.00 28.72 N +ATOM 820 CA VAL A 179 4.826 1.171 3.569 1.00 28.72 C +ATOM 821 C VAL A 179 3.381 0.836 3.237 1.00 28.72 C +ATOM 822 O VAL A 179 2.785 1.440 2.345 1.00 28.72 O +ATOM 823 CB VAL A 179 4.994 2.648 3.930 1.00 28.72 C +ATOM 824 CG1 VAL A 179 4.138 3.006 5.150 1.00 28.72 C +ATOM 825 CG2 VAL A 179 6.445 2.971 4.283 1.00 28.72 C +ATOM 826 N GLU A 180 2.795 -0.099 3.979 1.00 29.21 N +ATOM 827 CA GLU A 180 1.428 -0.578 3.769 1.00 29.21 C +ATOM 828 C GLU A 180 0.761 -0.889 5.107 1.00 29.21 C +ATOM 829 O GLU A 180 1.371 -1.490 5.985 1.00 29.21 O +ATOM 830 CB GLU A 180 1.420 -1.841 2.895 1.00 29.21 C +ATOM 831 CG GLU A 180 1.824 -1.581 1.435 1.00 29.21 C +ATOM 832 CD GLU A 180 1.857 -2.870 0.598 1.00 29.21 C +ATOM 833 OE1 GLU A 180 2.617 -2.908 -0.397 1.00 29.21 O +ATOM 834 OE2 GLU A 180 1.109 -3.809 0.948 1.00 29.21 O +END diff --git a/tests/files/restraints/GLU_ASP_renamed.cif b/tests/files/restraints/GLU_ASP_renamed.cif new file mode 100644 index 00000000..0140a35e --- /dev/null +++ b/tests/files/restraints/GLU_ASP_renamed.cif @@ -0,0 +1,406 @@ +data_comp_list +loop_ +_chem_comp.id +_chem_comp.three_letter_code +_chem_comp.name +_chem_comp.group +_chem_comp.number_atoms_all +_chem_comp.number_atoms_nh +_chem_comp.desc_level +GLU GLU "GLUTAMIC ACID" peptide 18 10 . +ASP ASP "ASPARTIC ACID" peptide 15 9 . + +data_comp_GLU +loop_ +_chem_comp_atom.comp_id +_chem_comp_atom.atom_id +_chem_comp_atom.type_symbol +_chem_comp_atom.type_energy +_chem_comp_atom.charge +_chem_comp_atom.x +_chem_comp_atom.y +_chem_comp_atom.z +GLU N N NT3 1 88.319 -7.751 -10.089 +GLU CA C CH1 0 87.677 -7.162 -11.296 +GLU C C C 0 88.359 -5.826 -11.640 +GLU O O O 0 88.389 -4.954 -10.744 +GLU CB C CH2 0 86.177 -6.950 -11.080 +GLU CG C CH2 0 85.390 -8.247 -10.908 +GLU CD C C 0 83.891 -8.048 -10.773 +GLU OE1 O O 0 83.281 -7.495 -11.711 +GLU OE2 O OC -1 83.334 -8.447 -9.728 +GLU OXT O OC -1 88.836 -5.708 -12.790 +GLU H H H 0 88.066 -7.298 -9.352 +GLU H2 H H 0 89.218 -7.717 -10.158 +GLU H3 H H 0 88.077 -8.615 -9.999 +GLU HAX H H 0 87.806 -7.788 -12.054 +GLU HBY H H 0 85.814 -6.460 -11.847 +GLU HBX H H 0 86.049 -6.394 -10.284 +GLU HGY H H 0 85.714 -8.714 -10.110 +GLU HGX H H 0 85.559 -8.827 -11.681 + +loop_ +_chem_comp_tree.comp_id +_chem_comp_tree.atom_id +_chem_comp_tree.atom_back +_chem_comp_tree.atom_forward +_chem_comp_tree.connect_type +GLU N n/a CA START +GLU H N . . +GLU H2 N . . +GLU H3 N . . +GLU CA N C . +GLU HAX CA . . +GLU CB CA CG . +GLU HBY CB . . +GLU HBX CB . . +GLU CG CB CD . +GLU HGY CG . . +GLU HGX CG . . +GLU CD CG OE2 . +GLU OE1 CD . . +GLU OE2 CD . . +GLU C CA . END +GLU O C . . +GLU OXT C . . + +loop_ +_chem_comp_acedrg.comp_id +_chem_comp_acedrg.atom_id +_chem_comp_acedrg.atom_type +GLU N N(CCCH)(H)3 +GLU CA C(CCHH)(NH3)(COO)(H) +GLU C C(CCHN)(O)2 +GLU O O(CCO) +GLU CB C(CCHH)(CCHN)(H)2 +GLU CG C(CCHH)(COO)(H)2 +GLU CD C(CCHH)(O)2 +GLU OE1 O(CCO) +GLU OE2 O(CCO) +GLU OXT O(CCO) +GLU H H(NCHH) +GLU H2 H(NCHH) +GLU H3 H(NCHH) +GLU HAX H(CCCN) +GLU HBY H(CCCH) +GLU HBX H(CCCH) +GLU HGY H(CCCH) +GLU HGX H(CCCH) + +loop_ +_chem_comp_bond.comp_id +_chem_comp_bond.atom_id_1 +_chem_comp_bond.atom_id_2 +_chem_comp_bond.type +_chem_comp_bond.aromatic +_chem_comp_bond.value_dist_nucleus +_chem_comp_bond.value_dist_nucleus_esd +_chem_comp_bond.value_dist +_chem_comp_bond.value_dist_esd +GLU N CA SINGLE n 1.487 0.0100 1.487 0.0100 +GLU CA C SINGLE n 1.538 0.0113 1.538 0.0113 +GLU CA CB SINGLE n 1.529 0.0100 1.529 0.0100 +GLU C O DOUBLE n 1.251 0.0183 1.251 0.0183 +GLU C OXT SINGLE n 1.251 0.0183 1.251 0.0183 +GLU CB CG SINGLE n 1.526 0.0100 1.526 0.0100 +GLU CG CD SINGLE n 1.518 0.0135 1.518 0.0135 +GLU CD OE1 DOUBLE n 1.249 0.0161 1.249 0.0161 +GLU CD OE2 SINGLE n 1.249 0.0161 1.249 0.0161 +GLU N H SINGLE n 1.018 0.0520 0.902 0.0102 +GLU N H2 SINGLE n 1.018 0.0520 0.902 0.0102 +GLU N H3 SINGLE n 1.018 0.0520 0.902 0.0102 +GLU CA HAX SINGLE n 1.092 0.0100 0.991 0.0200 +GLU CB HBY SINGLE n 1.092 0.0100 0.980 0.0168 +GLU CB HBX SINGLE n 1.092 0.0100 0.980 0.0168 +GLU CG HGY SINGLE n 1.092 0.0100 0.981 0.0172 +GLU CG HGX SINGLE n 1.092 0.0100 0.981 0.0172 + +loop_ +_chem_comp_angle.comp_id +_chem_comp_angle.atom_id_1 +_chem_comp_angle.atom_id_2 +_chem_comp_angle.atom_id_3 +_chem_comp_angle.value_angle +_chem_comp_angle.value_angle_esd +GLU CA N H 109.990 3.00 +GLU CA N H2 109.990 3.00 +GLU CA N H3 109.990 3.00 +GLU H N H2 109.032 3.00 +GLU H N H3 109.032 3.00 +GLU H2 N H3 109.032 3.00 +GLU N CA C 109.258 1.50 +GLU N CA CB 110.440 2.46 +GLU N CA HAX 108.387 1.58 +GLU C CA CB 111.059 3.00 +GLU C CA HAX 108.774 1.79 +GLU CB CA HAX 109.080 2.33 +GLU CA C O 117.148 1.60 +GLU CA C OXT 117.148 1.60 +GLU O C OXT 125.704 1.50 +GLU CA CB CG 113.294 1.61 +GLU CA CB HBY 108.677 1.74 +GLU CA CB HBX 108.677 1.74 +GLU CG CB HBY 108.696 2.80 +GLU CG CB HBX 108.696 2.80 +GLU HBY CB HBX 107.655 1.50 +GLU CB CG CD 114.140 3.00 +GLU CB CG HGY 108.968 1.50 +GLU CB CG HGX 108.968 1.50 +GLU CD CG HGY 108.472 1.50 +GLU CD CG HGX 108.472 1.50 +GLU HGY CG HGX 107.541 1.92 +GLU CG CD OE1 118.251 3.00 +GLU CG CD OE2 118.251 3.00 +GLU OE1 CD OE2 123.498 1.82 + +loop_ +_chem_comp_tor.comp_id +_chem_comp_tor.id +_chem_comp_tor.atom_id_1 +_chem_comp_tor.atom_id_2 +_chem_comp_tor.atom_id_3 +_chem_comp_tor.atom_id_4 +_chem_comp_tor.value_angle +_chem_comp_tor.value_angle_esd +_chem_comp_tor.period +GLU chi1 N CA CB CG -60.000 10.0 3 +GLU chi2 CA CB CG CD 180.000 10.0 3 +GLU chi3 CB CG CD OE1 180.000 10.0 6 +GLU sp3_sp3_1 C CA N H 180.000 10.0 3 +GLU sp2_sp3_1 O C CA N 0.000 10.0 6 + +loop_ +_chem_comp_chir.comp_id +_chem_comp_chir.id +_chem_comp_chir.atom_id_centre +_chem_comp_chir.atom_id_1 +_chem_comp_chir.atom_id_2 +_chem_comp_chir.atom_id_3 +_chem_comp_chir.volume_sign +GLU chir_1 CA N C CB positive + +loop_ +_chem_comp_plane_atom.comp_id +_chem_comp_plane_atom.plane_id +_chem_comp_plane_atom.atom_id +_chem_comp_plane_atom.dist_esd +GLU plan-1 C 0.020 +GLU plan-1 CA 0.020 +GLU plan-1 O 0.020 +GLU plan-1 OXT 0.020 +GLU plan-2 CD 0.020 +GLU plan-2 CG 0.020 +GLU plan-2 OE1 0.020 +GLU plan-2 OE2 0.020 + +loop_ +_pdbx_chem_comp_descriptor.comp_id +_pdbx_chem_comp_descriptor.type +_pdbx_chem_comp_descriptor.program +_pdbx_chem_comp_descriptor.program_version +_pdbx_chem_comp_descriptor.descriptor +GLU SMILES ACDLabs 12.01 O=C(O)C(N)CCC(=O)O +GLU SMILES_CANONICAL CACTVS 3.370 N[C@@H](CCC(O)=O)C(O)=O +GLU SMILES CACTVS 3.370 N[CH](CCC(O)=O)C(O)=O +GLU SMILES_CANONICAL "OpenEye OEToolkits" 1.7.0 C(CC(=O)O)[C@@H](C(=O)O)N +GLU SMILES "OpenEye OEToolkits" 1.7.0 C(CC(=O)O)C(C(=O)O)N +GLU InChI InChI 1.03 InChI=1S/C5H9NO4/c6-3(5(9)10)1-2-4(7)8/h3H,1-2,6H2,(H,7,8)(H,9,10)/t3-/m0/s1 +GLU InChIKey InChI 1.03 WHUUTDBJXJRKMK-VKHMYHEASA-N + +loop_ +_pdbx_chem_comp_description_generator.comp_id +_pdbx_chem_comp_description_generator.program_name +_pdbx_chem_comp_description_generator.program_version +_pdbx_chem_comp_description_generator.descriptor +GLU acedrg 278 "dictionary generator" +GLU acedrg_database 12 "data source" +GLU rdkit 2019.09.1 "Chemoinformatics tool" +GLU refmac5 5.8.0419 "optimization tool" + +data_comp_ASP +loop_ +_chem_comp_atom.comp_id +_chem_comp_atom.atom_id +_chem_comp_atom.type_symbol +_chem_comp_atom.type_energy +_chem_comp_atom.charge +_chem_comp_atom.x +_chem_comp_atom.y +_chem_comp_atom.z +ASP N N NT3 1 33.542 17.835 39.145 +ASP CA C CH1 0 34.991 17.614 38.867 +ASP C C C 0 35.178 17.161 37.413 +ASP O O O 0 36.268 17.442 36.867 +ASP CB C CH2 0 35.617 16.650 39.866 +ASP CG C C 0 34.932 15.299 40.037 +ASP OD1 O O 0 35.515 14.435 40.725 +ASP OD2 O OC -1 33.817 15.122 39.501 +ASP OXT O OC -1 34.232 16.542 36.875 +ASP H H H 0 33.408 17.940 40.031 +ASP H2 H H 0 33.045 17.139 38.857 +ASP H3 H H 0 33.263 18.581 38.722 +ASP HA H H 0 35.453 18.476 38.978 +ASP HBR H H 0 35.641 17.087 40.742 +ASP HBQ H H 0 36.543 16.483 39.592 + +loop_ +_chem_comp_tree.comp_id +_chem_comp_tree.atom_id +_chem_comp_tree.atom_back +_chem_comp_tree.atom_forward +_chem_comp_tree.connect_type +ASP N n/a CA START +ASP H N . . +ASP H2 N . . +ASP H3 N . . +ASP CA N C . +ASP HA CA . . +ASP CB CA CG . +ASP HBR CB . . +ASP HBQ CB . . +ASP CG CB OD2 . +ASP OD1 CG . . +ASP OD2 CG . . +ASP C CA . END +ASP O C . . +ASP OXT C . . + +loop_ +_chem_comp_acedrg.comp_id +_chem_comp_acedrg.atom_id +_chem_comp_acedrg.atom_type +ASP N N(CCCH)(H)3 +ASP CA C(CCHH)(NH3)(COO)(H) +ASP C C(CCHN)(O)2 +ASP O O(CCO) +ASP CB C(CCHN)(COO)(H)2 +ASP CG C(CCHH)(O)2 +ASP OD1 O(CCO) +ASP OD2 O(CCO) +ASP OXT O(CCO) +ASP H H(NCHH) +ASP H2 H(NCHH) +ASP H3 H(NCHH) +ASP HA H(CCCN) +ASP HBR H(CCCH) +ASP HBQ H(CCCH) + +loop_ +_chem_comp_bond.comp_id +_chem_comp_bond.atom_id_1 +_chem_comp_bond.atom_id_2 +_chem_comp_bond.type +_chem_comp_bond.aromatic +_chem_comp_bond.value_dist_nucleus +_chem_comp_bond.value_dist_nucleus_esd +_chem_comp_bond.value_dist +_chem_comp_bond.value_dist_esd +ASP N CA SINGLE n 1.490 0.0100 1.490 0.0100 +ASP CA C SINGLE n 1.533 0.0100 1.533 0.0100 +ASP CA CB SINGLE n 1.521 0.0100 1.521 0.0100 +ASP C O DOUBLE n 1.251 0.0183 1.251 0.0183 +ASP C OXT SINGLE n 1.251 0.0183 1.251 0.0183 +ASP CB CG SINGLE n 1.522 0.0100 1.522 0.0100 +ASP CG OD1 DOUBLE n 1.249 0.0161 1.249 0.0161 +ASP CG OD2 SINGLE n 1.249 0.0161 1.249 0.0161 +ASP N H SINGLE n 1.018 0.0520 0.902 0.0102 +ASP N H2 SINGLE n 1.018 0.0520 0.902 0.0102 +ASP N H3 SINGLE n 1.018 0.0520 0.902 0.0102 +ASP CA HA SINGLE n 1.092 0.0100 0.984 0.0200 +ASP CB HBR SINGLE n 1.092 0.0100 0.980 0.0165 +ASP CB HBQ SINGLE n 1.092 0.0100 0.980 0.0165 + +loop_ +_chem_comp_angle.comp_id +_chem_comp_angle.atom_id_1 +_chem_comp_angle.atom_id_2 +_chem_comp_angle.atom_id_3 +_chem_comp_angle.value_angle +_chem_comp_angle.value_angle_esd +ASP CA N H 109.990 3.00 +ASP CA N H2 109.990 3.00 +ASP CA N H3 109.990 3.00 +ASP H N H2 109.032 3.00 +ASP H N H3 109.032 3.00 +ASP H2 N H3 109.032 3.00 +ASP N CA C 109.258 1.50 +ASP N CA CB 111.400 1.50 +ASP N CA HA 108.387 1.58 +ASP C CA CB 112.421 3.00 +ASP C CA HA 108.774 1.79 +ASP CB CA HA 108.472 2.65 +ASP CA C O 117.148 1.60 +ASP CA C OXT 117.148 1.60 +ASP O C OXT 125.704 1.50 +ASP CA CB CG 115.436 1.50 +ASP CA CB HBR 108.799 3.00 +ASP CA CB HBQ 108.799 3.00 +ASP CG CB HBR 108.242 2.79 +ASP CG CB HBQ 108.242 2.79 +ASP HBR CB HBQ 107.976 2.66 +ASP CB CG OD1 117.985 1.50 +ASP CB CG OD2 117.985 1.50 +ASP OD1 CG OD2 124.031 1.82 + +loop_ +_chem_comp_tor.comp_id +_chem_comp_tor.id +_chem_comp_tor.atom_id_1 +_chem_comp_tor.atom_id_2 +_chem_comp_tor.atom_id_3 +_chem_comp_tor.atom_id_4 +_chem_comp_tor.value_angle +_chem_comp_tor.value_angle_esd +_chem_comp_tor.period +ASP chi1 N CA CB CG -60.000 10.0 3 +ASP chi2 CA CB CG OD1 180.000 10.0 6 +ASP sp3_sp3_1 C CA N H 180.000 10.0 3 +ASP sp2_sp3_1 O C CA N 0.000 10.0 6 + +loop_ +_chem_comp_chir.comp_id +_chem_comp_chir.id +_chem_comp_chir.atom_id_centre +_chem_comp_chir.atom_id_1 +_chem_comp_chir.atom_id_2 +_chem_comp_chir.atom_id_3 +_chem_comp_chir.volume_sign +ASP chir_1 CA N C CB positive + +loop_ +_chem_comp_plane_atom.comp_id +_chem_comp_plane_atom.plane_id +_chem_comp_plane_atom.atom_id +_chem_comp_plane_atom.dist_esd +ASP plan-1 C 0.020 +ASP plan-1 CA 0.020 +ASP plan-1 O 0.020 +ASP plan-1 OXT 0.020 +ASP plan-2 CB 0.020 +ASP plan-2 CG 0.020 +ASP plan-2 OD1 0.020 +ASP plan-2 OD2 0.020 + +loop_ +_pdbx_chem_comp_descriptor.comp_id +_pdbx_chem_comp_descriptor.type +_pdbx_chem_comp_descriptor.program +_pdbx_chem_comp_descriptor.program_version +_pdbx_chem_comp_descriptor.descriptor +ASP SMILES ACDLabs 12.01 O=C(O)CC(N)C(=O)O +ASP SMILES_CANONICAL CACTVS 3.370 N[C@@H](CC(O)=O)C(O)=O +ASP SMILES CACTVS 3.370 N[CH](CC(O)=O)C(O)=O +ASP SMILES_CANONICAL "OpenEye OEToolkits" 1.7.0 C([C@@H](C(=O)O)N)C(=O)O +ASP SMILES "OpenEye OEToolkits" 1.7.0 C(C(C(=O)O)N)C(=O)O +ASP InChI InChI 1.03 InChI=1S/C4H7NO4/c5-2(4(8)9)1-3(6)7/h2H,1,5H2,(H,6,7)(H,8,9)/t2-/m0/s1 +ASP InChIKey InChI 1.03 CKLJMWTZIZZHCS-REOHCLBHSA-N + +loop_ +_pdbx_chem_comp_description_generator.comp_id +_pdbx_chem_comp_description_generator.program_name +_pdbx_chem_comp_description_generator.program_version +_pdbx_chem_comp_description_generator.descriptor +ASP acedrg 278 "dictionary generator" +ASP acedrg_database 12 "data source" +ASP rdkit 2019.09.1 "Chemoinformatics tool" +ASP refmac5 5.8.0419 "optimization tool" diff --git a/tests/files/restraints/GLU_renamed.cif b/tests/files/restraints/GLU_renamed.cif new file mode 100644 index 00000000..bfec34e2 --- /dev/null +++ b/tests/files/restraints/GLU_renamed.cif @@ -0,0 +1,217 @@ +data_comp_list +loop_ +_chem_comp.id +_chem_comp.three_letter_code +_chem_comp.name +_chem_comp.group +_chem_comp.number_atoms_all +_chem_comp.number_atoms_nh +_chem_comp.desc_level +GLU GLU "GLUTAMIC ACID" peptide 18 10 . + +data_comp_GLU +loop_ +_chem_comp_atom.comp_id +_chem_comp_atom.atom_id +_chem_comp_atom.type_symbol +_chem_comp_atom.type_energy +_chem_comp_atom.charge +_chem_comp_atom.x +_chem_comp_atom.y +_chem_comp_atom.z +GLU N N NT3 1 88.319 -7.751 -10.089 +GLU CA C CH1 0 87.677 -7.162 -11.296 +GLU C C C 0 88.359 -5.826 -11.640 +GLU O O O 0 88.389 -4.954 -10.744 +GLU CB C CH2 0 86.177 -6.950 -11.080 +GLU CG C CH2 0 85.390 -8.247 -10.908 +GLU CD C C 0 83.891 -8.048 -10.773 +GLU OE1 O O 0 83.281 -7.495 -11.711 +GLU OE2 O OC -1 83.334 -8.447 -9.728 +GLU OXT O OC -1 88.836 -5.708 -12.790 +GLU H H H 0 88.066 -7.298 -9.352 +GLU H2 H H 0 89.218 -7.717 -10.158 +GLU H3 H H 0 88.077 -8.615 -9.999 +GLU HAX H H 0 87.806 -7.788 -12.054 +GLU HBY H H 0 85.814 -6.460 -11.847 +GLU HBX H H 0 86.049 -6.394 -10.284 +GLU HGY H H 0 85.714 -8.714 -10.110 +GLU HGX H H 0 85.559 -8.827 -11.681 + +loop_ +_chem_comp_tree.comp_id +_chem_comp_tree.atom_id +_chem_comp_tree.atom_back +_chem_comp_tree.atom_forward +_chem_comp_tree.connect_type +GLU N n/a CA START +GLU H N . . +GLU H2 N . . +GLU H3 N . . +GLU CA N C . +GLU HAX CA . . +GLU CB CA CG . +GLU HBY CB . . +GLU HBX CB . . +GLU CG CB CD . +GLU HGY CG . . +GLU HGX CG . . +GLU CD CG OE2 . +GLU OE1 CD . . +GLU OE2 CD . . +GLU C CA . END +GLU O C . . +GLU OXT C . . + +loop_ +_chem_comp_acedrg.comp_id +_chem_comp_acedrg.atom_id +_chem_comp_acedrg.atom_type +GLU N N(CCCH)(H)3 +GLU CA C(CCHH)(NH3)(COO)(H) +GLU C C(CCHN)(O)2 +GLU O O(CCO) +GLU CB C(CCHH)(CCHN)(H)2 +GLU CG C(CCHH)(COO)(H)2 +GLU CD C(CCHH)(O)2 +GLU OE1 O(CCO) +GLU OE2 O(CCO) +GLU OXT O(CCO) +GLU H H(NCHH) +GLU H2 H(NCHH) +GLU H3 H(NCHH) +GLU HAX H(CCCN) +GLU HBY H(CCCH) +GLU HBX H(CCCH) +GLU HGY H(CCCH) +GLU HGX H(CCCH) + +loop_ +_chem_comp_bond.comp_id +_chem_comp_bond.atom_id_1 +_chem_comp_bond.atom_id_2 +_chem_comp_bond.type +_chem_comp_bond.aromatic +_chem_comp_bond.value_dist_nucleus +_chem_comp_bond.value_dist_nucleus_esd +_chem_comp_bond.value_dist +_chem_comp_bond.value_dist_esd +GLU N CA SINGLE n 1.487 0.0100 1.487 0.0100 +GLU CA C SINGLE n 1.538 0.0113 1.538 0.0113 +GLU CA CB SINGLE n 1.529 0.0100 1.529 0.0100 +GLU C O DOUBLE n 1.251 0.0183 1.251 0.0183 +GLU C OXT SINGLE n 1.251 0.0183 1.251 0.0183 +GLU CB CG SINGLE n 1.526 0.0100 1.526 0.0100 +GLU CG CD SINGLE n 1.518 0.0135 1.518 0.0135 +GLU CD OE1 DOUBLE n 1.249 0.0161 1.249 0.0161 +GLU CD OE2 SINGLE n 1.249 0.0161 1.249 0.0161 +GLU N H SINGLE n 1.018 0.0520 0.902 0.0102 +GLU N H2 SINGLE n 1.018 0.0520 0.902 0.0102 +GLU N H3 SINGLE n 1.018 0.0520 0.902 0.0102 +GLU CA HAX SINGLE n 1.092 0.0100 0.991 0.0200 +GLU CB HBY SINGLE n 1.092 0.0100 0.980 0.0168 +GLU CB HBX SINGLE n 1.092 0.0100 0.980 0.0168 +GLU CG HGY SINGLE n 1.092 0.0100 0.981 0.0172 +GLU CG HGX SINGLE n 1.092 0.0100 0.981 0.0172 + +loop_ +_chem_comp_angle.comp_id +_chem_comp_angle.atom_id_1 +_chem_comp_angle.atom_id_2 +_chem_comp_angle.atom_id_3 +_chem_comp_angle.value_angle +_chem_comp_angle.value_angle_esd +GLU CA N H 109.990 3.00 +GLU CA N H2 109.990 3.00 +GLU CA N H3 109.990 3.00 +GLU H N H2 109.032 3.00 +GLU H N H3 109.032 3.00 +GLU H2 N H3 109.032 3.00 +GLU N CA C 109.258 1.50 +GLU N CA CB 110.440 2.46 +GLU N CA HAX 108.387 1.58 +GLU C CA CB 111.059 3.00 +GLU C CA HAX 108.774 1.79 +GLU CB CA HAX 109.080 2.33 +GLU CA C O 117.148 1.60 +GLU CA C OXT 117.148 1.60 +GLU O C OXT 125.704 1.50 +GLU CA CB CG 113.294 1.61 +GLU CA CB HBY 108.677 1.74 +GLU CA CB HBX 108.677 1.74 +GLU CG CB HBY 108.696 2.80 +GLU CG CB HBX 108.696 2.80 +GLU HBY CB HBX 107.655 1.50 +GLU CB CG CD 114.140 3.00 +GLU CB CG HGY 108.968 1.50 +GLU CB CG HGX 108.968 1.50 +GLU CD CG HGY 108.472 1.50 +GLU CD CG HGX 108.472 1.50 +GLU HGY CG HGX 107.541 1.92 +GLU CG CD OE1 118.251 3.00 +GLU CG CD OE2 118.251 3.00 +GLU OE1 CD OE2 123.498 1.82 + +loop_ +_chem_comp_tor.comp_id +_chem_comp_tor.id +_chem_comp_tor.atom_id_1 +_chem_comp_tor.atom_id_2 +_chem_comp_tor.atom_id_3 +_chem_comp_tor.atom_id_4 +_chem_comp_tor.value_angle +_chem_comp_tor.value_angle_esd +_chem_comp_tor.period +GLU chi1 N CA CB CG -60.000 10.0 3 +GLU chi2 CA CB CG CD 180.000 10.0 3 +GLU chi3 CB CG CD OE1 180.000 10.0 6 +GLU sp3_sp3_1 C CA N H 180.000 10.0 3 +GLU sp2_sp3_1 O C CA N 0.000 10.0 6 + +loop_ +_chem_comp_chir.comp_id +_chem_comp_chir.id +_chem_comp_chir.atom_id_centre +_chem_comp_chir.atom_id_1 +_chem_comp_chir.atom_id_2 +_chem_comp_chir.atom_id_3 +_chem_comp_chir.volume_sign +GLU chir_1 CA N C CB positive + +loop_ +_chem_comp_plane_atom.comp_id +_chem_comp_plane_atom.plane_id +_chem_comp_plane_atom.atom_id +_chem_comp_plane_atom.dist_esd +GLU plan-1 C 0.020 +GLU plan-1 CA 0.020 +GLU plan-1 O 0.020 +GLU plan-1 OXT 0.020 +GLU plan-2 CD 0.020 +GLU plan-2 CG 0.020 +GLU plan-2 OE1 0.020 +GLU plan-2 OE2 0.020 + +loop_ +_pdbx_chem_comp_descriptor.comp_id +_pdbx_chem_comp_descriptor.type +_pdbx_chem_comp_descriptor.program +_pdbx_chem_comp_descriptor.program_version +_pdbx_chem_comp_descriptor.descriptor +GLU SMILES ACDLabs 12.01 O=C(O)C(N)CCC(=O)O +GLU SMILES_CANONICAL CACTVS 3.370 N[C@@H](CCC(O)=O)C(O)=O +GLU SMILES CACTVS 3.370 N[CH](CCC(O)=O)C(O)=O +GLU SMILES_CANONICAL "OpenEye OEToolkits" 1.7.0 C(CC(=O)O)[C@@H](C(=O)O)N +GLU SMILES "OpenEye OEToolkits" 1.7.0 C(CC(=O)O)C(C(=O)O)N +GLU InChI InChI 1.03 InChI=1S/C5H9NO4/c6-3(5(9)10)1-2-4(7)8/h3H,1-2,6H2,(H,7,8)(H,9,10)/t3-/m0/s1 +GLU InChIKey InChI 1.03 WHUUTDBJXJRKMK-VKHMYHEASA-N + +loop_ +_pdbx_chem_comp_description_generator.comp_id +_pdbx_chem_comp_description_generator.program_name +_pdbx_chem_comp_description_generator.program_version +_pdbx_chem_comp_description_generator.descriptor +GLU acedrg 278 "dictionary generator" +GLU acedrg_database 12 "data source" +GLU rdkit 2019.09.1 "Chemoinformatics tool" +GLU refmac5 5.8.0419 "optimization tool" diff --git a/tests/fixtures/README.md b/tests/fixtures/README.md new file mode 100644 index 00000000..b146cd38 --- /dev/null +++ b/tests/fixtures/README.md @@ -0,0 +1,47 @@ +# Fixture ownership + +The root `tests/conftest.py` owns pytest options, markers, capability gating, +and the `pytest_plugins` registry. Put reusable setup in the modules below. +Keep a fixture in its test module when only that module needs it. + +| Module | Responsibility | Visibility / lifetime | +|---|---|---| +| `paths.py` | Repository, bundled-data and optional library paths | All tests; session | +| `files.py` | Sample paths and named compatibility pairs | All tests; sample paths session-scoped, extended pairs function-scoped; no loading | +| `devices.py` | Configured device, explicit backends, device parametrization | All tests; existing per-fixture scopes | +| `precision.py` | Comparison tolerances and CPU-double reference context | All tests; reference fixture restores state after each test | +| `objects.py` | Mutable models, data, scalers and restraints | All tests; fresh per test except explicitly shared bundles | +| `collections.py` | Paired difference-refinement datasets, models and scalers | All tests; fresh per test | +| `numerical.py` | Synthetic tensors and factories | Imported only by `unit/conftest.py`; function | +| `functional.py` | Read-only `shared_model_ft` | Imported only by `functional/conftest.py`; module | + +Fixtures are available without imports in tests. Import reusable helpers from +their defining module, never from the root `conftest.py`. Subtree +conftests import fixture functions explicitly; register shared plugins only at +the root so pytest also works when invoked from a subdirectory. + +Use `shared_*` only for read-only checks. They capture the package configuration +at module setup and may populate derived caches. Do not move them, change their +parameters, tables, masks or grids, backpropagate through them, or use them in +tests that switch global configuration. A target or scaler can mutate a model it +borrows, so a shared model must not be passed to such an operation. + +Tests that verify loading must execute a fresh loader, directly or through a +function-scoped fixture. Tests of mutation, +device movement, or empty caches use fresh objects. `loaded_model`, +`loaded_model_ft`, `loaded_reflection_data`, and their composed fixtures in `objects.py` provide +fresh mutable objects per test. The explicitly shared session bundles in that +module retain their documented ownership contracts. + +`compatibility_structure_pair` selects named slow cases from +`tests/helpers/structure_cases.py`. `compatibility_model` loads just that model; +`compatibility_model_and_data` adds observations only when needed. Skipped slow +cases do not load any structures. + +Use `cpu_double_precision()` to scope an explicit numerical reference, or request +`double_cpu` for a single test. The structure-factor package uses the same context +at package scope; both usages restore dtype, device, and density cutoff on exit. + +This separation preserves the existing numerical-factory allocation policy and +test-selection policy. Those policies are independent of fixture registration and +scope, and can be revised in their respective modules. diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py new file mode 100644 index 00000000..660fb98a --- /dev/null +++ b/tests/fixtures/__init__.py @@ -0,0 +1,5 @@ +"""Provide pytest fixtures grouped by responsibility. + +Root conftest registers shared plugins. Unit and functional conftests explicitly +import their scoped fixtures; this package deliberately re-exports none. +""" diff --git a/tests/fixtures/collections.py b/tests/fixtures/collections.py new file mode 100644 index 00000000..c0cf2927 --- /dev/null +++ b/tests/fixtures/collections.py @@ -0,0 +1,40 @@ +"""Build fresh paired collections for difference-refinement integration tests.""" + +import pytest +import torch + + +@pytest.fixture +def difference_models(loaded_reflection_data, sample_structure_pair): + """Return independent dark/light models and raw 1DAW dataset copies per test.""" + from torchref.cli._common import load_model + from torchref.io import DatasetCollection + from torchref.model import ModelCollection + + data = loaded_reflection_data + assert data.I is not None + models = [ + load_model( + str(sample_structure_pair["model"]), + max_res=2.05, + device=data.device, + verbose=0, + ) + for _ in range(2) + ] + with torch.no_grad(): + models[1].xyz.refinable_params += 0.2 + dc = DatasetCollection(device=data.device, verbose=0) + dc.add_dataset("dark", data, set_as_reference=True).add_dataset("light", data) + mc = ModelCollection(models, dark_key="dark", verbose=0) + mc.add_dark().add_timepoint("light", [0.78, 0.22]) + return dc, mc + + +@pytest.fixture +def difference_collection(difference_models): + """Add a fresh initialized model-to-data scaler to the paired collection.""" + from torchref.scaling import CollectionScaler + + dc, mc = difference_models + return dc, mc, CollectionScaler(dc, mc, verbose=0).initialize() diff --git a/tests/fixtures/devices.py b/tests/fixtures/devices.py new file mode 100644 index 00000000..579daeff --- /dev/null +++ b/tests/fixtures/devices.py @@ -0,0 +1,119 @@ +"""Provide explicit backend fixtures and the configured default device. + +Capability probes are shared with collection hooks and structure-factor cases. +Device parametrization is constructed at import so collection can see its marks. +""" + +import pytest +import torch + + +def _cuda_available() -> bool: + return torch.cuda.is_available() + + +def _mps_available() -> bool: + return hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + + +def _accelerator() -> "torch.device | None": + """The canonical accelerator this host can actually use, or ``None``. + + Indices are filled in (``cuda:0`` / ``mps:0``) so the value compares equal + to a device read back off a real tensor -- ``torch.device('mps')`` and + ``torch.device('mps:0')`` are *not* equal even though they name the same + physical device. + """ + if _cuda_available(): + return torch.device("cuda", torch.cuda.current_device()) + if _mps_available(): + return torch.device("mps", 0) + return None + + +@pytest.fixture(scope="session") +def cpu_device() -> torch.device: + """CPU torch device.""" + return torch.device("cpu") + + +@pytest.fixture(scope="session") +def gpu_device() -> torch.device: + """Select CUDA, then MPS, for tests marked ``gpu``. + + Skip if neither backend is available. Use ``cuda_device`` or ``mps_device`` + when the test exercises a backend-specific contract. + """ + accel = _accelerator() + if accel is None: + pytest.skip("No accelerator (CUDA or MPS) on this host") + return accel + + +@pytest.fixture(scope="session") +def cuda_device() -> torch.device: + """Canonical CUDA device for ``cuda``-marked tests. + + Deliberately unguarded. What runs is decided by the ``cuda`` marker in + :func:`pytest_collection_modifyitems` and nowhere else, so this fixture does + not re-check availability: on a host without CUDA the test is *meant* to + error with the real backend error rather than be quietly skipped here. + """ + return torch.device("cuda", 0) + + +@pytest.fixture(scope="session") +def mps_device() -> torch.device: + """Canonical MPS device for ``mps``-marked tests. + + Unguarded for the same reason as :func:`cuda_device` -- the ``mps`` marker + owns the decision. + """ + return torch.device("mps", 0) + + +# Built at import time so the ``gpu`` mark is attached during *collection*. +# Adding it later (e.g. via ``request.node.add_marker`` inside the fixture) is +# too late for ``pytest_collection_modifyitems`` to gate on. +_DEVICE_PARAMS = [pytest.param(torch.device("cpu"), id="cpu")] +_ACCELERATOR = _accelerator() +if _ACCELERATOR is not None: + _DEVICE_PARAMS.append( + pytest.param( + _ACCELERATOR, + id=_ACCELERATOR.type, + # Backend-specific mark, so a CUDA-less host skips the cuda leg and + # a non-Mac skips the mps leg, each with an accurate reason. + marks=getattr(pytest.mark, _ACCELERATOR.type), + ) + ) + + +@pytest.fixture(params=_DEVICE_PARAMS) +def any_device(request: pytest.FixtureRequest) -> torch.device: + """Every device this host can actually use, one test run per device. + + The CPU leg always runs. An available accelerator runs automatically and + carries its backend-specific marker. No accelerator leg is created on a + CPU-only host. + """ + return request.param + + +@pytest.fixture +def device(request: pytest.FixtureRequest) -> torch.device: + """Default test device. + + Uses the package-wide auto-detected default (``torchref.device.current``) + so tests run on whichever device the user's machine resolved to at + import time: cuda -> mps -> cpu. Tests marked ``@pytest.mark.cuda_only`` + are skipped when CUDA is not available. + """ + from torchref.config import get_default_device + + markers = {m.name for m in request.node.iter_markers()} + if "cuda_only" in markers and not torch.cuda.is_available(): + pytest.skip("Test requires CUDA") + if "gpu" in markers and not (_cuda_available() or _mps_available()): + pytest.skip("No GPU (CUDA or MPS) available") + return get_default_device() diff --git a/tests/fixtures/files.py b/tests/fixtures/files.py new file mode 100644 index 00000000..390ea678 --- /dev/null +++ b/tests/fixtures/files.py @@ -0,0 +1,89 @@ +"""Select sample paths and matching model/reflection pairs without loading them.""" + +from pathlib import Path + +import pytest + +from tests.helpers.structure_cases import EXTENDED_PAIR_CODES + + +@pytest.fixture( + params=[pytest.param(code, marks=pytest.mark.slow) for code in EXTENDED_PAIR_CODES] +) +def compatibility_structure_pair( + cif_dir: Path, mtz_dir: Path, request: pytest.FixtureRequest +) -> dict: + """Select one named extended crystal without loading its model or observations.""" + code = request.param + return { + "pdb_id": code, + "model": cif_dir / f"{code}.cif", + "reflections": mtz_dir / f"{code}.mtz", + } + + +@pytest.fixture(scope="session") +def sample_cif_file(cif_dir: Path) -> Path: + """Return a sample CIF file for testing.""" + cif_file = cif_dir / "1DAW.cif" + if cif_file.exists(): + return cif_file + # Try any available CIF file + cif_files = list(cif_dir.glob("*.cif")) + if cif_files: + return cif_files[0] + pytest.skip("No CIF files found in test data") + + +@pytest.fixture(scope="session") +def sample_mtz_file(mtz_dir: Path) -> Path: + """Return a sample MTZ file for testing.""" + mtz_file = mtz_dir / "1DAW.mtz" + if mtz_file.exists(): + return mtz_file + # Try any available MTZ file + mtz_files = list(mtz_dir.glob("*.mtz")) + if mtz_files: + return mtz_files[0] + pytest.skip("No MTZ files found in test data") + + +@pytest.fixture(scope="session") +def sample_pdb_file(pdb_dir: Path) -> Path: + """Return a sample PDB file for testing.""" + pdb_files = sorted(pdb_dir.glob("*.pdb")) + if not pdb_files: + pytest.skip("No PDB files found in test data directory") + return pdb_files[0] + + +@pytest.fixture(scope="session") +def sample_structure_factor_cif(cif_sf_dir: Path) -> Path: + """Return a sample structure factor CIF file.""" + sf_files = sorted(cif_sf_dir.glob("*.cif")) + if not sf_files: + pytest.skip("No structure factor CIF files found") + return sf_files[0] + + +@pytest.fixture(scope="session") +def sample_structure_pair(cif_dir: Path, mtz_dir: Path) -> dict[str, Path]: + """Return a matching pair of CIF model and MTZ reflections.""" + # Try to find matching files + pdb_id = "1DAW" + cif_file = cif_dir / f"{pdb_id}.cif" + mtz_file = mtz_dir / f"{pdb_id}.mtz" + + if cif_file.exists() and mtz_file.exists(): + return {"model": cif_file, "reflections": mtz_file} + + # Try to find any matching pair + cif_files = {f.stem: f for f in cif_dir.glob("*.cif")} + mtz_files = {f.stem: f for f in mtz_dir.glob("*.mtz")} + + common_ids = set(cif_files.keys()) & set(mtz_files.keys()) + if common_ids: + pdb_id = min(common_ids) + return {"model": cif_files[pdb_id], "reflections": mtz_files[pdb_id]} + + pytest.skip("No matching CIF/MTZ pairs found in test data") diff --git a/tests/fixtures/functional.py b/tests/fixtures/functional.py new file mode 100644 index 00000000..5ee204a4 --- /dev/null +++ b/tests/fixtures/functional.py @@ -0,0 +1,25 @@ +"""Share loaded objects within a functional module for read-only checks. + +These fixtures capture the configured dtype/device at module setup. Callers may +populate derived caches but must not change parameters, tables, grids, masks, +device, or configuration. Tests of loading, mutation, and empty caches construct +fresh objects instead. No loaded objects are shared across test modules. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from torchref.model import ModelFT + + +@pytest.fixture(scope="module") +def shared_model_ft(sample_cif_file: Path) -> ModelFT: + """Load a read-only Fourier model with a 2 Å resolution limit per module.""" + from torchref.model import ModelFT + + return ModelFT(max_res=2.0, verbose=0).load_cif(str(sample_cif_file)) diff --git a/tests/fixtures/numerical.py b/tests/fixtures/numerical.py new file mode 100644 index 00000000..c0d54fc3 --- /dev/null +++ b/tests/fixtures/numerical.py @@ -0,0 +1,195 @@ +"""Generate small synthetic numerical inputs for unit tests. + +Imported by the unit conftest only. Factories return fresh CPU tensors on every +call, using TorchRef's numeric dtypes. They reset the global NumPy random seed; +``random_seed`` also resets PyTorch's seed. Accelerator coverage requires an +explicit move by the caller under this allocation policy. +""" + +from collections.abc import Callable + +import numpy as np +import pytest +import torch + +from torchref.config import dtypes + + +@pytest.fixture +def random_seed() -> int: + """Set random seed for reproducibility.""" + seed = 42 + np.random.seed(seed) + torch.manual_seed(seed) + return seed + + +@pytest.fixture +def random_coordinates() -> Callable[..., torch.Tensor]: + """Return a factory for Cartesian coordinates of shape (n_atoms, 3) in Å.""" + + def _generate(n_atoms: int = 10, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + return torch.tensor(np.random.rand(n_atoms, 3) * 10, dtype=dtypes.float) + + return _generate + + +@pytest.fixture +def random_fractional_coordinates() -> Callable[..., torch.Tensor]: + """Return a factory for fractional coordinates (n_atoms, 3) in [0, 1).""" + + def _generate(n_atoms: int = 10, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + return torch.tensor(np.random.rand(n_atoms, 3), dtype=dtypes.float) + + return _generate + + +@pytest.fixture +def random_adp() -> Callable[..., torch.Tensor]: + """Return a factory for isotropic B-factors (n_atoms,) in [10, 60) Ų.""" + + def _generate(n_atoms: int = 10, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + return torch.tensor(np.random.rand(n_atoms) * 50 + 10, dtype=dtypes.float) + + return _generate + + +@pytest.fixture +def random_occupancies() -> Callable[..., torch.Tensor]: + """Return a factory for dimensionless occupancies (n_atoms,) in [0.5, 1).""" + + def _generate(n_atoms: int = 10, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + return torch.tensor(np.random.rand(n_atoms) * 0.5 + 0.5, dtype=dtypes.float) + + return _generate + + +@pytest.fixture +def mock_cell() -> torch.Tensor: + """Return an orthorhombic cell (6,), lengths in Å and angles in degrees.""" + return torch.tensor([50.0, 60.0, 70.0, 90.0, 90.0, 90.0], dtype=dtypes.float) + + +@pytest.fixture +def mock_cell_triclinic() -> torch.Tensor: + """Return a triclinic cell (6,), lengths in Å and angles in degrees.""" + return torch.tensor([40.0, 50.0, 60.0, 70.0, 80.0, 85.0], dtype=dtypes.float) + + +@pytest.fixture +def mock_hkl_indices() -> Callable[..., torch.Tensor]: + """Return a factory for floating HKL triples (n_kept, 3), excluding the origin. + + The output uses ``dtypes.float``; ``n_kept`` can be less than the requested + reflection count when the origin is sampled. + """ + + def _generate( + n_reflections: int = 100, max_index: int = 10, seed: int = 42 + ) -> torch.Tensor: + np.random.seed(seed) + h = np.random.randint(-max_index, max_index + 1, n_reflections) + k = np.random.randint(-max_index, max_index + 1, n_reflections) + l = np.random.randint(-max_index, max_index + 1, n_reflections) + # Exclude (0,0,0) + mask = ~((h == 0) & (k == 0) & (l == 0)) + h, k, l = h[mask], k[mask], l[mask] + return torch.tensor(np.stack([h, k, l], axis=1), dtype=dtypes.float) + + return _generate + + +@pytest.fixture +def mock_structure_factors() -> Callable[..., torch.Tensor]: + """Return a factory for complex structure factors (n_reflections,) in electrons.""" + + def _generate(n_reflections: int = 100, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + real = np.random.randn(n_reflections) * 100 + imag = np.random.randn(n_reflections) * 100 + return torch.tensor(real + 1j * imag, dtype=dtypes.complex) + + return _generate + + +@pytest.fixture +def mock_F_obs() -> Callable[..., torch.Tensor]: + """Return a factory for observed amplitudes (n_reflections,) in electrons.""" + + def _generate(n_reflections: int = 100, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + # Positive values with realistic distribution + return torch.tensor( + np.abs(np.random.randn(n_reflections) * 100) + 10, dtype=dtypes.float + ) + + return _generate + + +@pytest.fixture +def mock_F_sigma() -> Callable[..., torch.Tensor]: + """Return a factory for amplitude uncertainties (n_reflections,) in electrons.""" + + def _generate(n_reflections: int = 100, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + return torch.tensor( + np.abs(np.random.randn(n_reflections) * 5) + 1, dtype=dtypes.float + ) + + return _generate + + +@pytest.fixture +def mock_aniso_u() -> Callable[..., torch.Tensor]: + """Return a factory for Cartesian U tensors (n_atoms, 6) in Ų. + + Components are ordered U11, U22, U33, U12, U13, U23. + """ + + def _generate(n_atoms: int = 10, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + # Diagonal elements (positive) + u11 = np.random.rand(n_atoms) * 0.05 + 0.02 + u22 = np.random.rand(n_atoms) * 0.05 + 0.02 + u33 = np.random.rand(n_atoms) * 0.05 + 0.02 + # Off-diagonal elements (can be negative, smaller magnitude) + u12 = (np.random.rand(n_atoms) - 0.5) * 0.02 + u13 = (np.random.rand(n_atoms) - 0.5) * 0.02 + u23 = (np.random.rand(n_atoms) - 0.5) * 0.02 + return torch.tensor( + np.stack([u11, u22, u33, u12, u13, u23], axis=1), dtype=dtypes.float + ) + + return _generate + + +@pytest.fixture +def mock_scattering_factors() -> Callable[..., torch.Tensor]: + """Return a factory for scattering factors (n_reflections, n_atoms) in electrons.""" + + def _generate( + n_reflections: int = 100, n_atoms: int = 10, seed: int = 42 + ) -> torch.Tensor: + np.random.seed(seed) + # Decreasing with resolution (approximate) + return torch.tensor( + np.random.rand(n_reflections, n_atoms) * 5 + 1, dtype=dtypes.float + ) + + return _generate + + +@pytest.fixture +def mock_weights() -> Callable[..., torch.Tensor]: + """Return a factory for dimensionless weights (n_atoms, 1) summing to one.""" + + def _generate(n_atoms: int = 10, seed: int = 42) -> torch.Tensor: + np.random.seed(seed) + weights = np.random.rand(n_atoms) + return torch.tensor(weights / weights.sum(), dtype=dtypes.float).reshape(-1, 1) + + return _generate diff --git a/tests/fixtures/objects.py b/tests/fixtures/objects.py new file mode 100644 index 00000000..3eb7628e --- /dev/null +++ b/tests/fixtures/objects.py @@ -0,0 +1,152 @@ +"""Load fresh mutable models, reflection data, scalers, and restraints. + +Function-scoped fixtures isolate test mutations. The explicitly shared device +bundle caches one model per device and must be treated as read-only by callers. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import pytest +import torch + +if TYPE_CHECKING: + from torchref.io import ReflectionData + from torchref.model import Model, ModelFT + from torchref.scaling import Scaler + + +@pytest.fixture +def compatibility_model(compatibility_structure_pair: dict) -> Model: + """Load a fresh model for one slow compatibility case.""" + from torchref.model import Model + + path = compatibility_structure_pair["model"] + assert path.is_file() + return Model(verbose=0).load_cif(str(path)) + + +@pytest.fixture +def compatibility_model_and_data( + compatibility_model: Model, compatibility_structure_pair: dict +) -> dict: + """Load observations only for the single crystal used by the current pipeline case.""" + from torchref.io import ReflectionData + + path = compatibility_structure_pair["reflections"] + assert path.is_file() + return { + "model": compatibility_model, + "data": ReflectionData(verbose=0).load_mtz(str(path)), + } + + +@pytest.fixture +def loaded_model(sample_cif_file: Path) -> Model: + """Load a fresh mutable Model from the sample CIF file.""" + from torchref.model.model import Model + + model = Model() + model.load_cif(str(sample_cif_file)) + return model + + +@pytest.fixture +def loaded_model_ft(sample_cif_file: Path) -> ModelFT: + """Load a fresh mutable Fourier model with a 2 Å resolution limit.""" + from torchref.model import ModelFT + + return ModelFT(max_res=2.0, verbose=0).load_cif(str(sample_cif_file)) + + +@pytest.fixture +def loaded_reflection_data(sample_mtz_file: Path) -> ReflectionData: + """Load fresh mutable reflection data from the sample MTZ file.""" + from torchref.io import ReflectionData + + data = ReflectionData() + data.load_mtz(str(sample_mtz_file)) + return data + + +@pytest.fixture +def model_and_data(sample_structure_pair: dict[str, Path]) -> dict[str, Any]: + """Load a fresh matching model and reflection dataset.""" + from torchref.io import ReflectionData + from torchref.model.model import Model + + model = Model() + model.load_cif(str(sample_structure_pair["model"])) + + data = ReflectionData() + data.load_mtz(str(sample_structure_pair["reflections"])) + + return {"model": model, "data": data} + + +@pytest.fixture +def model_with_symmetry(loaded_model: Model) -> dict[str, Any]: + """Pair a fresh model with initialized symmetry.""" + from torchref.symmetry import SpaceGroup + + sg = SpaceGroup(loaded_model.spacegroup) + return {"model": loaded_model, "symmetry": sg} + + +@pytest.fixture +def initialized_scaler(model_and_data: dict[str, Any]) -> Scaler: + """Build a scaler around a fresh matching model and dataset.""" + from torchref.scaling.scaler import Scaler + + model = model_and_data["model"] + data = model_and_data["data"] + + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) + return scaler + + +@pytest.fixture +def model_with_restraints(loaded_model: Model) -> dict[str, Any]: + """Build restraints around a fresh model.""" + from torchref.topology.restraints import Restraints + + restraints = Restraints( + pdb=loaded_model.pdb, + xyz_fn=loaded_model.xyz, + vdw_radii_fn=loaded_model.get_vdw_radii, + verbose=0, + ) + restraints.build_restraints() + return {"model": loaded_model, "restraints": restraints} + + +@pytest.fixture(scope="session") +def _device_model_cache() -> dict: + """``{device_str: ModelFT}`` built at most once per device, per session.""" + return {} + + +@pytest.fixture +def device_model_bundle( + _device_model_cache: dict[str, ModelFT], pdb_dir: Path, any_device: torch.device +) -> dict[str, ModelFT]: + """Borrow a session-shared model on the requested device. + + Notes + ----- + Treat the model as read-only, including when a target borrows it. Moving a + target can move its model too; tests of movement need a fresh model. + """ + key = str(any_device) + if key not in _device_model_cache: + pdb = pdb_dir / "1DAW.pdb" + if not pdb.exists(): + pytest.skip("1DAW.pdb fixture not present") + from torchref.model import ModelFT + + _device_model_cache[key] = ModelFT(device=any_device, verbose=0).load_pdb( + str(pdb) + ) + return {"model": _device_model_cache[key]} diff --git a/tests/fixtures/paths.py b/tests/fixtures/paths.py new file mode 100644 index 00000000..16ac98d4 --- /dev/null +++ b/tests/fixtures/paths.py @@ -0,0 +1,68 @@ +"""Locate bundled test data and optional monomer-library installations.""" + +from pathlib import Path + +import pytest + + +@pytest.fixture(scope="session") +def tests_root() -> Path: + """Return the root of the test tree.""" + return Path(__file__).resolve().parents[1] + + +@pytest.fixture(scope="session") +def project_root() -> Path: + """Return the project root.""" + return Path(__file__).resolve().parents[2] + + +@pytest.fixture(scope="session") +def test_files_dir(tests_root: Path) -> Path: + """Return the bundled test-data directory.""" + return tests_root / "files" + + +@pytest.fixture(scope="session") +def cif_dir(test_files_dir: Path) -> Path: + """Return the model CIF directory.""" + return test_files_dir / "cif" + + +@pytest.fixture(scope="session") +def cif_sf_dir(test_files_dir: Path) -> Path: + """Return the structure-factor CIF directory.""" + return test_files_dir / "cif_sf" + + +@pytest.fixture(scope="session") +def mtz_dir(test_files_dir: Path) -> Path: + """Return the MTZ reflection directory.""" + return test_files_dir / "mtz" + + +@pytest.fixture(scope="session") +def pdb_dir(test_files_dir: Path) -> Path: + """Return the model PDB directory.""" + return test_files_dir / "pdb" + + +@pytest.fixture(scope="session") +def external_monomer_library(project_root: Path) -> Path: + """Return the optional external monomer-library path without checking it.""" + return project_root / "external_monomer_library" + + +@pytest.fixture(scope="session") +def monomer_library_path(project_root: Path) -> str: + """Get path to the monomer library as a string. + + Returns + ------- + str + Absolute path to the external_monomer_library directory. + """ + lib_path = project_root / "external_monomer_library" + if not lib_path.exists(): + pytest.skip("Monomer library not found") + return str(lib_path) diff --git a/tests/fixtures/precision.py b/tests/fixtures/precision.py new file mode 100644 index 00000000..6769c4de --- /dev/null +++ b/tests/fixtures/precision.py @@ -0,0 +1,51 @@ +"""Scope numerical reference configuration and expose comparison tolerances.""" + +from collections.abc import Iterator +from contextlib import contextmanager + +import pytest +import torch + +import torchref +from torchref.config import device, dtypes + + +@contextmanager +def cpu_double_precision() -> Iterator[None]: + """Temporarily select CPU float64/complex128 for numerical references. + + Notes + ----- + Mutate process-wide TorchRef defaults, not PyTorch factory defaults. Restore + float/complex dtype, device, and density cutoff even when the body raises. + Objects allocated inside the context retain their own dtype and device. + """ + original = dtypes.float, dtypes.complex, device.current + cutoff = torchref.sigma_cutoff_ed.value + dtypes.float = torch.float64 + dtypes.complex = torch.complex128 + device.current = torch.device("cpu") + try: + yield + finally: + dtypes.float, dtypes.complex, device.current = original + torchref.sigma_cutoff_ed.value = cutoff + + +@pytest.fixture +def double_cpu() -> Iterator[None]: + """Use CPU double precision for one test and restore configuration afterward.""" + with cpu_double_precision(): + yield + + +@pytest.fixture +def rtol() -> float: + """Relative tolerance for floating point comparisons.""" + return 1e-5 + + +@pytest.fixture +def atol() -> float: + """Absolute tolerance for floating point comparisons.""" + return 1e-8 diff --git a/tests/functional/af_trajectory_reference.json b/tests/functional/af_trajectory_reference.json new file mode 100644 index 00000000..1cc06d09 --- /dev/null +++ b/tests/functional/af_trajectory_reference.json @@ -0,0 +1,116 @@ +{ + "cycles": 2, + "spread_runs": 5, + "structures": { + "1VER": { + "trajectory": [ + [ + 0.378484, + 0.332842 + ], + [ + 0.336894, + 0.317116 + ], + [ + 0.330538, + 0.308256 + ], + [ + 0.312144, + 0.305218 + ] + ], + "spread": 7.1e-05, + "tolerance": 0.002 + }, + "6VHI": { + "trajectory": [ + [ + 0.332963, + 0.344154 + ], + [ + 0.287595, + 0.332427 + ], + [ + 0.282997, + 0.309342 + ], + [ + 0.269246, + 0.324271 + ] + ], + "spread": 0.000567, + "tolerance": 0.002 + }, + "1BYW": { + "trajectory": [ + [ + 0.396823, + 0.352145 + ], + [ + 0.356685, + 0.322779 + ], + [ + 0.359504, + 0.321488 + ], + [ + 0.336396, + 0.309443 + ] + ], + "spread": 0.000107, + "tolerance": 0.002 + }, + "6JZA": { + "trajectory": [ + [ + 0.416852, + 0.378224 + ], + [ + 0.36559, + 0.35986 + ], + [ + 0.366328, + 0.363104 + ], + [ + 0.329881, + 0.325672 + ] + ], + "spread": 0.004474, + "tolerance": 0.013422 + }, + "6SXW": { + "trajectory": [ + [ + 0.430794, + 0.4249 + ], + [ + 0.383688, + 0.430362 + ], + [ + 0.387311, + 0.431517 + ], + [ + 0.340972, + 0.407762 + ] + ], + "spread": 9.8e-05, + "tolerance": 0.002 + } + } +} diff --git a/tests/functional/conftest.py b/tests/functional/conftest.py index 58e2352e..ba6caf3b 100644 --- a/tests/functional/conftest.py +++ b/tests/functional/conftest.py @@ -1,6 +1,3 @@ -""" -Functional test fixtures. +"""Expose module-shared read-only fixtures to functional tests.""" -All shared fixtures (sample files, loaded models, scalers, restraints, etc.) -are defined in the root tests/conftest.py and are automatically available here. -""" +from tests.fixtures.functional import shared_model_ft # noqa: F401 diff --git a/tests/functional/test_af_trajectory.py b/tests/functional/test_af_trajectory.py new file mode 100644 index 00000000..b242acb6 --- /dev/null +++ b/tests/functional/test_af_trajectory.py @@ -0,0 +1,225 @@ +"""Refinement trajectories from AlphaFold starts, against a committed reference. + +An AlphaFold model is far from convergence, so an early-cycle change in restraint +weighting or connectivity shows up in the trajectory rather than being washed out by the +endpoint. That makes these five structures a sharper probe of a topology or restraint +change than a converged-structure comparison would be. + +TorchRef refinement is not reproducible run to run, so a deviation from the reference +means nothing on its own. Each structure therefore carries its **own** tolerance, +measured over :data:`SPREAD_RUNS` independent runs when the reference was written, and +committed alongside it. Sizing the tolerance from a couple of runs at test time does +not work: 6JZA is bimodal at this cycle count -- its trajectories land in one of two +basins about 0.0074 apart -- and two runs that happen to pick the same basin report a +spread 140 times too small. + +R-work and R-free are held to different bounds. R-work is reproducible to a few parts in +ten thousand, so it keeps the measured tolerance and is what catches a change that +actually moves the refinement. R-free is computed on the small free set and has a rare +second basin of its own, a few thousandths wide, that a handful of runs will usually +miss; it therefore carries :data:`RFREE_TOLERANCE_FLOOR`, wide enough to sit outside +that basin and still far inside the descent the trajectory shows. + +Regenerate deliberately, after a change meant to move these numbers:: + + ./.dev/bin/python tests/functional/test_af_trajectory.py +""" + +import json +from pathlib import Path + +import numpy as np +import pytest +import torch + +#: AlphaFold-start structures, chosen so all five actually descend under refinement and +#: the space groups span centred tetragonal, centred monoclinic, hexagonal and trigonal. +CODES = ["1VER", "6VHI", "1BYW", "6JZA", "6SXW"] + +#: Macro-cycles per trajectory. Two is enough for the descent to be visible while +#: keeping the whole test at a few seconds per structure. +CYCLES = 2 + +#: Independent runs used to size each structure's tolerance at regeneration time. Enough +#: to see a second basin if there is one; two is not. +SPREAD_RUNS = 5 + +#: Multiple of the measured spread a deviation may reach before it counts as a change. +SPREAD_MULTIPLE = 3.0 + +#: Tolerance floor for R-work, so a structure whose runs agree very closely is not held +#: to an unreasonably tight bound. +TOLERANCE_FLOOR = 0.002 + +#: Tolerance floor for R-free, which moves in discrete basins rather than jitter. Set +#: above the widest basin separation seen on this set and kept well inside every +#: structure's own R-work descent, so a change large enough to matter still fails. +#: ``test_tolerances_are_tight_enough_to_detect_something`` enforces the second half. +RFREE_TOLERANCE_FLOOR = 0.01 + +REFERENCE = Path(__file__).with_name("af_trajectory_reference.json") + + +def trajectory(pdb_path, mtz_path, cycles=CYCLES): + """Per-stage ``(r_work, r_free)`` series of a short refinement. + + Returns + ------- + list of tuple of float + Two entries per macro-cycle -- after scaling and after refinement. + """ + from torchref.refinement.lbfgs_refinement import LBFGSRefinement + + refinement = LBFGSRefinement( + data_file=str(mtz_path), + pdb=str(pdb_path), + verbose=0, + device=torch.device("cpu"), + ) + history = refinement.refine_everything(macro_cycles=cycles) + key = next(k for k in history if k.startswith("refinement_everything")) + + series = [] + for cycle in history[key]: + for stage in ("after_scaling", "after_refinement"): + metrics = cycle.get(stage) or {} + series.append((float(metrics["rwork"]), float(metrics["rfree"]))) + return series + + +def _max_deviation(a, b): + """Largest absolute difference between two trajectories, over both R-factors.""" + return max(max(abs(x[0] - y[0]), abs(x[1] - y[1])) for x, y in zip(a, b)) + + +def _deviations(a, b): + """Largest absolute difference between two trajectories, ``(R-work, R-free)``.""" + return ( + max(abs(x[0] - y[0]) for x, y in zip(a, b)), + max(abs(x[1] - y[1]) for x, y in zip(a, b)), + ) + + +def _rfree_tolerance(entry): + """The R-free bound: this structure's measured tolerance, floored.""" + return max(float(entry["tolerance"]), RFREE_TOLERANCE_FLOOR) + + +@pytest.fixture(scope="module") +def reference(): + """The committed reference trajectories and their tolerances.""" + if not REFERENCE.exists(): + pytest.fail( + f"{REFERENCE.name} is missing. Regenerate it with " + f"`./.dev/bin/python {Path(__file__).name}`." + ) + data = json.loads(REFERENCE.read_text()) + assert ( + data["cycles"] == CYCLES + ), f"reference was written for {data['cycles']} cycles, test runs {CYCLES}" + return data + + +@pytest.mark.integration +@pytest.mark.parametrize("code", CODES) +def test_af_trajectory_matches_reference(code, reference, test_files_dir): + """The trajectory sits inside this structure's own measured run-to-run spread.""" + pdb_path = test_files_dir / "pdb" / f"{code}_af.pdb" + mtz_path = test_files_dir / "mtz" / f"{code}.mtz" + assert pdb_path.exists(), f"{pdb_path.name} is not bundled" + assert mtz_path.exists(), f"{mtz_path.name} is not bundled" + + entry = reference["structures"][code] + expected = [tuple(point) for point in entry["trajectory"]] + rwork_tolerance = float(entry["tolerance"]) + rfree_tolerance = _rfree_tolerance(entry) + + observed = trajectory(pdb_path, mtz_path) + assert len(observed) == len( + expected + ), f"{code}: trajectory has {len(observed)} stages, reference has {len(expected)}" + + dev_work, dev_free = _deviations(observed, expected) + for label, deviation, tolerance in ( + ("R-work", dev_work, rwork_tolerance), + ("R-free", dev_free, rfree_tolerance), + ): + assert deviation <= tolerance, ( + f"{code}: {label} deviates from the reference by {deviation:.6f}, above its " + f"tolerance of {tolerance:.6f}.\nobserved={observed}\n" + f"reference={expected}" + ) + + +@pytest.mark.integration +def test_af_starts_actually_descend(reference): + """Each reference trajectory improves R-work, so a regression has signal to lose. + + A structure that barely moves under refinement cannot show a trajectory regression, + so the set is only useful while every member descends. + """ + for code in CODES: + series = reference["structures"][code]["trajectory"] + first, last = series[0][0], series[-1][0] + assert last < first - 0.02, ( + f"{code} only moves R-work from {first:.4f} to {last:.4f}; it is too flat " + f"to serve as a trajectory probe" + ) + + +@pytest.mark.integration +def test_tolerances_are_tight_enough_to_detect_something(reference): + """A tolerance so wide it would accept any change is not a test. + + Guards against a future regeneration quietly widening a bound until the structure + stops constraining anything. The bound has to stay well inside the descent the + trajectory itself shows. + """ + for code in CODES: + entry = reference["structures"][code] + series = entry["trajectory"] + descent = series[0][0] - series[-1][0] + # Checks the widest bound actually applied, not the stored one: flooring R-free + # would otherwise widen the real bound without this guard seeing it. + widest = max(float(entry["tolerance"]), _rfree_tolerance(entry)) + assert widest < descent / 4.0, ( + f"{code}: tolerance {widest:.4f} is not small against its " + f"own R-work descent of {descent:.4f}" + ) + + +def _write_reference(): + """Regenerate the reference, sizing each tolerance from independent runs.""" + root = Path(__file__).resolve().parents[1] / "files" + structures = {} + for code in CODES: + pdb_path = root / "pdb" / f"{code}_af.pdb" + mtz_path = root / "mtz" / f"{code}.mtz" + + runs = [trajectory(pdb_path, mtz_path) for _ in range(SPREAD_RUNS)] + mean = [ + tuple(float(np.mean([run[stage][i] for run in runs])) for i in (0, 1)) + for stage in range(len(runs[0])) + ] + spread = max(_max_deviation(run, mean) for run in runs) + tolerance = max(TOLERANCE_FLOOR, SPREAD_MULTIPLE * spread) + + structures[code] = { + "trajectory": [[round(w, 6), round(f, 6)] for w, f in mean], + "spread": round(spread, 6), + "tolerance": round(tolerance, 6), + } + print(f"{code}: spread={spread:.6f} tolerance={tolerance:.6f}", flush=True) + + REFERENCE.write_text( + json.dumps( + {"cycles": CYCLES, "spread_runs": SPREAD_RUNS, "structures": structures}, + indent=2, + ) + + "\n" + ) + print(f"wrote {REFERENCE}") + + +if __name__ == "__main__": + _write_reference() diff --git a/tests/functional/test_io_functional.py b/tests/functional/test_io_functional.py deleted file mode 100644 index 9cdef611..00000000 --- a/tests/functional/test_io_functional.py +++ /dev/null @@ -1,311 +0,0 @@ -""" -Functional tests for I/O operations. - -Tests file loading and data processing with real crystallographic data. -""" - -import pytest -import torch -import numpy as np - - -class TestCIFReadingFunctional: - """Functional tests for CIF file reading.""" - - @pytest.mark.integration - def test_load_multiple_cif_files(self, cif_dir): - """Test loading multiple CIF files successfully.""" - from torchref.model.model import Model - - cif_files = list(cif_dir.glob("*.cif")) - assert len(cif_files) > 0, "No CIF files found in test directory" - - for cif_file in cif_files: - model = Model() - model.load_cif(str(cif_file)) - - # Each file should load with atoms - n_atoms = model.xyz().shape[0] - assert n_atoms > 0, f"No atoms loaded from {cif_file}" - - # Should have cell parameters - assert model.cell is not None - assert len(model.cell) == 6 - - @pytest.mark.integration - def test_cif_atom_properties(self, sample_cif_file): - """Test that atom properties are correctly loaded from CIF.""" - from torchref.model.model import Model - - model = Model() - model.load_cif(str(sample_cif_file)) - - pdb = model.pdb - - # Check required columns exist - required_cols = ['x', 'y', 'z', 'element', 'resname', 'chainid', 'resseq'] - for col in required_cols: - assert col in pdb.columns or col.upper() in pdb.columns, f"Missing column: {col}" - - @pytest.mark.integration - def test_cif_element_types(self, sample_cif_file): - """Test that element types are properly assigned.""" - from torchref.model.model import Model - - model = Model() - model.load_cif(str(sample_cif_file)) - - elements = model.pdb['element'].unique() - - # Should have common protein elements - common_elements = ['C', 'N', 'O', 'S'] - found_any = any(elem in elements for elem in common_elements) - assert found_any, "No common elements found" - - -class TestMTZReadingFunctional: - """Functional tests for MTZ file reading.""" - - @pytest.mark.integration - def test_load_multiple_mtz_files(self, mtz_dir): - """Test loading multiple MTZ files successfully.""" - from torchref.io import ReflectionData - - mtz_files = list(mtz_dir.glob("*.mtz")) - assert len(mtz_files) > 0, "No MTZ files found in test directory" - - for mtz_file in mtz_files: - data = ReflectionData() - data.load_mtz(str(mtz_file)) - - # Each file should load with reflections - n_refl = data.hkl.shape[0] - assert n_refl > 0, f"No reflections loaded from {mtz_file}" - - # Should have cell parameters - assert data.cell is not None - - @pytest.mark.integration - def test_mtz_data_properties(self, sample_mtz_file): - """Test that MTZ data properties are correctly loaded.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Check HKL indices are integers or can be converted - hkl = data.hkl - assert hkl.shape[1] == 3, "HKL should have 3 columns" - - # Check F values are loaded - assert data.F is not None - assert data.F.shape[0] == hkl.shape[0] - - # Check sigma values - if hasattr(data, 'F_sigma') and data.F_sigma is not None: - assert data.F_sigma.shape[0] == hkl.shape[0] - - @pytest.mark.integration - def test_mtz_resolution_range(self, sample_mtz_file): - """Test that resolution range is computed correctly.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Check if resolution data is available - if hasattr(data, 'd') and data.d is not None: - d_min = data.d.min().item() - d_max = data.d.max().item() - - # Resolution should be positive - assert d_min > 0 - assert d_max > d_min - - # Typical protein data: 0.8 - 500 Å - assert d_min > 0.5 - assert d_max < 1000 - - -class TestSFCIFReadingFunctional: - """Functional tests for structure factor CIF reading.""" - - @pytest.mark.integration - def test_load_sf_cif(self, cif_sf_dir): - """Test loading structure factor CIF files.""" - from torchref.io import ReflectionData - - sf_files = list(cif_sf_dir.glob("*.cif")) - if not sf_files: - pytest.skip("No SF-CIF files found") - - for sf_file in sf_files: - data = ReflectionData() - try: - data.load_cif(str(sf_file)) - - # Should have loaded reflections - if data.hkl is not None: - assert data.hkl.shape[0] > 0 - except Exception as e: - # Some files may not be valid SF-CIF format - pass - - -class TestDataConsistencyFunctional: - """Test consistency between model and data files.""" - - @pytest.mark.integration - def test_cell_parameters_match(self, sample_structure_pair): - """Test that cell parameters match between model and reflections.""" - from torchref.model.model import Model - from torchref.io import ReflectionData - - model = Model() - model.load_cif(str(sample_structure_pair["model"])) - - data = ReflectionData() - data.load_mtz(str(sample_structure_pair["reflections"])) - - model_cell = model.cell - data_cell = data.cell - - if model_cell is not None and data_cell is not None: - # Convert to tensors if needed - if not isinstance(model_cell, torch.Tensor): - model_cell = torch.tensor(model_cell) - if not isinstance(data_cell, torch.Tensor): - data_cell = torch.tensor(data_cell) - - # Cell parameters should be similar (1% tolerance) - assert torch.allclose(model_cell.float(), data_cell.float(), rtol=0.01, atol=0.1) - - @pytest.mark.integration - def test_spacegroup_consistency(self, sample_structure_pair): - """Test that spacegroup is consistent.""" - from torchref.model.model import Model - from torchref.io import ReflectionData - - model = Model() - model.load_cif(str(sample_structure_pair["model"])) - - data = ReflectionData() - data.load_mtz(str(sample_structure_pair["reflections"])) - - # Both should have spacegroup defined - assert model.spacegroup is not None - - -class TestDataBinningFunctional: - """Test data binning operations.""" - - @pytest.mark.integration - def test_get_bins(self, sample_mtz_file): - """Test resolution binning of reflection data.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Get bins - bins, n_bins = data.get_bins(n_bins=10) - - assert bins is not None - assert bins.shape[0] == data.hkl.shape[0] - assert bins.min() >= 0 - assert bins.max() < n_bins - - @pytest.mark.integration - def test_mean_res_per_bin(self, sample_mtz_file): - """Test mean resolution per bin calculation.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Get bins first - bins, n_bins = data.get_bins(n_bins=10) - - # Get mean resolution per bin - if hasattr(data, 'mean_res_per_bin'): - mean_res = data.mean_res_per_bin() - - assert mean_res is not None - assert len(mean_res) == n_bins - - # Mean resolution should decrease with bin index (low res to high res) - # or increase (high res to low res) - depends on implementation - assert torch.all(torch.isfinite(mean_res)) - - -class TestFrenchWilsonFunctional: - """Test French-Wilson conversion with real data.""" - - @pytest.mark.integration - def test_french_wilson_applied(self, sample_mtz_file): - """Test that French-Wilson conversion is applied.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # After French-Wilson, F values should be non-negative - valid_F = data.F[~torch.isnan(data.F)] - - if len(valid_F) > 0: - # All valid F values should be >= 0 - assert torch.all(valid_F >= 0) - - -class TestRfreeHandlingFunctional: - """Test R-free flag handling.""" - - @pytest.mark.integration - def test_rfree_flags_loaded(self, sample_mtz_file): - """Test that R-free flags are loaded or generated.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Should have rfree attribute - if hasattr(data, 'rfree') and data.rfree is not None: - assert data.rfree.shape[0] == data.hkl.shape[0] - - # Should be boolean or can be converted to boolean - assert data.rfree.dtype == torch.bool or torch.all((data.rfree == 0) | (data.rfree == 1)) - - @pytest.mark.integration - def test_rfree_fraction(self, sample_mtz_file): - """Test R-free set fraction is reasonable.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - if hasattr(data, 'rfree') and data.rfree is not None: - # Work set mask (True for work, False for test) - work_fraction = data.rfree.float().mean().item() - - # Typically 90-95% work set, 5-10% test set - # So work_fraction should be 0.9-0.95 typically - assert 0.7 < work_fraction <= 1.0 - - -class TestMaskHandlingFunctional: - """Test reflection mask handling.""" - - @pytest.mark.integration - def test_masks_method(self, sample_mtz_file): - """Test masks() method returns valid mask.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - if hasattr(data, 'masks'): - mask = data.masks() - - assert mask is not None - assert mask.shape[0] == data.hkl.shape[0] - assert mask.dtype == torch.bool diff --git a/tests/functional/test_loss_weighting_functional.py b/tests/functional/test_loss_weighting_functional.py deleted file mode 100644 index f1bcb602..00000000 --- a/tests/functional/test_loss_weighting_functional.py +++ /dev/null @@ -1,198 +0,0 @@ -""" -Functional tests for loss weighting module. - -These tests exercise the loss weighting strategies with realistic data. -Updated to use the new component_weighting and LossState architecture. -""" -import pytest -import torch -import numpy as np -from unittest.mock import Mock - - -@pytest.mark.integration -class TestLossStateWeightingFunctional: - """Test LossState weighting functionality.""" - - def test_loss_state_add_and_get_weights(self): - """Test adding and getting weights from LossState.""" - from torchref.refinement.loss_state import LossState - - state = LossState() - state.set_weight('xray', 1.5) - state.set_weight('geometry', 0.7) - - assert state.get_weight('xray') == 1.5 - assert state.get_weight('geometry') == 0.7 - - def test_loss_state_hierarchical_weights(self): - """Test hierarchical weights in LossState.""" - from torchref.refinement.loss_state import LossState - - state = LossState() - state.set_weight('geometry', 2.0) - state.set_weight('geometry/bond', 3.0) - - # Effective weight should be product: 2.0 * 3.0 = 6.0 - effective = state.get_effective_weight('geometry/bond') - assert effective == 6.0 - - -@pytest.mark.integration -class TestWeightingMathOperations: - """Test mathematical operations with weights.""" - - def test_total_weighted_loss_from_state(self): - """Test computing total weighted loss from LossState via aggregate.""" - from torchref.refinement.loss_state import LossState - - state = LossState() - state.register_target('xray', lambda: torch.tensor(10.0)) - state.register_target('geometry', lambda: torch.tensor(5.0)) - state.register_target('adp', lambda: torch.tensor(2.0)) - - state.set_weight('xray', 1.0) - state.set_weight('geometry', 0.5) - state.set_weight('adp', 0.25) - - total = state.aggregate() - - # Expected: 10*1.0 + 5*0.5 + 2*0.25 = 10 + 2.5 + 0.5 = 13.0 - assert torch.isclose(total, torch.tensor(13.0)) - - -@pytest.mark.integration -class TestNLLXrayFunction: - """Test the NLL X-ray function used in weighting.""" - - def test_nll_xray_basic(self): - """Test basic NLL X-ray calculation.""" - from torchref.base.math_torch import nll_xray - - fobs = torch.tensor([100.0, 200.0, 300.0], dtype=torch.float32) - fcalc = torch.tensor([105.0, 195.0, 305.0], dtype=torch.float32) - sigma = torch.tensor([10.0, 15.0, 20.0], dtype=torch.float32) - - nll = nll_xray(fobs, fcalc, sigma) - - # nll returns per-reflection values - assert torch.all(torch.isfinite(nll)) - - def test_nll_decreases_with_better_fit(self): - """Test that NLL decreases as fit improves.""" - from torchref.base.math_torch import nll_xray - - fobs = torch.tensor([100.0], dtype=torch.float32) - sigma = torch.tensor([10.0], dtype=torch.float32) - - # Good fit - fcalc_good = torch.tensor([100.0], dtype=torch.float32) - nll_good = nll_xray(fobs, fcalc_good, sigma) - - # Bad fit - fcalc_bad = torch.tensor([150.0], dtype=torch.float32) - nll_bad = nll_xray(fobs, fcalc_bad, sigma) - - # Good fit should have lower NLL - assert nll_good < nll_bad - - -@pytest.mark.integration -class TestGradnormUtility: - """Test the gradnorm utility function.""" - - def test_gradnorm_basic(self): - """Test basic gradnorm calculation.""" - from torchref.utils.gradnorm import gradnorm - - # Create simple parameter - param = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) - - # Create loss - loss = param.sum() - - # Compute gradient norm - norm = gradnorm(loss, [param]) - - assert torch.isfinite(norm) - assert norm > 0 - - def test_gradnorm_with_multiple_params(self): - """Test gradnorm with multiple parameters.""" - from torchref.utils.gradnorm import gradnorm - - param1 = torch.tensor([1.0, 2.0], requires_grad=True) - param2 = torch.tensor([3.0, 4.0], requires_grad=True) - - loss = param1.sum() + param2.sum() - - norm = gradnorm(loss, [param1, param2]) - - assert torch.isfinite(norm) - assert norm > 0 - - -@pytest.mark.integration -class TestWeightingEdgeCases: - """Test edge cases in weighting.""" - - def test_zero_weight(self): - """Test zero weight (disabling a loss term).""" - from torchref.refinement.loss_state import LossState - - state = LossState() - state.register_target('adp', lambda: torch.tensor(100.0)) - state.set_weight('adp', 0.0) - - # Zero weight should effectively disable ADP term - total = state.aggregate() - assert torch.isclose(total, torch.tensor(0.0)) - - -@pytest.mark.integration -class TestLossAggregatorFunctional: - """Test LossAggregator functionality.""" - - def test_aggregator_basic(self): - """Test basic aggregator functionality (LossState.aggregate).""" - from torchref.refinement.loss_state import LossState - - state = LossState() - state.register_target('xray', lambda: torch.tensor(2.0)) - state.register_target('bond', lambda: torch.tensor(1.0)) - state.set_weight('xray', 1.0) - state.set_weight('bond', 0.5) - - total = state.aggregate() - - # Expected: 2.0 * 1.0 + 1.0 * 0.5 = 2.5 - expected = torch.tensor(2.5) - assert torch.isclose(total, expected) - - def test_loss_state_caches_losses(self): - """Test that LossState caches computed losses.""" - from torchref.refinement.loss_state import LossState - - call_count = [0] - def counting_target(): - call_count[0] += 1 - return torch.tensor(2.0) - - state = LossState() - state.register_target('xray', counting_target) - state.set_weight('xray', 1.0) - - # register_target probes the target once to walk the autograd graph; - # reset the counter so we measure only aggregate() invocations. - call_count[0] = 0 - - # First aggregation computes the loss - total1 = state.aggregate() - assert call_count[0] == 1 - - # Get cached loss doesn't recompute - cached = state.get_loss('xray') - assert cached is not None - assert torch.isclose(cached, torch.tensor(2.0)) - - diff --git a/tests/functional/test_model_ft_functional.py b/tests/functional/test_model_ft_functional.py index 3a430cc1..ce9b5e71 100644 --- a/tests/functional/test_model_ft_functional.py +++ b/tests/functional/test_model_ft_functional.py @@ -4,10 +4,9 @@ These tests exercise the ModelFT class with real crystallographic data, testing the FFT-based structure factor calculation pipeline. """ + import pytest import torch -import numpy as np -from pathlib import Path @pytest.mark.integration @@ -17,7 +16,7 @@ class TestModelFTInitialization: def test_modelft_empty_initialization(self): """Test empty ModelFT initialization.""" from torchref.model.model_ft import ModelFT - + model = ModelFT() assert model is not None assert model.max_res == 1.0 # Default @@ -29,282 +28,101 @@ def test_modelft_with_custom_resolution(self): model = ModelFT(max_res=1.5) assert model.max_res == 1.5 - def test_modelft_load_cif(self, sample_cif_file): - """Test loading a CIF file into ModelFT.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Verify basic properties - assert model.xyz() is not None - assert model.xyz().shape[0] > 0 - assert model.cell is not None - assert len(model.cell) == 6 - def test_modelft_has_gridsize(self, sample_cif_file): - """Test that ModelFT sets up gridsize after loading.""" + """The grid resolves from the loaded cell and space group on first read.""" from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Check gridsize is set - if model.gridsize is not None: - assert len(model.gridsize) == 3 - assert all(g > 0 for g in model.gridsize) - -@pytest.mark.integration -class TestModelFTParametrization: - """Test ModelFT parametrization with real structures.""" - - def test_parametrization_built(self, sample_cif_file): - """Test that parametrization is built after loading.""" - from torchref.model.model_ft import ModelFT - model = ModelFT(max_res=2.0, verbose=0) + assert model.gridsize is None # no crystal yet model.load_cif(str(sample_cif_file)) - - # Parametrization should be set - assert model.parametrization is not None - def test_scattering_factors_available(self, sample_cif_file): - """Test that scattering factors can be computed.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Should be able to access atom properties - xyz = model.xyz() - assert xyz is not None - assert xyz.dtype == torch.float32 or xyz.dtype == torch.float64 + assert model.gridsize is not None + assert len(model.gridsize) == 3 + assert all(g > 0 for g in model.gridsize) + assert model.xyz().shape[0] > 0 + assert model.parametrization + assert model.adp().shape == (len(model.xyz()),) + assert torch.all(model.adp() >= 0) -@pytest.mark.integration +@pytest.mark.integration class TestModelFTGridOperations: """Test ModelFT grid operations.""" - def test_setup_gridsize(self, sample_cif_file): - """Test grid size setup.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Setup gridsize - gridsize = model.setup_gridsize(max_res=2.0) - - assert gridsize is not None - assert len(gridsize) == 3 - assert all(g > 0 for g in gridsize) + def test_setup_grid(self, loaded_model_ft): + """An explicit grid size overrides the resolution-derived one.""" - def test_setup_grid(self, sample_cif_file): - """Test full grid setup.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Model should have grid setup - assert model.gridsize is not None or hasattr(model, 'map') + model = loaded_model_ft + derived = model.grid_shape + assert derived is not None and len(derived) == 3 + + model.setup_grid(gridsize=(24, 24, 24)) + assert model.grid_shape == (24, 24, 24) + assert model.explicit_gridsize == (24, 24, 24) + + model.explicit_gridsize = None + assert model.grid_shape == derived @pytest.mark.integration class TestModelFTRealSpaceMap: """Test ModelFT real space electron density map construction.""" - def test_get_real_space_grid(self, sample_cif_file): + def test_get_real_space_grid(self, loaded_model_ft): """Test getting real space grid.""" - from torchref.model.model_ft import ModelFT from torchref.base.math_torch import get_real_grid - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - if model.gridsize is not None: - # Get real space grid - grid = get_real_grid(model.cell, max_res=2.0, device='cpu') - - assert grid is not None - assert len(grid.shape) == 4 # Should be 4D (nx, ny, nz, 3) + model = loaded_model_ft -@pytest.mark.integration -class TestModelFTSymmetry: - """Test ModelFT symmetry operations.""" + assert model.gridsize is not None + grid = get_real_grid(model.cell, max_res=2.0, device=model.device) - def test_map_symmetry_available(self, sample_cif_file): - """Test map symmetry is available after loading.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Model should have spacegroup after loading - assert model.spacegroup is not None - - # Map symmetry can be created if gridsize is available - if model.gridsize is not None: - from torchref.symmetry.map_symmetry import MapSymmetry - - gridsize = tuple(model.gridsize.tolist()) - cell_params = model.cell - - map_sym = MapSymmetry(model.spacegroup, gridsize, cell_params) - assert map_sym is not None - - -@pytest.mark.integration -class TestModelFTStateDictFunctional: - """Test ModelFT state dict operations with real data.""" - - def test_save_and_load_state_dict(self, sample_cif_file, tmp_path): - """Test saving and loading state dict.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - original_xyz = model.xyz().clone() - - # Save state dict - state_dict = model.state_dict() - - # Create new model and load state - model2 = ModelFT(max_res=2.0, verbose=0) - - # We need to ensure proper initialization - # For now just verify state_dict works - assert state_dict is not None - assert len(state_dict) > 0 + assert grid is not None + assert len(grid.shape) == 4 # Should be 4D (nx, ny, nz, 3) + assert grid.device == model.xyz().device + assert grid.dtype == model.xyz().dtype @pytest.mark.integration -class TestModelFTForwardPass: - """Test ModelFT forward pass (structure factor calculation).""" - - def test_forward_method_exists(self, sample_cif_file): - """Test that forward method is available.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Check forward method exists - assert hasattr(model, 'forward') - - def test_build_map_method(self, sample_cif_file): - """Test build_map method if available.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=3.0, verbose=0) # Lower res for faster test - model.load_cif(str(sample_cif_file)) - - # Check build_map method - if hasattr(model, 'build_map'): - # Try to build map - try: - model.build_map() - assert model.map is not None - except Exception as e: - # May fail if missing dependencies - pytest.skip(f"build_map not available: {e}") - - -@pytest.mark.integration -class TestModelFTMultipleStructures: - """Test ModelFT with multiple structures.""" - - def test_modelft_multiple_structures(self, all_structure_pairs): - """Test ModelFT works with different structures.""" - from torchref.model.model_ft import ModelFT - - tested = 0 - for pair in all_structure_pairs[:3]: # Test first 3 - try: - model = ModelFT(max_res=3.0, verbose=0) - model.load_cif(str(pair["model"])) - - # Basic checks - assert model.xyz() is not None - assert model.xyz().shape[0] > 0 - - tested += 1 - except Exception as e: - # Some structures may fail to load - continue - - assert tested >= 1, "At least one structure should load" - +class TestModelFTSymmetry: + """Test ModelFT symmetry operations.""" -@pytest.mark.integration -class TestModelFTCaching: - """Test ModelFT caching mechanism.""" + def test_map_symmetry_available(self, shared_model_ft): + """Test map symmetry is available after loading.""" - def test_cache_initialization(self, sample_cif_file): - """Test that CachedForwardMixin cache starts empty.""" - from torchref.model.model_ft import ModelFT + model = shared_model_ft - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) + # Model should have spacegroup after loading + assert model.spacegroup is not None - # Mixin cache should start empty (lazily initialized) - assert getattr(model, "_fwd_cached_output", None) is None + # The map operator comes from the space group, keyed on the grid shape. + gridsize = model.grid_shape + assert gridsize is not None - def test_cache_usage(self, sample_cif_file): - """Test that cache can be used for computations.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Access xyz twice - should use caching - xyz1 = model.xyz() - xyz2 = model.xyz() - - # Should return same tensor - assert torch.allclose(xyz1, xyz2) + operator = model.spacegroup.map_operator(gridsize) + assert operator is not None + assert operator.map_shape == gridsize @pytest.mark.integration class TestModelFTCoordinateOperations: """Test ModelFT coordinate operations.""" - def test_cartesian_to_fractional(self, sample_cif_file): - """Test coordinate conversion.""" - from torchref.model.model_ft import ModelFT - from torchref.base.math_torch import cartesian_to_fractional_torch - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - xyz = model.xyz() - cell = model.cell - - # Convert to fractional - frac = cartesian_to_fractional_torch(xyz, cell.data) - - # Fractional coords should be bounded (mostly between 0 and 1) - assert frac.shape == xyz.shape - - def test_fractional_to_cartesian(self, sample_cif_file): + def test_fractional_to_cartesian(self, shared_model_ft): """Test fractional to cartesian conversion.""" - from torchref.model.model_ft import ModelFT from torchref.base.math_torch import ( cartesian_to_fractional_torch, - fractional_to_cartesian_torch + fractional_to_cartesian_torch, ) - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - + + model = shared_model_ft + xyz = model.xyz() cell = model.cell # Round trip conversion frac = cartesian_to_fractional_torch(xyz, cell.data) + assert frac.shape == xyz.shape xyz_back = fractional_to_cartesian_torch(frac, cell.data) # Should get back original coordinates (float32 roundtrip) @@ -312,32 +130,35 @@ def test_fractional_to_cartesian(self, sample_cif_file): @pytest.mark.integration -class TestModelFTAnisoHandling: - """Test ModelFT handling of anisotropic parameters.""" - - def test_access_aniso_atoms(self, sample_cif_file): - """Test accessing anisotropic atom information.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Check if aniso is available - if hasattr(model, 'get_aniso') or hasattr(model, 'aniso'): - # Structure has aniso - pass - - def test_isotropic_b_factors(self, sample_cif_file): - """Test accessing isotropic B-factors.""" - from torchref.model.model_ft import ModelFT - - model = ModelFT(max_res=2.0, verbose=0) - model.load_cif(str(sample_cif_file)) - - # Get B-factors (now accessed via adp()) - b_factors = model.adp() - - assert b_factors is not None - assert b_factors.shape[0] == model.xyz().shape[0] - # B-factors should be positive - assert torch.all(b_factors > 0) or torch.all(b_factors >= 0) +def test_forward_cache_contract( + loaded_model_ft, loaded_reflection_data, monkeypatch +) -> None: + """A model computes complex structure factors and caches only until invalidation.""" + from unittest.mock import Mock + + from torchref.config import caching, get_complex_dtype + + model = loaded_model_ft + hkl = loaded_reflection_data.hkl[:32] + monkeypatch.setattr(caching, "value", True) + forward = Mock(wraps=model.forward) + monkeypatch.setattr(model, "forward", forward) + assert getattr(model, "_fwd_cached_output", None) is None + + first = model(hkl) + assert first.shape == (len(hkl),) + assert first.dtype == get_complex_dtype() + assert first.device == hkl.device + assert torch.isfinite(first).all() + assert first.abs().sum() > 0 + assert model(hkl) is first + assert forward.call_count == 1 + + refreshed = model(hkl, recalc=True) + assert forward.call_count == 2 + assert refreshed is not first + # Accelerator reductions need not repeat bit-for-bit after recomputation. + relative_error = torch.linalg.vector_norm( + (refreshed - first).abs() + ) / torch.linalg.vector_norm(first.abs()) + assert relative_error < 256 * torch.finfo(first.real.dtype).eps diff --git a/tests/functional/test_restraints_functional.py b/tests/functional/test_restraints_functional.py index 815ce836..841022cc 100644 --- a/tests/functional/test_restraints_functional.py +++ b/tests/functional/test_restraints_functional.py @@ -15,16 +15,13 @@ class TestRestraintsBuildingFunctional: def test_build_restraints_from_cif(self, sample_cif_file): """Test building restraints from a real CIF file.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) restraints.build_restraints() @@ -36,40 +33,37 @@ def test_build_restraints_from_cif(self, sample_cif_file): def test_bond_restraints_built(self, sample_cif_file): """Test that bond restraints are built correctly.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) restraints.build_restraints() # Check bond restraints exist - assert 'bond' in restraints.restraints + assert "bond" in restraints.restraints # Check intra-residue bonds - if 'intra' in restraints.restraints['bond']: - bond_intra = restraints.restraints['bond']['intra'] - assert 'indices' in bond_intra - assert 'references' in bond_intra - assert 'sigmas' in bond_intra + if "intra" in restraints.restraints["bond"]: + bond_intra = restraints.restraints["bond"]["intra"] + assert "indices" in bond_intra + assert "references" in bond_intra + assert "sigmas" in bond_intra # Indices should be 2D with shape (N, 2) - indices = bond_intra['indices'] + indices = bond_intra["indices"] assert len(indices.shape) == 2 assert indices.shape[1] == 2 # References should match number of bonds - assert bond_intra['references'].shape[0] == indices.shape[0] - assert bond_intra['sigmas'].shape[0] == indices.shape[0] + assert bond_intra["references"].shape[0] == indices.shape[0] + assert bond_intra["sigmas"].shape[0] == indices.shape[0] # Bond lengths should be positive and reasonable (0.5-3.0 Å) - refs = bond_intra['references'] + refs = bond_intra["references"] assert torch.all(refs > 0.5) assert torch.all(refs < 3.0) @@ -77,65 +71,59 @@ def test_bond_restraints_built(self, sample_cif_file): def test_angle_restraints_built(self, sample_cif_file): """Test that angle restraints are built correctly.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) restraints.build_restraints() # Check angle restraints exist - assert 'angle' in restraints.restraints + assert "angle" in restraints.restraints - if 'intra' in restraints.restraints['angle']: - angle_intra = restraints.restraints['angle']['intra'] - assert 'indices' in angle_intra - assert 'references' in angle_intra - assert 'sigmas' in angle_intra + if "intra" in restraints.restraints["angle"]: + angle_intra = restraints.restraints["angle"]["intra"] + assert "indices" in angle_intra + assert "references" in angle_intra + assert "sigmas" in angle_intra # Indices should be 2D with shape (N, 3) - indices = angle_intra['indices'] + indices = angle_intra["indices"] assert len(indices.shape) == 2 assert indices.shape[1] == 3 # References should match number of angles - assert angle_intra['references'].shape[0] == indices.shape[0] + assert angle_intra["references"].shape[0] == indices.shape[0] @pytest.mark.integration def test_torsion_restraints_built(self, sample_cif_file): """Test that torsion restraints are built correctly.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) restraints.build_restraints() # Check torsion restraints exist - assert 'torsion' in restraints.restraints + assert "torsion" in restraints.restraints - if 'intra' in restraints.restraints['torsion']: - torsion_intra = restraints.restraints['torsion']['intra'] - assert 'indices' in torsion_intra - assert 'references' in torsion_intra - assert 'sigmas' in torsion_intra - assert 'periods' in torsion_intra + if "intra" in restraints.restraints["torsion"]: + torsion_intra = restraints.restraints["torsion"]["intra"] + assert "indices" in torsion_intra + assert "references" in torsion_intra + assert "sigmas" in torsion_intra + assert "periods" in torsion_intra # Indices should be 2D with shape (N, 4) - indices = torsion_intra['indices'] + indices = torsion_intra["indices"] assert len(indices.shape) == 2 assert indices.shape[1] == 4 @@ -143,29 +131,26 @@ def test_torsion_restraints_built(self, sample_cif_file): def test_plane_restraints_built(self, sample_cif_file): """Test that plane restraints are built correctly.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) restraints.build_restraints() # Check plane restraints exist - assert 'plane' in restraints.restraints + assert "plane" in restraints.restraints # Planes are grouped by atom count (3_atoms, 4_atoms, etc.) - plane_restraints = restraints.restraints['plane'] + plane_restraints = restraints.restraints["plane"] if len(list(plane_restraints.keys())) > 0: # Check at least one plane group exists for key, plane_group in plane_restraints.items(): - if 'indices' in plane_group: - indices = plane_group['indices'] + if "indices" in plane_group: + indices = plane_group["indices"] # Planes need at least 3 atoms if len(indices.shape) == 2: assert indices.shape[1] >= 3 @@ -178,21 +163,18 @@ class TestRestraintsDeviationsFunctional: def test_bond_deviations(self, sample_cif_file): """Test computing bond length deviations.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) restraints.build_restraints() # Compute bond deviations - if hasattr(restraints, 'bond_deviations'): + if hasattr(restraints, "bond_deviations"): deviations, sigmas = restraints.bond_deviations() assert torch.all(torch.isfinite(deviations)) @@ -205,21 +187,18 @@ def test_bond_deviations(self, sample_cif_file): def test_angle_deviations(self, sample_cif_file): """Test computing angle deviations.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) restraints.build_restraints() # Compute angle deviations - if hasattr(restraints, 'angle_deviations'): + if hasattr(restraints, "angle_deviations"): deviations, sigmas = restraints.angle_deviations() assert torch.all(torch.isfinite(deviations)) @@ -231,28 +210,20 @@ class TestRestraintsMultipleStructures: @pytest.mark.integration @pytest.mark.slow - def test_restraints_multiple_cif_files(self, cif_dir): - """Test building restraints for multiple CIF files.""" - from torchref.model.model import Model - from torchref.restraints import Restraints - - cif_files = list(cif_dir.glob("*.cif"))[:3] # First 3 structures - - for cif_file in cif_files: - model = Model() - model.load_cif(str(cif_file)) - - restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 - ) - restraints.build_restraints() + def test_restraints_multiple_cif_files(self, compatibility_model): + """Each extended crystal supplies bond and angle restraints.""" + from torchref.topology.restraints import Restraints - # Should have built restraints for each structure - assert 'bond' in restraints.restraints - assert 'angle' in restraints.restraints + model = compatibility_model + restraints = Restraints( + pdb=model.pdb, + xyz_fn=model.xyz, + vdw_radii_fn=model.get_vdw_radii, + verbose=0, + ) + restraints.build_restraints() + assert "bond" in restraints.restraints + assert "angle" in restraints.restraints class TestRestraintsDeviceHandling: @@ -262,22 +233,19 @@ class TestRestraintsDeviceHandling: def test_restraints_device_movement(self, sample_cif_file, cpu_device): """Test moving restraints to different devices.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model(device=cpu_device) model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) restraints.build_restraints() # Check that tensors are on the correct device - if 'bond' in restraints.restraints and 'intra' in restraints.restraints['bond']: - bond_indices = restraints.restraints['bond']['intra']['indices'] + if "bond" in restraints.restraints and "intra" in restraints.restraints["bond"]: + bond_indices = restraints.restraints["bond"]["intra"]["indices"] assert bond_indices.device == cpu_device @@ -288,16 +256,13 @@ class TestRestraintsCIFParsing: def test_cif_dict_loaded(self, sample_cif_file): """Test that CIF dictionary is loaded correctly.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) # CIF dict should be populated with residue restraints @@ -305,25 +270,25 @@ def test_cif_dict_loaded(self, sample_cif_file): assert len(restraints.cif_dict) > 0 # Should have standard amino acids - common_residues = ['ALA', 'GLY', 'VAL', 'LEU', 'ILE'] + common_residues = ["ALA", "GLY", "VAL", "LEU", "ILE"] for res in common_residues: if res in restraints.cif_dict: - assert 'bonds' in restraints.cif_dict[res] or 'angles' in restraints.cif_dict[res] + assert ( + "bonds" in restraints.cif_dict[res] + or "angles" in restraints.cif_dict[res] + ) @pytest.mark.integration def test_unique_residues_detected(self, sample_cif_file): """Test that unique residues are detected from model.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) restraints = Restraints( - pdb=model.pdb, - xyz_fn=model.xyz, - vdw_radii_fn=model.get_vdw_radii, - verbose=0 + pdb=model.pdb, xyz_fn=model.xyz, vdw_radii_fn=model.get_vdw_radii, verbose=0 ) # Should have detected unique residues diff --git a/tests/functional/test_scaler_functional.py b/tests/functional/test_scaler_functional.py index d1364665..e6c3b3a7 100644 --- a/tests/functional/test_scaler_functional.py +++ b/tests/functional/test_scaler_functional.py @@ -6,7 +6,23 @@ import pytest import torch -import numpy as np + + +@pytest.mark.integration +def test_scaler_crystal_compatibility(compatibility_model_and_data) -> None: + """Each extended crystal produces finite anisotropic scale corrections.""" + from torchref.scaling import Scaler + + model = compatibility_model_and_data["model"] + data = compatibility_model_and_data["data"] + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) + scaler.setup_anisotropy_correction() + assert scaler.s is not None + assert scaler.bins is not None + assert scaler.U is not None + correction = scaler.anisotropy_correction() + assert correction.shape == (len(data.hkl),) + assert torch.isfinite(correction).all() class TestScalerCreationFunctional: @@ -15,18 +31,18 @@ class TestScalerCreationFunctional: @pytest.mark.integration def test_scaler_full_initialization(self, sample_structure_pair): """Test full scaler initialization with model and data.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=20, verbose=0) - + # Check all components are initialized assert scaler.model is not None assert scaler._data is not None @@ -39,8 +55,8 @@ def test_scaler_full_initialization(self, sample_structure_pair): @pytest.mark.parametrize("nbins", [5, 10, 15, 20]) def test_scaler_with_different_nbins(self, sample_structure_pair, nbins): """Test scaler with different bin counts.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler model = Model() @@ -62,18 +78,18 @@ class TestScatteringVectorsFunctional: @pytest.mark.integration def test_scattering_vectors_shape(self, sample_structure_pair): """Test scattering vectors have correct shape.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) - + # s should have shape (n_reflections, 3) n_refl = data.hkl.shape[0] assert scaler.s.shape == (n_refl, 3) @@ -81,21 +97,21 @@ def test_scattering_vectors_shape(self, sample_structure_pair): @pytest.mark.integration def test_scattering_vectors_magnitude(self, sample_structure_pair): """Test scattering vector magnitudes are reasonable.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) - + # Calculate |s| = sin(theta)/lambda = 1/(2d) s_mag = torch.norm(scaler.s, dim=1) - + # For typical protein data: # Low resolution (d=100Å): |s| ~ 0.005 # High resolution (d=1Å): |s| ~ 0.5 @@ -109,26 +125,26 @@ class TestAnisotropyCorrectionFunctional: @pytest.mark.integration def test_anisotropy_setup_and_compute(self, sample_structure_pair): """Test setting up and computing anisotropy correction.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) scaler.setup_anisotropy_correction() - + # U parameters should exist - assert hasattr(scaler, 'U') + assert hasattr(scaler, "U") assert scaler.U.shape == (6,) # U11, U22, U33, U12, U13, U23 - + # Compute correction correction = scaler.anisotropy_correction() - + # Correction should be positive (exponential) assert correction.shape[0] == data.hkl.shape[0] assert torch.all(correction > 0) @@ -137,22 +153,22 @@ def test_anisotropy_setup_and_compute(self, sample_structure_pair): @pytest.mark.integration def test_anisotropy_correction_near_unity(self, sample_structure_pair): """Test anisotropy correction starts near unity with small U.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) scaler.setup_anisotropy_correction() - + # With small random U values, correction should be close to 1 correction = scaler.anisotropy_correction() - + # Most values should be between 0.5 and 2.0 for small U mean_correction = correction.mean().item() assert 0.5 < mean_correction < 2.0 @@ -164,20 +180,20 @@ class TestBinwiseBfactorFunctional: @pytest.mark.integration def test_setup_binwise_bfactor(self, sample_structure_pair): """Test setting up bin-wise B-factor parameters.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) scaler.setup_bin_wise_bfactor() - assert hasattr(scaler, 'bin_wise_bfactor') + assert hasattr(scaler, "bin_wise_bfactor") assert scaler.bin_wise_bfactor.shape == (10,) # Initially should be zeros assert torch.allclose( @@ -187,24 +203,24 @@ def test_setup_binwise_bfactor(self, sample_structure_pair): @pytest.mark.integration def test_binwise_bfactor_correction(self, sample_structure_pair): """Test computing bin-wise B-factor correction.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) scaler.setup_bin_wise_bfactor() - + # Set some non-zero B-factors scaler.bin_wise_bfactor.data = torch.linspace(0, 20, 10, device=scaler.device) - + correction = scaler.bin_wise_bfactor_correction() - + # Correction should have same length as reflections assert correction.shape[0] == data.hkl.shape[0] # Should be positive (exponential) @@ -218,35 +234,35 @@ class TestScalerStateDictFunctional: @pytest.mark.integration def test_save_and_load_state_dict(self, sample_structure_pair, tmp_path): """Test saving and loading scaler state.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + # Create scaler with some setup scaler1 = Scaler(model=model, data=data, nbins=10, verbose=0) scaler1.setup_anisotropy_correction() scaler1.setup_bin_wise_bfactor() - + # Modify parameters scaler1.U.data = torch.randn(6, device=scaler1.device) scaler1.bin_wise_bfactor.data = torch.randn(10, device=scaler1.device) - + # Save state state_path = tmp_path / "scaler_state.pt" torch.save(scaler1.state_dict(), state_path) - + # Create new scaler and load state scaler2 = Scaler(model=model, data=data, nbins=10, verbose=0) scaler2.setup_anisotropy_correction() scaler2.setup_bin_wise_bfactor() scaler2.load_state_dict(torch.load(state_path, weights_only=False)) - + # Parameters should match assert torch.allclose(scaler1.U, scaler2.U) assert torch.allclose(scaler1.bin_wise_bfactor, scaler2.bin_wise_bfactor) @@ -258,18 +274,18 @@ class TestScalerHKLPropertyFunctional: @pytest.mark.integration def test_hkl_property(self, sample_structure_pair): """Test that HKL property returns correct indices.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) - + # HKL from scaler should match data hkl = scaler.hkl assert hkl is not None @@ -283,43 +299,45 @@ class TestScalerDeviceOperationsFunctional: @pytest.mark.integration def test_scaler_cpu_operation(self, sample_structure_pair): """Test scaler works on CPU.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - - scaler = Scaler(model=model, data=data, nbins=10, verbose=0, device=torch.device('cpu')) + + scaler = Scaler( + model=model, data=data, nbins=10, verbose=0, device=torch.device("cpu") + ) scaler.setup_anisotropy_correction() - - assert scaler.device.type == 'cpu' - assert scaler.s.device.type == 'cpu' - assert scaler.U.device.type == 'cpu' + + assert scaler.device.type == "cpu" + assert scaler.s.device.type == "cpu" + assert scaler.U.device.type == "cpu" @pytest.mark.integration def test_scaler_cpu_method(self, sample_structure_pair): """Test scaler.cpu() method.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) scaler.setup_anisotropy_correction() scaler.cpu() - + # All tensors should be on CPU for param in scaler.parameters(): - assert param.device.type == 'cpu' + assert param.device.type == "cpu" class TestScalerUMatrixFunctional: @@ -328,23 +346,23 @@ class TestScalerUMatrixFunctional: @pytest.mark.integration def test_u_to_matrix_conversion(self, sample_structure_pair): """Test conversion from U parameters to 3x3 matrix.""" - from torchref.model.model import Model + from torchref.base.math_torch import U_to_matrix from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - from torchref.base.math_torch import U_to_matrix - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) scaler.setup_anisotropy_correction() - + # Convert U vector to matrix U_matrix = U_to_matrix(scaler.U) - + # Should be 3x3 assert U_matrix.shape == (3, 3) # Should be symmetric @@ -357,24 +375,24 @@ class TestScalerGradientsFunctional: @pytest.mark.integration def test_anisotropy_gradients(self, sample_structure_pair): """Test gradients flow through anisotropy correction.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) scaler.setup_anisotropy_correction() - + # Compute correction and loss correction = scaler.anisotropy_correction() loss = correction.sum() loss.backward() - + # U should have gradients assert scaler.U.grad is not None assert torch.all(torch.isfinite(scaler.U.grad)) @@ -382,56 +400,24 @@ def test_anisotropy_gradients(self, sample_structure_pair): @pytest.mark.integration def test_binwise_bfactor_gradients(self, sample_structure_pair): """Test gradients flow through bin-wise B-factor correction.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler - + model = Model() model.load_cif(str(sample_structure_pair["model"])) - + data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) - + scaler = Scaler(model=model, data=data, nbins=10, verbose=0) scaler.setup_bin_wise_bfactor() - + # Compute correction and loss correction = scaler.bin_wise_bfactor_correction() loss = correction.sum() loss.backward() - + # bin_wise_bfactor should have gradients assert scaler.bin_wise_bfactor.grad is not None assert torch.all(torch.isfinite(scaler.bin_wise_bfactor.grad)) - - -class TestScalerMultipleStructuresFunctional: - """Functional tests with multiple structures.""" - - @pytest.mark.integration - def test_scaler_with_different_structures(self, all_test_structures): - """Test scaler works with different crystal structures.""" - from torchref.scaling.scaler import Scaler - - tested = 0 - for struct in all_test_structures: - pdb_id = struct["pdb_id"] - model = struct["model"] - data = struct["data"] - - scaler = Scaler(model=model, data=data, nbins=10, verbose=0) - scaler.setup_anisotropy_correction() - - # Verify scaler is set up correctly - assert scaler.s is not None - assert scaler.bins is not None - assert scaler.U is not None - - correction = scaler.anisotropy_correction() - assert torch.all(torch.isfinite(correction)) - - tested += 1 - if tested >= 3: # Test first 3 structures - break - - assert tested >= 1, "No test structures with both CIF and MTZ found" diff --git a/tests/functional/test_targets_functional.py b/tests/functional/test_targets_functional.py index 5eb505e1..c3f2f70d 100644 --- a/tests/functional/test_targets_functional.py +++ b/tests/functional/test_targets_functional.py @@ -4,9 +4,9 @@ Tests target functions with real model and data objects. """ +import numpy as np import pytest import torch -import numpy as np class TestXrayTargetsFunctional: @@ -15,9 +15,9 @@ class TestXrayTargetsFunctional: @pytest.mark.integration def test_gaussian_nll_with_real_data(self, sample_structure_pair): """Test Gaussian NLL calculation with real reflection data.""" - from torchref.model.model import Model - from torchref.io import ReflectionData from torchref.base.math_torch import nll_xray + from torchref.io import ReflectionData + from torchref.model.model import Model model = Model() model.load_cif(str(sample_structure_pair["model"])) @@ -43,8 +43,8 @@ def test_gaussian_nll_with_real_data(self, sample_structure_pair): @pytest.mark.integration def test_least_squares_with_real_data(self, sample_structure_pair): """Test least squares calculation with real data.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model model = Model() model.load_cif(str(sample_structure_pair["model"])) @@ -74,9 +74,9 @@ class TestRfactorCalculationsFunctional: @pytest.mark.integration def test_rfactor_with_real_data(self, sample_structure_pair): """Test R-factor calculation with real reflection data.""" - from torchref.model.model import Model - from torchref.io import ReflectionData from torchref.base.math_torch import get_rfactors + from torchref.io import ReflectionData + from torchref.model.model import Model model = Model() model.load_cif(str(sample_structure_pair["model"])) @@ -110,9 +110,9 @@ def test_rfactor_with_real_data(self, sample_structure_pair): @pytest.mark.integration def test_bin_wise_rfactors(self, sample_structure_pair): """Test bin-wise R-factor calculation.""" - from torchref.model.model import Model - from torchref.io import ReflectionData from torchref.base.math_torch import bin_wise_rfactors + from torchref.io import ReflectionData + from torchref.model.model import Model model = Model() model.load_cif(str(sample_structure_pair["model"])) @@ -234,27 +234,6 @@ def test_angle_target_with_real_structure(self, sample_cif_file, external_monome assert torch.isfinite(loss) -class TestStructureFactorCalculationFunctional: - """Functional tests for structure factor calculation.""" - - @pytest.mark.integration - def test_fcalc_shape_matches_data(self, sample_structure_pair): - """Test that calculated structure factors have correct shape.""" - from torchref.model.model import Model - from torchref.io import ReflectionData - - model = Model() - model.load_cif(str(sample_structure_pair["model"])) - - data = ReflectionData() - data.load_mtz(str(sample_structure_pair["reflections"])) - - # Check if model has fcalc calculation method - if hasattr(model, 'calc_fcalc'): - fcalc = model.calc_fcalc(data) - - # Fcalc should have same number of reflections as data - assert fcalc.shape[0] == data.hkl.shape[0] class TestScalingWithRealData: @@ -263,8 +242,8 @@ class TestScalingWithRealData: @pytest.mark.integration def test_scaler_initialization_with_real_data(self, sample_structure_pair): """Test scaler initialization with real model and data.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler model = Model() @@ -285,8 +264,8 @@ def test_scaler_initialization_with_real_data(self, sample_structure_pair): @pytest.mark.integration def test_anisotropy_correction_values(self, sample_structure_pair): """Test that anisotropy correction produces reasonable values.""" - from torchref.model.model import Model from torchref.io import ReflectionData + from torchref.model.model import Model from torchref.scaling.scaler import Scaler model = Model() @@ -315,8 +294,8 @@ class TestMathFunctionsFunctional: @pytest.mark.integration def test_scattering_vectors_from_real_data(self, sample_structure_pair): """Test scattering vector calculation with real HKL and cell.""" - from torchref.io import ReflectionData from torchref.base.math_torch import get_scattering_vectors + from torchref.io import ReflectionData data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) @@ -333,11 +312,11 @@ def test_scattering_vectors_from_real_data(self, sample_structure_pair): @pytest.mark.integration def test_coordinate_transformations_with_real_cell(self, sample_cif_file): """Test coordinate transformations with real unit cell.""" - from torchref.model.model import Model from torchref.base.math_torch import ( cartesian_to_fractional_torch, - fractional_to_cartesian_torch + fractional_to_cartesian_torch, ) + from torchref.model.model import Model model = Model() model.load_cif(str(sample_cif_file)) @@ -439,8 +418,8 @@ class TestNLLFunctionsFunctional: @pytest.mark.integration def test_nll_xray_with_identical_data(self, sample_structure_pair): """Test NLL is minimal when Fobs equals Fcalc.""" - from torchref.io import ReflectionData from torchref.base.math_torch import nll_xray + from torchref.io import ReflectionData data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) @@ -463,8 +442,8 @@ def test_nll_xray_with_identical_data(self, sample_structure_pair): @pytest.mark.integration def test_nll_xray_increases_with_error(self, sample_structure_pair): """Test NLL increases as Fcalc differs from Fobs.""" - from torchref.io import ReflectionData from torchref.base.math_torch import nll_xray + from torchref.io import ReflectionData data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) @@ -494,8 +473,8 @@ def test_nll_xray_increases_with_error(self, sample_structure_pair): @pytest.mark.integration def test_nll_xray_lognormal(self, sample_structure_pair): """Test lognormal NLL calculation.""" - from torchref.io import ReflectionData from torchref.base.math_torch import nll_xray_lognormal + from torchref.io import ReflectionData data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) @@ -520,8 +499,9 @@ class TestRiceDistributionFunctional: @pytest.mark.integration def test_rice_nll_acentric(self, sample_structure_pair): """Test Rice NLL for acentric reflections.""" - from torchref.io import ReflectionData from torch.special import i0 + + from torchref.io import ReflectionData data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) @@ -592,8 +572,8 @@ def test_sigma_weighting(self, sample_structure_pair): @pytest.mark.integration def test_resolution_weighting(self, sample_structure_pair): """Test resolution-based weighting.""" - from torchref.io import ReflectionData from torchref.base.math_torch import get_scattering_vectors + from torchref.io import ReflectionData data = ReflectionData() data.load_mtz(str(sample_structure_pair["reflections"])) @@ -704,9 +684,9 @@ class TestCombinedLossFunctional: @pytest.mark.integration def test_xray_plus_geometry_loss(self, sample_structure_pair, external_monomer_library): """Test combining X-ray and geometry losses.""" - from torchref.model.model import Model - from torchref.io import ReflectionData from torchref.base.math_torch import nll_xray + from torchref.io import ReflectionData + from torchref.model.model import Model model = Model() model.load_cif(str(sample_structure_pair["model"])) diff --git a/tests/helpers/device_cases.py b/tests/helpers/device_cases.py index 7a7515b3..d2555971 100644 --- a/tests/helpers/device_cases.py +++ b/tests/helpers/device_cases.py @@ -56,51 +56,224 @@ class DeviceCase: ignore: tuple = field(default_factory=tuple) +def _riding_xyz(device): + """A bonded torsion group and a water exercise both orientation buffers.""" + import numpy as np + + from torchref.model.riding_xyz import RidingXYZTensor + from torchref.topology.hydrogens import HydrogenFrames + + xyz = torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.5, 0.0, 0.0], + [1.5, 1.5, 0.0], + [-0.6, 0.8, 0.0], + [-0.6, -0.8, 0.0], + [3.0, 0.0, 0.0], + [3.8, 0.0, 0.0], + [2.8, 0.8, 0.0], + ], + dtype=torch.float32, + ) + frames = HydrogenFrames( + h_row=np.array([3, 4, 6, 7]), + parent_row=np.array([0, 0, 5, 5]), + n1_row=np.array([1, 1, -1, -1]), + n2_row=np.array([2, 2, -1, -1]), + frame_valid=np.array([True, True, False, False]), + torsion_group=np.array([0, 0, -1, -1]), + rotation_group=np.array([-1, -1, 0, 0]), + ) + return RidingXYZTensor(xyz, frames, device=device) + + +def _symmetry(device): + """A bare Symmetry from an explicit operation list (no space group involved).""" + import torch as _torch + + from torchref.config import get_float_dtype + from torchref.symmetry import Symmetry + + dtype = get_float_dtype() + matrices = _torch.eye(3, dtype=dtype, device=device).unsqueeze(0).repeat(2, 1, 1) + matrices[1] = -matrices[1] + translations = _torch.zeros(2, 3, dtype=dtype, device=device) + return Symmetry(matrices=matrices, translations=translations) + + +def _map_symmetry_interpolation(device): + """The interpolating map operator, on a grid that forbids direct indexing. + + P212121 requires even dimensions, so an odd grid forces the interpolating + variant rather than the streaming one. + """ + from torchref.symmetry import SpaceGroup + from torchref.symmetry.map_symmetry_interpolation import ( + _MapSymmetryInterpolation, + ) + + return _MapSymmetryInterpolation(SpaceGroup(_SG, device=device), (15, 15, 15)) + + +def _edge_block(device): + """A small bond block, origin-sorted, built straight from index arrays.""" + import numpy as np + + from torchref.topology import EdgeBlock + + return EdgeBlock.from_origins( + {"intra": np.array([[0, 1], [1, 2], [2, 3]], dtype=np.int64)}, + 2, + "bond", + device=device, + ) + + +def _atom_graph(device): + """A four-atom chain: enough to exercise the edge blocks and the CSR adjacency.""" + import numpy as np + + from torchref.topology import EdgeBlock + from torchref.topology.atom_graph import AtomGraph + + def block(rows, arity, edge_type): + return EdgeBlock.from_origins( + {"intra": np.asarray(rows, dtype=np.int64).reshape(-1, arity)}, + arity, + edge_type, + device=device, + ) + + return AtomGraph( + name=np.array(["N", "CA", "C", "O"]), + element=np.array(["N", "C", "C", "O"]), + altloc=np.array([" ", " ", " ", " "]), + residue_of=torch.zeros(4, dtype=torch.int64, device=device), + bonds=block([[0, 1], [1, 2], [2, 3]], 2, "bond"), + angles=block([[0, 1, 2], [1, 2, 3]], 3, "angle"), + torsions=block([[0, 1, 2, 3]], 4, "torsion"), + chirals=block([[1, 0, 2, 3]], 4, "chiral"), + planes={3: block([[1, 2, 3]], 3, "plane")}, + ) + + +def _topology(device): + """The atom graph above under a one-residue sequence.""" + import numpy as np + + from torchref.topology.residue_graph import ResidueGraph + from torchref.topology.topology import Topology + + residues = ResidueGraph( + chain=np.array(["A"]), + resseq=np.array([1], dtype=np.int64), + icode=np.array([""]), + resname=np.array(["GLY"]), + template_key=np.array(["GLY"], dtype=object), + atom_start=np.array([0], dtype=np.int64), + atom_end=np.array([4], dtype=np.int64), + ) + return Topology(residues=residues, atoms=_atom_graph(device)) + + def _cell(device): from torchref.symmetry import Cell return Cell(_CELL, device=device) +def _ctx(d): + """A ModelContext whose cell and space group both live on ``d``.""" + from torchref.model.context import ModelContext + from torchref.symmetry import SpaceGroup + + return ModelContext(cell=_cell(d), spacegroup=SpaceGroup(_SG, device=d)) + + +def _sffft_with_grid(d): + """An SfFFT whose grid buffers exist, so the tensor walk reaches them.""" + from torchref.model.sf_fft import SfFFT + + sf = SfFFT(_ctx(d), max_res=2.0) + sf.ensure_grid() + return sf + + +def _dataset_scaler(device): + from pathlib import Path + + from torchref.io import ReflectionData + from torchref.scaling import DatasetScaler + + data = ReflectionData(device=device, verbose=0).load_mtz( + str(Path(__file__).parents[1] / "files" / "mtz" / "1DAW.mtz") + ) + return DatasetScaler({"a": data, "b": data}, device=device) + + +def _scaled_dataset(device): + from torchref.io import ScaledDataset + + scaler = _dataset_scaler(device) + return ScaledDataset(scaler.datasets["a"], scaler, "a") + + +def _dataset_scaling_target(device): + from torchref.refinement.targets import DatasetScalingTarget + + return DatasetScalingTarget(_dataset_scaler(device)) + + CASES: List[DeviceCase] = [ + DeviceCase("DatasetScaler", _dataset_scaler, "DatasetScaler"), + DeviceCase("ScaledDataset", _scaled_dataset, "ScaledDataset"), + DeviceCase("DatasetScalingTarget", _dataset_scaling_target, "DatasetScalingTarget"), + DeviceCase("EdgeBlock", _edge_block, "EdgeBlock"), + DeviceCase("AtomGraph", _atom_graph, "AtomGraph"), + DeviceCase("Topology", _topology, "Topology"), DeviceCase("Cell", _cell, "Cell"), DeviceCase( "SpaceGroup", - lambda d: __import__( - "torchref.symmetry", fromlist=["SpaceGroup"] - ).SpaceGroup(_SG, device=d), + lambda d: __import__("torchref.symmetry", fromlist=["SpaceGroup"]).SpaceGroup( + _SG, device=d + ), "SpaceGroup", ), # D4: device implied by the cell; the SpaceGroup used to be built from the # raw (None) device argument and land on the process default instead. DeviceCase( "SfFFT_from_cell", - lambda d: __import__( - "torchref.model.sf_fft", fromlist=["SfFFT"] - ).SfFFT(cell=_cell(d), spacegroup=_SG, max_res=2.0), + lambda d: __import__("torchref.model.sf_fft", fromlist=["SfFFT"]).SfFFT( + _ctx(d), max_res=2.0 + ), "SfFFT", ), - # D4: explicit device disagreeing with the supplied cell. + # D4: explicit device disagreeing with the supplied context. DeviceCase( "SfFFT_explicit_device", - lambda d: __import__( - "torchref.model.sf_fft", fromlist=["SfFFT"] - ).SfFFT(cell=_cell("cpu"), spacegroup=_SG, max_res=2.0, device=d), + lambda d: __import__("torchref.model.sf_fft", fromlist=["SfFFT"]).SfFFT( + _ctx("cpu"), max_res=2.0, device=d + ), + "SfFFT", + ), + # The grid buffers are derived on first use; this case has them resolved. + DeviceCase( + "SfFFT_with_grid", + _sffft_with_grid, "SfFFT", ), DeviceCase( "SfDS_from_cell", - lambda d: __import__( - "torchref.model.sf_ds", fromlist=["SfDS"] - ).SfDS(cell=_cell(d), spacegroup=_SG), + lambda d: __import__("torchref.model.sf_ds", fromlist=["SfDS"]).SfDS(_ctx(d)), "SfDS", ), # D1: tensor-free shells, whose tracker is the only thing to check. DeviceCase( "ScalerBase_empty", - lambda d: __import__( - "torchref.scaling", fromlist=["ScalerBase"] - ).ScalerBase(device=d), + lambda d: __import__("torchref.scaling", fromlist=["ScalerBase"]).ScalerBase( + device=d + ), "ScalerBase", tensor_free=True, ), @@ -139,6 +312,18 @@ def _cell(device): ), "MixedTensor", ), + DeviceCase( + "RidingXYZTensor_empty", + lambda d: __import__( + "torchref.model.riding_xyz", fromlist=["RidingXYZTensor"] + ).RidingXYZTensor(device=d), + "RidingXYZTensor", + ), + DeviceCase( + "RidingXYZTensor_populated", + _riding_xyz, + "RidingXYZTensor", + ), DeviceCase( "PositiveMixedTensor", lambda d: __import__( @@ -157,6 +342,34 @@ def _cell(device): ).OccupancyTensor(device=d), "OccupancyTensor", ), + DeviceCase( + "DisorderFieldTensor_empty", + lambda d: __import__( + "torchref.model.disorder_field", fromlist=["DisorderFieldTensor"] + ).DisorderFieldTensor(device=d), + "DisorderFieldTensor", + ), + DeviceCase( + "DisorderFieldTensor_populated", + lambda d: __import__( + "torchref.model.disorder_field", fromlist=["DisorderFieldTensor"] + ).DisorderFieldTensor( + initial_values=torch.full((8,), 20.0), + xyz_fn=__import__( + "torchref.model.parameter_wrappers", fromlist=["MixedTensor"] + ).MixedTensor( + torch.arange(24, dtype=torch.float32).reshape(8, 3), device=d + ), + n_nodes=3, + k_neighbors=2, + device=d, + ), + "DisorderFieldTensor", + # The coordinate accessor is borrowed: a ModuleReference is absent from ``.to()`` + # by design, so in isolation nobody moves the referent alongside the field. + # ``Model`` owns both and moves them together. + ignore=("->ref",), + ), DeviceCase( "RigidXYZTensor_empty", lambda d: __import__( @@ -165,34 +378,48 @@ def _cell(device): "RigidXYZTensor", ), DeviceCase( - "ReciprocalSymmetryGrid", - lambda d: __import__( - "torchref.symmetry", fromlist=["ReciprocalSymmetryGrid"] - ).ReciprocalSymmetryGrid(_SG, grid_shape=(16, 16, 16), device=d), - "ReciprocalSymmetryGrid", + "SpaceGroup", + lambda d: __import__("torchref.symmetry", fromlist=["SpaceGroup"]).SpaceGroup( + _SG, device=d + ), + "SpaceGroup", ), DeviceCase( - "MapSymmetryDirect", + "Symmetry", + _symmetry, + "Symmetry", + ), + DeviceCase( + "HydrogenTopology_empty", lambda d: __import__( - "torchref.symmetry", fromlist=["MapSymmetryDirect"] - ).MapSymmetryDirect( - _SG, map_shape=(16, 16, 16), cell_params=_CELL, device=d - ), - "MapSymmetryDirect", + "torchref.topology.riding", + fromlist=["HydrogenTopology"], + ).HydrogenTopology(device=d), + "HydrogenTopology", + # The builders attach every tensor later, so a fresh topology is a bare shell + # and only its tracker can be checked. + tensor_free=True, + ), + DeviceCase( + "_MapSymmetryInterpolation", + _map_symmetry_interpolation, + "_MapSymmetryInterpolation", + # ``symmetry`` is the group this operator was built from, not state it owns. + ignore=("symmetry",), ), DeviceCase( "TensorMasks", - lambda d: __import__( - "torchref.utils", fromlist=["TensorMasks"] - ).TensorMasks(device=d), + lambda d: __import__("torchref.utils", fromlist=["TensorMasks"]).TensorMasks( + device=d + ), "TensorMasks", tensor_free=True, ), DeviceCase( "ReflectionData_empty", - lambda d: __import__( - "torchref.io", fromlist=["ReflectionData"] - ).ReflectionData(device=d), + lambda d: __import__("torchref.io", fromlist=["ReflectionData"]).ReflectionData( + device=d + ), "ReflectionData", ), DeviceCase( @@ -250,6 +477,23 @@ class TargetDeviceCase: ).ADPLocalityTarget(b["model"]), "ADPLocalityTarget", ), + # Registered unconditionally and inert off field mode, so the plain bundle + # model is enough to exercise their device handling. + TargetDeviceCase( + "NodeLoadTarget", + lambda b, d: __import__( + "torchref.refinement.targets.adp.node_load", fromlist=["NodeLoadTarget"] + ).NodeLoadTarget(b["model"]), + "NodeLoadTarget", + ), + TargetDeviceCase( + "NodeSmoothnessTarget", + lambda b, d: __import__( + "torchref.refinement.targets.adp.node_smoothness", + fromlist=["NodeSmoothnessTarget"], + ).NodeSmoothnessTarget(b["model"]), + "NodeSmoothnessTarget", + ), # Owns no tensors at all -- the case that exercises the request-driven # tracker path rather than the owned-tensor path. TargetDeviceCase( @@ -265,6 +509,7 @@ class TargetDeviceCase: # Device-bearing classes deliberately not in CASES. Every entry needs a reason; # "hard to build" is a reason, "didn't get to it" is not. UNCOVERED: Dict[str, str] = { + "CollectionDifferenceSigmaDTarget": "needs a dataset collection", # --- abstract / mixin bases: never instantiated directly ----------------- "Target": "abstract base; covered through its concrete subclasses", "ModelTarget": "abstract base; needs a loaded model", @@ -286,20 +531,23 @@ class TargetDeviceCase: "ModelCollection": "needs several loaded models", "_SharedMixedModel": "internal view owned by ModelCollection", "Scaler": "needs a loaded model + data; covered in integration", - "RestraintsNew": "needs a model + monomer library", - "HydrogenTopology": "needs a built restraint topology", + "Restraints": "needs a model + monomer library", "FrenchWilson": "needs loaded intensities", "DatasetCollection": "needs several loaded datasets", "FcalcDataset": "needs computed structure factors", "Map": "needs data + model", "DifferenceMap": "needs two datasets + a model", "LBFGSRefinement": "full pipeline; covered in integration", - "MapSymmetry": "interpolation variant; needs a real map grid", + "ModelContext": "needs a loaded structure to hold a cell and space group; " + "covered through Model in tests/unit/model/test_model_state_dict_device.py", + "_MapSymmetryDirect": "stateless view over its Symmetry: recomputes index grids " + "per operation to keep peak memory O(grid), so it owns no tensors to move", "CholeskyMixedTensor": "needs a valid ADP tensor; shares MixedTensor's paths", "CollectionScaler": "needs a dataset collection", "CollectionDifferenceTarget": "needs a dataset collection", + "CollectionTwoMomentIntensityTarget": "needs a dataset collection", "CollectionMLTarget": "needs a dataset collection", - "CollectionRiceTarget": "needs a dataset collection", + "CollectionDifferenceIntensityTarget": "needs a dataset collection", "ADPSigdTarget": "needs a model with ADPs", "AngleTarget": "needs a model with restraints", "ChiralTarget": "needs a model with restraints", @@ -319,6 +567,7 @@ class TargetDeviceCase: "ScalerLogScaleTrendTarget": "needs a scaler", "ScalerURegularizationTarget": "needs a scaler", "NLLXrayTarget": "needs model + data + scaler", + "NLLIntensityXrayTarget": "needs model + data + scaler with intensities", "LeastSquaresXrayTarget": "needs model + data + scaler", "UnitWeightK1XrayTarget": "needs model + data + scaler", "SigmaAXrayTarget": "abstract base; needs model + data + scaler", @@ -331,8 +580,6 @@ class TargetDeviceCase: "PhaseInformedDifferenceTarget": "needs two datasets + phases", "RiceDifferenceTarget": "needs two datasets", "TaylorCorrectedDifferenceTarget": "needs two datasets", - "RigidTransform": "alignment helper; needs a coordinate set", - "RigidBodyRefinement": "experimental; needs model + data", } # Everything under torchref/experimental is out of scope for the conformance diff --git a/tests/helpers/dtype_inventory.py b/tests/helpers/dtype_inventory.py new file mode 100644 index 00000000..69848dae --- /dev/null +++ b/tests/helpers/dtype_inventory.py @@ -0,0 +1,153 @@ +"""Static inventory of hardcoded float/int/complex dtypes in the torchref source. + +The library resolves one dtype per category at import (``get_float_dtype`` / +``get_int_dtype`` / ``get_complex_dtype``) and every allocation on a live path is +expected to honour it. A literal ``torch.float32`` / ``torch.int64`` / +``torch.complex128`` baked into an allocation is a latent bug: MPS has no +float64, so a float64 config silently downcasts and a float32 config silently +upcasts -- neither raises, both corrupt results far from the cause. + +This module finds every guarded ``torch.`` reference by parsing the source +(AST, not regex, so dtypes named inside docstrings or comments do not count), +and reports whether each one carries an inline justification. The guard test in +``tests/unit/test_dtype_conformance.py`` turns that into a rule: outside a small +set of inherently-exempt modules, every hardcoded dtype must be justified with a +``# dtype-ok: `` marker on its own line or the comment block above it. + +``torch.bool`` is not guarded: a mask is categorical, not numeric precision. +""" + +from __future__ import annotations + +import ast +from pathlib import Path +from typing import List, NamedTuple + +__all__ = [ + "FLOAT_DTYPES", + "INT_DTYPES", + "COMPLEX_DTYPES", + "GUARDED_DTYPES", + "JUSTIFY_MARKER", + "EXEMPT_PREFIXES", + "DtypeUse", + "find_hardcoded_dtypes", + "is_exempt", +] + +# The dtypes that must not be hardcoded on a live path -- each has a config +# default (get_float_dtype / get_int_dtype / get_complex_dtype) that an +# allocation is meant to honour. ``torch.bool`` is deliberately excluded: a mask +# is categorical, not numeric precision, so pinning it is correct not a deviation. +FLOAT_DTYPES = frozenset( + {"float64", "float32", "float16", "double", "half", "bfloat16"} +) +INT_DTYPES = frozenset( + { + "int64", "int32", "int16", "int8", + "uint8", "uint16", "uint32", "uint64", + "long", "int", "short", "char", "byte", + } +) +COMPLEX_DTYPES = frozenset( + {"complex128", "complex64", "complex32", "cfloat", "cdouble", "chalf"} +) + +# Category lookup so a finding can say which config default it should use. +_CATEGORY = {name: "float" for name in FLOAT_DTYPES} +_CATEGORY.update({name: "int" for name in INT_DTYPES}) +_CATEGORY.update({name: "complex" for name in COMPLEX_DTYPES}) +GUARDED_DTYPES = frozenset(_CATEGORY) + +# A hardcoded float dtype is allowed when this marker appears on its line or the +# line immediately above it. The text after the colon is the required reason. +JUSTIFY_MARKER = "# dtype-ok:" + +# Module path prefixes (relative to the package parent, e.g. "torchref/...") +# where hardcoded float dtypes are inherent to the file's job and a per-line +# marker would be noise rather than signal: +# * triton kernels compile against explicit, static dtypes; +# * config.py *defines* the dtype maps the rest of the code reads; +# * scripts/ generate static on-disk tables offline, not model tensors. +EXEMPT_PREFIXES = ( + "torchref/base/targets/triton/", + "torchref/base/direct_summation/triton_ds.py", + "torchref/config.py", + "torchref/scripts/", +) + + +class DtypeUse(NamedTuple): + """One ``torch.`` reference (float, int, or complex) in the source.""" + + where: str # "relative/path.py:lineno" + rel_path: str # "relative/path.py" + lineno: int + dtype: str # e.g. "float64" + category: str # "float" | "int" | "complex" + line: str # the source line, stripped + justified: bool # carries JUSTIFY_MARKER on its line or the block above + + +def is_exempt(rel_path: str) -> bool: + """Whether ``rel_path`` is in an inherently-exempt module.""" + return any(rel_path.startswith(p) for p in EXEMPT_PREFIXES) + + +def _is_torch_dtype(node: ast.AST) -> str | None: + """Return the dtype name if ``node`` is a guarded ``torch.``, else None.""" + if ( + isinstance(node, ast.Attribute) + and node.attr in GUARDED_DTYPES + and isinstance(node.value, ast.Name) + and node.value.id == "torch" + ): + return node.attr + return None + + +def find_hardcoded_dtypes(package_root: Path) -> List[DtypeUse]: + """Every guarded ``torch.`` reference under ``package_root``. + + Uses the AST so references inside strings and comments are not counted, then + reads the raw source lines to decide whether each carries a justification. + """ + uses: List[DtypeUse] = [] + + for path in sorted(package_root.rglob("*.py")): + text = path.read_text(encoding="utf-8") + try: + tree = ast.parse(text, filename=str(path)) + except (SyntaxError, UnicodeDecodeError): # pragma: no cover + continue + lines = text.splitlines() + rel = str(path.relative_to(package_root.parent)) + + for node in ast.walk(tree): + dtype = _is_torch_dtype(node) + if dtype is None: + continue + lineno = node.lineno + this_line = lines[lineno - 1] if 0 < lineno <= len(lines) else "" + # A marker counts if it is on the reference's own line, or anywhere in + # the contiguous block of comment-only lines immediately above it -- so + # a multi-line justification works with the marker on any of its lines. + justified = JUSTIFY_MARKER in this_line + i = lineno - 2 + while not justified and i >= 0 and lines[i].strip().startswith("#"): + if JUSTIFY_MARKER in lines[i]: + justified = True + i -= 1 + uses.append( + DtypeUse( + where=f"{rel}:{lineno}", + rel_path=rel, + lineno=lineno, + dtype=dtype, + category=_CATEGORY[dtype], + line=this_line.strip(), + justified=justified, + ) + ) + + return uses diff --git a/tests/helpers/structure_cases.py b/tests/helpers/structure_cases.py new file mode 100644 index 00000000..74a84c28 --- /dev/null +++ b/tests/helpers/structure_cases.py @@ -0,0 +1,27 @@ +"""Name compatibility datasets explicitly so adding a file cannot grow test work silently. + +The quick reader contracts use 1DAW. The broader panel exercises deposited files +across crystal systems and file encodings in the slow tier. Pair-based pipeline +checks use trigonal 2DQ6 and body-centred tetragonal 3A5V in addition to their +separate 1DAW checks. +""" + +MODEL_CODES = ( + "1DAW", # C-centred monoclinic; quick reference structure. + "2DQ6", # Trigonal. + "3A5V", # Body-centred tetragonal. + "3E98", # Monoclinic screw axis. + "3GR5", # Hexagonal screw axis. + "3K7M", # Cubic. + "3VRJ", # Additional monoclinic deposition. + "4BX9", # Tetragonal screw axis. + "5BOV", # Triclinic P1. + "6G9X", # Orthorhombic. +) + +MTZ_CODES = MODEL_CODES + ("1AK5", "1BYW", "1VER", "6JZA", "6SXW", "6VHI") +SF_CIF_CODES = MODEL_CODES + ("7L84",) +EXTENDED_PAIR_CODES = ("2DQ6", "3A5V") +MODEL_CIF_FILES = tuple(f"{code}.cif" for code in MODEL_CODES) + ( + "test_ihm_ensemble.cif", +) diff --git a/tests/integration/alignment/__init__.py b/tests/integration/alignment/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration/alignment/benchmark_phaser_rotation_ranking.py b/tests/integration/alignment/benchmark_phaser_rotation_ranking.py new file mode 100644 index 00000000..ed4c13eb --- /dev/null +++ b/tests/integration/alignment/benchmark_phaser_rotation_ranking.py @@ -0,0 +1,442 @@ +#!/usr/bin/env python +""" +Phaser FRF-only baseline for the rotation-ranking benchmark. + +Mirrors `benchmark_rotation_ranking.py` but invokes Phaser in `MODE MR_FRF` +so we measure ONLY the rotation-function peak list — no rescore, no +translation. This isolates rotation-function quality from any +downstream-stage difference. + +For each (PDB, seed): + 1. Apply the same `R_true` that the torchref benchmark applies. + 2. Write rotated PDB. + 3. Run `phenix.phaser` MODE MR_FRF with PEAKS ROT NUMBER 500. + 4. Parse `SOLU TRIAL` lines from the .sol; each gives Euler + RFZ. + 5. Convert to rotation matrices, compute symmetry-orbit rank-of-truth + against R_true (modulo spacegroup), report rank + best angle. + 6. CSV output. + +The orbit-distance convention is the same as the torchref benchmark: + target = R_true (verified empirically against the Phaser-placed PDB + in the prior MR_AUTO sweep — if the convention + were R_true.T, the placement's err° would not be + 0.03° on 1DAW.) + +`module load phenix/phenix-1.20-4459` must be done in the surrounding +slurm script. +""" + +from __future__ import annotations + +import argparse +import csv +import math +import random +import re +import shutil +import subprocess +import time +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional + +import torch + +from torchref.experimental.alignment.frf.rotation_utils import ( + rotation_angular_distance_deg, + rotation_matrix_from_edmonds_euler, +) +from torchref.io.datasets.reflection_data import ReflectionData +from torchref.model import ModelFT + + +TEST_FILES = Path( + "/das/work/p17/p17490/Peter/Library/work_trees_torchref/fix_alignment/tests/files" +) +PAIRS = { + "1AK5": (TEST_FILES / "pdb" / "1AK5_with_H.pdb", TEST_FILES / "mtz" / "1AK5.mtz"), + "1DAW": (TEST_FILES / "pdb" / "1DAW.pdb", TEST_FILES / "mtz" / "1DAW.mtz"), + "2DQ6": (TEST_FILES / "pdb" / "2DQ6.pdb", TEST_FILES / "mtz" / "2DQ6.mtz"), + "3A5V": (TEST_FILES / "pdb" / "3A5V.pdb", TEST_FILES / "mtz" / "3A5V.mtz"), + "3E98": (TEST_FILES / "pdb" / "3E98.pdb", TEST_FILES / "mtz" / "3E98.mtz"), + "3GR5": (TEST_FILES / "pdb" / "3GR5.pdb", TEST_FILES / "mtz" / "3GR5.mtz"), + "3K7M": (TEST_FILES / "pdb" / "3K7M.pdb", TEST_FILES / "mtz" / "3K7M.mtz"), + "3VRJ": (TEST_FILES / "pdb" / "3VRJ.pdb", TEST_FILES / "mtz" / "3VRJ.mtz"), + "4BX9": (TEST_FILES / "pdb" / "4BX9.pdb", TEST_FILES / "mtz" / "4BX9.mtz"), + "6G9X": (TEST_FILES / "pdb" / "6G9X.pdb", TEST_FILES / "mtz" / "6G9X.mtz"), +} + + +@dataclass +class PhaserRotPeak: + """One row from Phaser's FRF peak list.""" + alpha_deg: float + beta_deg: float + gamma_deg: float + rfz: float + + +def _random_rotation(seed: int) -> torch.Tensor: + """Identical to torchref's benchmark — same seed → same R_true.""" + g = torch.Generator().manual_seed(int(seed)) + A = torch.randn(3, 3, generator=g, dtype=torch.float64) + Q, R = torch.linalg.qr(A) + Q = Q @ torch.diag(torch.sign(torch.diag(R))) + if torch.det(Q) < 0: + Q[:, 0] = -Q[:, 0] + return Q + + +def _discover_mtz_amplitude_labels(mtz_path: Path) -> tuple[str, str]: + import gemmi + mtz = gemmi.read_mtz_file(str(mtz_path)) + f_label, sigf_label = None, None + for c in mtz.columns: + if c.type == "F" and f_label is None: + f_label = c.label + elif c.type == "Q" and sigf_label is None: + sigf_label = c.label + if f_label and sigf_label: + break + if f_label is None or sigf_label is None: + raise RuntimeError(f"no F/SIGF columns in {mtz_path}") + return f_label, sigf_label + + +_SOLU_TRIAL_RE = re.compile( + r"SOLU\s+TRIAL\s+ENSEMBLE\s+\S+\s+" + r"EULER\s+([-+\d.]+)\s+([-+\d.]+)\s+([-+\d.]+)\s+" + r"RF\s+([-+\d.eE]+)\s+RFZ\s+([-+\d.eE]+)", + re.IGNORECASE, +) + + +def _parse_phaser_frf_peaks(sol_path: Path) -> List[PhaserRotPeak]: + """Extract FRF peaks from Phaser's .rlist file: + + SOLU TRIAL ENSEMBLE search EULER 39.764 69.564 211.710 RF 78.6 RFZ 10.25 + """ + if not sol_path.exists(): + return [] + text = sol_path.read_text() + out: List[PhaserRotPeak] = [] + for ln in text.splitlines(): + if "SOLU TRIAL" not in ln.upper(): + continue + m = _SOLU_TRIAL_RE.search(ln) + if m is None: + continue + a, b, g, _rf, rfz = (float(x) for x in m.groups()) + out.append(PhaserRotPeak(a, b, g, rfz)) + # Phaser writes peaks in descending-RFZ order. Re-sort defensively. + out.sort(key=lambda p: p.rfz, reverse=True) + return out + + +def _write_phaser_frf_kw( + work: Path, *, mtz_path: Path, rotated_pdb: Path, + f_label: str, sigf_label: str, root: str, + n_peaks: int, +) -> Path: + """Phaser keywords for FRF-only mode.""" + kw = work / "phaser_frf.kw" + # Phaser keyword grammar (from PEAK.cc): + # PEAKS ROT SELECT {SIG|NUM|PERCENT|ALL} + # PEAKS ROT CUTOFF (interpretation depends on SELECT) + # PEAKS ROT CLUSTER {ON|OFF} + # With `SELECT NUM` + `CUTOFF n`, Phaser keeps top-n peaks. Disable + # clustering so symmetry-equivalents aren't merged before we rank them. + kw.write_text(f"""TITLE FRF-only rotation ranking +MODE MR_FRF +HKLIN {mtz_path} +LABIN F={f_label} SIGF={sigf_label} +ENSEMBLE search PDB {rotated_pdb} IDENT 1.0 +COMPOSITION BY AVERAGE +SEARCH ENSEMBLE search +PEAKS ROT SELECT ALL +PEAKS ROT CLUSTER OFF +PEAKS ROT LEVEL 0 +ROOT {root} +""") + return kw + + +def _run_phaser_frf(work: Path, kw_path: Path, timeout_s: int) -> tuple[int, float]: + """Pipe keywords to `phenix.phaser` stdin (CCP4-keyword mode).""" + t0 = time.time() + keywords = kw_path.read_text() + try: + proc = subprocess.run( + ["phenix.phaser"], + cwd=str(work), + input=keywords, + capture_output=True, + text=True, + timeout=timeout_s, + ) + rc = proc.returncode + (work / "phaser.stdout").write_text(proc.stdout or "") + (work / "phaser.stderr").write_text(proc.stderr or "") + except subprocess.TimeoutExpired: + rc = -1 + return rc, time.time() - t0 + + +def _euler_zyz_to_R(alpha, beta, gamma) -> torch.Tensor: + """Edmonds ZYZ active rotation: R = Rz(α) · Ry(β) · Rz(γ). Vectorised.""" + a, b, g = (torch.as_tensor(x, dtype=torch.float64) for x in (alpha, beta, gamma)) + ca, sa = a.cos(), a.sin() + cb, sb = b.cos(), b.sin() + cg, sg = g.cos(), g.sin() + # R = Rz(a) Ry(b) Rz(g), 3x3 column-vector convention. + R = torch.stack([ + torch.stack([ca*cb*cg - sa*sg, -ca*cb*sg - sa*cg, ca*sb], dim=-1), + torch.stack([sa*cb*cg + ca*sg, -sa*cb*sg + ca*cg, sa*sb], dim=-1), + torch.stack([-sb*cg, sb*sg, cb ], dim=-1), + ], dim=-2) + return R + + +def orbit_rank( + peaks: List[PhaserRotPeak], + R_target: torch.Tensor, + sym_mats: torch.Tensor, + threshold_deg: float = 5.0, +) -> dict: + """Symmetry-orbit-aware rank, vectorised over all peaks at once.""" + if not peaks: + return { + "best_rank": -1, "best_ang_deg": float("inf"), + "best_rfz": 0.0, "top1_ang_deg": float("inf"), + "any_below_threshold": False, "n_peaks": 0, + } + R_target = R_target.cpu().to(torch.float64) + sym_mats = sym_mats.cpu().to(torch.float64) + + alphas = torch.tensor( + [math.radians(p.alpha_deg) for p in peaks], dtype=torch.float64, + ) + betas = torch.tensor( + [math.radians(p.beta_deg) for p in peaks], dtype=torch.float64, + ) + gammas = torch.tensor( + [math.radians(p.gamma_deg) for p in peaks], dtype=torch.float64, + ) + R_peaks = _euler_zyz_to_R(alphas, betas, gammas) # (N, 3, 3) + orbit = sym_mats @ R_target.unsqueeze(0) # (n_ops, 3, 3) + + # For each peak and each sym op: trace(R_peak · (S_k R_target)^T) / 2 + # Then arccos for angular distance; min over sym ops; argmin over peaks. + Rk_T = orbit.transpose(-1, -2) # (n_ops, 3, 3) + # Batched matrix-multiply: (N, 1, 3, 3) @ (1, n_ops, 3, 3) → (N, n_ops, 3, 3) + M = R_peaks.unsqueeze(1) @ Rk_T.unsqueeze(0) + tr = M.diagonal(offset=0, dim1=-2, dim2=-1).sum(-1) # (N, n_ops) + cos_a = ((tr - 1.0) / 2.0).clamp(-1.0, 1.0) + angles = cos_a.arccos() * (180.0 / math.pi) # (N, n_ops) + per_peak_min, _ = angles.min(dim=-1) # (N,) + + top1_ang = float(per_peak_min[0].item()) + # Find first peak below threshold; otherwise return global argmin. + below = (per_peak_min <= threshold_deg).nonzero(as_tuple=True)[0] + if below.numel() > 0: + rank = int(below[0].item()) + else: + rank = int(per_peak_min.argmin().item()) + best_ang = float(per_peak_min[rank].item()) + best_rfz = peaks[rank].rfz + + return { + "best_rank": rank, + "best_ang_deg": best_ang, + "best_rfz": best_rfz, + "top1_ang_deg": top1_ang, + "any_below_threshold": best_ang <= threshold_deg, + "n_peaks": len(peaks), + } + + +def run_one( + pdb_key: str, + seed: int, + *, + work_root: Path, + n_peaks: int, + threshold_deg: float, + timeout_s: int, + verbose: int, +) -> dict: + pdb_path, mtz_path = PAIRS[pdb_key] + print(f"\n=== PHASER FRF {pdb_key} seed={seed} ===", flush=True) + + t0 = time.time() + model = ModelFT(device=torch.device("cpu")).load_pdb(str(pdb_path)) + data = ReflectionData(device="cpu").load_mtz(str(mtz_path)) + sym_mats = data.spacegroup.matrices.to(dtype=torch.float64) + + R_true = _random_rotation(seed) + centroid = model.xyz().mean(0) + rotated = model.rotate(R_true.to(model.dtype_float), center=centroid) + + work = work_root / f"{pdb_key}_seed{seed}" + if work.exists(): + shutil.rmtree(work) + work.mkdir(parents=True) + rotated_pdb = work / "rotated.pdb" + rotated.write_pdb(str(rotated_pdb)) + + f_label, sigf_label = _discover_mtz_amplitude_labels(mtz_path) + root = "phaser_frf" + kw_path = _write_phaser_frf_kw( + work, mtz_path=mtz_path, rotated_pdb=rotated_pdb, + f_label=f_label, sigf_label=sigf_label, root=root, + n_peaks=n_peaks, + ) + rc, t_phaser = _run_phaser_frf(work, kw_path, timeout_s) + # Phaser FRF writes the rotation peak list to .rlist, not .sol. + # (.sol is reserved for full MR_AUTO solutions.) + rlist_path = work / f"{root}.rlist" + + peaks = _parse_phaser_frf_peaks(rlist_path) + if verbose: + print(f" rc={rc} n_peaks={len(peaks)} t={t_phaser:.1f}s", flush=True) + + # Test convention against R_true (the same target the torchref benchmark uses). + rank_R_true = orbit_rank(peaks, R_true, sym_mats, threshold_deg=threshold_deg) + # Also evaluate against R_true.T to check whether Phaser's Euler points + # the opposite direction. The smaller `top1_ang_deg` indicates the + # correct convention for this pipeline. + rank_R_true_T = orbit_rank(peaks, R_true.T, sym_mats, threshold_deg=threshold_deg) + + if verbose: + print( + f" against R_true: rank={rank_R_true['best_rank']:>4} " + f"ang={rank_R_true['best_ang_deg']:.2f}° " + f"top1ang={rank_R_true['top1_ang_deg']:.2f}°", + flush=True, + ) + print( + f" against R_true.T: rank={rank_R_true_T['best_rank']:>4} " + f"ang={rank_R_true_T['best_ang_deg']:.2f}° " + f"top1ang={rank_R_true_T['top1_ang_deg']:.2f}°", + flush=True, + ) + + # Pick the convention with smaller top1 angular distance as the + # "canonical" measurement for this row. + pick = ( + rank_R_true if rank_R_true["top1_ang_deg"] <= rank_R_true_T["top1_ang_deg"] + else rank_R_true_T + ) + convention = "R_true" if pick is rank_R_true else "R_true.T" + + return { + "pdb": pdb_key, + "seed": seed, + "spacegroup": str(data.spacegroup), + "phaser_exit": rc, + "n_peaks": len(peaks), + "convention_used": convention, + "rank": pick["best_rank"], + "ang_deg": pick["best_ang_deg"], + "rfz": pick["best_rfz"], + "top1_ang_deg": pick["top1_ang_deg"], + "rank_R_true": rank_R_true["best_rank"], + "ang_R_true_deg": rank_R_true["best_ang_deg"], + "rank_R_true_T": rank_R_true_T["best_rank"], + "ang_R_true_T_deg": rank_R_true_T["best_ang_deg"], + "phaser_time_s": t_phaser, + "total_time_s": time.time() - t0, + } + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--pdb", default=None, choices=sorted(PAIRS.keys())) + ap.add_argument("--seed", type=int, default=None, + help="Top-level seed; same semantics as the torchref " + "benchmark so the per-trial R_true matches.") + ap.add_argument("--n-trials", type=int, default=1) + ap.add_argument("--sweep", action="store_true") + ap.add_argument("--n-peaks", type=int, default=500, + help="PEAKS ROT NUMBER for Phaser FRF (default 500 to " + "match the torchref benchmark's n_rotation_peaks).") + ap.add_argument("--threshold-deg", type=float, default=5.0) + ap.add_argument("--timeout-s", type=int, default=1200) + ap.add_argument("--workdir", default=None) + ap.add_argument("--out-csv", default=None) + ap.add_argument("--verbose", type=int, default=1) + args = ap.parse_args() + + if args.seed is None: + args.seed = int(time.time()) + rng = random.Random(args.seed) + print(f"top-level seed = {args.seed}", flush=True) + + if args.sweep: + worklist = [(pdb, rng.randint(0, 10 ** 9)) + for pdb in sorted(PAIRS.keys()) + for _ in range(args.n_trials)] + else: + worklist = [] + for _ in range(args.n_trials): + pdb = args.pdb if args.pdb is not None else rng.choice(list(PAIRS.keys())) + worklist.append((pdb, rng.randint(0, 10 ** 9))) + + work_root = Path(args.workdir or f"/tmp/phaser_frf_{int(time.time())}") + work_root.mkdir(parents=True, exist_ok=True) + out_csv = Path(args.out_csv or f"phaser_frf_results_{int(time.time())}.csv") + + rows = [] + for pdb_key, trial_seed in worklist: + try: + r = run_one( + pdb_key, trial_seed, + work_root=work_root, + n_peaks=args.n_peaks, + threshold_deg=args.threshold_deg, + timeout_s=args.timeout_s, + verbose=args.verbose, + ) + rows.append(r) + except Exception as exc: + import traceback + traceback.print_exc() + print(f" TRIAL FAILED on {pdb_key} seed={trial_seed}: {exc!r}", + flush=True) + rows.append({"pdb": pdb_key, "seed": trial_seed, + "error": repr(exc)}) + + if rows: + cols = sorted({k for r in rows for k in r}) + with open(out_csv, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=cols) + writer.writeheader() + writer.writerows(rows) + print(f"\nwrote {len(rows)} rows to {out_csv}", flush=True) + + # Console summary + print("\n=== Phaser FRF summary (best rank / best ang°, both conventions) ===", + flush=True) + print( + f"{'pdb':>6} {'seed':>10} {'n_pk':>5} " + f"{'pick':>9} {'rank':>4} {'ang°':>6} {'RFZ':>6} " + f"{'R_true rank':>11} {'R_T rank':>9} {'t(s)':>6}", + flush=True, + ) + for r in rows: + if r.get("error"): + print(f"{r['pdb']:>6} {r['seed']:>10} FAILED: {r['error']}", + flush=True) + continue + print( + f"{r['pdb']:>6} {r['seed']:>10} {r['n_peaks']:>5} " + f"{r['convention_used']:>9} {r['rank']:>4} " + f"{r['ang_deg']:>6.2f} {r['rfz']:>6.2f} " + f"{r['rank_R_true']:>11} {r['rank_R_true_T']:>9} " + f"{r['phaser_time_s']:>6.1f}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/integration/alignment/profile_fit.py b/tests/integration/alignment/profile_fit.py new file mode 100644 index 00000000..896feb3e --- /dev/null +++ b/tests/integration/alignment/profile_fit.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python +""" +Profile `ModelFT.fit_to_data` to find where the time goes. + +Run from the repo root: + .venv/bin/python tests/integration/alignment/profile_fit.py [--pdb 1DAW] \ + [--n-rotation-candidates 3] [--n-translation-candidates 3] \ +Output: top-50 cumulative-time entries from cProfile + a custom per-stage timer +breakdown (rotation search, ML rescore, TF, local refine, joint refine, +final Scaler refit). +""" +from __future__ import annotations + +import argparse +import cProfile +import pstats +import time +from contextlib import contextmanager +from pathlib import Path + +import torch + +from torchref.experimental.alignment import align_model_to_data +from torchref.experimental.alignment.frf.rotation_utils import rotation_matrix_from_edmonds_euler +from torchref.io.datasets.reflection_data import ReflectionData +from torchref.model import ModelFT + + +TEST_FILES = Path("/das/work/p17/p17490/Peter/Library/work_trees_torchref/fix_alignment/tests/files") + +PAIRS = { + "1DAW": (TEST_FILES / "pdb" / "1DAW.pdb", TEST_FILES / "mtz" / "1DAW.mtz"), + "1AK5": (TEST_FILES / "pdb" / "1AK5_with_H.pdb", TEST_FILES / "mtz" / "1AK5.mtz"), + "3A5V": (TEST_FILES / "pdb" / "3A5V.pdb", TEST_FILES / "mtz" / "3A5V.mtz"), +} + + +_TIMINGS: dict[str, float] = {} + + +@contextmanager +def _stage(name: str): + t0 = time.time() + try: + yield + finally: + _TIMINGS[name] = _TIMINGS.get(name, 0.0) + (time.time() - t0) + print(f" [{name}] {time.time()-t0:.2f}s", flush=True) + + +def _patch_for_timing(): + """Wrap key fit_to_data stages so we get an inline breakdown.""" + from torchref.experimental.alignment import rotation_search, translation + from torchref import scaling + + originals = {} + + def wrap(module, attr, label): + original = getattr(module, attr) + originals[(module, attr)] = original + + def wrapper(*args, **kwargs): + with _stage(label): + return original(*args, **kwargs) + + setattr(module, attr, wrapper) + + wrap(rotation_search, "search_peaks", "rotation_search") + wrap(translation, "amplitude_translation_search", "amplitude_translation_search") + wrap(translation, "local_translation_refine", "local_translation_refine") + wrap(translation, "precompute_G_for_rotation", "precompute_G_for_rotation") + wrap(translation, "llg_translation_rescore", "llg_translation_rescore") + return originals + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--pdb", default="1DAW", choices=sorted(PAIRS.keys())) + ap.add_argument("--n-rotation-candidates", type=int, default=3) + ap.add_argument("--n-translation-candidates", type=int, default=3) + ap.add_argument("--top", type=int, default=40, help="top N cProfile entries") + args = ap.parse_args() + + pdb_path, mtz_path = PAIRS[args.pdb] + print(f"=== Profiling fit_to_data on {args.pdb} ===", flush=True) + + data = ReflectionData().load_mtz(str(mtz_path)) + canonical = ModelFT().load_pdb(str(pdb_path)) + canonical.spacegroup = "P 1" + R_true = rotation_matrix_from_edmonds_euler(0.6, 0.4, 1.2) + rotated_p = canonical.rotate( + R_true.to(canonical.dtype_float), center=canonical.xyz().mean(dim=0), + ) + t_true = torch.tensor([0.18, -0.07, 0.23], dtype=canonical.dtype_float) + perturbed = rotated_p.translate(t_true, fractional=True) + print(f" spacegroup={data.spacegroup}, n_atoms={canonical.xyz().shape[0]}, " + f"n_hkl={data.hkl.shape[0]}", flush=True) + + _patch_for_timing() + + profiler = cProfile.Profile() + t0 = time.time() + profiler.enable() + aligned = align_model_to_data( + perturbed, + data, + n_rotation_candidates=args.n_rotation_candidates, + n_translation_candidates=args.n_translation_candidates, + verbose=0, + ) + profiler.disable() + total = time.time() - t0 + + print(f"\n=== Stage breakdown (total {total:.2f}s) ===", flush=True) + other = total - sum(_TIMINGS.values()) + for name, t in sorted(_TIMINGS.items(), key=lambda kv: -kv[1]): + print(f" {name:40s} {t:8.2f}s ({100*t/total:5.1f}%)", flush=True) + print(f" {'(unattributed)':40s} {other:8.2f}s ({100*other/total:5.1f}%)", + flush=True) + + print(f"\n=== Top-{args.top} cProfile (cumulative time) ===", flush=True) + stats = pstats.Stats(profiler).sort_stats("cumulative") + stats.print_stats(args.top) + + print(f"\n=== Top-{args.top} cProfile (own time) ===", flush=True) + stats = pstats.Stats(profiler).sort_stats("tottime") + stats.print_stats(args.top) + + +if __name__ == "__main__": + main() diff --git a/tests/integration/alignment/run_phaser_mr.py b/tests/integration/alignment/run_phaser_mr.py new file mode 100644 index 00000000..d42eb79b --- /dev/null +++ b/tests/integration/alignment/run_phaser_mr.py @@ -0,0 +1,431 @@ +#!/usr/bin/env python +""" +Phaser-MR baseline for the same problem `run_random_pdb_fit.py` runs against +torchref's alignment pipeline. + +Flow: + 1. Pick a PDB / MTZ pair from `tests/files`. + 2. Load the model + data; apply the SAME random rotation R_true that + `run_random_pdb_fit.py` would (same seed → same R_true). + 3. Write the rotated model to a temporary PDB. + 4. Invoke `phenix.phaser` MODE MR_AUTO on (rotated_model, original_mtz). + 5. Parse Phaser's `.sum` for LLG / TFZ / R-factor. + 6. Load Phaser's placed PDB; compute angular distance to the canonical + model, modulo the spacegroup symmetry orbit (same metric as + `run_random_pdb_fit.py`'s `aligned-vs-canonical` line). + 7. Append a CSV row. + +The script does NOT compute torchref's Scaler R-work on Phaser's placement — +per the design decision, both pipelines report their OWN R-work and we +display them side-by-side. Differences between the columns are partly +scaler-vs-scaler effects, not just placement quality. + +Requires `module load phenix/phenix-1.20-4459` to be done in the surrounding +slurm script (PATH-resolved at subprocess.run time). +""" + +from __future__ import annotations + +import argparse +import csv +import math +import os +import random +import re +import shutil +import subprocess +import time +from pathlib import Path +from typing import Optional + +import torch + +from torchref.base.metrics.rfactor import rfactor_work_free +from torchref.experimental.alignment.frf.rotation_utils import rotation_angular_distance_deg +from torchref.io.datasets.reflection_data import ReflectionData +from torchref.model import ModelFT +from torchref.scaling import Scaler + + +# Same PAIRS as run_random_pdb_fit.py — keep in sync. +TEST_FILES = Path( + "/das/work/p17/p17490/Peter/Library/work_trees_torchref/fix_alignment/tests/files" +) +PAIRS = { + "1AK5": (TEST_FILES / "pdb" / "1AK5_with_H.pdb", TEST_FILES / "mtz" / "1AK5.mtz"), + "1DAW": (TEST_FILES / "pdb" / "1DAW.pdb", TEST_FILES / "mtz" / "1DAW.mtz"), + "2DQ6": (TEST_FILES / "pdb" / "2DQ6.pdb", TEST_FILES / "mtz" / "2DQ6.mtz"), + "3A5V": (TEST_FILES / "pdb" / "3A5V.pdb", TEST_FILES / "mtz" / "3A5V.mtz"), + "3E98": (TEST_FILES / "pdb" / "3E98.pdb", TEST_FILES / "mtz" / "3E98.mtz"), + "3GR5": (TEST_FILES / "pdb" / "3GR5.pdb", TEST_FILES / "mtz" / "3GR5.mtz"), + "3K7M": (TEST_FILES / "pdb" / "3K7M.pdb", TEST_FILES / "mtz" / "3K7M.mtz"), + "3VRJ": (TEST_FILES / "pdb" / "3VRJ.pdb", TEST_FILES / "mtz" / "3VRJ.mtz"), + "4BX9": (TEST_FILES / "pdb" / "4BX9.pdb", TEST_FILES / "mtz" / "4BX9.mtz"), + "6G9X": (TEST_FILES / "pdb" / "6G9X.pdb", TEST_FILES / "mtz" / "6G9X.mtz"), +} + + +def _random_rotation(seed: int) -> torch.Tensor: + """Identical to run_random_pdb_fit.py — must produce byte-equal R_true.""" + g = torch.Generator().manual_seed(int(seed)) + A = torch.randn(3, 3, generator=g, dtype=torch.float64) + Q, R = torch.linalg.qr(A) + Q = Q @ torch.diag(torch.sign(torch.diag(R))) + if torch.det(Q) < 0: + Q[:, 0] = -Q[:, 0] + return Q + + +def _kabsch_rotation(xyz_a: torch.Tensor, xyz_b: torch.Tensor) -> torch.Tensor: + """Identical to run_random_pdb_fit.py.""" + a = (xyz_a.detach() - xyz_a.detach().mean(0)).to(torch.float64) + b = (xyz_b.detach() - xyz_b.detach().mean(0)).to(torch.float64) + H = b.T @ a + U, _, Vt = torch.linalg.svd(H) + d = float(torch.sign(torch.det(Vt.T @ U.T))) + D = torch.diag(torch.tensor([1.0, 1.0, d], dtype=H.dtype, device=H.device)) + return Vt.T @ D @ U.T + + +def _discover_mtz_amplitude_labels(mtz_path: Path) -> tuple[str, str]: + """ + Read the MTZ column descriptions and return the first (F, SIGF) pair. + + Column types we look for: 'F' (structure-factor amplitude) and 'Q' + (standard deviation) per CCP4 MTZ column-type convention. + """ + import gemmi + mtz = gemmi.read_mtz_file(str(mtz_path)) + cols = mtz.columns + f_label, sigf_label = None, None + for c in cols: + if c.type == "F" and f_label is None: + f_label = c.label + elif c.type == "Q" and sigf_label is None: + sigf_label = c.label + if f_label and sigf_label: + break + if f_label is None or sigf_label is None: + raise RuntimeError( + f"could not find F/SIGF columns in {mtz_path}; " + f"saw {[(c.label, c.type) for c in cols]}" + ) + return f_label, sigf_label + + +_SOLU_RE = re.compile( + r"SOLU\s+SET\s+.*?RFZ=([-\d.eE+]+).*?TFZ=([-\d.eE+]+)" + r".*?LLG=([-\d.eE+]+).*?Rfac=([-\d.eE+%]+)?", + re.IGNORECASE | re.DOTALL, +) + + +def _parse_phaser_sol(sol_path: Path) -> dict: + """ + Extract LLG / TFZ / RFZ from Phaser's `.sol` solution file. + + Phaser writes one or more `SOLU SET` lines summarising each solution; + the top one is the highest-LLG. Format (Phaser 2.x): + SOLU SET RFZ=10.2 TFZ=8.9 PAK=1 LLG=260 TFZ==11.9 LLG=6176 TFZ==60.2 ... + + The double-equals (`TFZ==`, `LLG=...`) entries are post-refinement + scores; the first single-equals are pre-refinement. We report the + final refined LLG and TFZ (the last LLG / TFZ== occurrences). + + Phaser does NOT print an R-factor in the .sol — we compute one + downstream via torchref's Scaler on the placed PDB. + """ + if not sol_path.exists(): + return {"solved": False, "llg": None, "tfz": None, "rfz": None, + "n_solutions": 0} + text = sol_path.read_text() + solu_lines = [ln for ln in text.splitlines() if ln.strip().startswith("SOLU SET")] + n = len(solu_lines) + if n == 0: + return {"solved": False, "llg": None, "tfz": None, "rfz": None, + "n_solutions": 0} + + top = solu_lines[0] + # Final refined LLG / TFZ are the LAST occurrences on the line. + # (Phaser writes the pre-refinement then the refined version after.) + llg_matches = re.findall(r"LLG=([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)", top) + tfz_matches = re.findall( + r"TFZ=={0,1}([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)", top, + ) + rfz_matches = re.findall(r"RFZ=([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)", top) + llg = float(llg_matches[-1]) if llg_matches else None + tfz = float(tfz_matches[-1]) if tfz_matches else None + rfz = float(rfz_matches[-1]) if rfz_matches else None + + return { + "solved": True, + "llg": llg, "tfz": tfz, "rfz": rfz, + "n_solutions": n, + } + + +def _write_phaser_keywords( + work: Path, *, mtz_path: Path, rotated_pdb: Path, + f_label: str, sigf_label: str, root: str, +) -> Path: + """Build the Phaser keyword script. MODE MR_AUTO runs through to RNP.""" + kw = work / "phaser.kw" + kw.write_text(f"""TITLE torchref Phaser MR comparison +MODE MR_AUTO +HKLIN {mtz_path} +LABIN F={f_label} SIGF={sigf_label} +ENSEMBLE search PDB {rotated_pdb} IDENT 1.0 +COMPOSITION BY AVERAGE +SEARCH ENSEMBLE search NUM 1 +ROOT {root} +JOBS 1 +""") + return kw + + +def _run_phaser(work: Path, kw_path: Path, timeout_s: int, + verbose: int) -> tuple[int, float]: + """Invoke `phenix.phaser` with the keyword script piped to stdin. + + `phenix.phaser` with no command-line argument launches the CCP4-style + binary that reads keywords from stdin. Passing the keywords as a file + argument triggers the PHIL-parameter path, which doesn't understand + CCP4 syntax. + """ + t0 = time.time() + keywords = kw_path.read_text() + try: + proc = subprocess.run( + ["phenix.phaser"], + cwd=str(work), + input=keywords, + capture_output=True, + text=True, + timeout=timeout_s, + ) + rc = proc.returncode + # Always save stdout/stderr — Phaser's solution summary lands in + # stdout, and we want to keep a record of failed runs too. + (work / "phaser.stdout").write_text(proc.stdout or "") + (work / "phaser.stderr").write_text(proc.stderr or "") + except subprocess.TimeoutExpired: + rc = -1 + return rc, time.time() - t0 + + +def run(pdb_key: str, seed: int, *, work_root: Path, + verbose: int = 1, timeout_s: int = 3600) -> dict: + pdb_path, mtz_path = PAIRS[pdb_key] + print(f"\n=== PHASER {pdb_key} seed={seed} ===", flush=True) + + t0 = time.time() + device = torch.device("cpu") + # Load model on CPU — Phaser doesn't care, and we only need the rotated + # PDB on disk. The Scaler comparison at the end also runs on CPU. + model = ModelFT(device=device).load_pdb(str(pdb_path)) + data = ReflectionData(device=str(device)).load_mtz(str(mtz_path)) + sym_mats = data.spacegroup.matrices.to(dtype=torch.float64, device=device) + xyz_canonical = model.xyz().detach().clone().to(device) + + # Apply R_true (same random rotation `run_random_pdb_fit.py` applies). + R_true = _random_rotation(seed).to(device) + centroid = xyz_canonical.mean(0) + rotated = model.rotate( + R_true.to(model.dtype_float), center=centroid, + ) + + # Per-trial workspace + work = work_root / f"{pdb_key}_seed{seed}" + if work.exists(): + shutil.rmtree(work) + work.mkdir(parents=True) + rotated_pdb = work / "rotated.pdb" + rotated.write_pdb(str(rotated_pdb)) + + f_label, sigf_label = _discover_mtz_amplitude_labels(mtz_path) + if verbose: + print(f" MTZ labels: F={f_label}, SIGF={sigf_label}", flush=True) + + root = "phaser_out" + kw_path = _write_phaser_keywords( + work, mtz_path=mtz_path, rotated_pdb=rotated_pdb, + f_label=f_label, sigf_label=sigf_label, root=root, + ) + rc, phaser_time_s = _run_phaser(work, kw_path, timeout_s, verbose) + sol_path = work / f"{root}.sol" + placed_pdb = work / f"{root}.1.pdb" + + phaser = _parse_phaser_sol(sol_path) + if verbose: + print( + f" Phaser: rc={rc}, solved={phaser['solved']}, " + f"LLG={phaser['llg']}, TFZ={phaser['tfz']}, " + f"time={phaser_time_s:.1f}s", + flush=True, + ) + + # Angular distance + R-work on Phaser's placed model. + err_canonical_deg = float("nan") + rwork_torchref_scaler = float("nan") + rfree_torchref_scaler = float("nan") + if placed_pdb.exists(): + try: + placed = ModelFT(device=device).load_pdb(str(placed_pdb)) + # If atom counts disagree (Phaser may include only one ensemble + # member) we Kabsch on the common prefix. Phaser's placement + # preserves atom order. + xyz_placed = placed.xyz().detach().to(device) + n_match = min(xyz_placed.shape[0], xyz_canonical.shape[0]) + R_residual = _kabsch_rotation( + xyz_placed[:n_match], xyz_canonical[:n_match], + ) + err_canonical_deg = min( + rotation_angular_distance_deg( + R_residual.to(torch.float64), sym_mats[k], + ) + for k in range(sym_mats.shape[0]) + ) + except Exception as exc: + print(f" Kabsch on placed.pdb failed: {exc!r}", flush=True) + + # torchref-Scaler R-work on Phaser's placement (same Scaler the + # torchref pipeline uses for its end-of-pipeline R-work — apples + # to apples for the post-MR placement quality). + try: + placed = ModelFT(device=device).load_pdb(str(placed_pdb)) + scaler = Scaler(model=placed, data=data, nbins=20, verbose=0, + device=device) + with torch.no_grad(): + fcalc = placed(data.hkl).detach() + scaler.initialize(fcalc) + scaler.refine_lbfgs(fcalc=fcalc) + with torch.no_grad(): + rw, rf = rfactor_work_free(data, torch.abs(scaler.forward(fcalc))) + rwork_torchref_scaler = ( + rw.item() if hasattr(rw, "item") else float(rw) + ) + rfree_torchref_scaler = ( + rf.item() if hasattr(rf, "item") else float(rf) + ) + except Exception as exc: + print(f" torchref-Scaler on placed.pdb failed: {exc!r}", + flush=True) + + if verbose: + print( + f" Phaser placed-vs-canonical angular distance " + f"(mod {data.spacegroup}-symmetry): {err_canonical_deg:.2f}° " + f"R-work (torchref Scaler) = {rwork_torchref_scaler:.4f}", + flush=True, + ) + + return { + "pdb": pdb_key, + "seed": seed, + "spacegroup": str(data.spacegroup), + "phaser_exit": rc, + "phaser_solved": phaser["solved"], + "phaser_n_solutions": phaser["n_solutions"], + "phaser_llg": phaser["llg"], + "phaser_tfz": phaser["tfz"], + "phaser_rfz": phaser.get("rfz"), + "err_canonical_deg": err_canonical_deg, + "rwork_torchref_scaler": rwork_torchref_scaler, + "rfree_torchref_scaler": rfree_torchref_scaler, + "phaser_time_s": phaser_time_s, + "total_time_s": time.time() - t0, + } + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--pdb", default=None, choices=sorted(PAIRS.keys())) + ap.add_argument("--seed", type=int, default=None, + help="Same semantics as run_random_pdb_fit.py: top-level " + "rng seed that derives the per-trial seeds.") + ap.add_argument("--n-trials", type=int, default=1) + ap.add_argument("--sweep", action="store_true", + help="Iterate over all PDBs in PAIRS, n-trials per PDB.") + ap.add_argument("--timeout-s", type=int, default=3600, + help="Per-trial Phaser timeout in seconds (default 1h).") + ap.add_argument("--workdir", default=None, + help="Per-trial workspace root (default: /tmp/phaser_mr_).") + ap.add_argument("--out-csv", default=None, + help="CSV output path. Default: timestamped under cwd.") + ap.add_argument("--verbose", type=int, default=1) + args = ap.parse_args() + + if args.seed is None: + args.seed = int(time.time()) + rng = random.Random(args.seed) + print(f"top-level seed = {args.seed}", flush=True) + + if args.sweep: + worklist = [(pdb, rng.randint(0, 10 ** 9)) + for pdb in sorted(PAIRS.keys()) + for _ in range(args.n_trials)] + else: + worklist = [] + for _ in range(args.n_trials): + pdb = args.pdb if args.pdb is not None else rng.choice(list(PAIRS.keys())) + worklist.append((pdb, rng.randint(0, 10 ** 9))) + + work_root = Path(args.workdir or f"/tmp/phaser_mr_{int(time.time())}") + work_root.mkdir(parents=True, exist_ok=True) + print(f"work root: {work_root}", flush=True) + + out_csv = Path(args.out_csv or f"phaser_mr_results_{int(time.time())}.csv") + rows = [] + for pdb_key, trial_seed in worklist: + try: + r = run(pdb_key, trial_seed, work_root=work_root, + verbose=args.verbose, timeout_s=args.timeout_s) + rows.append(r) + except Exception as exc: + import traceback + traceback.print_exc() + print(f" TRIAL FAILED on {pdb_key} seed={trial_seed}: {exc!r}", + flush=True) + rows.append({ + "pdb": pdb_key, "seed": trial_seed, + "phaser_solved": False, "phaser_exit": -2, + "error": repr(exc), + }) + + if rows: + # Union of all columns across all rows. + cols = sorted({k for r in rows for k in r}) + with open(out_csv, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=cols) + writer.writeheader() + writer.writerows(rows) + print(f"\nwrote {len(rows)} rows to {out_csv}", flush=True) + + # Quick console summary + print("\n=== Phaser MR summary ===", flush=True) + print( + f"{'pdb':>6} {'seed':>10} {'solved':>6} {'LLG':>8} " + f"{'TFZ':>6} {'r_torchref':>10} {'err°':>7} {'time_s':>8}", + flush=True, + ) + for r in rows: + if r.get("error"): + print(f"{r['pdb']:>6} {r['seed']:>10} FAILED: {r['error']}", + flush=True) + continue + def _f(x, default=float('nan')): + return x if x is not None else default + print( + f"{r['pdb']:>6} {r['seed']:>10} " + f"{str(r.get('phaser_solved', False)):>6} " + f"{_f(r.get('phaser_llg')):>8.2f} " + f"{_f(r.get('phaser_tfz')):>6.2f} " + f"{_f(r.get('rwork_torchref_scaler')):>10.4f} " + f"{_f(r.get('err_canonical_deg')):>7.2f} " + f"{r.get('phaser_time_s', 0.0):>8.1f}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/integration/alignment/submit_sweep.sh b/tests/integration/alignment/submit_sweep.sh new file mode 100644 index 00000000..8a2c2f27 --- /dev/null +++ b/tests/integration/alignment/submit_sweep.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# Submit one SLURM job per (PDB × trial) — 11 PDBs × 3 trials = 33 jobs. +# +# Usage: ./submit_sweep.sh [N_TRIALS_PER_PDB] (default 3) + +set -euo pipefail +N=${1:-3} + +PDBS=(1AK5 1DAW 2DQ6 3A5V 3E98 3GR5 3K7M 3VRJ 4BX9 5BOV 6G9X) +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SUBMIT_DIR="$SCRIPT_DIR" +LOG_DIR="$(pwd)/sweep_logs_$(date +%Y%m%d_%H%M%S)" +mkdir -p "$LOG_DIR" +cd "$LOG_DIR" + +# Pick a base seed (fixed for reproducibility across the sweep; per-trial +# seed is base + trial_idx so each (pdb, trial) gets a distinct seed). +BASE_SEED=${BASE_SEED:-42} + +echo "Submitting $((${#PDBS[@]} * N)) jobs into $LOG_DIR" >&2 +JOBIDS=() +for pdb in "${PDBS[@]}"; do + for trial in $(seq 0 $((N - 1))); do + seed=$((BASE_SEED + trial * 1000003)) + jobid=$(sbatch --parsable --job-name="fit_${pdb}_t${trial}" \ + "$SUBMIT_DIR/sweep_slurm.sh" "$pdb" "$seed") + echo " $pdb trial $trial → seed=$seed jobid=$jobid" + JOBIDS+=("$jobid") + done +done + +echo +echo "Submitted ${#JOBIDS[@]} jobs. Track with:" +echo " squeue --user \$USER --jobs=$(IFS=,; echo "${JOBIDS[*]}")" +echo "Logs in: $LOG_DIR" diff --git a/tests/integration/alignment/sweep_slurm.sh b/tests/integration/alignment/sweep_slurm.sh new file mode 100644 index 00000000..97224c20 --- /dev/null +++ b/tests/integration/alignment/sweep_slurm.sh @@ -0,0 +1,27 @@ +#!/bin/bash +#SBATCH --job-name=fit_to_data +#SBATCH --output=slurm-%j.out +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=64G +#SBATCH --time=03:00:00 +#SBATCH --partition=day + +# Usage: sbatch sweep_slurm.sh +# E.g. sbatch sweep_slurm.sh 1AK5 12345 + +set -euo pipefail + +PDB_KEY="${1:?need PDB key}" +SEED="${2:?need seed}" + +REPO=/das/work/units/LBR-FEL/p17490/Peter/Library/work_trees_torchref/fix_alignment +PYTHON=/das/work/units/LBR-FEL/p17490/Peter/Library/work_trees_torchref/fix_alignment/.venv/bin/python + +cd "$REPO" +export PYTHONPATH="$REPO" +export TORCHREF_NUM_THREADS="${SLURM_CPUS_PER_TASK:-8}" + +echo "Running fit on PDB=$PDB_KEY seed=$SEED on $(hostname)" +$PYTHON tests/integration/alignment/run_random_pdb_fit.py \ + --pdb "$PDB_KEY" --seed "$SEED" --n-trials 1 --verbose 1 diff --git a/tests/integration/alignment/test_fit_to_data.py b/tests/integration/alignment/test_fit_to_data.py new file mode 100644 index 00000000..2df12120 --- /dev/null +++ b/tests/integration/alignment/test_fit_to_data.py @@ -0,0 +1,125 @@ +""" +Integration test for ``align_model_to_data``: end-to-end rotation search +returning a re-oriented ModelFT. + +Setup: +- F_obs: real 1DAW.mtz at its native C2 spacegroup. +- Search model: P1 copy of 1DAW.pdb whose atomic coordinates have been + rotated by a random R_true. + +Acceptance: after ``align_model_to_data``, the returned model's atom coordinates are +within 8° rotation distance (modulo C2 symmetry of F_obs) of the un-rotated +canonical orientation, for 5/5 random trials. +""" +from pathlib import Path + +import pytest +import torch + +from torchref.experimental.alignment import align_model_to_data +from torchref.experimental.alignment.frf.rotation_utils import rotation_angular_distance_deg +from torchref.io.datasets.reflection_data import ReflectionData +from torchref.model import ModelFT + + +TEST_FILES = Path(__file__).resolve().parents[2] / "files" +PDB_1DAW = TEST_FILES / "pdb" / "1DAW.pdb" +MTZ_1DAW = TEST_FILES / "mtz" / "1DAW.mtz" + + +def _load_p1_search_model() -> ModelFT: + """Load 1DAW and force spacegroup to P1 via the proper setter.""" + m = ModelFT().load_pdb(str(PDB_1DAW)) + m.spacegroup = "P 1" + return m + + +@pytest.fixture(scope="module") +def real_setup(): + """Real C2 F_obs + P1 search model factory.""" + data = ReflectionData().load_mtz(str(MTZ_1DAW)) + return data, _load_p1_search_model + + +def _random_rotation(seed: int) -> torch.Tensor: + g = torch.Generator().manual_seed(seed) + A = torch.randn(3, 3, generator=g, dtype=torch.float64) + Q, R = torch.linalg.qr(A) + Q = Q @ torch.diag(torch.sign(torch.diag(R))) + if torch.det(Q) < 0: + Q[:, 0] = -Q[:, 0] + return Q + + +def _best_alignment_rotation(xyz_a: torch.Tensor, xyz_b: torch.Tensor) -> torch.Tensor: + """ + Kabsch: return R that minimises ||xyz_a - xyz_b @ R.T|| (with both centred). + Used here to recover the effective rotation between two atom sets. + """ + a = xyz_a - xyz_a.mean(0) + b = xyz_b - xyz_b.mean(0) + H = b.T @ a + U, _, Vt = torch.linalg.svd(H) + d = torch.sign(torch.det(Vt.T @ U.T)) + D = torch.diag(torch.tensor([1.0, 1.0, d], dtype=H.dtype)) + R = Vt.T @ D @ U.T + return R + + +def _cartesian_symops(data) -> torch.Tensor: + """Point-group rotations as Cartesian matrices, ``B S B^-1``. + + ``spacegroup.matrices`` act on fractional coordinates. A Kabsch rotation is + Cartesian, and comparing the two directly is only right when the cell is + orthogonal and the operator diagonal -- in a trigonal cell two of the six + mates of a correct placement read as 30 and 21 degrees. + """ + B = data.cell.fractional_matrix.to(torch.float64) + S = data.spacegroup.matrices.to(torch.float64) + return B @ S @ torch.linalg.inv(B) + + +@pytest.mark.integration +@pytest.mark.slow +@pytest.mark.parametrize("trial", range(5)) +def test_fit_to_data_real_1daw(real_setup, trial): + """ + Apply a random R_true to a P1 search model, align it to the real data, and + verify the returned model is within 8° rotation distance of the + canonical orientation (modulo C2 symmetry of F_obs). + """ + data, make_model = real_setup + sym_mats = _cartesian_symops(data) + + canonical = make_model() + xyz_canonical = canonical.xyz().clone() + centroid = xyz_canonical.mean(0) + + R_true = _random_rotation(seed=5000 + trial) + search = canonical.rotate(R_true.to(canonical.dtype_float), center=centroid) + + aligned = align_model_to_data( + search, + data, + d_min=4.0, d_max=15.0, + n_shells=20, + n_rotation_peaks=200, + do_translation=False, # this test only checks rotation accuracy + verbose=0, + ) + + # The effective rotation between aligned.xyz() and xyz_canonical should + # be a C2 symmetry operator (i.e. nearly identity or 2-fold along b). + R_residual = _best_alignment_rotation( + aligned.xyz().to(torch.float64), xyz_canonical.to(torch.float64), + ) + # Compare R_residual to identity / each C2 op. + best_err = float("inf") + for k in range(sym_mats.shape[0]): + err = rotation_angular_distance_deg(R_residual, sym_mats[k]) + if err < best_err: + best_err = err + assert best_err < 8.0, ( + f"trial {trial}: aligned model is {best_err:.2f}° from a C2-equivalent " + f"canonical orientation" + ) diff --git a/tests/integration/alignment/test_fit_to_data_translation.py b/tests/integration/alignment/test_fit_to_data_translation.py new file mode 100644 index 00000000..3aa390f9 --- /dev/null +++ b/tests/integration/alignment/test_fit_to_data_translation.py @@ -0,0 +1,139 @@ +""" +Integration test for ``align_model_to_data`` with the translation search on. + +Setup: 1DAW.mtz (C2) F_obs + P1 search model. Apply a known rotation and +fractional translation, then ask the pipeline to recover both. Acceptance: the +returned model is within 8 degrees of canonical modulo the C2 symmetry, its +position is within 3 A of a symmetry image of canonical (modulo lattice +translations and the group's allowed origin shifts; y is polar and free), and +its scaled R-work is close to the deposited model's. +""" +from pathlib import Path + +import pytest +import torch + +from torchref.base.metrics.rfactor import rfactor_work_free +from torchref.experimental.alignment import align_model_to_data +from torchref.experimental.alignment.frf.rotation_utils import ( + rotation_angular_distance_deg, + rotation_matrix_from_edmonds_euler, +) +from torchref.io.datasets.reflection_data import ReflectionData +from torchref.model import ModelFT +from torchref.scaling import Scaler + + +TEST_FILES = Path(__file__).resolve().parents[2] / "files" +PDB_1DAW = TEST_FILES / "pdb" / "1DAW.pdb" +MTZ_1DAW = TEST_FILES / "mtz" / "1DAW.mtz" + + +def _scale_and_rwork(model: ModelFT, data) -> float: + scaler = Scaler(model=model, data=data, nbins=20, verbose=0) + fcalc = model(data.hkl) + scaler.initialize(fcalc) + scaler.refine_lbfgs(fcalc=fcalc) + rw, _ = rfactor_work_free(data, torch.abs(scaler.forward(fcalc))) + return rw.item() if hasattr(rw, "item") else float(rw) + + +def _wrap_frac(t: torch.Tensor) -> torch.Tensor: + """Wrap fractional coords into [-0.5, 0.5).""" + return (t + 0.5) % 1.0 - 0.5 + + +def _cartesian_symops(data) -> torch.Tensor: + """Point-group rotations as Cartesian matrices, ``B S B^-1``. + + ``spacegroup.matrices`` act on fractional coordinates. A Kabsch rotation is + Cartesian, and comparing the two directly is only right when the cell is + orthogonal and the operator diagonal -- in a trigonal cell two of the six + mates of a correct placement read as 30 and 21 degrees. + """ + B = data.cell.fractional_matrix.to(torch.float64) + S = data.spacegroup.matrices.to(torch.float64) + return B @ S @ torch.linalg.inv(B) + + +@pytest.mark.integration +@pytest.mark.slow +def test_fit_to_data_recovers_rotation_and_translation(): + data = ReflectionData().load_mtz(str(MTZ_1DAW)) + canonical = ModelFT().load_pdb(str(PDB_1DAW)) + canonical.spacegroup = "P 1" + + # Apply a known random rotation + fractional translation. + R_true = rotation_matrix_from_edmonds_euler(0.6, 0.4, 1.2) + R_apply = R_true.to(canonical.dtype_float) + # .copy() first: Model.rotate mutates in place, so rotating `canonical` + # directly would perturb the very reference this test compares against. + rotated = canonical.copy().rotate(R_apply, center=canonical.xyz().mean(dim=0)) + t_frac_true = torch.tensor([0.18, -0.07, 0.23], dtype=canonical.dtype_float) + perturbed = rotated.translate(t_frac_true, fractional=True) + + rwork_pre = _scale_and_rwork(perturbed, data) + + aligned = align_model_to_data( + perturbed, + data, + d_min=4.0, d_max=15.0, + n_shells=20, + n_rotation_peaks=200, + do_translation=True, + verbose=0, + ) + + # Recovered rotation (modulo C2). Compare canonical vs aligned via centroid. + xyz_canon = canonical.xyz().to(torch.float64) + xyz_aligned = aligned.xyz().to(torch.float64) + c_canon = xyz_canon.mean(dim=0) + c_aligned = xyz_aligned.mean(dim=0) + a = xyz_canon - c_canon + b = xyz_aligned - c_aligned + H = b.T @ a + U, _, Vt = torch.linalg.svd(H) + d = float(torch.sign(torch.det(Vt.T @ U.T))) + D = torch.diag(torch.tensor([1.0, 1.0, d], dtype=H.dtype)) + R_residual = Vt.T @ D @ U.T + sym_cart = _cartesian_symops(data) + errs = [rotation_angular_distance_deg(R_residual, sym_cart[k]) + for k in range(sym_cart.shape[0])] + k_best = min(range(len(errs)), key=errs.__getitem__) + best_rot_err = errs[k_best] + assert best_rot_err < 8.0, ( + f"residual rotation {best_rot_err:.2f}° > 8° gate" + ) + + # The translation, against the symmetry image whose rotation matched. In + # C2 the origin is free along y, the centring makes (1/2, 1/2, 0) a lattice + # vector, and (0, *, 1/2) is an allowed origin shift -- so x and z are each + # determined only modulo 1/2 and y not at all. A placement at the right + # orientation and 40 A from the true position used to pass this test. + B = data.cell.fractional_matrix.to(torch.float64) + Binv = torch.linalg.inv(B) + S = data.spacegroup.matrices.to(torch.float64) + T = data.spacegroup.translations.to(torch.float64) + c_a = Binv @ c_aligned + c_c = S[k_best] @ (Binv @ c_canon) + T[k_best] + delta = c_a - c_c + delta_xz = (delta + 0.25) % 0.5 - 0.25 + delta_xz[1] = 0.0 + trans_A = float((B @ delta_xz).norm()) + assert trans_A < 3.0, f"placed {trans_A:.1f} A from a symmetry image of canonical" + + # The crystallographic check that this is a valid solution is the R-factor + # of the scaled model. + # The translation function brings R-work close to the canonical-native + # reference (0.21 for 1DAW). The residual gap (~0.12) is from the + # rotation function's ~2° angular error — a separate refinement that's + # not part of the translation function. A 1.87° rotation residual on a + # 100Å molecule moves atoms by ~3Å, which costs ~0.12 in R-work even + # with the exactly correct translation. + rwork_post = _scale_and_rwork(aligned, data) + canonical_native = ModelFT().load_pdb(str(PDB_1DAW)) # native C2 + rwork_ref = _scale_and_rwork(canonical_native, data) + assert rwork_post < rwork_ref + 0.18, ( + f"R-work {rwork_post:.4f} > reference {rwork_ref:.4f} + 0.18 " + f"(pre-fit was {rwork_pre:.4f})" + ) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 8bdf5d61..424307e5 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -1,7 +1,6 @@ -""" -Integration test specific fixtures. -Integration tests use real file I/O and test the full pipeline. +"""Use shared fixtures registered by the root conftest for integration tests. -All shared fixtures (sample files, path fixtures, monomer library, etc.) -are defined in the root tests/conftest.py and are automatically available here. +Pipeline-specific fixtures belong in their consuming modules. Mutable loaded +objects from ``tests.fixtures.objects`` are function-scoped unless documented +as explicitly shared. """ diff --git a/tests/integration/test_adp_field_refinement.py b/tests/integration/test_adp_field_refinement.py new file mode 100644 index 00000000..4dca8c03 --- /dev/null +++ b/tests/integration/test_adp_field_refinement.py @@ -0,0 +1,398 @@ +"""A node-field ADP representation driven the way production drives it. + +The unit tests cover the payload arithmetic and the wiring. What they cannot cover is +whether the thing is usable: whether a refinement constructed with a field mode sizes +itself from the data, carries weights appropriate to a field rather than to the per-atom +representation it replaced, actually reduces R-work over real cycles, survives a +checkpoint round trip, and can be switched into and out of after setup. + +Every one of those was broken or absent at some point in this feature's life, and none +of them fails a unit test. +""" + +import pytest +import torch + +from torchref.model.disorder_field import ( + DisorderFieldTensor, + ModeCovariancePayload, + payload_code, +) +from torchref.refinement.base_refinement import DEFAULT_GROUP_WEIGHTS +from torchref.refinement.lbfgs_refinement import LBFGSRefinement + +MODE_SET = "rigid_dilation" + + +@pytest.fixture(scope="module") +def files(mtz_dir, pdb_dir): + return str(mtz_dir / "1DAW.mtz"), str(pdb_dir / "1DAW.pdb") + + +def _refinement(files, **kw): + mtz, pdb = files + return LBFGSRefinement(data_file=mtz, pdb=pdb, verbose=0, **kw) + + +@pytest.fixture(scope="module") +def field_refinement(files): + """Built the way a caller would: mode and mode set, no explicit node count.""" + return _refinement(files, adp_mode="field_aniso", adp_mode_set=MODE_SET) + + +# ---------------------------------------------------------------------------------- +# Setup: sized from the data, weighted for a field. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_construction_installs_the_mode_field(field_refinement): + ref = field_refinement + assert ref.model.adp_is_field + field = ref.model.adp_field + assert isinstance(field.payload, ModeCovariancePayload) + assert field.payload.mode_set == MODE_SET + + +@pytest.mark.integration +def test_node_count_comes_from_the_reflections_not_the_atoms(field_refinement): + """The whole point of putting the sizing on the refinement. + + The model's own default is one node per 25 atoms, which knows nothing about how much + data there is. Here the achieved ratio has to land near the requested one, and the + node count has to disagree with the atom-count rule (otherwise the test would pass + on a coincidence). + """ + ref = field_refinement + n_par = sum(p.numel() for p in ref.model.parameters_of_types(("adp", "u"))) + achieved = ref.data.work.n / n_par + assert 4.0 < achieved < 12.0, f"{achieved:.1f} work reflections per ADP parameter" + + atom_rule = max(4, round(len(ref.model.pdb) / 25.0)) + assert ref.model.adp_field.n_nodes != atom_rule + + +@pytest.mark.integration +@pytest.mark.parametrize("ratio", [3.5, 7.0, 15.0]) +def test_requested_ratio_is_honoured(files, ratio): + ref = _refinement( + files, adp_mode="field_aniso", adp_mode_set=MODE_SET, + reflections_per_adp_parameter=ratio, + ) + n_par = sum(p.numel() for p in ref.model.parameters_of_types(("adp", "u"))) + achieved = ref.data.work.n / n_par + # Node count is an integer, so the achieved ratio cannot match exactly; it must + # track, and it must be monotone in the request. + assert 0.6 * ratio < achieved < 1.7 * ratio, f"asked {ratio}, got {achieved:.1f}" + + +@pytest.mark.integration +def test_explicit_node_count_bypasses_the_budget(files): + ref = _refinement( + files, adp_mode="field_aniso", adp_mode_set=MODE_SET, n_nodes=11 + ) + assert ref.model.adp_field.n_nodes == 11 + + +@pytest.mark.integration +def test_field_mode_does_not_register_the_restraints_it_duplicates(field_refinement): + """``simu`` and ``locality`` penalise what the field enforces by construction. + + They must be absent from the component set, not present at weight zero: a zero is a + lever, and anyone adjusting the ``adp`` group weight for their own reasons would + silently re-enable a restraint that double-counts the parametrisation. + """ + components = field_refinement.adp_target.target_losses() + assert "simu" not in components + assert "locality" not in components + # What a field does need. + assert "node_load" in components + assert "node_smoothness" in components + assert "sigd" in components, "the marginal-B prior applies to either representation" + + # And the loss is NOT rebalanced for a field: the parametrisation is the constraint, + # so a field needs less regularisation than a per-atom model, not a reweighted + # version of the same priors. An earlier override also silently scaled the two + # adp/scaler_* terms, which have nothing to do with atomic ADPs. + weights = field_refinement.weighting() + for key, expected in DEFAULT_GROUP_WEIGHTS.items(): + assert weights.get(key) == expected, f"{key} diverged from the default" + + +@pytest.mark.integration +def test_per_atom_mode_does_not_register_the_node_targets(files): + """The converse: nothing node-shaped has anything to act on off field mode.""" + ref = _refinement(files, adp_mode="isotropic") + components = ref.adp_target.target_losses() + assert "simu" in components and "locality" in components + assert "node_load" not in components + assert "node_smoothness" not in components + + +@pytest.mark.integration +def test_switching_replaces_the_component_set_and_the_loss_state(files): + """A switch changes WHICH targets exist, so a cached LossState must not survive it.""" + ref = _refinement(files, adp_mode="isotropic") + before = set(ref.adp_target.target_losses()) + assert "simu" in before + # Force the LossState to exist so the switch has something stale to invalidate. + ref.complete_loss_state() + assert ref._loss_state is not None + + logger_before = ref.logger # binds a Logger to the state that is about to go + ref.set_adp_representation("field_aniso", mode_set=MODE_SET) + after = set(ref.adp_target.target_losses()) + # The Logger holds a reference to the LossState, so replacing the state without + # replacing the Logger would leave it recording into an object nothing else reads. + assert ref.logger is not logger_before + assert ref.logger.state is ref.loss_state + assert "simu" not in after and "node_load" in after + state = ref.complete_loss_state() + registered = set(state.targets) + assert not any(k.endswith("simu") or k.endswith("locality") for k in registered), ( + f"a stale component survived the switch: {sorted(registered)}" + ) + assert any(k.endswith("node_load") for k in registered) + + +@pytest.mark.integration +def test_switching_preserves_weights_it_does_not_own(files): + """A call about ADPs must not reset the caller's xray or geometry weights.""" + from torchref.refinement.weighting import ManualWeighting + + ref = _refinement(files, adp_mode="isotropic") + custom = {**ref.weighting(), "xray": 2.5, "geometry": 0.35} + ref.weighting = ManualWeighting(custom) + + ref.set_adp_representation("field_aniso", mode_set=MODE_SET) + weights = ref.weighting() + assert weights["xray"] == 2.5, "xray weight was clobbered by an ADP call" + assert weights["geometry"] == 0.35 + # Nothing is rebalanced, so the caller's adp weight survives too. + assert weights["adp"] == custom["adp"] + + ref.set_adp_representation("isotropic") + assert ref.weighting()["xray"] == 2.5 + + +@pytest.mark.integration +def test_per_atom_mode_keeps_the_per_atom_weights(files): + ref = _refinement(files, adp_mode="isotropic") + assert ref.weighting()["adp"] == DEFAULT_GROUP_WEIGHTS["adp"] + + +# ---------------------------------------------------------------------------------- +# It has to actually refine. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_refinement_reduces_rwork_and_stays_finite(files): + """Two ADP-only cycles on real data. Nothing here may be NaN and R-work must fall.""" + ref = _refinement(files, adp_mode="field_aniso", adp_mode_set=MODE_SET) + rw0, rf0 = (float(v) for v in ref.get_rfactor()) + for _ in range(2): + ref.refine_scaler() + ref.refine_adp() + rw1, rf1 = (float(v) for v in ref.get_rfactor()) + + for name, v in (("Rwork", rw1), ("Rfree", rf1)): + assert v == v, f"{name} is NaN" + assert 0.0 < v < 0.7, f"{name} = {v:.4f} is not a plausible R-factor" + assert rw1 < rw0 + 1e-6, f"R-work rose: {rw0:.4f} -> {rw1:.4f}" + + u6 = ref.model.adp_u6().detach() + assert torch.isfinite(u6).all() + ev = torch.linalg.eigvalsh(_u6_to_matrix(u6)) + assert float(ev.min()) > 0.0, "an atom went non-positive-definite during refinement" + + +@pytest.mark.integration +def test_full_refine_moves_coordinates_without_staling_the_field(files): + """``refine()`` is scaler -> xyz -> ADP, so the field is read at moved coordinates. + + The field borrows the coordinate accessor rather than taking coordinates as an + argument, so its forward cache has to fold them into its key. If it does not, the + ADPs silently come from wherever the atoms used to be. + """ + ref = _refinement(files, adp_mode="field_aniso", adp_mode_set=MODE_SET) + xyz0 = ref.model.xyz().detach().clone() + u0 = ref.model.adp_u6().detach().clone() + + ref.refine(macro_cycles=1) + + xyz1 = ref.model.xyz().detach() + u1 = ref.model.adp_u6().detach() + assert not torch.allclose(xyz0, xyz1), "coordinates did not move, test proves nothing" + assert torch.isfinite(u1).all() + assert not torch.allclose(u0, u1), "ADPs unchanged after xyz moved -- stale cache" + + +def _u6_to_matrix(u6): + M = torch.zeros(u6.shape[0], 3, 3, dtype=u6.dtype) + M[:, 0, 0], M[:, 1, 1], M[:, 2, 2] = u6[:, 0], u6[:, 1], u6[:, 2] + M[:, 0, 1] = M[:, 1, 0] = u6[:, 3] + M[:, 0, 2] = M[:, 2, 0] = u6[:, 4] + M[:, 1, 2] = M[:, 2, 1] = u6[:, 5] + return M + + +# ---------------------------------------------------------------------------------- +# Switching after setup, which the model alone documents as unsupported. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_switch_into_and_out_of_field_mode_after_setup(files): + ref = _refinement(files, adp_mode="isotropic") + assert not ref.model.adp_is_field + before = float(ref.get_rfactor()[0]) + + applied = ref.set_adp_representation("field_aniso", mode_set=MODE_SET) + assert ref.model.adp_is_field + assert applied["n_nodes"] >= 2 + assert "simu" not in ref.adp_target.target_losses() + ref.refine_adp() + assert torch.isfinite(torch.as_tensor(float(ref.get_rfactor()[0]))) + + ref.set_adp_representation("isotropic") + assert not ref.model.adp_is_field + # Leaving must put the per-atom weights back, or the next stage is misweighted. + assert ref.weighting()["adp"] == DEFAULT_GROUP_WEIGHTS["adp"] + ref.refine_adp() + after = float(ref.get_rfactor()[0]) + assert after == after and 0.0 < after < 0.7 + assert before == before + + +@pytest.mark.integration +def test_mode_set_on_an_isotropic_field_is_rejected(files): + ref = _refinement(files, adp_mode="isotropic") + with pytest.raises(ValueError, match="field_aniso"): + ref.set_adp_representation("field", mode_set="rigid") + + +# ---------------------------------------------------------------------------------- +# Checkpoints. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_payload_identity_survives_the_state_dict(field_refinement): + """``state_dict`` holds tensors only, so the payload has to be encoded as one. + + Without the code, a restore infers the payload from the slot and rebuilds a + constant-U field: wrong storage width, wrong ADPs. + """ + sd = field_refinement.model.state_dict() + key = [k for k in sd if k.endswith("payload_code")] + assert key, f"no payload code in the state dict: {sorted(sd)[:8]}..." + assert int(sd[key[0]]) == payload_code(field_refinement.model.adp_field.payload) + + +@pytest.mark.integration +def test_model_state_dict_round_trip_rebuilds_the_same_field(field_refinement): + """A restore must rebuild the payload the code names, not one inferred from the slot.""" + from torchref.model.model import Model + + src = field_refinement.model + # create_from_state_dict restores the values itself; a further load_state_dict + # would be strict against restraint and metadata keys it never builds. + restored = Model.create_from_state_dict(src.state_dict(), verbose=0) + field = restored.adp_field + assert field is not None, "restored model has no field at all" + assert isinstance(field.payload, ModeCovariancePayload) + assert field.payload.mode_set == MODE_SET + assert field.node_shape == src.adp_field.node_shape + assert torch.allclose( + restored.adp_u6().detach(), src.adp_u6().detach(), atol=1e-5 + ) + + +@pytest.mark.integration +def test_model_copy_round_trip_on_a_bare_model(): + """``copy()`` carries the payload and the borrowed accessor, not a deep copy of them. + + Deliberately on a model that has never been handed to a Refinement --- see + :func:`test_copy_after_refinement_setup_is_broken_for_every_representation` for why. + """ + from torchref.model.model import Model + + import os + + here = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + src = Model(verbose=0) + src.load_pdb(os.path.join(here, "files", "pdb", "1DAW.pdb")) + src.set_adp_mode("field_aniso", n_nodes=8, k_neighbors=8, mode_set=MODE_SET) + clone = src.copy() + assert isinstance(clone.adp_field.payload, ModeCovariancePayload) + assert clone.adp_field.payload.mode_set == MODE_SET + assert torch.allclose(clone.adp_u6().detach(), src.adp_u6().detach()) + assert clone.adp_field.refinable_params is not src.adp_field.refinable_params + + +@pytest.mark.integration +@pytest.mark.xfail( + reason="PRE-EXISTING and representation-independent: once a Model has been through " + "Refinement setup, a cache somewhere holds a graph-attached tensor and deepcopy " + "refuses it. Measured identically for adp_mode isotropic, anisotropic and " + "field_aniso, and a bare model copies fine, so the node field is not the cause -- " + "it means no refinement of any kind can currently be checkpointed by copy().", + raises=RuntimeError, + strict=True, +) +@pytest.mark.integration +def test_copy_after_refinement_setup_is_broken_for_every_representation(field_refinement): + field_refinement.model.copy() + + +# ---------------------------------------------------------------------------------- +# The CLI, which is where a flag that does nothing hides best. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_cli_flags_reach_the_refinement(): + """A null control on the plumbing: the parsed values must arrive, not the defaults. + + Five CLI flags in this codebase were once silently no-ops. The check is not that the + argument parses but that a non-default value changes what the refinement does. + """ + import argparse + + from torchref.cli._common import add_adp_mode_arg + + parser = argparse.ArgumentParser() + add_adp_mode_arg(parser) + args = parser.parse_args( + ["--adp-mode", "field_aniso", "--adp-mode-set", "affine", + "--reflections-per-adp-parameter", "3.5", "--adp-nodes", "17"] + ) + assert args.adp_mode == "field_aniso" + assert args.adp_mode_set == "affine" + assert args.reflections_per_adp_parameter == 3.5 + assert args.adp_nodes == 17 + + defaults = parser.parse_args([]) + assert defaults.adp_mode == "isotropic" + assert defaults.adp_mode_set is None + assert defaults.adp_nodes is None + + +@pytest.mark.integration +def test_cli_field_mode_end_to_end(files, tmp_path): + """The parsed flags, through the real constructor, produce a real field.""" + mtz, pdb = files + ref = LBFGSRefinement( + data_file=mtz, pdb=pdb, verbose=0, + adp_mode="field_aniso", adp_mode_set="affine", n_nodes=9, + ) + assert ref.model.adp_field.payload.mode_set == "affine" + assert ref.model.adp_field.n_nodes == 9 + ref.refine_adp() + out = tmp_path / "out.pdb" + ref.model.update_pdb() + ref.model.write_pdb(str(out)) + assert out.exists() and out.stat().st_size > 0 + text = out.read_text() + assert "ANISOU" in text, "an anisotropic field must write ANISOU records" diff --git a/tests/integration/test_batched_component_fcalcs.py b/tests/integration/test_batched_component_fcalcs.py new file mode 100644 index 00000000..904e45d9 --- /dev/null +++ b/tests/integration/test_batched_component_fcalcs.py @@ -0,0 +1,130 @@ +"""Batched structure factors preserve mixed-model values and Friedel phases.""" + +import pytest +import torch + +from torchref.config import get_default_device, get_float_dtype + +pytestmark = pytest.mark.integration + + +class TestBatchedMatchesTheLoop: + def test_component_stack_matches_per_model_structure_factors( + self, difference_models + ): + dc, mc = difference_models + data = dc["dark"] + + stacked = dc.component_structure_factors(mc, recalc=True) + assert stacked.shape == (mc.n_base_models, len(data.hkl)) + + for k, model in enumerate(mc.base_models): + reference = data.structure_factors(model, recalc=False) + assert torch.equal( + stacked[k], reference + ), f"component {k} differs from data.structure_factors" + + def test_mixture_matches_the_per_timepoint_forward(self, difference_models): + """A batched contraction agrees with each mixed-model forward.""" + dc, mc = difference_models + data = dc["dark"] + + stacked = dc.component_structure_factors(mc, recalc=True) + mixed = mc.mix_component_fcalcs(stacked, mc.get_fractions_matrix()) + assert mixed.shape == (len(mc), len(data.hkl)) + + for row, key in enumerate(mc.keys()): + reference = data.structure_factors(mc[key], recalc=False) + assert torch.allclose( + mixed[row], reference, rtol=1e-6, atol=1e-6 + ), f"timepoint {key!r} differs from its own mixed forward" + + def test_compute_all_fcalc_agrees_on_the_signed_index(self, difference_models): + """``compute_all_fcalc`` takes the caller's indices verbatim, so handed the + signed ones it must reproduce the Friedel-corrected mixture up to the + conjugation that ``component_structure_factors`` applies.""" + dc, mc = difference_models + data = dc["dark"] + + direct = mc.compute_all_fcalc(data._hkl_for_sf(), recalc=True) + corrected = data.conjugate_friedel(direct) + + stacked = dc.component_structure_factors(mc, recalc=False) + mixed = mc.mix_component_fcalcs(stacked, mc.get_fractions_matrix()) + + assert torch.allclose(corrected, mixed, rtol=1e-6, atol=1e-6) + + +@pytest.fixture +def flagged_pair(difference_models): + """Pair deposited models with both signed Miller-index conventions.""" + from torchref import ReflectionData + from torchref.io.datasets.collection import DatasetCollection + + dc_ref, mc = difference_models + src = dc_ref["dark"] + + hkl = src.hkl.clone() + half = torch.zeros(len(hkl), dtype=torch.bool, device=get_default_device()) + half[::2] = True + hkl[half] = -hkl[half] + + data = ReflectionData.from_tensors( + hkl=hkl, + F=src.F.clone(), + F_sigma=src.F_sigma.clone(), + cell=src.cell, + spacegroup=src.spacegroup, + rfree_flags=src.rfree_flags.clone(), + device=get_default_device(), + verbose=0, + ) + + assert data.friedel_flags.any() and (~data.friedel_flags).any() + dc = DatasetCollection(verbose=0, device=get_default_device()) + dc.add_dataset("dark", data, set_as_reference=True) + return dc, mc + + +class TestConventionIsNotSkipped: + + def test_component_stack_is_conjugated_where_flagged(self, flagged_pair): + """``component_structure_factors`` must apply the conjugation, not skip it.""" + dc, mc = flagged_pair + data = dc["dark"] + + stacked = dc.component_structure_factors(mc, recalc=True) + signed = mc.compute_component_fcalcs(data._hkl_for_sf(), recalc=False) + flagged = data.friedel_flags + assert torch.equal(stacked[:, flagged], signed[:, flagged].conj()) + assert torch.equal(stacked[:, ~flagged], signed[:, ~flagged]) + assert not torch.allclose(stacked[:, flagged], signed[:, flagged]) + + +class TestContraction: + def test_weights_matrix_is_applied_row_wise(self, difference_models): + """A transposed einsum would still return the right shape when T == K.""" + dc, mc = difference_models + stacked = dc.component_structure_factors(mc, recalc=True) + + w = torch.tensor( + [[1.0, 0.0], [0.0, 1.0]], + device=get_default_device(), + dtype=get_float_dtype(), + ) + mixed = mc.mix_component_fcalcs(stacked, w) + + assert torch.equal(mixed[0], stacked[0]) + assert torch.equal(mixed[1], stacked[1]) + + def test_gradient_flows_through_the_contraction(self, difference_models): + dc, mc = difference_models + mc.unfreeze_all_fractions() + stacked = dc.component_structure_factors(mc, recalc=True) + w = mc.get_fractions_matrix() + + mc.mix_component_fcalcs(stacked, w).abs().sum().backward() + + grad = mc._activation_logit.grad + assert grad is not None + assert torch.isfinite(grad).all() diff --git a/tests/integration/test_cli_ded_weights.py b/tests/integration/test_cli_ded_weights.py new file mode 100644 index 00000000..8679564b --- /dev/null +++ b/tests/integration/test_cli_ded_weights.py @@ -0,0 +1,299 @@ +"""The registered difference weights through the CLIs. + +Pinned: ``torchref.difference-map`` writes ``DF`` with one mean-one weight column per +scheme and ``KSCALE``; ``torchref.mtz2map`` builds the weighted map from those columns +and the electrons map is the volume-normalised synthesis on the absolute scale; +``torchref.validate-ded`` reports every scheme side by side and records a fallback; +``torchref.difference-refine`` runs with the ``difference_sd`` row and reports the fit. +""" + +import json +import os +import subprocess +import sys + +import numpy as np +import pytest + +pytestmark = [pytest.mark.integration, pytest.mark.slow] + +DIFF_COLUMNS = { + "Fo_dark": "SFAmplitude", + "SIGFo_dark": "Stddev", + "Fo_light": "SFAmplitude", + "SIGFo_light": "Stddev", + "DF": "SFAmplitude", + "SIGDF": "Stddev", + "PHDELWT": "Phase", + "W_IVW": "Weight", + "W_SD": "Weight", + "KSCALE": "MTZReal", + "Fc_dark": "SFAmplitude", + "FreeR_flag_dark": "MTZInt", + "FreeR_flag_light": "MTZInt", +} + + +@pytest.fixture(scope="module") +def pair(mtz_dir, pdb_dir, tmp_path_factory): + """A dark/light pair from 1DAW with a perturbed light state. + + The light amplitudes carry an added difference proportional to ``F`` with a + resolution-dependent power, so the sigma_D fit has signal to find; the dark set + keeps the deposited values. The light model is the dark one shifted by 0.2 A. + """ + import torch + + from torchref import ReflectionData + from torchref.config import get_int_dtype + + mtz = mtz_dir / "1DAW.mtz" + pdb = pdb_dir / "1DAW.pdb" + assert mtz.is_file() and pdb.is_file() + out = tmp_path_factory.mktemp("ded_weights_cli") + + data = ReflectionData(device="cpu", verbose=0).load_mtz(str(mtz)) + n = len(data) + idx = torch.arange(n, dtype=get_int_dtype()) + dark = data.__select__(idx < int(n * 0.97)) + dark.write_mtz(str(out / "dark.mtz")) + + sel = data.__select__(idx >= int(n * 0.03)) + g = torch.Generator().manual_seed(11) + f = sel.F + dss = 1.0 / sel.resolution**2 + change = 0.08 * f * torch.exp(-2.0 * dss) * torch.randn(len(sel), generator=g) + light = ReflectionData.from_tensors( + hkl=sel.hkl, + F=(f + change).clamp(min=0.0), + F_sigma=sel.F_sigma, + cell=sel.cell, + spacegroup=sel.spacegroup, + rfree_flags=sel.rfree_flags, + device="cpu", + verbose=0, + ) + light.write_mtz(str(out / "light.mtz")) + + lines = [] + for line in pdb.read_text().splitlines(): + if line.startswith(("ATOM", "HETATM")): + x = float(line[30:38]) + 0.2 + line = line[:30] + f"{x:8.3f}" + line[38:] + lines.append(line) + (out / "light.pdb").write_text("\n".join(lines) + "\n") + return {"dir": out, "pdb": pdb, "light_pdb": out / "light.pdb"} + + +def _run(project_root, module, *argv, timeout=1800): + script = project_root / "torchref" / "cli" / module + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root) + os.pathsep + env.get("PYTHONPATH", "") + proc = subprocess.run( + [sys.executable, str(script), *map(str, argv)], + capture_output=True, + text=True, + timeout=timeout, + env=env, + ) + assert proc.returncode == 0, ( + f"{module} failed ({proc.returncode})\nstdout tail:\n{proc.stdout[-2000:]}" + f"\nstderr tail:\n{proc.stderr[-3000:]}" + ) + return proc + + +@pytest.fixture(scope="module") +def diff_mtz(project_root, pair): + out = pair["dir"] / "diff.mtz" + _run( + project_root, + "difference_map.py", + "-dm", + pair["pdb"], + "-dsf", + pair["dir"] / "dark.mtz", + "-lsf", + pair["dir"] / "light.mtz", + "--dmin", + "2.2", + "--device", + "cpu", + "--ded-weight", + "sigma_d", + "-v", + "1", + "-o", + out, + ) + return out + + +def _read(path): + import reciprocalspaceship as rs + + return rs.read_mtz(str(path)) + + +def test_difference_map_writes_df_weights_and_scale(diff_mtz): + df = _read(diff_mtz) + assert {c: str(df.dtypes[c]) for c in df.columns} == DIFF_COLUMNS + for col in ("W_IVW", "W_SD"): + w = df[col].to_numpy().astype(float) + assert np.isfinite(w).all() and (w >= 0).all() + assert abs(w.mean() - 1.0) < 1e-4 + assert (df["KSCALE"].to_numpy().astype(float) > 0).all() + # The sigma_D weights favour the strong reflections, inverse variance does not. + f = df["Fo_dark"].to_numpy().astype(float) + w_sd = df["W_SD"].to_numpy().astype(float) + strong = f > np.median(f) + assert w_sd[strong].mean() > w_sd[~strong].mean() + + +def test_mtz2map_builds_the_weighted_and_electron_maps(project_root, pair, diff_mtz): + import gemmi + + out = pair["dir"] + _run( + project_root, + "mtz2map.py", + "-sf", + diff_mtz, + "-csf", + "DF", + "-cw", + "W_SD", + "-cphi", + "PHDELWT", + "--device", + "cpu", + "-o", + out / "sd_sigma.ccp4", + ) + _run( + project_root, + "mtz2map.py", + "-sf", + diff_mtz, + "-csf", + "DF", + "-cw", + "W_SD", + "-cphi", + "PHDELWT", + "--units", + "raw", + "--device", + "cpu", + "-o", + out / "sd_raw.ccp4", + ) + _run( + project_root, + "mtz2map.py", + "-sf", + diff_mtz, + "-csf", + "DF", + "-cw", + "W_SD", + "-cphi", + "PHDELWT", + "--units", + "electrons", + "--device", + "cpu", + "-o", + out / "sd_e.ccp4", + ) + sigma = np.array(gemmi.read_ccp4_map(str(out / "sd_sigma.ccp4")).grid, copy=False) + raw = np.array(gemmi.read_ccp4_map(str(out / "sd_raw.ccp4")).grid, copy=False) + electrons = np.array(gemmi.read_ccp4_map(str(out / "sd_e.ccp4")).grid, copy=False) + assert abs(sigma.std() - 1.0) < 1e-3 and abs(sigma.mean()) < 1e-3 + # Same map up to normalisation: the correlation is one. + assert np.corrcoef(sigma.ravel(), raw.ravel())[0, 1] > 0.9999 + # Dividing by the per-reflection KSCALE reshapes the map slightly, so electrons is + # highly but not perfectly correlated with the sigma map. + assert 0.9 < np.corrcoef(sigma.ravel(), electrons.ravel())[0, 1] < 0.9999 + ratio = electrons.std() / raw.std() + assert np.isfinite(ratio) and ratio > 0 + + +def test_validate_ded_reports_every_scheme(project_root, pair): + out = pair["dir"] / "val" + proc = _run( + project_root, + "validate_ded.py", + "-dsf", + pair["dir"] / "dark.mtz", + "-lsf", + pair["dir"] / "light.mtz", + "-dm", + pair["pdb"], + "-lm", + pair["light_pdb"], + "--fraction", + "0.3", + "--dmin", + "2.2", + "--device", + "cpu", + "--ded-weight", + "sigma_d", + "-v", + "1", + "-o", + out, + ) + results = json.loads((out / "validate_ded_results.json").read_text()) + assert results["weights"]["requested"] == "sigma_d" + assert results["weights"]["applied"] in ("sigma_d", "inverse_variance") + assert set(results["by_weight"]) == {"none", "inverse_variance", "sigma_d"} + for entry in results["by_weight"].values(): + assert np.isfinite(entry["reciprocal_cc_overall"]) + assert "full_cell" in entry["realspace_correlation"] + headline = results["by_weight"][results["weights"]["applied"]] + assert results["reciprocal_cc_overall"] == pytest.approx( + headline["reciprocal_cc_overall"], abs=1e-3 + ) + assert "weights " in proc.stdout and "sigma_d" in proc.stdout + + +def test_difference_refine_runs_the_sigma_d_row(project_root, pair): + out = pair["dir"] / "refine" + _run( + project_root, + "collection_difference_refine.py", + "-dm", + pair["pdb"], + "-lm", + pair["light_pdb"], + "-dsf", + pair["dir"] / "dark.mtz", + "-lsf", + pair["dir"] / "light.mtz", + "--fraction", + "0.25", + "--difference-target", + "difference_sd", + "--n-cycles", + "1", + "--n-steps", + "1", + "--max-iter", + "3", + "--dmin", + "2.2", + "--device", + "cpu", + "--verbose", + "0", + "-o", + out, + ) + summaries = list(out.glob("*_summary.json")) + assert len(summaries) == 1 + results = json.loads(summaries[0].read_text())["results"] + assert results["ded_weights"]["scheme"] == "inverse_variance" + assert results["ded_weights"]["applied"] == "inverse_variance" + assert "gamma" in results["ded_weights"]["sigma_d"] diff --git a/tests/integration/test_cli_hydrogens.py b/tests/integration/test_cli_hydrogens.py new file mode 100644 index 00000000..f29b2235 --- /dev/null +++ b/tests/integration/test_cli_hydrogens.py @@ -0,0 +1,107 @@ +"""Exercise hydrogen opt-in from CLI parsing through deposited-model loading.""" + +import sys +from pathlib import Path + +import pytest + +from torchref.refinement.base_refinement import Refinement +from torchref.refinement.lbfgs_refinement import LBFGSRefinement + + +class _ModelLoaded(Exception): + """Stop after real model loading, before scaling and optimization.""" + + +@pytest.mark.integration +@pytest.mark.parametrize("add_hydrogens", [False, True]) +@pytest.mark.parametrize("model_format", ["pdb", "cif"]) +def test_cli_hydrogen_generation_is_opt_in( + test_files_dir: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + add_hydrogens: bool, + model_format: str, +) -> None: + """The CLI flag generates missing hydrogens for PDB and mmCIF inputs.""" + from torchref.cli import refine + + loaded = [] + + def stop_after_loading(refinement: Refinement) -> None: + loaded.append(refinement.model) + raise _ModelLoaded + + monkeypatch.setattr(Refinement, "_sync_model_cell_to_data", stop_after_loading) + argv = [ + "torchref.refine", + "-m", + str(test_files_dir / model_format / f"1DAW.{model_format}"), + "-sf", + str(test_files_dir / "mtz" / "1DAW.mtz"), + "-o", + str(tmp_path / "refined"), + "-v", + "0", + ] + if add_hydrogens: + argv.append("--add-hydrogens") + monkeypatch.setattr(sys, "argv", argv) + + with pytest.raises(_ModelLoaded): + refine.main() + + (model,) = loaded + assert model.ctx.add_hydrogens is add_hydrogens + assert len(model.pdb) > 0 + n_hydrogens = int(model.pdb["element"].str.strip().eq("H").sum()) + assert (n_hydrogens > 0) is add_hydrogens + + +@pytest.mark.integration +def test_cli_generates_from_the_user_cif( + test_files_dir: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """``--cif`` reaches the model before it loads, so generation reads it.""" + from torchref.cli import refine + + loaded = [] + + def stop_after_loading(refinement: Refinement) -> None: + loaded.append(refinement.model) + raise _ModelLoaded + + monkeypatch.setattr(Refinement, "_sync_model_cell_to_data", stop_after_loading) + cif = str(test_files_dir / "restraints" / "GLU_renamed.cif") + argv = [ + "torchref.refine", + "-m", + str(test_files_dir / "pdb" / "1DAW.pdb"), + "-sf", + str(test_files_dir / "mtz" / "1DAW.mtz"), + "-o", + str(tmp_path / "refined"), + "-v", + "0", + "--add-hydrogens", + "--cif", + cif, + ] + monkeypatch.setattr(sys, "argv", argv) + with pytest.raises(_ModelLoaded): + refine.main() + + (model,) = loaded + registered = model.ctx.cif_path + assert (registered if isinstance(registered, list) else [registered]) == [cif] + pdb = model.pdb + glu_h = (pdb["resname"].str.strip() == "GLU") & (pdb["element"].str.strip() == "H") + assert {"HAX", "HBX", "HBY", "HGX", "HGY"} <= set(pdb.loc[glu_h, "name"].str.strip()) + + +@pytest.mark.unit +@pytest.mark.parametrize("add_hydrogens", [False, True]) +def test_empty_refinement_preserves_hydrogen_setting(add_hydrogens: bool) -> None: + """An empty refinement shell forwards the setting to its model too.""" + refinement = LBFGSRefinement(verbose=0, add_hydrogens=add_hydrogens) + assert refinement.model.ctx.add_hydrogens is add_hydrogens diff --git a/tests/integration/test_cli_two_moment_mtz.py b/tests/integration/test_cli_two_moment_mtz.py new file mode 100644 index 00000000..01270bae --- /dev/null +++ b/tests/integration/test_cli_two_moment_mtz.py @@ -0,0 +1,384 @@ +"""CLI output column types, phase conventions and two-moment numerical consistency.""" + +import json +import os +import subprocess +import sys + +import pytest + +pytestmark = [pytest.mark.integration, pytest.mark.slow] + + +# The default set, with a light model supplied (which the refinement CLI always does). +# H/K/L are the index, so they are not among ``.columns``. +DEFAULT_COLUMNS = { + "Fo_dark": "SFAmplitude", + "SIGFo_dark": "Stddev", + "Fo_light": "SFAmplitude", + "SIGFo_light": "Stddev", + "DF": "SFAmplitude", + "SIGDF": "Stddev", + # The difference map: DF on the dark phases, one weight column per registered + # scheme, and the observed-to-model scale. + "PHDELWT": "Phase", + "W_IVW": "Weight", + "W_SD": "Weight", + "KSCALE": "MTZReal", + "Fc_dark": "SFAmplitude", + # The mixed model, and the extrapolated map to refine against. + "FC": "SFAmplitude", + "PHIC": "Phase", + "FEXT": "SFAmplitude", + "SIGFEXT": "Stddev", + "FWT": "SFAmplitude", + "PHWT": "Phase", +} +FLAG_COLUMNS = {"FreeR_flag_dark", "FreeR_flag_light"} + +TWO_MOMENT_COLUMNS = { + "DELFWT_corr": "SFAmplitude", + "Fo_light_corr": "SFAmplitude", + "SIGFo_light_corr": "Stddev", + "DF_corr": "SFAmplitude", + "SIGDF_corr": "Stddev", + "DDF": "SFAmplitude", +} + +# What ``--all-columns`` adds on top, given a light model. +ALL_COLUMNS_EXTRA = { + "2mDFop-DFc": "SFAmplitude", + "mDFop-DFc": "SFAmplitude", + "PHIC_diff": "Phase", + "DFc": "SFAmplitude", + "DFc_phased": "SFAmplitude", + "FEXT_PHASED": "SFAmplitude", + "SIGFEXT_PHASED": "Stddev", + "2FEXT_PHASED-Fc": "SFAmplitude", + "FEXT_PHASED-Fc": "SFAmplitude", + "PHFEXT_PHASED": "Phase", + "FEXT_SCALAR": "SFAmplitude", + "SIGFEXT_SCALAR": "Stddev", + "2FEXT_SCALAR-Fc": "SFAmplitude", + "FEXT_SCALAR-Fc": "SFAmplitude", + "PHFEXT_SCALAR": "Phase", +} + +# And what it adds again once the two-moment model is on. +ALL_COLUMNS_TWO_MOMENT_EXTRA = { + "Io_light": "Intensity", + "SIGIo_light": "Stddev", + "Ic_light_coh": "Intensity", + "Ic_light_2mom": "Intensity", + "IVAR_ALPHA": "Intensity", + "W_2MOM": "Weight", + "2mDFop-DFc_corr": "SFAmplitude", + "mDFop-DFc_corr": "SFAmplitude", +} + +FRACTION = 0.25 +LAMBDA_TWIN = 0.2 + + +@pytest.fixture(scope="module") +def cli_script(project_root): + script = project_root / "torchref" / "cli" / "collection_difference_refine.py" + assert script.is_file() + return script + + +@pytest.fixture(scope="module") +def intensity_pair(mtz_dir, pdb_dir, tmp_path_factory): + """Write a dark/light I/SIGI pair from deposited 1DAW observations.""" + import torch + + from torchref import ReflectionData + from torchref.config import get_int_dtype + + mtz = mtz_dir / "1DAW.mtz" + pdb = pdb_dir / "1DAW.pdb" + assert mtz.is_file() and pdb.is_file() + + data = ReflectionData(device="cpu", verbose=0).load_mtz(str(mtz)) + assert data.I is not None + + out = tmp_path_factory.mktemp("two_moment_cli") + n = len(data) + idx = torch.arange(n, dtype=get_int_dtype(), device=data.device) + # Slightly different reflection sets, as a real dark/light pair would be. + data.__select__(idx < int(n * 0.97)).write_mtz(str(out / "dark.mtz")) + data.__select__(idx >= int(n * 0.03)).write_mtz(str(out / "light.mtz")) + return {"dir": out, "pdb": pdb} + + +def _run(cli_script, pair, outdir, *extra): + env = dict(os.environ) + # The installed torchref may point at a different checkout; make the subprocess + # import the tree under test. + root = str(cli_script.parents[2]) + env["PYTHONPATH"] = root + os.pathsep + env.get("PYTHONPATH", "") + + cmd = [ + sys.executable, + str(cli_script), + "-dm", + str(pair["pdb"]), + "-lm", + str(pair["pdb"]), + "-dsf", + str(pair["dir"] / "dark.mtz"), + "-lsf", + str(pair["dir"] / "light.mtz"), + "--fraction", + str(FRACTION), + "--n-cycles", + "1", + "--n-steps", + "1", + "--max-iter", + "3", + "--dmin", + "2.2", + "-o", + str(outdir), + "--device", + "cpu", + "--verbose", + "0", + *extra, + ] + proc = subprocess.run(cmd, capture_output=True, text=True, timeout=1800, env=env) + assert ( + proc.returncode == 0 + ), f"CLI failed ({proc.returncode})\nstderr tail:\n{proc.stderr[-3000:]}" + prefix = f"fractions_{round((1 - FRACTION) * 100)}_{round(FRACTION * 100)}_" + return outdir / f"{prefix}difference_data.mtz", outdir / f"{prefix}summary.json" + + +@pytest.fixture(scope="module") +def baseline_mtz(cli_script, intensity_pair, tmp_path_factory): + outdir = tmp_path_factory.mktemp("baseline") + return _run(cli_script, intensity_pair, outdir) + + +@pytest.fixture(scope="module") +def two_moment_mtz(cli_script, intensity_pair, tmp_path_factory): + outdir = tmp_path_factory.mktemp("two_moment") + return _run( + cli_script, + intensity_pair, + outdir, + "--two-moment", + "--lambda-twin", + str(LAMBDA_TWIN), + ) + + +@pytest.fixture(scope="module") +def two_moment_all_mtz(cli_script, intensity_pair, tmp_path_factory): + """Everything on. The value-consistency tests below need the diagnostic columns, + which is exactly what ``--all-columns`` is for.""" + outdir = tmp_path_factory.mktemp("two_moment_all") + return _run( + cli_script, + intensity_pair, + outdir, + "--two-moment", + "--lambda-twin", + str(LAMBDA_TWIN), + "--all-columns", + ) + + +def _read(path): + import reciprocalspaceship as rs + + return rs.read_mtz(str(path)) + + +@pytest.mark.parametrize( + "fixture,extra", + [ + ("baseline_mtz", {}), + ("two_moment_mtz", TWO_MOMENT_COLUMNS), + ( + "two_moment_all_mtz", + {**TWO_MOMENT_COLUMNS, **ALL_COLUMNS_EXTRA, **ALL_COLUMNS_TWO_MOMENT_EXTRA}, + ), + ], +) +def test_column_layout_and_types(request, fixture, extra): + """Each CLI mode writes exactly its documented columns with MTZ types.""" + frame = _read(request.getfixturevalue(fixture)[0]) + expected = {**DEFAULT_COLUMNS, **extra} + assert set(frame.columns) == set(expected) | FLAG_COLUMNS + for name, dtype in expected.items(): + assert frame.dtypes[name].name == dtype + assert all(hasattr(dtype, "mtztype") for dtype in frame.dtypes) + + +def test_map_coefficients_are_grouped_into_described_datasets(two_moment_all_mtz): + """``FWT``/``PHWT`` keep the label Coot auto-opens but sit in the + ``extrapolated_light`` dataset, and every dataset has a history line saying what it + holds -- the file, not the label, says which map a standard name is.""" + import gemmi + + mtz = gemmi.read_mtz_file(str(two_moment_all_mtz[0])) + names = {ds.id: ds.dataset_name for ds in mtz.datasets} + where = {col.label: names[col.dataset_id] for col in mtz.columns} + + assert where["FWT"] == where["PHWT"] == where["FEXT"] == "extrapolated_light" + assert where["FEXT_PHASED"] == where["FEXT_SCALAR"] == "extrapolated_light" + assert where["DF"] == where["PHDELWT"] == where["W_SD"] == "difference" + assert where["FC"] == where["PHIC"] == "light_model" + assert where["DF_corr"] == "two_moment" + assert where["H"] == where["Fo_dark"] == where["FreeR_flag_dark"] == "observed" + assert all(ds.crystal_name == "torchref" for ds in mtz.datasets) + + assert all(len(line) <= 80 for line in mtz.history) + for name in set(names.values()): + assert any(line.startswith(f"{name}:") for line in mtz.history), name + + +class TestDefaultLayout: + + def test_the_difference_columns_carry_df_and_the_registered_weights( + self, baseline_mtz + ): + """``DF`` must be ``Fo_light - Fo_dark``, ``W_IVW`` the mean-normalised inverse + variance and ``W_SD`` a mean-one weight -- the constructions + ``torchref.validate-ded`` correlates against. If these ever diverge, the map + built from the file stops being the map the validation reports on, which is how + the output drifted from the science before.""" + import numpy as np + + df = _read(baseline_mtz[0]) + dfo = df["Fo_light"].to_numpy().astype(float) - df["Fo_dark"].to_numpy().astype( + float + ) + sig = np.sqrt( + df["SIGFo_dark"].to_numpy().astype(float) ** 2 + + df["SIGFo_light"].to_numpy().astype(float) ** 2 + ) + w = 1 / np.maximum(sig, 0.1 * np.median(sig)) ** 2 + w = w / w.mean() + + got_df = df["DF"].to_numpy().astype(float) + scale = max(float(np.abs(dfo).max()), 1e-30) + assert np.abs(got_df - dfo).max() / scale < 1e-5 + got_w = df["W_IVW"].to_numpy().astype(float) + assert np.abs(got_w - w).max() / max(float(np.abs(w).max()), 1e-30) < 1e-4 + w_sd = df["W_SD"].to_numpy().astype(float) + assert np.isfinite(w_sd).all() and abs(w_sd.mean() - 1.0) < 1e-4 + assert (df["KSCALE"].to_numpy().astype(float) > 0).all() + + # And the phase is the dark model's, not the mixed model's. + assert not np.allclose( + df["PHDELWT"].to_numpy().astype(float), + df["PHIC"].to_numpy().astype(float), + ) + + +class TestTwoMomentLayout: + + def test_the_corrected_difference_map_pairs_with_the_same_phases( + self, two_moment_mtz + ): + """``DELFWT_corr`` is the corrected difference on the *same* dark phases, so it + is opened against ``PHDELWT`` and carries the selected weight scheme, the + inverse-variance weights ``W_IVW`` by default.""" + import numpy as np + + df = _read(two_moment_mtz[0]) + w = df["W_IVW"].to_numpy().astype(float) + + expected = df["DF_corr"].to_numpy().astype(float) * w + got = df["DELFWT_corr"].to_numpy().astype(float) + scale = max(float(np.abs(expected).max()), 1e-30) + assert np.abs(got - expected).max() / scale < 1e-5 + + +class TestTwoMomentValuesAreConsistent: + def test_ivar_alpha_is_sigma_sq_times_the_squared_difference( + self, two_moment_all_mtz + ): + """The variance column must be the quantity it claims, not a rescaling of it.""" + import numpy as np + + mtz, summary = two_moment_all_mtz + df = _read(mtz) + results = json.loads(summary.read_text())["results"] + + sigma_sq = results["sigma_alpha_sq"] + dfc = df["DFc_phased"].to_numpy().astype(float) + ivar = df["IVAR_ALPHA"].to_numpy().astype(float) + + expected = sigma_sq * dfc**2 + scale = max(float(np.abs(expected).max()), 1e-30) + assert np.abs(ivar - expected).max() / scale < 1e-5 + + def test_the_two_moment_intensity_exceeds_the_coherent_one_by_the_variance( + self, two_moment_all_mtz + ): + """``Ic_2mom - Ic_coh`` must equal ``IVAR_ALPHA``, to whatever precision float32 + leaves after the cancellation.""" + import numpy as np + + df = _read(two_moment_all_mtz[0]) + coh = df["Ic_light_coh"].to_numpy().astype(float) + two = df["Ic_light_2mom"].to_numpy().astype(float) + ivar = df["IVAR_ALPHA"].to_numpy().astype(float) + + # Absolute error float32 can leave in the difference of two intensities. + eps32 = float(np.finfo(np.float32).eps) + floor = eps32 * np.maximum(np.abs(coh), np.abs(two)) + residual = np.abs((two - coh) - ivar) + + assert (residual <= 4.0 * floor + 1e-12).all(), ( + f"recovered variance term differs from IVAR_ALPHA by more than float32 " + f"cancellation allows: worst {np.max(residual / (floor + 1e-30)):.1f} ulp" + ) + # The variance term has no sign: it can only add. + assert (two >= coh - 4.0 * floor).all() + + def test_the_weight_is_the_contamination_ratio(self, two_moment_all_mtz): + """``W_2MOM`` must be ``sigma_I**2 / (sigma_I**2 + IVAR_ALPHA)``.""" + import numpy as np + + df = _read(two_moment_all_mtz[0]) + w = df["W_2MOM"].to_numpy().astype(float) + sig = df["SIGIo_light"].to_numpy().astype(float) + ivar = df["IVAR_ALPHA"].to_numpy().astype(float) + + assert (w > 0).all() and (w <= 1.0 + 1e-6).all() + expected = sig**2 / np.maximum(sig**2 + ivar, 1e-12) + assert np.allclose(w, expected, rtol=1e-5, atol=1e-7) + # Anti-vacuity for the formula: the contamination must not be identically zero, + # or the ratio above is trivially 1 and proves nothing. + assert (ivar > 0).any() + + def test_the_correction_moves_the_difference_amplitudes(self, two_moment_mtz): + """DDF is the diagnostic; if it were identically zero the whole column set + would be decorative.""" + import numpy as np + + df = _read(two_moment_mtz[0]) + ddf = df["DDF"].to_numpy().astype(float) + assert np.count_nonzero(ddf) > 0.5 * len(ddf) + # Subtracting a positive contamination lowers the light amplitude on average. + assert ddf.mean() < 0.0 + + def test_summary_reports_the_activation_moments(self, two_moment_mtz): + _, summary = two_moment_mtz + results = json.loads(summary.read_text())["results"] + for key in ("alpha_mean", "lambda_twin", "sigma_alpha_sq"): + assert key in results, f"summary is missing {key}" + assert results["lambda_twin"] == pytest.approx(LAMBDA_TWIN) + + def test_summary_reports_the_shrinkage_diagnostics(self, two_moment_mtz): + """``tau_sq`` and mean ``w(h)`` say whether the default extrapolated map is + over-shrunk, so they belong in the summary rather than only in a print.""" + _, summary = two_moment_mtz + results = json.loads(summary.read_text())["results"] + assert "tau_sq" in results and "w_shrinkage_mean" in results + assert 0.0 < results["w_shrinkage_mean"] <= 1.0 diff --git a/tests/integration/test_collection_joint_scale_fit.py b/tests/integration/test_collection_joint_scale_fit.py new file mode 100644 index 00000000..d03c2d01 --- /dev/null +++ b/tests/integration/test_collection_joint_scale_fit.py @@ -0,0 +1,79 @@ +"""Joint model-to-data scaling objectives and shared parameter ownership.""" + +import inspect + +import pytest +import torch + +pytestmark = pytest.mark.integration + + +@pytest.fixture +def collection(loaded_model_ft, loaded_reflection_data): + """One structural model paired with independent dark/timepoint observations.""" + from torchref.io import DatasetCollection + from torchref.model import ModelCollection + + data = loaded_reflection_data + dc = DatasetCollection(device=data.device, verbose=0) + dc.add_dataset("dark", data, set_as_reference=True).add_dataset("t1", data) + mc = ModelCollection([loaded_model_ft], dark_key="dark", verbose=0) + mc.add_dark().add_timepoint("t1", [1.0]) + return dc, mc + + +def _fresh_scaler(collection): + from torchref.scaling.collection_scaler import CollectionScaler + + dc, mc = collection + return CollectionScaler(dc, mc, verbose=0).initialize() + + +def test_it_offers_exactly_the_selectable_objectives(): + from torchref.scaling.collection_scaler import CollectionScaler + from torchref.scaling.scaler_base import DEFAULT_SCALE_TARGET, SCALE_TARGETS + + sig = inspect.signature(CollectionScaler.refine_lbfgs_joint) + assert sig.parameters["scale_target"].default == DEFAULT_SCALE_TARGET + assert ( + DEFAULT_SCALE_TARGET == "ls" + ), "the joint fit's default must track the single-dataset one" + assert "nll" in SCALE_TARGETS and "ml_noalpha" in SCALE_TARGETS + + +def test_unknown_objective_fails_closed(collection): + scaler = _fresh_scaler(collection) + with pytest.raises(ValueError, match="scale_target must be one of"): + scaler.refine_lbfgs_joint(scale_target="nll_i") + + +@pytest.mark.parametrize("scale_target", ["ls", "nll", "ml_noalpha"]) +def test_every_objective_fits_finite_parameters(collection, scale_target): + """Every selectable row must drive the joint fit to finite parameters.""" + scaler = _fresh_scaler(collection) + m = scaler.refine_lbfgs_joint( + nsteps=2, max_iter=20, verbose=False, scale_target=scale_target + ) + for p in scaler.parameters(): + assert torch.isfinite(p).all(), f"{scale_target}: non-finite scale parameter" + assert m["rwork"] and all(0.0 < r < 1.0 for r in m["rwork"]), m["rwork"] + + +def test_the_dataset_view_shares_the_parents_parameters(collection): + """A row's scaler must be a view, not a copy.""" + from torchref.scaling.collection_scaler import _DatasetScalerView + + scaler = _fresh_scaler(collection) + dc, mc = collection + fracs = mc[mc.dark_key].fractions.detach() + view = _DatasetScalerView(scaler, fracs) + + # No parameters of its own -- only the bound fractions buffer. + assert list(view.parameters()) == [] + assert view.device == scaler.device + # And it routes through the parent's mixed-solvent path. + with torch.no_grad(): + fcalc = mc[mc.dark_key](dc[mc.dark_key].hkl) + got = view(fcalc) + want = scaler.forward_mixed(fcalc, fracs) + assert torch.equal(got, want) diff --git a/tests/integration/test_collection_scaler_batched.py b/tests/integration/test_collection_scaler_batched.py new file mode 100644 index 00000000..bc6eb50c --- /dev/null +++ b/tests/integration/test_collection_scaler_batched.py @@ -0,0 +1,114 @@ +"""Batched scaling preserves individual results, affinity and solvent caches.""" + +import pytest +import torch + +from torchref.config import get_default_device, get_float_dtype + +pytestmark = pytest.mark.integration + + +def _weights(alpha: float) -> torch.Tensor: + """Single-row activation weights ``[[1 - a, a]]``.""" + return torch.tensor( + [[1.0 - alpha, alpha]], device=get_default_device(), dtype=get_float_dtype() + ) + + +class TestBatchedMatchesUnbatched: + def test_every_row_matches_forward_mixed(self, difference_collection): + dc, mc, scaler = difference_collection + components = dc.component_structure_factors(mc, recalc=True) + + w = torch.tensor( + [[1.0, 0.0], [0.78, 0.22], [0.3, 0.7]], + device=get_default_device(), + dtype=get_float_dtype(), + ) + fcalc = mc.mix_component_fcalcs(components, w) + + batched = scaler.forward_batched(fcalc, w) + assert batched.shape == fcalc.shape + + for i in range(w.shape[0]): + single = scaler.forward_mixed(fcalc[i], w[i]) + assert torch.allclose( + batched[i], single, rtol=1e-6, atol=1e-6 + ), f"batched row {i} disagrees with forward_mixed" + + def test_solvent_stack_rows_are_the_per_component_solvents( + self, difference_collection + ): + """A transposed or misordered stack would still have the right shape.""" + _, mc, scaler = difference_collection + stack = scaler.compute_component_solvent_raw() + assert stack.shape == (mc.n_base_models, len(scaler.hkl)) + assert stack.is_complex() + for k in range(mc.n_base_models): + assert torch.equal(stack[k], scaler._get_component_f_sol_raw(k)) + + +class TestAffineInTheMixingWeights: + def test_secant_in_alpha_equals_the_scaled_jacobian(self, difference_collection): + """The property the two-moment forward model rests on.""" + dc, mc, scaler = difference_collection + components = dc.component_structure_factors(mc, recalc=True) + + solvent = scaler.compute_component_solvent_raw() + assert not torch.allclose(solvent[0], solvent[1]) + a1, a2 = 0.60, 0.10 + w1, w2 = _weights(a1), _weights(a2) + jac = torch.tensor( + [[-1.0, 1.0]], device=get_default_device(), dtype=get_float_dtype() + ) # d/da of [1 - a, a] + + s1 = scaler.forward_batched(mc.mix_component_fcalcs(components, w1), w1) + s2 = scaler.forward_batched(mc.mix_component_fcalcs(components, w2), w2) + deriv = scaler.forward_batched(mc.mix_component_fcalcs(components, jac), jac) + + secant = s1 - s2 + expected = (a1 - a2) * deriv + + # Subtraction roundoff scales with its operands, not the smaller secant. + scale = (s1.abs() + s2.abs() + expected.abs()).max() + tolerance = 16 * torch.finfo(s1.real.dtype).eps * scale + assert (secant - expected).abs().max() <= tolerance + + def test_scaling_is_linear_in_the_structure_factors(self, difference_collection): + """The other half of affinity: doubling F_calc at fixed weights doubles the + F_calc-dependent part, leaving the solvent offset behind.""" + dc, mc, scaler = difference_collection + components = dc.component_structure_factors(mc, recalc=True) + w = _weights(0.22) + fcalc = mc.mix_component_fcalcs(components, w) + + s1 = scaler.forward_batched(fcalc, w) + s2 = scaler.forward_batched(2.0 * fcalc, w) + zero = scaler.forward_batched(torch.zeros_like(fcalc), w) + + # (S(2F) - S(0)) == 2 * (S(F) - S(0)) + lhs, rhs = s2 - zero, 2.0 * (s1 - zero) + rel = (lhs - rhs).abs().max() / rhs.abs().max() + assert rel < 1e-5 + + +class TestSolventCacheIsNotPoisoned: + def test_batched_calls_leave_the_cache_alone(self, difference_collection): + """Two batched calls with different weights, then a plain one.""" + dc, mc, scaler = difference_collection + components = dc.component_structure_factors(mc, recalc=True) + w = _weights(0.22) + jac = torch.tensor( + [[-1.0, 1.0]], device=get_default_device(), dtype=get_float_dtype() + ) + fcalc = mc.mix_component_fcalcs(components, w) + + before = scaler.forward_mixed(fcalc[0], w[0]).clone() + + scaler.forward_batched(fcalc, w) + scaler.forward_batched(mc.mix_component_fcalcs(components, jac), jac) + + after = scaler.forward_mixed(fcalc[0], w[0]) + assert torch.allclose( + after, before, rtol=1e-6, atol=1e-6 + ), "a batched call changed what a later forward_mixed returns" diff --git a/tests/integration/test_collection_stack_accessors.py b/tests/integration/test_collection_stack_accessors.py new file mode 100644 index 00000000..b4e66412 --- /dev/null +++ b/tests/integration/test_collection_stack_accessors.py @@ -0,0 +1,80 @@ +"""Collection row selection, partitions and error handling.""" + +import pytest +import torch + +pytestmark = pytest.mark.integration + + +@pytest.fixture +def collection(loaded_reflection_data): + """Two independent observation sets for selection and partition checks.""" + from torchref.io import DatasetCollection + + data = loaded_reflection_data + return ( + DatasetCollection(device=data.device, verbose=0) + .add_dataset("dark", data, set_as_reference=True) + .add_dataset("light", data) + ) + + +class TestThreeWayMasks: + @pytest.mark.parametrize("use_set", ["work", "free", "val"]) + def test_rows_match_the_per_dataset_subset(self, collection, use_set): + attr = {"work": "work", "free": "free", "val": "validation"}[use_set] + stacked = collection.stack_masks(use_set=use_set) + for row, key in enumerate(collection.keys()): + assert torch.equal(stacked[row], getattr(collection[key], attr).mask) + + def test_the_three_subsets_partition_the_valid_reflections(self, collection): + work = collection.stack_masks(use_set="work") + free = collection.stack_masks(use_set="free") + val = collection.stack_masks(use_set="val") + + assert not (work & free).any() + assert not (work & val).any() + assert not (free & val).any() + + for row, key in enumerate(collection.keys()): + valid = collection[key].masks().to(torch.bool) + assert torch.equal(work[row] | free[row] | val[row], valid) + + def test_a_validation_set_is_carved_out_of_free_not_work(self, collection): + """The 3-way behaviour a 2-way flag array cannot reproduce.""" + data = collection["dark"] + free_before = int(collection.stack_masks(use_set="free")[0].sum()) + data.generate_validation_set(val_fraction_of_free=0.5, seed=0) + + free_after = int(collection.stack_masks(use_set="free")[0].sum()) + val_after = int(collection.stack_masks(use_set="val")[0].sum()) + + assert val_after > 0 + assert free_after < free_before + assert free_after + val_after == pytest.approx(free_before, abs=1) + + def test_an_unknown_subset_name_is_rejected(self, collection): + with pytest.raises(ValueError, match="use_set must be"): + collection.stack_masks(use_set="test") + + +class TestSelectionAndErrors: + def test_keys_argument_selects_and_orders_the_rows(self, collection): + both = collection.stack_F_obs() + one = collection.stack_F_obs(keys=["light"]) + assert one.shape[0] == 1 + assert torch.equal(one[0], both[1]) + + reversed_ = collection.stack_F_obs(keys=["light", "dark"]) + assert torch.equal(reversed_[0], both[1]) + assert torch.equal(reversed_[1], both[0]) + + def test_unknown_key_is_rejected(self, collection): + with pytest.raises(KeyError, match="Unknown dataset keys"): + collection.stack_F_obs(keys=["nope"]) + + def test_centric_flags_are_shared_and_hkl_shaped(self, collection): + centric = collection.get_centric_flags() + assert centric is not None + assert centric.shape == (len(collection.hkl),) + assert centric.dtype == torch.bool diff --git a/tests/integration/test_collection_target_characterisation.py b/tests/integration/test_collection_target_characterisation.py new file mode 100644 index 00000000..7a16dc29 --- /dev/null +++ b/tests/integration/test_collection_target_characterisation.py @@ -0,0 +1,156 @@ +"""Collection targets consume live scaled data, selected subsets and summed losses.""" + +import pytest +import torch + +from torchref.config import get_default_device, get_float_dtype + +pytestmark = pytest.mark.integration + +# Allow float32 differences from threaded structure-factor reductions. +LOSS_RTOL = 1e-4 + + +@pytest.fixture +def collection(difference_collection): + """Install live observation corrections on a fresh collection.""" + dc, mc, scaler = difference_collection + dc.scale(nsteps=1) + return dc, mc, scaler + + +def _targets(dc, mc, scaler): + from torchref.refinement.targets import ( + CollectionDifferenceIntensityTarget, + CollectionDifferenceTarget, + CollectionMLTarget, + ) + + return { + "difference": CollectionDifferenceTarget(dc, mc, scaler=scaler, verbose=0), + "difference_i": CollectionDifferenceIntensityTarget( + dc, mc, scaler=scaler, verbose=0 + ), + "ml": CollectionMLTarget(dc, mc, scaler=scaler, verbose=0), + } + + +class TestObservedAmplitudesAreScaled: + """The loss must move when a dataset's own scale moves.""" + + @pytest.mark.parametrize("name", ["difference", "difference_i", "ml"]) + def test_loss_responds_to_the_datasets_own_log_scale(self, collection, name): + dc, mc, scaler = collection + target = _targets(dc, mc, scaler)[name] + + before = target.forward().item() + with torch.no_grad(): + dc.scaler.raw_parameters[1, 0] += 0.25 + target.maintenance() if hasattr(target, "maintenance") else None + after = target.forward().item() + + rel = abs(after - before) / abs(before) + assert rel > 1e-3, ( + f"{name}: changing the light dataset's log_scale moved the loss by only " + f"{rel:.2e}; the target is reading raw amplitudes, not the scaled ones" + ) + + +class TestSubsetSelectionIsThreeWay: + """Work / free / validation, with validation carved out of both.""" + + @pytest.mark.parametrize("use_set", ["work", "free"]) + def test_loss_is_restricted_to_the_selected_subset(self, collection, use_set): + """Work and free are different sizes here, so a target that ignored + ``use_set`` would return the same number for both.""" + from torchref.refinement.targets import CollectionDifferenceTarget + + dc, mc, scaler = collection + target = CollectionDifferenceTarget( + dc, mc, scaler=scaler, use_set=use_set, verbose=0 + ) + assert target.use_set == use_set + n = target._n_reflections() + expected = sum( + (dc[k].work if use_set == "work" else dc[k].free).n for k in target._keys() + ) + assert n == expected + + +class TestLossesAreSummedNotAveraged: + """A summed X-ray term grows with the data; a meaned one does not.""" + + def test_adding_a_dataset_grows_the_absolute_loss( + self, collection, loaded_reflection_data + ): + """Summed loss grows in proportion to the number of datasets.""" + from torchref.refinement.targets import CollectionMLTarget + + dc, mc, scaler = collection + target_before = CollectionMLTarget(dc, mc, scaler=scaler, verbose=0) + n_before = len(target_before._keys()) + one = target_before.forward().item() + + dc.add_dataset("light2", loaded_reflection_data) + mc.add_timepoint("light2", mc["light"].fractions.detach().tolist()) + target_after = CollectionMLTarget(dc, mc, scaler=scaler, verbose=0) + n_after = len(target_after._keys()) + two = target_after.forward().item() + + assert (n_before, n_after) == (2, 3) + ratio = two / one + # Refitting shared beta on pooled free reflections shifts the per-row loss. + assert ratio == pytest.approx(n_after / n_before, rel=0.15), ( + f"{n_after} datasets gave {ratio:.3f}x the loss of {n_before}; a summed " + f"target should scale with the count and a meaned one stay near 1.0" + ) + + +class TestReportedNumbers: + """The shape of what ``get_rfactor`` / ``stats`` promise, plus reproducibility.""" + + @pytest.mark.parametrize("name", ["difference", "difference_i", "ml"]) + def test_forward_is_finite_and_reproducible(self, collection, name): + dc, mc, scaler = collection + target = _targets(dc, mc, scaler)[name] + first = target.forward().item() + second = target.forward().item() + assert torch.isfinite( + torch.tensor(first, device=get_default_device(), dtype=get_float_dtype()) + ) + assert second == pytest.approx(first, rel=LOSS_RTOL) + + def test_rfactor_shape_and_range(self, collection): + dc, mc, scaler = collection + target = _targets(dc, mc, scaler)["difference"] + + rf = target.get_rfactor() + assert set(rf) == {"per_dataset", "rwork_pct", "rfree_pct"} + assert set(rf["per_dataset"]) == {"dark", "light"} + for key, (rwork, rfree) in rf["per_dataset"].items(): + assert 0.0 < rwork < 1.0, f"{key} rwork out of range: {rwork}" + assert 0.0 < rfree < 1.0, f"{key} rfree out of range: {rfree}" + assert set(rf["rwork_pct"]) == {"p10", "p25", "p50", "p75", "p90"} + + def test_stats_reports_the_median_of_the_per_dataset_distribution(self, collection): + dc, mc, scaler = collection + target = _targets(dc, mc, scaler)["difference"] + + stats = target.stats() + rf = target.get_rfactor() + for key in ("loss", "n", "rwork", "rfree"): + assert key in stats, f"missing stat: {key}" + assert stats["rwork"].value == pytest.approx( + rf["rwork_pct"]["p50"], rel=LOSS_RTOL + ) + assert stats["n"].value == target._n_reflections() + + def test_gradient_reaches_the_light_model(self, collection): + dc, mc, scaler = collection + target = _targets(dc, mc, scaler)["difference"] + + target.forward().backward() + grad = mc.base_models[1].xyz.refinable_params.grad + assert grad is not None + assert torch.isfinite(grad).all() + assert grad.abs().max() > 0 diff --git a/tests/integration/test_crystfel_hkl.py b/tests/integration/test_crystfel_hkl.py new file mode 100644 index 00000000..4dc19c88 --- /dev/null +++ b/tests/integration/test_crystfel_hkl.py @@ -0,0 +1,52 @@ +"""CrystFEL parsing, intensity conversion and alignment on real split-half excerpts.""" + +import pytest +import torch + +from torchref import DatasetCollection, ReflectionData +from torchref.config import get_default_device + +pytestmark = pytest.mark.integration + +CELL = [14.97, 18.85, 18.89, 89.4, 84.9, 67.8] + + +@pytest.fixture +def halves(test_files_dir): + return [ + ReflectionData(device=get_default_device(), verbose=0).load_crystfel_hkl( + str(test_files_dir / "hkl" / f"dark_half{i}.hkl"), + cell=CELL, + spacegroup="P 1", + ) + for i in (1, 2) + ] + + +@pytest.mark.parametrize("index", [0, 1]) +def test_observations_and_metadata(halves, index, test_files_dir): + """Preserve negative intensities, uncertainties, caller metadata and valid rows.""" + data = halves[index] + rows = (test_files_dir / "hkl" / f"dark_half{index+1}.hkl").read_text().splitlines() + assert rows[-1] == "End of reflections" + assert len(data.hkl) == len(rows) - 4 + assert data.I.shape == data.I_sigma.shape == data.F.shape == (len(data),) + assert torch.isfinite(data.I_sigma).all() and (data.I_sigma >= 0).all() + assert (data.I < 0).any() + assert (data.F >= 0).all() and data._FrenchWilson is not None + torch.testing.assert_close(data.cell.data, data.cell.data.new_tensor(CELL)) + assert data.spacegroup.number == 1 + + +def test_partial_overlap_alignment_preserves_sources(halves): + """Align overlapping split halves on their union without mutating either input.""" + a, b = halves + original = [data.hkl.clone() for data in halves] + sets = [{tuple(row) for row in data.hkl.tolist()} for data in halves] + assert sets[0] != sets[1] and sets[0] & sets[1] + dc = DatasetCollection(verbose=0, device=get_default_device()) + dc.add_dataset("a", a).add_dataset("b", b) + assert len(dc) == len(sets[0] | sets[1]) + assert dc.stack_I_obs().shape == (2, len(dc)) + for data, hkl in zip(halves, original): + assert torch.equal(data.hkl, hkl) diff --git a/tests/integration/test_dataset_scaler.py b/tests/integration/test_dataset_scaler.py new file mode 100644 index 00000000..ec858edc --- /dev/null +++ b/tests/integration/test_dataset_scaler.py @@ -0,0 +1,371 @@ +"""Joint observed-data scaling and live dataset access on deposited reflections.""" + +import copy +import math + +import pytest +import torch + +from torchref import DatasetCollection, ScaledDataset +from torchref.base.targets.dataset_scaling import dataset_scaling_loss +from torchref.config import get_default_device, get_int_dtype +from torchref.scaling import DatasetScaler + +pytestmark = pytest.mark.integration + + +def clone(data): + return data.__select__( + torch.arange(len(data), device=data.device, dtype=get_int_dtype()) + ) + + +def collection(data, factors): + dc = DatasetCollection(device=data.device, verbose=0) + for i, factor in enumerate(factors): + raw = clone(data) + raw.F *= factor + raw.F_sigma *= factor + if raw.I is not None: + raw.I *= factor**2 + raw.I_sigma *= factor**2 + dc.add_dataset(str(i), raw) + return dc + + +@pytest.mark.parametrize("factors", [(1, 1), (1, 1000), (1, 2, 8)]) +def test_known_scales_are_centered_and_sources_stay_raw( + loaded_reflection_data, factors +): + before = loaded_reflection_data.F.clone() + dc = collection(loaded_reflection_data, factors).scale(nsteps=2) + expected = torch.tensor( + factors, dtype=dc.scaler.raw_parameters.dtype, device=get_default_device() + ).log() + expected = expected.mean() - expected + torch.testing.assert_close( + dc.scaler.corrections[:, 0], expected, atol=1e-4, rtol=1e-4 + ) + assert dc.scaler.corrections[:, 1:].abs().max() < 1e-4 + assert torch.equal(loaded_reflection_data.F, before) + assert not hasattr(loaded_reflection_data, "parameters") + assert not hasattr(loaded_reflection_data, "log_scale") + assert not hasattr(loaded_reflection_data, "U_aniso") + for key, _ in dc: + assert isinstance(dc[key], ScaledDataset) + torch.testing.assert_close(dc[key].F, dc["0"].F, rtol=1e-4, atol=1e-4) + + +def test_known_anisotropy_is_recovered(loaded_reflection_data): + dc = collection(loaded_reflection_data, (1, 1, 1)) + reference = DatasetScaler(dc.datasets) + truth = torch.tensor( + [ + [0.3, 0.2, -0.1, 0.1, 0.03, -0.02, 0.05], + [-0.2, -0.1, 0.2, -0.05, -0.02, 0.01, -0.03], + [-0.1, -0.1, -0.1, -0.05, -0.01, 0.01, -0.02], + ], + dtype=reference.raw_parameters.dtype, + device=get_default_device(), + ) + for i, ds in enumerate(dc.values()): + distortion = (reference.design(ds.hkl) @ truth[i]).exp() + ds.F /= distortion + ds.F_sigma /= distortion + dc.scale(nsteps=4) + torch.testing.assert_close(dc.scaler.corrections, truth, atol=2e-3, rtol=2e-3) + + +def test_live_access_scales_both_sigmas_and_all_entrypoints(loaded_reflection_data): + dc = collection(loaded_reflection_data, (1, 1)).scale(nsteps=1) + data = dc["0"] + with torch.no_grad(): + dc.scaler.raw_parameters[0, 0] = 2 * math.log(2) + dc.scaler.raw_parameters[0, 1] = 0.1 + correction = 2 * torch.exp( + 0.05 * (data.hkl[:, 0] / dc.scaler.hkl_scale[0]).square() + ) + data.generate_validation_set(val_fraction_of_free=0.5, seed=0) + for name, power, subset_attr, stack in [ + ("F", 1, "F", dc.stack_F_obs), + ("F_sigma", 1, "sigF", dc.stack_F_sigma), + ("I", 2, "I", dc.stack_I_obs), + ("I_sigma", 2, "sigI", dc.stack_I_sigma), + ]: + raw = getattr(data, name + "_raw") + actual = getattr(data, name) + torch.testing.assert_close(actual, raw * correction**power) + torch.testing.assert_close(stack()[0], actual) + for kind in ("work", "free", "validation"): + subset = getattr(data, kind) + torch.testing.assert_close( + getattr(subset, subset_attr), actual[subset.mask] + ) + torch.testing.assert_close( + getattr(subset, subset_attr + "_raw"), raw[subset.mask] + ) + torch.testing.assert_close(dc(mask=False)["0"][1], data.F) + torch.testing.assert_close(data.get_corrected_data(), (data.F, data.F_sigma)) + torch.testing.assert_close(data.get_corrected_intensities(), (data.I, data.I_sigma)) + dc.scaler.requires_grad_(True) + for _ in range(2): + dc.scaler.zero_grad() + data.work.F.sum().backward() + assert dc.scaler.raw_parameters.grad[1].abs().sum() > 0 + dc.scaler.requires_grad_(False) + + +def test_selection_copy_alignment_and_independent_collections(loaded_reflection_data): + a = collection(loaded_reflection_data, (1, 4)).scale(nsteps=1) + b = collection(loaded_reflection_data, (1, 9)).scale(nsteps=1) + view = a["0"] + selected = view.__select__( + torch.arange( + 0, len(view), 3, device=get_default_device(), dtype=get_int_dtype() + ) + ) + torch.testing.assert_close(selected.F, view.F[::3]) + assert selected.scaler is a.scaler + assert copy.deepcopy(view).scaler is a.scaler + aligned = view.copy().validate_hkl(view.hkl.flip(0)) + torch.testing.assert_close(aligned.F, view.F.flip(0)) + assert a.scaler is not b.scaler + assert not torch.allclose(a["0"].F, b["0"].F) + scaler = a.scaler + a.scale(nsteps=1) + assert a.scaler is scaler + a.add_dataset("extra", loaded_reflection_data) + assert a.scaler is None + a.scale(nsteps=1) + assert a.scaler is not scaler + torch.testing.assert_close(view.F, view.F_raw * 2) + + +@pytest.mark.usefixtures("double_cpu") +def test_two_dataset_loss_matches_propagated_variance_and_gradients( + loaded_reflection_data, +): + f = torch.stack( + (loaded_reflection_data.F[:128], loaded_reflection_data.F[:128] * 1.1) + ).double() + sigma = torch.stack( + (loaded_reflection_data.F_sigma[:128], loaded_reflection_data.F_sigma[:128] * 3) + ).double() + log_k = torch.zeros_like(f, requires_grad=True) + mask = torch.isfinite(f) & torch.isfinite(sigma) & (sigma > 0) + actual = dataset_scaling_loss(f, sigma, log_k, mask) + common = mask.all(dim=0) + expected = 0.5 * ((f[0] - f[1]).square() / sigma.square().sum(dim=0))[common].sum() + torch.testing.assert_close(actual, expected) + assert torch.autograd.gradcheck( + lambda p: dataset_scaling_loss(f, sigma, p - p.mean(dim=0), mask), + (log_k,), + fast_mode=True, + ) + noisier = dataset_scaling_loss(f, sigma * 2, log_k, mask) + torch.testing.assert_close(noisier, actual / 4) + + +def test_missing_and_invalid_observations_have_finite_gradients(loaded_reflection_data): + f = torch.stack( + (loaded_reflection_data.F[:128], loaded_reflection_data.F[:128] * 1.1) + ) + sigma = torch.stack( + (loaded_reflection_data.F_sigma[:128], loaded_reflection_data.F_sigma[:128]) + ) + mask = torch.isfinite(f) & torch.isfinite(sigma) & (sigma > 0) + mask[:, :10] = False + f[:, :10] = float("nan") + sigma[:, :10] = 0 + log_k = torch.zeros_like(f, requires_grad=True) + loss = dataset_scaling_loss(f, sigma, log_k, mask) + loss.backward() + assert torch.isfinite(loss) + assert torch.isfinite(log_k.grad).all() + assert torch.equal(log_k.grad[:, :10], torch.zeros_like(log_k.grad[:, :10])) + + +def test_free_and_validation_changes_do_not_affect_fit(loaded_reflection_data): + results = [] + for filler in (3, 900): + dc = collection(loaded_reflection_data, (1, 2)) + ds = dc["1"] + mismatched = ds.work.indices[:60] + ds.rfree_flags[mismatched[:30]] = False + ds.validation_flags = torch.zeros_like(ds.rfree_flags, dtype=torch.bool) + ds.validation_flags[mismatched[30:]] = True + held_out = ~ds.work.mask + ds.F[held_out] = filler + ds.F_sigma[held_out] = filler + dc.scale(nsteps=2) + assert torch.equal(dc.hkl, dc.scaler.hkl) + assert not dc.scaler.fit_mask[:, mismatched].any() + results.append(dc.scaler.raw_parameters.detach().clone()) + assert torch.equal(*results) + + +def test_permutation_and_partial_overlap_chain(loaded_reflection_data): + n = len(loaded_reflection_data) + sources = { + "a": clone(loaded_reflection_data).__select__( + torch.arange(n // 2, device=get_default_device(), dtype=get_int_dtype()) + ), + "b": clone(loaded_reflection_data), + "c": clone(loaded_reflection_data).__select__( + torch.arange(n // 2, n, device=get_default_device(), dtype=get_int_dtype()) + ), + } + sources["b"].F *= 2 + sources["b"].F_sigma *= 2 + sources["c"].F *= 4 + sources["c"].F_sigma *= 4 + results = [] + for order in [("a", "b", "c"), ("c", "a", "b")]: + dc = DatasetCollection(device=get_default_device(), verbose=0) + for key in order: + dc.add_dataset(key, sources[key]) + dc.scale(nsteps=2) + results.append( + {k: dc.scaler.corrections[dc.scaler.keys.index(k)] for k in order} + ) + for key in sources: + torch.testing.assert_close( + results[0][key], results[1][key], atol=1e-4, rtol=1e-4 + ) + with pytest.raises(ValueError, match="disconnected"): + DatasetScaler({k: sources[k] for k in ("a", "c")}) + with pytest.raises(ValueError, match="identify"): + DatasetScaler( + { + "a": loaded_reflection_data.__select__( + torch.arange(6, device=get_default_device(), dtype=get_int_dtype()) + ), + "b": loaded_reflection_data.__select__( + torch.arange(6, device=get_default_device(), dtype=get_int_dtype()) + ), + } + ) + + +def test_checkpoint_and_mtz_export_preserve_observations( + loaded_reflection_data, tmp_path +): + import reciprocalspaceship as rs + + dc = collection(loaded_reflection_data, (1, 4)).scale(nsteps=1) + path = tmp_path / "collection.pt" + dc.save_state(path) + restored = DatasetCollection.load_state(path, device=get_default_device()) + assert restored["0"].scaler is restored["1"].scaler is restored.scaler + torch.testing.assert_close(restored.stack_F_obs(), dc.stack_F_obs()) + path = tmp_path / "view.pt" + dc["0"].save_state(path) + restored_view = ScaledDataset.load_state(path, device=get_default_device()) + torch.testing.assert_close(restored_view.I, dc["0"].I) + path = tmp_path / "scaled.mtz" + dc["0"].write_mtz(str(path)) + exported = rs.read_mtz(str(path)) + assert len(exported) == len(dc["0"]) + import numpy as np + + for column, attribute in [ + ("F-obs", "F"), + ("SIGF-obs", "F_sigma"), + ("I-obs", "I"), + ("SIGI-obs", "I_sigma"), + ]: + np.testing.assert_allclose( + exported[column].to_numpy(dtype=float), + getattr(dc["0"], attribute).detach().cpu().numpy(), + rtol=1e-6, + ) + + +def test_bijvoet_observations_keep_distinct_identities(loaded_reflection_data): + """Canonical duplicate HKLs retain separate signed observations and sigmas.""" + data = loaded_reflection_data.__select__( + torch.arange( + 256, device=get_default_device(), dtype=get_int_dtype() + ).repeat_interleave(2) + ) + data.friedel_merged = False + data.friedel_flags = ( + torch.arange(len(data), device=get_default_device(), dtype=get_int_dtype()) % 2 + == 1 + ) + data.hkl_anomalous = torch.where(data.friedel_flags[:, None], -data.hkl, data.hkl) + data.F[data.friedel_flags] *= 1.25 + dc = collection(data, (1, 4)).scale(nsteps=1) + assert len(dc) == len(data) + lookup = {tuple(h.tolist()): f for h, f in zip(data.hkl_anomalous, data.F)} + expected = torch.stack([lookup[tuple(h.tolist())] for h in dc["0"].hkl_anomalous]) + torch.testing.assert_close(dc["0"].F_raw, expected) + torch.testing.assert_close(dc["0"].F, dc["1"].F, rtol=1e-4, atol=1e-4) + assert dc["0"].friedel_flags.sum() == 256 + selected = ( + dc["0"] + .copy() + .validate_hkl(dc.hkl.flip(0), identity_hkl=dc["0"].hkl_anomalous.flip(0)) + ) + torch.testing.assert_close(selected.F, dc["0"].F.flip(0)) + + +def test_raw_dataset_excludes_deprecated_interfaces(loaded_reflection_data): + """Observation storage exposes subset views without optimization methods.""" + for name in ( + "log_scale", + "U_aniso", + "parameters", + "setup_scale", + "setup_anisotropy", + "compute_e_values", + "get_radial_shells", + "get_work_set", + "get_test_set", + "get_rfree_masks", + "_masked_unpack", + "get_mask", + ): + assert not hasattr(loaded_reflection_data, name) + assert not callable(loaded_reflection_data) + + +def test_ded_context_consumes_collection_scaled_views( + loaded_reflection_data, monkeypatch +): + """DED preparation uses corrected amplitudes after installing scaled members.""" + from torchref.cli import validate_ded + + dark, light = clone(loaded_reflection_data), clone(loaded_reflection_data) + light.F *= 4 + light.F_sigma *= 4 + inputs = {"dark": dark, "light": light} + monkeypatch.setattr( + validate_ded, "load_reflection_data", lambda path, **kwargs: inputs[path] + ) + context = validate_ded.setup_ded_context( + "dark", "light", dmin=2.2, device=get_default_device() + ) + dc = context["collection"] + assert context["data_dark"] is dc["dark"] + assert context["data_light"] is dc["light"] + torch.testing.assert_close(dc["dark"].F, dc["light"].F, rtol=1e-4, atol=1e-4) + relative_difference = ( + context["w_dfo"].norm() / dc["dark"].F[context["refl_mask"]].norm() + ) + assert relative_difference < 1e-5 + + +def test_missing_intensity_access(loaded_reflection_data): + """Raw and scaled views expose missing columns and reject intensity-only reads.""" + dc = collection(loaded_reflection_data, (1, 2)) + dc["0"].I = dc["0"].I_sigma = None + raw = dc["0"] + dc.scale(nsteps=1) + for data in (raw, dc["0"]): + with pytest.raises(ValueError, match="No intensities"): + data.get_corrected_intensities() + for name in ("I", "sigI", "I_raw", "sigI_raw"): + assert getattr(data.work, name) is None + with pytest.raises(ValueError, match="'0'"): + dc.stack_I_obs() diff --git a/tests/integration/test_device_mixin.py b/tests/integration/test_device_mixin.py index b538e100..01a1d2d5 100644 --- a/tests/integration/test_device_mixin.py +++ b/tests/integration/test_device_mixin.py @@ -4,8 +4,10 @@ The primary acceptance test exercises a CPU -> GPU -> CPU round-trip on a ``ModelFT`` instance: structure factors are recomputed at each leg and checked for the expected device placement and numerical agreement. This -covers the ``__dict__`` walk (Cell, SfFFT, anomalous cache), the always- -invalidate cache policy, and the post-move ``_rebuild_sf_indices`` hook. +covers the ``__dict__`` walk (Cell, SfFFT, anomalous cache) and the always- +invalidate cache policy. The SF index partition needs no post-move hook: it is +keyed on its inputs and recomputed on access, so a device move invalidates it +by changing the fingerprint rather than by being told to. """ from __future__ import annotations @@ -13,6 +15,7 @@ import pytest import torch + def _load_model_ft(pdb_file, mtz_file): """Helper: load a ModelFT and matching reflection data on CPU. @@ -45,7 +48,7 @@ def test_modelft_cpu_gpu_cpu_sf_round_trip(sample_pdb_file, sample_mtz_file): match the original CPU result. """ model, data = _load_model_ft(sample_pdb_file, sample_mtz_file) - hkl, *_ = data() + hkl = data.hkl # ---- CPU leg --------------------------------------------------------- assert hkl.device.type == "cpu", "test setup: hkl should start on CPU" @@ -113,7 +116,7 @@ def test_modelft_cpu_only_recompute_after_to(sample_pdb_file, sample_mtz_file): it should match the first call. """ model, data = _load_model_ft(sample_pdb_file, sample_mtz_file) - hkl, *_ = data() + hkl = data.hkl fcalc_before = model(hkl).detach().clone() model.to("cpu") # idempotent move diff --git a/tests/integration/test_dtype_config_float64.py b/tests/integration/test_dtype_config_float64.py index 942cca31..dccb37cb 100644 --- a/tests/integration/test_dtype_config_float64.py +++ b/tests/integration/test_dtype_config_float64.py @@ -11,18 +11,18 @@ import torch - @pytest.mark.unit def test_translation_phases_complex_dtype_float64(double_cpu): - """compute_translation_phases must honor the configured complex dtype.""" - from torchref.base.reciprocal.symmetry import compute_translation_phases + """Symmetry.phase_factors must honor the configured complex dtype.""" + from torchref.symmetry import SpaceGroup - hkl = torch.tensor([[1.0, 0.0, 0.0], [2.0, 1.0, 0.0], [0.0, 0.0, 3.0]]) - translations = torch.tensor([[0.0, 0.0, 0.0], [0.5, 0.5, 0.0]]) + # P21 gives two operations, one carrying a half translation. + sym = SpaceGroup("P 21") + hkl = torch.tensor([[1, 0, 0], [2, 1, 0], [0, 0, 3]]) - phases = compute_translation_phases(hkl, translations) + phases = sym.phase_factors(hkl) - # Was complex64 (float32 hardcode); under float64 config must be complex128. + # Must not narrow to complex64 under a float64 configuration. assert phases.dtype == torch.complex128 assert phases.shape == (2, 3) assert torch.isfinite(phases.real).all() @@ -43,7 +43,7 @@ def test_scaler_binwise_mean_intensity_float64(double_cpu, sample_structure_pair scaler = Scaler(model=model, data=data, nbins=10, verbose=0) - hkl = data()[0] + hkl = data.hkl fcalc = model(hkl) assert fcalc.dtype == torch.complex128 @@ -58,11 +58,11 @@ def test_scaler_binwise_mean_intensity_float64(double_cpu, sample_structure_pair @pytest.mark.integration def test_occupancy_floor_density_matmul_float64(double_cpu, sample_structure_pair): """compute_density_at_positions hardcoded hkl.T.float(); matmul raised under float64.""" - from torchref.io import ReflectionData - from torchref.model.model_ft import ModelFT from torchref.experimental.targets.occupancy_floor_diagnostic import ( OccupancyFloorDiagnostic, ) + from torchref.io import ReflectionData + from torchref.model.model_ft import ModelFT model = ModelFT() model.load_cif(str(sample_structure_pair["model"])) @@ -74,7 +74,7 @@ def test_occupancy_floor_density_matmul_float64(double_cpu, sample_structure_pai positions = model.cell.cartesian_to_fractional(model.xyz()) assert positions.dtype == torch.float64 - hkl = data()[0] + hkl = data.hkl diagnostic = OccupancyFloorDiagnostic(model_dark=model, model_light=model) # Pre-fix this raised: float64 positions @ float32 hkl.T. diff --git a/tests/integration/test_fcalc_add_noise.py b/tests/integration/test_fcalc_add_noise.py new file mode 100644 index 00000000..69326550 --- /dev/null +++ b/tests/integration/test_fcalc_add_noise.py @@ -0,0 +1,152 @@ +"""Unbiased noisy intensities, uncertainty propagation and reproducible draws.""" + +import pytest +import torch + +from torchref.config import get_default_device, get_float_dtype, get_int_dtype + +pytestmark = pytest.mark.integration + + +@pytest.fixture +def fcalc_scene(loaded_reflection_data): + """Use deposited 1DAW amplitudes with deterministic phases for noisy draws.""" + from torchref.io import FcalcDataset + + data = loaded_reflection_data + result = FcalcDataset( + hkl=data.hkl.clone(), + cell=data.cell, + spacegroup=data.spacegroup, + device=data.device, + ) + phase = torch.linspace(-2.0, 2.0, len(data), dtype=data.F.dtype, device=data.device) + result.set_fcalc(data.F * torch.exp(1j * phase)) + return result + + +class TestNegativesSurvive: + def test_some_intensities_come_out_negative(self, fcalc_scene): + noisy = fcalc_scene.add_noise(sigma_mul=0.5, seed=3, verbose=False) + assert noisy.I is not None, "add_noise did not retain intensities" + assert bool((noisy.I < 0).any()), ( + "no negative intensities at 50% multiplicative noise -- either the scene has " + "no weak reflections or the intensities were clamped" + ) + + def test_the_amplitude_is_clamped_but_the_intensity_is_not(self, fcalc_scene): + """The asymmetry is deliberate, so it is asserted rather than assumed.""" + noisy = fcalc_scene.add_noise(sigma_mul=0.5, seed=3, verbose=False) + assert bool((noisy.fcalc_amp >= 0).all()) + negative = noisy.I < 0 + assert bool(negative.any()) + # Where the intensity is negative the amplitude is floored at zero, so the two + # cannot agree -- which is exactly the information a clamp would have destroyed. + assert torch.allclose( + noisy.fcalc_amp[negative], torch.zeros_like(noisy.fcalc_amp[negative]) + ) + + @staticmethod + def _weak(noisy, truth): + """The subset a clamp at zero can touch: reflections within 2 sigma of zero.""" + return truth < 2.0 * noisy.I_sigma + + def test_the_intensity_is_unbiased_on_the_weak_reflections(self, fcalc_scene): + """The property the clamp breaks, as a bound on the mean of the weak tail.""" + noisy = fcalc_scene.add_noise( + sigma_lin=200.0, sigma_mul=0.0, seed=5, verbose=False + ) + truth = fcalc_scene.fcalc_amp**2 + weak = self._weak(noisy, truth) + assert int(weak.sum()) > 50, "too few weak reflections to say anything" + + assert (noisy.I[weak] < 0).any() + residual = (noisy.I - truth)[weak] + sem = float(noisy.I_sigma[weak].pow(2).sum().sqrt() / int(weak.sum())) + bias = float(residual.mean()) + assert abs(bias) < 4.0 * sem, ( + f"weak-reflection intensity bias {bias:.4g} exceeds 4 sigma " + f"({4 * sem:.4g}); the intensities are being clamped or otherwise skewed" + ) + + +class TestSigmaAndHalves: + def test_sigma_of_the_mean_is_the_single_draw_sigma_over_root_two( + self, fcalc_scene + ): + a = fcalc_scene.add_noise(sigma_mul=0.2, seed=11, verbose=False) + b = fcalc_scene.add_noise(sigma_mul=0.2, seed=12, verbose=False) + # Same model, same noise scale: the reported sigma is a property of the model, + # not of the draw, so it must be identical across seeds. + assert torch.allclose(a.I_sigma, b.I_sigma) + + expected = torch.sqrt( + torch.tensor(0.2, device=get_default_device(), dtype=get_float_dtype()) ** 2 + * fcalc_scene.fcalc_amp**4 + ) / (2.0**0.5) + assert torch.allclose(a.I_sigma, expected, rtol=1e-5) + + def test_amplitude_sigma_uses_the_true_amplitude(self, fcalc_scene): + """Propagating against the noisy amplitude would warp sigma per draw and break + inverse-variance weighting downstream.""" + a = fcalc_scene.add_noise(sigma_mul=0.2, seed=11, verbose=False) + b = fcalc_scene.add_noise(sigma_mul=0.2, seed=99, verbose=False) + assert torch.allclose(a.fobs_sigma, b.fobs_sigma) + + def test_different_seeds_give_different_draws(self, fcalc_scene): + a = fcalc_scene.add_noise(sigma_mul=0.2, seed=1, verbose=False) + b = fcalc_scene.add_noise(sigma_mul=0.2, seed=2, verbose=False) + assert not torch.allclose(a.I, b.I) + + def test_the_same_seed_reproduces(self, fcalc_scene): + a = fcalc_scene.add_noise(sigma_mul=0.2, seed=7, verbose=False) + b = fcalc_scene.add_noise(sigma_mul=0.2, seed=7, verbose=False) + assert torch.equal(a.I, b.I) + + def test_phases_are_untouched(self, fcalc_scene): + """Only the amplitude is perturbed; the phase is the model's.""" + noisy = fcalc_scene.add_noise(sigma_mul=0.2, seed=4, verbose=False) + strong = fcalc_scene.fcalc_amp > 1.0 + assert torch.allclose( + noisy.fcalc_phase[strong], fcalc_scene.fcalc_phase[strong], atol=1e-5 + ) + + def test_the_source_dataset_is_not_modified(self, fcalc_scene): + before = fcalc_scene.fcalc_amp.clone() + fcalc_scene.add_noise(sigma_mul=0.3, seed=8, verbose=False) + assert torch.equal(fcalc_scene.fcalc_amp, before) + assert fcalc_scene.I is None + + +class TestReferenceDriven: + def test_sigmas_are_grafted_from_the_reference( + self, fcalc_scene, loaded_reflection_data + ): + """Use the reference's measured uncertainties for both independent draws.""" + noisy = fcalc_scene.add_noise( + reference=loaded_reflection_data, seed=1, verbose=False + ) + torch.testing.assert_close( + noisy.I_sigma, loaded_reflection_data.I_sigma / (2.0**0.5) + ) + + def test_a_mismatched_reference_is_rejected( + self, fcalc_scene, loaded_reflection_data + ): + ref = loaded_reflection_data.__select__( + torch.arange( + 1, + len(loaded_reflection_data), + device=loaded_reflection_data.device, + dtype=get_int_dtype(), + ) + ) + with pytest.raises(ValueError, match="does not match"): + fcalc_scene.add_noise(reference=ref, verbose=False) + + def test_a_reference_without_sigmas_is_rejected(self, fcalc_scene): + from types import SimpleNamespace + + bad = SimpleNamespace(I_sigma=None, hkl=fcalc_scene.hkl) + with pytest.raises(ValueError, match="I_sigma is None"): + fcalc_scene.add_noise(reference=bad, verbose=False) diff --git a/tests/integration/test_intensity_observable.py b/tests/integration/test_intensity_observable.py new file mode 100644 index 00000000..7764ff18 --- /dev/null +++ b/tests/integration/test_intensity_observable.py @@ -0,0 +1,107 @@ +"""Intensity targets consume measured observations and their uncertainties.""" + +import pytest +import torch + +pytestmark = pytest.mark.integration + + +@pytest.fixture +def refinement(sample_structure_pair): + """Fit a fresh 1DAW refinement for each mutable target check.""" + from torchref import LBFGSRefinement + + ref = LBFGSRefinement( + data_file=str(sample_structure_pair["reflections"]), + pdb=str(sample_structure_pair["model"]), + target_mode="ml", + verbose=0, + ) + ref.get_scales() + return ref + + +def _t(refinement, mode, use_set="work"): + from torchref.refinement.targets.xray.factory import create_xray_target + + return create_xray_target( + data=refinement.reflection_data, + model=refinement.model, + scaler=refinement.scaler, + mode=mode, + use_set=use_set, + ) + + +def test_the_row_is_selectable_and_reads_intensities(refinement): + """``nll_i`` comes out of the factory and its ``get_data`` returns the I columns.""" + t = _t(refinement, "nll_i") + obs, calc, sigma, centric, sub = t.get_data() + data = refinement.reflection_data + + torch.testing.assert_close(obs, data.work.I) + torch.testing.assert_close(sigma, data.work.sigI) + # The model is the SQUARED scaled amplitude, not the amplitude. + torch.testing.assert_close(calc, sub.select(t.get_F_calc_scaled(recalc=False) ** 2)) + assert obs.shape == calc.shape == sigma.shape == (sub.n,) + assert centric.shape == (sub.n,) + + +def test_the_intensity_model_is_the_squared_scaled_amplitude(refinement): + """``get_I_calc_scaled`` squares the SCALED amplitude, not the raw one.""" + t = _t(refinement, "nll_i") + with torch.no_grad(): + amp = t.get_F_calc_scaled(recalc=False) + inten = t.get_I_calc_scaled(recalc=False) + torch.testing.assert_close(inten, amp**2, rtol=1e-6, atol=1e-6) + + +def test_rfactors_stay_on_amplitudes(refinement): + """An intensity row reports the SAME R-factors as an amplitude row.""" + r_i = _t(refinement, "nll_i").get_rfactor() + r_a = _t(refinement, "nll").get_rfactor() + assert r_i == pytest.approx(r_a, abs=1e-9) + + +def test_the_loss_is_finite_differentiable_and_summed(refinement): + t = _t(refinement, "nll_i") + loss = t.forward() + assert torch.isfinite(loss) and loss.ndim == 0 + loss.backward() + grads = [ + p.grad + for p in refinement.model.parameters() + if p.requires_grad and p.grad is not None + ] + assert grads, "no gradient reached the model" + assert all(torch.isfinite(g).all() for g in grads) + refinement.model.zero_grad(set_to_none=True) + + +@pytest.mark.parametrize("use_set", ["work", "free"]) +def test_a_reflections_residual_does_not_depend_on_the_arrays_length( + refinement, use_set +): + """``residuals()`` restricted to a subset must equal ``forward()`` on that subset.""" + t = _t(refinement, "nll_i", use_set=use_set) + sub = t._subset() + with torch.no_grad(): + fwd = t.forward() + summed = t.residuals().index_select(0, sub.indices).sum() + torch.testing.assert_close(summed, fwd, rtol=1e-6, atol=1e-6) + + +def test_missing_intensities_raise_at_construction_not_at_forward(refinement): + """LossState probes ``forward()`` at registration, so a missing column has to be + caught in ``__init__`` or it surfaces from deep inside setup with no mention of why. + """ + import copy + + data = copy.copy(refinement.reflection_data) + data.I = None + from torchref.refinement.targets.xray import NLLIntensityXrayTarget + + with pytest.raises(ValueError, match="dataset carries none"): + NLLIntensityXrayTarget( + data=data, model=refinement.model, scaler=refinement.scaler + ) diff --git a/tests/integration/test_io_cif.py b/tests/integration/test_io_cif.py index 6f944587..2bd47fa1 100644 --- a/tests/integration/test_io_cif.py +++ b/tests/integration/test_io_cif.py @@ -6,116 +6,37 @@ import pytest import torch -from pathlib import Path - -class TestCIFLoading: - """Tests for loading CIF model files.""" - - @pytest.mark.integration - def test_load_model_cif(self, sample_cif_file): - """Test loading a real CIF model file.""" - from torchref.model.model import Model - - model = Model() - model.load_cif(str(sample_cif_file)) - - # Basic checks - use xyz().shape[0] for atom count - n_atoms = model.xyz().shape[0] - assert n_atoms > 0 - assert hasattr(model, 'xyz') - assert hasattr(model, 'adp') - assert hasattr(model, 'occupancy') - - @pytest.mark.integration - def test_model_atom_counts(self, sample_cif_file): - """Test that model has consistent atom counts.""" - from torchref.model.model import Model - - model = Model() - model.load_cif(str(sample_cif_file)) - - # All arrays should have same number of atoms - n_atoms = model.xyz().shape[0] - assert model.xyz().shape[0] == n_atoms - assert model.adp().shape[0] == n_atoms - assert model.occupancy().shape[0] == n_atoms - - @pytest.mark.integration - def test_model_cell_parameters(self, sample_cif_file): - """Test that model has valid cell parameters.""" - from torchref.model.model import Model - - model = Model() - model.load_cif(str(sample_cif_file)) - - # Cell should have 6 parameters - assert len(model.cell) == 6 - # All cell parameters should be positive - assert all(p > 0 for p in model.cell[:3].tolist()) # a, b, c - # Angles should be reasonable (0-180) - assert all(0 < p <= 180 for p in model.cell[3:].tolist()) # alpha, beta, gamma - - @pytest.mark.integration - def test_model_spacegroup(self, sample_cif_file): - """Test that model has a valid spacegroup.""" - from torchref.model.model import Model - - model = Model() - model.load_cif(str(sample_cif_file)) - - # Spacegroup should be set (can be string or gemmi.SpaceGroup) - assert model.spacegroup is not None - # Check it can be converted to string representation - assert len(str(model.spacegroup)) > 0 - - @pytest.mark.integration - def test_model_element_types(self, sample_cif_file): - """Test that element types are recognized.""" - from torchref.model.model import Model - - model = Model() - model.load_cif(str(sample_cif_file)) - - # Should have pdb DataFrame with element column - assert hasattr(model, 'pdb') - assert 'element' in model.pdb.columns - # Elements should be strings like 'C', 'N', 'O', etc. - elements = set(model.pdb['element'].unique()) - common_elements = {'C', 'N', 'O', 'S', 'H', 'CA', 'MG', 'ZN', 'FE'} - # At least some elements should be recognized - assert len(elements.intersection(common_elements)) > 0 or len(elements) > 0 - - -class TestMultipleCIFFiles: - """Tests that load multiple CIF files.""" - - @pytest.mark.integration - @pytest.mark.slow - def test_load_all_test_structures(self, all_cif_files): - """Test loading all available test structures.""" - from torchref.model.model import Model - - loaded = 0 - errors = [] - - for cif_file in all_cif_files: - try: - model = Model() - model.load_cif(str(cif_file)) - n_atoms = model.xyz().shape[0] - assert n_atoms > 0 - loaded += 1 - except Exception as e: - errors.append((cif_file.name, str(e))) - - # Report - print(f"\nLoaded {loaded}/{len(all_cif_files)} structures") - if errors: - print(f"Errors: {errors}") - - # Should load at least most structures - assert loaded > 0 +from torchref.config import canonical_device, get_default_device, get_float_dtype + + +@pytest.mark.integration +def test_cif_loading_contract(loaded_model, sample_cif_file) -> None: + """A deposited CIF supplies aligned atomic tensors and its crystal metadata.""" + import gemmi + + model = loaded_model + reference = gemmi.read_structure(str(sample_cif_file)) + xyz, adp, occupancy = model.xyz(), model.adp(), model.occupancy() + assert xyz.shape == (len(model.pdb), 3) + assert len(xyz) > 0 + assert adp.shape == occupancy.shape == (len(xyz),) + for tensor in (xyz, adp, occupancy, model.cell.data): + assert tensor.dtype == get_float_dtype() + assert canonical_device(tensor.device) == canonical_device(get_default_device()) + assert torch.isfinite(tensor).all() + assert torch.all(adp >= 0) + assert {"x", "y", "z", "element", "resname", "chainid", "resseq"} <= set( + model.pdb.columns + ) + assert {"C", "N", "O"} <= set(model.pdb.element) + torch.testing.assert_close( + model.cell.data, xyz.new_tensor(reference.cell.parameters) + ) + assert ( + model.spacegroup.number + == gemmi.find_spacegroup_by_name(reference.spacegroup_hm).number + ) class TestCIFSaving: @@ -125,22 +46,26 @@ class TestCIFSaving: def test_save_and_reload_cif(self, sample_cif_file, tmp_path): """Test saving a model to CIF and reloading it.""" from torchref.model.model import Model - + # Load original model1 = Model() model1.load_cif(str(sample_cif_file)) n_atoms1 = model1.xyz().shape[0] - + # Save to temp file using write_pdb (CIF saving may not exist) output_path = tmp_path / "test_output.pdb" model1.write_pdb(str(output_path)) - + assert output_path.exists() - - # Reload - model2 = Model() + + # add_hydrogens=False on reload: what is under test is whether the written + # file round-trips, not whether generation reruns. Regenerating on reload can + # legitimately differ, because ``write_pdb`` does not emit LINK records -- so a + # metal-coordinated nitrogen comes back with a free valence and takes a hydrogen + # it did not have before. + model2 = Model(add_hydrogens=False) model2.load_pdb(str(output_path)) n_atoms2 = model2.xyz().shape[0] - + # Compare atom counts assert n_atoms2 == n_atoms1 diff --git a/tests/integration/test_io_reflections.py b/tests/integration/test_io_reflections.py index 79642253..cb9e9a1b 100644 --- a/tests/integration/test_io_reflections.py +++ b/tests/integration/test_io_reflections.py @@ -6,106 +6,72 @@ import pytest import torch -from pathlib import Path - -class TestMTZLoading: - """Tests for loading MTZ reflection files.""" - - @pytest.mark.integration - def test_load_mtz_file(self, sample_mtz_file): - """Test loading a real MTZ file.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Should have reflections loaded - assert hasattr(data, 'hkl') - assert hasattr(data, 'F') - assert data.hkl is not None - - @pytest.mark.integration - def test_mtz_reflection_counts(self, sample_mtz_file): - """Test that MTZ has consistent reflection counts.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - n_refl = data.hkl.shape[0] - assert n_refl > 0 - - # F should match hkl count - if data.F is not None: - assert data.F.shape[0] == n_refl - - @pytest.mark.integration - def test_mtz_hkl_indices(self, sample_mtz_file): - """Test HKL indices are valid integers.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # HKL should have 3 columns - assert data.hkl.shape[1] == 3 - - # Should contain integer-like values - hkl_rounded = torch.round(data.hkl) - assert torch.allclose(data.hkl, hkl_rounded) - - @pytest.mark.integration - def test_mtz_cell_parameters(self, sample_mtz_file): - """Test that MTZ has valid cell parameters.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - if hasattr(data, 'cell') and data.cell is not None: - assert len(data.cell) == 6 - assert all(c > 0 for c in data.cell[:3].tolist()) - - @pytest.mark.integration - def test_mtz_spacegroup(self, sample_mtz_file): - """Test that MTZ has a valid spacegroup.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Spacegroup should be set (can be string or gemmi.SpaceGroup) - assert data.spacegroup is not None - # Check it can be converted to string representation - assert len(str(data.spacegroup)) > 0 - - @pytest.mark.integration - def test_mtz_sigma_values(self, sample_mtz_file): - """Test that sigma values are loaded.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - if hasattr(data, 'F_sigma') and data.F_sigma is not None: - assert data.F_sigma.shape[0] == data.F.shape[0] - # Check that non-NaN sigma values are positive - valid_sigma = data.F_sigma[~torch.isnan(data.F_sigma)] - if len(valid_sigma) > 0: - assert torch.all(valid_sigma > 0) - - @pytest.mark.integration - def test_mtz_rfree_flags(self, sample_mtz_file): - """Test that R-free flags are loaded or generated.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Should have rfree_flags (loaded or generated) - if hasattr(data, 'rfree_flags') and data.rfree_flags is not None: - assert data.rfree_flags.shape[0] == data.hkl.shape[0] +from torchref.config import ( + canonical_device, + get_default_device, + get_float_dtype, + get_int_dtype, +) + + +@pytest.mark.integration +def test_mtz_loading_contract(loaded_reflection_data, sample_mtz_file) -> None: + """MTZ loading supplies aligned observations, masks and crystal metadata.""" + import gemmi + + data = loaded_reflection_data + reference = gemmi.read_mtz_file(str(sample_mtz_file)) + n = len(data.hkl) + assert n > 0 + assert data.hkl.shape == (n, 3) + assert data.hkl.dtype == get_int_dtype() + assert canonical_device(data.hkl.device) == canonical_device(get_default_device()) + for tensor in (data.F, data.F_sigma, data.resolution): + assert tensor.shape == (n,) + assert tensor.dtype == get_float_dtype() + assert canonical_device(tensor.device) == canonical_device(get_default_device()) + mask = data.masks() + assert mask.shape == (n,) + assert mask.dtype == torch.bool + assert mask.any() + assert torch.isfinite(data.F[mask]).all() + assert torch.all(data.F[mask] >= 0) + assert torch.isfinite(data.F_sigma[mask]).all() + assert torch.all(data.F_sigma[mask] > 0) + assert torch.isfinite(data.resolution).all() + assert torch.all(data.resolution > 0) + assert data.rfree_flags.shape == (n,) + assert data.rfree_flags.dtype == torch.bool + assert data.rfree_flags.any() and (~data.rfree_flags).any() + assert 0.7 < data.rfree_flags.to(get_float_dtype()).mean().item() < 1.0 + torch.testing.assert_close( + data.cell.data, data.F.new_tensor(reference.cell.parameters) + ) + assert data.spacegroup.number == reference.spacegroup.number + + +@pytest.mark.integration +def test_resolution_bins(loaded_reflection_data) -> None: + """Every bin mean equals the mean d-spacing of its unmasked reflections.""" + data = loaded_reflection_data + bins, n_bins = data.get_bins(n_bins=10) + assert bins.shape == (len(data.hkl),) + assert n_bins > 0 + assert bins.min() >= 0 and bins.max() < n_bins + groups = [(bins == i) & data.masks() for i in range(n_bins)] + assert all(group.any() for group in groups) + expected = torch.stack([data.resolution[group].mean() for group in groups]) + torch.testing.assert_close(data.mean_res_per_bin(), expected) + + +@pytest.mark.integration +def test_structure_pair_consistency(model_and_data) -> None: + """Matching model and reflection files describe the same crystal.""" + model, data = model_and_data["model"], model_and_data["data"] + assert len(model.xyz()) > 0 and len(data.hkl) > 0 + torch.testing.assert_close(model.cell.data, data.cell.data, rtol=0.01, atol=0.1) + assert model.spacegroup.number == data.spacegroup.number class TestSFCIFLoading: @@ -115,95 +81,27 @@ class TestSFCIFLoading: def test_load_sf_cif(self, sample_structure_factor_cif): """Test loading a structure factor CIF file.""" from torchref.io import ReflectionData - + data = ReflectionData() data.load_cif(str(sample_structure_factor_cif)) - - assert data.hkl is not None + + assert data.hkl.shape[0] > 0 + assert data.hkl.shape[1] == 3 class TestReflectionDataProperties: """Tests for computed properties of reflection data.""" - @pytest.mark.integration - def test_resolution_calculation(self, sample_mtz_file): - """Test resolution can be calculated from loaded data.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Should have resolution attribute - if hasattr(data, 'resolution') and data.resolution is not None: - assert torch.all(data.resolution > 0) - assert torch.all(torch.isfinite(data.resolution)) - - @pytest.mark.integration - def test_wilson_b_factor(self, sample_mtz_file): - """Test Wilson B-factor is calculated.""" - from torchref.io import ReflectionData - - data = ReflectionData() - data.load_mtz(str(sample_mtz_file)) - - # Wilson B should be calculated during loading - if hasattr(data, 'wilson_b') and data.wilson_b is not None: - assert data.wilson_b > 0 - @pytest.mark.integration def test_data_device_movement(self, sample_mtz_file, cpu_device): """Test moving reflection data to different devices.""" from torchref.io import ReflectionData - + data = ReflectionData() data.load_mtz(str(sample_mtz_file)) - + # Move to device data = data.to(cpu_device) - - # Tensors should be on correct device - if data.hkl is not None: - assert data.hkl.device == cpu_device - if data.F is not None: - assert data.F.device == cpu_device - - -class TestMatchingDataPairs: - """Tests using matching model and reflection data.""" - @pytest.mark.integration - def test_load_structure_pair(self, sample_structure_pair): - """Test loading matching model and reflection data.""" - from torchref.model.model import Model - from torchref.io import ReflectionData - - model = Model() - model.load_cif(str(sample_structure_pair["model"])) - - data = ReflectionData() - data.load_mtz(str(sample_structure_pair["reflections"])) - - # Both should load successfully - n_atoms = model.xyz().shape[0] - assert n_atoms > 0 - assert data.hkl is not None - - @pytest.mark.integration - def test_cell_consistency(self, sample_structure_pair): - """Test that model and data have consistent cell parameters.""" - from torchref.model.model import Model - from torchref.io import ReflectionData - - model = Model() - model.load_cif(str(sample_structure_pair["model"])) - - data = ReflectionData() - data.load_mtz(str(sample_structure_pair["reflections"])) - - # Cell parameters should be similar (may have small differences) - if hasattr(data, 'cell') and data.cell is not None: - model_cell = torch.tensor(model.cell) - data_cell = torch.tensor(data.cell) - - # Allow 1% tolerance for cell parameters - assert torch.allclose(model_cell, data_cell, rtol=0.01, atol=0.1) + assert data.hkl.device == cpu_device + assert data.F.device == cpu_device diff --git a/tests/integration/test_model_operations.py b/tests/integration/test_model_operations.py index 49b3a48c..2db27886 100644 --- a/tests/integration/test_model_operations.py +++ b/tests/integration/test_model_operations.py @@ -289,7 +289,12 @@ def test_model_roundtrip_pdb(self, sample_cif_file, tmp_path): output_path = tmp_path / "output.pdb" model1.write_pdb(str(output_path)) - model2 = Model() + # add_hydrogens=False on reload: what is under test is whether the written + # file round-trips, not whether generation reruns. Regenerating on reload can + # legitimately differ, because ``write_pdb`` does not emit LINK records -- so a + # metal-coordinated nitrogen comes back with a free valence and takes a hydrogen + # it did not have before. + model2 = Model(add_hydrogens=False) model2.load_pdb(str(output_path)) n_atoms2 = model2.xyz().shape[0] diff --git a/tests/integration/test_refinement_pipeline.py b/tests/integration/test_refinement_pipeline.py index 4f385578..94cb9258 100644 --- a/tests/integration/test_refinement_pipeline.py +++ b/tests/integration/test_refinement_pipeline.py @@ -44,7 +44,7 @@ def test_setup_refinement_components(self, sample_structure_pair): def test_restraints_from_model(self, sample_cif_file): """Test building restraints from a loaded model.""" from torchref.model.model import Model - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints model = Model() model.load_cif(str(sample_cif_file)) diff --git a/tests/integration/test_rigid_body_isolation.py b/tests/integration/test_rigid_body_isolation.py new file mode 100644 index 00000000..97ca1587 --- /dev/null +++ b/tests/integration/test_rigid_body_isolation.py @@ -0,0 +1,132 @@ +"""refine_rigid_body must not disturb the refinement it is called on. + +Every cutoff rebinds the step's Refinement to a resolution-truncated data view, +which rebuilds the scaler and every target and drops the loss state. Those +rebuilds are needed for the x-ray target, whose data and scaler genuinely +change; they are collateral for the ADP and geometry targets, which are built +from the model alone and are dropped from the loss state before the optimizer +runs. `RigidBodyRefinementStep` therefore runs against a sandbox clone, and +these tests pin that the caller sees none of it. +""" +import pytest +import torch + +from torchref import LBFGSRefinement + + +@pytest.fixture(scope="module") +def refinement(mtz_dir, pdb_dir): + def build(): + return LBFGSRefinement( + data_file=str(mtz_dir / "3E98.mtz"), + pdb=str(pdb_dir / "3E98.pdb"), + device=torch.device("cpu"), + verbose=0, + ) + return build + + +def test_component_restraint_config_survives(refinement): + """A sigma set on a target must still be set afterwards. + + Regression: `_init_targets` rebuilt `TotalADPTarget` per cutoff with no + restraint parameters, so this silently reverted to the ADPSimilarityTarget + default of 2.0 and refinement continued at a restraint weight nobody chose. + """ + ref = refinement() + ref.adp_target["simu"].simu_sigma = 0.25 + ref.get_scales() + + ref.refine_rigid_body(iterations_per_step=10) + + assert ref.adp_target["simu"].simu_sigma == pytest.approx(0.25) + + +def test_custom_loss_state_weight_survives(refinement): + """A weight registered on the LossState must still be registered afterwards. + + `adp/simu` is deliberately a key absent from DEFAULT_GROUP_WEIGHTS: a key + that is present gets visibly overwritten, while a custom one silently + disappears and its target falls back to the group weight. + """ + ref = refinement() + ref.get_scales() + ref.complete_loss_state().set_weight("adp/simu", 0.77) + + ref.refine_rigid_body(iterations_per_step=10) + + assert ref.complete_loss_state().weights.get("adp/simu") == pytest.approx(0.77) + + +def test_targets_and_data_are_not_replaced(refinement): + """Object identity, not just values -- a caller may hold its own references.""" + ref = refinement() + ref.get_scales() + adp, geometry, data = ref.adp_target, ref.geometry_target, ref.reflection_data + + ref.refine_rigid_body(iterations_per_step=10) + + assert ref.adp_target is adp + assert ref.geometry_target is geometry + assert ref.reflection_data is data + + +def test_resolution_range_survives_a_coarse_only_cutoff_list(refinement): + """The caller's resolution range must be what it was, not the last cutoff's. + + `cut_res` masks in place and returns `self`, so a `cutoffs` list ending above + the native d_min can leave the caller truncated. Object identity does not + catch it -- the data object is the same one throughout. + """ + ref = refinement() + data = ref.reflection_data + ref.get_scales() + n_before = int(data.masks().sum()) + n_work_before = int(data.work.mask.sum()) + d_min_before = data.get_max_res() + + # Deliberately coarse-only, and deliberately not ending at the native d_min. + ref.refine_rigid_body(iterations_per_step=5, cutoffs=[6.0, 4.0]) + + assert int(data.masks().sum()) == n_before + assert int(data.work.mask.sum()) == n_work_before + assert data.get_max_res() == pytest.approx(d_min_before) + + +def test_a_caller_supplied_resolution_limit_is_not_widened(refinement): + """A refinement built with `max_res` keeps that limit across a rigid-body run.""" + ref = refinement() + ref.reflection_data.cut_res(highres=3.5) + data = ref.reflection_data + ref.get_scales() + n_before = int(data.masks().sum()) + + ref.refine_rigid_body(iterations_per_step=5) + + assert int(data.masks().sum()) == n_before + assert data.get_max_res() >= 3.5 + + +def test_refined_coordinates_still_reach_the_caller(refinement): + """The sandbox shares the model, so the whole point still has to work.""" + ref = refinement() + ref.get_scales() + before = ref.model.xyz().detach().clone() + + ref.refine_rigid_body(iterations_per_step=30) + + shift = (ref.model.xyz().detach() - before).norm(dim=-1) + assert float(shift.max()) > 0.0, "rigid body moved nothing" + + +def test_refinement_is_usable_afterwards(refinement): + """A normal macrocycle must still run against the caller's own objects.""" + ref = refinement() + ref.get_scales() + + ref.refine_rigid_body(iterations_per_step=10) + ref.refine_scaler() + ref.refine_adp() + + rwork, rfree = ref.get_rfactor() + assert 0.0 < rwork < 1.0 and 0.0 < rfree < 1.0 diff --git a/tests/integration/test_sfds_device.py b/tests/integration/test_sfds_device.py index 7a6ac700..8a1d0806 100644 --- a/tests/integration/test_sfds_device.py +++ b/tests/integration/test_sfds_device.py @@ -29,8 +29,12 @@ def _atoms(device, n=8): @pytest.mark.integration def test_sfds_same_device_cpu(): """Sanity: hkl already on the module device works and stays on it.""" + from torchref.model.context import ModelContext + from torchref.symmetry import SpaceGroup + cell = Cell(_CELL, device="cpu") - sf = SfDS(cell, spacegroup="P212121").to("cpu") + ctx = ModelContext(cell=cell, spacegroup=SpaceGroup("P212121", device="cpu")) + sf = SfDS(ctx).to("cpu") xyz, adp, occ, A, B = _atoms("cpu") hkl = torch.randint(-6, 7, (50, 3)).float() F, _ = sf.compute_structure_factors(hkl, xyz, adp, occ, A, B) @@ -42,9 +46,13 @@ def test_sfds_same_device_cpu(): @pytest.mark.integration def test_sfds_hkl_on_different_device(): """hkl on CPU while the module + atoms are on CUDA must not crash.""" + from torchref.model.context import ModelContext + from torchref.symmetry import SpaceGroup + cuda = torch.device("cuda") cell = Cell(_CELL, device=cuda) - sf = SfDS(cell, spacegroup="P212121").to(cuda) + ctx = ModelContext(cell=cell, spacegroup=SpaceGroup("P212121", device=cuda)) + sf = SfDS(ctx).to(cuda) xyz, adp, occ, A, B = _atoms(cuda) hkl_cpu = torch.randint(-6, 7, (50, 3)).float() # deliberately on CPU diff --git a/tests/integration/test_structure_compatibility.py b/tests/integration/test_structure_compatibility.py new file mode 100644 index 00000000..2ad00c10 --- /dev/null +++ b/tests/integration/test_structure_compatibility.py @@ -0,0 +1,101 @@ +"""Exercise explicitly named extra structure files in the slow compatibility tier.""" + +import pytest +import torch + +from tests.helpers.structure_cases import ( + EXTENDED_PAIR_CODES, + MODEL_CIF_FILES, + MTZ_CODES, + SF_CIF_CODES, +) +from torchref.config import ( + canonical_device, + get_default_device, + get_float_dtype, + get_int_dtype, +) + +pytestmark = pytest.mark.integration + + +@pytest.mark.parametrize( + "directory, expected", + [ + ("cif", MODEL_CIF_FILES), + ("mtz", tuple(f"{code}.mtz" for code in MTZ_CODES)), + ("cif_sf", tuple(f"{code}-sf.cif" for code in SF_CIF_CODES)), + ], + ids=["models", "mtz", "sf-cif"], +) +def test_compatibility_inventory(test_files_dir, directory, expected) -> None: + """Every bundled input has an explicit quick or extended coverage assignment.""" + suffix = ".mtz" if directory == "mtz" else ".cif" + actual = {path.name for path in (test_files_dir / directory).glob(f"*{suffix}")} + assert actual == set(expected) + + +@pytest.mark.slow +@pytest.mark.parametrize( + "filename", [name for name in MODEL_CIF_FILES if name != "1DAW.cif"] +) +def test_model_cif_compatibility(cif_dir, filename) -> None: + """Each extra CIF loads atoms and finite symmetry operators on the default device.""" + from torchref.model import Model + + path = cif_dir / filename + assert path.is_file() + model = Model(verbose=0).load_cif(str(path)) + xyz = model.xyz() + assert xyz.shape == (len(model.pdb), 3) + assert len(xyz) > 0 + assert xyz.dtype == get_float_dtype() + assert canonical_device(xyz.device) == canonical_device(get_default_device()) + assert torch.isfinite(xyz).all() + assert model.cell.data.shape == (6,) + assert model.spacegroup.matrices.shape[0] > 0 + assert torch.isfinite(model.spacegroup.matrices).all() + + +@pytest.mark.slow +@pytest.mark.parametrize("code", EXTENDED_PAIR_CODES) +def test_modelft_cif_compatibility(cif_dir, code) -> None: + """Fourier models initialize their scattering parametrization in distinct crystals.""" + from torchref.model import ModelFT + + path = cif_dir / f"{code}.cif" + assert path.is_file() + model = ModelFT(max_res=3.0, verbose=0).load_cif(str(path)) + assert len(model.xyz()) > 0 + assert model.parametrization + assert all(size > 0 for size in model.grid_shape) + + +@pytest.mark.slow +@pytest.mark.parametrize( + "directory, filename, loader", + [("mtz", f"{code}.mtz", "load_mtz") for code in MTZ_CODES if code != "1DAW"] + + [ + ("cif_sf", f"{code}-sf.cif", "load_cif") + for code in SF_CIF_CODES + if code != "1DAW" + ], + ids=[f"mtz-{code}" for code in MTZ_CODES if code != "1DAW"] + + [f"sf-cif-{code}" for code in SF_CIF_CODES if code != "1DAW"], +) +def test_reflection_file_compatibility( + test_files_dir, directory, filename, loader +) -> None: + """Every named MTZ/SF-CIF must load reflections; one success cannot mask another failure.""" + from torchref.io import ReflectionData + + path = test_files_dir / directory / filename + assert path.is_file() + data = ReflectionData(verbose=0) + getattr(data, loader)(str(path)) + assert data.hkl.shape == (len(data.hkl), 3) + assert len(data.hkl) > 0 + assert data.hkl.dtype == get_int_dtype() + assert canonical_device(data.hkl.device) == canonical_device(get_default_device()) + assert data.cell.data.shape == (6,) + assert data.F.shape == (len(data.hkl),) diff --git a/tests/integration/test_symmetry_integration.py b/tests/integration/test_symmetry_integration.py index 9ccedc52..a39ce804 100644 --- a/tests/integration/test_symmetry_integration.py +++ b/tests/integration/test_symmetry_integration.py @@ -6,7 +6,6 @@ import pytest import torch -from pathlib import Path class TestSpaceGroupInitialization: @@ -106,8 +105,8 @@ class TestSpaceGroupDevice: @pytest.mark.integration def test_spacegroup_default_device(self): """Test SpaceGroup matrices land on the configured default device.""" - from torchref.symmetry import SpaceGroup from torchref.config import get_default_device + from torchref.symmetry import SpaceGroup sg = SpaceGroup("P 21 21 21") @@ -140,7 +139,7 @@ def test_expand_coordinates(self, sample_cif_file): # The model should be able to generate symmetry mates # Check if there's an expand method - if hasattr(sg, 'expand') or hasattr(sg, 'expand_atoms'): + if hasattr(sg, "expand") or hasattr(sg, "expand_atoms"): expanded = sg.expand(xyz) assert expanded.shape[0] >= xyz.shape[0] @@ -148,23 +147,6 @@ def test_expand_coordinates(self, sample_cif_file): class TestSpacegroupVariants: """Tests for different spacegroup conventions.""" - @pytest.mark.integration - @pytest.mark.parametrize("sg_name", [ - "P 1", # Triclinic - "P 21", # Monoclinic - "P 21 21 21", # Orthorhombic - "P 43 21 2", # Tetragonal - "P 3 2 1", # Trigonal - "P 6 2 2", # Hexagonal - "P 2 3", # Cubic - ]) - def test_common_spacegroups(self, sg_name): - """Test loading common spacegroups.""" - from torchref.symmetry import SpaceGroup - - sg = SpaceGroup(sg_name) - assert sg.matrices is not None - @pytest.mark.integration def test_spacegroup_name_variations(self): """Test that different spacegroup name formats work.""" @@ -181,25 +163,6 @@ def test_spacegroup_name_variations(self): class TestSpaceGroupWithData: """Tests for SpaceGroup with real crystallographic data.""" - @pytest.mark.integration - def test_spacegroup_with_multiple_structures(self, cif_dir): - """Test SpaceGroup for multiple structures.""" - from torchref.model.model import Model - from torchref.symmetry import SpaceGroup - - cif_files = list(cif_dir.glob("*.cif"))[:3] - - for cif_file in cif_files: - model = Model() - model.load_cif(str(cif_file)) - - sg = SpaceGroup(model.spacegroup) - - # Should have valid matrices - assert sg.matrices is not None - assert sg.matrices.shape[0] >= 1 - assert torch.all(torch.isfinite(sg.matrices)) - @pytest.mark.integration def test_spacegroup_consistent_with_cell(self, sample_cif_file): """Test that SpaceGroup is consistent with unit cell.""" @@ -221,18 +184,21 @@ class TestMapSymmetry: """Tests for map symmetry operations.""" @pytest.mark.integration - def test_map_symmetry_initialization(self, sample_cif_file): - """Test map symmetry initialization.""" + def test_symmetrize_map_round_trip(self, sample_cif_file): + """Symmetrizing a map goes through the space group and preserves shape.""" + import torch + from torchref.model.model import Model - from torchref.symmetry.map_symmetry import MapSymmetry model = Model() model.load_cif(str(sample_cif_file)) - # Check if MapSymmetry can be initialized - try: - map_sym = MapSymmetry(model.spacegroup, model.cell) - assert map_sym is not None - except (TypeError, AttributeError): - # May not support all initialization patterns - pass + sg = model.spacegroup + shape = sg.suggest_grid_size((16, 16, 16)) + density = torch.rand(shape, device=sg.device, dtype=sg.dtype) + + symmetrized = sg.symmetrize_map(density) + + assert symmetrized.shape == density.shape + # The operator is cached for this shape and dropped on a device move. + assert sg.map_operator(shape) is sg.map_operator(shape) diff --git a/tests/integration/test_two_moment_intensity.py b/tests/integration/test_two_moment_intensity.py new file mode 100644 index 00000000..9a4c2b1e --- /dev/null +++ b/tests/integration/test_two_moment_intensity.py @@ -0,0 +1,314 @@ +"""Two-moment target limits, gradients, invalid observations and weight calibration.""" + +import pytest +import torch + +pytestmark = pytest.mark.integration + + +def _target(dc, mc, scaler, **kw): + from torchref.refinement.targets import CollectionTwoMomentIntensityTarget + + return CollectionTwoMomentIntensityTarget(dc, mc, scaler=scaler, verbose=0, **kw) + + +class TestCoherentLimit: + def test_lambda_zero_reduces_to_the_squared_mean(self, difference_collection): + dc, mc, scaler = difference_collection + mc.set_lambda_twin(0.0) + target = _target(dc, mc, scaler) + + model = target.intensity_model(recalc=True) + + rows = target._row_indices(target._keys()) + weights = mc.fractions_matrix()[rows] + components = dc.component_structure_factors(mc, recalc=False) + mean = scaler.forward_batched( + mc.mix_component_fcalcs(components, weights), weights + ) + assert torch.equal(model, mean.abs() ** 2) + + def test_lambda_zero_skips_nonfinite_derivatives( + self, difference_collection, monkeypatch + ): + """Zero dispersion must not evaluate a potentially non-finite derivative.""" + dc, mc, scaler = difference_collection + mc.set_lambda_twin(0.0) + weights = mc.fractions_matrix() + monkeypatch.setattr(mc, "fractions_matrix", lambda: weights) + + def forbidden(): + raise AssertionError("coherent prediction evaluated its variance branch") + + monkeypatch.setattr(mc, "activation_jacobian", forbidden) + assert torch.isfinite(_target(dc, mc, scaler).forward()) + + def test_full_dispersion_matches_incoherent_intensity(self, difference_collection): + """Fully dark or fully lit crystals mix in intensity, including solvent.""" + dc, mc, scaler = difference_collection + mc.set_lambda_twin(1.0) + target = _target(dc, mc, scaler) + actual = target.intensity_model(recalc=True) + components = dc.component_structure_factors(mc, recalc=False) + basis = torch.eye( + mc.n_base_models, device=components.device, dtype=components.real.dtype + ) + pure = scaler.forward_batched(components, basis) + expected = mc.fractions_matrix() @ pure.abs().square() + # Complex mixture sums lose relative precision near solvent cancellation. + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=1e-5) + + def test_a_nonzero_lambda_changes_the_prediction(self, difference_collection): + """Anti-vacuity: the variance branch must actually do something.""" + dc, mc, scaler = difference_collection + mc.set_lambda_twin(0.0) + coherent = _target(dc, mc, scaler).intensity_model(recalc=True) + + mc.set_lambda_twin(0.5) + dispersed = _target(dc, mc, scaler).intensity_model(recalc=True) + + assert not torch.allclose(coherent, dispersed) + # Strictly positive: |dF|^2 has no sign. + assert bool((dispersed >= coherent - 1e-6).all()) + + +class TestForwardModelStructure: + def test_the_variance_term_is_sigma_sq_times_the_scaled_jacobian( + self, difference_collection + ): + dc, mc, scaler = difference_collection + mc.set_lambda_twin(0.4) + target = _target(dc, mc, scaler) + total = target.intensity_model(recalc=True) + + rows = target._row_indices(target._keys()) + components = dc.component_structure_factors(mc, recalc=False) + weights = mc.fractions_matrix()[rows] + jacobian = mc.activation_jacobian()[rows] + + mean = scaler.forward_batched( + mc.mix_component_fcalcs(components, weights), weights + ) + deriv = scaler.forward_batched( + mc.mix_component_fcalcs(components, jacobian), jacobian + ) + expected = mean.abs() ** 2 + mc.sigma_alpha_sq * deriv.abs() ** 2 + assert torch.allclose(total, expected, rtol=1e-6) + + def test_the_reference_row_carries_no_variance(self, difference_collection): + """The dark's Jacobian row is exactly zero, so its prediction is coherent + regardless of the dispersion -- a dark dataset holds no activation information. + """ + dc, mc, scaler = difference_collection + keys = _target(dc, mc, scaler)._keys() + assert keys[0] == "dark" + + mc.set_lambda_twin(0.0) + coherent = _target(dc, mc, scaler).intensity_model(recalc=True)[0] + mc.set_lambda_twin(0.9) + dispersed = _target(dc, mc, scaler).intensity_model(recalc=True)[0] + + assert torch.allclose(coherent, dispersed, rtol=1e-6) + + def test_shape_follows_the_fitted_keys(self, difference_collection): + dc, mc, scaler = difference_collection + target = _target(dc, mc, scaler) + model = target.intensity_model(recalc=True) + assert model.shape == (len(target._keys()), len(dc.hkl)) + + +class TestLossAndReporting: + def test_forward_is_finite_and_positive(self, difference_collection): + dc, mc, scaler = difference_collection + loss = _target(dc, mc, scaler).forward() + assert torch.isfinite(loss) + assert loss.numel() == 1 + + def test_gradient_reaches_the_light_model(self, difference_collection): + dc, mc, scaler = difference_collection + target = _target(dc, mc, scaler) + target.forward().backward() + grad = mc.base_models[1].xyz.refinable_params.grad + assert grad is not None and torch.isfinite(grad).all() + assert float(grad.abs().max()) > 0 + + def test_gradient_reaches_the_dispersion_when_refinable( + self, difference_collection + ): + dc, mc, scaler = difference_collection + mc.set_lambda_twin(0.3, refinable=True) + _target(dc, mc, scaler).forward().backward() + grad = mc._lambda_logit.grad + assert grad is not None and torch.isfinite(grad).all() + assert float(grad.abs()) > 0 + + def test_rfactor_uses_the_two_moment_amplitude(self, difference_collection): + dc, mc, scaler = difference_collection + target = _target(dc, mc, scaler) + + rf = target.get_rfactor() + assert set(rf) == {"per_dataset", "rwork_pct", "rfree_pct"} + assert set(rf["per_dataset"]) == set(target._keys()) + for key, (rwork, rfree) in rf["per_dataset"].items(): + assert 0.0 < rwork < 2.0, f"{key}: {rwork}" + assert 0.0 < rfree < 2.0, f"{key}: {rfree}" + + def test_stats_report_the_activation_moments(self, difference_collection): + dc, mc, scaler = difference_collection + mc.set_lambda_twin(0.25) + stats = _target(dc, mc, scaler).stats() + for key in ( + "alpha_mean", + "lambda_twin", + "sigma_alpha_sq", + "alpha_sd", + "dI_frac", + "rwork", + "rfree", + "loss", + ): + assert key in stats, f"missing stat: {key}" + assert stats["alpha_mean"].value == pytest.approx(0.22, abs=1e-4) + assert stats["lambda_twin"].value == pytest.approx(0.25, abs=1e-4) + assert stats["dI_frac"].value > 0.0 + + def test_di_frac_is_zero_in_the_coherent_limit(self, difference_collection): + """The stat that distinguishes "refined to zero" from "never refined".""" + dc, mc, scaler = difference_collection + mc.set_lambda_twin(0.0) + assert _target(dc, mc, scaler).stats()["dI_frac"].value == 0.0 + + @pytest.mark.parametrize("use_set", ["work", "free"]) + def test_subset_selection_is_honoured(self, difference_collection, use_set): + dc, mc, scaler = difference_collection + target = _target(dc, mc, scaler, use_set=use_set) + assert target.use_set == use_set + expected = sum( + (dc[k].work if use_set == "work" else dc[k].free).n for k in target._keys() + ) + assert target._n_reflections() == expected + + +class TestIntensityRequirement: + def test_construction_fails_without_intensities(self, difference_models): + """Reject absent intensity columns before evaluating the loss.""" + from torchref.refinement.targets import CollectionTwoMomentIntensityTarget + + dc, mc = difference_models + for data in dc.values(): + data.I = data.I_sigma = None + with pytest.raises(ValueError, match="I/SIGI"): + CollectionTwoMomentIntensityTarget(dc, mc, verbose=0) + + +class TestNonFiniteObservations: + """Real reflection files carry non-finite intensities, and they must not reach the + gradient.""" + + def test_a_nan_observation_does_not_poison_the_gradient( + self, difference_collection + ): + dc, mc, scaler = difference_collection + data = dc["light"] + with torch.no_grad(): + data.I[5] = float("nan") + data.I[11] = float("inf") + + target = _target(dc, mc, scaler) + loss = target.forward() + assert torch.isfinite(loss), "loss went non-finite" + + loss.backward() + grad = mc.base_models[1].xyz.refinable_params.grad + assert grad is not None + assert torch.isfinite(grad).all(), ( + "non-finite observations reached the gradient; every optimizer step " + "would be rejected and the model would not move" + ) + + def test_a_nan_sigma_does_not_poison_the_gradient(self, difference_collection): + dc, mc, scaler = difference_collection + data = dc["light"] + with torch.no_grad(): + data.I_sigma[7] = float("nan") + + target = _target(dc, mc, scaler) + loss = target.forward() + loss.backward() + grad = mc.base_models[1].xyz.refinable_params.grad + assert torch.isfinite(loss) and torch.isfinite(grad).all() + + def test_the_bad_reflections_are_excluded_not_absorbed(self, difference_collection): + """They must drop out of the sum, not contribute a large finite penalty -- + otherwise the loss depends on how many reflections the file happened to reject. + """ + dc, mc, scaler = difference_collection + data = dc["light"] + target = _target(dc, mc, scaler) + baseline = target.forward().item() + with torch.no_grad(): + data.I[3] = float("nan") + with_nan = _target(dc, mc, scaler).forward().item() + + # One reflection out of tens of thousands: the loss should drop slightly, not + # jump by a penalty term. + assert with_nan <= baseline + assert abs(with_nan - baseline) / baseline < 1e-2 + + +class TestWeightCalibration: + """Intensities are squared amplitudes, so this target's gradient is orders of + magnitude away from the amplitude target beside it. Left uncalibrated it swamps the + geometry restraints and buys R-free by moving the model further than the data + supports.""" + + def test_calibration_equalises_the_gradient_norms(self, difference_collection): + dc, mc, scaler = difference_collection + from torchref.refinement.targets import CollectionDifferenceTarget + + params = [p for p in mc.base_models[1].parameters() if p.requires_grad] + diff = CollectionDifferenceTarget(dc, mc, scaler=scaler, verbose=0) + target = _target(dc, mc, scaler) + + target.calibrate_base_weight(diff, params) + + def gnorm(t): + g = torch.autograd.grad(t.forward(), params, allow_unused=True) + return sum(float((x**2).sum()) for x in g if x is not None) ** 0.5 + + assert gnorm(target) == pytest.approx(gnorm(diff), rel=0.05) + + def test_the_ratio_argument_scales_the_result(self, difference_collection): + dc, mc, scaler = difference_collection + from torchref.refinement.targets import CollectionDifferenceTarget + + params = [p for p in mc.base_models[1].parameters() if p.requires_grad] + diff = CollectionDifferenceTarget(dc, mc, scaler=scaler, verbose=0) + + a = _target(dc, mc, scaler) + b = _target(dc, mc, scaler) + wa = a.calibrate_base_weight(diff, params, ratio=1.0) + wb = b.calibrate_base_weight(diff, params, ratio=0.25) + assert wb == pytest.approx(0.25 * wa, rel=1e-3) + + def test_base_weight_scales_the_loss_on_the_work_set(self, difference_collection): + dc, mc, scaler = difference_collection + one = _target(dc, mc, scaler, base_weight=1.0).forward().item() + three = _target(dc, mc, scaler, base_weight=3.0).forward().item() + assert three == pytest.approx(3.0 * one, rel=1e-5) + + def test_the_free_set_value_is_left_unweighted(self, difference_collection): + """The free-set number is a diagnostic and has to stay comparable across + weightings.""" + dc, mc, scaler = difference_collection + one = _target(dc, mc, scaler, use_set="free", base_weight=1.0).forward().item() + five = _target(dc, mc, scaler, use_set="free", base_weight=5.0).forward().item() + assert five == pytest.approx(one, rel=1e-6) + + def test_calibration_needs_refinable_parameters(self, difference_collection): + dc, mc, scaler = difference_collection + from torchref.refinement.targets import CollectionDifferenceTarget + + diff = CollectionDifferenceTarget(dc, mc, scaler=scaler, verbose=0) + with pytest.raises(ValueError, match="No refinable parameters"): + _target(dc, mc, scaler).calibrate_base_weight(diff, []) diff --git a/tests/unit/alignment/__init__.py b/tests/unit/alignment/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/alignment/test_anisotropy_fit.py b/tests/unit/alignment/test_anisotropy_fit.py new file mode 100644 index 00000000..26848e0d --- /dev/null +++ b/tests/unit/alignment/test_anisotropy_fit.py @@ -0,0 +1,146 @@ +"""The overall-anisotropy fit, on data whose anisotropy is known by construction. + +Synthetic Wilson intensities: acentric reflections are exponentially +distributed about their shell mean, centric ones follow a chi-squared with one +degree of freedom, and the anisotropy enters as ``exp(-2 pi^2 s.U.s)`` on the +mean. Feeding that in and asking for U back is the only way to separate a +correct fit from one that merely returns something plausible. + +The first test is the one that matters: **isotropic data must give back zero +anisotropy**. Fitting the same relation in log space instead is biased -- +``E[ln(I/)]`` is ``-gamma``, not zero -- and with no constant term that +offset can only be absorbed by the quadratic form. It comes back as tens of +square Angstrom of anisotropy that is not in the data. +""" + +import math + +import pytest +import torch + +from torchref.experimental.alignment.sh import ( + assign_shells, + equal_count_shell_edges, + fit_overall_anisotropy, +) + +pytestmark = pytest.mark.unit + +B_PER_U = 8.0 * math.pi ** 2 + + +def _synthetic(U_true, n=40000, seed=0, centric_fraction=0.0): + """Wilson intensities carrying exactly ``U_true``, on a 4-15 A shell.""" + g = torch.Generator().manual_seed(seed) + smag = 1.0 / (4.0 + 11.0 * torch.rand(n, generator=g, dtype=torch.float64)) + ct = 2 * torch.rand(n, generator=g, dtype=torch.float64) - 1 + phi = 2 * math.pi * torch.rand(n, generator=g, dtype=torch.float64) + st = (1 - ct * ct).clamp(min=0).sqrt() + s = torch.stack([smag * st * torch.cos(phi), + smag * st * torch.sin(phi), + smag * ct], dim=-1) + + # Shell mean falls off with resolution, times the anisotropic term. + sigma_shell = torch.exp(-20.0 * smag * smag) + aniso = torch.exp(-2.0 * (math.pi ** 2) * torch.einsum( + "ni,ij,nj->n", s, U_true.to(torch.float64), s)) + mean_I = sigma_shell * aniso + + centric = torch.rand(n, generator=g, dtype=torch.float64) < centric_fraction + # Acentric: I = mean * Exp(1). Centric: I = mean * chi^2_1. + e = -torch.log(torch.rand(n, generator=g, dtype=torch.float64).clamp(min=1e-300)) + z = torch.randn(n, generator=g, dtype=torch.float64) ** 2 + I = mean_I * torch.where(centric, z, e) + F = I.clamp(min=0).sqrt() + + edges, _ = equal_count_shell_edges(smag, 20) + return F, s, assign_shells(smag, edges), centric + + +def _spread_B(U): + ev = torch.linalg.eigvalsh(U.to(torch.float64)) * B_PER_U + return float(ev[2] - ev[0]) + + +#: Reflections used by the isotropic-data tests. The estimator's own scatter on +#: the B spread is about 7 A^2 at this count, and falls as 1/sqrt(n). +_N_ISO = 40000 + +#: Threshold for "no anisotropy detected" at ``_N_ISO``: above the estimator's +#: measured scatter (median 6.9, max 8.2 over six seeds) with margin. This is a +#: noise floor, not a bias tolerance -- see the scaling test below. +_ISO_TOLERANCE_B = 12.0 + + +@pytest.mark.parametrize("centric_fraction", [0.0, 0.15]) +def test_isotropic_data_gives_no_anisotropy(centric_fraction): + """Zero anisotropy in, nothing but estimation noise out.""" + F, s, idx, cen = _synthetic( + torch.zeros(3, 3), n=_N_ISO, seed=1, centric_fraction=centric_fraction) + U = fit_overall_anisotropy(F, s, idx, cen, P=20) + assert _spread_B(U) < _ISO_TOLERANCE_B, ( + f"isotropic data produced {_spread_B(U):.1f} A^2 of B anisotropy, " + f"beyond this estimator's scatter at n={_N_ISO}" + ) + + +def test_the_isotropic_residual_is_noise_not_bias(): + """The spurious spread must shrink as 1/sqrt(n), not plateau. + + This is the real test of the fit's centring. A biased estimator -- for + instance one regressing ``ln(I/)`` with no constant term, where + ``E[ln(I/)] = -gamma`` has to be absorbed by the quadratic form -- gives + a spread that stays put as reflections are added. An unbiased one averages + it away. + """ + def median_spread(n): + vals = [] + for k in range(3): + F, s, idx, cen = _synthetic(torch.zeros(3, 3), n=n, seed=100 + k) + vals.append(_spread_B(fit_overall_anisotropy(F, s, idx, cen, P=20))) + return sorted(vals)[1] + + coarse, fine = median_spread(20000), median_spread(320000) + # 16x the reflections should buy about 4x, i.e. well over 2x even allowing + # for the scatter of a 3-seed median. + assert fine < coarse / 2.0, ( + f"B spread went {coarse:.2f} -> {fine:.2f} A^2 for 16x the reflections; " + f"an unbiased fit should shrink roughly 4x, a biased one not at all" + ) + + +def test_a_known_tensor_is_recovered(): + """Uniaxial anisotropy of a realistic size comes back to within a few A^2.""" + B_true = torch.diag(torch.tensor([-15.0, -15.0, 30.0], dtype=torch.float64)) + U_true = B_true / B_PER_U + F, s, idx, cen = _synthetic(U_true, n=60000, seed=2) + U = fit_overall_anisotropy(F, s, idx, cen, P=20) + B_fit = U.to(torch.float64) * B_PER_U + # The isotropic part is a gauge -- the per-shell normalisation removes it -- + # so compare the deviatoric parts. + dev = lambda M: M - torch.eye(3, dtype=torch.float64) * torch.diagonal(M).mean() + err = (dev(B_fit) - dev(B_true)).abs().max().item() + assert err < 6.0, f"recovered B off by {err:.1f} A^2:\n{B_fit}" + + +def test_zero_amplitudes_do_not_dominate(): + """A handful of vanishing amplitudes must not steer the fit. + + The earlier version clamped them to 1e-30 and took a logarithm, turning each + into a residual of about -69 in an unweighted least squares. + """ + F, s, idx, cen = _synthetic(torch.zeros(3, 3), seed=3) + clean = fit_overall_anisotropy(F, s, idx, cen, P=20) + F2 = F.clone() + F2[::500] = 0.0 + spiked = fit_overall_anisotropy(F2, s, idx, cen, P=20) + assert abs(_spread_B(spiked) - _spread_B(clean)) < 3.0, ( + f"zeroing 0.2% of amplitudes moved the fit from " + f"{_spread_B(clean):.2f} to {_spread_B(spiked):.2f} A^2" + ) + + +def test_too_few_reflections_returns_zero(): + F, s, idx, cen = _synthetic(torch.zeros(3, 3), n=60, seed=4) + U = fit_overall_anisotropy(F, s, idx, cen, P=20, min_count=20) + assert torch.equal(U, torch.zeros(3, 3, dtype=U.dtype)) diff --git a/tests/unit/alignment/test_conjugate_contraction.py b/tests/unit/alignment/test_conjugate_contraction.py new file mode 100644 index 00000000..3e6ce842 --- /dev/null +++ b/tests/unit/alignment/test_conjugate_contraction.py @@ -0,0 +1,91 @@ +"""The radial contraction has to *materialise* its conjugate. + +``torch.conj`` does not conjugate anything: it returns a view carrying a +conjugate *bit*, and every consumer is expected to honour it. MPS's batched +complex matmul does not -- and the contraction in +:func:`~torchref.experimental.alignment.frf.data_mr.cross_correlate_xi` is an +``einsum`` that lowers to exactly that. The failure is silent and total: the +unconjugated values are contracted instead, which on 1DAW moved ``xi`` by 173% +of its own peak magnitude, reordered the entire rotation-function peak list, and +pushed the true orientation from rank 0 out of the top 200 -- while the top +score moved by only 0.1%, so nothing looked wrong. + +Elementwise ops, ``where``, ``index_add`` and 2-D ``matmul`` all honour the bit; +only the batched path drops it. That is narrow enough that the guard has to be +the value of the contraction itself, checked on whatever device this host has, +against a reference computed on the host in double. +""" + +import pytest +import torch + +from torchref.config import get_complex_dtype, get_default_device +from torchref.experimental.alignment.frf.data_mr import cross_correlate_xi +from torchref.experimental.alignment.frf.types import BesselSHCoefficients + +pytestmark = pytest.mark.unit + +L, N_RADIAL = 6, 8 + + +def _coeffs(seed, device, dtype): + """A filled ``(N_radial, L, 2L-1)`` coefficient block.""" + g = torch.Generator().manual_seed(seed) + real = torch.randn(N_RADIAL, L, 2 * L - 1, generator=g, dtype=torch.float64) + imag = torch.randn(N_RADIAL, L, 2 * L - 1, generator=g, dtype=torch.float64) + c = torch.complex(real, imag) + return BesselSHCoefficients( + coeffs=c.to(device=device, dtype=dtype), L=L, bessel_h_scale=40.0, + ) + + +def test_the_contraction_conjugates_the_calc_side(): + """On this host's device, against the same sum done on the host in double. + + A dropped conjugation is not a small error -- it changes the sign of every + imaginary part in one operand -- so the bar can be tight without being + brittle about float32 rounding. + """ + device, cplx = get_default_device(), get_complex_dtype() + obs, calc = _coeffs(0, device, cplx), _coeffs(1, device, cplx) + got = cross_correlate_xi(obs, calc) + + host_obs = BesselSHCoefficients( + coeffs=obs.coeffs.cpu().to(torch.complex128), L=L, bessel_h_scale=40.0) + host_calc = BesselSHCoefficients( + coeffs=calc.coeffs.cpu().to(torch.complex128), L=L, bessel_h_scale=40.0) + ref = torch.einsum( + "rln,rlm->lmn", + host_obs.coeffs, + torch.conj(host_calc.coeffs).resolve_conj(), + ) + + err = (got.cpu().to(torch.complex128) - ref).abs().max() + scale = ref.abs().max() + assert float(err / scale) < 1e-5, ( + f"contraction is {float(err / scale):.2e} away from the host double " + f"reference on {device}; a dropped conjugation shows up here as O(1)" + ) + + +def test_dropping_the_conjugation_would_be_caught(): + """The guard above has to be able to see the failure it exists for. + + Contracting the unconjugated coefficients is what a lost conjugate bit + produces, so that has to land far outside the tolerance -- otherwise the + test would pass on the broken path too. + """ + device, cplx = get_default_device(), get_complex_dtype() + obs, calc = _coeffs(0, device, cplx), _coeffs(1, device, cplx) + ref = torch.einsum( + "rln,rlm->lmn", + obs.coeffs.cpu().to(torch.complex128), + torch.conj(calc.coeffs.cpu().to(torch.complex128)).resolve_conj(), + ) + unconjugated = torch.einsum( + "rln,rlm->lmn", + obs.coeffs.cpu().to(torch.complex128), + calc.coeffs.cpu().to(torch.complex128), + ) + rel = float((unconjugated - ref).abs().max() / ref.abs().max()) + assert rel > 0.1, f"the two differ by only {rel:.2e}; this guard is blind" diff --git a/tests/unit/alignment/test_likelihood_convention.py b/tests/unit/alignment/test_likelihood_convention.py new file mode 100644 index 00000000..7a85fb56 --- /dev/null +++ b/tests/unit/alignment/test_likelihood_convention.py @@ -0,0 +1,86 @@ +"""The translation likelihood's variance convention, pinned as a pdf. + +A likelihood is only right up to what its variance argument means, and that is +exactly the sort of thing that survives code review and unit tests written +against the implementation rather than against the distribution. The alignment +package carried its own Rice and Woolfson for a while, parameterised by the +*amplitude* variance, and handed both branches the same number -- which put +acentrics at twice their intended variance on 90-95% of reflections, for as long +as nobody integrated the thing. + +These tests integrate it. They assert the property the pipeline actually depends +on: at ``D = 0`` the likelihood must believe `` = 1``, because +:class:`~torchref.scaling.WilsonNormaliser` makes `` = 1`` an identity of +the fit that produced ``E_obs``. Anything else means the model and the data +disagree about the scale of the very quantity being compared. +""" +import pytest +import torch + +from torchref.base.targets.xray_likelihoods import rice_per_refl + +pytestmark = pytest.mark.unit + + +def _pdf_moments(logp, F, dF): + """Norm and second moment of ``exp(logp)`` treated as a density in ``F``.""" + p = torch.exp(logp) + norm = float((p * dF).sum()) + return norm, float((p * F**2 * dF).sum()) / norm + + +@pytest.fixture(scope="module") +def grid(): + F = torch.linspace(1e-6, 12.0, 200001, dtype=torch.float64) + return F, float(F[1] - F[0]) + + +@pytest.mark.parametrize("centric", [False, True]) +def test_unit_sigma_means_unit_second_moment(grid, centric): + """At Sigma = 1 and no model, the likelihood expects = 1. + + This is the property the whole sigma_A path rests on. Both branches must + satisfy it from the SAME Sigma -- that is what makes a single complex + variance the right parameterisation and an amplitude variance the wrong one. + """ + F, dF = grid + ll = -rice_per_refl(F, torch.zeros_like(F), torch.ones_like(F), + torch.full_like(F, centric, dtype=torch.bool)) + norm, m2 = _pdf_moments(ll, F, dF) + assert norm == pytest.approx(1.0, rel=1e-4), "not a normalised density" + assert m2 == pytest.approx(1.0, rel=1e-4), ( + f"{'centric' if centric else 'acentric'} branch expects = {m2:.4f} " + f"at Sigma = 1; the observations have = 1 by construction" + ) + + +@pytest.mark.parametrize("sigma", [0.25, 0.5, 2.0]) +@pytest.mark.parametrize("centric", [False, True]) +def test_second_moment_tracks_sigma(grid, sigma, centric): + """ = Sigma with no model, for both branches. Fixes the scale, not just the shape.""" + F, dF = grid + ll = -rice_per_refl(F, torch.zeros_like(F), torch.full_like(F, sigma), + torch.full_like(F, centric, dtype=torch.bool)) + norm, m2 = _pdf_moments(ll, F, dF) + assert norm == pytest.approx(1.0, rel=1e-4) + assert m2 == pytest.approx(sigma, rel=1e-4) + + +@pytest.mark.parametrize("centric", [False, True]) +def test_second_moment_with_a_model_present(grid, centric): + """With a model, = Sigma + Fc^2 -- the signal adds to the noise. + + The sigma_A likelihood is evaluated at ``Fc = D E_calc`` and + ``Sigma = 1 - D^2``, so on data normalised to `` = 1`` this gives + `` = 1`` for every D. That invariance is why D is identifiable at all, + and it fails if the two branches disagree about what the variance means. + """ + F, dF = grid + D, E_calc = 0.6, 1.0 + Fc, Sigma = D * E_calc, 1.0 - D * D + ll = -rice_per_refl(F, torch.full_like(F, Fc), torch.full_like(F, Sigma), + torch.full_like(F, centric, dtype=torch.bool)) + norm, m2 = _pdf_moments(ll, F, dF) + assert norm == pytest.approx(1.0, rel=1e-4) + assert m2 == pytest.approx(Sigma + Fc**2, rel=1e-4) + assert m2 == pytest.approx(1.0, rel=1e-4), "D must not change the expected " diff --git a/tests/unit/alignment/test_patterson_translation.py b/tests/unit/alignment/test_patterson_translation.py new file mode 100644 index 00000000..d5545ef2 --- /dev/null +++ b/tests/unit/alignment/test_patterson_translation.py @@ -0,0 +1,136 @@ +"""Unit tests for the fast translation function. + +``fast_translation_function`` accumulates the Crowther-Blow coefficients of a +normalised, weighted ``E_obs^2`` against the candidate's normalised +``|E_calc(h, t)|^2`` and inverts one FFT. With the search model at canonical +positions and ``F_obs`` derived from a translated copy of the same model, the +top peak (or one of the top three) must land at ``-t_true`` modulo an allowed +origin shift of the space group -- the translation that would bring the search +model into agreement with the observed data -- and the likelihood must prefer +that peak. + +``TranslationObs`` carries the observed side. It is built once here, as the +pipeline builds it once per run, because normalisation and weighting are +properties of the observations and do not change when the model moves. +""" +from pathlib import Path + +import numpy as np +import pytest +import torch + +from torchref.experimental.alignment.translation import ( + TranslationObs, + fast_translation_function, + llg_at_translations, + prepare_candidate, +) +from torchref.io.datasets.reflection_data import ReflectionData +from torchref.model import ModelFT + + +TEST_FILES = Path(__file__).resolve().parents[2] / "files" +PDB_1DAW = TEST_FILES / "pdb" / "1DAW.pdb" +MTZ_1DAW = TEST_FILES / "mtz" / "1DAW.mtz" + + +@pytest.fixture(scope="module") +def setup(): + canonical = ModelFT().load_pdb(str(PDB_1DAW)) + data = ReflectionData().load_mtz(str(MTZ_1DAW)) + # The pipeline's default window: the rotation search's 15-4 A. + # The resolution arithmetic runs on the host in double -- `.cpu()` before + # the widening, since a backend without float64 cannot hold the wide copy -- + # and only the resulting boolean goes back to where the data live. + rec = data.cell.reciprocal_basis_matrix.cpu().to(torch.float64) + s = (data.hkl.cpu().to(torch.float64) @ rec).norm(dim=-1) + window = ((s >= 1.0 / 15.0) & (s <= 1.0 / 4.0)).to(data.hkl.device) + mask = data.get_valid_mask() & window + return canonical, data, mask + + +def _search(canonical, data, mask, t_true): + """Peaks and their likelihoods for a search model translated by ``t_true``.""" + with torch.no_grad(): + F_obs = canonical(data.hkl[mask]).abs() + model_p1 = canonical.copy() + model_p1.max_res = 4.0 / 1.5 + model_p1.spacegroup = "P 1" + if t_true is not None: + model_p1 = model_p1.translate( + torch.tensor(t_true, dtype=canonical.dtype_float), fractional=True, + ) + obs = TranslationObs.build(F_obs, data.hkl[mask], data.spacegroup, data.cell) + cand = prepare_candidate(model_p1, obs, data.spacegroup, data.cell) + _, peaks = fast_translation_function( + obs, cand, data.cell, grid_spacing_A=4.0 / 3.0, n_peaks=3, + cluster_radius_A=4.0, + ) + assert len(peaks) > 0 + llg = llg_at_translations( + obs, cand, + torch.as_tensor(np.stack([p.translation for p in peaks]), dtype=torch.float64), + ) + return peaks, llg + + +def _xz_dist_to_origin_class(t: np.ndarray, t_true: np.ndarray) -> float: + """Distance of ``t + t_true`` from an allowed origin in C2, x and z only. + + C2's origin is free along y; the centring makes (1/2, 1/2, 0) a lattice + vector and (0, *, 1/2) an allowed shift, so x and z are each determined + only modulo 1/2. + """ + d = np.array([t[0] + t_true[0], t[2] + t_true[2]]) + d = (d + 0.25) % 0.5 - 0.25 + return float(np.linalg.norm(d)) + + +@pytest.mark.unit +@pytest.mark.slow +def test_fast_tf_zero_translation(setup): + """Un-translated model: the likelihood's pick sits at an origin-equivalent.""" + canonical, data, mask = setup + peaks, llg = _search(canonical, data, mask, None) + best = peaks[int(llg.argmax())] + dist = _xz_dist_to_origin_class(best.translation, np.zeros(3)) + assert dist < 0.03, ( + f"likelihood pick misses an origin-equivalent by {dist:.3f}; " + f"peaks: {[p.translation.round(3).tolist() for p in peaks]}" + ) + + +@pytest.mark.unit +@pytest.mark.slow +def test_fast_tf_recovers_known_translation(setup): + """A model translated by t_true: the likelihood's pick is at -t_true (mod origins).""" + canonical, data, mask = setup + t_true = np.array([0.18, -0.07, 0.23]) + peaks, llg = _search(canonical, data, mask, t_true) + best = peaks[int(llg.argmax())] + dist = _xz_dist_to_origin_class(best.translation, t_true) + assert dist < 0.03, ( + f"likelihood pick misses -t_true by {dist:.3f}; " + f"peaks: {[p.translation.round(3).tolist() for p in peaks]}" + ) + # The fast map's own top peak should already be the right one here; the + # likelihood is the arbiter when it is not. + assert _xz_dist_to_origin_class(peaks[0].translation, t_true) < 0.03 + + +@pytest.mark.unit +@pytest.mark.slow +def test_e_calc_is_normalised(setup): + """```` is one to within the fit's tolerance, for a placed candidate.""" + canonical, data, mask = setup + with torch.no_grad(): + F_obs = canonical(data.hkl[mask]).abs() + model_p1 = canonical.copy() + model_p1.max_res = 4.0 / 1.5 + model_p1.spacegroup = "P 1" + obs = TranslationObs.build(F_obs, data.hkl[mask], data.spacegroup, data.cell) + cand = prepare_candidate(model_p1, obs, data.spacegroup, data.cell) + # The normalisation already carries eps: E is per unit of eps*Sigma_calc. + E2 = cand.e_calc(torch.zeros(3, dtype=torch.float64)) ** 2 + mean_e2 = float(E2.mean()) + assert abs(mean_e2 - 1.0) < 0.15, mean_e2 diff --git a/tests/unit/alignment/test_peak_finder_symmetry.py b/tests/unit/alignment/test_peak_finder_symmetry.py new file mode 100644 index 00000000..55200c5b --- /dev/null +++ b/tests/unit/alignment/test_peak_finder_symmetry.py @@ -0,0 +1,58 @@ +"""Rotation-function peaks are one per orientation, not one per symmetry mate. + +The greedy SO(3) suppression treats ``R`` and its point-group mates ``R R_g`` as +the same peak when the Cartesian symmetry rotations are supplied. The group +composes on the **right** -- measured on real peak lists, where the left orbit +finds no coincident pairs and the right orbit finds every mate -- so the test +pins that side too: a left-composed copy must survive as a distinct peak. +""" +import math + +import pytest +import torch + +from torchref.experimental.alignment.frf.peak_finder import find_rotation_peaks +from torchref.experimental.alignment.frf.rotation_utils import ( + edmonds_euler_from_rotation_matrix, + rotation_matrix_from_edmonds_euler, +) +from torchref.experimental.alignment.frf.types import AdaptiveRotationFunction + +pytestmark = pytest.mark.unit + + +def _rz(deg: float) -> torch.Tensor: + c, s = math.cos(math.radians(deg)), math.sin(math.radians(deg)) + return torch.tensor([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=torch.float64) + + +def _arf(rotations, values) -> AdaptiveRotationFunction: + eul = torch.tensor([edmonds_euler_from_rotation_matrix(R) for R in rotations], + dtype=torch.float64) + n = eul.shape[0] + return AdaptiveRotationFunction( + alphas=eul[:, 0], betas=eul[:, 1], gammas=eul[:, 2], + values=torch.tensor(values, dtype=torch.float64), + beta_starts=torch.tensor([0, n]), beta_grid=torch.tensor([0.0]), + grid_sampling_deg=3.0, + ) + + +def test_symmetry_mates_collapse_to_one_peak_and_the_side_is_right(): + sym_cart = torch.stack([_rz(0.0), _rz(90.0), _rz(180.0), _rz(270.0)]) # 4 about z + R1 = rotation_matrix_from_edmonds_euler(0.3, 0.7, 1.1) + R_right = R1 @ _rz(90.0) # a mate: the group acts on the right + R_left = _rz(90.0) @ R1 # not a mate of R1 for a generic R1 + R3 = rotation_matrix_from_edmonds_euler(2.0, 1.3, 0.4) + arf = _arf([R1, R_right, R_left, R3], [10.0, 9.0, 8.5, 8.0]) + + plain = find_rotation_peaks(arf, n_peaks=10, sigma_threshold=-50.0, nms_radius_deg=6.0) + assert len(plain) == 4, "without symmetry every sample is its own peak" + + dedup = find_rotation_peaks(arf, n_peaks=10, sigma_threshold=-50.0, + nms_radius_deg=6.0, sym_cart=sym_cart) + scores = sorted(p.score for p in dedup) + assert scores == [8.0, 8.5, 10.0], ( + "the right-composed mate (9.0) must be suppressed and the " + f"left-composed copy (8.5) kept; got {scores}" + ) diff --git a/tests/unit/alignment/test_phaser_modelprep.py b/tests/unit/alignment/test_phaser_modelprep.py new file mode 100644 index 00000000..7d59aa1e --- /dev/null +++ b/tests/unit/alignment/test_phaser_modelprep.py @@ -0,0 +1,153 @@ +"""Unit tests for the three Phaser model-prep helpers added to +`torchref.experimental.alignment.frf.preprocessing`: + +- ``bulk_solvent_factor`` (Babinet, Phaser solTerm.h:9) +- ``oeffner_vrms`` (Oeffner empirical, rms_estimate.cc:37) +- ``fit_relative_wilson_b`` (Wilson-B regression, EnsemblePDB.cc:793-851) +""" +from __future__ import annotations + +import math + +import pytest +import torch + +from torchref.experimental.alignment.frf.preprocessing import ( + bulk_solvent_factor, + fit_relative_wilson_b, + oeffner_vrms, +) + + +# ----------------------------------------------------------------------------- +# bulk_solvent_factor +# ----------------------------------------------------------------------------- + + +def test_bulk_solvent_factor_low_res_suppressed(): + """At s→0 the Babinet term → 1 − fsol ≈ 0.05 (default fsol=0.95).""" + s = torch.tensor([1e-6, 1e-5]) # essentially zero + val = bulk_solvent_factor(s, fsol=0.95, bsol=300.0) + assert torch.allclose(val, torch.tensor([0.05, 0.05]), atol=1e-3) + + +def test_bulk_solvent_factor_high_res_unaffected(): + """At large s the exponential vanishes → term → 1.""" + s = torch.tensor([1.0, 2.0]) # very high res + val = bulk_solvent_factor(s, fsol=0.95, bsol=300.0) + assert torch.allclose(val, torch.ones_like(s), atol=1e-6) + + +def test_bulk_solvent_factor_sigA_min_clamp(): + """Clamp to sigA_min when Babinet would go below.""" + s = torch.tensor([1e-10]) + val = bulk_solvent_factor(s, fsol=0.999, bsol=300.0, sigA_min=0.05) + # 1 - 0.999·1 = 0.001 < sigA_min=0.05 → clamped to 0.05 + assert val.item() == pytest.approx(0.05, abs=1e-6) + + +def test_bulk_solvent_factor_monotonic_in_s(): + """Strictly non-decreasing in |s|.""" + s = torch.linspace(0.001, 1.0, 100) + val = bulk_solvent_factor(s) + diffs = val[1:] - val[:-1] + assert (diffs >= -1e-10).all() + + +# ----------------------------------------------------------------------------- +# oeffner_vrms +# ----------------------------------------------------------------------------- + + +def test_oeffner_vrms_small_perfect_model(): + """N_res = 125 (lower clamp), identity = 1 (perfect): + vrms = 0.0569 · (173 + 125)^(1/3) · 1 ≈ 0.379 Å. + """ + val = oeffner_vrms(125, identity=1.0) + expected = 0.0569 * (173 + 125) ** (1.0 / 3.0) + assert val == pytest.approx(expected, abs=1e-6) + assert 0.35 < val < 0.45 + + +def test_oeffner_vrms_large_perfect_model(): + """N_res = 1500 (upper clamp): vrms = 0.0569 · 1673^(1/3) ≈ 0.675 Å.""" + val = oeffner_vrms(1500, identity=1.0) + expected = 0.0569 * (173 + 1500) ** (1.0 / 3.0) + assert val == pytest.approx(expected, abs=1e-6) + assert 0.6 < val < 0.7 + + +def test_oeffner_vrms_clamps_inputs(): + """N_res < 125 clamps to 125; N_res > 1500 clamps to 1500.""" + low = oeffner_vrms(50, identity=1.0) + expected_125 = oeffner_vrms(125, identity=1.0) + assert low == pytest.approx(expected_125) + high = oeffner_vrms(5000, identity=1.0) + expected_1500 = oeffner_vrms(1500, identity=1.0) + assert high == pytest.approx(expected_1500) + + +def test_oeffner_vrms_identity_dependence(): + """Lower identity → larger vrms (exp(C·(1-ident)) grows).""" + v_perfect = oeffner_vrms(500, identity=1.0) + v_30pct = oeffner_vrms(500, identity=0.3) + assert v_30pct > v_perfect + # exp(1.52 · 0.7) ≈ 2.9× larger + assert v_30pct / v_perfect == pytest.approx(math.exp(1.52 * 0.7), rel=1e-6) + + +# ----------------------------------------------------------------------------- +# fit_relative_wilson_b +# ----------------------------------------------------------------------------- + + +def test_fit_relative_wilson_b_recovers_synthetic_B(): + """Synthetic: F_calc = F_obs · exp(-B_true · s² / 4) should recover B_true. + + log(/) = log(exp(B·s²/2)) = B/2 · s². Slope = B/2, and + `WilsonB = -2·slope` per Phaser EnsemblePDB.cc:850 means recovery is + NEGATIVE: the fit returns -B_true (because we put the B on F_calc, the + regression sees Σ_N/Σ_P = exp(+B·s²/2) → positive slope → returns -B). + + To test recovery of a POSITIVE B (model more disordered than data), we + multiply F_obs by exp(-B·s²/4) instead → Σ_N/Σ_P = exp(-B·s²/2) + → negative slope → returns positive B. + """ + torch.manual_seed(0) + N = 5000 + s = torch.linspace(0.05, 0.5, N, dtype=torch.float64) # 2 - 20 Å range + F_obs_base = (1.0 + 0.05 * torch.randn(N, dtype=torch.float64)).abs() + 0.1 + F_calc = F_obs_base.clone() + # Apply B = +10 to F_obs (obs more disordered): Σ_N/Σ_P = exp(-10·s²/2) + B_true = 10.0 + F_obs = F_obs_base * torch.exp(-B_true * s * s / 4.0) + val = fit_relative_wilson_b(F_obs, F_calc, s, n_shells=20) + assert val == pytest.approx(B_true, abs=1.5) + + +def test_fit_relative_wilson_b_zero_when_matched(): + """F_obs ≈ F_calc → B ≈ 0.""" + torch.manual_seed(0) + N = 5000 + s = torch.linspace(0.05, 0.5, N, dtype=torch.float64) + F = (1.0 + 0.05 * torch.randn(N, dtype=torch.float64)).abs() + 0.1 + val = fit_relative_wilson_b(F, F.clone(), s, n_shells=20) + assert abs(val) < 1.0 # Should be near zero modulo discretisation + + +def test_fit_relative_wilson_b_clamp(): + """Clamp guard: extreme synthetic case clamps to ±clamp_b.""" + N = 5000 + s = torch.linspace(0.05, 0.5, N, dtype=torch.float64) + F_obs = torch.ones(N, dtype=torch.float64) * 1e-10 # vanishing obs + F_calc = torch.ones(N, dtype=torch.float64) + val = fit_relative_wilson_b(F_obs, F_calc, s, n_shells=20, clamp_b=50.0) + assert -50.0 <= val <= 50.0 + + +def test_fit_relative_wilson_b_handles_sparse_data(): + """Too few shells contribute (all at low res) → returns 0 (no fit).""" + s = torch.linspace(0.001, 0.005, 20, dtype=torch.float64) # all > 200 Å + F = torch.ones(20, dtype=torch.float64) + val = fit_relative_wilson_b(F, F.clone(), s, n_shells=10) + assert val == 0.0 diff --git a/tests/unit/alignment/test_pipeline_aliasing.py b/tests/unit/alignment/test_pipeline_aliasing.py new file mode 100644 index 00000000..e08e753a --- /dev/null +++ b/tests/unit/alignment/test_pipeline_aliasing.py @@ -0,0 +1,134 @@ +"""Guards for two aliasing traps in the molecular-replacement pipeline. + +`Model.rotate` and `Model.translate` mutate in place and return ``self``. Every +call site that treats them as returning a fresh model therefore has to +``.copy()`` first. `MolecularReplacementPipeline._make_rotated` is called once +per rotation candidate off the same ``self.model``, so without the copy +candidate *k+1* is evaluated at an orientation composed on top of candidate *k* +and ``self.model`` is destroyed along the way. + +`Model.spacegroup` is a property, but ``SpaceGroup`` is an ``nn.Module``: +assigning a SpaceGroup *object* is intercepted by ``nn.Module.__setattr__``, +stored in ``_modules`` under the property's own name, and the setter never runs. +The pipeline builds P1 copies for its dense-transform stages, so a silent no-op +there means the "P1 search model" still carries the crystal symmetry. + +Both are pinned here rather than only in the slow integration tests, which +``--run-slow`` gates off by default. +""" + +import math + +import pytest +import torch + +pytestmark = pytest.mark.unit + + +@pytest.fixture(scope="module") +def small_model(pdb_dir): + from torchref.model import ModelFT + + p = pdb_dir / "1DAW.pdb" + if not p.exists(): + pytest.skip("1DAW.pdb not available") + return ModelFT(verbose=0).load_pdb(str(p)) + + +def _peak(alpha, beta, gamma): + from torchref.experimental.alignment.frf.types import RotationPeak + + return RotationPeak(alpha=alpha, beta=beta, gamma=gamma, score=1.0, sigma=1.0) + + +def test_rotate_mutates_in_place_and_returns_self(small_model): + """The premise. If this ever changes, the copies below stop being necessary.""" + m = small_model.copy() + before = m.xyz().clone() + R = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=m.dtype_float, + ) + out = m.rotate(R) + assert out is m, "rotate no longer returns self" + assert not torch.allclose(m.xyz(), before), "rotate no longer mutates in place" + + +def test_make_rotated_leaves_the_search_model_untouched(small_model): + """The pipeline's own model must survive candidate generation.""" + from torchref.experimental.alignment.pipeline import MolecularReplacementPipeline + + pipe = object.__new__(MolecularReplacementPipeline) + pipe.model = small_model + reference = small_model.xyz().clone() + + rotated, _ = pipe._make_rotated(_peak(0.3, 0.7, 1.1)) + + assert rotated is not pipe.model + assert torch.allclose(pipe.model.xyz(), reference), ( + "_make_rotated mutated the pipeline's search model" + ) + + +def test_successive_candidates_do_not_compound(small_model): + """Candidate k+1 must not be rotated on top of candidate k.""" + from torchref.experimental.alignment.pipeline import MolecularReplacementPipeline + + pipe = object.__new__(MolecularReplacementPipeline) + pipe.model = small_model + + p1 = _peak(0.3, 0.7, 1.1) + p2 = _peak(2.0, 1.3, 0.4) + + first, _ = pipe._make_rotated(p1) + second, _ = pipe._make_rotated(p2) + + # A fresh pipeline that only ever sees p2 is the ground truth for p2. + solo = object.__new__(MolecularReplacementPipeline) + solo.model = small_model.copy() + expected, _ = solo._make_rotated(p2) + + assert torch.allclose(second.xyz(), expected.xyz(), atol=1e-5), ( + "the second candidate depends on the first -- rotations are compounding" + ) + assert not torch.allclose(first.xyz(), second.xyz()), ( + "the two candidates are identical; the peaks chosen do not discriminate" + ) + + +def test_spacegroup_name_assignment_works(small_model): + """The supported form: pass the space-group NAME.""" + m = small_model.copy() + m.spacegroup = "P 1" + assert m.spacegroup.number == 1 + assert int(m.spacegroup.matrices.shape[0]) == 1 + + +def test_spacegroup_object_assignment_now_takes_effect(small_model): + """The trap this used to pin is gone, fixed at the root rather than avoided. + + ``Model.spacegroup`` is a property, but ``SpaceGroup`` is an ``nn.Module``, so + ``model.spacegroup = sg_object`` used to be intercepted by + ``nn.Module.__setattr__``, filed under ``_modules["spacegroup"]`` with the + setter never running -- the assignment silently did nothing, and the name was + then a registered child module, so the *correct* string assignment afterwards + raised ``TypeError``. Call sites worked around it by passing a name string. + + The space group now lives on ``ModelContext``, which is deliberately a + dataclass and not an ``nn.Module``, so there is nothing to intercept. Both + forms work and neither registers a submodule. + """ + from torchref.symmetry import SpaceGroup + + m = small_model.copy() + assert m.spacegroup.number != 1, "1DAW should not already be P1" + + m.spacegroup = SpaceGroup("P 1") + assert m.spacegroup.number == 1, "object assignment did not take effect" + assert "spacegroup" not in m._modules, ( + "the space group was registered as a child module -- the interception " + "this test exists for has come back" + ) + + # The string form must still work afterwards, which is what used to raise. + m.spacegroup = "P 21 21 21" + assert m.spacegroup.number == 19 diff --git a/tests/unit/alignment/test_radial_truncation.py b/tests/unit/alignment/test_radial_truncation.py new file mode 100644 index 00000000..94b88c5a --- /dev/null +++ b/tests/unit/alignment/test_radial_truncation.py @@ -0,0 +1,143 @@ +"""Pin the SH-Bessel radial band to Phaser's per-``l`` size. + +Phaser allocates the ``Elmn`` array with ``nmax = (lmax - l + 2) / 2`` radial +terms for each even ``l`` and runs ``n`` from 1 to ``nmax`` +(``DataMR.cc:894-896``). The band therefore narrows as ``l`` rises -- for +``lmax = 76`` it is 38 terms at ``l = 2`` and a single term at ``l = 76`` -- so +the high-``l`` bands cannot carry more radial detail than the reflection set +supports. + +``bessel_sh_expand`` allocates a flat ``(N_radial, L, 2L-1)`` array, where +``N_radial`` is Phaser's *widest* band (the one at ``l = 2``). That shape is +easy to misread as "every ``l`` carries ``N_radial`` radial terms"; it does not. +The ``(l, n) -> u = l + 2n + 1`` index build populates only +``n_l = (lmax_even - l)//2 + 1`` terms per ``l``, which is exactly Phaser's +``nmax``, so the truncation is already in force through the allocated support. + +These tests pin that invariant against the formula, so the agreement is +asserted rather than inferred from the array shape. + +They run under ``double_cpu`` because the band's *width* and its *populated +extent* only coincide in double precision. At the working precision the radial +weight ``sqrt(2u+1) j_u(x)/x`` underflows for high ``u`` at small ``x`` -- see +``test_the_working_precision_truncates_the_band_further`` -- so counting +non-zeros there measures float32's exponent range, not Phaser's formula. +""" + +import pytest +import torch + +from torchref.experimental.alignment.frf.data_mr import bessel_sh_expand + + +def _expand(L: int, n_points: int = 900): + """Expand a fixed pseudo-random point set; returns ``(N_radial, L, 2L-1)``.""" + g = torch.Generator().manual_seed(5) + s = torch.randn(n_points, 3, generator=g, dtype=torch.float64) + s = s / s.norm(dim=-1, keepdim=True) * ( + 0.05 + 0.15 * torch.rand(n_points, 1, generator=g, dtype=torch.float64) + ) + intensity = torch.randn(n_points, generator=g, dtype=torch.float64) + out = bessel_sh_expand(s, intensity, L=L, bessel_h_scale=30.0) + return out.coeffs.detach().cpu() + + +def _phaser_nmax(lmax_even: int, l: int) -> int: + """``(lmax - l + 2) / 2`` -- Phaser's radial band size for one ``l``.""" + return (lmax_even - l + 2) // 2 + + +@pytest.mark.parametrize("L", [21, 41, 67]) +def test_radial_band_matches_phaser_width(L, double_cpu): + """Non-zero radial indices at each ``l`` must stop at Phaser's ``nmax``. + + Our ``n`` index is 0-based against Phaser's 1-based, so the condition is + ``n < nmax(l)``. + """ + c = _expand(L) + N_radial, _, _ = c.shape + lmax_even = L - 1 if (L - 1) % 2 == 0 else L - 2 + + for l in range(2, lmax_even + 1, 2): + nz = (c[:, l, :].abs() > 0).any(dim=-1).nonzero().flatten() + if nz.numel() == 0: + continue # legitimately empty band + allowed = _phaser_nmax(lmax_even, l) + assert int(nz.max()) < allowed, ( + f"L={L} l={l}: radial index {int(nz.max())} exceeds Phaser's " + f"nmax={allowed}" + ) + + +@pytest.mark.parametrize("L", [21, 41, 67]) +def test_band_narrows_to_a_single_term_at_lmax(L, double_cpu): + """The top band keeps exactly one radial term, the widest keeps them all.""" + lmax_even = L - 1 if (L - 1) % 2 == 0 else L - 2 + assert _phaser_nmax(lmax_even, lmax_even) == 1 + c = _expand(L) + assert _phaser_nmax(lmax_even, 2) == c.shape[0], ( + "N_radial should equal Phaser's widest band, at l=2" + ) + top = (c[:, lmax_even, :].abs() > 0).any(dim=-1).nonzero().flatten() + if top.numel(): + assert int(top.max()) == 0, "l=lmax must retain only the n=0 term" + + +@pytest.mark.parametrize("L", [21, 41, 67]) +def test_band_is_exactly_phasers_width_not_merely_bounded(L, double_cpu): + """The populated band must *equal* Phaser's ``nmax(l)``, not just fit inside. + + A bound alone would also pass for an expansion that silently drops radial + terms it should keep, which would cost resolution at low ``l``. + """ + c = _expand(L) + lmax_even = L - 1 if (L - 1) % 2 == 0 else L - 2 + for l in range(2, lmax_even + 1, 2): + nz = (c[:, l, :].abs() > 0).any(dim=-1).nonzero().flatten() + assert nz.numel(), f"L={L} l={l}: band is empty" + assert int(nz.max()) + 1 == _phaser_nmax(lmax_even, l), ( + f"L={L} l={l}: {int(nz.max()) + 1} radial terms, " + f"Phaser has {_phaser_nmax(lmax_even, l)}" + ) + + +def test_allocated_width_exceeds_the_populated_band(double_cpu): + """The array is wider than the support at every l above the first. + + This is the fact that makes the array shape misleading, and the reason the + tests above assert the support rather than ``coeffs.shape``. + """ + L = 67 + lmax_even = L - 1 if (L - 1) % 2 == 0 else L - 2 + c = _expand(L) + N_radial = c.shape[0] + assert N_radial == _phaser_nmax(lmax_even, 2) + assert _phaser_nmax(lmax_even, lmax_even) == 1 + assert N_radial > _phaser_nmax(lmax_even, lmax_even), ( + "allocated width should exceed the top band's single term" + ) + + +def test_the_working_precision_truncates_the_band_further(): + """At float32 the tail of the radial band is not small, it is *absent*. + + ``j_u(x)`` for ``u >> x`` is genuinely negligible -- at ``x = 1.9`` every + ``u >= 33`` is below float32's smallest normal -- so at the working precision + the populated band is narrower than Phaser's ``nmax(l)`` wherever the shell's + Bessel argument is small. That is a property of the engine as shipped, not a + defect, and it is pinned here so it is not rediscovered as a regression: the + structural tests above deliberately run in double precision, which would + otherwise leave the shipped configuration untested. + """ + L = 67 + c = _expand(L) # working precision, no double_cpu + lmax_even = L - 1 if (L - 1) % 2 == 0 else L - 2 + populated = (c[:, 2, :].abs() > 0).any(dim=-1).nonzero().flatten() + assert populated.numel(), "l=2 band is entirely empty" + got = int(populated.max()) + 1 + phaser = _phaser_nmax(lmax_even, 2) + assert got < phaser, ( + f"expected the float32 tail to underflow: got {got} terms, Phaser has " + f"{phaser}. If this now matches, the working precision widened and the " + f"structural tests above no longer need double_cpu." + ) diff --git a/tests/unit/alignment/test_rotation_search.py b/tests/unit/alignment/test_rotation_search.py new file mode 100644 index 00000000..f5930f1a --- /dev/null +++ b/tests/unit/alignment/test_rotation_search.py @@ -0,0 +1,209 @@ +"""The rotation search's public contract: three inputs, and one convention. + +``rotation_search(model, data, model_error_A)`` is the whole surface. What a +caller most easily gets wrong is not the arguments but the *sense* of the +returned rotation -- whether to apply ``R`` or ``R.T`` to the coordinates. The +round-trip test below settles that operationally rather than by reading a +docstring: it rotates a model by a known matrix, searches, applies the inverse +of a returned solution, and requires the model back where it started, modulo the +crystal's rotational symmetry. + +These run on real data (1DAW, C2) because the search needs a real Patterson; +1DAW is the small fast case. +""" + +import math + +import pytest +import torch + +pytestmark = [pytest.mark.unit, pytest.mark.slow] + +#: The search is scored on whether truth is inside the candidate window the +#: placement search carries, not on being rank 0. +TOP_N = 20 + + +@pytest.fixture(scope="module") +def case(pdb_dir, mtz_dir): + from torchref.io.datasets.reflection_data import ReflectionData + from torchref.model import ModelFT + + pdb, mtz = pdb_dir / "1DAW.pdb", mtz_dir / "1DAW.mtz" + if not (pdb.exists() and mtz.exists()): + pytest.skip("1DAW not available") + model = ModelFT(verbose=0).load_pdb(str(pdb)) + data = ReflectionData(verbose=0).load_mtz(str(mtz)) + return model, data + + +def _rotation(seed: int) -> torch.Tensor: + """Haar-uniform SO(3) via QR with the sign correction.""" + g = torch.Generator().manual_seed(seed) + a = torch.randn(3, 3, generator=g, dtype=torch.float64) + q, r = torch.linalg.qr(a) + return q * torch.sign(torch.diagonal(r)).unsqueeze(0) + + +def _angle_deg(a: torch.Tensor, b: torch.Tensor) -> float: + a = a.detach().cpu().to(torch.float64) + b = b.detach().cpu().to(torch.float64) + tr = torch.diagonal(a @ b.T).sum().item() + return math.degrees(math.acos(max(-1.0, min(1.0, (tr - 1.0) / 2.0)))) + + +def _sym_cartesian(data) -> torch.Tensor: + """Space-group rotations as Cartesian operators, on the CPU. + + ``rotations`` is always CPU float64 while the data's tensors may sit on an + accelerator, and the Miller-index matrices have to be carried into the + Cartesian basis before they can be compared with a rotation of coordinates. + """ + from torchref.experimental.alignment.sh import hkl_symops_to_cartesian + + # `.cpu()` before the widening: the data may sit on an accelerator that + # cannot hold a float64 tensor at all, and the widening is exact on the host. + return hkl_symops_to_cartesian( + data.spacegroup.matrices.cpu().to(torch.float64), + data.cell.reciprocal_basis_matrix.cpu().to(torch.float64), + ) + + +def _kabsch(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Rotation taking ``a`` onto ``b``, both centred. CPU float64.""" + a = a.detach().cpu().to(torch.float64) + b = b.detach().cpu().to(torch.float64) + x = a - a.mean(0) + y = b - b.mean(0) + u, _, vt = torch.linalg.svd(y.T @ x) + d = torch.sign(torch.linalg.det(vt.T @ u.T)) + return vt.T @ torch.diag(torch.tensor([1.0, 1.0, d], dtype=torch.float64)) @ u.T + + +def _search_rotated(case, seed, model_error_A=0.8, n_peaks=200): + from torchref.experimental.alignment import rotation_search + + model, data = case + R_true = _rotation(seed) + rotated = model.copy().rotate( + R_true.to(model.dtype_float), center=model.xyz().mean(0), + ) + return rotated, data, R_true, rotation_search( + rotated, data, model_error_A, n_peaks=n_peaks, + ) + + +def test_returns_the_documented_shapes(case): + from torchref.experimental.alignment import RotationSolutions + + _, _, _, sol = _search_rotated(case, seed=11, n_peaks=50) + assert isinstance(sol, RotationSolutions) + n = len(sol) + assert n > 0 + assert sol.rotations.shape == (n, 3, 3) + assert sol.rotations.dtype == torch.float64 + for name in ("scores", "z_scores"): + assert getattr(sol, name).shape == (n,) + assert sol.euler_zyz.shape == (n, 3) + assert sol.lmax > 0 and sol.d_min > 0 + assert sol.model_error_A == pytest.approx(0.8) + # Best first, on the standardised scale. + z = sol.z_scores + assert torch.all(z[:-1] >= z[1:] - 1e-9) + + +def test_rotations_are_rotations(case): + _, _, _, sol = _search_rotated(case, seed=11, n_peaks=50) + R = sol.rotations + eye = torch.eye(3, dtype=torch.float64).expand_as(R) + assert torch.allclose(R @ R.transpose(-1, -2), eye, atol=1e-9) + assert torch.allclose(torch.linalg.det(R), + torch.ones(len(sol), dtype=torch.float64), atol=1e-9) + + +def test_euler_and_matrix_agree(case): + """``euler_zyz`` is the same orientation as ``rotations``, not a variant.""" + from torchref.experimental.alignment.frf.rotation_utils import ( + rotation_matrix_from_edmonds_euler, + ) + + _, _, _, sol = _search_rotated(case, seed=11, n_peaks=20) + for i in range(min(5, len(sol))): + a, b, g = sol.euler_zyz[i].tolist() + assert torch.allclose(rotation_matrix_from_edmonds_euler(a, b, g), + sol.rotations[i], atol=1e-12) + + +def test_a_solution_inverts_the_applied_rotation(case): + """The convention, settled by algebra: ``rotations[i]`` is ``S . R_true``. + + ``R_true`` is the rotation applied to the model's coordinates to build the + search model, so a returned solution composed with its inverse must leave a + symmetry operator behind. That fixes the sense of the returned matrix + without appealing to the docstring. If it ever inverts, the placement stage + silently searches the wrong orientation. + """ + _, data, R_true, sol = _search_rotated(case, seed=11) + sym_cart = _sym_cartesian(data) + best = min( + min(_angle_deg(sol.rotations[i] @ R_true.T, S) for S in sym_cart) + for i in range(min(TOP_N, len(sol))) + ) + assert best < 5.0, ( + f"no solution in the top {TOP_N} composes with R_true^-1 to within 5 " + f"degrees of a symmetry operator; closest was {best:.2f} degrees. " + f"Either the search failed on 1DAW or the convention flipped." + ) + + +def test_applying_the_transpose_places_the_model(case): + """The documented usage, at the level of coordinates. + + ``model.rotate(rotations[i].T)`` is what the docstring tells a caller to do. + Doing it to the search model must superpose it back onto the unrotated model + up to a symmetry operation -- checked through the actual rotate() call, so + the test would catch a mismatch between the docstring and the maths. + """ + rotated, data, _, sol = _search_rotated(case, seed=11) + model, _ = case + reference = model.xyz() + centre = rotated.xyz().mean(0) + sym_cart = _sym_cartesian(data) + + best = None + for i in range(min(TOP_N, len(sol))): + placed = rotated.copy().rotate( + sol.rotations[i].T.to(rotated.dtype_float).contiguous(), center=centre, + ) + residual = _kabsch(placed.xyz(), reference) + ang = min(_angle_deg(residual, S) for S in sym_cart) + best = ang if best is None else min(best, ang) + assert best is not None and best < 5.0, ( + f"applying rotations[i].T did not superpose the model onto the " + f"unrotated reference for any of the top {TOP_N}; closest residual was " + f"{best:.2f} degrees from a symmetry operator" + ) + + +def test_model_error_changes_the_result(case): + """``model_error_A`` must reach the engine. + + It sets the sigma_A fall-off, so it decides how much the high-resolution + terms count. The previous entry point accepted a coordinate error and then + overwrote it with an empirical estimate from the atom count, so the caller's + value was silently discarded; this asserts it is not. + """ + _, _, _, tight = _search_rotated(case, seed=11, model_error_A=0.2, n_peaks=20) + _, _, _, loose = _search_rotated(case, seed=11, model_error_A=2.5, n_peaks=20) + assert tight.model_error_A != loose.model_error_A + assert not torch.allclose(tight.z_scores[:5], loose.z_scores[:5], atol=1e-6), ( + "model_error_A did not change the rotation function" + ) + + +def test_uninitialised_model_is_rejected(): + from torchref.experimental.alignment import rotation_search + from torchref.model import ModelFT + + with pytest.raises(RuntimeError, match="no coordinates"): + rotation_search(ModelFT(verbose=0), None, 0.8) diff --git a/tests/unit/alignment/test_rotation_search_dtype_device.py b/tests/unit/alignment/test_rotation_search_dtype_device.py new file mode 100644 index 00000000..b6aff525 --- /dev/null +++ b/tests/unit/alignment/test_rotation_search_dtype_device.py @@ -0,0 +1,178 @@ +"""The rotation search takes its working precision from config. + +**Precision is `torchref.config`'s, not an argument's.** The expansion used to +be handed ``compute_dtype=torch.complex64`` by one call site, which made the +fused CPU Legendre kernel reachable only through that argument -- its +``BackendTable`` row gates on ``dtypes=(torch.float32,)``, so dropping the +argument silently routed to the portable path. Now the default *is* float32, so +the gate matches by construction and flipping the config flips the engine. + +Device resolution -- that both inputs are read, rather than whichever one the +code happens to touch first -- lives in ``tests/unit/utils/test_device_resolution.py``, +which exercises ``resolve_device`` directly instead of through a loaded case. +""" + +import pytest +import torch + +from torchref.config import dtypes + +pytestmark = pytest.mark.unit + + +def _tiny_reflections(n=4000, seed=0): + """A P1 shell of reflections wide enough to bin into 20 Wilson shells.""" + g = torch.Generator().manual_seed(seed) + s_mag = 0.07 + 0.18 * torch.rand(n, generator=g, dtype=torch.float64) + theta = torch.acos(2 * torch.rand(n, generator=g, dtype=torch.float64) - 1) + phi = 6.283185307179586 * torch.rand(n, generator=g, dtype=torch.float64) + s_vec = torch.stack( + [s_mag * torch.sin(theta) * torch.cos(phi), + s_mag * torch.sin(theta) * torch.sin(phi), + s_mag * torch.cos(theta)], dim=-1, + ) + F = torch.randn(n, generator=g, dtype=torch.float64).abs() + 0.1 + centric = torch.zeros(n, dtype=torch.bool) + return s_vec, F, centric + + +@pytest.mark.parametrize("float_dtype,want_complex", [ + (torch.float32, torch.complex64), + (torch.float64, torch.complex128), +]) +def test_expansion_follows_the_configured_float_dtype(float_dtype, want_complex): + from torchref.experimental.alignment.frf.data_mr import bessel_sh_expand + + s_vec, F, _ = _tiny_reflections() + original = dtypes.float, dtypes.complex + try: + dtypes.float = float_dtype + dtypes.complex = want_complex + # `s_vec` stays float64 on purpose -- the clustering keys need that + # resolution -- so this also pins that the OUTPUT dtype is not inherited + # from the input. + out = bessel_sh_expand(s_vec, F, L=12, bessel_h_scale=40.0) + finally: + dtypes.float, dtypes.complex = original + assert out.coeffs.dtype == want_complex, ( + f"expansion returned {out.coeffs.dtype} with dtypes.float={float_dtype}" + ) + + +def test_the_back_half_runs_at_the_configured_complex_dtype(): + """The whole chain -- expansion, radial sum, Wigner contraction, FFT -- is + one dtype, the configured one, so a device without float64 runs it as is. + + The radial sum used to be accumulated one step wider. Measured, narrowing it + moved scores by 1e-4 relative and reordered the deep peak list while leaving + the top peak unchanged; the placement search now consumes only the top few + distinct orientations and single precision recovers every pose on the + benchmark panel, so the wider accumulator is gone and nothing downstream + re-decides the width. + """ + from torchref.experimental.alignment.frf.data_mr import ( + bessel_sh_expand, cross_correlate_xi, + ) + from torchref.experimental.alignment.frf.sitelist_ang import ( + build_dense_map_per_beta, + ) + from torchref.experimental.alignment.frf.wigner_d import ( + wigner_contraction_per_beta, + ) + + s_vec, F, _ = _tiny_reflections() + original = dtypes.float, dtypes.complex + try: + dtypes.float, dtypes.complex = torch.float32, torch.complex64 + c = bessel_sh_expand(s_vec, F, L=12, bessel_h_scale=40.0) + xi = cross_correlate_xi(c, c) + finally: + dtypes.float, dtypes.complex = original + + assert c.coeffs.dtype == torch.complex64, "expansion should be at config dtype" + assert xi.dtype == torch.complex64, ( + f"the radial accumulation left the configured dtype: {xi.dtype}" + ) + # Everything downstream follows xi rather than re-deciding. + betas = torch.linspace(0.0, 3.0, 5, dtype=torch.float64) + S = wigner_contraction_per_beta(xi, betas) + assert S.dtype == xi.dtype, f"Wigner did not follow xi: {S.dtype}" + M = build_dense_map_per_beta(xi, betas, fft_size=48) + assert M.dtype == xi.dtype, f"FFT did not follow xi: {M.dtype}" + + +def test_wigner_blocks_carry_no_float64_to_the_device(): + """The eigendecomposition is precision-critical but belongs on the host. + + What reaches the device is the ``d^l`` blocks, bounded in [-1, 1], at the + working precision -- so an accelerator without float64 is not excluded. + """ + from torchref.experimental.alignment.frf import wigner_d + + wigner_d.clear_wigner_d_cache() + betas = torch.linspace(0.0, 3.0, 4, dtype=torch.float64) + blocks = wigner_d._wigner_d_blocks( + 8, betas, torch.device("cpu"), torch.float32, + ) + assert all(b.dtype == torch.float32 for b in blocks) + # d^l(0) = I, the cheapest correctness check on the blocks themselves. + for l, b in enumerate(blocks, start=1): + eye = torch.eye(2 * l + 1, dtype=torch.float32) + assert torch.allclose(b[0], eye, atol=1e-5), f"d^{l}(0) is not I" + + +def test_double_is_a_device_capability_not_a_constant(): + """Where precision is load-bearing, the width comes from the device. + + Two places in the engine need more than the working precision for reasons + that are not preferences: the radial accumulation cancels (see + ``cross_correlate_xi``) and the Bessel ladder's unnormalised intermediates + reach 1e157. Both used to hardcode float64, which is an error on a backend + that has none. They now ask, so that a device without float64 gets the + working dtype -- and the accuracy that implies -- instead of a crash. + """ + from torchref.config import (supports_double, widest_complex_dtype, + widest_float_dtype) + + assert supports_double("cpu") + assert widest_float_dtype("cpu") is torch.float64 + assert widest_complex_dtype("cpu") is torch.complex128 + + # The no-float64 branch cannot be exercised on a host without such a device, + # so pin the policy itself: exactly the backends known to lack float64, and + # the fallback is the configured working dtype rather than a second constant. + from torchref.config import _NO_DOUBLE_DEVICE_TYPES + assert "mps" in _NO_DOUBLE_DEVICE_TYPES + assert not supports_double(torch.device("mps")) + original = dtypes.float, dtypes.complex + try: + dtypes.float, dtypes.complex = torch.float32, torch.complex64 + assert widest_float_dtype(torch.device("mps")) is torch.float32 + assert widest_complex_dtype(torch.device("mps")) is torch.complex64 + dtypes.float, dtypes.complex = torch.float64, torch.complex128 + assert widest_float_dtype(torch.device("mps")) is torch.float64 + finally: + dtypes.float, dtypes.complex = original + + +def test_the_radial_accumulation_follows_the_configuration_not_the_input(): + """Coefficients that arrive wider than the configured dtype are brought to it. + + The width of the accumulation is a property of the run's configuration -- on + a device without float64 it is the only width there is -- not of whatever a + caller happened to hand in. + """ + from torchref.experimental.alignment.frf.data_mr import cross_correlate_xi + from torchref.experimental.alignment.frf.types import BesselSHCoefficients + + c = BesselSHCoefficients( + coeffs=torch.zeros((2, 5, 9), dtype=torch.complex128), # dtype-ok: deliberately wider than the configuration, to see it brought back + L=5, bessel_h_scale=20.0, + ) + original = dtypes.float, dtypes.complex + try: + dtypes.float, dtypes.complex = torch.float32, torch.complex64 + out = cross_correlate_xi(c, c) + finally: + dtypes.float, dtypes.complex = original + assert out.dtype == torch.complex64 diff --git a/tests/unit/alignment/test_sh.py b/tests/unit/alignment/test_sh.py new file mode 100644 index 00000000..c9c15244 --- /dev/null +++ b/tests/unit/alignment/test_sh.py @@ -0,0 +1,139 @@ +"""Leaf mathematics for the rotation function: shell binning and the Legendre reference. + +Two independent things, both pinned because something downstream trusts them. + +**Shell binning.** Two consumers deriving their own equal-count edges from the +same ``|s|`` is how boundary reflections end up in different shells depending on +which stage asked. The edges are computed once and the index passed down; these +tests pin the round trip that makes that safe. + +**``_bar_legendre_recurrence``.** Production never calls it -- the Bessel-SH +expansion runs the same recurrence inside its kernels. It exists as the +*independent* implementation that +``tests/unit/frf_separate/test_bessel_sh_grouping.py`` builds a slow reference +expansion on, to check the fused one against something other than itself. That +only works if the reference is itself trustworthy, which is what the scipy +comparison here is for: it used to reach this recurrence through ``evaluate_ylm``, +and that wrapper is gone. +""" +import math + +import numpy as np +import pytest +import torch + +from torchref.experimental.alignment.sh import ( + _bar_legendre_recurrence, + assign_shells, + equal_count_shell_edges, +) + + +# --------------------------------------------------------------------------- +# The Legendre reference the FRF expansion is checked against +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +@pytest.mark.parametrize("L", [4, 9]) +def test_bar_legendre_matches_scipy(L): + """bar_P_l^m(x) = sqrt[(2l+1)/(4pi) (l-m)!/(l+m)!] |P_l^m(x)|, against scipy. + + scipy's ``lpmv`` carries the Condon-Shortley phase and this recurrence does + not, so the comparison is on magnitude -- which is the convention the + docstring states and the kernels implement. + """ + scipy_special = pytest.importorskip("scipy.special") + + theta = torch.tensor([0.3, 1.1, 2.0, 2.9], dtype=torch.float64) + cos_t, sin_t = torch.cos(theta), torch.sin(theta) + bar_P = _bar_legendre_recurrence(cos_t, sin_t, L) # (M, L, L) + + x = cos_t.numpy() + for l in range(L): + for m in range(l + 1): + norm = math.sqrt( + (2 * l + 1) / (4 * math.pi) + * math.factorial(l - m) / math.factorial(l + m) + ) + expected = norm * np.abs(scipy_special.lpmv(m, l, x)) + np.testing.assert_allclose( + bar_P[:, l, m].abs().numpy(), expected, atol=1e-11, rtol=1e-11, + err_msg=f"l={l}, m={m}", + ) + + +@pytest.mark.unit +def test_bar_legendre_pole_values(): + """At the north pole (cos theta = 1): zero for m > 0, sqrt((2l+1)/4pi) at m = 0. + + The pole is where the recurrence is most fragile -- sin(theta) = 0 kills the + sectoral seed, and everything above it comes from the vertical step. + """ + L = 5 + theta = torch.tensor([0.0], dtype=torch.float64) + bar_P = _bar_legendre_recurrence(torch.cos(theta), torch.sin(theta), L) + for l in range(L): + for m in range(1, l + 1): + assert abs(bar_P[0, l, m].item()) < 1e-14, f"l={l}, m={m}" + np.testing.assert_allclose( + bar_P[0, l, 0].item(), math.sqrt((2 * l + 1) / (4 * math.pi)), + atol=1e-12, + ) + + +# --------------------------------------------------------------------------- +# Shell binning +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_shell_assignment_round_trip(): + """The bins really are equal-count, and every reflection lands in one. + + ``equal_count_shell_edges`` returns ``(edges, centers)`` -- the second value + is the shell mid-points, not the occupancies -- so equal-count is asserted + on the assignment rather than read off the return. + """ + torch.manual_seed(0) + s = torch.rand(1000, dtype=torch.float64) * 0.4 + 0.05 + n_shells = 12 + edges, centers = equal_count_shell_edges(s, n_shells) + idx = assign_shells(s, edges) + + assert edges.shape == (n_shells + 1,) + assert centers.shape == (n_shells,) + assert int(idx.min()) >= 0 and int(idx.max()) < n_shells + counts = torch.bincount(idx, minlength=n_shells) + assert int(counts.sum()) == s.numel(), "a reflection fell outside every shell" + # 1000 into 12 cannot divide evenly; equal-count means within one of ideal. + assert int(counts.max()) - int(counts.min()) <= 1, counts + # Centres must sit inside their own shell. + assert bool(((centers > edges[:-1]) & (centers < edges[1:])).all()) + + +@pytest.mark.unit +def test_shells_are_monotone_in_resolution(): + """A shell is a resolution range, so the bins must not interleave.""" + torch.manual_seed(1) + s = torch.rand(500, dtype=torch.float64) * 0.3 + 0.02 + edges, _ = equal_count_shell_edges(s, 8) + idx = assign_shells(s, edges) + hi = torch.stack([s[idx == b].max() for b in range(8) if (idx == b).any()]) + assert bool((hi[1:] >= hi[:-1]).all()) + + +@pytest.mark.unit +def test_assignment_is_stable_under_a_subset(): + """Slicing rows must not move a reflection's shell; rebuilding edges would. + + This is the property the "assign once, pass it down" rule rests on: given + the SAME edges, a subset of the reflections lands in the same bins it did in + the full set. + """ + torch.manual_seed(2) + s = torch.rand(600, dtype=torch.float64) * 0.3 + 0.02 + edges, _ = equal_count_shell_edges(s, 10) + full = assign_shells(s, edges) + sub = torch.arange(0, 600, 3) + torch.testing.assert_close(assign_shells(s[sub], edges), full[sub]) diff --git a/tests/unit/alignment/test_symmetry_conventions.py b/tests/unit/alignment/test_symmetry_conventions.py new file mode 100644 index 00000000..41570390 --- /dev/null +++ b/tests/unit/alignment/test_symmetry_conventions.py @@ -0,0 +1,264 @@ +"""Guard the reciprocal-space symmetry convention used by the FRF. + +Reciprocal space transforms as ``h' = h·R`` (equivalently ``Rᵀ·h``), the rule +:meth:`SpaceGroup.expand_reciprocal` implements. The alignment package re-derives +symmetry expansions inline in several places, and each of those inline copies +once used ``R·h`` instead. + +The reason that survived so long is worth stating, because it is what makes +these tests necessary rather than merely nice: **the two conventions agree +whenever the symmetry matrices are orthogonal**, which they are in every +monoclinic, orthorhombic, tetragonal and cubic setting. Only in a trigonal or +hexagonal basis, where ``S·Sᵀ ≠ I``, do they diverge -- and there the wrong +convention silently builds an orbit that mixes non-equivalent reflections, +writing conflicting ``|F|`` onto one Miller index. + +So every test here is parametrised over a space group whose matrices are *not* +orthogonal. A test that only covers P2₁2₁2₁ or P4₃2₁2 cannot fail. +""" + +from pathlib import Path + +import pytest +import torch + +from torchref.experimental.alignment.sh import ( + hkl_symops_to_cartesian, + symmetrize_anisotropy, +) +from torchref.symmetry import SpaceGroup + + +#: Space groups spanning both regimes. ``non_orthogonal`` flags the settings +#: whose rotation matrices are not orthogonal in their own basis -- the only +#: ones that can discriminate ``h·R`` from ``R·h``. +SPACEGROUPS = [ + pytest.param("P 31 2 1", True, id="P3121-trigonal"), + pytest.param("P 65 2 2", True, id="P6522-hexagonal"), + pytest.param("P 63", True, id="P63-hexagonal"), + pytest.param("P 4 3 2", False, id="P432-cubic"), + pytest.param("P 21 21 2", False, id="P21212-orthorhombic"), + pytest.param("C 1 2 1", False, id="C2-monoclinic"), +] + + +def _cell_for(hm: str): + """A cell consistent with the space group's lattice constraints.""" + if hm.startswith("P 3") or hm.startswith("P 6"): + return (60.0, 60.0, 95.0, 90.0, 90.0, 120.0) + if hm.startswith("P 4 3") or hm.startswith("P 21 21"): + return (70.0, 70.0, 70.0, 90.0, 90.0, 90.0) if "4 3" in hm else ( + 45.0, 55.0, 65.0, 90.0, 90.0, 90.0) + return (80.0, 40.0, 60.0, 90.0, 104.0, 90.0) + + +def _reciprocal_basis(cell): + """``B`` from TorchRef's own :class:`Cell`, so the convention matches by + construction rather than by a hand-rolled duplicate.""" + from torchref.symmetry import Cell + + return Cell(list(cell)).reciprocal_basis_matrix.detach().cpu().to(torch.float64) + + +@pytest.mark.parametrize("hm, non_orthogonal", SPACEGROUPS) +def test_symop_orthogonality_flags_the_hard_cases(hm, non_orthogonal): + """The premise of every other test here: which settings can discriminate. + + If this ever reports a trigonal/hexagonal group as orthogonal, the tests + below stop testing anything. + """ + S = SpaceGroup(hm).matrices.detach().cpu().to(torch.float64) + eye = torch.eye(3, dtype=torch.float64) + err = max(float((S[k] @ S[k].T - eye).abs().max()) for k in range(S.shape[0])) + assert (err > 0.5) == non_orthogonal, ( + f"{hm}: max|S·Sᵀ-I| = {err:.3f}, expected " + f"{'non-orthogonal' if non_orthogonal else 'orthogonal'} matrices" + ) + + +@pytest.mark.parametrize("hm, non_orthogonal", SPACEGROUPS) +def test_cartesian_symops_are_rotations(hm, non_orthogonal): + """``hkl_symops_to_cartesian`` must return genuine rotations. + + A Cartesian symmetry operator is orthogonal with determinant +1 by + definition. Conjugating with the untransposed ``S`` does not produce one in + a non-orthogonal basis -- it returned matrices with orthogonality error 5.33 + for P 3₁ 2 1 and P 6₅ 2 2, which is what let a wrong anisotropy + symmetrisation through undetected. + """ + del non_orthogonal + S = SpaceGroup(hm).matrices.detach().cpu().to(torch.float64) + B = _reciprocal_basis(_cell_for(hm)) + R = hkl_symops_to_cartesian(S, B).detach().cpu().to(torch.float64) + + eye = torch.eye(3, dtype=torch.float64) + for k in range(R.shape[0]): + # 1e-6 is float64 noise through the matrix inverse; the defect this + # guards against produced an error of 5.33. + assert torch.allclose(R[k] @ R[k].T, eye, atol=1e-6), ( + f"{hm} op {k} is not orthogonal:\n{R[k]}" + ) + assert float(torch.det(R[k])) == pytest.approx(1.0, abs=1e-6), ( + f"{hm} op {k} has determinant {float(torch.det(R[k])):.6f}, not +1" + ) + + +@pytest.mark.parametrize("hm, non_orthogonal", SPACEGROUPS) +def test_unroll_matches_expand_reciprocal(hm, non_orthogonal): + """Any inline symmetry expansion must agree with the shared helper. + + ``expand_reciprocal`` is the single definition of the convention, and it + returns the operator axis first -- so the einsum contraction + ``"kji,nj->kni"`` matches it directly. ``"kij,nj->kni"`` is the bug and + differs here for the non-orthogonal settings. + """ + del non_orthogonal + sg = SpaceGroup(hm) + S = sg.matrices.detach().cpu().to(torch.float64) + g = torch.Generator().manual_seed(7) + hkl = torch.randint(-12, 13, (200, 3), generator=g).to(torch.float64) + + reference = sg.expand_reciprocal(hkl).detach().cpu().to(torch.float64) # (ops, N, 3) + unrolled = torch.einsum("kji,nj->kni", S, hkl) # (ops, N, 3) + assert torch.allclose(unrolled, reference, atol=1e-9), ( + f"{hm}: inline unroll disagrees with SpaceGroup.expand_reciprocal" + ) + + +@pytest.mark.parametrize("hm, non_orthogonal", SPACEGROUPS) +def test_wrong_convention_changes_the_orbit(hm, non_orthogonal): + """``R·h`` yields a different *orbit* from ``h·R`` iff ``S`` is non-orthogonal. + + Per operator the two always differ unless ``S`` is symmetric, so the honest + statement is about the orbit as a set. For an orthogonal group ``Sᵀ = S⁻¹`` + is itself a group member, so the set is unchanged and the bug is invisible; + in a hexagonal basis ``Sᵀ`` leaves the group and the orbit genuinely moves. + + This is what gives the other tests their teeth: it asserts the defect is + observable at all for these settings, so a regression cannot hide behind a + benchmark built only from orthogonal lattices. + """ + S = SpaceGroup(hm).matrices.detach().cpu().to(torch.float64) + g = torch.Generator().manual_seed(11) + hkl = torch.randint(-12, 13, (200, 3), generator=g).to(torch.float64) + + def orbit_keys(contraction): + o = torch.einsum(contraction, S, hkl).round().to(torch.long) # (ops, N, 3) + k = ((o[..., 0] + 64) * 256 + (o[..., 1] + 64)) * 256 + (o[..., 2] + 64) + return torch.sort(k, dim=0).values # per-reflection set + + same_orbits = torch.equal(orbit_keys("kji,nj->kni"), orbit_keys("kij,nj->kni")) + assert (not same_orbits) == non_orthogonal, ( + f"{hm}: expected the orbits to " + f"{'differ' if non_orthogonal else 'coincide'} between conventions" + ) + + +@pytest.mark.parametrize("hm", ["P 31 2 1", "P 65 2 2"]) +def test_symmetrised_anisotropy_obeys_the_lattice(hm): + """With a 3- or 6-fold along c, Cartesian ``U`` must be ``diag(a, a, c)``. + + The previous conjugation returned a ``U`` *more* anisotropic than its input + (diag 0.90/1.30/0.60 became 1.58/3.70/0.60 with off-diagonals of 1.83), + which no averaging over a point group can do. + """ + S = SpaceGroup(hm).matrices.detach().cpu().to(torch.float64) + B = _reciprocal_basis(_cell_for(hm)) + R = hkl_symops_to_cartesian(S, B).detach().cpu().to(torch.float64) + + U = torch.tensor([[0.90, 0.10, 0.05], + [0.10, 1.30, -0.07], + [0.05, -0.07, 0.60]], dtype=torch.float64) + U_sym = symmetrize_anisotropy(U, R).detach().cpu().to(torch.float64) + + off = U_sym - torch.diag(U_sym.diagonal()) + assert float(off.abs().max()) < 1e-6, f"{hm}: U not diagonal:\n{U_sym}" + assert float(U_sym[0, 0]) == pytest.approx(float(U_sym[1, 1]), abs=1e-6), ( + f"{hm}: U11 != U22 ({U_sym[0,0]:.6f} vs {U_sym[1,1]:.6f})" + ) + # c is unconstrained by the axis, so it must survive untouched. + assert float(U_sym[2, 2]) == pytest.approx(0.60, abs=1e-6) + # An average cannot exceed the input's range. + assert float(U_sym[0, 0]) == pytest.approx(1.10, abs=1e-6) + + +@pytest.mark.parametrize("hm, non_orthogonal", SPACEGROUPS) +def test_epsilon_uses_the_row_vector_convention(hm, non_orthogonal): + """The multiplicity counts ops fixing ``h``, which needs ``h·R``. + + Reflections on a symmetry axis must come out with multiplicity > 1; with + the wrong convention the wrong reflections are flagged. + """ + del non_orthogonal + sg = SpaceGroup(hm) + S = sg.matrices.detach().cpu().to(torch.float64) + g = torch.Generator().manual_seed(3) + hkl = torch.randint(-9, 10, (400, 3), generator=g) + + eps = sg.epsilon(hkl, friedel=False).detach().cpu() + assert int(eps.min()) >= 1 + # Recompute independently through the shared helper. + ref = sg.expand_reciprocal(hkl).detach().cpu() # (ops, N, 3) + same = (ref == hkl.to(torch.int64)).all(dim=-1) + assert torch.equal(eps.to(torch.long), same.sum(dim=0).clamp(min=1)), ( + f"{hm}: epsilon(friedel=False) disagrees with expand_reciprocal" + ) + + +@pytest.mark.parametrize("hm, non_orthogonal", SPACEGROUPS) +def test_symmetry_unroll_stays_within_the_true_orbit(hm, non_orthogonal): + """Every emitted position must be a genuine symmetry mate of its input. + + This exercises a real call site rather than the contraction in isolation. + Under the wrong convention the emitted positions leave the true orbit for a + non-orthogonal setting, which is what let two inequivalent reflections land + on one Miller index carrying different ``|F|``. + + Against ``SpaceGroup.expand_hkl``, the shared helper. The alignment package + had its own ``epsilon_aware_unroll`` doing the same orbit dedup -- measured + to emit ``n_ops/epsilon(h)`` distinct mates, exactly as this does -- and it + was deleted once nothing in production called it. Guarding the shared one is + worth more than guarding the copy was. + """ + del non_orthogonal + sg = SpaceGroup(hm) + g = torch.Generator().manual_seed(19) + hkl = torch.randint(-9, 10, (150, 3), generator=g) + + unrolled, asu_idx, _ = sg.expand_hkl(hkl, include_friedel=False) + unrolled = unrolled.detach().cpu().to(torch.long) + asu_idx = asu_idx.detach().cpu().to(torch.long) + + # true orbit of each parent, via the shared helper + orbit = sg.expand_reciprocal(hkl).detach().cpu().to(torch.long) # (ops, N, 3) + for i in range(0, unrolled.shape[0], 17): # stride: full sweep is redundant + parent = int(asu_idx[i]) + mates = orbit[:, parent] # (ops, 3) + assert (mates == unrolled[i]).all(dim=-1).any(), ( + f"{hm}: emitted {unrolled[i].tolist()} is not in the orbit of " + f"{hkl[parent].tolist()}" + ) + + +def test_no_column_convention_survives_in_the_alignment_package(): + """No ``S·h`` symmetry contraction may reappear in the alignment package. + + The four defects this module guards were four *copies* of one rule. The + other tests pin the rule; this one pins the absence of new copies, which is + the failure mode that actually occurred -- the shared helper was corrected + while the inline duplicates were not. + """ + import re + + pkg = Path(__file__).resolve().parents[3] / "torchref" / "experimental" / "alignment" + # `kij` contracted against an hkl-like index is the column convention. + pattern = re.compile(r'einsum\(\s*["\']k?ij,\s*nj->') + offenders = [] + for path in sorted(pkg.rglob("*.py")): + for lineno, line in enumerate(path.read_text().splitlines(), 1): + if pattern.search(line): + offenders.append(f"{path.relative_to(pkg.parent.parent.parent)}:{lineno}: {line.strip()}") + assert not offenders, ( + "reciprocal-space symmetry expansion must use h·R (\"kji,nj->\"), not " + "R·h (\"kij,nj->\"):\n " + "\n ".join(offenders) + ) diff --git a/tests/unit/alignment/test_translation_obs.py b/tests/unit/alignment/test_translation_obs.py new file mode 100644 index 00000000..5eb1b20f --- /dev/null +++ b/tests/unit/alignment/test_translation_obs.py @@ -0,0 +1,188 @@ +"""The observed side of the translation search, and what it guarantees. + +``TranslationObs`` exists so that "what is E_obs here" has exactly one answer +across the whole run. Three separate answers used to live in this path -- one in +the coarse search, one in the likelihood, one hand-rolled in the pipeline -- and +they disagreed, which meant the stage that chose a peak and the stage that +re-ranked it were not scoring the same quantity. + +These tests pin the properties that claim rests on: the normalisation is the +shared Wilson fit, the weight is a real per-reflection weight rather than a +per-shell one (per-shell weights cancel in a correlation), and the whole object +is invariant to the units the amplitudes arrive in. +""" + +import pytest +import torch + +from torchref.experimental.alignment.translation import TranslationObs +from torchref.scaling import WilsonNormaliser +from torchref.symmetry.cell import Cell +from torchref.symmetry.spacegroup import SpaceGroup + + +def _case(n=4000, seed=0, sg="P 21 21 21"): + """A synthetic reflection set with a realistic Wilson falloff.""" + g = torch.Generator().manual_seed(seed) + # Pinned to CPU: this host has an accelerator, and Cell/SpaceGroup would + # land there while the synthetic hkl below stays on the host. + cell = Cell([61.0, 72.0, 83.0, 90.0, 90.0, 90.0], device="cpu") + spacegroup = SpaceGroup(sg, device="cpu") + # Miller indices on a coarse block, origin removed. + rng = torch.arange(-9, 10) + h, k, l = torch.meshgrid(rng, rng, rng, indexing="ij") + hkl = torch.stack([h.flatten(), k.flatten(), l.flatten()], dim=-1) + hkl = hkl[(hkl.abs().sum(dim=-1) > 0)] + hkl = hkl[torch.randperm(hkl.shape[0], generator=g)[:n]] + + s_mag = (hkl.to(torch.float64) + @ cell.reciprocal_basis_matrix.to(torch.float64)).norm(dim=-1) + # Exponential intensities on a Wilson curve, so = 1 is reachable. + Sigma = 3000.0 * torch.exp(-2.0 * 25.0 * s_mag ** 2) + I = Sigma * -torch.rand(n, generator=g, dtype=torch.float64).clamp(min=1e-9).log() + F = I.sqrt() + sig_F = 0.05 * F + 0.01 * F.mean() + return F, sig_F, hkl, spacegroup, cell, s_mag + + +@pytest.mark.unit +def test_normalisation_is_the_shared_wilson_fit(): + """E_obs is WilsonNormaliser's E, not a private per-shell mean.""" + F, _, hkl, sg, cell, _ = _case() + obs = TranslationObs.build(F, hkl, sg, cell) + + direct = WilsonNormaliser( + obs.F_obs * obs.F_obs, obs.s_mag, eps=obs.eps, centric=obs.centric, + n_coeff=6, + ) + torch.testing.assert_close(obs.E_obs, direct.E.to(obs.E_obs.dtype)) + + +@pytest.mark.unit +def test_mean_e_squared_is_one(): + """ = 1 is an identity of the Gamma fit, so it holds to fit precision. + + k-weighted, because that is the score equation the intercept solves: + sum_h k_h (I_h/mu_h - 1) = 0 with k = 1 acentric, 1/2 centric. + """ + F, _, hkl, sg, cell, _ = _case() + obs = TranslationObs.build(F, hkl, sg, cell) + # Reduced on the host in double: the identity is asserted to 1e-6 over 4000 + # reflections and the accumulation should not be what limits that -- but the + # widening has to happen after the readback, since a backend without float64 + # cannot hold the wide copy. + k = torch.where(obs.centric.cpu(), 0.5, 1.0).to(torch.float64) + E2 = obs.E_obs.detach().cpu().to(torch.float64) ** 2 + mean_e2 = (k * E2).sum() / k.sum() + assert abs(float(mean_e2) - 1.0) < 1e-6, mean_e2 + + +@pytest.mark.unit +def test_e_obs_is_invariant_to_the_amplitude_scale(): + """Rescaling every amplitude must not change E. It is an ABSOLUTE normaliser. + + Exact in the model -- a common factor lands entirely in Sigma's intercept -- + but the fit is IRLS in the configured float dtype and stops at a relative + tolerance, so the bar is that tolerance rather than machine epsilon. + Measured 1.2e-4 worst case over 4000 reflections at a 7.5x rescale; 1e-3 + catches a genuine scale dependence without chasing the solver. + """ + F, sig_F, hkl, sg, cell, _ = _case() + base = TranslationObs.build(F, hkl, sg, cell, sig_F=sig_F) + scaled = TranslationObs.build(7.5 * F, hkl, sg, cell, sig_F=7.5 * sig_F) + torch.testing.assert_close(base.E_obs, scaled.E_obs, rtol=1e-3, atol=1e-3) + # F/sigma is unchanged by a common factor, so the weight must be too. + torch.testing.assert_close(base.weight, scaled.weight, rtol=1e-3, atol=1e-3) + + +@pytest.mark.unit +def test_weight_varies_within_a_shell(): + """The part of the weight that is not gauge is the part that varies within a shell. + + A weight constant inside a resolution shell is a per-shell weight, and a + correlation absorbs those -- which is the whole reason twelve E conventions + moved the rotation function's truth rank by nothing. So this asserts the + thing that makes weighting worth doing at all, not merely that a weight + exists. + """ + F, sig_F, hkl, sg, cell, _ = _case() + # Give two reflections at the SAME resolution very different sigmas. + obs = TranslationObs.build(F, hkl, sg, cell, sig_F=sig_F) + + from torchref.experimental.alignment.sh import (assign_shells, + equal_count_shell_edges) + + edges, _ = equal_count_shell_edges(obs.s_mag, 20) + shell_idx = assign_shells(obs.s_mag, edges).clamp(min=0) + within = [] + for shell in range(20): + w = obs.weight[shell_idx == shell] + if w.numel() > 20: + within.append(float(w.std() / w.mean().clamp(min=1e-30))) + assert within, "no populated shells" + assert min(within) > 1e-3, ( + f"weight is effectively constant within shells (max rel. spread " + f"{max(within):.2e}); it would be absorbed by the correlation" + ) + + +@pytest.mark.unit +def test_weight_is_uniform_without_sigmas(): + """No sigmas, no weight. The varying half of the weight IS the measurement term.""" + F, _, hkl, sg, cell, _ = _case() + obs = TranslationObs.build(F, hkl, sg, cell, sig_F=None) + assert torch.allclose(obs.weight, torch.ones_like(obs.weight)) + + +@pytest.mark.unit +def test_weight_is_normalised_to_mean_one(): + """So the score's scale does not depend on how the sigmas happened to be scaled. + + To a few ulps of the working dtype, not to a fixed 1e-9: the normalisation + divides by this very mean, so what is left is the rounding of a 4000-term + reduction, and at float32 one ulp is already 1.2e-7. The old fixed bound + passed or failed on the reduction order alone -- it held on MPS and missed + by exactly one ulp on the CPU. + """ + F, sig_F, hkl, sg, cell, _ = _case() + obs = TranslationObs.build(F, hkl, sg, cell, sig_F=sig_F) + tol = 8 * torch.finfo(obs.weight.dtype).eps + assert abs(float(obs.weight.mean()) - 1.0) < tol + + +@pytest.mark.unit +@pytest.mark.parametrize("sg_name", ["P 1", "P 21 21 21", "C 1 2 1", "P 43 21 2"]) +def test_epsilon_and_centricity_come_from_the_spacegroup(sg_name): + """Both reach the fit, and epsilon is the friedel=False count. + + Wilson's = eps*Sigma counts the operations mapping h to itself, which add + coherently and set the mean. The Friedel-folded count answers a different + question -- it describes the distribution's shape, which enters as the Gamma + shape via centricity, separately. Applying the wrong one here shifts the + normalisation of every axial reflection. + """ + F, _, hkl, sg, cell, _ = _case(sg=sg_name) + obs = TranslationObs.build(F, hkl, sg, cell) + hkl_l = obs.hkl.round().to(torch.int64) + torch.testing.assert_close( + obs.eps, sg.epsilon(hkl_l, friedel=False).to(obs.eps.dtype).clamp(min=1.0), + ) + torch.testing.assert_close(obs.centric, sg.is_centric(hkl_l).to(torch.bool)) + + +@pytest.mark.unit +def test_coefficient_is_the_rotation_functions_score_equation(): + """The fast search's coefficient is LERF1's intensity times sigma_A^2. + + One score equation for both searches: ``cw (E^2 - 1) w`` is what the + rotation function expands, and ``sigma_A^2`` is its calc-side weight. + """ + from torchref.experimental.alignment.frf.preprocessing import ( + build_lerf1_intensity, eterm_sigma_a) + + F, sig_F, hkl, sg, cell, _ = _case() + obs = TranslationObs.build(F, hkl, sg, cell, sig_F=sig_F, delta_vrms_A=0.8) + expected = (build_lerf1_intensity(obs.E_obs, obs.centric, weight=obs.weight) + * eterm_sigma_a(obs.s_mag, 0.8) ** 2) + torch.testing.assert_close(obs.coeff, expected) + torch.testing.assert_close(obs.sigma_a, eterm_sigma_a(obs.s_mag, 0.8)) diff --git a/tests/unit/alignment/test_wigner_d_cache.py b/tests/unit/alignment/test_wigner_d_cache.py new file mode 100644 index 00000000..f9890b1a --- /dev/null +++ b/tests/unit/alignment/test_wigner_d_cache.py @@ -0,0 +1,105 @@ +"""The memoised small-d blocks must not change what the contraction returns. + +The blocks ``d^l(β)`` depend only on the bandwidth and the β grid, never on the +data, so hoisting them out of the per-call loop is pure reuse. That makes the +cache a correctness risk in exactly one way: a stale entry served for the wrong +``(L, betas)``. These tests pin the key, the identity of the result, and the +one-entry footprint bound. +""" + +import math + +import pytest +import torch + +from torchref.experimental.alignment.frf.wigner_d import ( + _WIGNER_D_CACHE, + _wigner_d_blocks, + clear_wigner_d_cache, + wigner_contraction_per_beta, +) + + +def _betas(n, step_deg=3.0): + return torch.arange(n, dtype=torch.float64) * step_deg * (math.pi / 180.0) + + +def _xi(L, seed=0): + g = torch.Generator().manual_seed(seed) + dim = 2 * L - 1 + return torch.randn(L, dim, dim, generator=g, dtype=torch.float64).to( + torch.complex128 + ) + + +@pytest.mark.unit +def test_the_cached_call_is_bit_identical(): + """A cache hit must give exactly the first call's answer, not merely close.""" + clear_wigner_d_cache() + L, betas = 9, _betas(12) + xi = _xi(L) + first = wigner_contraction_per_beta(xi, betas) + second = wigner_contraction_per_beta(xi, betas) + assert torch.equal(first, second) + + +@pytest.mark.unit +def test_different_data_at_the_same_bandwidth_still_differs(): + """Guard the premise: the cache holds β geometry, not the data. + + Without this, a cache keyed too loosely -- or one that memoised the whole + result -- would pass the identity test above by returning a stale answer. + """ + clear_wigner_d_cache() + L, betas = 9, _betas(12) + a = wigner_contraction_per_beta(_xi(L, seed=0), betas) + b = wigner_contraction_per_beta(_xi(L, seed=1), betas) + assert not torch.allclose(a, b) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "L2, n_beta2", [(9, 15), (11, 12)], ids=["other-beta-grid", "other-bandwidth"] +) +def test_a_different_key_is_not_served_the_cached_blocks(L2, n_beta2): + """Changing either the bandwidth or the β grid must rebuild.""" + clear_wigner_d_cache() + ref_blocks = _wigner_d_blocks(9, _betas(12), torch.device("cpu"), torch.float64) + got = _wigner_d_blocks(L2, _betas(n_beta2), torch.device("cpu"), torch.float64) + assert len(got) == L2 - 1 + assert got[0].shape[0] == n_beta2 + assert got is not ref_blocks + + +@pytest.mark.unit +def test_the_cache_holds_one_entry(): + """The blocks are hundreds of MB at production bandwidths, so they are not + allowed to accumulate across keys.""" + clear_wigner_d_cache() + for L in (7, 9, 11): + _wigner_d_blocks(L, _betas(12), torch.device("cpu"), torch.float64) + assert len(_WIGNER_D_CACHE) == 1 + clear_wigner_d_cache() + assert len(_WIGNER_D_CACHE) == 0 + + +@pytest.mark.unit +def test_the_blocks_are_the_wigner_small_d_matrices(): + """Anchor the cached quantity against d^l(β) computed from the definition. + + ``d^l(0) = I`` and ``d^l(β)`` is orthogonal for every β; both follow from + ``d^l(β) = exp(-i β J_y)`` and neither holds for a mis-shaped or + mis-transposed block. + """ + clear_wigner_d_cache() + L = 7 + betas = torch.tensor([0.0, 0.4, 1.7, 3.0], dtype=torch.float64) + blocks = _wigner_d_blocks(L, betas, torch.device("cpu"), torch.float64) + for l, d in enumerate(blocks, start=1): + sz = 2 * l + 1 + assert d.shape == (betas.numel(), sz, sz) + torch.testing.assert_close(d[0], torch.eye(sz, dtype=d.dtype)) + for k in range(betas.numel()): + torch.testing.assert_close( + d[k] @ d[k].transpose(-1, -2), torch.eye(sz, dtype=d.dtype) + ) diff --git a/tests/unit/base/test_canonical_sphere_cpu.py b/tests/unit/base/test_canonical_sphere_cpu.py index 943baac1..1522bb41 100644 --- a/tests/unit/base/test_canonical_sphere_cpu.py +++ b/tests/unit/base/test_canonical_sphere_cpu.py @@ -74,12 +74,6 @@ def _cell(beta_deg, dtype=torch.float32, dims=(48, 40, 34), abc=(28.0, 24.0, 20. return f64.to(dtype), torch.linalg.inv(f64).to(dtype), dims, f64 -def _voxel_size(f64, dims): - """``voxel_size`` as ``sf_fft`` derives it; unused by the kernels, still in the - ``build_electron_density`` signature.""" - return (f64.norm(dim=0) / torch.tensor(dims, dtype=torch.float64)).float() - - def _iso_atoms(f64, n=36, dtype=torch.float32, seed=0): g = torch.Generator().manual_seed(seed) z = torch.tensor([6, 7, 8, 16]).repeat(n // 4 + 1)[:n] @@ -242,8 +236,7 @@ def test_fused_float64_is_exact(): # AUTO vs EAGER through the real dispatch: no accelerator needed # =========================================================================== -def _build(pin, dims, frac, inv_frac, voxel, dtype, iso=None, aniso=None): - rsg = torch.zeros(*dims, 3, dtype=dtype) # shape only; no kernel reads its values +def _build(pin, dims, frac, inv_frac, dtype, iso=None, aniso=None): xi, ai, oi, Ai, Bi = iso if iso is not None else _empty_iso(dtype) kw = {} if aniso is not None: @@ -251,7 +244,8 @@ def _build(pin, dims, frac, inv_frac, voxel, dtype, iso=None, aniso=None): kw = dict(xyz_aniso=xa, u_aniso=ua, occ_aniso=oa, A_aniso=Aa, B_aniso=Ba) with (use_portable() if pin else contextlib.nullcontext()): return build_electron_density( - rsg, xi, ai, oi, Ai, Bi, inv_frac, frac, voxel, dtype=dtype, **kw) + dims, torch.device("cpu"), xi, ai, oi, Ai, Bi, inv_frac, frac, + dtype=dtype, **kw) @pytest.mark.parametrize("beta", _BETAS) @@ -259,10 +253,9 @@ def _build(pin, dims, frac, inv_frac, voxel, dtype, iso=None, aniso=None): ids=["float32", "float64"]) def test_auto_matches_eager_iso(beta, dtype): frac, inv_frac, dims, f64 = _cell(beta, dtype=dtype) - voxel = _voxel_size(f64, dims) atoms = _iso_atoms(f64, dtype=dtype) - ref = _build(True, dims, frac, inv_frac, voxel, dtype, iso=atoms) - got = _build(False, dims, frac, inv_frac, voxel, dtype, iso=atoms) + ref = _build(True, dims, frac, inv_frac, dtype, iso=atoms) + got = _build(False, dims, frac, inv_frac, dtype, iso=atoms) tol = _F32_TOL if dtype is torch.float32 else 1e-12 assert _rel_l2(got, ref) < tol @@ -270,10 +263,9 @@ def test_auto_matches_eager_iso(beta, dtype): @pytest.mark.parametrize("beta", _BETAS) def test_auto_matches_eager_aniso(beta): frac, inv_frac, dims, f64 = _cell(beta) - voxel = _voxel_size(f64, dims) atoms = _aniso_atoms(f64) - ref = _build(True, dims, frac, inv_frac, voxel, torch.float32, aniso=atoms) - got = _build(False, dims, frac, inv_frac, voxel, torch.float32, aniso=atoms) + ref = _build(True, dims, frac, inv_frac, torch.float32, aniso=atoms) + got = _build(False, dims, frac, inv_frac, torch.float32, aniso=atoms) assert _rel_l2(got, ref) < _F32_TOL @@ -282,7 +274,6 @@ def test_auto_matches_eager_gradients(kind): """Direction *and* magnitude: a kernel returning ``2 * grad`` is perfectly parallel, so cosine alone cannot catch it.""" frac, inv_frac, dims, f64 = _cell(100.0, dtype=torch.float64, dims=(32, 28, 24)) - voxel = _voxel_size(f64, dims) w = torch.randn(dims, generator=torch.Generator().manual_seed(7), dtype=torch.float64) if kind == "iso": @@ -293,7 +284,7 @@ def test_auto_matches_eager_gradients(kind): def grads(pin): x, pp, o = (t.clone().requires_grad_() for t in (xyz, p, occ)) pack = (x, pp, o, A, B) - dm = _build(pin, dims, frac, inv_frac, voxel, torch.float64, + dm = _build(pin, dims, frac, inv_frac, torch.float64, **({"iso": pack} if kind == "iso" else {"aniso": pack})) (dm * w).sum().backward() return x.grad, pp.grad, o.grad @@ -350,7 +341,7 @@ def recording(*args, **kwargs): return real(*args, **kwargs) monkeypatch.setattr(sphere_splat, "add_isotropic_cpu_sphere_var", recording) - _build(False, dims, frac, inv_frac, _voxel_size(f64, dims), dtype, + _build(False, dims, frac, inv_frac, dtype, iso=_iso_atoms(f64, dtype=dtype)) assert calls, f"the default did not reach the fused CPU splat for {dtype}" @@ -358,22 +349,20 @@ def recording(*args, **kwargs): def test_empty_atom_sets(): """A structure with no isotropic (or no anisotropic) atoms must not crash.""" frac, inv_frac, dims, f64 = _cell(90.0) - voxel = _voxel_size(f64, dims) - only_aniso = _build(False, dims, frac, inv_frac, voxel, torch.float32, + only_aniso = _build(False, dims, frac, inv_frac, torch.float32, aniso=_aniso_atoms(f64)) assert torch.isfinite(only_aniso).all() and float(only_aniso.abs().sum()) > 0 - both_empty = _build(False, dims, frac, inv_frac, voxel, torch.float32) + both_empty = _build(False, dims, frac, inv_frac, torch.float32) assert float(both_empty.abs().sum()) == 0.0 def test_density_map_accumulates_not_overwrites(): """Both passes add into one map, so the aniso pass must not clobber the iso one.""" frac, inv_frac, dims, f64 = _cell(90.0) - voxel = _voxel_size(f64, dims) iso, aniso = _iso_atoms(f64), _aniso_atoms(f64) - a = _build(False, dims, frac, inv_frac, voxel, torch.float32, iso=iso) - b = _build(False, dims, frac, inv_frac, voxel, torch.float32, aniso=aniso) - both = _build(False, dims, frac, inv_frac, voxel, torch.float32, + a = _build(False, dims, frac, inv_frac, torch.float32, iso=iso) + b = _build(False, dims, frac, inv_frac, torch.float32, aniso=aniso) + both = _build(False, dims, frac, inv_frac, torch.float32, iso=iso, aniso=aniso) assert _rel_l2(both, a + b) < 1e-6 @@ -407,20 +396,19 @@ class of index and factor errors that cross-backend parity cannot see because bo """ frac, inv_frac, dims, f64 = _cell(beta, dtype=torch.float64) xyz, adp, occ, A, B = _iso_atoms(f64, n=24, dtype=torch.float64) - voxel = _voxel_size(f64, dims) - grid = torch.zeros(*dims, 3, dtype=torch.float64) u_sph = torch.zeros(xyz.shape[0], 6, dtype=torch.float64) u_sph[:, :3] = (adp / (8.0 * math.pi**2)).unsqueeze(1) with (use_portable() if pin else contextlib.nullcontext()): iso_map = build_electron_density( - grid, xyz, adp, occ, A, B, inv_frac, frac, voxel, dtype=torch.float64 + dims, torch.device("cpu"), xyz, adp, occ, A, B, inv_frac, frac, + dtype=torch.float64 ) aniso_map = build_electron_density( - grid, + dims, torch.device("cpu"), xyz[:0], adp[:0], occ[:0], A[:0], B[:0], - inv_frac, frac, voxel, + inv_frac, frac, xyz_aniso=xyz, u_aniso=u_sph, occ_aniso=occ, A_aniso=A, B_aniso=B, dtype=torch.float64, ) @@ -465,20 +453,20 @@ def test_fused_kernel_is_thread_invariant(n_threads): frac, inv_frac, dims, f64 = _cell(115.0, dtype=torch.float32) xyz, adp, occ, A, B = _iso_atoms(f64, n=96, dtype=torch.float32, seed=7) - voxel = _voxel_size(f64, dims) - grid = torch.zeros(*dims, 3, dtype=torch.float32) original = torch.get_num_threads() try: torch.set_num_threads(1) with contextlib.nullcontext(): ref = build_electron_density( - grid, xyz, adp, occ, A, B, inv_frac, frac, voxel, dtype=torch.float32 + dims, torch.device("cpu"), xyz, adp, occ, A, B, inv_frac, frac, + dtype=torch.float32 ) torch.set_num_threads(n_threads) with contextlib.nullcontext(): got = build_electron_density( - grid, xyz, adp, occ, A, B, inv_frac, frac, voxel, dtype=torch.float32 + dims, torch.device("cpu"), xyz, adp, occ, A, B, inv_frac, frac, + dtype=torch.float32 ) finally: torch.set_num_threads(original) diff --git a/tests/unit/base/test_intensity_likelihoods.py b/tests/unit/base/test_intensity_likelihoods.py new file mode 100644 index 00000000..df1353cd --- /dev/null +++ b/tests/unit/base/test_intensity_likelihoods.py @@ -0,0 +1,96 @@ +"""Intensity targets read measured intensities and propagate their uncertainties.""" + +import math + +import pytest +import torch + +from torchref.base.targets.xray_likelihoods import ( + VAR_FLOOR, + amplitude_var_from_sigma_obs, + floor_sigma_obs, + gaussian_per_refl, + intensity_var_from_sigma_obs, + nll_per_refl, +) +from torchref.config import get_default_device, get_float_dtype + + +@pytest.mark.unit +def test_the_shared_gaussian_reproduces_the_amplitude_one_bitwise(): + """``nll_per_refl`` is ``gaussian_per_refl`` on ``|F_calc|``, exactly.""" + options = dict(dtype=get_float_dtype(), device=get_default_device()) + obs = torch.linspace(0.1, 100.0, 5000, **options) + calc = torch.linspace(-100.0, 100.0, 5000, **options) + var = amplitude_var_from_sigma_obs(torch.linspace(0.1, 10.0, 5000, **options)) + assert torch.equal( + nll_per_refl(obs, calc, var), gaussian_per_refl(obs, calc.abs(), var) + ) + + +@pytest.mark.unit +def test_the_absolute_variance_floor_is_opt_out_and_matters(rtol): + """``VAR_FLOOR`` is a distortion, not a safeguard, once the builder has floored + sigma.""" + sigma = torch.full( + (256,), 1e-3, dtype=get_float_dtype(), device=get_default_device() + ) + var = intensity_var_from_sigma_obs(sigma) # 1e-6, comfortably above VAR_FLOOR + obs = torch.zeros(256, dtype=get_float_dtype(), device=get_default_device()) + model = torch.full( + (256,), 1e-4, dtype=get_float_dtype(), device=get_default_device() + ) + assert torch.equal( + gaussian_per_refl(obs, model, var, var_floor=0.0), + gaussian_per_refl(obs, model, var, var_floor=VAR_FLOOR), + ), "the floor must be inert when the variance is above it" + + # Below it, the two differ -- and by a lot, not by an ulp. + tiny = torch.full( + (256,), 1e-6, dtype=get_float_dtype(), device=get_default_device() + ) # var = 1e-12 << VAR_FLOOR + var_tiny = intensity_var_from_sigma_obs(tiny) + free = gaussian_per_refl(obs, model, var_tiny, var_floor=0.0) + clamped = gaussian_per_refl(obs, model, var_tiny, var_floor=VAR_FLOOR) + assert not torch.allclose(free, clamped) + # `free` is the honest one: it uses the variance the builder actually produced. + expected = ( + 0.5 * (1e-4) ** 2 / 1e-12 + 0.5 * math.log(1e-12) + 0.5 * math.log(2 * math.pi) + ) + assert free[0].item() == pytest.approx(expected, rel=rtol) + + +@pytest.mark.unit +def test_the_intensity_sigma_floor_respects_the_fitted_subset(): + """``mask`` restricts the median, because unfitted rows carry filler.""" + # The first 20 fitted rows are BELOW the fitted median's floor, so the floor is what + # they come back as -- which is the only way to observe which median was used. + sigma = torch.cat( + [ + torch.full( + (20,), 0.01, device=get_default_device(), dtype=get_float_dtype() + ), # fitted, and below floor either way + torch.full( + (80,), 10.0, device=get_default_device(), dtype=get_float_dtype() + ), # fitted, sets the fitted median + torch.full( + (900,), 1e6, device=get_default_device(), dtype=get_float_dtype() + ), # NOT fitted: filler + ] + ) + mask = torch.cat( + [ + torch.ones(100, dtype=torch.bool, device=get_default_device()), + torch.zeros(900, dtype=torch.bool, device=get_default_device()), + ] + ) + masked = floor_sigma_obs(sigma, mask, abs_floor=1e-12) + unmasked = floor_sigma_obs(sigma, None, abs_floor=1e-12) + assert masked[:20].min().item() == pytest.approx(1.0) # floor = 10 * 0.1 + assert unmasked[:20].min().item() == pytest.approx( + 1e5 + ) # floor = 1e6 * 0.1, swamped + # An explicit floor overrides the median entirely -- the set-independent path. + assert floor_sigma_obs(sigma, mask, floor=0.5)[:20].min().item() == pytest.approx( + 0.5 + ) diff --git a/tests/unit/base/test_local_frame.py b/tests/unit/base/test_local_frame.py new file mode 100644 index 00000000..32e28e84 --- /dev/null +++ b/tests/unit/base/test_local_frame.py @@ -0,0 +1,94 @@ +"""Local-frame placement: exact inverse, exact gradients, and a safe fallback. + +A riding hydrogen is stored as an offset in a frame built from three atoms. What has +to hold is that the offset round-trips through placement exactly, that a force on the +placed point reaches only the three frame atoms, and that a frame the geometry cannot +define is recognised rather than used. +""" + +import pytest +import torch + +from torchref.base.coordinates.local_frame import ( + frame_is_degenerate, + local_frame_coordinates, + place_local_frame, +) + + +def _frames(n: int, dtype=torch.float64): + generator = torch.Generator().manual_seed(7) + p = torch.rand(n, 3, generator=generator, dtype=dtype) * 10 + n1 = p + torch.randn(n, 3, generator=generator, dtype=dtype) + n2 = p + torch.randn(n, 3, generator=generator, dtype=dtype) + point = p + torch.randn(n, 3, generator=generator, dtype=dtype) + return p, n1, n2, point + + +@pytest.mark.unit +def test_local_coordinates_invert_placement(): + """A point expressed in its frame and placed again lands where it started.""" + p, n1, n2, point = _frames(64) + valid = torch.ones(64, dtype=torch.bool) + local = local_frame_coordinates(p, n1, n2, point) + back = place_local_frame(p, n1, n2, local, valid, point - p) + assert torch.allclose(back, point, atol=1e-12) + + +@pytest.mark.unit +def test_offset_is_invariant_under_rigid_motion(): + """Rotating and translating the three frame atoms carries the point along.""" + p, n1, n2, point = _frames(16) + local = local_frame_coordinates(p, n1, n2, point) + angle = torch.tensor(0.7, dtype=torch.float64) + rotation = torch.tensor( + [ + [torch.cos(angle), -torch.sin(angle), 0.0], + [torch.sin(angle), torch.cos(angle), 0.0], + [0.0, 0.0, 1.0], + ], + dtype=torch.float64, + ) + shift = torch.tensor([1.0, -2.0, 3.0], dtype=torch.float64) + moved = [x @ rotation.T + shift for x in (p, n1, n2, point)] + valid = torch.ones(16, dtype=torch.bool) + placed = place_local_frame(moved[0], moved[1], moved[2], local, valid, point - p) + assert torch.allclose(placed, moved[3], atol=1e-12) + + +@pytest.mark.unit +def test_gradients_are_exact_and_reach_only_the_frame_atoms(): + """Autograd through the frame matches finite differences.""" + p, n1, n2, point = _frames(6) + local = local_frame_coordinates(p, n1, n2, point) + valid = torch.ones(6, dtype=torch.bool) + rigid = point - p + + def place(pp, a, b): + return place_local_frame(pp, a, b, local, valid, rigid) + + leaves = tuple(x.clone().requires_grad_() for x in (p, n1, n2)) + assert torch.autograd.gradcheck(place, leaves, eps=1e-6, atol=1e-6) + + +@pytest.mark.unit +def test_invalid_frames_fall_back_to_rigid_translation(): + """Where the frame is flagged invalid the point simply follows its parent.""" + p, n1, n2, point = _frames(8) + local = torch.zeros(8, 3, dtype=torch.float64) + valid = torch.zeros(8, dtype=torch.bool) + rigid = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float64).expand(8, 3) + placed = place_local_frame(p, n1, n2, local, valid, rigid) + assert torch.allclose(placed, p + rigid) + + +@pytest.mark.unit +def test_degenerate_frames_are_detected(): + """Collinear or collapsed reference bonds are flagged, healthy frames are not.""" + p = torch.zeros(3, 3, dtype=torch.float64) + n1 = torch.tensor([[1.0, 0.0, 0.0]] * 3, dtype=torch.float64) + n2 = torch.tensor( + [[0.0, 1.0, 0.0], [2.0, 0.0, 0.0], [1e-5, 0.0, 0.0]], dtype=torch.float64 + ) + flagged = frame_is_degenerate(p, n1, n2) + assert flagged.tolist() == [False, True, True] diff --git a/tests/unit/base/test_loss.py b/tests/unit/base/test_loss.py new file mode 100644 index 00000000..eff9ffcb --- /dev/null +++ b/tests/unit/base/test_loss.py @@ -0,0 +1,36 @@ +"""Pin the amplitude-metric Gaussian likelihood's value and reduction contract.""" + +import math + +import pytest +import torch + +from torchref.base.metrics.loss import nll_xray, nll_xray_mean, nll_xray_sum +from torchref.config import get_default_device, get_float_dtype + +pytestmark = pytest.mark.unit + + +def test_gaussian_nll_value_and_reduction() -> None: + """The NLL includes its normalization and sums over reflections.""" + obs = torch.tensor( + [10.0, 20.0, 30.0], dtype=get_float_dtype(), device=get_default_device() + ) + sigma = obs.new_tensor([1.0, 2.0, 4.0]) + calc = obs + sigma + expected = obs.new_tensor(1.5 + math.log(8.0) + 1.5 * math.log(2.0 * math.pi)) + + torch.testing.assert_close(nll_xray(obs, calc, sigma), expected) + torch.testing.assert_close(nll_xray_sum(obs, calc, sigma), expected) + torch.testing.assert_close(nll_xray_mean(obs, calc, sigma), expected / obs.numel()) + + +def test_gaussian_nll_penalizes_amplitude_error() -> None: + """A one-sigma residual adds one half per reflection to the perfect-fit NLL.""" + obs = torch.tensor( + [10.0, 20.0, 30.0], dtype=get_float_dtype(), device=get_default_device() + ) + sigma = torch.ones_like(obs) + good = nll_xray(obs, obs, sigma) + bad = nll_xray(obs, obs + sigma, sigma) + torch.testing.assert_close(bad - good, obs.new_tensor(1.5)) diff --git a/tests/unit/base/test_target_values.py b/tests/unit/base/test_target_values.py new file mode 100644 index 00000000..bf27f0a6 --- /dev/null +++ b/tests/unit/base/test_target_values.py @@ -0,0 +1,127 @@ +"""Compare restraint kernels with host references on deposited Cartesian coordinates.""" + +import math + +import numpy as np +import pytest +import torch + +from torchref.base.targets._common import EPS +from torchref.base.targets.adp import adp_simu_math +from torchref.base.targets.angle import angle_math +from torchref.base.targets.bond import bond_math +from torchref.base.targets.chiral import chiral_math +from torchref.base.targets.planarity import planarity_math +from torchref.base.targets.xray_ls import ls_xray_loss_math +from torchref.config import get_default_device, get_float_dtype, get_int_dtype + +pytestmark = pytest.mark.unit + + +@pytest.fixture(scope="module") +def deposited_atoms(sample_cif_file): + """Return detached Cartesian coordinates (Å) and isotropic B-factors (Ų).""" + from torchref.model import Model + + model = Model(verbose=0) + model.load_cif(str(sample_cif_file)) + return model.xyz().detach().clone(), model.adp().detach().clone() + + +def _indices(rows, device): + return torch.tensor(rows, dtype=get_int_dtype(), device=device) + + +def _gaussian_sum(residual, sigma): + return np.sum( + 0.5 * (residual / sigma) ** 2 + np.log(sigma) + 0.5 * math.log(2 * math.pi) + ) + + +def test_bond_value(deposited_atoms) -> None: + """Bond lengths enter a summed Gaussian NLL in Å.""" + xyz, _ = deposited_atoms + host = xyz[:4].cpu().numpy().astype(np.float64) + idx = _indices([[0, 1], [2, 3]], xyz.device) + refs = xyz.new_tensor([1.4, 1.5]) + sigma = xyz.new_tensor([0.1, 0.2]) + # The kernel regularizes squared distance to keep coincident-atom gradients finite. + distance = np.sqrt(np.sum((host[[0, 2]] - host[[1, 3]]) ** 2, axis=1) + EPS) + expected = _gaussian_sum(distance - refs.cpu().numpy(), sigma.cpu().numpy()) + torch.testing.assert_close( + bond_math(xyz, idx, refs, sigma), xyz.new_tensor(expected) + ) + + +def test_angle_value(deposited_atoms) -> None: + """Angles and their restraint sigmas enter the NLL in radians.""" + import gemmi + + xyz, _ = deposited_atoms + positions = [gemmi.Position(*row) for row in xyz[:4].cpu().tolist()] + angles = np.array( + [gemmi.calculate_angle(*positions[:3]), gemmi.calculate_angle(*positions[1:4])] + ) + idx = _indices([[0, 1, 2], [1, 2, 3]], xyz.device) + refs = xyz.new_tensor([1.8, 2.0]) + sigma = xyz.new_tensor([0.1, 0.2]) + expected = _gaussian_sum(angles - refs.cpu().numpy(), sigma.cpu().numpy()) + torch.testing.assert_close( + angle_math(xyz, idx, refs, sigma), xyz.new_tensor(expected) + ) + + +def test_chiral_value(deposited_atoms) -> None: + """The signed scalar triple product, without a 1/6 factor, sets chirality.""" + xyz, _ = deposited_atoms + host = xyz[:4].cpu().numpy().astype(np.float64) + volume = np.linalg.det(host[1:] - host[0]) + idx = _indices([[0, 1, 2, 3]], xyz.device) + refs = xyz.new_tensor([2.0]) + sigma = xyz.new_tensor([0.5]) + expected = _gaussian_sum(volume - 2.0, 0.5) + torch.testing.assert_close( + chiral_math(xyz, idx, refs, sigma), xyz.new_tensor(expected) + ) + + +def test_planarity_value(deposited_atoms) -> None: + """The plane penalty sums signed-distance Gaussian NLLs over its atoms.""" + xyz, _ = deposited_atoms + host = xyz[:5].cpu().numpy().astype(np.float64) + centered = host - host.mean(axis=0) + _, _, vh = np.linalg.svd(centered, full_matrices=False) + distances = centered @ vh[-1] + idx = _indices([[0, 1, 2, 3, 4]], xyz.device) + sigma = xyz.new_full((1, 5), 0.2) + expected = _gaussian_sum(distances, 0.2) + torch.testing.assert_close( + planarity_math(xyz, [(idx, sigma)]), xyz.new_tensor(expected) + ) + + +def test_simu_value(deposited_atoms) -> None: + """SIMU penalizes differences of deposited isotropic B-factors in Ų.""" + _, b = deposited_atoms + host = b[:4].cpu().numpy().astype(np.float64) + idx = _indices([[0, 1], [2, 3]], b.device) + expected = _gaussian_sum(host[[0, 2]] - host[[1, 3]], 2.0) + torch.testing.assert_close( + adp_simu_math(b, idx, b.new_tensor(2.0)), b.new_tensor(expected) + ) + + +@pytest.mark.parametrize("weighting, expected", [("sigma", 6.5), ("unit", 20.0)]) +def test_least_squares_value_and_mask(weighting: str, expected: float) -> None: + """Least squares sums half squared amplitude errors using the selected weights.""" + obs = torch.tensor( + [10.0, 20.0, 30.0], dtype=get_float_dtype(), device=get_default_device() + ) + calc = -obs - obs.new_tensor([2.0, 6.0, 50.0]) + sigma = obs.new_tensor([1.0, 2.0, 5.0]) + mask = torch.tensor([True, True, False], device=obs.device) + loss = ls_xray_loss_math(obs, calc, sigma, mask, weighting=weighting) + torch.testing.assert_close(loss, obs.new_tensor(expected)) + torch.testing.assert_close( + ls_xray_loss_math(obs, obs, sigma, weighting=weighting), obs.new_zeros(()) + ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 96ce096c..ecc8cdde 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,148 +1,18 @@ -""" -Unit test specific fixtures. -Unit tests should NOT use real file I/O - use mocks or minimal in-memory data. -""" -import pytest -import torch -import numpy as np - -from torchref.config import dtypes - - -@pytest.fixture -def random_seed(): - """Set random seed for reproducibility.""" - seed = 42 - np.random.seed(seed) - torch.manual_seed(seed) - return seed - - -@pytest.fixture -def random_coordinates(): - """Generate random atomic coordinates.""" - def _generate(n_atoms: int = 10, seed: int = 42): - np.random.seed(seed) - return torch.tensor(np.random.rand(n_atoms, 3) * 10, dtype=dtypes.float) - return _generate - - -@pytest.fixture -def random_fractional_coordinates(): - """Generate random fractional coordinates (0-1 range).""" - def _generate(n_atoms: int = 10, seed: int = 42): - np.random.seed(seed) - return torch.tensor(np.random.rand(n_atoms, 3), dtype=dtypes.float) - return _generate - - -@pytest.fixture -def random_adp(): - """Generate random ADPs (atomic displacement parameters, reasonable range 10-60 Ų).""" - def _generate(n_atoms: int = 10, seed: int = 42): - np.random.seed(seed) - return torch.tensor(np.random.rand(n_atoms) * 50 + 10, dtype=dtypes.float) - return _generate - - -@pytest.fixture -def random_occupancies(): - """Generate random occupancies (0-1 range).""" - def _generate(n_atoms: int = 10, seed: int = 42): - np.random.seed(seed) - return torch.tensor(np.random.rand(n_atoms) * 0.5 + 0.5, dtype=dtypes.float) - return _generate - - -@pytest.fixture -def mock_cell(): - """Mock cell parameters [a, b, c, alpha, beta, gamma].""" - return torch.tensor([50.0, 60.0, 70.0, 90.0, 90.0, 90.0], dtype=dtypes.float) - - -@pytest.fixture -def mock_cell_triclinic(): - """Mock triclinic cell parameters.""" - return torch.tensor([40.0, 50.0, 60.0, 70.0, 80.0, 85.0], dtype=dtypes.float) - - -@pytest.fixture -def mock_hkl_indices(): - """Generate mock HKL indices.""" - def _generate(n_reflections: int = 100, max_index: int = 10, seed: int = 42): - np.random.seed(seed) - h = np.random.randint(-max_index, max_index + 1, n_reflections) - k = np.random.randint(-max_index, max_index + 1, n_reflections) - l = np.random.randint(-max_index, max_index + 1, n_reflections) - # Exclude (0,0,0) - mask = ~((h == 0) & (k == 0) & (l == 0)) - h, k, l = h[mask], k[mask], l[mask] - return torch.tensor(np.stack([h, k, l], axis=1), dtype=dtypes.float) - return _generate - - -@pytest.fixture -def mock_structure_factors(): - """Generate mock structure factors (complex).""" - def _generate(n_reflections: int = 100, seed: int = 42): - np.random.seed(seed) - real = np.random.randn(n_reflections) * 100 - imag = np.random.randn(n_reflections) * 100 - return torch.tensor(real + 1j * imag, dtype=dtypes.complex) - return _generate - - -@pytest.fixture -def mock_F_obs(): - """Generate mock observed structure factor amplitudes.""" - def _generate(n_reflections: int = 100, seed: int = 42): - np.random.seed(seed) - # Positive values with realistic distribution - return torch.tensor(np.abs(np.random.randn(n_reflections) * 100) + 10, dtype=dtypes.float) - return _generate - - -@pytest.fixture -def mock_F_sigma(): - """Generate mock sigma values for F_obs.""" - def _generate(n_reflections: int = 100, seed: int = 42): - np.random.seed(seed) - return torch.tensor(np.abs(np.random.randn(n_reflections) * 5) + 1, dtype=dtypes.float) - return _generate - - -@pytest.fixture -def mock_aniso_u(): - """Generate mock anisotropic U tensor components [U11, U22, U33, U12, U13, U23].""" - def _generate(n_atoms: int = 10, seed: int = 42): - np.random.seed(seed) - # Diagonal elements (positive) - u11 = np.random.rand(n_atoms) * 0.05 + 0.02 - u22 = np.random.rand(n_atoms) * 0.05 + 0.02 - u33 = np.random.rand(n_atoms) * 0.05 + 0.02 - # Off-diagonal elements (can be negative, smaller magnitude) - u12 = (np.random.rand(n_atoms) - 0.5) * 0.02 - u13 = (np.random.rand(n_atoms) - 0.5) * 0.02 - u23 = (np.random.rand(n_atoms) - 0.5) * 0.02 - return torch.tensor(np.stack([u11, u22, u33, u12, u13, u23], axis=1), dtype=dtypes.float) - return _generate - - -@pytest.fixture -def mock_scattering_factors(): - """Generate mock scattering factors.""" - def _generate(n_reflections: int = 100, n_atoms: int = 10, seed: int = 42): - np.random.seed(seed) - # Decreasing with resolution (approximate) - return torch.tensor(np.random.rand(n_reflections, n_atoms) * 5 + 1, dtype=dtypes.float) - return _generate - - -@pytest.fixture -def mock_weights(): - """Generate mock weights for atoms.""" - def _generate(n_atoms: int = 10, seed: int = 42): - np.random.seed(seed) - weights = np.random.rand(n_atoms) - return torch.tensor(weights / weights.sum(), dtype=dtypes.float).reshape(-1, 1) - return _generate +"""Expose synthetic fixtures only to the unit-test subtree.""" + +from tests.fixtures.numerical import ( # noqa: F401 + mock_aniso_u, + mock_cell, + mock_cell_triclinic, + mock_F_obs, + mock_F_sigma, + mock_hkl_indices, + mock_scattering_factors, + mock_structure_factors, + mock_weights, + random_adp, + random_coordinates, + random_fractional_coordinates, + random_occupancies, + random_seed, +) diff --git a/tests/unit/frf_separate/__init__.py b/tests/unit/frf_separate/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/frf_separate/test_bessel_rescale.py b/tests/unit/frf_separate/test_bessel_rescale.py new file mode 100644 index 00000000..200b277d --- /dev/null +++ b/tests/unit/frf_separate/test_bessel_rescale.py @@ -0,0 +1,128 @@ +"""What the Bessel ladder's on-the-fly rescaling is allowed to cost: nothing. + +Miller's downward recurrence for ``j_u(x)`` seeds at an arbitrary magnitude and +renormalises at the end, so the intermediate ladder is inflated by whatever the +true ``j_{n_start}(x)`` happens to be -- 1e157 at the FRF's low-resolution end. +That is fine in float64 and overflows float32 for every ``x`` below about 35, +which is most of the resolution range, so the recurrence rescales as it goes. + +The rescale factor is a power of two on purpose: dividing by one decrements the +exponent and leaves the mantissa alone, so the table must come out **bitwise** +equal to the un-rescaled ladder rather than merely close to it. That is the +property these tests pin, against two independent references -- the previous +implementation, transcribed inline, and ``scipy.special.spherical_jn``. +""" + +import math + +import pytest +import torch + +from torchref.experimental.alignment.frf.data_mr import ( + _BESSEL_RESCALE_EXP, + spherical_bessel_table, +) + +pytestmark = pytest.mark.unit + +#: Bessel arguments spanning the FRF's range. ``bessel_h_scale = lmax_even * +#: d_min_eff``, so ``x`` runs from ``bessel_h_scale / d_max`` (~1.3 at +#: d_max = 100 A) up to exactly ``lmax_even`` at the high-resolution limit. +_X_VALUES = [1.257, 1.885, 3.0, 5.0, 10.0, 20.0, 40.0, 64.0] + +_U_MAX = 65 # lmax_even + 1 at the shipped LMAX_CAP = 64 +_N_EXTRA = 25 + + +def _unrescaled_reference(x, u_max, n_extra=_N_EXTRA): + """The recurrence as it stood before rescaling, transcribed verbatim. + + Kept as a literal copy rather than a call into the module: the point is to + compare against the *previous* arithmetic, so it must not track any later + edit to the production function. + """ + x64 = x.to(torch.float64) + safe_x = x64.clamp(min=1e-30) + inv_x = 1.0 / safe_x + n_start = max(u_max + n_extra, u_max + 2) + j_high = torch.zeros_like(x64) + j_mid = torch.ones_like(x64) + j_table = torch.zeros((u_max + 1, *x64.shape), dtype=torch.float64) + peak = torch.zeros_like(x64) + for n in range(n_start, 0, -1): + j_low = (2.0 * n + 1.0) * inv_x * j_mid - j_high + if n - 1 <= u_max: + j_table[n - 1] = j_low + peak = torch.maximum(peak, j_low.abs()) + j_high = j_mid + j_mid = j_low + true_j0 = torch.sin(x64) * inv_x + true_j0 = torch.where(x64 < 1e-30, torch.ones_like(x64), true_j0) + computed_j0 = j_table[0] + safe_j0 = torch.where( + computed_j0.abs() < 1e-30, torch.ones_like(computed_j0), computed_j0, + ) + j_table = j_table * (true_j0 / safe_j0).unsqueeze(0) + perm = list(range(1, j_table.dim())) + [0] + return j_table.permute(*perm).contiguous(), peak + + +def test_rescaling_is_bit_identical_to_the_unrescaled_ladder(): + x = torch.tensor(_X_VALUES, dtype=torch.float64) + ref, _ = _unrescaled_reference(x, _U_MAX) + got = spherical_bessel_table(x, _U_MAX) + assert got.shape == ref.shape + assert torch.equal(got, ref), ( + "rescaling perturbed the table; max relative deviation " + f"{((got - ref).abs() / ref.abs().clamp(min=1e-300)).max().item():.3e}" + ) + + +def test_the_rescale_branch_is_actually_exercised(): + """Guard against a vacuous equality test. + + If the ladder never crossed the threshold the comparison above would pass + while testing nothing, so assert the un-rescaled ladder really does run away + -- and past float32's ceiling, which is the reason the rescaling exists. + """ + x = torch.tensor(_X_VALUES, dtype=torch.float64) + _, peak = _unrescaled_reference(x, _U_MAX) + threshold = 2.0 ** _BESSEL_RESCALE_EXP + f32_max = torch.finfo(torch.float32).max + crossed = int((peak > threshold).sum()) + over_f32 = int((peak > f32_max).sum()) + assert crossed >= len(_X_VALUES) - 2, ( + f"only {crossed} of {len(_X_VALUES)} arguments cross 2**" + f"{_BESSEL_RESCALE_EXP}; the rescale path is nearly dead" + ) + assert over_f32 >= 5, ( + f"only {over_f32} arguments overflow float32 (peaks: " + f"{[f'{v:.2e}' for v in peak.tolist()]})" + ) + + +def test_matches_scipy_spherical_jn(): + """Independent oracle, over the range where float64 carries the answer.""" + scipy_special = pytest.importorskip("scipy.special") + x = torch.tensor(_X_VALUES, dtype=torch.float64) + got = spherical_bessel_table(x, _U_MAX) + for i, xv in enumerate(_X_VALUES): + for u in range(0, _U_MAX + 1): + want = float(scipy_special.spherical_jn(u, xv)) + # Below ~1e-290 the reference itself is at the edge of float64, and + # the FRF flushes anything under float32's smallest normal anyway. + if abs(want) < 1e-30: + continue + mine = float(got[i, u]) + assert math.isclose(mine, want, rel_tol=1e-9, abs_tol=1e-300), ( + f"j_{u}({xv}) = {mine:.12e}, scipy says {want:.12e}" + ) + + +def test_batched_shape_and_dtype_round_trip(): + x = torch.rand(7, 3, dtype=torch.float64) * 60.0 + 1.3 + got = spherical_bessel_table(x, 20) + assert got.shape == (7, 3, 21) + assert got.dtype == torch.float64 + got32 = spherical_bessel_table(x.to(torch.float32), 20) + assert got32.dtype == torch.float32 diff --git a/tests/unit/frf_separate/test_bessel_sh_grouping.py b/tests/unit/frf_separate/test_bessel_sh_grouping.py new file mode 100644 index 00000000..44a83ccb --- /dev/null +++ b/tests/unit/frf_separate/test_bessel_sh_grouping.py @@ -0,0 +1,301 @@ +"""What the SH-Bessel expansion's grouping is allowed to cost. + +`bessel_sh_expand` does not sum over reflections one at a time. It groups them +by ``(|s|, cos theta)`` -- both factors are constant within a group -- and sums +over groups, which is what makes the cost tractable at L = 100. Two properties +make that safe, and both are cheap to lose silently: + +* the negative-m half of the result is redundant, exactly, so only half is + computed and the rest mirrored; +* the group representative is the group *mean*, so the grouping's error stays + far below what the rest of the chain contributes. + +The reference for the second is an expansion with a grouping key so fine that +every reflection is its own group -- i.e. the ungrouped sum. Comparing against +the previous implementation instead would only show that two approximations +agree with each other. + +Everything here runs at the configured precision on the configured device, so +the expansion is exercised where it actually runs. That fixes what the budgets +can say. At float32 the difference between two groupings is dominated by the +accumulation's own rounding, not by the grouping, and that floor is not +portable: the three grouping cases cost 5.5e-08 to 1.5e-07 on CPU, 4.2e-07 to +7.4e-07 on this machine's MPS, and up to 2.1e-06 on the project's MPS CI, +because the backends reduce in a different order and MPS varies further by GPU +and torch version. So a single working-precision budget covers them all -- +see ``_WORKING_PRECISION_TOLERANCE`` for where the number comes from and what +it costs. +""" + +import math + +import pytest +import torch + +import torchref.experimental.alignment.frf.data_mr as dm +from torchref.config import get_default_device, get_float_dtype +from torchref.experimental.alignment.frf.data_mr import bessel_sh_expand + +pytestmark = pytest.mark.unit + +#: A key this fine puts every reflection in its own group: the exact sum. +_UNGROUPED = 10 ** 16 + +#: What any of these comparisons is allowed to cost at the working precision. +#: +#: Measured, worst case over the assertions below: 1.9e-06 on CPU float32 and +#: 1.4e-06 on this machine's MPS, with the grouping cases reaching 2.1e-06 on +#: the project's MPS CI. 1e-04 clears the worst observed figure by ~50x, which is +#: ample for the backend spread -- the same case differs 3x between two MPS +#: devices, and CPU is worse than MPS on the lattice. +#: +#: It still sits far below the accuracy of what feeds it. The sphere/voxel +#: discretisation upstream of the structure factors is itself good to about +#: 5e-03 rel L2 (``tests/unit/base/test_canonical_sphere_cpu.py``), so a budget +#: here of 1e-04 is ~50x tighter than the input it is expanding. Downstream is +#: the same story: rotation-function scores moved 1.8e-07 to 2.5e-04 relative +#: across the float32 migration and the peak lists on 1DAW, 3K7M, 2DQ6 and 4BX9 +#: came back in the same order. +#: +#: What it gives up, and this is the real cost: at this width the grouping +#: comparison notices neither a 10x coarser key (1.6e-06) nor a 100x one +#: (4.1e-05). Only 1000x (1.3e-03) trips it. These are smoke tests for gross +#: breakage on the device the code actually runs on, not tripwires for a +#: degraded key -- ``test_matches_an_independent_direct_summation`` is the one +#: that would still catch a dropped term, because it compares against a +#: reference rather than against the expansion at another grouping. Given the +#: 5e-03 upstream, a key error of 4.1e-05 is not a correctness problem anyway; +#: it would be a performance-versus-accuracy choice made by accident. +_WORKING_PRECISION_TOLERANCE = 1e-4 + +#: Alias kept for the grouping cases, which is what the assertion message names. +_GROUPING_TOLERANCE = _WORKING_PRECISION_TOLERANCE + +#: The direct-summation check compares against a host-double reference rather +#: than against the expansion at another grouping, so it carries the working +#: precision's whole error. Measured 8.1e-07; same budget, same reasoning. +_DIRECT_SUM_TOLERANCE = _WORKING_PRECISION_TOLERANCE + + +def _to_working(s, I, dtype=None, device=None): + """Place a host-double set at the configured precision and device. + + Drawn in double and cast once, rather than generated at the working dtype, + so the *set* is the same whatever precision it is evaluated in -- otherwise + a tolerance measured at one dtype is not comparable with the same tolerance + at another, because the sample moved too. + """ + dtype = get_float_dtype() if dtype is None else dtype + device = get_default_device() if device is None else device + return (s.to(device=device, dtype=dtype), I.to(device=device, dtype=dtype)) + + +def _random_set(seed, n, dtype=None, device=None): + g = torch.Generator().manual_seed(seed) + s = torch.randn(n, 3, generator=g, dtype=torch.float64) + s = s / s.norm(dim=-1, keepdim=True) * ( + 0.07 + 0.18 * torch.rand(n, 1, generator=g, dtype=torch.float64)) + I = torch.randn(n, generator=g, dtype=torch.float64) + return _to_working(s, I, dtype, device) + + +def _grid_set(k=8, step=0.013, dtype=None, device=None): + """A lattice, where |s| degeneracy is exact and the grouping pays most.""" + idx = torch.arange(-k, k + 1, dtype=torch.float64) + a, b, c = torch.meshgrid(idx, idx, idx, indexing="ij") + s = torch.stack([a.reshape(-1), b.reshape(-1), c.reshape(-1)], dim=-1) * step + s = s[s.norm(dim=-1) > 1e-9] + g = torch.Generator().manual_seed(4) + I = torch.randn(s.shape[0], generator=g, dtype=torch.float64) + return _to_working(s, I, dtype, device) + + +@pytest.fixture +def ungrouped(): + """Run the expansion with grouping effectively disabled.""" + def run(s, I, **kw): + ks, kc = dm._GROUP_SCALE_S, dm._GROUP_SCALE_COS + dm._GROUP_SCALE_S = dm._GROUP_SCALE_COS = _UNGROUPED + try: + return bessel_sh_expand(s, I, **kw).coeffs + finally: + dm._GROUP_SCALE_S, dm._GROUP_SCALE_COS = ks, kc + return run + + +@pytest.mark.parametrize("L,scale", [(17, 30.0), (33, 48.0), (65, 64.0)]) +def test_grouping_error_stays_far_below_the_reference_implementation( + L, scale, ungrouped): + """The grouped sum must track the ungrouped one.""" + s, I = _random_set(seed=21, n=5000) + ref = ungrouped(s, I, L=L, bessel_h_scale=scale) + got = bessel_sh_expand(s, I, L=L, bessel_h_scale=scale).coeffs + rel = (got - ref).abs().max().item() / max(ref.abs().max().item(), 1e-300) + assert rel < _GROUPING_TOLERANCE, ( + f"L={L}: grouping cost {rel:.2e} relative, over the {_GROUPING_TOLERANCE:.0e} " + f"budget. A coarser key, or a group representative that is not the mean, " + f"will do this." + ) + + +def test_a_lattice_groups_without_loss(ungrouped): + """On a lattice the degeneracy is exact, so the grouping is free.""" + s, I = _grid_set() + ref = ungrouped(s, I, L=65, bessel_h_scale=64.0) + got = bessel_sh_expand(s, I, L=65, bessel_h_scale=64.0).coeffs + rel = (got - ref).abs().max().item() / max(ref.abs().max().item(), 1e-300) + # "Loss-free" is a float64 statement: on a lattice the |s| degeneracy is + # exact, so in double this lands at 1e-15. At float32 the accumulation floor + # is 1.9e-06 (CPU) / 1.4e-06 (MPS) and swamps it, so what is asserted here + # is that the lattice case is no worse than any other -- the exactness claim + # is only visible in double. + assert rel < _WORKING_PRECISION_TOLERANCE, ( + f"lattice grouping is not loss-free: {rel:.2e}" + ) + + +@pytest.mark.parametrize("L", [17, 33, 65]) +def test_negative_m_is_the_conjugate_of_positive_m(L): + """``c[n,l,-m] = (-1)^m conj(c[n,l,+m])``, which is why only half is summed. + + The intensity, the radial weight and the Legendre factor are all real, and + P_{l,|m|} does not distinguish +m from -m, so m enters only through the + azimuthal phase. If this ever fails, the mirrored half of the array is + wrong and the rotation function is being fed a non-Hermitian Patterson. + """ + s, I = _random_set(seed=5, n=3000) + c = bessel_sh_expand(s, I, L=L, bessel_h_scale=48.0).coeffs + for m in range(1, L): + pos = c[:, :, (L - 1) + m] + neg = c[:, :, (L - 1) - m] + assert torch.equal(neg, ((-1.0) ** m) * pos.conj()), f"m={m} mirror broken" + + +def test_the_group_representative_is_the_mean_not_a_member(): + """A scatter assignment leaves an arbitrary member; the mean is centred. + + Two reflections inside one key bin, placed either side of the bin centre: + with a mean representative the expansion is symmetric under swapping which + one comes first in the array, with a last-writer-wins representative it is + not. + """ + L, scale = 33, 48.0 + eps = 0.4 / dm._GROUP_SCALE_S # comfortably inside one bin + base = torch.tensor([[0.10, 0.03, 0.05]], dtype=torch.float64) + unit = base / base.norm() + r = base.norm() + a = unit * (r - eps) + b = unit * (r + eps) + I = torch.tensor([1.0, 1.0], dtype=torch.float64) + + ab, I_w = _to_working(torch.cat([a, b]), I) + ba, _ = _to_working(torch.cat([b, a]), I) + fwd = bessel_sh_expand(ab, I_w, L=L, bessel_h_scale=scale).coeffs + rev = bessel_sh_expand(ba, I_w, L=L, bessel_h_scale=scale).coeffs + rel = (fwd - rev).abs().max().item() / max(fwd.abs().max().item(), 1e-300) + # Not loosened with the others: the mean of two numbers does not depend on + # their order in floating point either, so this measures 0.0 exactly at + # float32 and float64 alike. A budget here would only hide a real failure. + assert rel == 0.0, ( + f"reordering two reflections in the same bin changed the result by " + f"{rel:.2e}; the representative is order-dependent, so it is not the mean" + ) + + +def _reference_expansion(s_vec, intensity, L, bessel_h_scale): + """A slow, obvious expansion, built from independently-checked parts. + + Sums over reflections one at a time, taking the Legendre factor from + ``sh._bar_legendre_recurrence`` -- which ``tests/unit/alignment/test_sh.py`` + pins against ``scipy`` -- and the radial factor from + ``spherical_bessel_table``. No grouping, no shells, no fused loop. + + This exists because the other tests in this file compare + ``bessel_sh_expand`` against *itself* at a different grouping resolution, so + a term dropped from the sum appears identically on both sides and cancels. + One did: a refactor formed the products that feed the shell accumulation + before writing the sectoral (m = l) entry of the Legendre row, silently + losing that entry for every even l. Every test here passed, and the only + signal was a benchmark truth rank moving from 8 to 13. + """ + from torchref.experimental.alignment.frf.data_mr import spherical_bessel_table + from torchref.experimental.alignment.sh import _bar_legendre_recurrence + + lmax = L - 1 + lmax_even = lmax if lmax % 2 == 0 else lmax - 1 + N_radial = (lmax_even - 2) // 2 + 1 + u_max = lmax_even + 1 + + smag = s_vec.norm(dim=-1).clamp(min=1e-30) + cos_t = (s_vec[:, 2] / smag).clamp(-1.0, 1.0) + sin_t = (1.0 - cos_t * cos_t).clamp(min=0.0).sqrt() + phi = torch.atan2(s_vec[:, 1], s_vec[:, 0]) + + barP = _bar_legendre_recurrence(cos_t, sin_t, L) # (M, L, L) + x = (bessel_h_scale * smag).clamp(min=1e-30) + j = spherical_bessel_table(x, u_max) # (M, u_max+1) + + out = torch.zeros((N_radial, L, 2 * L - 1), dtype=torch.complex128) + for l in range(2, lmax_even + 1, 2): + for n in range((lmax_even - l) // 2 + 1): + u = l + 2 * n + 1 + radial = math.sqrt(2 * u + 1) * j[:, u] / x + for m in range(-l, l + 1): + # Y_lm = barP_{l,|m|} * C(m, phi), with C = (-1)^m e^{i m phi} + # for m >= 0 and e^{i m phi} for m < 0; the expansion uses the + # conjugate. + sign = (-1.0) ** m if m >= 0 else 1.0 + phase = torch.polar(torch.ones_like(phi), -m * phi) + term = (intensity * radial * barP[:, l, abs(m)] * sign) * phase + out[n, l, (L - 1) + m] = term.sum() + return out + + +@pytest.mark.parametrize("L", [9, 13]) +def test_matches_an_independent_direct_summation(L): + """The whole expansion, against a reference that shares no code with it. + + The reference stays in host double while the expansion runs at the working + precision, so this is the one test here that measures the expansion's true + accuracy rather than its self-consistency. That is also why its budget is + the widest: it carries the working precision's whole error, not just the + grouping's share of it. + """ + s64, I64 = _random_set(seed=31, n=120, dtype=torch.float64, device="cpu") + ref = _reference_expansion(s64, I64, L, 24.0) + got = bessel_sh_expand(*_random_set(seed=31, n=120), + L=L, bessel_h_scale=24.0).coeffs + assert got.shape == ref.shape + # `.cpu()` before widening: complex128 cannot be materialised on a backend + # that has no double. + got = got.cpu().to(ref.dtype) + rel = (got - ref).abs().max().item() / max(ref.abs().max().item(), 1e-300) + assert rel < _DIRECT_SUM_TOLERANCE, ( + f"L={L}: differs from a direct summation by {rel:.2e}" + ) + + +def test_the_antipodal_copy_would_only_double_the_result(): + """Why the expansion no longer concatenates ``-s``. + + Only even ``l`` are computed, and ``Y_lm(-s_hat) = (-1)^l Y_lm(s_hat)``, so + for even ``l`` the antipodal copy contributes exactly what the original does. + The intensity is duplicated verbatim and ``|s|`` is unchanged, so the whole + coefficient array doubles and nothing about its *shape* changes -- which is + why removing it rescales the rotation function by four and reorders nothing. + + Pinned here rather than argued in a comment: if a future change made odd ``l`` + carry signal, this equality would break and the removal would need revisiting. + """ + s, I = _random_set(seed=17, n=300) + single = bessel_sh_expand(s, I, L=17, bessel_h_scale=30.0).coeffs + doubled = bessel_sh_expand( + torch.cat([s, -s], dim=0), torch.cat([I, I], dim=0), + L=17, bessel_h_scale=30.0, + ).coeffs + scale = single.abs().max().clamp(min=1e-300) + rel = ((doubled - 2.0 * single).abs().max() / scale).item() + assert rel < _WORKING_PRECISION_TOLERANCE, ( + f"the antipodal copy is not an exact factor of two: {rel:.2e} relative. " + f"Removing it was justified on that being exact." + ) diff --git a/tests/unit/frf_separate/test_invariants.py b/tests/unit/frf_separate/test_invariants.py new file mode 100644 index 00000000..be941667 --- /dev/null +++ b/tests/unit/frf_separate/test_invariants.py @@ -0,0 +1,169 @@ +"""Tier 1 invariant tests for frf_separate. + +Mathematical truths that must hold regardless of input data. Failure +here means a sign / aliasing / normalisation bug in the FFT or the +adaptive sample list. +""" +from __future__ import annotations + +import math + +import pytest +import torch + +from torchref.experimental.alignment.frf.sitelist_ang import ( + build_adaptive_sample_list, + build_dense_map_per_beta, + evaluate_rotation_function, +) +from torchref.symmetry.symmetry import find_fft_friendly_size +from torchref.experimental.alignment.frf.wigner_d import ( + _wigner_d_blocks, + wigner_contraction_per_beta, +) + + +def _make_xi(L: int, seed: int = 42) -> torch.Tensor: + """ξ_{lmn} built as an outer product of two complex SH coefficient sets + that each derive from a real-valued function on S² (so each has the + symmetry ``c_{l,-m} = (-1)^m · conj(c_{l, m})``). The resulting ξ has + the symmetry required for a real-valued RF without us needing to + derive the index relationship by hand. + """ + torch.manual_seed(seed) + # Two random complex SH coefficient vectors with the real-function symmetry. + def _real_sh_coeffs(L): + c = torch.zeros((L, 2 * L - 1), dtype=torch.complex128) + for l in range(2, L, 2): + # m = 0 entries must be real. + c[l, L - 1] = torch.randn(1, dtype=torch.float64).item() + for m in range(1, l + 1): + v = torch.randn(2, dtype=torch.float64) + z = complex(v[0].item(), v[1].item()) + c[l, L - 1 + m] = z + c[l, L - 1 - m] = ((-1) ** m) * complex(z.real, -z.imag) + return c + + c_obs = _real_sh_coeffs(L) + c_calc = _real_sh_coeffs(L) + # ξ_{l, m, n} = c_obs_{l, n} · conj(c_calc_{l, m}) + xi = torch.einsum("ln,lm->lmn", c_obs, torch.conj(c_calc)) + return xi + + +def test_fft_size_is_five_smooth(): + """The dense-map grid must factor into 2, 3 and 5 only. + + Uses the shared ``find_fft_friendly_size``; the alignment package's own + ``adjust_gridding`` was deleted after being measured identical to it for + every n from 1 to 4000 at the only setting the FRF ever called it with. + """ + assert find_fft_friendly_size(180) == 180 # 180 = 4·45 = 4·9·5 + assert find_fft_friendly_size(7) == 8 # rounds up to the next 5-smooth + assert find_fft_friendly_size(243) == 243 # 3^5 + assert find_fft_friendly_size(1) == 1 + + +def test_sample_count_matches_so3_measure(): + """Phaser's adaptive sample count is ≈ (720·360·n_β)/(π·Δ²).""" + Δ = 5.0 + _, _, _, _, beta_grid = build_adaptive_sample_list(Δ) + n_beta = beta_grid.shape[0] + # Manually count by re-building (cheap). + alphas, _, _, _, _ = build_adaptive_sample_list(Δ) + n_samples = alphas.shape[0] + expected = (720.0 * 360.0 * n_beta) / (math.pi * Δ * Δ) + # 10 % tolerance because of dedup + integer rounding. + assert 0.85 * expected <= n_samples <= 1.15 * expected, ( + f"got {n_samples}, expected ~{expected:.0f}" + ) + + +def test_polar_caps_are_one_dimensional(): + """At β=0 only α + γ matters → samples lie on the diagonal α = γ.""" + alphas, betas, gammas, beta_starts, beta_grid = build_adaptive_sample_list(5.0) + i0, i1 = int(beta_starts[0].item()), int(beta_starts[1].item()) + # β=0 slice: α should equal γ for every sample. + assert torch.allclose(alphas[i0:i1], gammas[i0:i1], atol=1e-9) + + +def test_identity_rotation_value(): + """At (α=0, β=0, γ=0): RF = Σ_{l, m} ξ_{l, m, m}. + + The first sample (β=0, p=0) has (α, γ) = (0, 0), so the interpolated + map value at that sample equals RF(0, 0, 0) / N² (the implicit + inverse-FFT normalisation). + """ + L = 8 + xi = _make_xi(L) + Δ = 10.0 + arf = evaluate_rotation_function(xi, grid_sampling_deg=Δ) + + # Expected RF(0,0,0): Σ_l Σ_m ξ_{l, m, m} since d^l_{m,n}(0) = δ_{m,n}. + expected = 0.0 + for l in range(2, L, 2): + for m in range(-l, l + 1): + expected += xi[l, L - 1 + m, L - 1 + m].real.item() + + # The map at (α=γ=0) corresponds to sample index 0 (β=0, p=0). + # fft2 (no normalisation) gives RF directly. + measured = arf.values[0].item() + assert abs(measured - expected) / max(abs(expected), 1e-9) < 1e-5, ( + f"measured={measured}, expected={expected}" + ) + + +def test_real_output_from_hermitian_xi(): + """RF must be real (imaginary residue at numerical-noise level only).""" + L = 8 + xi = _make_xi(L) + Δ = 10.0 + bmax = int(math.ceil(180.0 / Δ)) + N = find_fft_friendly_size(2 * max(bmax, 2 * L - 1)) + _, _, _, _, beta_grid = build_adaptive_sample_list(Δ) + M = build_dense_map_per_beta(xi, beta_grid, N) + # Imaginary part divided by typical magnitude should be < 1e-10. + typ = M.real.abs().mean() + imag_residue = M.imag.abs().mean() + assert (imag_residue / max(typ.item(), 1e-30)) < 1e-9, ( + f"imag/real = {imag_residue.item()/typ.item():.2e}" + ) + + +def test_wigner_contraction_symmetry(): + """At β = π/2 the small-d satisfies d^l_{m,n}(π/2) = (-1)^{l+m} d^l_{m,-n}(π/2).""" + L = 8 + betas = torch.tensor([math.pi / 2], dtype=torch.float64) + blocks = _wigner_d_blocks(L, betas, torch.device("cpu"), torch.float64) + for l in range(2, L, 2): + d = blocks[l - 1][0] # (2l+1, 2l+1) + for m in range(-l, l + 1): + for n in range(-l, l + 1): + lhs = d[l + m, l + n].item() + rhs = ((-1) ** (l + m)) * d[l + m, l - n].item() + assert abs(lhs - rhs) < 1e-10, ( + f"l={l} m={m} n={n}: {lhs} vs {rhs}" + ) + + +def test_beta_reflection_identity(): + """d^l_{m,n}(π − β) = (-1)^{l+m} d^l_{m,-n}(β) (small-d β-reflection).""" + L = 6 + beta = 0.37 + betas = torch.tensor([beta, math.pi - beta], dtype=torch.float64) + blocks = _wigner_d_blocks(L, betas, torch.device("cpu"), torch.float64) + for l in range(2, L, 2): + d = blocks[l - 1] # (2, 2l+1, 2l+1) + for m in range(-l, l + 1): + for n in range(-l, l + 1): + lhs = d[1, l + m, l + n].item() # d(π-β) + rhs = ((-1) ** (l + m)) * d[0, l + m, l - n].item() # (-1)^(l+m) d(β)|n→-n + assert abs(lhs - rhs) < 1e-10 + + +def test_zero_xi_gives_zero_rf(): + """Trivial sanity: empty input → empty output.""" + L = 6 + xi = torch.zeros((L, 2 * L - 1, 2 * L - 1), dtype=torch.complex128) + arf = evaluate_rotation_function(xi, grid_sampling_deg=10.0) + assert arf.values.abs().max().item() < 1e-15 diff --git a/tests/unit/frf_separate/test_legendre_kernel.py b/tests/unit/frf_separate/test_legendre_kernel.py new file mode 100644 index 00000000..0df4d5e8 --- /dev/null +++ b/tests/unit/frf_separate/test_legendre_kernel.py @@ -0,0 +1,106 @@ +"""The fused Legendre/shell kernel against the portable reference. + +Two things need pinning. The kernel must agree with the torch reference -- it is +selected automatically wherever it builds, so a divergence would silently change +every rotation search on that host. And it must **refuse** float64 rather than +accept it: it reads every array through a raw ``float*``, so a float64 buffer +would be reinterpreted as twice as many float32s, not converted. + +Agreement is to a float32 tolerance, not bit-exact, and deliberately so: the +kernel accumulates cluster-by-cluster within a shell while ``index_add_`` +accumulates over the whole chunk, and in single precision a different summation +order is a different answer. Measured 4e-7 to 1e-5 relative, which sits between +the grouping's own error and Phaser's cos(theta) bucketing at ~2e-5. +""" + +import pytest +import torch + +from torchref.experimental.alignment.frf.kernels import portable +from torchref.experimental.alignment.frf.kernels.cpu import legendre_shell as fused +from torchref.experimental.alignment.sh import legendre_recurrence_coefficients + +pytestmark = pytest.mark.unit + +#: Summation order in single precision, nothing more. Set above the measured +#: 1e-5 worst case with margin; a real divergence is orders larger, because a +#: dropped term changes whole rows rather than their last digits. +_TOL = 1e-4 + + +def _case(L, n_clusters, n_shells, seed): + """Random input in the layout the kernel requires: shells sorted.""" + g = torch.Generator().manual_seed(seed) + cos_t = 2 * torch.rand(n_clusters, generator=g, dtype=torch.float32) - 1 + sin_t = (1 - cos_t * cos_t).clamp(min=0).sqrt() + Dr = torch.randn(n_clusters, L, generator=g, dtype=torch.float32) + Di = torch.randn(n_clusters, L, generator=g, dtype=torch.float32) + shell = torch.sort( + torch.randint(0, n_shells, (n_clusters,), generator=g))[0] + a, b, sect = legendre_recurrence_coefficients( + L, torch.float32, torch.device("cpu")) + n_even = (L - 1 if (L - 1) % 2 == 0 else L - 2) // 2 + return dict(shape=(n_even, n_shells, L), args=(cos_t, sin_t, Dr, Di, shell, + a, b, sect)) + + +def _run(fn, case): + Tr = torch.zeros(case["shape"], dtype=torch.float32) + Ti = torch.zeros_like(Tr) + fn(Tr, Ti, *case["args"]) + return Tr, Ti + + +@pytest.mark.parametrize("L,n_clusters,n_shells", [(13, 500, 40), + (65, 4000, 300), + (101, 3000, 250)]) +def test_fused_agrees_with_portable(L, n_clusters, n_shells): + if not fused.available(): + pytest.skip(f"fused kernel unavailable: {fused.why_unavailable()}") + case = _case(L, n_clusters, n_shells, seed=4) + ref_r, ref_i = _run(portable.legendre_shell_accumulate, case) + got_r, got_i = _run(fused.legendre_shell_accumulate, case) + for name, ref, got in (("real", ref_r, got_r), ("imag", ref_i, got_i)): + rel = (got - ref).abs().max().item() / max(ref.abs().max().item(), 1e-30) + assert rel < _TOL, f"L={L} {name} part differs by {rel:.2e}" + + +def test_float64_is_refused_not_reinterpreted(): + """A float64 caller must raise, naming the dtype.""" + if not fused.available(): + pytest.skip(f"fused kernel unavailable: {fused.why_unavailable()}") + L, n_clusters, n_shells = 9, 20, 5 + n_even = (L - 1) // 2 + d = torch.float64 + with pytest.raises(RuntimeError, match="float32 only"): + fused.legendre_shell_accumulate( + torch.zeros(n_even, n_shells, L, dtype=d), + torch.zeros(n_even, n_shells, L, dtype=d), + torch.zeros(n_clusters, dtype=d), torch.zeros(n_clusters, dtype=d), + torch.zeros(n_clusters, L, dtype=d), + torch.zeros(n_clusters, L, dtype=d), + torch.zeros(n_clusters, dtype=torch.long), + torch.zeros(L, L, dtype=d), torch.zeros(L, L, dtype=d), + torch.zeros(L, dtype=d)) + + +def test_shell_offsets_partition_the_clusters(): + """The kernel's work split: contiguous, complete, and in shell order.""" + shell = torch.tensor([0, 0, 2, 2, 2, 5], dtype=torch.long) + off = fused.shell_offsets(shell, 6) + assert off.tolist() == [0, 2, 2, 5, 5, 5, 6] + assert int(off[-1]) == shell.numel() + + +def test_the_dispatch_prefers_the_fused_kernel_when_it_builds(): + """Whatever the table selects is what the expansion runs.""" + from torchref.experimental.alignment.frf._backends import LEGENDRE_BACKENDS + from torchref.utils.backends import select + + probe = [torch.zeros(1, 1, 9, dtype=torch.float32)] * 6 + chosen = select(LEGENDRE_BACKENDS, probe).name + expected = "cpu_fused" if fused.available() else "portable" + assert chosen == expected, ( + f"table chose {chosen!r} but the fused kernel is " + f"{'available' if fused.available() else 'unavailable'}" + ) diff --git a/tests/unit/frf_separate/test_synthetic.py b/tests/unit/frf_separate/test_synthetic.py new file mode 100644 index 00000000..70b882e3 --- /dev/null +++ b/tests/unit/frf_separate/test_synthetic.py @@ -0,0 +1,178 @@ +"""Tier 2 synthetic golden-input tests for frf_separate. + +End-to-end: feed a known-rotation pair into the engine and check the top +peak is at the right Euler. +""" +from __future__ import annotations + +import math + +import pytest +import torch + +from torchref.experimental.alignment.frf.api import FastRotationFunction + + +def _random_rotation_matrix(seed: int) -> torch.Tensor: + """Uniform random SO(3) rotation as a 3×3 matrix (Edmonds ZYZ).""" + g = torch.Generator().manual_seed(seed) + # Shoemake's quaternion method + u = torch.rand(3, generator=g, dtype=torch.float64) + q0 = math.sqrt(1 - u[0]) * math.sin(2 * math.pi * u[1]) + q1 = math.sqrt(1 - u[0]) * math.cos(2 * math.pi * u[1]) + q2 = math.sqrt(u[0]) * math.sin(2 * math.pi * u[2]) + q3 = math.sqrt(u[0]) * math.cos(2 * math.pi * u[2]) + q = torch.tensor([q0, q1, q2, q3], dtype=torch.float64) + w, x, y, z = q.tolist() + return torch.tensor( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], + [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], + [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], + ], + dtype=torch.float64, + ) + + +def _euler_to_matrix_edmonds_zyz(a: float, b: float, g: float) -> torch.Tensor: + ca, sa = math.cos(a), math.sin(a) + cb, sb = math.cos(b), math.sin(b) + cg, sg = math.cos(g), math.sin(g) + return torch.tensor([ + [ca*cb*cg - sa*sg, -ca*cb*sg - sa*cg, ca*sb], + [sa*cb*cg + ca*sg, -sa*cb*sg + ca*cg, sa*sb], + [-sb*cg, sb*sg, cb], + ], dtype=torch.float64) + + +def _so3_angular_distance_deg(R1: torch.Tensor, R2: torch.Tensor) -> float: + tr = torch.einsum("ij,ij->", R1, R2).item() + cos_t = max(min((tr - 1.0) * 0.5, 1.0), -1.0) + return math.acos(cos_t) * 180.0 / math.pi + + +def _make_random_reflections(seed: int, n: int = 800) -> tuple: + """Build a P1 reflection set on a uniform sphere in [s_min, s_max]. + + Returns (s_vec, F, centric) torch tensors. + """ + g = torch.Generator().manual_seed(seed) + # Uniform on shell |s| ∈ [0.07, 0.25] (~4-15 Å) + s_mag = 0.07 + 0.18 * torch.rand(n, generator=g, dtype=torch.float64) + # Uniform on sphere + theta = torch.acos(2 * torch.rand(n, generator=g, dtype=torch.float64) - 1) + phi = 2 * math.pi * torch.rand(n, generator=g, dtype=torch.float64) + s_vec = torch.stack( + [s_mag * torch.sin(theta) * torch.cos(phi), + s_mag * torch.sin(theta) * torch.sin(phi), + s_mag * torch.cos(theta)], dim=-1, + ) + # Random positive amplitudes (Wilson-distributed, roughly) + F = torch.randn(n, generator=g, dtype=torch.float64).abs() + 0.1 + centric = torch.zeros(n, dtype=torch.bool) + return s_vec, F, centric + + +def _search(s_obs, F_obs, centric_obs, s_calc, F_calc, *, sym_mats, + n_peaks, sigma_threshold, **kw): + """Construct the engine and score one model, as `rotation_search` does.""" + engine = FastRotationFunction(s_obs, F_obs, centric_obs, sym_mats, **kw) + return engine.score_model(s_calc, F_calc, n_peaks=n_peaks, + sigma_threshold=sigma_threshold) + + +@pytest.mark.parametrize("seed", [0, 1, 2, 7]) +def test_synthetic_rotation_recovery(seed: int): + """Apply a known rotation to obs; FRF top peak must recover it (within Δ).""" + torch.manual_seed(seed) + s_calc, F_calc, centric = _make_random_reflections(seed=seed, n=2000) + R_truth = _random_rotation_matrix(seed=seed + 100) + + # s_obs = R_truth · s_calc (apply rotation on the same intensities) + s_obs = (R_truth @ s_calc.T).T + F_obs = F_calc.clone() + + sym_mats = torch.eye(3, dtype=torch.float64).unsqueeze(0) # P1 + + arf, peaks = _search( + s_obs, F_obs, centric, s_calc, F_calc, + sym_mats=sym_mats, L=16, d_min=4.0, d_max=15.0, + delta_vrms_A=0.5, grid_sampling_deg=5.0, + n_peaks=20, sigma_threshold=-100.0, + ) + + assert len(peaks) > 0, "no peaks returned" + # Convention (phaser_frf.py:1140): the peak Euler R satisfies s_calc = R · s_obs. + # We set s_obs = R_truth · s_calc, so the expected peak is at R_truth^{-1}. + R_expected = R_truth.T + # Top peak might not be exactly top-1; allow top-3 (Phaser bracket). + best_err = min( + _so3_angular_distance_deg( + _euler_to_matrix_edmonds_zyz(p.alpha, p.beta, p.gamma), + R_expected, + ) + for p in peaks[:3] + ) + assert best_err < 15.0, ( + f"seed={seed}: best-of-top-3 {best_err:.2f}° from truth; " + f"top peak Euler " + f"({math.degrees(peaks[0].alpha):.1f}, {math.degrees(peaks[0].beta):.1f}, " + f"{math.degrees(peaks[0].gamma):.1f})" + ) + + +def test_api_returns_correct_types(): + """Type check on the drop-in API.""" + from torchref.experimental.alignment.frf.types import AdaptiveRotationFunction, RotationPeak + + s_calc, F_calc, centric = _make_random_reflections(seed=0, n=500) + sym_mats = torch.eye(3, dtype=torch.float64).unsqueeze(0) + arf, peaks = _search( + s_calc, F_calc, centric, s_calc, F_calc, + sym_mats=sym_mats, L=8, d_min=4.0, d_max=15.0, + grid_sampling_deg=10.0, n_peaks=5, sigma_threshold=-100.0, + ) + assert isinstance(arf, AdaptiveRotationFunction) + assert len(peaks) > 0 + assert all(isinstance(p, RotationPeak) for p in peaks) + + +def test_search_is_bit_reproducible_single_threaded(): + """The engine itself is deterministic: no RNG, no order-dependent atomics. + + Repeat runs at the default thread count are NOT bit-identical -- the + structure-factor reduction is float32 and its parallel summation order + varies, which on a real high-symmetry case (3GR5, P 6_5 2 2) perturbs peak + scores by ~5e-8 relative and reorders roughly a dozen of 500 peaks. Truth + rank was stable there, but nothing guarantees it for two peaks that close. + + Pinning one thread removes the only source of variation, so any future + difference under this test is a real change in the maths. It is also the + configuration to use when comparing peak lists across a refactor. + """ + import torch as _torch + + s_calc, F_calc, centric = _make_random_reflections(seed=3, n=800) + sym_mats = _torch.eye(3, dtype=_torch.float64).unsqueeze(0) + kwargs = dict( + sym_mats=sym_mats, L=12, d_min=4.0, d_max=15.0, delta_vrms_A=0.5, + grid_sampling_deg=8.0, n_peaks=25, sigma_threshold=-100.0, + ) + + n_threads = _torch.get_num_threads() + _torch.set_num_threads(1) + try: + _, first = _search(s_calc, F_calc, centric, s_calc, F_calc, **kwargs) + _, second = _search(s_calc, F_calc, centric, s_calc, F_calc, **kwargs) + finally: + _torch.set_num_threads(n_threads) + + assert len(first) == len(second) > 0 + for i, (a, b) in enumerate(zip(first, second)): + assert a.alpha == b.alpha and a.beta == b.beta and a.gamma == b.gamma, ( + f"peak {i} moved between two identical single-threaded runs" + ) + assert a.score == b.score, ( + f"peak {i} score changed by {abs(a.score - b.score):.3g} between " + f"two identical single-threaded runs" + ) diff --git a/tests/unit/io/test_anomalous_output.py b/tests/unit/io/test_anomalous_output.py index e110a6e2..b32a2e68 100644 --- a/tests/unit/io/test_anomalous_output.py +++ b/tests/unit/io/test_anomalous_output.py @@ -16,7 +16,6 @@ import reciprocalspaceship as rs import torch -from torchref.base.french_wilson import is_centric_from_hkl from torchref.io.datasets.reflection_data import ReflectionData from torchref.model.model_ft import ModelFT @@ -68,7 +67,7 @@ def test_hkl_for_sf_signs(self, anomalous_data): def test_centrics_not_flagged(self, anomalous_data): d = anomalous_data - centric = is_centric_from_hkl(d.hkl, d.spacegroup) + centric = d.spacegroup.is_centric(d.hkl) assert not bool((d.friedel_flags & centric).any()) def test_hkl_for_sf_fallback(self): diff --git a/tests/unit/io/test_anomalous_reader.py b/tests/unit/io/test_anomalous_reader.py index 2fc0692d..d72e0f8e 100644 --- a/tests/unit/io/test_anomalous_reader.py +++ b/tests/unit/io/test_anomalous_reader.py @@ -15,7 +15,6 @@ import reciprocalspaceship as rs import torch -from torchref.base.french_wilson import is_centric_from_hkl from torchref.io.datasets.reflection_data import ReflectionData from torchref.model.model_ft import ModelFT @@ -81,7 +80,7 @@ def test_centrics_not_duplicated(anomalous_two_column_mtz): path, _ = anomalous_two_column_mtz d = ReflectionData(verbose=0) d.load_mtz(path) - centric = is_centric_from_hkl(d.hkl, d.spacegroup) + centric = d.spacegroup.is_centric(d.hkl) # Centric reflections obey Friedel's law and must appear exactly once each. canon = [tuple(h) for h in d.hkl[centric].tolist()] assert len(canon) == len(set(canon)) @@ -128,7 +127,11 @@ def test_generated_rfree_shared_across_mates(anomalous_two_column_mtz): assert bool(d.friedel_flags.any()) d.regenerate_rfree_flags(force=True, seed=0) - assert d.rfree_source == "Generated (resolution-binned, ASU-grouped)" + # The seed is part of the provenance: "generated" without it names a draw + # nobody can reproduce. + assert d.rfree_source == ( + "Generated (resolution-binned, ASU-grouped, seed 0)" + ) assert bool((d.rfree_flags == 0).any()) # a free set actually exists assert _mixed_partition_groups(d) == [] diff --git a/tests/unit/io/test_hkl_convention.py b/tests/unit/io/test_hkl_convention.py index 6a08ba34..44636ac4 100644 --- a/tests/unit/io/test_hkl_convention.py +++ b/tests/unit/io/test_hkl_convention.py @@ -57,7 +57,10 @@ def anomalous_data(mtz_dir, tmp_path): def _model(pdb_dir, data): - m = ModelFT(verbose=0, max_res=2.0) + # strip_H: what is under test is the phase convention, and the absolute check + # compares against a gemmi calculation that calls ``remove_hydrogens``. Letting + # torchref generate hydrogens would have it computing a different structure. + m = ModelFT(verbose=0, max_res=2.0, strip_H=True) m.load_pdb(str(pdb_dir / f"{CODE}.pdb")) m.cell, m.spacegroup = data.cell, data.spacegroup return m diff --git a/tests/unit/io/test_multicomponent_restraints.py b/tests/unit/io/test_multicomponent_restraints.py index 780d3414..65d344ab 100644 --- a/tests/unit/io/test_multicomponent_restraints.py +++ b/tests/unit/io/test_multicomponent_restraints.py @@ -18,8 +18,8 @@ import pytest from torchref.io.cif_readers import RestraintCIFReader -from torchref.restraints.library import get_library_manager -from torchref.restraints.restraints_helper import ( +from torchref.topology.monomer.library import get_library_manager +from torchref.topology.monomer.cif import ( split_data_blocks, validate_restraint_data, ) @@ -230,7 +230,7 @@ class TestChiralitySpellings: """The CCP4 library writes both ``positive`` and the truncated ``positiv``.""" def test_short_spellings_are_not_dropped(self): - from torchref.restraints.builders_fast import PreprocessedCIF + from torchref.topology.builders import PreprocessedCIF chirals = pd.DataFrame( { diff --git a/tests/unit/io/test_refinement_header.py b/tests/unit/io/test_refinement_header.py new file mode 100644 index 00000000..9da010e0 --- /dev/null +++ b/tests/unit/io/test_refinement_header.py @@ -0,0 +1,440 @@ +"""Regression tests for the refinement output header. + +The writer used to carry an input PDB's header through verbatim and then append +its own ``REMARK 3``, so a refined file asserted two different refinements at +once -- the previous program's R-factors first, ours several hundred lines +later. The selection was also inverted: TITLE, AUTHOR and REMARK were kept +(including the superseded statistics) while SEQRES, SSBOND, LINK, CISPEP, +EXPDTA, COMPND, DBREF and FORMUL were dropped, i.e. the chemistry refinement +does not invalidate. + +What is asserted here is the resulting contract: exactly one refinement block +and it is ours; records that describe the crystal survive; records that describe +a superseded refinement do not; and the chain of programs applied to the model +is preserved through mmCIF's ``_software`` loop rather than by hoarding the +previous program's output. +""" + +import pandas as pd +import pytest + +from torchref.io import cif, pdb +from torchref.io.metadata import RefinementMetadata + +# 3GR5 was refined with REFMAC 5.1.24 and carries a full deposition header: +# 420 lines including REMARK 2/3/500, JRNL, AUTHOR, SEQRES, SSBOND and SITE. +INPUT_PDB = "tests/files/pdb/3GR5.pdb" + +#: PDB record order, abridged to the records this writer can emit. The format +#: mandates this sequence; TITLE used to be written *after* REMARK 900. +RECORD_ORDER = [ + "TITLE", "COMPND", "SOURCE", "KEYWDS", "EXPDTA", "MDLTYP", "AUTHOR", + "REMARK", "DBREF", "DBREF1", "DBREF2", "SEQADV", "SEQRES", "MODRES", + "HET", "HETNAM", "HETSYN", "FORMUL", "SSBOND", "LINK", "CISPEP", "SITE", +] + + +def _atom_df(): + """Two atoms whose coordinates exercise trailing-zero formatting. + + ``18.3`` and ``31.51`` are the values that exposed the writer using + ``str()`` semantics on a rounded float instead of ``%8.3f``. + """ + df = pd.DataFrame( + { + "ATOM": ["ATOM", "HETATM"], + "serial": [1, 2], + "name": ["CA", "O"], + "altloc": ["", ""], + "resname": ["LEU", "HOH"], + "chainid": ["A", "A"], + "resseq": [1, 2], + "icode": ["", ""], + "x": [-7.223, -9.22], + "y": [29.982, 31.51], + "z": [18.3, 20.364], + "occupancy": [1.0, 1.0], + "tempfactor": [95.4, 30.0], + "element": ["C", "O"], + "charge": [0, 0], + "anisou_flag": [False, False], + "u11": [0.0, 0.0], + "u22": [0.0, 0.0], + "u33": [0.0, 0.0], + "u12": [0.0, 0.0], + "u13": [0.0, 0.0], + "u23": [0.0, 0.0], + } + ) + df.attrs["cell"] = [90.645, 90.645, 133.422, 90.0, 90.0, 120.0] + df.attrs["spacegroup"] = "P 65 2 2" + return df + + +def _refined_metadata(): + """Input header plus this refinement's statistics, as the writer builds it.""" + meta = RefinementMetadata.from_pdb_file(INPUT_PDB) + ours = RefinementMetadata( + program_version="0.7.0", + target_function="ML", + optimizer="1 MACROCYCLE, SEPARATE, ISOTROPIC ADP, SCALE TARGET NLL", + r_work=0.2083, + r_free=0.2471, + percent_free=9.9, + n_reflections_all=20942, + n_reflections_test=2063, + resolution_high=2.05, + resolution_low=18.69, + starting_model=INPUT_PDB, + rfree_selection="MTZReader FreeR", + ) + return meta.merge(ours) + + +def _header_lines(): + return _refined_metadata().render_pdb_header().splitlines() + + +def _remark_number(line): + try: + return int(line[7:10]) + except ValueError: + return None + + +# ====================================================================== # +# One refinement block, and it is ours +# ====================================================================== # + + +@pytest.mark.unit +def test_exactly_one_refinement_block(): + header = "\n".join(_header_lines()) + assert header.count("REMARK 3 REFINEMENT.") == 1 + assert "TORCHREF" in header + + +@pytest.mark.unit +@pytest.mark.parametrize("program", ["REFMAC", "PHENIX", "BUSTER", "CNS"]) +def test_no_foreign_program_is_credited(program): + """The input was refined by REFMAC; the output must not say so anywhere.""" + assert program not in "\n".join(_header_lines()) + + +@pytest.mark.unit +def test_superseded_statistics_are_dropped(): + """REMARK 2, 3 and 500 describe a resolution and a model we replaced.""" + numbers = {_remark_number(line) for line in _header_lines() + if line.startswith("REMARK")} + assert 2 not in numbers # resolution: regenerated + assert 500 not in numbers # geometry outliers of old coordinates + assert 3 in numbers # ours, and the only one + + +@pytest.mark.unit +def test_input_remark_3_is_not_carried_through(): + """The specific inversion this module exists to prevent.""" + meta = RefinementMetadata.from_pdb_file(INPUT_PDB) + assert not any(_remark_number(r) == 3 for r in meta.passthrough_pdb_remarks) + # ... while the input genuinely has one, so the assertion has teeth. + with open(INPUT_PDB) as handle: + assert any(line.startswith("REMARK 3") for line in handle) + + +# ====================================================================== # +# Attribution +# ====================================================================== # + + +@pytest.mark.unit +def test_authors_and_journal_are_not_inherited(): + """Both credit the deposition, not this refinement.""" + meta = RefinementMetadata.from_pdb_file(INPUT_PDB) + assert meta.authors == [] + lines = _header_lines() + assert not any(line.startswith(("AUTHOR", "JRNL")) for line in lines) + + +@pytest.mark.unit +def test_explicitly_set_authors_are_written(): + """Not inheriting authors must not stop --authors from working.""" + meta = _refined_metadata() + meta.authors = ["A.Person", "B.Other"] + lines = meta.render_pdb_header().splitlines() + author = [line for line in lines if line.startswith("AUTHOR")] + assert author and "A.Person" in author[0] + + +# ====================================================================== # +# Records that survive +# ====================================================================== # + + +@pytest.mark.unit +@pytest.mark.parametrize( + "record", ["COMPND", "SOURCE", "KEYWDS", "EXPDTA", "DBREF", "SEQADV", + "SEQRES", "HETNAM", "FORMUL", "SSBOND", "SITE"] +) +def test_structural_records_are_carried_through(record): + """Refinement moves atoms; it does not change the sequence or chemistry.""" + with open(INPUT_PDB) as handle: + expected = sum(1 for line in handle if line.startswith(record)) + assert expected > 0, f"{record} absent from the fixture" + written = sum(1 for line in _header_lines() if line.startswith(record)) + assert written == expected + + +@pytest.mark.unit +def test_secondary_structure_is_not_emitted(): + """Nothing here computes HELIX/SHEET, so carrying them can only mislead.""" + lines = _header_lines() + assert not any(line.startswith(("HELIX", "SHEET")) for line in lines) + + +@pytest.mark.unit +def test_records_are_in_mandated_order(): + seen = [] + for line in _header_lines(): + record = line[:6].strip() + if record and (not seen or seen[-1] != record): + seen.append(record) + assert all(r in RECORD_ORDER for r in seen), [ + r for r in seen if r not in RECORD_ORDER + ] + positions = [RECORD_ORDER.index(r) for r in seen] + assert positions == sorted(positions), seen + + +@pytest.mark.unit +def test_remarks_ascend_with_ours_at_three(): + numbers = [ + _remark_number(line) for line in _header_lines() + if line.startswith("REMARK") + ] + assert numbers == sorted(numbers) + assert 3 in numbers + + +# ====================================================================== # +# Free text is the author's, never generated +# ====================================================================== # + + +@pytest.mark.unit +def test_no_remarks_section_without_author_text(): + assert "OTHER REFINEMENT REMARKS" not in "\n".join(_header_lines()) + + +@pytest.mark.unit +def test_author_text_is_written_when_supplied(): + meta = _refined_metadata() + meta.output_remarks = "Re-refined for the benchmark.\n\nSecond paragraph." + header = meta.render_pdb_header() + assert "REMARK 3 OTHER REFINEMENT REMARKS:" in header + assert "Re-refined for the benchmark." in header + assert "Second paragraph." in header + + +@pytest.mark.unit +def test_free_set_provenance_is_reported(): + """R-free from a different test set is not comparable; say which one.""" + header = "\n".join(_header_lines()) + assert "FREE R VALUE TEST SET SELECTION" in header + assert "MTZReader FreeR" in header + + +@pytest.mark.unit +def test_remark_3_colons_line_up_within_each_block(): + """Colons used to be ragged, because alignment relied on caller padding. + + Two columns exist by design, which is what REFMAC does too: a narrow one + for the identification lines (``PROGRAM :``) and a wide one for the + statistics (``RESOLUTION RANGE HIGH (ANGSTROMS) :``). Each must be + internally consistent, including on wrapped continuation lines. + """ + columns = [ + line.index(" : ") for line in _header_lines() + if line.startswith("REMARK 3 ") and " : " in line + ] + assert set(columns) == {24, 46}, sorted(set(columns)) + + +@pytest.mark.unit +def test_header_lines_fit_the_format(): + assert [line for line in _header_lines() if len(line) > 80] == [] + + +@pytest.mark.unit +def test_long_identification_values_wrap_rather_than_overflow(): + """Naming cycles, mode, ADP model and scale target overruns column 80.""" + meta = _refined_metadata() + meta.optimizer = ( + "12 MACROCYCLES, EVERYTHING, FIELD_ANISO ADP, SCALE TARGET ML_NOALPHA, " + "RIGID BODY 5 ITERATIONS" + ) + lines = meta.render_pdb_header().splitlines() + assert [line for line in lines if len(line) > 80] == [] + optimizer = [line for line in lines if "OPTIMIZER" in line] + assert len(optimizer) == 1 # one head line ... + head = lines.index(optimizer[0]) + assert lines[head + 1].startswith("REMARK 3 : ") + # ... and the value survives the wrap intact. + joined = " ".join( + line.split(":", 1)[1].strip() + for line in lines[head:head + 3] + if " : " in line + ) + assert "RIGID BODY 5 ITERATIONS" in joined + + +# ====================================================================== # +# Coordinates +# ====================================================================== # + + +@pytest.mark.unit +def test_numeric_columns_keep_their_decimals(tmp_path): + """``18.3`` must be written `` 18.300`` and ``95.4`` as `` 95.40``. + + Both fields were formatted as ``str()`` of a rounded float rather than with + an explicit precision, so trailing zeros vanished. + """ + out = tmp_path / "out.pdb" + pdb.write(_atom_df(), str(out)) + atoms = [ + line for line in out.read_text().splitlines() + if line.startswith(("ATOM", "HETATM")) + ] + assert atoms + for line in atoms: + assert len(line) == 80, repr(line) + # x, y, z: %8.3f + for start in (30, 38, 46): + field = line[start:start + 8] + assert len(field.split(".")[1]) == 3, repr(field) + # occupancy and B: %6.2f + for start in (54, 60): + field = line[start:start + 6] + assert len(field.split(".")[1]) == 2, repr(field) + assert " 95.40" in atoms[0] + + +# ====================================================================== # +# The software chain -- mmCIF's only room for prior work +# ====================================================================== # + + +@pytest.mark.unit +def test_software_is_a_loop_with_an_ordinal(): + cats = _refined_metadata().render_cif_categories() + software = cats["_software"] + assert software["_software.pdbx_ordinal"] == ["1"] + assert software["_software.name"] == ["TORCHREF"] + + +@pytest.mark.unit +def test_software_ordinal_increments_across_refinements(tmp_path): + """Refining a refined file must append to the chain, not replace it.""" + first = tmp_path / "first.cif" + cif.write_model(_atom_df(), str(first), metadata=_refined_metadata()) + + # Read it back the way a second refinement would, then write again. + carried = RefinementMetadata.from_cif_file(str(first)) + assert len(carried.software_chain) == 1 + + second_meta = carried.merge( + RefinementMetadata(program_version="0.7.0", optimizer="2 MACROCYCLES") + ) + second = tmp_path / "second.cif" + cif.write_model(_atom_df(), str(second), metadata=second_meta) + + chain = RefinementMetadata.from_cif_file(str(second)).software_chain + assert [e.get("pdbx_ordinal") for e in chain] == ["1", "2"] + + +@pytest.mark.unit +def test_previous_refinement_description_survives_the_round_trip(tmp_path): + """The chain is worthless if it cannot say what each program did.""" + out = tmp_path / "one.cif" + meta = _refined_metadata() + meta.optimizer = "7 MACROCYCLES, SEPARATE, ISOTROPIC ADP" + cif.write_model(_atom_df(), str(out), metadata=meta) + + chain = RefinementMetadata.from_cif_file(str(out)).software_chain + assert "7 MACROCYCLES" in chain[0]["description"] + + +@pytest.mark.unit +def test_loop_values_with_spaces_survive_a_round_trip(tmp_path): + """An unquoted loop cell silently splits into extra columns when re-read.""" + out = tmp_path / "quoted.cif" + meta = _refined_metadata() + meta.optimizer = "MANY WORDS, WITH COMMAS, AND SPACES" + cif.write_model(_atom_df(), str(out), metadata=meta) + + chain = RefinementMetadata.from_cif_file(str(out)).software_chain + assert len(chain) == 1 + assert chain[0]["description"].endswith("AND SPACES") + + +@pytest.mark.unit +def test_input_refine_statistics_are_not_carried_into_cif(tmp_path): + """The mmCIF form of the duplicated REMARK 3.""" + source = tmp_path / "prior.cif" + prior = RefinementMetadata(program="REFMAC", r_work=0.213, r_free=0.251) + cif.write_model(_atom_df(), str(source), metadata=prior) + + carried = RefinementMetadata.from_cif_file(str(source)) + assert "_refine" not in carried.passthrough_cif_categories + assert carried.r_work is None + # The prior program is remembered as a link in the chain, not as statistics. + assert [e["name"] for e in carried.software_chain] == ["REFMAC"] + + +@pytest.mark.unit +def test_starting_model_is_recorded_as_an_accession(tmp_path): + cats = _refined_metadata().render_cif_categories() + initial = cats["_pdbx_initial_refinement_model"] + assert initial["_pdbx_initial_refinement_model.accession_code"] == "3GR5" + assert initial["_pdbx_initial_refinement_model.type"] == "experimental model" + assert cats["_refine"]["_refine.pdbx_starting_model"] == INPUT_PDB + + +@pytest.mark.unit +def test_unaccessionable_starting_model_is_named_not_guessed(): + meta = _refined_metadata() + meta.starting_model = "/tmp/my_working_model_v3.pdb" + initial = meta.render_cif_categories()["_pdbx_initial_refinement_model"] + assert "accession_code" not in " ".join(initial) + assert initial["_pdbx_initial_refinement_model.details"] == ( + "my_working_model_v3.pdb" + ) + + +# ====================================================================== # +# Annotating a file is not refining it +# ====================================================================== # + + +@pytest.mark.unit +def test_annotation_keeps_the_existing_refinement(): + """``add-metadata`` adds a title; it does not re-refine. + + Nothing supersedes the input's REMARK 3 or AUTHOR records in that case, so + dropping them would discard statistics and credit that are still accurate. + """ + meta = RefinementMetadata.from_pdb_file( + INPUT_PDB, supersede_refinement=False + ) + numbers = {_remark_number(r) for r in meta.passthrough_pdb_remarks} + assert {2, 3, 500} <= numbers + assert meta.authors # REFMAC-era depositors keep their credit + + +@pytest.mark.unit +def test_refinement_output_supersedes_by_default(): + """The default is the refinement case: the old block goes.""" + meta = RefinementMetadata.from_pdb_file(INPUT_PDB) + numbers = {_remark_number(r) for r in meta.passthrough_pdb_remarks} + assert not ({2, 3, 500} & numbers) + assert meta.authors == [] diff --git a/tests/unit/io/test_reflection_accessor.py b/tests/unit/io/test_reflection_accessor.py index 366e7d83..57d15fd8 100644 --- a/tests/unit/io/test_reflection_accessor.py +++ b/tests/unit/io/test_reflection_accessor.py @@ -58,14 +58,12 @@ def test_validation_carves_out_of_work_and_free(self, data_1daw): (work_before | free_before) - val ) # nothing new appears - def test_F_matches_legacy_masking(self, data_1daw): + def test_F_matches_valid_partition_masking(self, data_1daw): + """Subset amplitudes apply both validity and work/free partition masks.""" d = data_1daw valid = d.masks().to(torch.bool) - _, F, _, rfree = d() # legacy call - F_data = F.get_data() if hasattr(F, "get_data") else F - vmask = F.get_mask() if hasattr(F, "get_mask") else valid - assert torch.allclose(d.work.F, F_data[vmask & rfree.bool()]) - assert torch.allclose(d.free.F, F_data[vmask & ~rfree.bool()]) + assert torch.allclose(d.work.F, d.F[valid & d.rfree_flags.bool()]) + assert torch.allclose(d.free.F, d.F[valid & ~d.rfree_flags.bool()]) def test_select_aligns_full_array(self, data_1daw): d = data_1daw diff --git a/tests/unit/io/test_reflection_data_reindex.py b/tests/unit/io/test_reflection_data_reindex.py index 07a58cc2..8b8e5058 100644 --- a/tests/unit/io/test_reflection_data_reindex.py +++ b/tests/unit/io/test_reflection_data_reindex.py @@ -49,7 +49,7 @@ def test_carries_all_per_reflection_fields(self): # ``light`` lacks the last 10%; the reference grid lacks the first 10%, # so the two sets genuinely differ (each has reflections the other lacks). light = _synthetic(grid[: int(n * 0.9)], seed=1) - ref_hkl = grid[int(n * 0.1):].clone() + ref_hkl = grid[int(n * 0.1) :].clone() assert len(light.hkl_anomalous) == len(light.hkl) # sane before light.validate_hkl(ref_hkl) @@ -146,34 +146,3 @@ def test_forward_finite_with_different_reflection_sets(self, pdb_dir, mtz_dir): target = CollectionDifferenceTarget(dc, mc, scaler=scaler, verbose=0) loss = target.forward() assert torch.isfinite(loss) - - -@pytest.mark.unit -def test_reindex_preserves_non_per_reflection_u_aniso(): - """``U_aniso`` must be exempt by *name*, not by a shape coincidence. - - The reindexer decides "is this per-reflection?" by ``shape[0] == n_hkl``. - That heuristic collides when the dataset happens to have exactly as many - reflections as the field is long -- ``U_aniso`` is ``(6,)``, so a - 6-reflection dataset would have it gathered and reordered as if it were - per-reflection data. - """ - # Exactly 6 reflections: the same length as U_aniso. - hkl = torch.tensor( - [[1, 0, 1], [0, 1, 1], [0, 0, 1], [1, 1, 1], [1, 0, 2], [0, 1, 2]], - dtype=torch.int32, - ) - data = _synthetic(hkl) - u_aniso = torch.tensor([0.1, 0.2, 0.3, 0.01, 0.02, 0.03]) - data.U_aniso = u_aniso.clone().to(device=data.device) - assert len(data.hkl) == data.U_aniso.shape[0] == 6, "precondition: lengths collide" - - # Reindex onto a different HKL ordering/size; U_aniso must not follow. - ref_hkl = torch.tensor( - [[0, 1, 2], [1, 0, 2], [1, 1, 1], [0, 0, 1], [0, 1, 1], [1, 0, 1], [2, 0, 1]], - dtype=torch.int32, - ) - data.validate_hkl(ref_hkl.to(data.device)) - - assert data.U_aniso.shape == (6,), f"U_aniso reshaped to {tuple(data.U_aniso.shape)}" - assert torch.allclose(data.U_aniso.cpu(), u_aniso), "U_aniso values were permuted" diff --git a/tests/unit/io/test_select_reflection_data.py b/tests/unit/io/test_select_reflection_data.py index a3db5a8f..8c69449d 100644 --- a/tests/unit/io/test_select_reflection_data.py +++ b/tests/unit/io/test_select_reflection_data.py @@ -24,9 +24,6 @@ def _make_reflection_data(n=20): rd.phase = torch.rand(n, dtype=torch.float32) * 6.28 rd.fom = torch.rand(n, dtype=torch.float32) - # Non-per-reflection tensor (should NOT be indexed) - rd.U_aniso = torch.rand(6, dtype=torch.float32) - # Cell and spacegroup rd.cell = Cell( torch.tensor([50.0, 60.0, 70.0, 90.0, 90.0, 90.0]), @@ -87,21 +84,6 @@ def test_permutation(self): torch.testing.assert_close(sel.phase, rd.phase[perm]) -class TestNonMatchingTensors: - """Tensors whose first dim != n_refl should be cloned, not indexed.""" - - def test_u_aniso_copied(self): - rd = _make_reflection_data(20) - mask = torch.zeros(20, dtype=torch.bool) - mask[:10] = True - - sel = rd[mask] - - # U_aniso has shape (6,), not (20,), so it should be copied as-is - torch.testing.assert_close(sel.U_aniso, rd.U_aniso) - assert sel.U_aniso.shape == (6,) - - class TestCellAndSpacegroup: """Cell is cloned; spacegroup is copied by reference.""" diff --git a/tests/unit/io/test_struct_conn_links.py b/tests/unit/io/test_struct_conn_links.py new file mode 100644 index 00000000..2e9bc234 --- /dev/null +++ b/tests/unit/io/test_struct_conn_links.py @@ -0,0 +1,74 @@ +"""``_struct_conn`` rows become the LINK-record table the PDB reader produces.""" + +import numpy as np +import pytest + +from torchref.io.cif_readers import ModelCIFReader +from torchref.io.pdb import LINK_COLUMNS + + +@pytest.mark.unit +def test_3e98_peptide_links_to_selenomethionine(cif_dir): + links = ModelCIFReader(str(cif_dir / "3E98.cif")).links + assert list(links.columns) == list(LINK_COLUMNS) + assert len(links) == 8 + assert set(links["name1"]) == {"C"} and set(links["name2"]) == {"N"} + assert set(links["resname2"]) <= {"MSE", "ARG", "ASP"} + assert (links["altloc1"] == "").all() and (links["icode1"] == "").all() + assert links["resseq1"].dtype.kind == "i" + assert np.isfinite(links["length"]).all() + + +@pytest.mark.unit +def test_1daw_metal_contacts_are_kept(cif_dir): + links = ModelCIFReader(str(cif_dir / "1DAW.cif")).links + assert len(links) == 14 + magnesium = links[(links["resname2"] == "MG") | (links["resname1"] == "MG")] + assert len(magnesium) > 0 + pairs = set(zip(links["name1"], links["resname1"], links["resseq1"])) + assert ("OD2", "ASP", 175) in pairs + + +@pytest.mark.unit +def test_disulfides_are_left_to_distance_detection(cif_dir): + links = ModelCIFReader(str(cif_dir / "3A5V.cif")).links + assert len(links) == 12 + assert "SG" not in set(links["name1"]) | set(links["name2"]) + + +@pytest.mark.unit +def test_file_without_struct_conn_gives_empty_table(tmp_path): + minimal = """\ +data_test +_cell.length_a 10.0 +_cell.length_b 10.0 +_cell.length_c 10.0 +_cell.angle_alpha 90.0 +_cell.angle_beta 90.0 +_cell.angle_gamma 90.0 +_symmetry.space_group_name_H-M 'P 1' +loop_ +_atom_site.group_PDB +_atom_site.id +_atom_site.type_symbol +_atom_site.label_atom_id +_atom_site.label_alt_id +_atom_site.label_comp_id +_atom_site.label_asym_id +_atom_site.label_seq_id +_atom_site.pdbx_PDB_ins_code +_atom_site.Cartn_x +_atom_site.Cartn_y +_atom_site.Cartn_z +_atom_site.occupancy +_atom_site.B_iso_or_equiv +_atom_site.auth_seq_id +_atom_site.auth_asym_id +ATOM 1 N N . ALA A 1 ? 0.0 0.0 0.0 1.0 20.0 1 A +ATOM 2 C CA . ALA A 1 ? 1.5 0.0 0.0 1.0 20.0 1 A +""" + path = tmp_path / "no_links.cif" + path.write_text(minimal) + links = ModelCIFReader(str(path)).links + assert len(links) == 0 + assert list(links.columns) == list(LINK_COLUMNS) diff --git a/tests/unit/io/test_wilson_outlier_masks.py b/tests/unit/io/test_wilson_outlier_masks.py index 201bdf2a..23f01b26 100644 --- a/tests/unit/io/test_wilson_outlier_masks.py +++ b/tests/unit/io/test_wilson_outlier_masks.py @@ -126,10 +126,10 @@ def test_planted_zingers_are_rejected(): @pytest.mark.unit -def test_absent_measurements_are_sanity_not_outliers(): +def test_absent_measurements_are_sanity_not_outliers(mtz_dir): """A row with no measurement is not an improbable observation, and counting it as one is what made the old report meaningless.""" - data = ReflectionData(verbose=0).load_mtz("tests/files/mtz/6G9X.mtz") + data = ReflectionData(verbose=0).load_mtz(str(mtz_dir / "6G9X.mtz")) absent = ~data.masks["sanity_F"] assert int(absent.sum()) > 20000, "6G9X carries a large absent population" @@ -140,8 +140,8 @@ def test_absent_measurements_are_sanity_not_outliers(): @pytest.mark.unit -def test_intensity_path_keeps_french_wilsons_guard_under_its_own_key(): - data = ReflectionData(verbose=0).load_mtz("tests/files/mtz/4BX9.mtz") +def test_intensity_path_keeps_french_wilsons_guard_under_its_own_key(mtz_dir): + data = ReflectionData(verbose=0).load_mtz(str(mtz_dir / "4BX9.mtz")) assert data.I is not None, "4BX9 should load via the intensity path" assert ReflectionData.FRENCH_WILSON_MASK_KEY in data.masks @@ -164,11 +164,11 @@ def test_intensity_path_keeps_french_wilsons_guard_under_its_own_key(): "name", ["1DAW", "2DQ6", "3A5V", "3E98", "3GR5", "3K7M", "3VRJ", "4BX9", "5BOV", "6G9X"], ) -def test_deposited_structures_lose_almost_nothing(name, pdb_dir): +def test_deposited_structures_lose_almost_nothing(name, mtz_dir): """Deposited data has already been through processing and merging; a criterion that rejects percent-level populations of it is mis-calibrated, not perceptive.""" - data = ReflectionData(verbose=0).load_mtz(f"tests/files/mtz/{name}.mtz") + data = ReflectionData(verbose=0).load_mtz(str(mtz_dir / f"{name}.mtz")) measured = int(data.masks["sanity_F"].sum()) rejected = int((~data.masks[ReflectionData.WILSON_MASK_KEY]).sum()) @@ -177,7 +177,7 @@ def test_deposited_structures_lose_almost_nothing(name, pdb_dir): @pytest.mark.unit -def test_flagged_reflections_show_no_directional_bias(): +def test_flagged_reflections_show_no_directional_bias(mtz_dir): """The regression that catches a lost anisotropy correction. 1DAW diffracts about four times more strongly along ``h*`` than ``l*``. A @@ -185,7 +185,7 @@ def test_flagged_reflections_show_no_directional_bias(): the strong one, so the flagged set piles up along ``h*`` -- 52 of 56 with mean ``|h|`` nearly twice the dataset's, before the correction existed. """ - data = ReflectionData(verbose=0).load_mtz("tests/files/mtz/1DAW.mtz") + data = ReflectionData(verbose=0).load_mtz(str(mtz_dir / "1DAW.mtz")) _, flagged = _plant_zingers(data, n=400, seed=13) assert len(flagged) > 100, "the planted population must be found first" diff --git a/tests/unit/maps/test_ded_weights.py b/tests/unit/maps/test_ded_weights.py new file mode 100644 index 00000000..60abe281 --- /dev/null +++ b/tests/unit/maps/test_ded_weights.py @@ -0,0 +1,128 @@ +"""The registered difference-coefficient weight schemes. + +Pinned: the three schemes exist with their MTZ column names; ``none`` is flat; +``inverse_variance`` has mean one and floors a zero sigma; ``sigma_d`` gives strong +reflections more weight than weak ones within a shell where inverse variance cannot; +and an all-noise input falls back to inverse variance with a warning that names why. +""" + +import pytest +import torch + +from tests.unit.refinement.test_sigma_d import synth_diff +from torchref.maps.ded_weights import ( + DEFAULT_SCHEME, + SCHEMES, + WEIGHT_COLUMNS, + DedWeightFallbackWarning, + all_ded_weights, + compute_ded_weights, + normalise_mean_one, +) +from torchref.symmetry import SpaceGroup + + +def _inputs(n=20000, sig_frac=1.0, device="cpu"): + d = synth_diff(n=n, sig_frac=sig_frac, device=device) + g = torch.Generator().manual_seed(5) + hkl = torch.randint(-20, 21, (n, 3), generator=g).to(device) + cell = torch.tensor([40.0, 50.0, 60.0, 90.0, 90.0, 90.0], device=device) + return d, hkl, cell, SpaceGroup("P 1", device=device) + + +@pytest.mark.unit +def test_registry_is_consistent(): + assert DEFAULT_SCHEME in SCHEMES + assert set(WEIGHT_COLUMNS) == set(SCHEMES) - {"none"} + with pytest.raises(ValueError): + compute_ded_weights( + "bogus", + delta_obs=torch.zeros(3), + sigma_diff=torch.ones(3), + hkl=torch.zeros(3, 3), + cell=torch.ones(6), + spacegroup=None, + ) + + +@pytest.mark.unit +def test_normalise_mean_one_handles_nonfinite_and_zero(): + # Non-finite entries drop to zero and count in the mean, so the column mean is one + # however many reflections carry weight. + w = normalise_mean_one(torch.tensor([1.0, 3.0, float("nan"), float("inf")])) + assert torch.allclose(w, torch.tensor([1.0, 3.0, 0.0, 0.0])) + assert w.mean() == pytest.approx(1.0) + half = normalise_mean_one(torch.tensor([0.0, 0.0, 2.0, 6.0])) + assert torch.allclose(half, torch.tensor([0.0, 0.0, 1.0, 3.0])) + z = normalise_mean_one(torch.zeros(4)) + assert torch.equal(z, torch.zeros(4)) + + +@pytest.mark.unit +def test_none_and_inverse_variance(any_device): + d, hkl, cell, sg = _inputs(n=2000, device=any_device) + sig = d["sigma_diff"].clone() + sig[0] = 0.0 + kw = { + "delta_obs": d["delta_obs"], + "sigma_diff": sig, + "hkl": hkl, + "cell": cell, + "spacegroup": sg, + } + flat = compute_ded_weights("none", **kw) + assert torch.equal(flat.weights, torch.ones_like(sig)) + ivw = compute_ded_weights("inverse_variance", **kw) + assert ivw.applied == "inverse_variance" + assert abs(float(ivw.weights.mean()) - 1.0) < 1e-5 + # The zero sigma is floored, so it carries the largest finite weight. + assert torch.isfinite(ivw.weights).all() + assert ( + float(ivw.weights[0]) == float(ivw.weights.max()) > float(ivw.weights[1:].max()) + ) + assert ivw.weights.device == d["delta_obs"].device + + +@pytest.mark.unit +def test_sigma_d_favours_strong_reflections_where_inverse_variance_cannot(any_device): + d, hkl, cell, sg = _inputs(device=any_device) + kw = { + "delta_obs": d["delta_obs"], + "sigma_diff": d["sigma_diff"], + "hkl": hkl, + "cell": cell, + "spacegroup": sg, + "f_dark": d["f_dark"], + } + every = all_ded_weights(**kw) + assert set(every) == set(SCHEMES) + sd = every["sigma_d"] + assert sd.applied == "sigma_d" and abs(float(sd.weights.mean()) - 1.0) < 1e-4 + assert 0.8 < sd.diagnostics["gamma"] < 1.2 + assert sd.diagnostics["n_shell"] > 10 and "shells" in sd.diagnostics + # Within the highest-resolution tenth, the strongest reflections carry more weight. + order = torch.argsort(d["d_star_sq"])[-2000:] + f, w = d["f_dark"][order], sd.weights[order] + strong, weak = f > f.median(), f <= f.median() + assert float(w[strong].mean()) > 1.5 * float(w[weak].mean()) + ivw = every["inverse_variance"].weights[order] + assert abs(float(ivw[strong].mean()) - float(ivw[weak].mean())) < 1e-4 + + +@pytest.mark.unit +def test_all_noise_falls_back_to_inverse_variance_with_a_warning(): + d, hkl, cell, sg = _inputs(n=5000, sig_frac=50.0) + kw = { + "delta_obs": d["delta_obs"], + "sigma_diff": d["sigma_diff"] * 1.2, + "hkl": hkl, + "cell": cell, + "spacegroup": sg, + "f_dark": d["f_dark"], + } + with pytest.warns(DedWeightFallbackWarning, match="inverse-variance"): + sd = compute_ded_weights("sigma_d", **kw) + assert sd.scheme == "sigma_d" and sd.applied == "inverse_variance" + assert "fallback_reason" in sd.diagnostics + ivw = compute_ded_weights("inverse_variance", **kw) + assert torch.allclose(sd.weights, ivw.weights) diff --git a/tests/unit/maps/test_map_units.py b/tests/unit/maps/test_map_units.py new file mode 100644 index 00000000..414e77c8 --- /dev/null +++ b/tests/unit/maps/test_map_units.py @@ -0,0 +1,51 @@ +"""Map units: the electrons-per-cubic-Angstrom synthesis. + +Pinned: ``units="electrons"`` is the ``1/N``-normalised FFT rescaled by ``N / V``, i.e. +``(1/V) sum_h F(h) exp(-2 pi i h.x)``; the default is unchanged; an unknown unit is +rejected; a ``DifferenceMap`` accepts a per-reflection scale and the same units. +""" + +import pytest +import torch + +from torchref.io import ReflectionData +from torchref.maps import DifferenceMap, Map +from torchref.model.model_ft import ModelFT + + +@pytest.fixture(scope="module") +def model_ft_and_data(sample_structure_pair): + model = ModelFT() + model.load_cif(str(sample_structure_pair["model"])) + data = ReflectionData() + data.load_mtz(str(sample_structure_pair["reflections"])) + return model, data + + +@pytest.mark.unit +def test_electrons_is_the_volume_normalised_synthesis(model_ft_and_data): + model, data = model_ft_and_data + normalized = Map(data, model, map_type="Fcalc").calculate() + electrons = Map(data, model, map_type="Fcalc", units="electrons").calculate() + volume = data.cell.volume.to(normalized.dtype) + assert torch.allclose( + electrons, normalized * (normalized.numel() / volume), rtol=1e-5, atol=1e-6 + ) + + +@pytest.mark.unit +def test_unknown_units_are_rejected(model_ft_and_data): + model, data = model_ft_and_data + with pytest.raises(ValueError, match="units must be one of"): + Map(data, model, units="e/A3") + + +@pytest.mark.unit +def test_difference_map_scale_and_units(model_ft_and_data): + model, data = model_ft_and_data + plain = DifferenceMap(data, data, model).calculate() + scale = torch.full((len(data),), 2.0, dtype=plain.dtype, device=plain.device) + scaled = DifferenceMap(data, data, model, scale=scale, units="electrons") + out = scaled.calculate() + assert out.shape == plain.shape and torch.isfinite(out).all() + assert scaled.units == "electrons" and scaled.scale is scale diff --git a/tests/unit/model/test_adp_field_mode.py b/tests/unit/model/test_adp_field_mode.py new file mode 100644 index 00000000..14a8c1d1 --- /dev/null +++ b/tests/unit/model/test_adp_field_mode.py @@ -0,0 +1,217 @@ +"""``set_adp_mode("field")``: a third ADP representation on the existing switch. + +The switch is a *conversion*, not a freeze, so entering field mode has to fit the node +values to the B it replaces and leaving has to materialise them back per atom. The +inert-until-selected property matters too: adding the mode must not perturb a model +that never asks for it. +""" + +import math + +import pytest +import torch + +from torchref.model.disorder_field import DisorderFieldTensor +from torchref.model.model import Model +from torchref.model.parameter_wrappers import PositiveMixedTensor + + +@pytest.fixture(scope="module") +def pdb_path(pdb_dir): + return str(pdb_dir / "3GR5.pdb") + + +def _model(pdb_path): + model = Model(verbose=0) + model.load_pdb(pdb_path) + return model + + +@pytest.mark.unit +def test_isotropic_mode_is_untouched(pdb_path): + """The default representation is unchanged by the new branch existing.""" + model = _model(pdb_path) + model.set_adp_mode("isotropic") + assert isinstance(model.adp, PositiveMixedTensor) + assert not model.adp_is_field + assert model.adp().shape == (len(model.pdb),) + + +@pytest.mark.unit +def test_unknown_mode_still_raises(pdb_path): + model = _model(pdb_path) + with pytest.raises(ValueError, match="field"): + model.set_adp_mode("nonsense") + + +@pytest.mark.unit +def test_entering_field_mode_replaces_the_wrapper(pdb_path): + """``adp`` becomes a field whose parameter count is set by nodes, not atoms.""" + model = _model(pdb_path) + n_atoms = len(model.pdb) + model.set_adp_mode("field", n_nodes=40, k_neighbors=8) + + assert model.adp_is_field + assert isinstance(model.adp, DisorderFieldTensor) + assert model.adp.n_nodes == 40 + assert model.adp().shape == (n_atoms,) + # The whole point: far fewer refinable numbers than atoms. + assert int(model.adp.refinable_params.numel()) < n_atoms + + +@pytest.mark.unit +def test_field_tracks_the_b_it_replaced(pdb_path): + """The fit is against the deposited B, so it must resemble it, not restart from flat.""" + model = _model(pdb_path) + before = model.adp().detach().clone() + model.set_adp_mode("field", n_nodes=64, k_neighbors=8) + after = model.adp().detach() + + spread = (before - before.mean()).pow(2).mean().sqrt() + resid = (after - before).pow(2).mean().sqrt() + assert resid < 0.6 * spread, f"rmse {resid:.2f} vs B spread {spread:.2f}" + assert bool((after > 0).all()) + + +@pytest.mark.unit +def test_more_nodes_track_more_closely(pdb_path): + """Node count is the accuracy dial, end to end through the switch.""" + errors = [] + for n in (4, 32, 200): + model = _model(pdb_path) + before = model.adp().detach().clone() + model.set_adp_mode("field", n_nodes=n, k_neighbors=8) + errors.append(float((model.adp().detach() - before).pow(2).mean().sqrt())) + assert errors[0] > errors[1] > errors[2], errors + + +@pytest.mark.unit +def test_leaving_field_mode_materialises_per_atom(pdb_path): + """Round trip out of field mode gives back a per-atom wrapper holding its values.""" + model = _model(pdb_path) + model.set_adp_mode("field", n_nodes=64, k_neighbors=8) + field_values = model.adp().detach().clone() + + model.set_adp_mode("isotropic") + + assert not model.adp_is_field + assert isinstance(model.adp, PositiveMixedTensor) + assert torch.allclose(model.adp().detach(), field_values, atol=1e-5) + + +@pytest.mark.unit +def test_field_mode_collapses_anisotropic_atoms_first(pdb_path): + """Entering from anisotropic goes through B_eq rather than dropping the U.""" + model = _model(pdb_path) + model.set_adp_mode("anisotropic") + assert bool(model.aniso_flag.any()) + + with torch.no_grad(): + U = model.u().detach() + beq = (8.0 * math.pi**2 / 3.0) * (U[:, 0] + U[:, 1] + U[:, 2]) + + model.set_adp_mode("field", n_nodes=200, k_neighbors=8) + + assert not bool(model.aniso_flag.any()), "field mode is isotropic in this stage" + got = model.adp().detach() + finite = torch.isfinite(beq) + spread = (beq[finite] - beq[finite].mean()).pow(2).mean().sqrt() + resid = (got[finite] - beq[finite]).pow(2).mean().sqrt() + assert resid < spread, "field ignored the equivalent isotropic B" + + +@pytest.mark.unit +def test_sf_indices_and_flags_stay_consistent(pdb_path): + """Everything keyed off the iso/aniso split is refreshed, not left stale.""" + model = _model(pdb_path) + model.set_adp_mode("field", n_nodes=32, k_neighbors=8) + + assert not bool(model.aniso_flag.any()) + assert not bool(model.pdb["anisou_flag"].to_numpy().any()) + assert int(model._iso_indices.numel()) == len(model.pdb) + assert bool(model._aniso_is_empty) + + +@pytest.mark.unit +def test_get_iso_and_adp_u6_work_unmodified(pdb_path): + """The two consumers that read ``adp()`` need no knowledge of the field.""" + model = _model(pdb_path) + model.set_adp_mode("field", n_nodes=32, k_neighbors=8) + + xyz, adp, occ = model.get_iso() + assert xyz.shape[0] == adp.shape[0] == occ.shape[0] == len(model.pdb) + + u6 = model.adp_u6() + assert u6.shape == (len(model.pdb), 6) + expected = (adp / (8.0 * math.pi**2)).detach() + assert torch.allclose(u6[:, 0].detach(), expected, atol=1e-6) + assert torch.allclose(u6[:, 3:].detach(), torch.zeros_like(u6[:, 3:]), atol=1e-12) + + +@pytest.mark.unit +def test_gradient_flows_to_the_nodes_through_adp_u6(pdb_path): + """The path an ADP restraint would take reaches the node parameters.""" + model = _model(pdb_path) + model.set_adp_mode("field", n_nodes=32, k_neighbors=8) + model.adp_u6().sum().backward() + grad = model.adp.refinable_params.grad + assert grad is not None and bool(grad.abs().sum() > 0) + + +@pytest.mark.unit +def test_refine_adp_sees_the_node_parameters(pdb_path): + """``parameters_of_types`` needs no change: the field is still the ``adp`` type.""" + model = _model(pdb_path) + model.set_adp_mode("field", n_nodes=32, k_neighbors=8) + params = model.parameters_of_types(["adp"]) + assert len(params) == 1 + assert params[0] is model.adp.refinable_params + # [log B, log sigma, dx, dy, dz] -- positions are refinable by default. + assert params[0].shape == (32, 5) + assert model.adp.refines_positions + + +@pytest.mark.unit +def test_model_copy_repoints_the_accessor(pdb_path): + """A copied field must read the COPY's coordinates, not the original's. + + ``copy()`` carries the borrowed accessor by reference, so without the re-point in + ``Model.copy`` the two models share coordinates and moving one changes the other's + ADPs. + """ + model = _model(pdb_path) + model.set_adp_mode("field", n_nodes=32, k_neighbors=8) + clone = model.copy() + + assert clone.adp._xyz_fn.module is clone.xyz + assert clone.adp._xyz_fn.module is not model.xyz + before = model.adp().detach().clone() + + with torch.no_grad(): + clone.xyz.refinable_params[: len(model.pdb) // 2, 0] += 5.0 + model.adp.reset_forward_cache() + + assert torch.allclose(model.adp().detach(), before, atol=1e-9), ( + "moving the copy changed the original's ADPs -- accessor was not re-pointed" + ) + + +@pytest.mark.unit +def test_state_dict_round_trip_in_field_mode(pdb_path): + """``create_from_state_dict`` rebuilds a field when the saved ``adp`` was one.""" + model = _model(pdb_path) + model.set_adp_mode("field", n_nodes=24, k_neighbors=6) + expected = model.adp().detach().clone() + + sd = { + key: (value.clone() if torch.is_tensor(value) else value) + for key, value in model.state_dict().items() + } + # Restored onto the model's own device: a restore builds on CPU and moves only + # when asked, and comparing a CPU-evaluated field against a device-evaluated one + # would be measuring backend arithmetic, not the round trip. + restored = Model.create_from_state_dict(sd, device=model.device, verbose=0) + + assert restored.adp_is_field + assert restored.adp.n_nodes == 24 + assert torch.allclose(restored.adp().detach(), expected, atol=1e-6) diff --git a/tests/unit/model/test_adp_field_mode_sets.py b/tests/unit/model/test_adp_field_mode_sets.py new file mode 100644 index 00000000..0f33465b --- /dev/null +++ b/tests/unit/model/test_adp_field_mode_sets.py @@ -0,0 +1,139 @@ +"""``set_adp_mode("field_aniso", mode_set=...)``: installing a displacement-mode field. + +The wiring, not the payload arithmetic --- that is +``test_mode_covariance_payload.py``. What matters here is that a mode set reaches the +model through the existing switch, lands in the ``u`` slot like any anisotropic field, +and starts where the constant-U field starts, so entering the parametrisation is not +itself a change to the model. +""" + +import pytest +import torch + +from torchref.model.disorder_field import ( + MODE_SETS, + DisorderFieldTensor, + ModeCovariancePayload, +) +from torchref.model.model import Model + + +@pytest.fixture(scope="module") +def pdb_path(pdb_dir): + return str(pdb_dir / "3GR5.pdb") + + +def _field_model(pdb_path, mode_set=None, n_nodes=8): + model = Model(verbose=0) + model.load_pdb(pdb_path) + model.set_adp_mode( + "field_aniso", n_nodes=n_nodes, k_neighbors=n_nodes, mode_set=mode_set + ) + return model + + +@pytest.mark.unit +@pytest.mark.parametrize("mode_set", sorted(MODE_SETS)) +def test_mode_set_installs_into_the_u_slot(pdb_path, mode_set): + """A mode field is an anisotropic field: same slot, same downstream consumers.""" + model = _field_model(pdb_path, mode_set) + assert model.adp_is_field + assert isinstance(model.u, DisorderFieldTensor) + assert isinstance(model.u.payload, ModeCovariancePayload) + assert model.u.payload.mode_set == mode_set + assert bool(model.aniso_flag.all()) + # The per-atom surface the structure-factor path uses. + u6 = model.adp_u6() + assert u6.shape == (len(model.pdb), 6) + assert torch.isfinite(u6).all() + + +@pytest.mark.unit +@pytest.mark.parametrize( + "mode_set,per_node", [("constant", 10), ("rigid", 25), ("rigid_dilation", 32), ("affine", 82)] +) +def test_parameter_count_is_payload_plus_sigma_plus_offset(pdb_path, mode_set, per_node): + """Storage is [payload | log sigma | 3 offset], so the ladder costs 10/25/32/82.""" + model = _field_model(pdb_path, mode_set, n_nodes=8) + assert model.u.refinable_params.numel() == 8 * per_node + + +@pytest.mark.unit +@pytest.mark.parametrize("mode_set", sorted(MODE_SETS)) +def test_entering_the_mode_starts_at_the_constant_u_field(pdb_path, mode_set): + """Every rung seeds its translation block from the same solve and floors the rest. + + So a freshly installed mode field must give essentially the constant-U field's ADPs. + That is what makes entering this parametrisation safe: the starting R-factor is one + already known, and refinement can only move away from it. + """ + base = _field_model(pdb_path, "constant").adp_u6().detach() + got = _field_model(pdb_path, mode_set).adp_u6().detach() + # Gradient modes start at the Cholesky floor, epsilon^2, which is ~1e-6 A^2 against + # a U of order 0.2 -- present but far below anything observable. + assert torch.allclose(got, base, atol=1e-5) + + +@pytest.mark.unit +def test_positive_definite_per_atom(pdb_path): + """Every atom's U must be PD or the anisotropic B-matrix inverse blows up.""" + model = _field_model(pdb_path, "affine", n_nodes=6) + u6 = model.adp_u6().detach() + M = torch.zeros(u6.shape[0], 3, 3, dtype=u6.dtype) + M[:, 0, 0], M[:, 1, 1], M[:, 2, 2] = u6[:, 0], u6[:, 1], u6[:, 2] + M[:, 0, 1] = M[:, 1, 0] = u6[:, 3] + M[:, 0, 2] = M[:, 2, 0] = u6[:, 4] + M[:, 1, 2] = M[:, 2, 1] = u6[:, 5] + assert float(torch.linalg.eigvalsh(M).min()) > 0.0 + + +@pytest.mark.unit +def test_gradient_reaches_the_node_parameters(pdb_path): + """Through the zero-argument forward, which is the path refinement actually uses.""" + model = _field_model(pdb_path, "rigid") + model.adp_u6().sum().backward() + grad = model.u.refinable_params.grad + assert grad is not None and float(grad.abs().sum()) > 0 + + +@pytest.mark.unit +def test_copy_round_trips_a_mode_field(pdb_path): + """``Model.copy`` shares no storage but must keep the payload and the accessor.""" + model = _field_model(pdb_path, "rigid_dilation") + clone = model.copy() + assert isinstance(clone.u.payload, ModeCovariancePayload) + assert clone.u.payload.mode_set == "rigid_dilation" + assert torch.allclose(clone.adp_u6().detach(), model.adp_u6().detach()) + assert clone.u.refinable_params is not model.u.refinable_params + # The accessor must point at the COPY's coordinates, not the original's. The + # perturbation has to be non-rigid: a field whose nodes are atom centroids is + # translation-invariant by construction, so shifting every atom would change + # nothing and prove nothing. + original = model.adp_u6().detach().clone() + with torch.no_grad(): + clone.xyz.refinable_params[: len(clone.pdb) // 2] += 3.0 + clone.u.reset_forward_cache() + assert not torch.allclose(clone.adp_u6().detach(), original) + # ...and the original must be untouched by it. + model.u.reset_forward_cache() + assert torch.allclose(model.adp_u6().detach(), original) + + +@pytest.mark.unit +def test_mode_set_is_rejected_on_the_isotropic_field(pdb_path): + """There is no isotropic form of a displacement-mode covariance.""" + model = Model(verbose=0) + model.load_pdb(pdb_path) + with pytest.raises(ValueError, match="no isotropic form"): + model.set_adp_mode("field", n_nodes=8, mode_set="rigid") + + +@pytest.mark.unit +def test_leaving_field_mode_materialises_per_atom(pdb_path): + """The conversion out reads the per-atom U, so it works for any payload.""" + model = _field_model(pdb_path, "affine", n_nodes=6) + before = model.adp_u6().detach().clone() + model.set_adp_mode("anisotropic") + assert not model.adp_is_field + kept = model.aniso_flag + assert torch.allclose(model.adp_u6().detach()[kept], before[kept], atol=1e-6) diff --git a/tests/unit/model/test_copy.py b/tests/unit/model/test_copy.py new file mode 100644 index 00000000..96fafa90 --- /dev/null +++ b/tests/unit/model/test_copy.py @@ -0,0 +1,116 @@ +"""A copy must be usable and independent, however the state gets there. + +Two things in ``copy()`` are neither buffers nor parameter wrappers, so the +buffer and module loops do not carry them, and each is handled by the model +rather than by ``copy()`` itself: + +* the **iso/aniso partition** (``_iso_indices``, ``_aniso_indices`` and the two + fast-path flags) is derived on access and keyed on ``aniso_flag``'s identity. + It used to be rebuilt eagerly, which is what made this fragile: a copy is + constructed, *then* has its context replaced and its buffers cloned, so + eagerly-built indices described the wrong ``aniso_flag`` -- and silently, since + a stale partition gathers the wrong atoms rather than raising. +* the **space group**, which lives on ``ModelContext``. That is deliberately a + dataclass and not an ``nn.Module``, so assigning one cannot be intercepted by + ``nn.Module.__setattr__`` and land in ``_modules`` under the property's own + name. The copy must own its space group, not alias the original's. + +These assert the *outcome* -- a copy whose partition is right and whose space +group is its own -- so they keep their teeth regardless of which mechanism +delivers it. + +A third pair of tests here covered ``ModelFT.copy(build_grid=)``, which skipped +building a real-space grid that the cell/spacegroup setters immediately replace. +That option is gone: ``real_space_grid`` is legacy -- the density splat +reconstructs voxel positions from ``frac_matrix`` and never reads it -- so the +waste is being removed where it is produced rather than worked around here. + +``4BX9`` is used because it carries ``ANISOU`` records, so the partition is +genuinely mixed (220 isotropic, 9973 anisotropic) rather than all-isotropic, +where an empty partition would still pass the fast path. +""" + +import pytest +import torch + +_MIXED_ADP_PDB = "4BX9.pdb" + + +@pytest.fixture(scope="module") +def mixed_adp_path(pdb_dir): + p = pdb_dir / _MIXED_ADP_PDB + if not p.exists(): + pytest.skip(f"{_MIXED_ADP_PDB} not available") + return str(p) + + +def _load(cls, path): + return cls(verbose=0).load_pdb(path) + + +@pytest.mark.unit +@pytest.mark.parametrize("cls_name", ["Model", "ModelFT"]) +def test_copy_has_a_usable_iso_aniso_partition(cls_name, mixed_adp_path): + """``get_iso``/``get_aniso`` must work on the copy and agree with the source.""" + import torchref.model as tm + + cls = getattr(tm, cls_name) + m = _load(cls, mixed_adp_path) + c = m.copy() + + # Mutating the original's flags after the copy must not reach the copy, and + # must be picked up by the original -- the property that eager rebuilding + # could not give us. + for attr in ("_iso_indices", "_aniso_indices", "_iso_covers_all", + "_aniso_is_empty"): + assert hasattr(c, attr), f"{cls_name}.copy() dropped {attr}" + + assert torch.equal(c._iso_indices, m._iso_indices) + assert torch.equal(c._aniso_indices, m._aniso_indices) + assert c._iso_covers_all == m._iso_covers_all + assert c._aniso_is_empty == m._aniso_is_empty + + # ``ModelFT`` appends the per-atom form-factor tables to the same tuple, so + # unpack positionally rather than by a fixed arity. + iso = c.get_iso() + aniso = c.get_aniso() + xyz_i, occ_i = iso[0], iso[2] + xyz_a, u_a, occ_a = aniso[0], aniso[1], aniso[2] + assert xyz_i.shape[0] == m._iso_indices.numel() + assert xyz_a.shape[0] == m._aniso_indices.numel() + assert u_a.shape[-1] == 6 + assert occ_i.shape[0] == xyz_i.shape[0] + assert occ_a.shape[0] == xyz_a.shape[0] + + +@pytest.mark.unit +@pytest.mark.parametrize("cls_name", ["Model", "ModelFT"]) +def test_the_partition_is_mixed_so_the_test_has_teeth(cls_name, mixed_adp_path): + """Guard the premise: an all-isotropic model would not exercise the split.""" + import torchref.model as tm + + m = _load(getattr(tm, cls_name), mixed_adp_path) + assert m._iso_indices.numel() > 0 + assert m._aniso_indices.numel() > 0 + assert not m._iso_covers_all + assert not m._aniso_is_empty + + +@pytest.mark.unit +@pytest.mark.parametrize("cls_name", ["Model", "ModelFT"]) +def test_copy_owns_its_spacegroup(cls_name, mixed_adp_path): + """The copy carries the space group without aliasing or double-registering.""" + import torchref.model as tm + + m = _load(getattr(tm, cls_name), mixed_adp_path) + c = m.copy() + + assert c.spacegroup is not None + assert str(c.spacegroup) == str(m.spacegroup) + # Own object: `.to(device)` on the copy must not move the original's matrices. + assert c.spacegroup is not m.spacegroup + # The space group is context state, not a submodule: nothing may register it + # under the property's own name, which is how the original bug manifested. + assert "spacegroup" not in c._modules + stray = [k for k in c.state_dict() if k.startswith("spacegroup.")] + assert stray == [], f"copy registered a second space group: {stray}" diff --git a/tests/unit/model/test_create_from_state_dict.py b/tests/unit/model/test_create_from_state_dict.py index 1e8ae822..4fd0231c 100644 --- a/tests/unit/model/test_create_from_state_dict.py +++ b/tests/unit/model/test_create_from_state_dict.py @@ -69,3 +69,33 @@ def test_create_from_state_dict_aniso_u_roundtrip(pdb_dir, cls_name): # At least some atoms are anisotropic → finite, non-trivial u values present. assert torch.isfinite(u_fresh).any() assert torch.allclose(u_fresh, u_restored, equal_nan=True) + + +@pytest.mark.unit +def test_altloc_pairs_survive_state_dict_round_trip(pdb_dir): + """``altloc_pairs`` must reach the state dict and come back. + + It lives on the model's context rather than the model, so a defensive + ``hasattr(self, "altloc_pairs")`` in ``state_dict`` silently substituted an empty + list -- losing the alternative-conformation grouping on every save without + failing anything. + """ + from torchref.model import ModelFT + + cpu = torch.device("cpu") + model = ModelFT() + model.load_pdb(str(pdb_dir / "7L84.pdb")) # carries alternative conformations + model.to(cpu) + + assert model.ctx.altloc_pairs, "fixture should have alternative conformations" + + sd = model.state_dict() + assert sd["altloc_pairs"], "altloc groups must reach the state dict" + + restored = ModelFT.create_from_state_dict(sd, device=cpu, verbose=0) + + assert len(restored.ctx.altloc_pairs) == len(model.ctx.altloc_pairs) + for got, want in zip(restored.ctx.altloc_pairs, model.ctx.altloc_pairs): + assert len(got) == len(want) + for g, w in zip(got, want): + assert torch.equal(g, w) diff --git a/tests/unit/model/test_disorder_field.py b/tests/unit/model/test_disorder_field.py new file mode 100644 index 00000000..dcb8f11d --- /dev/null +++ b/tests/unit/model/test_disorder_field.py @@ -0,0 +1,464 @@ +"""The node-field ADP wrapper: weights, positivity, cache correctness, lifecycle. + +Two properties here are load-bearing rather than cosmetic. The weights must be a +normalised mixture over each atom's candidate nodes, because that is what makes the +per-atom B a convex combination of positive node values and therefore positive without +a clamp. And the forward cache must notice that the coordinates moved: the wrapper reads +them through an injected accessor rather than a call argument, so +``CachedForwardMixin``'s own fingerprint cannot see them and +``DisorderFieldTensor._fingerprint_state`` has to fold them in. +""" + +import math + +import pytest +import torch + +from torchref.model.disorder_field import ( + DisorderFieldTensor, + build_neighbor_list, + farthest_point_anchors, +) +from torchref.model.parameter_wrappers import MixedTensor + + +@pytest.fixture +def coords(): + """A compact 3-D blob of atoms on a deterministic lattice.""" + g = torch.arange(6, dtype=torch.float64) + x, y, z = torch.meshgrid(g, g, g, indexing="ij") + return torch.stack([x.reshape(-1), y.reshape(-1), z.reshape(-1)], dim=1) * 1.7 + + +@pytest.fixture +def target_b(coords): + """A smooth B field with a real spatial gradient for the fit to chase.""" + return 20.0 + 1.5 * coords[:, 0] + 0.8 * coords[:, 2] + + +def _field(coords, target_b, **kw): + xyz = MixedTensor(coords.clone(), name="xyz") + kw.setdefault("n_nodes", 12) + kw.setdefault("k_neighbors", 6) + return DisorderFieldTensor( + initial_values=target_b, xyz_fn=xyz, dtype=torch.float64, **kw + ), xyz + + +@pytest.mark.unit +def test_anchor_selection_is_deterministic(coords): + """Same coordinates, same anchors -- no RNG anywhere in placement.""" + a = farthest_point_anchors(coords, 10) + b = farthest_point_anchors(coords, 10) + assert torch.equal(a, b) + assert a.shape[0] == 10 + assert a.dtype == torch.int64 + # Anchors are atom indices, and distinct. + assert int(a.max()) < coords.shape[0] + assert torch.unique(a).shape[0] == a.shape[0] + + +@pytest.mark.unit +def test_neighbor_list_is_nearest_first(coords): + """``build_neighbor_list`` returns each atom's k nearest nodes, closest first.""" + node_pos = coords[farthest_point_anchors(coords, 8)] + nl = build_neighbor_list(coords, node_pos, 4) + assert nl.shape == (coords.shape[0], 4) + d = torch.cdist(coords, node_pos) + gathered = torch.gather(d, 1, nl) + assert bool((gathered.diff(dim=1) >= -1e-12).all()), "not sorted by distance" + assert torch.equal(nl[:, 0], d.argmin(dim=1)) + + +@pytest.mark.unit +def test_weights_are_a_normalised_mixture(coords, target_b): + """Rows sum to one and are non-negative -- what makes the output positive.""" + field, _ = _field(coords, target_b) + W = field.weights() + assert W.shape == (coords.shape[0], 6) + assert bool((W >= 0).all()) + assert torch.allclose(W.sum(dim=1), torch.ones(coords.shape[0], dtype=W.dtype)) + + +@pytest.mark.unit +def test_output_is_per_atom_and_positive(coords, target_b): + """Public space is per-atom even though storage is per-node.""" + field, _ = _field(coords, target_b) + out = field() + assert out.shape == (coords.shape[0],) + assert field.shape == (coords.shape[0],) + assert field.node_shape == (12, 2) + assert bool((out > 0).all()) + assert bool(torch.isfinite(out).all()) + + +@pytest.mark.unit +def test_single_node_flat_kernel_is_a_constant(coords): + """K=1 reproduces a constant B exactly -- the analytic control. + + With one node every atom's weight vector is ``[1.0]`` whatever the distance, so the + field degenerates to a single scalar and must return it uniformly. + """ + b = torch.full((coords.shape[0],), 37.5, dtype=torch.float64) + field, _ = _field(coords, b, n_nodes=1, k_neighbors=1) + out = field() + assert torch.allclose(out, b, atol=1e-9), f"got spread {out.min()}..{out.max()}" + + +@pytest.mark.unit +def test_fit_tracks_a_smooth_gradient(coords, target_b): + """A 12-node field on a linear B ramp beats the best constant by a wide margin.""" + field, _ = _field(coords, target_b) + resid = (field() - target_b).pow(2).mean().sqrt() + constant = (target_b - target_b.mean()).pow(2).mean().sqrt() + assert resid < 0.25 * constant, f"rmse {resid:.3f} vs constant {constant:.3f}" + + +@pytest.mark.unit +def test_more_nodes_fit_better(coords, target_b): + """Reconstruction improves monotonically with node count on a smooth target.""" + errors = [] + for n in (2, 8, 32): + field, _ = _field(coords, target_b, n_nodes=n, k_neighbors=min(6, n)) + errors.append(float((field() - target_b).detach().pow(2).mean().sqrt())) + assert errors[0] > errors[1] > errors[2], errors + + +@pytest.mark.unit +def test_fingerprint_sees_the_coordinates(coords, target_b): + """The load-bearing test for the injected accessor, and it proves its own point. + + ``forward()`` takes no arguments, so the coordinates reach it through ``_xyz_fn``. + The inherited fingerprint covers only this module's own parameters and buffers, so + it cannot see them -- asserted directly here, by checking the base implementation + does NOT change while the override does. Without + ``DisorderFieldTensor._fingerprint_state`` the cache would serve a B computed at + coordinates that have since moved. + """ + from torchref.utils.caching import CachedForwardMixin + + field, xyz = _field(coords, target_b) + base_before = CachedForwardMixin._fingerprint_state(field) + full_before = field._fingerprint_state() + + with torch.no_grad(): + xyz.refinable_params[:8, 0] += 4.0 + + assert CachedForwardMixin._fingerprint_state(field) == base_before, ( + "the field's own parameters and buffers did not move, so the inherited " + "fingerprint is blind to this change -- which is why the override exists" + ) + assert field._fingerprint_state() != full_before, "override missed the coordinates" + + +@pytest.mark.unit +def test_cache_returns_a_fresh_value_after_a_non_rigid_move(coords, target_b): + """A change of relative geometry must reach the output, not a stale cache. + + The perturbation has to be non-rigid: the field is translation-invariant by + construction, so shifting every atom equally moves the nodes with them and is + *correctly* a no-op. + """ + field, xyz = _field(coords, target_b) + before = field().clone() + + with torch.no_grad(): + xyz.refinable_params[: coords.shape[0] // 2, 0] += 4.0 + + after = field() + assert not torch.allclose(before, after), "stale cache: coordinates were ignored" + + +@pytest.mark.unit +def test_rigid_translation_leaves_the_field_unchanged(coords, target_b): + """The invariance that makes the previous test need a non-rigid perturbation.""" + field, xyz = _field(coords, target_b) + before = field().clone() + + with torch.no_grad(): + xyz.refinable_params += 9.0 + + assert torch.allclose(field(), before, atol=1e-9) + + +@pytest.mark.unit +def test_node_positions_follow_the_coordinates(coords, target_b): + """Node positions are derived from the atoms, so a rigid shift carries them along.""" + field, xyz = _field(coords, target_b) + before = field.node_positions().clone() + + with torch.no_grad(): + xyz.refinable_params += 2.5 + + after = field.node_positions() + assert torch.allclose(after - before, torch.full_like(before, 2.5), atol=1e-9) + + +@pytest.mark.unit +def test_gradient_reaches_nodes_and_coordinates(coords, target_b): + """Both channels are live through the plain zero-arg forward path.""" + field, xyz = _field(coords, target_b) + field().sum().backward() + assert field.refinable_params.grad is not None + assert bool(field.refinable_params.grad.abs().sum() > 0) + assert xyz.refinable_params.grad is not None + assert bool(xyz.refinable_params.grad.abs().sum() > 0) + + +@pytest.mark.unit +def test_gradcheck_on_both_channels(coords, target_b): + """Analytic gradients match finite differences, for node values and coordinates. + + Checked through :meth:`DisorderFieldTensor.evaluate`, which is the field's + arithmetic without the accessor or the cache. Rebinding ``refinable_params`` inside + a gradcheck closure would not work: ``nn.Parameter(p)`` is a fresh leaf, so the + graph back to ``p`` is severed and there is nothing to check. + """ + small = coords[:20].clone() + field, _ = _field(small, target_b[:20], n_nodes=3, k_neighbors=3) + + raw = field.node_values().detach().clone().requires_grad_(True) + xyz_in = small.detach().clone().requires_grad_(True) + + assert torch.autograd.gradcheck( + field.evaluate, (xyz_in, raw), eps=1e-6, atol=1e-5 + ) + + +@pytest.mark.unit +def test_list_adequacy_invariant_is_exposed(coords, target_b): + """The smallest candidate weight is reported, and shrinks as k grows.""" + tight, _ = _field(coords, target_b, n_nodes=16, k_neighbors=2) + loose, _ = _field(coords, target_b, n_nodes=16, k_neighbors=12) + assert loose.smallest_candidate_weight() < tight.smallest_candidate_weight() + assert 0.0 <= loose.smallest_candidate_weight() <= 1.0 + + +@pytest.mark.unit +def test_node_load_is_in_node_space_and_conserves_total_weight(coords, target_b): + """Load must be scattered into node space, not summed over the candidate axis. + + ``weights()`` is ``(n_atoms, k)`` over CANDIDATES, so ``weights().sum(0)`` is a + length-k vector of per-slot totals with no meaning -- a trap worth pinning, because + it silently returns a plausible-looking tensor of the wrong length. + """ + field, _ = _field(coords, target_b, n_nodes=12, k_neighbors=6) + load = field.node_load() + + assert load.shape == (12,), "load must be per node, not per candidate slot" + assert field.weights().sum(dim=0).shape == (6,), "the trap this method avoids" + # Rows of W sum to 1, so the total load is exactly the atom count. + assert torch.allclose( + load.sum(), torch.tensor(float(coords.shape[0]), dtype=load.dtype) + ) + assert bool((load >= 0).all()) + + +@pytest.mark.unit +def test_rebuild_neighbor_list_is_explicit_and_refreshes(coords, target_b): + """Membership only changes when the caller asks; the rebuild then takes effect.""" + field, xyz = _field(coords, target_b) + original = field.neighbor_list.clone() + + with torch.no_grad(): + xyz.refinable_params[:, 0] += 40.0 # move atoms far past the nodes + + assert torch.equal(field.neighbor_list, original), "list changed without a rebuild" + field.rebuild_neighbor_list() + assert field.neighbor_list.shape == original.shape + assert bool(torch.isfinite(field()).all()) + + +@pytest.mark.unit +def test_atom_space_mask_collapses_to_node_space(coords, target_b): + """Masks arrive in ATOM space and are collapsed with OR onto the nodes.""" + field, _ = _field(coords, target_b) + n_atoms = coords.shape[0] + + mask = torch.zeros(n_atoms, dtype=torch.bool) + mask[0] = True + field.update_refinable_mask(mask) + + assert field.refinable_mask.shape == (field.n_nodes,) + served = field.neighbor_list[0] + assert bool(field.refinable_mask[served].all()), "atom 0's nodes must be refinable" + assert int(field.refinable_mask.sum()) == int(torch.unique(served).numel()) + + +@pytest.mark.unit +def test_wrong_sized_mask_is_rejected(coords, target_b): + """A node-space mask passed as atom space (or vice versa) is an error, not a guess.""" + field, _ = _field(coords, target_b) + with pytest.raises(ValueError, match="Atom-space mask"): + field.update_refinable_mask(torch.ones(field.n_nodes, dtype=torch.bool)) + with pytest.raises(ValueError, match="Node-space mask"): + field.update_refinable_mask( + torch.ones(coords.shape[0], dtype=torch.bool), in_node_space=True + ) + + +@pytest.mark.unit +def test_freezing_all_nodes_leaves_the_output_intact(coords, target_b): + """Repartitioning moves values between storage halves without changing them.""" + field, _ = _field(coords, target_b) + before = field().clone() + field.update_refinable_mask( + torch.zeros(field.n_nodes, dtype=torch.bool), in_node_space=True + ) + assert int(field.get_refinable_count()) == 0 + assert torch.allclose(field(), before, atol=1e-12) + + +@pytest.mark.unit +def test_per_atom_assignment_is_refused(coords, target_b): + """K nodes cannot represent arbitrary per-atom values, so writing is not silent.""" + field, _ = _field(coords, target_b) + with pytest.raises(NotImplementedError, match="per-atom assignment"): + field[0] = 42.0 + + +@pytest.mark.unit +def test_refit_moves_the_field_to_a_new_target(coords, target_b): + """``refit`` is the representable alternative to per-atom assignment.""" + field, _ = _field(coords, target_b) + new_target = target_b * 0.5 + 5.0 + field.refit(new_target) + resid = (field() - new_target).pow(2).mean().sqrt() + baseline = (new_target - new_target.mean()).pow(2).mean().sqrt() + assert resid < 0.25 * baseline + + +@pytest.mark.unit +def test_copy_is_independent_but_shares_the_accessor(coords, target_b): + """Parameters are copied; the coordinate accessor is deliberately shared. + + Deep-copying the accessor is what breaks ``Restraints.copy()`` today: ``deepcopy`` + walks the borrowed wrapper whose cache can hold a graph-attached tensor. + """ + field, xyz = _field(coords, target_b) + clone = field.copy() + + assert clone._xyz_fn.module is xyz, "accessor must reference the same wrapper" + assert clone.refinable_params.data_ptr() != field.refinable_params.data_ptr() + assert torch.allclose(clone(), field()) + + with torch.no_grad(): + clone.refinable_params[:, 0] += 1.0 + clone.reset_forward_cache() + field.reset_forward_cache() + assert not torch.allclose(clone(), field()), "copy is not independent" + + +@pytest.mark.unit +def test_copy_survives_an_evaluated_forward(coords, target_b): + """``copy()`` works after ``forward()`` has populated the accessor's cache.""" + field, _ = _field(coords, target_b) + field() # populate xyz's forward cache with a graph-attached tensor + clone = field.copy() + assert torch.allclose(clone(), field()) + + +@pytest.mark.unit +def test_state_dict_excludes_the_accessor_and_round_trips(coords, target_b): + """The callable is not state; everything else survives a save/load exactly. + + Restored the way ``Model.create_from_state_dict`` does it: build a wrapper with real + values to fix the shapes and masks, then let ``load_state_dict`` overwrite them. The + empty shell exists for shape-less construction, not as a load target. + """ + field, xyz = _field(coords, target_b) + # Cloned deliberately: state_dict() hands back detached REFERENCES, so a later + # in-place edit of refinable_params would silently rewrite the saved state too. + sd = {key: value.clone() for key, value in field.state_dict().items()} + assert not any("xyz_fn" in key for key in sd), sd.keys() + expected = field().clone() + + restored, _ = _field(coords, target_b) + with torch.no_grad(): + restored.refinable_params[:, 0] += 0.75 # move it off the saved state + restored.reset_forward_cache() + assert not torch.allclose(restored(), expected), "perturbation did not take" + + restored.load_state_dict(sd) + restored.reset_forward_cache() + + assert torch.allclose(restored(), expected, atol=1e-12) + + +@pytest.mark.unit +def test_empty_shell_needs_no_accessor(coords): + """The ``load_state_dict`` entry point constructs without coordinates.""" + # Pinned to CPU: every other case here inherits CPU from the tensors it is + # handed, but the empty shell resolves ``device=None`` to ``device.current``, + # and float64 does not exist on MPS. + shell = DisorderFieldTensor(dtype=torch.float64, device="cpu") + assert shell.neighbor_list is None + assert shell.shape == (0,) + + +@pytest.mark.unit +def test_accessor_is_required_when_values_are_given(coords, target_b): + """Nodes cannot be placed without coordinates, and that fails loudly.""" + with pytest.raises(ValueError, match="xyz_fn"): + DisorderFieldTensor(initial_values=target_b, dtype=torch.float64) + + +@pytest.mark.unit +def test_accessor_is_not_a_submodule_or_buffer(coords, target_b): + """It must stay out of module traversal, or a device move duplicates the graph.""" + field, xyz = _field(coords, target_b) + from torchref.utils.utils import ModuleReference + + assert isinstance(field._xyz_fn, ModuleReference) + assert xyz not in list(field.modules()) + assert not any(m is xyz for _, m in field.named_modules()) + # xyz's parameters must not appear among the field's own. + field_ptrs = {p.data_ptr() for p in field.parameters()} + assert xyz.refinable_params.data_ptr() not in field_ptrs + + +@pytest.mark.unit +def test_device_round_trip_keeps_buffers_together(coords, target_b): + """A ``.to()`` moves the node storage and the index buffers as one.""" + field, _ = _field(coords, target_b) + before = field().clone() + field.to(torch.device("cpu")) + assert field.neighbor_list.device.type == "cpu" + assert field.anchor_atom.device.type == "cpu" + assert torch.allclose(field(), before, atol=1e-12) + + +@pytest.mark.unit +def test_float32_and_float64_both_work(coords, target_b): + """The field works in either dtype without requiring one.""" + for dtype in (torch.float32, torch.float64): + xyz = MixedTensor(coords.to(dtype).clone(), name="xyz") + field = DisorderFieldTensor( + initial_values=target_b.to(dtype), + xyz_fn=xyz, + n_nodes=8, + k_neighbors=4, + dtype=dtype, + ) + out = field() + assert out.dtype == dtype + assert bool(torch.isfinite(out).all()) + + +@pytest.mark.unit +def test_ragged_anchor_neighbourhoods_average_their_atoms(coords, target_b): + """A node anchored on several atoms sits at their centroid.""" + xyz = MixedTensor(coords.clone(), name="xyz") + anchor_atom = torch.tensor([0, 1, 2, 10, 11], dtype=torch.int64) + anchor_node = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int64) + field = DisorderFieldTensor( + initial_values=target_b, + xyz_fn=xyz, + k_neighbors=2, + anchor_rows=(anchor_atom, anchor_node), + dtype=torch.float64, + ) + pos = field.node_positions() + assert pos.shape == (2, 3) + assert torch.allclose(pos[0], coords[[0, 1, 2]].mean(dim=0)) + assert torch.allclose(pos[1], coords[[10, 11]].mean(dim=0)) diff --git a/tests/unit/model/test_hydrogen_default.py b/tests/unit/model/test_hydrogen_default.py new file mode 100644 index 00000000..5203331c --- /dev/null +++ b/tests/unit/model/test_hydrogen_default.py @@ -0,0 +1,304 @@ +"""Keep deposited hydrogens by default and generate missing ones only on request. + +The interesting cases are the partially-hydrogenated file, which has to be topped up per +parent rather than left alone, and the per-atom buffers that are cached lazily and go +stale the moment the atom set grows. +""" + +from pathlib import Path + +import numpy as np +import pytest + +from torchref.model.model import Model +from torchref.model.model_ft import ModelFT + + +def _elements(model): + return model.pdb["element"].astype(str).str.strip().values + + +def _counts(model): + elements = _elements(model) + n_h = int((elements == "H").sum()) + return len(model.pdb), n_h + + +@pytest.mark.unit +@pytest.mark.parametrize("model_class", [Model, ModelFT]) +@pytest.mark.parametrize("filename", ["1DAW.pdb", "1AK5_with_H.pdb", "7L84.pdb"]) +def test_default_preserves_deposited_atoms( + pdb_dir: Path, + filename: str, + monkeypatch: pytest.MonkeyPatch, + model_class: type[Model], +) -> None: + """Default loading neither generates hydrogens nor removes deposited ones.""" + from torchref.io.pdb import PDBReader + + path = pdb_dir / filename + deposited, _, _ = PDBReader(verbose=0).read(str(path))() + + def unexpected_generation(self: Model) -> None: + pytest.fail("Default loading must not generate hydrogens") + + monkeypatch.setattr(Model, "_add_missing_hydrogens", unexpected_generation) + model = model_class(verbose=0).load_pdb(str(path)) + + np.testing.assert_array_equal(_elements(model), deposited["element"].str.strip()) + + +@pytest.mark.unit +def test_default_cif_load_does_not_generate_hydrogens( + cif_dir: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """mmCIF loading also leaves missing hydrogens absent by default.""" + + def unexpected_generation(self: Model) -> None: + pytest.fail("Default mmCIF loading must not generate hydrogens") + + monkeypatch.setattr(Model, "_add_missing_hydrogens", unexpected_generation) + model = Model(verbose=0).load_cif(str(cif_dir / "1DAW.cif")) + total, n_h = _counts(model) + assert total > 0 + assert n_h == 0 + + +@pytest.mark.unit +def test_context_defaults_to_no_hydrogen_generation() -> None: + """A standalone model context leaves hydrogen generation disabled.""" + from torchref.model.context import ModelContext + + assert ModelContext().add_hydrogens is False + + +@pytest.mark.unit +def test_a_file_without_hydrogens_gets_them(pdb_dir): + """1DAW ships none, so every hydrogen here is generated.""" + model = Model(verbose=0, add_hydrogens=True) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + + total, n_h = _counts(model) + assert n_h > 0 + heavy = total - n_h + assert ( + 0.7 < n_h / heavy < 1.3 + ), f"{n_h} hydrogens on {heavy} heavy atoms is not a plausible ratio" + + +@pytest.mark.unit +def test_a_partially_hydrogenated_file_is_topped_up(pdb_dir): + """1AK5 ships 675 hydrogens on 2582 heavy atoms, where full is roughly 2500. + + Generation is decided per parent -- the plan proposes only a hydrogen the template + names and the model lacks -- so a file that already has some still gets the rest. A + does-the-table-contain-any test would have left this structure as deposited. + """ + kept = Model(verbose=0, add_hydrogens=False) + kept.load_pdb(str(pdb_dir / "1AK5_with_H.pdb")) + _, n_kept = _counts(kept) + + topped = Model(verbose=0, add_hydrogens=True) + topped.load_pdb(str(pdb_dir / "1AK5_with_H.pdb")) + _, n_topped = _counts(topped) + + assert n_kept > 0, "1AK5_with_H is supposed to ship some hydrogens" + assert ( + n_topped > n_kept * 2 + ), f"only {n_topped} hydrogens after top-up, against {n_kept} in the file" + + +@pytest.mark.unit +def test_strip_H_still_removes_everything(pdb_dir): + """The opt-out is unaffected: no hydrogen survives, generated or deposited.""" + for name in ("1DAW.pdb", "7L84.pdb"): + model = Model(verbose=0, strip_H=True, add_hydrogens=True) + model.load_pdb(str(pdb_dir / name)) + _, n_h = _counts(model) + assert n_h == 0, f"{name} kept {n_h} hydrogens under strip_H" + + +@pytest.mark.unit +def test_add_hydrogens_false_keeps_the_file_as_it_is(pdb_dir): + """Generation off, stripping off: exactly what the reader produced.""" + model = Model(verbose=0, add_hydrogens=False) + model.load_pdb(str(pdb_dir / "7L84.pdb")) + total, n_h = _counts(model) + assert n_h > 0, "7L84 ships hydrogens, so they should have been kept" + + generated = Model(verbose=0, add_hydrogens=True) + generated.load_pdb(str(pdb_dir / "7L84.pdb")) + assert _counts(generated)[0] >= total + + +@pytest.mark.unit +def test_per_atom_buffers_are_rebuilt_for_the_new_atom_set(pdb_dir): + """Every lazily-cached per-atom buffer matches the table after generation. + + These are guarded by ``hasattr`` and returned as-is once built, which was safe only + while an atom-set change always produced a fresh model. Generating hydrogens in + place left the van der Waals radii at the heavy-atom count while the pair list + indexed the full set, and the non-bonded build raised ``IndexError``. + """ + model = Model(verbose=0, add_hydrogens=True) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + n_atoms = len(model.pdb) + + assert model.get_vdw_radii().shape[0] == n_atoms + assert model.Z.shape[0] == n_atoms + + radii = model.get_vdw_radii().detach().cpu().numpy() + assert np.isfinite(radii).all() + is_h = _elements(model) == "H" + assert is_h.any() + assert np.allclose(radii[is_h], 1.20), "hydrogens did not get a hydrogen radius" + + +@pytest.mark.unit +def test_restraints_build_over_the_hydrogenated_model(pdb_dir): + """Restraints cover the hydrogens, and each carries exactly one bond.""" + model = Model(verbose=0, add_hydrogens=True) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + restraints = model.restraints + + elements = _elements(model) + is_h = elements == "H" + assert is_h.any() + bonds = restraints.restraints["bond"]["all"]["indices"].cpu().numpy() + involves_h = is_h[bonds[:, 0]] | is_h[bonds[:, 1]] + assert int(involves_h.sum()) == int(is_h.sum()) + + vdw = restraints.restraints["vdw"]["indices"] + assert int(vdw.max()) < len( + model.pdb + ), "the non-bonded pair list indexes past the end of the atom table" + + +@pytest.mark.unit +def test_riding_hydrogens_are_not_placed_when_real_ones_exist(pdb_dir): + """The riding stand-in goes quiet once the model carries hydrogens. + + Riding hydrogens approximate the sterics of hydrogens the model does not have. + Placing them alongside real ones would put phantom atoms in the structure that push + real ones around -- and they would not even be the hydrogens the generator declined, + because the riding builder counts bonded neighbours by distance while the generator + reads them off the bond graph. + """ + model = Model(verbose=0, add_hydrogens=True) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + restraints = model.restraints + + assert restraints.h_topo is not None + assert restraints.h_topo.n_hydrogens == 0 + + stripped = Model(verbose=0, strip_H=True) + stripped.load_pdb(str(pdb_dir / "1DAW.pdb")) + assert ( + stripped.restraints.h_topo.n_hydrogens > 0 + ), "with hydrogens absent the riding stand-in should still be built" + + +# --- The user's restraint dictionary is the one that hydrogenates ------------------ + +RENAMED_GLU_H = {"HAX", "HBX", "HBY", "HGX", "HGY"} + + +@pytest.fixture +def renamed_glu_cif(test_files_dir): + """A GLU dictionary whose side-chain hydrogens carry names the library lacks.""" + return str(test_files_dir / "restraints" / "GLU_renamed.cif") + + +def _glu_hydrogen_names(model): + pdb = model.pdb + glu_h = (pdb["resname"].astype(str).str.strip() == "GLU") & ( + pdb["element"].astype(str).str.strip() == "H" + ) + return set(pdb.loc[glu_h, "name"].astype(str).str.strip()) + + +@pytest.mark.unit +def test_generation_reads_the_cif_given_at_construction(pdb_dir, renamed_glu_cif): + """A dictionary passed to the constructor overrides the library for generation. + + The names prove which dictionary was read, and the bond degree proves the generated + hydrogens are the ones the restraints know: a hydrogen generated from one + dictionary and restrained by another has no bond edge at all. + """ + model = Model(verbose=0, add_hydrogens=True, cif_path=renamed_glu_cif) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + assert model.ctx.cif_path == renamed_glu_cif + names = _glu_hydrogen_names(model) + assert RENAMED_GLU_H <= names + assert not {"HA", "HB2", "HB3", "HG2", "HG3"} & names + + atoms = model.restraints.topology.atoms + is_h = atoms.is_hydrogen.cpu().numpy() + degree = atoms.degree().cpu().numpy() + assert (degree[is_h] > 0).all(), "generated hydrogens without a bond restraint" + + +@pytest.mark.unit +def test_derived_models_keep_the_restraint_cif(pdb_dir, renamed_glu_cif): + """hydrogenate, strip_hydrogens and select all carry the dictionary along.""" + model = Model(verbose=0, add_hydrogens=False, cif_path=renamed_glu_cif) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + + hydrogenated = model.hydrogenate() + assert hydrogenated.ctx.cif_path == renamed_glu_cif + assert RENAMED_GLU_H <= _glu_hydrogen_names(hydrogenated) + assert "GLU" in hydrogenated.restraints.cif_dict + template_h = set( + hydrogenated.restraints.cif_dict["GLU"]["atoms"]["atom_id"].astype(str).str.strip() + ) + assert RENAMED_GLU_H <= template_h + + assert hydrogenated.strip_hydrogens().ctx.cif_path == renamed_glu_cif + assert model.select("resname GLU").ctx.cif_path == renamed_glu_cif + + +@pytest.mark.unit +def test_state_dict_round_trips_the_restraint_cif(pdb_dir, renamed_glu_cif): + model = Model(verbose=0, cif_path=renamed_glu_cif) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + restored = Model.create_from_state_dict(model.state_dict(), verbose=0) + assert restored.ctx.cif_path == renamed_glu_cif + + +@pytest.mark.unit +def test_load_model_registers_the_cif_before_loading(pdb_dir, renamed_glu_cif): + """The shared CLI loader generates from the user dictionary too.""" + from torchref.cli._common import load_model + + model = load_model( + str(pdb_dir / "1DAW.pdb"), verbose=0, cif=renamed_glu_cif, add_hydrogens=True + ) + assert model.ctx.cif_path == renamed_glu_cif + assert RENAMED_GLU_H <= _glu_hydrogen_names(model) + + +@pytest.mark.unit +def test_generation_reads_every_compound_of_a_multi_block_cif(pdb_dir, test_files_dir): + """A dictionary with several ``data_comp_`` blocks hydrogenates each of its compounds. + + Multi-compound dictionaries once restrained only their last block; generation now + reads the same dictionary, so both renamed sets must appear and every generated + hydrogen must carry a bond edge. + """ + cif = test_files_dir / "restraints" / "GLU_ASP_renamed.cif" + blocks = [l for l in cif.read_text().splitlines() if l.startswith("data_comp_")] + assert len(blocks) == 3, blocks # comp_list + GLU + ASP: the fixture is really multi-block + + model = Model(verbose=0, add_hydrogens=True, cif_path=str(cif)) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + pdb = model.pdb + is_h = pdb["element"].astype(str).str.strip() == "H" + resname = pdb["resname"].astype(str).str.strip() + names = pdb["name"].astype(str).str.strip() + assert RENAMED_GLU_H <= set(names[is_h & (resname == "GLU")]) + assert {"HBQ", "HBR"} <= set(names[is_h & (resname == "ASP")]) + assert not {"HB2", "HB3"} & set(names[is_h & (resname == "ASP")]) + + atoms = model.restraints.topology.atoms + degree = atoms.degree().cpu().numpy() + assert (degree[atoms.is_hydrogen.cpu().numpy()] > 0).all() diff --git a/tests/unit/model/test_hydrogen_mode.py b/tests/unit/model/test_hydrogen_mode.py new file mode 100644 index 00000000..c7676a5a --- /dev/null +++ b/tests/unit/model/test_hydrogen_mode.py @@ -0,0 +1,109 @@ +"""Switching a loaded model between riding and free hydrogens. + +The switch replaces the coordinate wrapper and nothing else: coordinates are +unchanged, the heavy-atom refinable set carries over, hydrogens leave or rejoin the +refinable set, the restraints keep reading the live wrapper, and the mode survives +copies, selections and state dicts. +""" + +import pytest +import torch + +from torchref.model.model import Model +from torchref.model.model_ft import ModelFT +from torchref.model.riding_xyz import RidingXYZTensor + + +@pytest.fixture +def free_model(pdb_dir): + model = Model(verbose=0, add_hydrogens=True) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + return model + + +def _n_h(model): + return int((model.pdb["element"].str.strip() == "H").sum()) + + +@pytest.mark.unit +def test_switch_to_riding_keeps_coordinates_and_drops_hydrogen_parameters(free_model): + model = free_model + before = model.xyz().detach().clone() + n_free = model.parameters_of_types(("xyz",))[0].shape[0] + model.set_hydrogen_mode("riding") + assert model.hydrogen_mode == "riding" + assert isinstance(model.xyz, RidingXYZTensor) + assert torch.allclose(model.xyz(), before, atol=1e-4) + assert model.parameters_of_types(("xyz",))[0].shape[0] == n_free - _n_h(model) + assert model.xyz.n_hydrogens == _n_h(model) + + +@pytest.mark.unit +def test_switch_back_to_free_restores_per_atom_wrapper(free_model): + model = free_model + model.set_hydrogen_mode("riding") + model.set_hydrogen_mode("free") + assert model.hydrogen_mode == "free" + assert not isinstance(model.xyz, RidingXYZTensor) + assert model.parameters_of_types(("xyz",))[0].shape[0] == len(model.pdb) + + +@pytest.mark.unit +def test_restraints_read_the_installed_wrapper(free_model): + model = free_model + restraints = model.restraints + model.set_hydrogen_mode("riding") + assert restraints._xyz_fn is model.xyz + with torch.no_grad(): + model.xyz.refinable_params.add_(0.1) + assert torch.equal(restraints.xyz(), model.xyz()) + + +@pytest.mark.unit +def test_frozen_heavy_atoms_stay_frozen_across_the_switch(free_model): + model = free_model + mask = torch.zeros(len(model.pdb), dtype=torch.bool, device=model.device) + mask[: len(model.pdb) // 2] = True + model.xyz.update_refinable_mask(mask) + model.set_hydrogen_mode("riding") + full = model.xyz.full_refinable_mask + heavy = ~torch.as_tensor((model.pdb["element"].str.strip() == "H").values, device=full.device) + assert torch.equal(full[heavy], mask[heavy]) + + +@pytest.mark.unit +def test_mode_survives_copy_select_and_shake(free_model): + model = free_model + model.set_hydrogen_mode("riding") + dup = model.copy() + assert dup.hydrogen_mode == "riding" and isinstance(dup.xyz, RidingXYZTensor) + assert torch.equal(dup.xyz(), model.xyz()) + sub = model.select("resseq 10:40") + assert isinstance(sub.xyz, RidingXYZTensor) + assert sub.xyz.shape[0] == len(sub.pdb) + model.shake_coords(0.05) + assert isinstance(model.xyz, RidingXYZTensor) + assert model.xyz.shape[0] == len(model.pdb) + + +@pytest.mark.unit +@pytest.mark.parametrize("model_class", [Model, ModelFT]) +def test_riding_mode_round_trips_through_state_dict(pdb_dir, model_class): + kwargs = {"max_res": 3.0} if model_class is ModelFT else {} + model = model_class(verbose=0, add_hydrogens=True, **kwargs) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + model.set_hydrogen_mode("riding") + with torch.no_grad(): + model.xyz.refinable_params[0].add_(0.25) + state = model.state_dict() + restored = model_class.create_from_state_dict(state, device=model.device) + assert restored.hydrogen_mode == "riding" + assert isinstance(restored.xyz, RidingXYZTensor) + assert torch.allclose(restored.xyz(), model.xyz(), atol=1e-5) + assert restored.xyz.get_refinable_count() == model.xyz.get_refinable_count() + + +@pytest.mark.unit +def test_none_mode_is_refused(free_model): + with pytest.raises(ValueError): + free_model.set_hydrogen_mode("none") diff --git a/tests/unit/model/test_hydrogens_in_xray.py b/tests/unit/model/test_hydrogens_in_xray.py new file mode 100644 index 00000000..ff709a7d --- /dev/null +++ b/tests/unit/model/test_hydrogens_in_xray.py @@ -0,0 +1,114 @@ +"""The ``hydrogens_in_xray`` flag: what it gates, and where it has to survive. + +It decides only whether hydrogen rows reach the structure-factor gathers. Restraints +never consult it, the solvent mask never includes hydrogens, and the setting must +follow the model through copies, selections, strips and state dicts. +""" + +import numpy as np +import pytest +import torch + +from torchref.model.context import ModelContext +from torchref.model.model import Model +from torchref.model.model_ft import ModelFT + + +@pytest.fixture(scope="module") +def with_hydrogens(pdb_dir): + """7L84 keeps its deposited hydrogens.""" + model = Model(verbose=0) + model.load_pdb(str(pdb_dir / "7L84.pdb")) + return model + + +def _n_h(model): + return int((model.pdb["element"].str.strip() == "H").sum()) + + +@pytest.mark.unit +def test_default_is_on(): + assert ModelContext().hydrogens_in_xray is True + assert Model(verbose=0).hydrogens_in_xray is True + + +@pytest.mark.unit +def test_partition_covers_hydrogens_only_when_on(with_hydrogens): + """Off drops exactly the hydrogen rows from the isotropic gather.""" + model = with_hydrogens + n_atoms, n_h = len(model.pdb), _n_h(model) + assert n_h > 0 + + def n_in_fcalc(): + return model.get_iso()[0].shape[0] + model.get_aniso()[0].shape[0] + + model.hydrogens_in_xray = True + assert n_in_fcalc() == n_atoms + model.hydrogens_in_xray = False + assert n_in_fcalc() == n_atoms - n_h + model.hydrogens_in_xray = True + assert n_in_fcalc() == n_atoms + + +@pytest.mark.unit +def test_off_matches_a_stripped_model(pdb_dir): + """Excluding hydrogens from Fcalc equals computing Fcalc without them.""" + full = ModelFT(verbose=0, max_res=2.5, hydrogens_in_xray=False) + full.load_pdb(str(pdb_dir / "7L84.pdb")) + heavy = ModelFT(verbose=0, max_res=2.5, strip_H=True) + heavy.load_pdb(str(pdb_dir / "7L84.pdb")) + grid = torch.arange(-3, 4) + hkl = torch.cartesian_prod(grid, grid, grid) + hkl = hkl[(hkl != 0).any(dim=1)].to(full.device) + with torch.no_grad(): + f_full = full(hkl) + f_heavy = heavy(hkl.to(heavy.device)) + assert torch.allclose(f_full, f_heavy, rtol=1e-4, atol=1e-3) + + +@pytest.mark.unit +def test_setting_survives_copy_select_and_strip(with_hydrogens): + model = with_hydrogens.copy() + model.hydrogens_in_xray = False + assert model.copy().hydrogens_in_xray is False + assert model.select("chain A").hydrogens_in_xray is False + assert model.strip_hydrogens().hydrogens_in_xray is False + + +@pytest.mark.unit +@pytest.mark.parametrize("model_class", [Model, ModelFT]) +def test_setting_round_trips_through_state_dict(pdb_dir, model_class): + kwargs = {"max_res": 3.0} if model_class is ModelFT else {} + model = model_class(verbose=0, hydrogens_in_xray=False, **kwargs) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + state = model.state_dict() + restored = model_class.create_from_state_dict(state, device=model.device) + assert restored.hydrogens_in_xray is False + state.pop("hydrogens_in_xray", None) + legacy = model_class.create_from_state_dict(state, device=model.device) + assert legacy.hydrogens_in_xray is True + + +@pytest.mark.unit +def test_deprecated_alias_is_inverted_and_warns(): + model = Model(verbose=0) + with pytest.warns(DeprecationWarning): + model.exclude_H_from_sf = True + assert model.hydrogens_in_xray is False + with pytest.warns(DeprecationWarning): + assert model.exclude_H_from_sf is True + + +@pytest.mark.unit +def test_solvent_mask_ignores_hydrogens(pdb_dir): + """The bulk-solvent mask is the same with and without hydrogen rows.""" + from torchref.scaling.solvent import SolventModel + + full = ModelFT(verbose=0, max_res=2.5) + full.load_pdb(str(pdb_dir / "7L84.pdb")) + heavy = ModelFT(verbose=0, max_res=2.5, strip_H=True) + heavy.load_pdb(str(pdb_dir / "7L84.pdb")) + assert _n_h(full) > 0 and _n_h(heavy) == 0 + mask_full = SolventModel(full, verbose=0).get_solvent_mask() + mask_heavy = SolventModel(heavy, verbose=0).get_solvent_mask() + assert torch.equal(mask_full, mask_heavy) diff --git a/tests/unit/model/test_mode_covariance_payload.py b/tests/unit/model/test_mode_covariance_payload.py new file mode 100644 index 00000000..db133145 --- /dev/null +++ b/tests/unit/model/test_mode_covariance_payload.py @@ -0,0 +1,336 @@ +"""The mode-covariance node payload, of which TLS is one member. + +A node stores the covariance of the displacement modes it represents rather than an ADP, +so the ADP an atom receives depends on where it sits inside the node's region. Two +properties have to hold before any of it is worth measuring: the rigid mode set must +reproduce the textbook TLS expression exactly, and every mode set must stay +positive-semidefinite at every displacement, because an indefinite U makes the +structure-factor FFT return NaN. +""" + +import math + +import pytest +import torch + +from torchref.model.disorder_field import ( + MODE_SETS, + AnisotropicPayload, + ModeCovariancePayload, +) +from torchref.model.parameter_wrappers import ( + chol_param_count, + psd_to_raw, + raw6_to_u6, + raw_to_cholesky, + u6_to_matrix, +) + +DTYPE = torch.float64 + + +def _u6_to_mat(u6): + return u6_to_matrix(u6) + + +def _random_sigma(q, k=1, scale=0.05, seed=0): + """A random PD ``(k, q, q)`` covariance.""" + g = torch.Generator().manual_seed(seed) + A = torch.randn(k, q, q, generator=g, dtype=DTYPE) * scale + return A @ A.transpose(-1, -2) + 1e-3 * torch.eye(q, dtype=DTYPE) + + +# ---------------------------------------------------------------------------------- +# Sizes. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.unit +@pytest.mark.parametrize( + "mode_set,q,width", + [("constant", 3, 6), ("rigid", 6, 21), ("rigid_dilation", 7, 28), ("affine", 12, 78)], +) +def test_mode_set_sizes(mode_set, q, width): + """The ladder is 6 / 21 / 28 / 78 parameters; TLS is the 21 (20 determinable).""" + p = ModeCovariancePayload(mode_set) + assert p.q == q + assert p.width == width == chol_param_count(q) + assert p.out_width == 6 + + +@pytest.mark.unit +def test_unknown_mode_set_is_rejected(): + with pytest.raises(ValueError, match="Unknown mode set"): + ModeCovariancePayload("librational_whimsy") + + +@pytest.mark.unit +def test_mode_sets_are_nested(): + """Each rung must contain the one below, or the ladder is not a ladder.""" + keys = ["constant", "rigid", "rigid_dilation", "affine"] + for lo, hi in zip(keys, keys[1:]): + assert set(MODE_SETS[lo]).issubset(set(MODE_SETS[hi])) + assert ModeCovariancePayload(lo).q < ModeCovariancePayload(hi).q + + +# ---------------------------------------------------------------------------------- +# The TLS identity. If this fails nothing downstream is trustworthy. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_rigid_mode_set_is_textbook_tls(): + """``Psi Sigma Psi^T`` with translations and rotations IS ``T + AS + S^T A^T + A L A^T``. + + ``A`` is the matrix whose columns are ``e_i x r``, so ``A lambda = lambda x r``. With + ``Sigma`` blocked as ``[[T, S^T], [S, L]]`` over ``c = (t, lambda)`` the expansion is + the classical TLS expression, which is the claim that makes this payload a + generalisation of TLS rather than something merely similar. + """ + payload = ModeCovariancePayload("rigid") + sigma = _random_sigma(6, k=1, seed=3)[0] + T, L, S = sigma[:3, :3], sigma[3:, 3:], sigma[3:, :3] + + torch.manual_seed(11) + for r in torch.randn(20, 3, dtype=DTYPE) * 5.0: + # Columns of A are e_i x r. + A = torch.stack( + [ + torch.tensor([0.0, -r[2], r[1]], dtype=DTYPE), + torch.tensor([r[2], 0.0, -r[0]], dtype=DTYPE), + torch.tensor([-r[1], r[0], 0.0], dtype=DTYPE), + ], + dim=1, + ) + expected = T + A @ S + (A @ S).T + A @ L @ A.T + + Psi = payload.modes(r) # (3, 6) + got = Psi @ sigma @ Psi.T + assert torch.allclose(got, expected, atol=1e-10), f"r={r.tolist()}" + + +@pytest.mark.unit +def test_rotation_generators_are_antisymmetric(): + """``A`` must be antisymmetric, which is what makes the rigid set a rigid motion.""" + payload = ModeCovariancePayload("rigid") + r = torch.tensor([1.3, -2.1, 0.7], dtype=DTYPE) + A = payload.modes(r)[:, 3:] + assert torch.allclose(A, -A.T, atol=1e-12) + + +@pytest.mark.unit +def test_trace_s_is_the_flat_direction(): + """TLS has exactly one unobservable combination: adding to ``tr S`` must be free. + + Shifting ``S -> S + c I`` leaves ``U(r)`` unchanged for every ``r``, because + ``A(cI) + (A cI)^T = c(A + A^T) = 0`` for antisymmetric ``A``. + """ + payload = ModeCovariancePayload("rigid") + sigma = _random_sigma(6, k=1, seed=5)[0] + shifted = sigma.clone() + shifted[3:, :3] += 0.01 * torch.eye(3, dtype=DTYPE) + shifted[:3, 3:] += 0.01 * torch.eye(3, dtype=DTYPE) + + torch.manual_seed(2) + for r in torch.randn(10, 3, dtype=DTYPE) * 4.0: + Psi = payload.modes(r) + assert torch.allclose(Psi @ sigma @ Psi.T, Psi @ shifted @ Psi.T, atol=1e-12) + + +# ---------------------------------------------------------------------------------- +# Positive-semidefiniteness, the property the whole construction exists for. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.unit +@pytest.mark.parametrize("mode_set", sorted(MODE_SETS)) +def test_psd_at_extreme_displacement(mode_set): + """PSD for any ``Sigma`` and any ``r``, including far outside the node's region. + + This is what an arbitrary polynomial in ``r`` cannot promise: it goes indefinite + somewhere, and somewhere is the edge of the region where the weights have not yet + decayed. + """ + payload = ModeCovariancePayload(mode_set) + torch.manual_seed(7) + raw = torch.randn(4, payload.width, dtype=DTYPE) * 3.0 + L = raw_to_cholesky(raw, payload.q, payload.epsilon) + sigma = L @ L.transpose(-1, -2) + + for scale in (0.0, 1e-3, 1.0, 50.0, 1000.0): + r = torch.randn(6, 3, dtype=DTYPE) * scale + Psi = payload.modes(r) # (6, 3, q) + U = Psi @ sigma[:, None] @ Psi.transpose(-1, -2)[None] # (4, 6, 3, 3) + ev = torch.linalg.eigvalsh(U) + assert float(ev.min()) >= -1e-9, f"{mode_set} indefinite at |r|~{scale}" + + +@pytest.mark.unit +@pytest.mark.parametrize("scale", [0.1, 1.0, 3.0, 6.0]) +def test_sigma_is_pd_for_any_parameter_value(scale): + """Cholesky storage: no parameter value can make the covariance indefinite. + + Judged against the matrix norm, not against zero. The diagonal is ``exp(x)``, so a + wide spread of parameters gives ``Sigma`` a huge dynamic range and ``eigvalsh`` + returns the small eigenvalues with an error set by the large ones -- a float64 + property of the eigensolver, not of the parametrisation. + """ + payload = ModeCovariancePayload("affine") + torch.manual_seed(1) + raw = torch.randn(8, payload.width, dtype=DTYPE) * scale + ev = torch.linalg.eigvalsh(payload.sigma(raw)) + tol = 1e-10 * ev.abs().max(dim=-1, keepdim=True).values + assert bool((ev > -tol).all()), f"min eigenvalue {float(ev.min()):.3e} at scale {scale}" + + +# ---------------------------------------------------------------------------------- +# Agreement with the payload it generalises. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_constant_mode_set_matches_anisotropic_payload(): + """``"constant"`` is the same model as :class:`AnisotropicPayload`. + + Both store a 3x3 Cholesky factor and hand every atom the same U, so given the same + raw parameters they must produce the same tensor. That makes the constant rung the + inertness guard: a change here that moved it would be a change to the existing + anisotropic field. + """ + eps = 1e-3 + mode = ModeCovariancePayload("constant", epsilon=eps) + aniso = AnisotropicPayload(epsilon=eps) + torch.manual_seed(4) + raw = torch.randn(5, 6, dtype=DTYPE) + + xyz = torch.randn(9, 3, dtype=DTYPE) * 10.0 + node_pos = torch.randn(5, 3, dtype=DTYPE) * 10.0 + nl = torch.randint(0, 5, (9, 3)) + + got = mode.contributions(raw, xyz, node_pos, nl) + expected = aniso.contributions(raw, xyz, node_pos, nl) + assert got.shape == expected.shape == (9, 3, 6) + assert torch.allclose(got, expected, atol=1e-12) + + +@pytest.mark.unit +def test_raw_to_cholesky_matches_the_unrolled_three_by_three(): + """The general helper must agree with the unrolled 3x3 pair it generalises.""" + eps = 1e-3 + torch.manual_seed(6) + raw = torch.randn(12, 6, dtype=DTYPE) + L = raw_to_cholesky(raw, 3, eps) + got = L @ L.transpose(-1, -2) + expected = _u6_to_mat(raw6_to_u6(raw, eps)) + assert torch.allclose(got, expected, atol=1e-12) + + +@pytest.mark.unit +def test_cholesky_round_trip(): + """``psd_to_raw`` inverts ``raw_to_cholesky`` for a PD matrix.""" + eps = 1e-4 + for q in (3, 6, 12): + sigma = _random_sigma(q, k=5, scale=0.3, seed=q) + raw = psd_to_raw(sigma, eps) + L = raw_to_cholesky(raw, q, eps) + assert torch.allclose(L @ L.transpose(-1, -2), sigma, atol=1e-8), f"q={q}" + + +# ---------------------------------------------------------------------------------- +# Fit and magnitude. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.unit +@pytest.mark.parametrize("mode_set", sorted(MODE_SETS)) +def test_fit_seeds_translations_and_floors_the_rest(mode_set): + """Entering the parametrisation must start as the equivalent constant-U field.""" + payload = ModeCovariancePayload(mode_set, epsilon=1e-3) + torch.manual_seed(8) + n_atoms, k_nodes = 40, 4 + xyz = torch.randn(n_atoms, 3, dtype=DTYPE) * 8.0 + node_pos = torch.randn(k_nodes, 3, dtype=DTYPE) * 8.0 + nl = torch.randint(0, k_nodes, (n_atoms, 2)) + w_dense = torch.rand(n_atoms, k_nodes, dtype=DTYPE) + w_dense = w_dense / w_dense.sum(dim=1, keepdim=True) + target_b = 10.0 + 20.0 * torch.rand(n_atoms, dtype=DTYPE) + + raw = payload.fit(target_b, w_dense, 1e-3, xyz, node_pos, nl) + assert raw.shape == (k_nodes, payload.width) + sigma = payload.sigma(raw) + + # Translation block carries the fit; every gradient mode sits at the floor. + assert float(sigma[:, :3, :3].diagonal(dim1=-2, dim2=-1).min()) > 1e-4 + if payload.q > 3: + grad = sigma[:, 3:, 3:] + assert float(grad.abs().max()) < 1e-5, "gradient modes did not start at the floor" + + +@pytest.mark.unit +def test_log_magnitude_is_b_eq_of_the_translation_block(): + payload = ModeCovariancePayload("affine") + torch.manual_seed(9) + raw = torch.randn(6, payload.width, dtype=DTYPE) * 0.5 + T = payload.sigma(raw)[:, :3, :3] + expected = torch.log( + ((8.0 * math.pi**2 / 3.0) * T.diagonal(dim1=-2, dim2=-1).sum(-1)).clamp(min=1e-6) + ) + assert torch.allclose(payload.log_magnitude(raw), expected, atol=1e-12) + + +# ---------------------------------------------------------------------------------- +# Gradients. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.unit +@pytest.mark.parametrize("mode_set", ["rigid", "affine"]) +def test_gradcheck_contributions(mode_set): + """Gradient w.r.t. the node parameters and the coordinates. + + ``node_pos`` is held constant here rather than differentiated: the conditioning + length scale is a detached median over node positions, so gradcheck's numerical + derivative would pick up a path the analytic one deliberately cuts. That cut is the + subject of :func:`test_length_scale_carries_no_gradient`; the gradient that reaches + a node's position through the displacement ``r`` is covered here by ``xyz``, which + enters the same subtraction with the opposite sign. + """ + payload = ModeCovariancePayload(mode_set) + torch.manual_seed(10) + xyz = (torch.randn(7, 3, dtype=DTYPE) * 4.0).requires_grad_(True) + node_pos = torch.randn(3, 3, dtype=DTYPE) * 4.0 + nl = torch.randint(0, 3, (7, 2)) + raw = (torch.randn(3, payload.width, dtype=DTYPE) * 0.3).requires_grad_(True) + + assert torch.autograd.gradcheck( + lambda p_, x: payload.contributions(p_, x, node_pos, nl), + (raw, xyz), + eps=1e-6, + atol=1e-7, + ) + + +@pytest.mark.unit +def test_length_scale_carries_no_gradient(): + """The conditioning length scale is a stop-gradient, on purpose. + + It is the median nearest-neighbour node distance, whose true derivative is supported + on whichever single node pair happens to be at the median -- an artifact of the + layout, not a direction any optimiser should follow. Cutting it also stops the + optimiser rescaling its own modes by spreading the nodes out. Node position still + gets its real gradient through the displacement ``r``. + """ + payload = ModeCovariancePayload("affine") + torch.manual_seed(12) + xyz = torch.randn(20, 3, dtype=DTYPE) * 5.0 + node_pos = (torch.randn(4, 3, dtype=DTYPE) * 5.0).requires_grad_(True) + nl = torch.randint(0, 4, (20, 3)) + raw = torch.randn(4, payload.width, dtype=DTYPE) * 0.3 + + payload.contributions(raw, xyz, node_pos, nl).sum().backward() + assert node_pos.grad is not None + assert float(node_pos.grad.abs().sum()) > 0, "no gradient reaches node positions at all" + + # The scale itself must be a plain float, not something carrying a graph. + lam = payload._length_scale(node_pos) + assert isinstance(lam, float) diff --git a/tests/unit/model/test_model.py b/tests/unit/model/test_model.py index 1781a89a..da9a28bb 100644 --- a/tests/unit/model/test_model.py +++ b/tests/unit/model/test_model.py @@ -20,7 +20,7 @@ def test_model_empty_initialization(self): model = Model() - assert model.initialized == False + assert model.ctx.initialized is False assert model.pdb is None assert model.xyz is None assert model.adp is None @@ -54,12 +54,13 @@ def test_model_custom_dtype(self): @pytest.mark.unit def test_model_strip_h_default(self): - """Test strip_H defaults to True.""" + """Hydrogen stripping and generation are both opt-in.""" from torchref.model.model import Model - + model = Model() - - assert model.strip_H == True + + assert model.ctx.strip_H is False + assert model.ctx.add_hydrogens is False @pytest.mark.unit def test_model_bool_uninitialized(self): @@ -116,3 +117,40 @@ def test_get_selection_mask_uninitialized_raises(self): with pytest.raises(RuntimeError, match="uninitialized"): model.get_selection_mask("chain A") + + +@pytest.mark.unit +def test_dropped_rows_leave_a_positional_index(pdb_dir, tmp_path): + """A model losing atoms to the NaN drop must still index its own tensors. + + ``load`` derives the ``index`` column from the DataFrame index, and every + consumer uses it to address length-N per-atom tensors positionally. Dropping rows + without reindexing leaves gaps, so the largest value exceeds N-1 and + ``_create_occupancy_groups`` walks off the end of ``initial_occ``. Roughly one + PDB-REDO entry in six carries an atom with no coordinates or no B and hit this. + """ + import pandas as pd + + from torchref.model.model import Model + + src = Model(verbose=0) + src.load_pdb(str(pdb_dir / "3GR5.pdb")) + df = src.pdb.copy() + n_before = len(df) + + # Blank the B of a few interior atoms so the dropna removes them. + victims = [5, 100, 500] + df.loc[victims, "tempfactor"] = float("nan") + cell = src.cell.data.cpu().numpy() + sg = src.spacegroup + + model = Model(verbose=0) + model.load(lambda: (df, cell, sg), add_hydrogens=False) + + assert len(model.pdb) == n_before - len(victims) + idx = model.pdb["index"].to_numpy() + assert idx.min() == 0 + assert idx.max() == len(model.pdb) - 1, "index must stay positional after a drop" + assert sorted(idx) == list(range(len(model.pdb))) + # The occupancy grouping is what actually indexed past the end. + assert model.occupancy().shape[0] == len(model.pdb) diff --git a/tests/unit/model/test_model_collection_fractions.py b/tests/unit/model/test_model_collection_fractions.py new file mode 100644 index 00000000..0ccc3a8b --- /dev/null +++ b/tests/unit/model/test_model_collection_fractions.py @@ -0,0 +1,468 @@ +"""Population fractions, shared ownership, constraints and gradients.""" + +import pytest +import torch +from torch import nn + +from torchref.config import ( + get_complex_dtype, + get_default_device, + get_float_dtype, + get_int_dtype, +) + + +class _StubModel(nn.Module): + """Minimal stand-in for ``ModelFT`` for fraction bookkeeping.""" + + def __init__(self, seed: int): + super().__init__() + self.anchor = nn.Parameter( + torch.zeros(1, device=get_default_device(), dtype=get_float_dtype()) + ) + self._seed = seed + + @property + def device(self): + return self.anchor.device + + @property + def dtype_float(self): + return self.anchor.dtype + + def forward(self, hkl, recalc: bool = False): + # Distinct per model, and a function of the parameter so a gradient can reach it. + n = hkl.shape[0] + base = torch.arange( + 1, n + 1, dtype=self.anchor.dtype, device=self.anchor.device + ) + amp = (base * float(self._seed + 1)) + self.anchor + return amp.to(get_complex_dtype()) + + +@pytest.fixture +def two_model_collection(): + """A dark + one-timepoint collection over two distinguishable stub models.""" + from torchref.model.model_collection import ModelCollection + + mc = ModelCollection([_StubModel(0), _StubModel(1)], dark_key="dark", verbose=0) + mc.add_dark() + mc.add_timepoint("light", [0.7, 0.3]) + return mc + + +@pytest.fixture +def hkl(): + return torch.tensor( + [[1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 1]], + device=get_default_device(), + dtype=get_int_dtype(), + ) + + +class TestPopulationFactorisation: + @pytest.mark.unit + def test_fractions_are_the_activation_times_the_branching( + self, two_model_collection + ): + mc = two_model_collection + alpha = mc.alpha_mean + expected = torch.stack([1.0 - alpha, alpha * mc.branching()[0][0]]) + assert torch.allclose(mc["light"].fractions, expected) + + @pytest.mark.unit + @pytest.mark.parametrize("f", [0.01, 0.22, 0.3, 0.5, 0.99]) + def test_requested_fractions_round_trip(self, f): + """What you pass to ``add_timepoint`` is what ``fractions`` reports back.""" + from torchref.model.model_collection import ModelCollection + + mc = ModelCollection([_StubModel(0), _StubModel(1)], verbose=0) + mc.add_timepoint("t", [1.0 - f, f]) + assert mc["t"].fractions[1].item() == pytest.approx(f, abs=1e-6) + assert mc["t"].fractions.sum().item() == pytest.approx(1.0, abs=1e-6) + + @pytest.mark.unit + def test_the_reference_row_is_exactly_e_ref(self, two_model_collection): + """The dark is the alpha = 0 evaluation, so its excited fraction is exactly zero + -- not a clamp floor.""" + dark = two_model_collection["dark"] + assert dark.fractions[1].item() == 0.0 + assert dark.fractions[0].item() == 1.0 + + @pytest.mark.unit + def test_fraction_dtype_and_device_follow_the_base_models(self, any_device): + from torchref.model.model_collection import ModelCollection + + models = [_StubModel(0).to(any_device), _StubModel(1).to(any_device)] + mc = ModelCollection(models, verbose=0) + mc.add_timepoint("t", [0.6, 0.4]) + assert mc._activation_logit.dtype == models[0].dtype_float + assert mc._activation_logit.device == models[0].device + + +class TestValidation: + @pytest.mark.unit + def test_fractions_must_sum_to_one(self): + from torchref.model.model_collection import ModelCollection + + mc = ModelCollection([_StubModel(0), _StubModel(1)], verbose=0) + with pytest.raises(ValueError, match="sum to 1"): + mc.add_timepoint("t", [0.5, 0.2]) + + @pytest.mark.unit + def test_fraction_count_must_match_the_model_count(self): + from torchref.model.model_collection import ModelCollection + + mc = ModelCollection([_StubModel(0), _StubModel(1)], verbose=0) + with pytest.raises(ValueError, match="must match"): + mc.add_timepoint("t", [1.0]) + + @pytest.mark.unit + def test_duplicate_timepoint_names_are_rejected(self, two_model_collection): + with pytest.raises(ValueError, match="already exists"): + two_model_collection.add_timepoint("light", [0.5, 0.5]) + + +class TestFreezing: + @pytest.mark.unit + def test_dark_is_frozen_and_timepoints_are_not(self, two_model_collection): + mc = two_model_collection + # Frozen by default: population refinement is opt-in. + assert mc._activation_logit.requires_grad is False + assert mc.fraction_parameters() == [ + mc._activation_logit, + mc._branching_logits[0], + ] + + @pytest.mark.unit + def test_freeze_and_unfreeze_flip_the_flag(self, two_model_collection): + mixed = two_model_collection["light"] + mixed.freeze_fractions() + assert mixed.collection._activation_logit.requires_grad is False + mixed.unfreeze_fractions() + assert mixed.collection._activation_logit.requires_grad is True + + @pytest.mark.unit + def test_the_reference_carries_no_population_parameter(self, two_model_collection): + """The reference is the alpha = 0 evaluation, not a row with pinned logits, + so there is nothing of its own to freeze or refine.""" + mc = two_model_collection + mc.unfreeze_all_fractions() + assert "dark" not in mc._branching_rows + assert "light" in mc._branching_rows + assert mc._activation_logit.requires_grad is True + + @pytest.mark.unit + def test_freeze_all_freezes_every_timepoint(self, two_model_collection): + mc = two_model_collection + mc.freeze_all_fractions() + assert all(not p.requires_grad for p in mc.fraction_parameters()) + + +class TestOverride: + @pytest.mark.unit + def test_override_replaces_fractions_and_clears_back(self, two_model_collection): + mixed = two_model_collection["light"] + forced = torch.tensor( + [0.1, 0.9], device=get_default_device(), dtype=get_float_dtype() + ) + + mixed.set_fraction_override(forced) + assert mixed.fractions is forced + + mixed.clear_fraction_override() + assert torch.allclose( + mixed.fractions, mixed.collection.fractions_matrix()[mixed._index] + ) + + @pytest.mark.unit + def test_override_reaches_the_forward(self, two_model_collection, hkl): + """The override is the point at which an external (kinetic) population enters + the structure-factor sum, so it has to change ``forward``, not just the property. + """ + mixed = two_model_collection["light"] + before = mixed(hkl, recalc=True) + + mixed.set_fraction_override( + torch.tensor( + [0.1, 0.9], device=get_default_device(), dtype=get_float_dtype() + ) + ) + after = mixed(hkl, recalc=True) + + assert not torch.allclose(before, after) + + @pytest.mark.unit + def test_override_carries_gradient(self, two_model_collection, hkl): + """Gradients must flow through the override to whatever produced it.""" + mixed = two_model_collection["light"] + forced = torch.tensor( + [0.4, 0.6], + requires_grad=True, + device=get_default_device(), + dtype=get_float_dtype(), + ) + mixed.set_fraction_override(forced) + + mixed(hkl, recalc=True).abs().sum().backward() + + assert forced.grad is not None + assert torch.isfinite(forced.grad).all() + + +class TestCollectionLevelViews: + @pytest.mark.unit + def test_fractions_matrix_rows_follow_insertion_order(self, two_model_collection): + mc = two_model_collection + matrix = mc.get_fractions_matrix() + assert matrix.shape == (2, 2) + for row, key in enumerate(mc.keys()): + assert torch.allclose(matrix[row], mc[key].fractions) + + @pytest.mark.unit + def test_timepoint_names_excludes_the_dark_key(self, two_model_collection): + mc = two_model_collection + assert mc.dark_key == "dark" + assert mc.timepoint_names == ["light"] + assert mc.keys() == ["dark", "light"] + + @pytest.mark.unit + def test_base_models_are_shared_not_copied(self, two_model_collection): + mc = two_model_collection + assert mc.n_base_models == 2 + for i in range(2): + assert mc["dark"].models[i] is mc["light"].models[i] + assert mc["dark"].models[i] is mc.base_models[i] + + @pytest.mark.unit + def test_a_timepoint_owns_only_its_fractions(self, two_model_collection): + """``_SharedMixedModel`` holds the base models in a plain list, not a + ``ModuleList``, so a timepoint must not re-register their parameters. + """ + mixed = two_model_collection["light"] + assert ( + list(mixed.parameters()) == [] + ), "a timepoint view registered a parameter of its own" + + @pytest.mark.unit + def test_shared_base_parameters_are_counted_once(self, two_model_collection): + """Two timepoints over two shared models: two base parameters plus one set of + fractions each. Double-registration would make the optimizer step a base model + once per timepoint. + """ + mc = two_model_collection + params = list(mc.parameters()) + # two base anchors + activation + lambda + one branching row + assert len(params) == 2 + 3 + + base_anchors = [m.anchor for m in mc.base_models] + for anchor in base_anchors: + assert sum(p is anchor for p in params) == 1 + + +class TestForwardAndGradient: + @pytest.mark.unit + def test_forward_is_the_fraction_weighted_sum_of_the_parts( + self, two_model_collection, hkl + ): + mixed = two_model_collection["light"] + parts = mixed.get_individual_fcalc(hkl, recalc=True) + w = mixed.fractions + + expected = w[0] * parts[0] + w[1] * parts[1] + assert torch.allclose(mixed(hkl, recalc=True), expected) + + @pytest.mark.unit + def test_gradient_reaches_the_activation(self, two_model_collection, hkl): + mc = two_model_collection + mc.unfreeze_all_fractions() + mixed = mc["light"] + mixed(hkl, recalc=True).abs().sum().backward() + + grad = mc._activation_logit.grad + assert grad is not None + assert torch.isfinite(grad).all() + + @pytest.mark.unit + def test_the_reference_contributes_no_activation_gradient( + self, two_model_collection, hkl + ): + """The dark dataset carries no activation information, so its row must be + exactly e_ref with no path back to the shared parameter.""" + mc = two_model_collection + mc.unfreeze_all_fractions() + mc["dark"](hkl, recalc=True).abs().sum().backward() + assert ( + mc._activation_logit.grad is None + or float(mc._activation_logit.grad.abs().max()) == 0.0 + ) + + +class TestSharedActivation: + """One activation serves every timepoint; only the branching varies with time.""" + + @pytest.mark.unit + def test_a_second_timepoint_may_rebranch_at_the_same_activation(self): + """Three components, two timepoints, same 30% activated but split differently + between the two excited states.""" + from torchref.model.model_collection import ModelCollection + + mc = ModelCollection([_StubModel(i) for i in range(3)], verbose=0) + mc.add_dark() + mc.add_timepoint("early", [0.7, 0.3, 0.0]) + mc.add_timepoint("late", [0.7, 0.0, 0.3]) + + assert float(mc.alpha_mean) == pytest.approx(0.3, abs=1e-5) + assert torch.allclose( + mc["early"].fractions, + torch.tensor( + [0.7, 0.3, 0.0], device=get_default_device(), dtype=get_float_dtype() + ), + atol=1e-5, + ) + assert torch.allclose( + mc["late"].fractions, + torch.tensor( + [0.7, 0.0, 0.3], device=get_default_device(), dtype=get_float_dtype() + ), + atol=1e-5, + ) + + @pytest.mark.unit + def test_a_conflicting_activation_is_rejected_not_projected(self): + """A silent least-squares projection here would produce populations nobody + asked for, so this raises and names the escape hatch.""" + from torchref.model.model_collection import ModelCollection + + mc = ModelCollection([_StubModel(0), _StubModel(1)], verbose=0) + mc.add_dark() + mc.add_timepoint("early", [0.7, 0.3]) + + with pytest.raises(ValueError, match="set_fraction_override"): + mc.add_timepoint("late", [0.5, 0.5]) + + @pytest.mark.unit + def test_adding_the_reference_after_a_timepoint_leaves_activation_alone(self): + """A pure-reference row carries no activation information.""" + from torchref.model.model_collection import ModelCollection + + mc = ModelCollection([_StubModel(0), _StubModel(1)], verbose=0) + mc.add_timepoint("light", [0.78, 0.22]) + mc.add_dark() + assert float(mc.alpha_mean) == pytest.approx(0.22, abs=1e-5) + + +class TestActivationJacobian: + @pytest.mark.unit + def test_rows_sum_to_zero_and_the_reference_row_vanishes( + self, two_model_collection + ): + """Fractions stay on the simplex, so the derivative is tangent to it; and the + reference does not move with the activation at all.""" + mc = two_model_collection + jac = mc.activation_jacobian() + + assert jac.shape == (len(mc), mc.n_base_models) + assert torch.allclose( + jac.sum(dim=1), + torch.zeros(len(mc), device=get_default_device(), dtype=get_float_dtype()), + atol=1e-6, + ) + assert torch.equal( + jac[0], + torch.zeros( + mc.n_base_models, device=get_default_device(), dtype=get_float_dtype() + ), + ) + assert float(jac[1][0]) == pytest.approx(-1.0) + + @pytest.mark.unit + def test_fractions_matrix_is_e_ref_plus_alpha_times_the_jacobian( + self, two_model_collection + ): + mc = two_model_collection + e_ref = torch.zeros( + mc.n_base_models, device=get_default_device(), dtype=get_float_dtype() + ) + e_ref[0] = 1.0 + expected = e_ref.unsqueeze(0) + mc.alpha_mean * mc.activation_jacobian() + assert torch.allclose(mc.fractions_matrix(), expected) + + @pytest.mark.unit + def test_the_mixture_is_exactly_linear_in_the_activation( + self, two_model_collection + ): + """The property the second moment rests on: the secant equals the derivative, + so there is no truncation term anywhere downstream.""" + mc = two_model_collection + jac = mc.activation_jacobian() + + mc.set_activation(0.6) + w1 = mc.fractions_matrix().clone() + mc.set_activation(0.1) + w2 = mc.fractions_matrix().clone() + + assert torch.allclose(w1 - w2, (0.6 - 0.1) * jac, atol=1e-6) + + +class TestActivationDispersion: + @pytest.mark.unit + def test_lambda_is_exactly_zero_by_default(self, two_model_collection): + """Exactly, not approximately: sigmoid can never return 0, so a fixed float is + the only way to reproduce the coherent single-moment model.""" + mc = two_model_collection + assert float(mc.lambda_twin) == 0.0 + assert float(mc.sigma_alpha_sq) == 0.0 + + @pytest.mark.unit + def test_lambda_is_not_a_live_parameter_until_asked_for(self, two_model_collection): + mc = two_model_collection + + def _present(): + # Identity, not ``in``: ``==`` on tensors is elementwise. + return any(p is mc._lambda_logit for p in mc.fraction_parameters()) + + assert not _present() + + mc.set_lambda_twin(0.3, refinable=True) + assert _present() + assert mc._lambda_logit.requires_grad is True + + @pytest.mark.unit + @pytest.mark.parametrize("lam", [0.0, 0.25, 0.5, 1.0]) + def test_the_variance_bound_holds_by_construction(self, two_model_collection, lam): + mc = two_model_collection + mc.set_lambda_twin(lam) + alpha = float(mc.alpha_mean) + bound = alpha * (1.0 - alpha) + + # The bound holds algebraically; the tolerance is float32 ulp, since the two + # sides reach alpha (1 - alpha) by different arithmetic. + sigma_sq = float(mc.sigma_alpha_sq) + assert 0.0 <= sigma_sq <= bound * (1.0 + 1e-6) + assert sigma_sq == pytest.approx(bound * lam, rel=1e-5) + + @pytest.mark.unit + def test_lambda_one_saturates_the_bound(self, two_model_collection): + """The fully incoherent limit: every crystal either fully activated or dark.""" + mc = two_model_collection + mc.set_lambda_twin(1.0) + alpha = float(mc.alpha_mean) + assert float(mc.sigma_alpha_sq) == pytest.approx( + alpha * (1.0 - alpha), rel=1e-5 + ) + + @pytest.mark.unit + @pytest.mark.parametrize("bad", [-0.1, 1.1]) + def test_out_of_range_lambda_is_rejected(self, two_model_collection, bad): + with pytest.raises(ValueError, match=r"\[0, 1\]"): + two_model_collection.set_lambda_twin(bad) + + @pytest.mark.unit + def test_refined_lambda_stays_strictly_interior(self, two_model_collection): + """Once refinable it is a sigmoid, so it can approach but never reach the + bounds -- which is why the fixed path exists.""" + mc = two_model_collection + mc.set_lambda_twin(0.0, refinable=True) + value = float(mc.lambda_twin.detach()) + assert 0.0 < value < 1.0 diff --git a/tests/unit/model/test_riding_orientations.py b/tests/unit/model/test_riding_orientations.py new file mode 100644 index 00000000..633af89f --- /dev/null +++ b/tests/unit/model/test_riding_orientations.py @@ -0,0 +1,276 @@ +"""Refinable hydrogen orientations preserve geometry and expose force gradients.""" + +import numpy as np +import pytest +import torch +from torchref import Model, ModelFT +from torchref.base.coordinates.local_frame import rotate_vectors +from torchref.config import get_float_dtype +from torchref.model.riding_xyz import RidingXYZTensor +from torchref.topology.hydrogens import _place_group, _template + + +@pytest.fixture(scope="module") +def oriented_model(pdb_dir): + """A deposited protein with explicitly generated protein and water hydrogens.""" + with torch.random.fork_rng(): + torch.manual_seed(19) + model = Model(verbose=0, device="cpu", add_hydrogens=True) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + model.set_hydrogen_mode("riding") + return model + + +@pytest.fixture +def two_groups(oriented_model): + """One methyl and one water, retaining the methyl's heavy frame atoms.""" + full = oriented_model.xyz + methyl = next( + int(g) + for g in torch.unique(full.torsion_group).tolist() + if g >= 0 + and int((full.torsion_group == g).sum()) == 3 + and oriented_model.pdb.iloc[ + int(full.parent_row[full.torsion_group == g][0]) + ].element.strip() + == "C" + ) + water = int(full.rotation_group[full.rotation_group >= 0][0]) + selected = (full.torsion_group == methyl) | (full.rotation_group == water) + keep = torch.zeros(full.shape[0], dtype=torch.bool) + for rows in (full.h_row, full.parent_row, full.n1_row, full.n2_row): + chosen = rows[selected] + keep[chosen[chosen >= 0]] = True + return full.select_rows(keep) + + +@pytest.mark.unit +def test_orientation_groups_preserve_initial_positions(oriented_model): + """Zero rotations reproduce the deposited/generated table without atom changes.""" + model = oriented_model + xyz = torch.as_tensor(model.pdb[["x", "y", "z"]].values, dtype=model.dtype_float) + assert torch.allclose(model.xyz(), xyz, atol=1e-4) + assert model.xyz.torsions.shape[0] > 0 + waters = model.pdb[model.pdb.resname.str.strip() == "HOH"] + assert len(waters) > 0 + assert model.xyz.rotations.shape == ( + int((waters.element.str.strip() == "O").sum()), + 3, + ) + + +@pytest.mark.unit +def test_water_initialization_is_seeded_and_preserves_existing_direction( + oriented_model, +): + """Water references are reproducible and completing an O–H pair preserves its angle.""" + model = oriented_model + xyz = model.xyz().detach().numpy() + first = int(model.xyz._rotation_first[0]) + parent = int(model.xyz.parent_row[first]) + h_rows = model.xyz.h_row[model.xyz.rotation_group == 0].numpy() + names = model.pdb.name.str.strip().to_numpy() + template = _template(model.restraints.cif_dict, "HOH") + h_names = list(names[h_rows]) + lengths = np.linalg.norm(xyz[h_rows] - xyz[parent], axis=1) + + def place(seed, present): + with torch.random.fork_rng(): + torch.manual_seed(seed) + return _place_group( + template, + names[parent], + xyz[parent], + np.empty((0, 3)), + 0, + h_names, + lengths, + present, + xyz, + [], + ) + + first_reference = place(1, {names[parent]: parent}) + assert np.array_equal(first_reference, place(1, {names[parent]: parent})) + assert not np.allclose(first_reference, place(2, {names[parent]: parent})) + completed = place(3, {names[parent]: parent, h_names[0]: int(h_rows[0])}) + assert np.allclose(completed[0], xyz[h_rows[0]], atol=1e-5) + assert np.allclose( + np.linalg.norm(completed[0] - completed[1]), + np.linalg.norm(first_reference[0] - first_reference[1]), + atol=1e-5, + ) + + +@pytest.mark.unit +def test_rotation_preserves_water_and_methyl_geometry(two_groups): + """Shared rotations preserve internal distances and methyl bond angles.""" + w = two_groups + before = w().detach() + with torch.no_grad(): + w.torsions.refinable_params.fill_(0.8) + w.rotations.refinable_params.copy_( + w.rotations.refinable_params.new_tensor([[0.4, -0.2, 0.7]]) + ) + after = w() + assert torch.equal(before[w.base_row], after[w.base_row]) + assert not torch.allclose(before[w.h_row], after[w.h_row]) + for group in (w.torsion_group, w.rotation_group): + rows = w.h_row[group >= 0] + parent = w.parent_row[group >= 0][0:1] + rows = torch.cat((parent, rows)) + assert torch.allclose( + torch.cdist(before[rows], before[rows]), + torch.cdist(after[rows], after[rows]), + atol=1e-5, + ) + rows = w.h_row[w.torsion_group >= 0] + neighbour = w.n1_row[w.torsion_group >= 0] + assert torch.allclose( + (before[rows] - before[neighbour]).norm(dim=-1), + (after[rows] - after[neighbour]).norm(dim=-1), + atol=1e-5, + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("angle", [0.0, 0.4]) +def test_orientation_gradients_match_finite_differences(two_groups, double_cpu, angle): + """The deposited methyl/water coordinate map has correct orientation gradients.""" + w = two_groups.to(dtype=get_float_dtype()) + base = w._storage_values().detach().requires_grad_() + torsion = w.torsions().detach().fill_(angle).requires_grad_() + rotation = w.rotations().detach().fill_(angle).requires_grad_() + assert torch.autograd.gradcheck(w.evaluate, (base, torsion, rotation), atol=2e-6) + vectors = w.rigid_offset[w.rotation_group >= 0].detach() + r = vectors.new_full(vectors.shape, angle, requires_grad=True) + assert torch.autograd.gradgradcheck(lambda x: rotate_vectors(vectors, x), (r,)) + + +@pytest.mark.unit +def test_xyz_optimizer_receives_and_updates_orientations(two_groups): + """An orientation-only selection delivers both leaves to the xyz optimizer.""" + w = two_groups + w.fix_all() + w.refine(w.h_row) + model = Model(device="cpu", verbose=0) + model.xyz = w + leaves = model.parameters_of_types(("xyz",)) + assert {id(p) for p in leaves} == {id(p) for p in w.optimization_parameters()} + target = w.evaluate( + w._storage_values(), w.torsions().detach() + 0.2, w.rotations().detach() + 0.2 + ).detach() + optimizer = torch.optim.SGD(leaves, lr=0.05) + before = w._storage_values().detach().clone() + losses = [] + for _ in range(8): + optimizer.zero_grad() + loss = (w() - target).square().sum() + losses.append(float(loss.detach())) + loss.backward() + optimizer.step() + assert losses[-1] < losses[0] / 2 + assert torch.equal(w._storage_values(), before) + assert w.torsions.refinable_params.abs().sum() > 0 + assert w.rotations.refinable_params.abs().sum() > 0 + + +@pytest.mark.unit +def test_orientation_mutation_invalidates_forward_cache(two_groups): + """Both orientation leaves participate in the coordinate cache fingerprint.""" + w = two_groups + before = w() + with torch.no_grad(): + w.torsions.refinable_params.add_(0.1) + changed = w() + assert changed is not before + with torch.no_grad(): + w.rotations.refinable_params.add_(0.1) + assert w() is not changed + + +@pytest.mark.unit +def test_copy_selection_and_checkpoint_preserve_orientations(two_groups): + """Nonzero rotations survive copy, subset, and empty-shell state restoration.""" + w = two_groups + with torch.no_grad(): + w.torsions.refinable_params.fill_(0.3) + w.rotations.refinable_params.fill_(-0.2) + expected = w().detach() + copied = w.copy() + assert torch.equal(copied(), expected) + assert torch.equal(copied.rotations(), w.rotations()) + keep = torch.ones(w.shape[0], dtype=torch.bool) + keep[w.h_row[w.torsion_group >= 0]] = False + subset = w.select_rows(keep) + assert torch.allclose(subset(), expected[keep], atol=1e-5) + restored = RidingXYZTensor(device="cpu") + restored.load_state_dict(w.state_dict()) + assert torch.equal(restored(), expected) + with torch.no_grad(): + copied.rotations.refinable_params.add_(0.1) + assert torch.equal(w(), expected) + + +@pytest.mark.unit +def test_checkpoint_without_orientation_metadata(two_groups): + """Coordinate-only checkpoints load with fixed hydrogen orientations.""" + state = { + name: value + for name, value in two_groups.state_dict().items() + if name not in ("torsion_group", "rotation_group", "virtual_reference") + and not name.startswith(("torsions.", "rotations.")) + } + restored = RidingXYZTensor(device="cpu") + restored.load_state_dict(state) + assert torch.allclose(restored(), two_groups(), atol=1e-5) + assert restored.torsions.shape == (0,) + assert restored.rotations.shape == (0, 3) + + +@pytest.mark.unit +@pytest.mark.parametrize("model_class", [Model, ModelFT]) +def test_model_checkpoint_restores_rotated_groups(oriented_model, model_class): + """Model checkpoints restore orientation values, masks, and frame references.""" + source = oriented_model.copy() + with torch.no_grad(): + source.xyz.torsions.refinable_params.fill_(0.3) + source.xyz.rotations.refinable_params.fill_(0.2) + source.xyz.rotations.fix(torch.arange(0, source.xyz.rotations.shape[0], 2)) + state = source.state_dict() + restored = model_class.create_from_state_dict(state, device="cpu") + assert torch.allclose(restored.xyz(), source.xyz(), atol=1e-5) + assert torch.equal( + restored.xyz.rotations.refinable_mask, source.xyz.rotations.refinable_mask + ) + + +@pytest.mark.unit +def test_torsion_without_second_reference_still_rotates(two_groups): + """A fixed reference direction completes a methyl frame with one heavy bond.""" + frames = two_groups.hydrogen_frames() + frames.n2_row[frames.torsion_group >= 0] = -1 + frames.frame_valid[frames.torsion_group >= 0] = False + w = RidingXYZTensor(two_groups().detach(), frames) + initial = w().detach() + with torch.no_grad(): + w.torsions.refinable_params.fill_(0.6) + rows = w.h_row[w.torsion_group >= 0] + assert not torch.allclose(w()[rows], initial[rows]) + assert torch.allclose( + torch.cdist(w()[rows], w()[rows]), + torch.cdist(initial[rows], initial[rows]), + atol=1e-5, + ) + + +@pytest.mark.unit +def test_orientation_forward_runs_on_requested_device(two_groups, any_device): + """Coordinate and orientation gradients stay on the requested backend.""" + w = two_groups.to(any_device) + output = w() + output.square().sum().backward() + assert output.device == any_device + for p in w.optimization_parameters(): + assert p.device == any_device + assert p.grad is not None and torch.isfinite(p.grad).all() diff --git a/tests/unit/model/test_riding_water_completion.py b/tests/unit/model/test_riding_water_completion.py new file mode 100644 index 00000000..418b75af --- /dev/null +++ b/tests/unit/model/test_riding_water_completion.py @@ -0,0 +1,184 @@ +"""Riding-mode water completion respects the model's hydrogen generation setting.""" + +import pytest +import torch + +from torchref import Model, ModelFT + + +@pytest.fixture(scope="module") +def heavy_model(pdb_dir): + """Deposited 1DAW loaded without hydrogen generation.""" + return Model(device="cpu", verbose=0, strip_H=True, add_hydrogens=False).load_pdb( + str(pdb_dir / "1DAW.pdb") + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("model_class", [Model, ModelFT]) +def test_riding_completes_only_waters_and_preserves_live_atoms( + heavy_model, model_class +): + """Water completion preserves current coordinates, ADPs, selections and links.""" + model = model_class(device="cpu", verbose=0, add_hydrogens=False, strip_H=False) + model.load( + lambda: (heavy_model.pdb.copy(), heavy_model.cell.data, heavy_model.spacegroup) + ) + with torch.no_grad(): + model.xyz.refinable_params.add_(0.25) + adp = model.adp().detach().clone() + mask = torch.arange(len(model.pdb)) % 2 == 0 + model.xyz.update_refinable_mask(mask) + model.adp.update_refinable_mask(mask) + model.occupancy.freeze_all() + before = model.xyz().detach().clone() + links = model.ctx.links + hkl = torch.tensor([[1, 0, 0], [0, 1, 0], [1, 1, 1]]) + if isinstance(model, ModelFT): + model(hkl) + model.ctx.add_hydrogens = True + returned = model.set_hydrogen_mode("riding") + assert returned is model + is_h = torch.as_tensor(model.pdb.element.str.strip().eq("H").to_numpy()) + is_water = model.pdb.resname.str.strip().eq("HOH") + assert int(is_h.sum()) == 2 * int( + heavy_model.pdb.resname.str.strip().eq("HOH").sum() + ) + assert is_water[is_h.numpy()].all() + assert torch.allclose(model.xyz()[~is_h], before, atol=1e-5) + assert torch.allclose(model.adp()[~is_h], adp) + assert torch.equal(model.xyz.full_refinable_mask[~is_h], mask) + assert torch.equal(model.adp.refinable_mask[~is_h], mask) + assert not model.occupancy.get_refinable_atoms().any() + assert model.ctx.links is links + assert model.xyz.rotations.shape[0] == int(is_h.sum()) // 2 + assert model.restraints.xyz().shape == model.xyz.shape + if isinstance(model, ModelFT): + assert torch.isfinite(model(hkl)).all() + restored = model_class.create_from_state_dict(model.state_dict(), device="cpu") + assert torch.allclose(restored.xyz(), model.xyz(), atol=1e-5) + + +@pytest.mark.unit +def test_partial_water_is_completed_without_moving_existing_hydrogen(heavy_model): + """A deposited oxygen with one supplied hydrogen receives just its missing partner.""" + water = ( + heavy_model.pdb[heavy_model.pdb.resname.str.strip().eq("HOH")].iloc[:1].copy() + ) + model = heavy_model._new_model_from_df(water, strip_H=False) + model.ctx.add_hydrogens = True + model.set_hydrogen_mode("riding") + model.update_pdb() + partial = model._new_model_from_df(model.pdb.iloc[:2].copy(), strip_H=False) + before = partial.xyz().detach().clone() + partial.set_hydrogen_mode("riding") + assert len(partial.pdb) == 2 + assert torch.allclose(partial.xyz(), before, atol=1e-5) + partial.ctx.add_hydrogens = True + partial.set_hydrogen_mode("riding") + assert len(partial.pdb) == 3 + assert torch.allclose(partial.xyz()[:2], before, atol=1e-5) + with torch.no_grad(): + partial.xyz.rotations.refinable_params.fill_(0.3) + wrapper = partial.xyz + coords = partial.xyz().detach().clone() + partial.set_hydrogen_mode("riding") + assert partial.xyz is wrapper + assert torch.equal(partial.xyz(), coords) + partial.set_hydrogen_mode("free").set_hydrogen_mode("riding") + assert len(partial.pdb) == 3 + assert torch.allclose(partial.xyz(), coords, atol=1e-5) + + +@pytest.mark.unit +def test_supplied_frames_are_remapped_when_waters_are_completed(heavy_model): + """Explicit frames for the original atom table coexist with generated water frames.""" + water = ( + heavy_model.pdb[heavy_model.pdb.resname.str.strip().eq("HOH")].iloc[:2].copy() + ) + model = heavy_model._new_model_from_df(water, strip_H=False) + frames = model.hydrogen_frames() + model.ctx.add_hydrogens = True + model.set_hydrogen_mode("riding", frames=frames) + assert len(model.pdb) == 6 + assert model.xyz.n_hydrogens == 4 + assert model.xyz.rotations.shape == (2, 3) + + +@pytest.mark.unit +def test_water_completion_preserves_adp_field(heavy_model): + """Adding water hydrogens retains the node parametrization and its parameters.""" + model = heavy_model._new_model_from_df(heavy_model.pdb.copy(), strip_H=False) + model.set_adp_mode("field", n_nodes=8, k_neighbors=4) + field = model.adp + values = field().detach().clone() + parameters = field.refinable_params + model.ctx.add_hydrogens = True + model.set_hydrogen_mode("riding") + heavy = torch.as_tensor(~model.pdb.element.str.strip().eq("H").to_numpy()) + assert model.adp is field + assert field.refinable_params is parameters + assert torch.allclose(field()[heavy], values, atol=1e-5) + assert field().shape == (len(model.pdb),) + + +@pytest.mark.integration +@pytest.mark.parametrize("generate", [False, True]) +def test_refinement_targets_follow_completed_atom_table(pdb_dir, mtz_dir, generate): + """Refinement adds water H and refreshes targets only when generation is enabled.""" + from torchref.refinement.base_refinement import Refinement + + refinement = Refinement( + pdb=str(pdb_dir / "1DAW.pdb"), + data_file=str(mtz_dir / "1DAW.mtz"), + device="cpu", + verbose=0, + max_res=3.0, + add_hydrogens=False, + ) + previous = refinement.adp_target + n_atoms = len(refinement.model.pdb) + refinement.model.ctx.add_hydrogens = generate + refinement.set_hydrogen_mode("riding") + assert (refinement.adp_target is not previous) == generate + assert (len(refinement.model.pdb) > n_atoms) == generate + geometry = refinement.geometry_target() + assert torch.isfinite(geometry) + geometry.backward() + if generate: + gradient = refinement.model.xyz.rotations.refinable_params.grad + assert gradient is not None and torch.isfinite(gradient).all() + + +@pytest.mark.unit +@pytest.mark.parametrize("model_class", [Model, ModelFT]) +@pytest.mark.parametrize("explicit_frames", [False, True]) +def test_disabled_generation_keeps_atom_table( + heavy_model, model_class, explicit_frames +): + """Riding mode leaves oxygen-only waters untouched when add_hydrogens is False.""" + model = model_class(device="cpu", verbose=0, add_hydrogens=False) + model.load( + lambda: (heavy_model.pdb.copy(), heavy_model.cell.data, heavy_model.spacegroup) + ) + coordinates = model.xyz().detach().clone() + table = model.pdb.copy() + adp, occupancy = model.adp, model.occupancy + frames = model.hydrogen_frames() if explicit_frames else None + model.set_hydrogen_mode("riding", frames=frames) + assert model.pdb.equals(table) + assert torch.equal(model.xyz(), coordinates) + assert model.adp is adp and model.occupancy is occupancy + assert model.xyz.n_hydrogens == 0 + assert model.xyz.rotations.shape == (0, 3) + + +@pytest.mark.unit +def test_stripping_prevents_water_completion(heavy_model): + """The stripping preference takes precedence even when generation is enabled.""" + model = heavy_model._new_model_from_df(heavy_model.pdb.copy(), strip_H=True) + model.ctx.add_hydrogens = True + n_atoms = len(model.pdb) + model.set_hydrogen_mode("riding") + assert len(model.pdb) == n_atoms + assert model.xyz.n_hydrogens == 0 diff --git a/tests/unit/model/test_riding_xyz.py b/tests/unit/model/test_riding_xyz.py new file mode 100644 index 00000000..ba5b9414 --- /dev/null +++ b/tests/unit/model/test_riding_xyz.py @@ -0,0 +1,213 @@ +"""The riding coordinate wrapper: full-space face, heavy-atom storage. + +Pinned here: the placed hydrogens are reproduced from the stored rows, a force on a +hydrogen reaches only the atoms that carry it, every public method speaks full atom +space while the single refinable leaf holds heavy rows only, and the wrapper survives +the conversions the model needs (to and from a plain per-atom wrapper, subsets, +copies, state dicts). +""" + +import numpy as np +import pytest +import torch + +from torchref.model.model import Model +from torchref.model.parameter_wrappers import MixedTensor +from torchref.model.riding_xyz import RidingXYZTensor +from torchref.topology.hydrogens import HydrogenFrames, hydrogen_frames + + +@pytest.fixture(scope="module") +def hydrogenated(pdb_dir): + """1DAW with generated hydrogens, its frames, and its full coordinate table.""" + model = Model(verbose=0, add_hydrogens=True) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + frames = hydrogen_frames(model.restraints.topology) + return model, frames, model.xyz().detach() + + +def _tolerance(dtype): + return 1e-4 if dtype == torch.float32 else 1e-9 + + +@pytest.mark.unit +def test_reconstructs_placed_hydrogens(hydrogenated): + model, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + out = riding() + assert out.shape == xyz.shape + assert riding.shape == tuple(xyz.shape) + assert float((out - xyz).norm(dim=1).max()) < _tolerance(xyz.dtype) + assert riding.n_hydrogens == frames.n_hydrogens + assert riding.refinable_params.shape[0] == xyz.shape[0] - frames.n_hydrogens + + +@pytest.mark.unit +def test_coordinate_leaf_holds_heavy_rows_with_separate_orientations(hydrogenated): + """Heavy positions and shared orientations have separate optimizer leaves.""" + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + leaves = list(riding.parameters()) + assert len(leaves) == 3 + assert leaves[0].shape == (xyz.shape[0] - frames.n_hydrogens, 3) + assert leaves[1] is riding.torsions.refinable_params + assert leaves[2] is riding.rotations.refinable_params + assert riding.full_refinable_mask.sum() == leaves[0].shape[0] + assert not riding.full_refinable_mask[riding.h_row].any() + + +@pytest.mark.unit +def test_gradient_through_hydrogens_lands_on_frame_atoms(hydrogenated): + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + out = riding() + k = 5 + h_row = int(riding.h_row[k]) + (grad,) = torch.autograd.grad(out[h_row].sum(), riding.refinable_params) + touched = set(torch.nonzero(grad.abs().sum(1) > 0).flatten().tolist()) + expected = { + int(riding._parent_bidx[k]), + int(riding._n1_bidx[k]), + int(riding._n2_bidx[k]), + } + assert touched == expected + + +@pytest.mark.unit +def test_evaluate_matches_finite_differences(): + """The pure map from stored to full rows has exact gradients (float64).""" + xyz = torch.tensor( + [[0.0, 0.0, 0.0], [1.5, 0.0, 0.0], [1.5, 1.5, 0.1], [-0.6, 0.8, 0.2], [-0.6, -0.8, 0.0]], + dtype=torch.float64, + ) + frames = HydrogenFrames( + h_row=np.array([3, 4]), + parent_row=np.array([0, 0]), + n1_row=np.array([1, 1]), + n2_row=np.array([2, 2]), + frame_valid=np.array([True, True]), + ) + riding = RidingXYZTensor(xyz, frames) + base = riding.refinable_params.detach().clone().requires_grad_() + assert torch.autograd.gradcheck(riding.evaluate, (base,), eps=1e-6, atol=1e-6) + + +@pytest.mark.unit +def test_masks_are_full_space_and_ignore_hydrogen_rows(hydrogenated): + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + n = xyz.shape[0] + mask = torch.zeros(n, dtype=torch.bool) + mask[:100] = True + riding.update_refinable_mask(mask) + heavy_first = int((~torch.isin(torch.arange(100), riding.h_row.cpu())).sum()) + assert riding.get_refinable_count() == heavy_first + assert riding.full_refinable_mask.sum() == heavy_first + riding.fix_all() + assert riding.get_refinable_count() == 0 + assert riding.refinable_params.numel() == 0 + riding.refine_all() + assert riding.get_refinable_count() == n - frames.n_hydrogens + riding.fix(torch.arange(10)) + assert riding.full_refinable_mask[:10].sum() == 0 + + +@pytest.mark.unit +def test_frozen_parent_keeps_its_hydrogens_still(hydrogenated): + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + riding.fix_all() + before = riding().detach().clone() + with torch.no_grad(): + riding.refinable_params.add_(1.0) # empty leaf: nothing moves + assert torch.equal(riding(), before) + + +@pytest.mark.unit +def test_rigid_motion_preserves_local_offsets(hydrogenated): + """Writing a rotated table keeps every hydrogen riding at the same offset.""" + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + offsets = riding.local_offset.clone() + angle = torch.tensor(0.4, dtype=xyz.dtype) + rot = torch.tensor( + [[torch.cos(angle), -torch.sin(angle), 0.0], [torch.sin(angle), torch.cos(angle), 0.0], [0.0, 0.0, 1.0]], + dtype=xyz.dtype, device=xyz.device, + ) + moved = xyz @ rot.T + torch.tensor([3.0, -1.0, 2.0], dtype=xyz.dtype, device=xyz.device) + riding[:] = moved + assert float((riding() - moved).norm(dim=1).max()) < 10 * _tolerance(xyz.dtype) + assert torch.allclose(riding.local_offset, offsets, atol=10 * _tolerance(xyz.dtype)) + + +@pytest.mark.unit +def test_assigning_a_hydrogen_row_becomes_a_new_offset(hydrogenated): + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + h = int(riding.h_row[0]) + target = xyz[h] + torch.tensor([0.3, -0.2, 0.1], dtype=xyz.dtype, device=xyz.device) + riding[h] = target + assert float((riding()[h] - target).norm()) < 10 * _tolerance(xyz.dtype) + # Heavy rows untouched. + assert float((riding()[riding.base_row] - xyz[riding.base_row]).norm(dim=1).max()) < 10 * _tolerance(xyz.dtype) + + +@pytest.mark.unit +def test_round_trip_with_plain_wrapper(hydrogenated): + _, frames, xyz = hydrogenated + plain = MixedTensor(xyz.clone(), name="xyz") + riding = RidingXYZTensor.from_mixed_tensor(plain, frames) + back = riding.to_mixed_tensor() + assert isinstance(back, MixedTensor) and not isinstance(back, RidingXYZTensor) + assert float((back() - xyz).norm(dim=1).max()) < _tolerance(xyz.dtype) + assert back.refinable_mask.all() + + +@pytest.mark.unit +def test_select_rows_frees_a_hydrogen_whose_parent_is_cut(hydrogenated): + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + keep = torch.ones(xyz.shape[0], dtype=torch.bool) + parent = int(riding.parent_row[0]) + keep[parent] = False + sub = riding.select_rows(keep) + assert sub.shape[0] == xyz.shape[0] - 1 + assert sub.n_hydrogens == frames.n_hydrogens - int((riding.parent_row == parent).sum()) + expected = xyz[keep.to(xyz.device)] + assert float((sub() - expected).norm(dim=1).max()) < _tolerance(xyz.dtype) + + +@pytest.mark.unit +def test_copy_is_independent_and_exact(hydrogenated): + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + dup = riding.copy() + assert torch.equal(dup(), riding()) + assert torch.equal(dup.local_offset, riding.local_offset) + with torch.no_grad(): + dup.refinable_params.add_(1.0) + assert not torch.equal(dup(), riding()) + + +@pytest.mark.unit +def test_state_dict_round_trip(hydrogenated): + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + state = riding.state_dict() + assert "h_row" in state and "local_offset" in state + # As the model restore does: a placeholder of the right shape, values from the dict. + placeholder = RidingXYZTensor(torch.zeros_like(xyz), frames) + placeholder.load_state_dict(state) + assert placeholder.shape == riding.shape + assert torch.equal(placeholder(), riding()) + + +@pytest.mark.unit +def test_forward_is_cached_until_parameters_move(hydrogenated): + _, frames, xyz = hydrogenated + riding = RidingXYZTensor(xyz, frames) + first = riding() + assert riding() is first + with torch.no_grad(): + riding.refinable_params[0, 0] += 0.5 + assert riding() is not first diff --git a/tests/unit/model/test_rigid_xyz.py b/tests/unit/model/test_rigid_xyz.py index d7e76531..4628c183 100644 --- a/tests/unit/model/test_rigid_xyz.py +++ b/tests/unit/model/test_rigid_xyz.py @@ -40,16 +40,18 @@ def test_known_transform(self, fresh_modelft): dtype = xyz.dtype # Pick the first chain (index 0); apply a transform only to that chain. + # euler_angles is in Angstrom-scaled units, so scale going in and read + # the physical angle back through rotation_radians. ang_vec = torch.tensor([0.0, 0.05, 0.0], dtype=dtype, device=device) trans_vec = torch.tensor([0.2, -0.3, 0.5], dtype=dtype, device=device) with torch.no_grad(): xyz.euler_angles.zero_() xyz.translations.zero_() - xyz.euler_angles[0] = ang_vec + xyz.euler_angles[0] = ang_vec * xyz.angle_scale[0] xyz.translations[0] = trans_vec center = xyz.chain_centers[0] - R = rotation_matrix_euler_xyz(ang_vec) + R = rotation_matrix_euler_xyz(xyz.rotation_radians[0]) atom_chain = xyz.chain_indices chain0_mobile = (atom_chain == 0) & xyz.mobile_mask non_mobile = ~xyz.mobile_mask @@ -118,9 +120,10 @@ def test_bake_preserves_forward_and_zeros_params(self, fresh_modelft): with torch.no_grad(): xyz.euler_angles.zero_() xyz.translations.zero_() - xyz.euler_angles[0] = ang + xyz.euler_angles[0] = ang * xyz.angle_scale[0] xyz.translations[0] = trans before = xyz().detach().clone() + scale_before = xyz.angle_scale.clone() xyz.bake() @@ -134,6 +137,9 @@ def test_bake_preserves_forward_and_zeros_params(self, fresh_modelft): assert diff_original < 1e-4 assert torch.all(xyz.euler_angles == 0).item() assert torch.all(xyz.translations == 0).item() + # Rigid motion preserves the radius of gyration, so the scale is + # unchanged; if it drifted, later angles would mean something else. + assert torch.allclose(xyz.angle_scale, scale_before, rtol=1e-6) # Chain centers should have moved by the translation on chain 0. # Centroid is mass-weighted (atomic Z) over MOBILE atoms; reconstruct @@ -147,6 +153,87 @@ def test_bake_preserves_forward_and_zeros_params(self, fresh_modelft): diff_center = (xyz.chain_centers[0] - expected_center0).abs().max().item() assert diff_center < 1e-4 + @pytest.mark.unit + def test_angle_scale_matches_explicit_radius_of_gyration(self, fresh_modelft): + """``angle_scale`` is the per-chain RMS radius about the rotation centre, + over mobile atoms, floored at 1 A.""" + fresh_modelft.use_rigid_xyz() + xyz = fresh_modelft.xyz + + for c in range(xyz.n_chains): + sel = (xyz.chain_indices == c) & xyz.mobile_mask + d = xyz.original_xyz[sel] - xyz.chain_centers[c] + expected = d.pow(2).sum(dim=1).mean().sqrt() + assert torch.allclose(xyz.angle_scale[c], expected, rtol=1e-5) + assert torch.all(xyz.angle_scale >= 1.0), "scale must be floored at 1 A" + + @pytest.mark.unit + def test_angle_scale_tracks_non_rigid_update_fixed_values(self, fresh_modelft): + """``update_fixed_values`` recomputes the scale, not just the centres. + + It accepts any coordinates, not only the rigid re-pose ``bake()`` supplies, + and a rigid re-pose leaves the radius of gyration unchanged. + """ + fresh_modelft.use_rigid_xyz() + xyz = fresh_modelft.xyz + before = xyz.angle_scale.clone() + + # A pure dilation is not a rigid motion, so the radius doubles. + with torch.no_grad(): + centers = xyz.chain_centers[xyz.chain_indices] + inflated = centers + 2.0 * (xyz.original_xyz - centers) + xyz.update_fixed_values(inflated) + + assert torch.allclose(xyz.angle_scale, 2.0 * before, rtol=1e-4) + + @pytest.mark.unit + def test_unit_angle_scale_reproduces_unscaled_behaviour(self, fresh_modelft): + """``angle_scale`` of ones gives the unscaled parametrization exactly.""" + fresh_modelft.use_rigid_xyz() + xyz = fresh_modelft.xyz + dtype, device = xyz.dtype, xyz.device + ang = torch.tensor([0.02, -0.03, 0.04], dtype=dtype, device=device) + + with torch.no_grad(): + xyz.angle_scale.fill_(1.0) + xyz.euler_angles.zero_() + xyz.euler_angles[0] = ang + xyz.reset_forward_cache() + + center = xyz.chain_centers[0] + R = rotation_matrix_euler_xyz(ang) + sel = (xyz.chain_indices == 0) & xyz.mobile_mask + expected = (xyz.original_xyz[sel] - center) @ R.T + center + with torch.no_grad(): + assert torch.allclose(xyz()[sel], expected, atol=1e-4) + + @pytest.mark.unit + def test_rotation_radians_is_the_descaled_angle(self, fresh_modelft): + """``rotation_radians`` descales the angle, and ``copy()`` carries both.""" + fresh_modelft.use_rigid_xyz() + xyz = fresh_modelft.xyz + with torch.no_grad(): + xyz.euler_angles.copy_(torch.randn_like(xyz.euler_angles) * 0.1) + + assert torch.allclose( + xyz.rotation_radians, xyz.euler_angles / xyz.angle_scale.unsqueeze(1) + ) + clone = xyz.copy() + assert torch.allclose(clone.angle_scale, xyz.angle_scale) + assert torch.allclose(clone.euler_angles, xyz.euler_angles) + with torch.no_grad(): + assert torch.allclose(clone(), xyz(), atol=1e-5) + + # An overridden scale must be carried, not re-derived from geometry: + # the angles are copied raw, so a re-derived scale changes the pose. + with torch.no_grad(): + xyz.angle_scale.fill_(1.0) + xyz.reset_forward_cache() + clone2 = xyz.copy() + assert torch.allclose(clone2.angle_scale, torch.ones_like(xyz.angle_scale)) + with torch.no_grad(): + assert torch.allclose(clone2(), xyz(), atol=1e-5) + @pytest.mark.unit def test_restore_commit_bakes_transform(self, fresh_modelft): fresh_modelft.use_rigid_xyz() diff --git a/tests/unit/model/test_sf_grid_key.py b/tests/unit/model/test_sf_grid_key.py new file mode 100644 index 00000000..c43b71df --- /dev/null +++ b/tests/unit/model/test_sf_grid_key.py @@ -0,0 +1,212 @@ +"""The FFT grid is derived from the model's context and cached on a value key.""" + +import pytest +import torch + +from torchref.config import dtypes +from torchref.model import ModelFT +from torchref.model.context import ModelContext +from torchref.model.sf_fft import SfFFT +from torchref.symmetry import Cell, SpaceGroup + + +@pytest.fixture(scope="module") +def pdb_path(pdb_dir): + path = pdb_dir / "1DAW.pdb" + if not path.exists(): + pytest.skip("1DAW.pdb fixture not present") + return str(path) + + +def _model(pdb_path, **kwargs) -> ModelFT: + return ModelFT(max_res=2.5, verbose=0, device="cpu", **kwargs).load_pdb(pdb_path) + + +def _hkl(model, n=64): + gen = torch.Generator().manual_seed(0) + return torch.randint(-8, 9, (n, 3), generator=gen).to( + dtype=dtypes.int, device=model.device + ) + + +def _count_calls(monkeypatch, cls, name): + counter = {"n": 0} + original = getattr(cls, name) + + def wrapped(self, *args, **kwargs): + counter["n"] += 1 + return original(self, *args, **kwargs) + + monkeypatch.setattr(cls, name, wrapped) + return counter + + +@pytest.mark.unit +@pytest.mark.parametrize("strip_H", [False, True]) +def test_one_engine_and_one_spacegroup_per_load(pdb_path, monkeypatch, strip_H): + import torchref.model.model_ft as model_ft_module + + engines = _count_calls(monkeypatch, model_ft_module.SfFFT, "__init__") + spacegroups = _count_calls(monkeypatch, SpaceGroup, "__init__") + + model = ModelFT(max_res=2.5, verbose=0, device="cpu", strip_H=strip_H) + model.load_pdb(pdb_path) + assert engines["n"] == 1 + assert spacegroups["n"] == 1 + + # Later crystal changes are followed by the key, not by rebuilding the engine. + model.cell = model.cell.clone() + model.max_res = 3.0 + assert model.grid_shape is not None + assert engines["n"] == 1 + + +@pytest.mark.unit +def test_engine_reads_the_model_context(pdb_path): + model = _model(pdb_path) + + def bound(m): + return ( + m.fft.ctx is m.ctx + and m.fft.cell is m.ctx.cell + and m.fft.spacegroup is m.ctx.spacegroup + ) + + assert bound(model) + + copied = model.copy() + assert bound(copied) + assert copied.ctx is not model.ctx + assert copied.fft.cell is not model.fft.cell + + selected = model.select("all") + assert bound(selected) + + restored = ModelFT.create_from_state_dict(model.state_dict(), device="cpu", verbose=0) + assert bound(restored) + assert restored.grid_shape == model.grid_shape + + +@pytest.mark.unit +def test_explicit_gridsize_survives_every_path(pdb_path): + explicit = (64, 32, 24) + model = _model(pdb_path, gridsize=explicit) + assert model.grid_shape == explicit + + model.cell = model.cell.clone() + assert model.grid_shape == explicit + + assert model.copy().grid_shape == explicit + assert model.select("all").grid_shape == explicit + + restored = ModelFT.create_from_state_dict(model.state_dict(), device="cpu", verbose=0) + assert restored.explicit_gridsize == explicit + assert restored.grid_shape == explicit + + +@pytest.mark.unit +def test_grid_follows_its_key(pdb_path): + model = _model(pdb_path) + shape0 = model.grid_shape + key0 = model.grid_key + ptr0 = model.fft.gridsize.data_ptr() + + # Same values, different object: nothing to rebuild. + model.cell = model.cell.clone() + assert model.grid_key == key0 + assert model.fft.gridsize.data_ptr() == ptr0 + + scale = torch.tensor([1.25, 1.25, 1.25, 1.0, 1.0, 1.0]) + model.cell = Cell(model.cell.data * scale, dtype=model.dtype_float, device="cpu") + assert model.grid_key != key0 + assert model.grid_shape != shape0 + assert all(n > m for n, m in zip(model.grid_shape, shape0)) + + model.max_res = 4.0 + coarse = model.grid_shape + assert all(n < m for n, m in zip(coarse, model.grid_shape)) is False + model.max_res = 2.5 + assert all(n > m for n, m in zip(model.grid_shape, coarse)) + + +@pytest.mark.unit +def test_first_fcalc_after_cell_reassignment_uses_the_late_path(pdb_path, monkeypatch): + model = _model(pdb_path) + hkl = _hkl(model) + with torch.no_grad(): + f0 = model(hkl).clone() + + model.cell = model.cell.clone() + model.reset_cache() + assert model.fft.late_symmetry_compatible is True + + symmetrised = _count_calls(monkeypatch, SpaceGroup, "symmetrize_map") + with torch.no_grad(): + f1 = model(hkl) + assert symmetrised["n"] == 0 + assert torch.allclose(f1, f0) + + +@pytest.mark.unit +def test_forward_cache_invalidates_on_a_spacegroup_change(pdb_path): + model = _model(pdb_path) + hkl = _hkl(model) + with torch.no_grad(): + f0 = model(hkl) + assert model(hkl) is f0 # cached + assert model._fwd_cached_state_fp[-1] == model.grid_key + + model.spacegroup = "P 1 2 1" # same cell, centring dropped + f1 = model(hkl) + assert f1 is not f0 + assert not torch.allclose(f1, f0) + assert model._fwd_cached_state_fp[-1] == model.grid_key + + +@pytest.mark.unit +def test_to_moves_the_shared_context_once(pdb_path, monkeypatch): + model = _model(pdb_path) + cell = model.ctx.cell + resets = {"n": 0} + original = Cell.reset_cache + + def counting(self): + if self is cell: + resets["n"] += 1 + return original(self) + + monkeypatch.setattr(Cell, "reset_cache", counting) + model.to(torch.device("cpu")) + assert resets["n"] == 1 + + +@pytest.mark.unit +def test_engine_without_a_crystal(): + sf = SfFFT(ModelContext(), max_res=1.0) + assert sf.gridsize is None + assert sf.grid_shape is None + assert sf.grid_key is None + + empty = torch.zeros((0, 3)) + with pytest.raises(RuntimeError, match="no cell or space group"): + sf.compute_structure_factors( + torch.zeros((1, 3), dtype=dtypes.int), + empty, torch.zeros(0), torch.zeros(0), torch.zeros((0, 5)), torch.zeros((0, 5)), + ) + + +@pytest.mark.unit +def test_legacy_gridsize_is_adopted_only_when_it_differs(pdb_path): + model = _model(pdb_path) + + same = model.state_dict() + same["_fft.gridsize"] = torch.tensor(model.grid_shape) + restored = ModelFT.create_from_state_dict(same, device="cpu", verbose=0) + assert restored.explicit_gridsize is None + assert restored.grid_shape == model.grid_shape + + different = model.state_dict() + different["_fft.gridsize"] = torch.tensor([64, 32, 24]) + restored = ModelFT.create_from_state_dict(different, device="cpu", verbose=0) + assert restored.explicit_gridsize == (64, 32, 24) + assert restored.grid_shape == (64, 32, 24) diff --git a/tests/unit/monomer/__init__.py b/tests/unit/monomer/__init__.py new file mode 100644 index 00000000..b7a88075 --- /dev/null +++ b/tests/unit/monomer/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the monomer-dictionary layer under torchref.topology.""" diff --git a/tests/unit/restraints/test_link_modifications.py b/tests/unit/monomer/test_link_modifications.py similarity index 96% rename from tests/unit/restraints/test_link_modifications.py rename to tests/unit/monomer/test_link_modifications.py index 11b29228..63b73854 100644 --- a/tests/unit/restraints/test_link_modifications.py +++ b/tests/unit/monomer/test_link_modifications.py @@ -11,12 +11,12 @@ import pandas as pd import pytest -from torchref.restraints.modifications import ( +from torchref.topology.monomer.modifications import ( apply_modifications, link_modifications, read_mod_definitions, ) -from torchref.restraints.restraints_helper import read_link_definitions +from torchref.topology.monomer.cif import read_link_definitions TOL = 1e-3 @@ -224,10 +224,16 @@ def test_peptide_modifications_carry_the_linked_backbone_targets(): def _built(pdb_path, strip_H=True): - """Build a model's restraints and return ``(model, table accessor)``.""" + """Build a model's restraints and return ``(model, table accessor)``. + + ``add_hydrogens=False``: these tests read the restraint targets of the hydrogens the + file carries. Generating more would add a chain-terminal ``CA-N-H``, which correctly + keeps the free-amino-acid target of 109.6 degrees rather than the linked 118.7 and so + is outside what they assert. + """ from torchref import Model - model = Model(verbose=0, strip_H=strip_H) + model = Model(verbose=0, strip_H=strip_H, add_hydrogens=False) model.load_pdb(str(pdb_path)) return model, model.restraints.restraints diff --git a/tests/unit/restraints/test_restraints.py b/tests/unit/monomer/test_restraints.py similarity index 96% rename from tests/unit/restraints/test_restraints.py rename to tests/unit/monomer/test_restraints.py index 1ec929cf..97ea6c1a 100644 --- a/tests/unit/restraints/test_restraints.py +++ b/tests/unit/monomer/test_restraints.py @@ -1,5 +1,5 @@ """ -Unit tests for torchref.restraints.restraints +Unit tests for torchref.topology.restraints Tests the Restraints class for geometry restraints. Note: Unit tests use mock data, not real file I/O. @@ -17,7 +17,7 @@ class TestRestraintsInitialization: @pytest.mark.unit def test_restraints_empty_init(self): """Test Restraints can be initialized empty.""" - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints restraints = Restraints() @@ -26,7 +26,7 @@ def test_restraints_empty_init(self): @pytest.mark.unit def test_restraints_is_nn_module(self): """Restraints should be a nn.Module.""" - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints restraints = Restraints() @@ -35,7 +35,7 @@ def test_restraints_is_nn_module(self): @pytest.mark.unit def test_restraints_verbose_setting(self): """Test verbosity setting.""" - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints restraints = Restraints(verbose=2) diff --git a/tests/unit/refinement/test_amber_target.py b/tests/unit/refinement/test_amber_target.py index 343ac601..1b88a35a 100644 --- a/tests/unit/refinement/test_amber_target.py +++ b/tests/unit/refinement/test_amber_target.py @@ -1,79 +1,368 @@ -""" -Unit tests for the single-molecule -:class:`~torchref.experimental.targets.amber_target.AmberTarget`. - -Uses a ligand-free protein (7L84) so the standard OpenMM ``Modeller`` path is -exercised — no antechamber/tleap (AmberTools) required, only OpenMM. This -closes the previous coverage gap where the single-molecule target had no unit -test (only manual scripts), which is how it was able to silently rot while the -ensemble targets were developed. -""" +"""AMBER evaluates the model's complete coordinates and returns atomic gradients.""" import os +import numpy as np +import pandas as pd import pytest import torch +from torchref.config import get_int_dtype +from torchref.experimental.targets.amber_target import AmberTarget, _OpenMMAMBERFunction from torchref.model.model import Model -from torchref.experimental.targets.amber_target import AmberTarget -# Ligand-free protein → standard OpenMM Modeller path; needs OpenMM (+ pdbfixer -# from the same [amber] extra), but no AmberTools. Gated centrally in conftest. pytestmark = pytest.mark.openmm - -# Ligand-free protein → standard Modeller path, no antechamber needed. TEST_PDB = os.path.join( os.path.dirname(__file__), "..", "..", "files", "pdb", "7L84.pdb" ) @pytest.fixture(scope="module") -def heavy_model() -> Model: - """Heavy-atom, single-conformation model (OpenMM adds H internally).""" - return Model(verbose=0, strip_H=True).load_pdb(TEST_PDB).strip_altlocs() +def protein(): + """A deposited, hydrogenated protein with methyl orientation parameters.""" + model = Model(verbose=0, device="cpu", add_hydrogens=True).load_pdb(TEST_PDB) + model = model.strip_altlocs() + model.set_hydrogen_mode("riding") + return model + + +@pytest.fixture(scope="module") +def target(protein): + """Build the expensive AMBER context once per module.""" + return AmberTarget(model=protein, verbose=0) + + +def test_build_preserves_model(protein): + """Target construction neither changes model atoms nor replaces coordinates.""" + before = protein.pdb.copy(deep=True) + xyz = protein.xyz + positions = xyz().detach().clone() + target = AmberTarget(model=protein) + pd.testing.assert_frame_equal(protein.pdb, before) + assert protein.xyz is xyz + assert torch.equal(protein.xyz(), positions) + assert target._n_omm_atoms == target._n_model_atoms == len(protein.pdb) + assert np.array_equal(np.sort(target._model_to_omm), np.arange(len(protein.pdb))) + h1 = int(np.flatnonzero(protein.pdb.name.str.strip().eq("H1"))[0]) + assert h1 in protein.xyz.h_row.tolist() + + +def test_forward_and_orientation_gradients(target, protein): + """AMBER hydrogen forces reach the model's methyl torsion parameters.""" + loss = target.forward() + leaves = protein.xyz.optimization_parameters() + gradients = torch.autograd.grad(loss, leaves, allow_unused=True) + assert loss.shape == () and torch.isfinite(loss) + for leaf, gradient in zip(leaves, gradients): + if leaf.numel(): + assert gradient is not None and torch.isfinite(gradient).all() + torsion_grad = torch.autograd.grad( + target.forward(), protein.xyz.torsions.refinable_params + )[0] + assert torsion_grad.abs().sum() > 0 + + +def test_every_position_and_force_has_model_order(target, protein): + """All H positions are supplied live, with force sign, units and normalization.""" + import openmm.unit as unit + + xyz = protein.xyz().detach().clone().requires_grad_() + h_rows = torch.tensor(np.flatnonzero(protein.pdb.element.str.strip().eq("H"))) + with torch.no_grad(): + xyz[h_rows] += xyz.new_tensor([0.003, -0.002, 0.001]) + energy = target._energy(xyz) + gradient = torch.autograd.grad(energy, xyz)[0] + state = target._context.getState(getPositions=True, getForces=True, getEnergy=True) + pos = np.asarray(state.getPositions(asNumpy=True).value_in_unit(unit.nanometer)) + np.testing.assert_allclose( + pos[target._model_to_omm], xyz.detach().numpy() * 0.1, atol=1e-7 + ) + forces = np.asarray( + state.getForces(asNumpy=True).value_in_unit( + unit.kilojoules_per_mole / unit.nanometer + ) + ) + scale = np.minimum( + 10000 / np.maximum(np.linalg.norm(forces, axis=1, keepdims=True), 1e-10), 1 + ) + expected = ( + -forces[target._model_to_omm] * scale[target._model_to_omm] * 0.1 / len(xyz) + ) + np.testing.assert_allclose(gradient.numpy(), expected, rtol=3e-6, atol=1e-5) + assert gradient[h_rows].abs().sum() > 0 + expected_energy = state.getPotentialEnergy().value_in_unit( + unit.kilojoules_per_mole + ) / len(xyz) + assert energy.item() == pytest.approx(expected_energy, rel=1e-6) + + +def test_permuted_positions_and_gradients(target, protein): + """The cached inverse map handles arbitrary model/OpenMM atom permutations.""" + xyz = protein.xyz().detach().clone().requires_grad_() + reference = target._energy(xyz) + reference_grad = torch.autograd.grad(reference, xyz)[0] + perm = torch.arange(len(xyz) - 1, -1, -1) + inverse = torch.argsort(perm) + original = target._omm_to_model.clone() + try: + target._omm_to_model = inverse[original].to(get_int_dtype()) + permuted = xyz.detach()[perm].requires_grad_() + loss = target._energy(permuted) + grad = torch.autograd.grad(loss, permuted)[0] + assert torch.allclose(loss, reference) + assert torch.allclose(grad, reference_grad[perm]) + finally: + target._omm_to_model = original + + +def test_missing_hydrogens_are_not_added(): + """Disabled generation leaves a heavy-only model untouched on AMBER rejection.""" + model = ( + Model(verbose=0, device="cpu", strip_H=True, add_hydrogens=False) + .load_pdb(TEST_PDB) + .strip_altlocs() + ) + before = model.pdb.copy(deep=True) + wrapper = model.xyz + with pytest.raises(ValueError, match="Prepare missing atoms"): + AmberTarget(model=model) + pd.testing.assert_frame_equal(model.pdb, before) + assert model.xyz is wrapper @pytest.fixture(scope="module") -def target(heavy_model) -> AmberTarget: - """Built once — the OpenMM system construction is the expensive part.""" - return AmberTarget(model=heavy_model, verbose=0) - - -def test_build_populates_state(target, heavy_model): - """Construction builds a usable OpenMM context + atom map + H tables.""" - assert target._context is not None - assert target._system is not None - assert target._n_model_atoms == len(heavy_model.pdb) - assert target._n_omm_atoms >= target._n_model_atoms # H added by Modeller - # For the single molecule the chemistry model IS the model. - assert target._chem_model is target._model - # H-attachment tables were built (this is a protein → has H). - assert target._h_idx is not None and target._h_idx.size > 0 - - -def test_forward_finite_energy(target): - """forward() returns a finite scalar energy.""" - e = target.forward() - assert e.shape == () - assert torch.isfinite(e).item() - - -def test_backward_flows_to_xyz(heavy_model): - """Energy gradient propagates to the model's xyz parameters and is finite.""" - target = AmberTarget(model=heavy_model, verbose=0) - heavy_model.xyz.refinable_params.grad = None - e = target.forward() - e.backward() - g = heavy_model.xyz.refinable_params.grad - assert g is not None - assert g.shape == (len(heavy_model.pdb), 3) - assert torch.isfinite(g).all().item() - assert g.abs().sum().item() > 0.0 # non-trivial forces - - -def test_default_charge_method_is_gas(): - """The unified base defaults to the (robust) Gasteiger charge method.""" - import inspect - - sig = inspect.signature(AmberTarget.__init__) - assert sig.parameters["charge_method"].default == "gas" +def water_target(pdb_dir): + """Two nearby deposited waters with TorchRef-generated rotating hydrogens.""" + model = Model(verbose=0, device="cpu", add_hydrogens=False).load_pdb( + str(pdb_dir / "1DAW.pdb") + ) + waters = model.pdb[model.pdb.resname.str.strip().eq("HOH")].copy() + coords = waters[["x", "y", "z"]].to_numpy() + distances = np.linalg.norm(coords[:, None] - coords[None, :], axis=-1) + np.fill_diagonal(distances, np.inf) + i, j = np.unravel_index(np.argmin(distances), distances.shape) + model = model._new_model_from_df(waters.iloc[sorted([i, j])].copy(), strip_H=False) + model.ctx.add_hydrogens = True + with torch.random.fork_rng(): + torch.manual_seed(42) + model.set_hydrogen_mode("riding") + return AmberTarget(model=model, normalize_by_atoms=False) + + +def test_water_rotation_changes_amber_energy(water_target): + """Rotating water H changes AMBER energy while oxygens remain fixed.""" + target = water_target + model = target._model + rotations = model.xyz.rotations.refinable_params + before = rotations.detach().clone() + positions = model.xyz().detach().clone() + try: + initial_energy = target.forward().detach() + with torch.no_grad(): + rotations[0] += rotations.new_tensor([0.2, -0.3, 0.4]) + energy = target.forward() + grad = torch.autograd.grad(energy, rotations)[0] + assert not torch.allclose(energy, initial_energy) + assert torch.isfinite(grad).all() and grad.abs().sum() > 0 + assert torch.equal( + model.xyz()[model.xyz.base_row], positions[model.xyz.base_row] + ) + finally: + with torch.no_grad(): + rotations.copy_(before) + + +def test_unclipped_water_rotation_derivative(water_target): + """AMBER's force bridge agrees with an orientation finite difference in float32.""" + target = water_target + wrapper = target._model.xyz + base = wrapper._storage_values().detach() + torsions = wrapper.torsions().detach() + rotation = wrapper.rotations().detach().requires_grad_() + + def energy(rot): + xyz = wrapper.evaluate(base, torsions, rot) + return _OpenMMAMBERFunction.apply( + target._compose_full_omm_xyz(xyz), target._context, float("inf") + ) + + gradient = torch.autograd.grad(energy(rotation), rotation)[0] + direction = gradient.detach() / gradient.norm() + step = 1e-3 + finite = ( + energy(rotation + step * direction) - energy(rotation - step * direction) + ) / (2 * step) + assert finite.item() == pytest.approx( + (gradient * direction).sum().item(), rel=0.015, abs=0.1 + ) + + +def test_atom_count_change_requires_rebuild(target, protein): + """A target rejects coordinates from an expanded or reduced atom table.""" + with pytest.raises(ValueError, match="Atom count changed"): + target._energy(protein.xyz()[:-1]) + + +@pytest.mark.parametrize("ligand_instances", [False, True]) +def test_gaff2_mapping_preserves_residue_instances(water_target, ligand_instances): + """Residue and atom permutations cannot exchange two identical molecules.""" + import openmm as mm + import openmm.app as app + + model = water_target._model + source = list(water_target._topology.residues()) + keys = list( + dict.fromkeys( + model.pdb[["chainid", "resseq", "icode"]].itertuples(index=False, name=None) + ) + ) + topology = app.Topology() + chain = topology.addChain("renumbered") + positions = [] + expected = [] + for residue in reversed(source): + dest = topology.addResidue(residue.name, chain, str(len(positions) + 10)) + mapped = {} + for atom in reversed(list(residue.atoms())): + mapped[atom.index] = topology.addAtom(atom.name, atom.element, dest) + row = int(water_target._omm_to_model[atom.index]) + expected.append(row) + positions.append(model.xyz()[row].detach().numpy() * 0.1) + for a, b in water_target._topology.bonds(): + if a.index in mapped and b.index in mapped: + topology.addBond(mapped[a.index], mapped[b.index]) + candidate = AmberTarget() + candidate._chem_model = model + candidate._topology = topology + candidate._system = mm.System() + for _ in expected: + candidate._system.addParticle(1) + candidate._tleap_residue_map = True + candidate._gaff2_residue_keys = list(reversed(keys)) if ligand_instances else [] + candidate._tleap_pos_nm = np.asarray(positions) + if ligand_instances: + candidate._tleap_pos_nm += 1000 + candidate._build_atom_map() + assert candidate._omm_to_model.tolist() == expected + + +def test_duplicate_particle_mapping_is_rejected(water_target): + """A duplicate source index cannot silently send two forces to one atom.""" + candidate = AmberTarget() + candidate._chem_model = water_target._model + candidate._topology = water_target._topology + candidate._system = water_target._system + candidate._source_model_rows = np.zeros(water_target._n_model_atoms, dtype=np.int32) + with pytest.raises(ValueError, match="not one-to-one"): + candidate._build_atom_map() + + +def test_supercell_gather_keeps_live_hydrogens(water_target): + """The ensemble path transfers all transformed hydrogen rows unchanged.""" + from torchref.experimental.ensemble.quasi_crystal_amber import ( + QuasiCrystalAmberTarget, + ) + + candidate = QuasiCrystalAmberTarget.__new__(QuasiCrystalAmberTarget) + torch.nn.Module.__init__(candidate) + n_atoms = water_target._n_model_atoms + candidate._n_members = 2 + candidate._n_model_per_member = n_atoms + candidate._omm_to_model = torch.arange(n_atoms - 1, -1, -1, dtype=get_int_dtype()) + xyz = water_target._model.xyz().detach().unsqueeze(0).repeat(2, 1, 1) * 0.1 + xyz[1] += 0.3 + xyz.requires_grad_() + result = candidate._compose_full_omm_xyz(xyz).reshape(2, n_atoms, 3) + assert torch.equal(result.flip(1), xyz) + weights = torch.arange(result.numel(), dtype=xyz.dtype).reshape_as(result) + gradient = torch.autograd.grad((result * weights).sum(), xyz)[0] + assert torch.equal(gradient, weights.flip(1)) + + +@pytest.mark.parametrize("use_reference_dtype", [False, True]) +def test_bridge_preserves_input_dtype(water_target, use_reference_dtype, double_cpu): + """The OpenMM boundary preserves configured reference and primary dtypes.""" + from torchref.config import get_float_dtype + + dtype = ( + get_float_dtype() if use_reference_dtype else water_target._model.xyz().dtype + ) + xyz = water_target._model.xyz().detach().to(dtype=dtype).requires_grad_() + loss = water_target._energy(xyz) + gradient = torch.autograd.grad(loss, xyz)[0] + assert loss.dtype == gradient.dtype == xyz.dtype + assert torch.isfinite(gradient).all() + + +def test_bridge_follows_coordinate_device(water_target, any_device): + """Coordinates and gradients stay on the caller's device across the CPU bridge.""" + xyz = water_target._model.xyz().detach().to(any_device).requires_grad_() + loss = water_target._energy(xyz) + gradient = torch.autograd.grad(loss, xyz)[0] + assert loss.device == gradient.device == xyz.device + assert torch.isfinite(gradient).all() + + +def test_torchref_hydrogenation_prepares_compatible_protein(): + """TorchRef can prepare all model hydrogens before AMBER construction.""" + model = ( + Model(verbose=0, device="cpu", strip_H=True, add_hydrogens=False) + .load_pdb(TEST_PDB) + .strip_altlocs() + .hydrogenate() + ) + model.set_hydrogen_mode("riding") + positions = model.xyz().detach().clone() + target = AmberTarget(model=model) + assert target._n_omm_atoms == len(model.pdb) + assert torch.equal(model.xyz(), positions) + assert torch.isfinite(target.forward()) + + +def test_context_initialization_uses_live_model(water_target, monkeypatch): + """Backend template coordinates cannot replace the model's live positions.""" + candidate = AmberTarget() + candidate._chem_model = water_target._model + candidate._source_model_rows = water_target._omm_to_model.cpu().numpy() + stale_positions = np.zeros_like(water_target._pos_buf) + monkeypatch.setattr( + candidate, + "_build_omm_system", + lambda params: (water_target._system, water_target._topology, stale_positions), + ) + captured = [] + monkeypatch.setattr( + candidate, "_build_context", lambda positions: captured.append(positions.copy()) + ) + candidate._build() + expected = ( + water_target._compose_full_omm_xyz(water_target._model.xyz()) + .detach() + .cpu() + .numpy() + ) + np.testing.assert_array_equal(captured[0], expected) + np.testing.assert_array_equal(candidate._pos_buf, expected) + + +def test_partial_terminal_hydrogens_preserve_h1_alias(protein): + """Completing an existing terminal H1 adds H2/H3 without an equivalent H.""" + pdb = protein.pdb + first = pdb.iloc[0] + residue = ( + (pdb.chainid == first.chainid) + & (pdb.resseq == first.resseq) + & (pdb.icode == first.icode) + ) + missing = residue & pdb.name.str.strip().isin(["H2", "H3"]) + partial = protein._new_model_from_df(pdb.loc[~missing].copy(), strip_H=False) + prepared = partial.hydrogenate() + first_residue = prepared.pdb[ + (prepared.pdb.chainid == first.chainid) & (prepared.pdb.resseq == first.resseq) + ] + names = set(first_residue.name.str.strip()) + assert {"H1", "H2", "H3"} <= names + assert "H" not in names + assert len(prepared.pdb) == len(protein.pdb) + target = AmberTarget(model=prepared) + assert torch.isfinite(target.forward()) diff --git a/tests/unit/refinement/test_collection_sigma_d_target.py b/tests/unit/refinement/test_collection_sigma_d_target.py new file mode 100644 index 00000000..dddb57d2 --- /dev/null +++ b/tests/unit/refinement/test_collection_sigma_d_target.py @@ -0,0 +1,126 @@ +"""The ``difference_sd`` collection row on a real dark/light pair. + +Pinned on 1DAW with a 0.2 A shifted light model: the loss is finite, gradients reach +the light model through ``dF_calc`` only, the sigma_D estimate is owned by the target, +fitted on free reflections of the timepoint row, cached across forwards and cleared +by ``maintenance()``, and the fit summary reaches ``stats()``. +""" + +import pytest +import torch + +from torchref.refinement.model_error_estimation import sigma_d as sigma_d_module +from torchref.refinement.model_error_estimation.sigma_d import SigmaDEstimator +from torchref.refinement.targets.collection import ( + CollectionDifferenceSigmaDTarget, + CollectionSigmaDLossInputs, +) + +pytestmark = pytest.mark.integration + + +@pytest.fixture +def target(loaded_reflection_data, sample_structure_pair): + """A dark/light collection whose light amplitudes carry a resolution-dependent + difference proportional to ``F``, so the sigma_D coupling is not zero, and a light + model shifted by 0.2 A.""" + from torchref import ReflectionData + from torchref.cli._common import load_model + from torchref.io import DatasetCollection + from torchref.model import ModelCollection + from torchref.scaling import CollectionScaler + + data = loaded_reflection_data + g = torch.Generator().manual_seed(11) + f = data.F + dss = 1.0 / data.resolution**2 + change = ( + 0.08 * f * torch.exp(-2.0 * dss) * torch.randn(len(data), generator=g).to(f) + ) + light = ReflectionData.from_tensors( + hkl=data.hkl, + F=(f + change).clamp(min=0.0), + F_sigma=data.F_sigma, + cell=data.cell, + spacegroup=data.spacegroup, + rfree_flags=data.rfree_flags, + device=str(data.device), + verbose=0, + ) + models = [ + load_model( + str(sample_structure_pair["model"]), + max_res=2.05, + device=data.device, + verbose=0, + ) + for _ in range(2) + ] + with torch.no_grad(): + models[1].xyz.refinable_params += 0.2 + dc = DatasetCollection(device=data.device, verbose=0) + dc.add_dataset("dark", data, set_as_reference=True).add_dataset("light", light) + mc = ModelCollection(models, dark_key="dark", verbose=0) + mc.add_dark().add_timepoint("light", [0.78, 0.22]) + scaler = CollectionScaler(dc, mc, verbose=0).initialize() + return dc, mc, CollectionDifferenceSigmaDTarget(dc, mc, scaler=scaler) + + +def test_forward_is_finite_and_owns_its_estimator(target): + _dc, _mc, t = target + assert isinstance(t._sigma_d, SigmaDEstimator) + loss = t.forward() + assert torch.isfinite(loss) + ctx = t._loss_inputs() + assert isinstance(ctx, CollectionSigmaDLossInputs) + assert ctx.alpha.shape == ctx.beta_model.shape == (ctx.obs.shape[1],) + assert not ctx.alpha.requires_grad and not ctx.beta_model.requires_grad + assert (ctx.beta_model >= 0).all() + shells = t._sigma_d.shells + assert shells.has_model and not shells.all_zero and not shells.degenerate + + +def test_gradient_reaches_the_light_model(target): + _dc, mc, t = target + light = mc.base_models[1] + light.zero_grad(set_to_none=True) + t.forward().backward() + grads = [p.grad for p in light.parameters() if p.grad is not None] + assert grads and any(torch.isfinite(g).all() and g.abs().sum() > 0 for g in grads) + + +def test_estimate_is_cached_until_maintenance(target): + _dc, _mc, t = target + t.forward() + assert t._sigma_d._cache is not None + first = t._sigma_d._cache + t.forward() + assert t._sigma_d._cache is first + t.maintenance() + assert t._sigma_d._cache is None + + +def test_fit_uses_free_reflections_of_the_timepoint_row(target, monkeypatch): + dc, _mc, t = target + seen = {} + real = sigma_d_module.estimate_sigma_d + + def spy(delta_obs, sigma_diff, epsilon, d_star_sq, f_dark, fit_mask, **kw): + seen["fit_mask"] = fit_mask.clone() + seen["n"] = delta_obs.numel() + return real(delta_obs, sigma_diff, epsilon, d_star_sq, f_dark, fit_mask, **kw) + + monkeypatch.setattr(sigma_d_module, "estimate_sigma_d", spy) + t.forward() + n_hkl = dc.hkl.shape[0] + assert seen["n"] == n_hkl + free = dc["light"].free.mask.to(seen["fit_mask"].device) + assert bool((seen["fit_mask"] & ~free).sum() == 0) + assert int(seen["fit_mask"].sum()) > 0 + + +def test_stats_carry_the_fit_summary(target): + _dc, _mc, t = target + t.forward() + stats = t.stats() + assert "sigma_d_gamma" in stats and "sigma_d_tau" in stats diff --git a/tests/unit/refinement/test_collection_taxonomy.py b/tests/unit/refinement/test_collection_taxonomy.py new file mode 100644 index 00000000..c0956b82 --- /dev/null +++ b/tests/unit/refinement/test_collection_taxonomy.py @@ -0,0 +1,88 @@ +"""Collection target registry classes, observables and invalid selections.""" + +import inspect + +import pytest + + +@pytest.mark.unit +def test_each_row_has_its_own_class(): + """One class per selectable row, and the test and the table must agree on the set.""" + from torchref.refinement.targets.collection import COLLECTION_XRAY_TARGETS + from torchref.refinement.targets.collection.intensity import ( + CollectionTwoMomentIntensityTarget, + ) + from torchref.refinement.targets.collection.xray import ( + CollectionDifferenceIntensityTarget, + CollectionDifferenceSigmaDTarget, + CollectionDifferenceTarget, + CollectionMLTarget, + ) + + expected = { + "difference": CollectionDifferenceTarget, + "difference_i": CollectionDifferenceIntensityTarget, + "difference_sd": CollectionDifferenceSigmaDTarget, + "two_moment": CollectionTwoMomentIntensityTarget, + "ml": CollectionMLTarget, + } + assert set(expected) == set( + COLLECTION_XRAY_TARGETS.names + ), "table and test disagree on the rows" + for name, cls in expected.items(): + assert COLLECTION_XRAY_TARGETS.by_name(name).target_cls is cls, name + + seen = {} + for name in COLLECTION_XRAY_TARGETS.names: + cls = COLLECTION_XRAY_TARGETS.by_name(name).target_cls + assert cls not in seen, f"{name} and {seen[cls]} share {cls.__name__}" + seen[cls] = name + + +@pytest.mark.unit +def test_the_observable_is_declared_not_passed(): + """The spec's claim and the class's own attribute must agree.""" + from torchref.refinement.targets.collection import ( + COLLECTION_XRAY_TARGETS, + CollectionXrayTargetSpec, + ) + from torchref.refinement.targets.collection.xray import CollectionDifferenceTarget + + by_obs = {} + for spec in COLLECTION_XRAY_TARGETS.specs: + assert spec.observable in ("amplitude", "intensity"), spec.name + assert getattr(spec.target_cls, "observable", "amplitude") == spec.observable + by_obs.setdefault(spec.observable, []).append(spec.name) + assert ( + "observable" not in inspect.signature(spec.target_cls.__init__).parameters + ) + + assert set(by_obs["intensity"]) == {"difference_i", "two_moment"} + assert set(by_obs["amplitude"]) == {"difference", "difference_sd", "ml"} + + with pytest.raises(ValueError, match="observable"): + CollectionXrayTargetSpec( + name="bogus", + target_cls=CollectionDifferenceTarget, + doc="", + observable="intensity", + ) + + +@pytest.mark.unit +def test_there_is_no_intensity_rice_row(): + """Rice is amplitude-only by nature, so the axis is not square.""" + from torchref.refinement.targets.collection import COLLECTION_XRAY_TARGETS + from torchref.refinement.targets.collection.xray import CollectionMLTarget + + for spec in COLLECTION_XRAY_TARGETS.specs: + if spec.observable == "intensity": + assert not issubclass(spec.target_cls, CollectionMLTarget), spec.name + + +@pytest.mark.unit +def test_unknown_rows_fail_closed(): + from torchref.refinement.targets.collection import COLLECTION_XRAY_TARGETS + + with pytest.raises(ValueError, match="Unknown collection X-ray target"): + COLLECTION_XRAY_TARGETS.by_name("no_such_row") diff --git a/tests/unit/refinement/test_loss_state.py b/tests/unit/refinement/test_loss_state.py index c6126464..dba1b464 100644 --- a/tests/unit/refinement/test_loss_state.py +++ b/tests/unit/refinement/test_loss_state.py @@ -8,6 +8,66 @@ import torch +@pytest.mark.unit +@pytest.mark.parametrize("mode", ["empty", "disabled", "disabled_compilable"]) +@pytest.mark.parametrize("log_values", [False, True]) +def test_zero_aggregate_uses_configured_dtype_and_device( + mode: str, log_values: bool +) -> None: + """An aggregate without active targets is a configured scalar zero.""" + from torchref.config import get_default_device, get_float_dtype + from torchref.refinement.loss_state import LossState + + state = LossState() + if mode != "empty": + + def disabled_target(): + pytest.fail("A zero-weight target must not be evaluated") + + state.register_target( + "geometry/bond", + disabled_target, + compile=mode == "disabled_compilable", + probe=False, + ) + state.set_weight("geometry", 0.0) + state.compile_aggregate() + + total = state.aggregate(log_values=log_values) + + expected = torch.zeros((), dtype=get_float_dtype(), device=get_default_device()) + torch.testing.assert_close(total, expected) + assert state._losses == {} + assert state.history == ([{"total": 0.0}] if log_values else []) + + +@pytest.mark.unit +@pytest.mark.parametrize("compiled", [False, True]) +def test_aggregate_ignores_torch_default_dtype( + monkeypatch: pytest.MonkeyPatch, compiled: bool +) -> None: + """Eager and compiled sums use TorchRef's dtype, not PyTorch's default.""" + from torchref.config import device, dtypes, get_float_dtype + from torchref.refinement.loss_state import LossState + + monkeypatch.setattr(device, "current", torch.device("cpu")) + monkeypatch.setattr(dtypes, "float", torch.float32) + state = LossState() + value = torch.tensor(2.0, dtype=get_float_dtype(), device=state.device) + state.register_target("geometry/bond", lambda: value, compile=compiled) + state.set_weight("geometry", 3.0) + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.float64) + if compiled: + state.compile_aggregate(backend="eager") + total = state.aggregate() + finally: + torch.set_default_dtype(previous_dtype) + + torch.testing.assert_close(total, value * 3.0) + + class TestLossStateBasic: """Tests for basic LossState functionality.""" @@ -44,7 +104,9 @@ def test_register_target(self): from torchref.refinement.loss_state import LossState state = LossState() - target_fn = lambda: torch.tensor(1.0) + + def target_fn(): + return torch.tensor(1.0) result = state.register_target("geometry/bond", target_fn) @@ -121,7 +183,7 @@ def test_prefix_with_hierarchical_weighting(self): total = state.aggregate() # Expected: 0.5 * 1.0 + 1.0 * 2.0 = 2.5 - assert torch.isclose(total, torch.tensor(2.5)) + assert total.item() == pytest.approx(2.5) class TestWeightManagement: @@ -136,6 +198,7 @@ def test_set_weight(self): result = state.set_weight("geometry", 0.5) assert state.weights["geometry"] == 0.5 + assert state.get_weight("geometry") == 0.5 assert result is state # Method chaining @pytest.mark.unit @@ -184,12 +247,11 @@ def test_get_effective_weight_hierarchical(self): from torchref.refinement.loss_state import LossState state = LossState() - state.set_weight("geometry", 0.5) - state.set_weight("geometry/bond", 2.0) + state.set_weight("geometry", 2.0) + state.set_weight("geometry/bond", 3.0) - # geometry/bond -> geometry (0.5) * geometry/bond (2.0) = 1.0 effective = state.get_effective_weight("geometry/bond") - assert effective == 1.0 + assert effective == 6.0 @pytest.mark.unit def test_get_effective_weight_missing_intermediate(self): @@ -207,6 +269,22 @@ def test_get_effective_weight_missing_intermediate(self): class TestAggregation: """Tests for loss aggregation.""" + @pytest.mark.unit + def test_zero_weight(self): + """A zero weight contributes zero to the aggregate.""" + from torchref.config import get_default_device, get_float_dtype + from torchref.refinement.loss_state import LossState + + value = torch.tensor( + 100.0, dtype=get_float_dtype(), device=get_default_device() + ) + state = LossState() + state.register_target("adp", lambda: value) + state.set_weight("adp", 0.0) + total = state.aggregate() + assert total.ndim == 0 + assert total.item() == 0.0 + @pytest.mark.unit def test_aggregate_simple(self): """Test simple aggregation.""" @@ -221,7 +299,7 @@ def test_aggregate_simple(self): total = state.aggregate(log_values=False) # 2.0 * 1.0 + 1.0 * 0.5 = 2.5 - assert torch.isclose(total, torch.tensor(2.5)) + assert total.item() == pytest.approx(2.5) @pytest.mark.unit def test_aggregate_hierarchical(self): @@ -240,7 +318,7 @@ def test_aggregate_hierarchical(self): # geometry/bond: 1.0 * 0.5 * 2.0 = 1.0 # geometry/angle: 2.0 * 0.5 * 1.0 = 1.0 # total = 2.0 - assert torch.isclose(total, torch.tensor(2.0)) + assert total.item() == pytest.approx(2.0) @pytest.mark.unit def test_aggregate_default_weights(self): @@ -255,20 +333,30 @@ def test_aggregate_default_weights(self): total = state.aggregate(log_values=False) # 2.0 * 1.0 + 1.0 * 1.0 = 3.0 - assert torch.isclose(total, torch.tensor(3.0)) + assert total.item() == pytest.approx(3.0) @pytest.mark.unit def test_aggregate_caches_losses(self): """Test that aggregate caches computed losses.""" + from torchref.config import get_default_device, get_float_dtype from torchref.refinement.loss_state import LossState state = LossState() - state.register_target("xray", lambda: torch.tensor(2.0)) + value = torch.tensor(2.0, dtype=get_float_dtype(), device=get_default_device()) + calls = 0 - state.aggregate(log_values=False) + def target(): + nonlocal calls + calls += 1 + return value - loss = state.get_loss("xray") - assert torch.isclose(loss, torch.tensor(2.0)) + state.register_target("xray", target) + # Registration probes the autograd graph; count only subsequent evaluations. + calls = 0 + state.aggregate(log_values=False) + assert calls == 1 + torch.testing.assert_close(state.get_loss("xray"), value) + assert calls == 1 class TestHistoryLogging: diff --git a/tests/unit/refinement/test_loss_weighting.py b/tests/unit/refinement/test_loss_weighting.py index 1f392822..d314e624 100644 --- a/tests/unit/refinement/test_loss_weighting.py +++ b/tests/unit/refinement/test_loss_weighting.py @@ -1,55 +1,6 @@ -""" -Unit tests for LossState weight handling. - -Covers the retained ``LossState`` weight API (``set_weight`` / -``get_effective_weight`` / ``aggregate``). The standalone weighting -schemes were removed; refinement now aggregates at uniform weight by -default, with explicit per-target/group multipliers set via the -``LossState`` weight dict. -""" +"""Pin refinement's default group weights; LossState owns weight arithmetic.""" import pytest -import torch - - -class TestLossStateWeights: - """Tests for the LossState weight dict (hierarchical multipliers).""" - - @pytest.mark.unit - def test_hierarchical_weights_multiply(self): - """Test that hierarchical weights multiply in get_effective_weight.""" - from torchref.refinement.loss_state import LossState - - state = LossState() - # Set group weight - state.set_weight('geometry', 2.0) - # Set component weight - state.set_weight('geometry/bond', 3.0) - - # Effective weight should multiply: 2.0 * 3.0 = 6.0 - effective = state.get_effective_weight('geometry/bond') - assert effective == 6.0 - - -class TestTotalLossFromState: - """Tests for computing total loss from LossState.""" - - @pytest.mark.unit - def test_total_weighted_loss(self): - """Test computing total weighted loss from state via aggregate.""" - from torchref.refinement.loss_state import LossState - - state = LossState() - state.register_target('xray', lambda: torch.tensor(2.0)) - state.register_target('bond', lambda: torch.tensor(1.0)) - state.set_weight('xray', 1.0) - state.set_weight('bond', 0.5) - - total = state.aggregate(log_values=False) - - # Expected: 2.0 * 1.0 + 1.0 * 0.5 = 2.5 - expected = torch.tensor(2.5) - assert torch.isclose(total, expected) class TestDefaultGroupWeights: @@ -68,4 +19,10 @@ def test_default_group_weights_values(self): # Sub-weight on the SIGD prior; 1.0 leaves it at the adp group weight # pending the R_free scan. Weights multiply down the path. "adp/sigd": 1.0, + # Node load balancing, inert unless the ADPs are a node field. Above the + # group weight because it bars a degenerate direction rather than + # competing with the data. + "adp/node_load": 10.0, + # Magnitude prior on node values, off pending measurement. + "adp/node_smoothness": 0.0, } diff --git a/tests/unit/refinement/test_ml_sigmaa.py b/tests/unit/refinement/test_ml_sigmaa.py index eb2190ec..acf6dbfc 100644 --- a/tests/unit/refinement/test_ml_sigmaa.py +++ b/tests/unit/refinement/test_ml_sigmaa.py @@ -153,17 +153,15 @@ class TestCollectionTargetsRelocated: def test_exported_from_refinement_targets(self): from torchref.refinement.targets import ( # noqa: F401 + COLLECTION_XRAY_TARGETS, + CollectionDifferenceIntensityTarget, CollectionDifferenceTarget, CollectionMLTarget, - CollectionRiceTarget, MultiModelADPTarget, MultiModelGeometryTarget, ) def test_kinetic_backcompat_reexports(self): - from torchref.experimental.kinetic.targets import ( # noqa: F401 - CollectionRiceTarget, - ) from torchref.experimental.kinetic.targets import CollectionMLTarget as KinCML from torchref.experimental.kinetic.targets import ( # noqa: F401 KineticPriorTarget, @@ -173,13 +171,30 @@ def test_kinetic_backcompat_reexports(self): assert RefCML is KinCML - def test_collection_ml_base_weight(self): + def test_collection_ml_has_its_maintenance_hook(self): + """``LossState`` calls it after each step block to drop the shared beta.""" from torchref.refinement.targets import CollectionMLTarget - assert CollectionMLTarget.DEFAULT_BASE_WEIGHT == 10.0 - # maintenance hook present (resets the target's own shared beta) assert hasattr(CollectionMLTarget, "maintenance") + def test_the_sigma_obs_in_a_rice_sigma_row_is_gone(self): + """``CollectionRiceTarget`` set ``beta = sigma_obs**2``. + + That pairs a measurement sigma with a Rice ``Sigma``, asserting an isotropic + *complex* error where ``sigma_obs`` carries no phase at all. The single-dataset + table refuses to offer such a row + (``test_nll_beta.py::test_rice_with_sigma_obs_is_not_offered``); the collection + table now agrees, and ``CollectionMLTarget`` -- one shared Luzzati beta -- is the + absolute channel instead. + """ + import torchref.refinement.targets as T + from torchref.refinement.targets.collection import COLLECTION_XRAY_TARGETS + + assert not hasattr(T, "CollectionRiceTarget") + assert "rice" not in COLLECTION_XRAY_TARGETS.names + with pytest.raises(ValueError, match="Unknown collection X-ray target"): + COLLECTION_XRAY_TARGETS.by_name("rice") + @pytest.mark.unit class TestBetaMath: diff --git a/tests/unit/refinement/test_nll_beta.py b/tests/unit/refinement/test_nll_beta.py index 59eb120a..832765e2 100644 --- a/tests/unit/refinement/test_nll_beta.py +++ b/tests/unit/refinement/test_nll_beta.py @@ -216,6 +216,7 @@ def test_each_mode_has_its_own_class(): MLNoAlphaXrayTarget, MLXrayTarget, NLLBetaXrayTarget, + NLLIntensityXrayTarget, NLLXrayTarget, UnitWeightK1XrayTarget, ) @@ -227,6 +228,7 @@ def test_each_mode_has_its_own_class(): "ml_full": MLFullXrayTarget, "nll_beta": NLLBetaXrayTarget, "nll": NLLXrayTarget, + "nll_i": NLLIntensityXrayTarget, "ls": LeastSquaresXrayTarget, "ls_wunit_k1": UnitWeightK1XrayTarget, } @@ -244,6 +246,60 @@ def test_each_mode_has_its_own_class(): seen[cls] = name +def test_the_observable_is_a_row_property_not_a_flag(): + """The observable is declared by the spec AND by the class, and they must agree. + + Same argument as ``test_mean_centring_is_intrinsic_to_the_spec_not_a_flag``: a + constructor kwarg would be a runtime branch, and would be silently dropped by the + construction sites that bypass ``Refinement._xray_target_kwargs`` (the ensemble + refinement has three of them). A row cannot be selected without selecting its class. + """ + import inspect + + from torchref.refinement.targets.xray import NLLIntensityXrayTarget, NLLXrayTarget + from torchref.refinement.targets.xray._specs import XRAY_TARGETS, XrayTargetSpec + + by_obs = {} + for spec in XRAY_TARGETS.specs: + assert spec.observable in ("amplitude", "intensity"), spec.name + # The spec's claim and the class's own declaration must match -- a row advertising + # intensities while reading `sub.F` would be wrong by 2|F| with nothing to catch it. + assert getattr(spec.target_cls, "observable", "amplitude") == spec.observable + by_obs.setdefault(spec.observable, []).append(spec.name) + + assert "nll_i" in by_obs["intensity"], "the intensity row went missing from the table" + assert "nll" in by_obs["amplitude"] + + # Not a constructor flag anywhere. + for cls in (NLLXrayTarget, NLLIntensityXrayTarget): + assert "observable" not in inspect.signature(cls.__init__).parameters + + # The spec rejects a class/spec mismatch rather than trusting either side. + with pytest.raises(ValueError, match="observable"): + XrayTargetSpec( + name="bogus", target_cls=NLLXrayTarget, doc="", observable="intensity" + ) + + +def test_rice_has_no_intensity_row(): + """Rice is amplitude-only *by nature*, so the observable axis is not square. + + Rice and the folded normal are distributions of an amplitude; the intensity analogue is + the exponential / chi-square_1 Wilson distribution, a different primitive. If a future + row pairs a Rice class with ``observable="intensity"`` it is a modelling error, not a + new feature -- so pin it here rather than discovering it from a bad refinement. + """ + from torchref.refinement.targets.xray import RiceXrayTarget, SigmaAXrayTarget + from torchref.refinement.targets.xray._specs import XRAY_TARGETS + + for spec in XRAY_TARGETS.specs: + if spec.observable != "intensity": + continue + assert not issubclass(spec.target_cls, (RiceXrayTarget, SigmaAXrayTarget)), ( + f"{spec.name} pairs an amplitude distribution with intensities" + ) + + def test_only_the_estimator_backed_rows_own_an_estimator(): """``nll`` must not pay for a model-error estimate it does not use. diff --git a/tests/unit/refinement/test_node_load_target.py b/tests/unit/refinement/test_node_load_target.py new file mode 100644 index 00000000..6ec6a3e8 --- /dev/null +++ b/tests/unit/refinement/test_node_load_target.py @@ -0,0 +1,204 @@ +"""The node load-balancing barrier. + +Two properties carry the design. It must be **one-sided** -- an abandoned node is +penalised, an over-loaded one is not -- because the symmetric choice (maximising load +entropy) is optimal at uniform load and would flatten the multiscale kernel-width spread +that a working field genuinely has. And it must act through the *weights* only, so its +gradient reaches node positions and widths but never the node values: it removes the +opportunity to place an extreme B rather than penalising the B. +""" + +import pytest +import torch + +from torchref.model.model import Model +from torchref.refinement.targets.adp import NodeLoadTarget + + +@pytest.fixture(scope="module") +def pdb_path(pdb_dir): + return str(pdb_dir / "3GR5.pdb") + + +def _field_model(pdb_path, n_nodes=32, **kw): + model = Model(verbose=0) + model.load_pdb(pdb_path) + model.set_adp_mode("field", n_nodes=n_nodes, k_neighbors=8, **kw) + return model + + +@pytest.mark.unit +def test_inert_outside_field_mode(pdb_path): + """Registered unconditionally, so it must cost nothing on the per-atom path.""" + model = Model(verbose=0) + model.load_pdb(pdb_path) + model.set_adp_mode("isotropic") + target = NodeLoadTarget(model) + assert float(target()) == 0.0 + assert target.stats()["node_load_active"].value == 0.0 + + +@pytest.mark.unit +def test_balanced_field_is_barely_penalised(pdb_path): + """A freshly fitted field has near-even load, so the barrier starts near its floor.""" + model = _field_model(pdb_path) + target = NodeLoadTarget(model) + rel = target._relative_load().detach() + # Mean relative load is 1 by construction. + assert torch.allclose(rel.mean(), torch.ones((), dtype=rel.dtype), atol=1e-6) + per_node = float(target()) / rel.numel() + assert per_node < 0.5, f"per-node penalty {per_node:.3f} on a balanced field" + + +@pytest.mark.unit +def test_abandoning_a_node_raises_the_penalty(pdb_path): + """Collapsing one node's kernel starves it, and the barrier must notice.""" + model = _field_model(pdb_path) + target = NodeLoadTarget(model) + before = float(target()) + + # Narrow one node far below the others: it loses every softmax contest except + # against an atom sitting on top of it. + with torch.no_grad(): + model.adp.refinable_params[0, 1] -= 6.0 + model.adp.reset_forward_cache() + + after = float(target()) + rel = target._relative_load().detach() + assert rel.min() < 0.25, "the node was not actually starved" + assert after > before + 1.0, f"barrier missed it: {before:.3f} -> {after:.3f}" + + +@pytest.mark.unit +def test_penalty_is_one_sided(pdb_path): + """Over-loading a node must not be penalised; only abandonment is. + + This is what separates the barrier from a load-entropy term, which is optimal at + uniform load and would push back on a legitimately broad node. + """ + model = _field_model(pdb_path) + target = NodeLoadTarget(model) + baseline = float(target()) + + with torch.no_grad(): + model.adp.refinable_params[0, 1] += 3.0 # widen one node -> it gains load + model.adp.reset_forward_cache() + widened = float(target()) + rel = target._relative_load().detach() + + assert float(rel.max()) > 1.5, "the node did not actually gain load" + # Widening one node necessarily takes load from others, so the total may rise a + # little; what must not happen is the over-loaded node itself being charged. + per_node_change = (widened - baseline) / rel.numel() + assert per_node_change < 0.5, ( + f"over-loading was penalised like abandonment ({per_node_change:.3f}/node)" + ) + + +@pytest.mark.unit +def test_gradient_reaches_geometry_but_not_values(pdb_path): + """Acts on the weights: positions and widths get gradient, node B does not.""" + model = _field_model(pdb_path) + target = NodeLoadTarget(model) + target().backward() + + grad = model.adp.refinable_params.grad + assert grad is not None + # Columns are [log B, log sigma, dx, dy, dz]. + assert float(grad[:, 0].abs().sum()) == pytest.approx(0.0, abs=1e-12), ( + "the barrier must not push the node VALUES" + ) + assert float(grad[:, 1].abs().sum()) > 0, "no gradient to the kernel widths" + assert float(grad[:, 2:5].abs().sum()) > 0, "no gradient to the node positions" + + +@pytest.mark.unit +def test_position_gradient_absent_when_positions_are_fixed(pdb_path): + """With positions fixed the barrier can only act through the widths.""" + model = _field_model(pdb_path, refine_node_positions=False) + target = NodeLoadTarget(model) + target().backward() + grad = model.adp.refinable_params.grad + assert grad.shape[1] == 2 + assert float(grad[:, 1].abs().sum()) > 0 + + +@pytest.mark.unit +def test_registered_in_the_total_adp_target(pdb_path): + """Reachable under the weight path 'adp/node_load'.""" + from torchref.refinement.targets.combined import TotalADPTarget + + model = _field_model(pdb_path) + total = TotalADPTarget(model, verbose=0) + assert "node_load" in total.target_losses() + assert torch.isfinite(torch.as_tensor(float(total["node_load"]()))) + + +@pytest.mark.unit +def test_default_weight_exists_for_the_component(pdb_path): + """A component with no weight entry would silently inherit the group weight.""" + from torchref.refinement.base_refinement import DEFAULT_GROUP_WEIGHTS + + assert "adp/node_load" in DEFAULT_GROUP_WEIGHTS + assert DEFAULT_GROUP_WEIGHTS["adp/node_load"] > 0 + + +# ---------------------------------------------------------------------------------- +# Payload independence. The barrier was written against the isotropic payload, whose +# storage is five columns wide, and the tests above poke column 1 by hand because that +# is where its log sigma sits. A displacement-mode field is 25 to 82 columns wide with +# log sigma near the end, so "acts through the weights, never the values" has to be +# re-established rather than assumed to carry over. +# ---------------------------------------------------------------------------------- + + +@pytest.mark.unit +@pytest.mark.parametrize( + "mode_set,width", [(None, 6), ("rigid", 21), ("rigid_dilation", 28), ("affine", 78)] +) +def test_barrier_prices_geometry_not_values_for_every_payload(pdb_path, mode_set, width): + """Zero gradient on the payload, non-zero on kernel width and node position. + + The invariant that separates this from a magnitude prior: it removes the opportunity + to place an extreme ADP rather than penalising the ADP. A payload-width bug would + show up here as gradient leaking into the payload columns. + """ + model = Model(verbose=0) + model.load_pdb(pdb_path) + model.set_adp_mode( + "field_aniso", n_nodes=24, k_neighbors=12, mode_set=mode_set + ) + field = model.adp_field + assert field.payload.width == width + assert field.node_shape[1] == width + 4 # payload | log sigma | 3 offset + + NodeLoadTarget(model)().backward() + grad = field.refinable_params.grad + assert grad is not None + assert float(grad[:, :width].abs().sum()) == pytest.approx(0.0, abs=1e-12), ( + "the barrier reached the node VALUES" + ) + assert float(grad[:, width].abs().sum()) > 0, "no gradient to the kernel width" + assert float(grad[:, width + 1 : width + 4].abs().sum()) > 0, ( + "no gradient to the node positions" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("mode_set", ["rigid", "rigid_dilation", "affine"]) +def test_barrier_still_catches_a_starved_node_on_a_mode_field(pdb_path, mode_set): + model = Model(verbose=0) + model.load_pdb(pdb_path) + model.set_adp_mode("field_aniso", n_nodes=24, k_neighbors=12, mode_set=mode_set) + field = model.adp_field + target = NodeLoadTarget(model) + before = float(target()) + + # log sigma is the column after the payload, wherever the payload ends. + with torch.no_grad(): + field.refinable_params[0, field.payload.width] -= 6.0 + field.reset_forward_cache() + + load = field.node_load().detach() + assert float((load / load.mean()).min()) < 0.25, "the node was not actually starved" + assert float(target()) > before + 1.0 diff --git a/tests/unit/refinement/test_shells.py b/tests/unit/refinement/test_shells.py new file mode 100644 index 00000000..c29a83d5 --- /dev/null +++ b/tests/unit/refinement/test_shells.py @@ -0,0 +1,127 @@ +"""Shell helpers shared by sigma_A and sigma_D. + +Pinned: the helpers ``sigma_a`` re-imports are the same objects ``_shells`` defines, so a +fit through either module reduces identically; ``equal_count_shells`` reproduces the +shell construction written out in ``estimate_beta``; the line shrinkage takes the line +outright when the scatter is below the noise, passes through with too few shells, +honours a slope clamp, and treats a shell without a finite variance as undetermined. +""" + +import math + +import pytest +import torch + +from torchref.refinement.model_error_estimation import _shells, sigma_a +from torchref.refinement.model_error_estimation._shells import ( + dl_shrink_to_line, + equal_count_shells, + interp_in_dss, + segsum, +) + + +@pytest.mark.unit +def test_sigma_a_uses_the_shared_helpers(): + assert sigma_a._segsum is _shells.segsum + assert sigma_a._interp_in_dss is _shells.interp_in_dss + assert sigma_a._segment_layout is _shells.segment_layout + + +@pytest.mark.unit +@pytest.mark.parametrize("n,per_bin", [(2000, 140), (37, 140), (9, 140), (5000, 500)]) +def test_equal_count_shells_matches_estimate_beta_construction(n, per_bin, any_device): + """The construction ``estimate_beta`` writes out inline, reproduced independently.""" + g = torch.Generator().manual_seed(3) + dss = (torch.rand(n, generator=g) * 0.3 + 0.02).to(any_device) + min_bins, min_per_bin = 5, 40 + order, seg, seg_lengths, n_bins = equal_count_shells( + dss, per_bin=per_bin, min_bins=min_bins, min_per_bin=min_per_bin + ) + # Oracle: the four lines of estimate_beta. + ref_order = torch.argsort(dss, stable=True) + n_by_count = max(1, n // per_bin) + n_cap = max(1, n // min_per_bin) + ref_bins = max(n_by_count, min(min_bins, n_cap)) + ref_seg = (torch.arange(n, device=dss.device) * ref_bins) // n + assert n_bins == ref_bins + assert torch.equal(order, ref_order) + assert torch.equal(seg, ref_seg) + assert torch.equal(seg_lengths, torch.bincount(ref_seg, minlength=ref_bins)) + assert int(seg_lengths.sum()) == n + assert int(seg_lengths.max() - seg_lengths.min()) <= 1 + + +@pytest.mark.unit +def test_segsum_and_interp_round_trip(any_device): + lengths = torch.tensor([3, 2, 4], device=any_device) + x = torch.arange(9, dtype=torch.float32, device=any_device) + assert torch.equal( + segsum(x, lengths), torch.tensor([3.0, 7.0, 26.0], device=any_device) + ) + bin_dss = torch.tensor([0.1, 0.2, 0.3], device=any_device) + vals = torch.tensor([1.0, 3.0, 5.0], device=any_device) + grid = torch.tensor([0.0, 0.15, 0.25, 0.5], device=any_device) + out = interp_in_dss(grid, bin_dss, vals) + assert torch.allclose(out, torch.tensor([1.0, 2.0, 4.0, 5.0], device=any_device)) + + +@pytest.mark.unit +def test_shrink_passes_through_with_fewer_than_four_shells(): + y = torch.tensor([1.0, 2.0, 3.0]) + var = torch.ones(3) + x = torch.tensor([0.1, 0.2, 0.3]) + out, w, tau_sq, a, b = dl_shrink_to_line(y, var, x) + assert torch.equal(out, y) + assert torch.equal(w, torch.zeros(3)) + assert float(tau_sq) == 0.0 and math.isnan(a) and math.isnan(b) + + +@pytest.mark.unit +def test_shrink_takes_the_line_when_scatter_is_below_noise(): + x = torch.linspace(0.05, 0.35, 8, dtype=torch.float64) + line = 2.0 - 3.0 * x + g = torch.Generator().manual_seed(1) + var = torch.full((8,), 0.04, dtype=torch.float64) + y = line + 0.01 * torch.randn(8, generator=g, dtype=torch.float64) + out, w, tau_sq, a, b = dl_shrink_to_line(y, var, x) + assert float(tau_sq) == 0.0 + assert torch.allclose(w, torch.ones(8, dtype=torch.float64)) + assert torch.allclose(out, a + b * x) + assert abs(a - 2.0) < 0.05 and abs(b + 3.0) < 0.3 + + +@pytest.mark.unit +def test_shrink_keeps_a_real_departure(): + x = torch.linspace(0.05, 0.35, 10, dtype=torch.float64) + y = 1.0 - 2.0 * x + y[4] += 3.0 # one shell far off the line, far beyond its own variance + var = torch.full((10,), 1e-4, dtype=torch.float64) + out, w, tau_sq, _, _ = dl_shrink_to_line(y, var, x) + assert float(tau_sq) > 0.0 + assert float(w[4]) < 0.01 + assert abs(float(out[4] - y[4])) < 0.05 + + +@pytest.mark.unit +def test_shrink_slope_clamp_is_honoured(): + x = torch.linspace(0.05, 0.35, 8, dtype=torch.float64) + y = 1.0 + 4.0 * x + var = torch.full((8,), 0.01, dtype=torch.float64) + _, _, _, _, b = dl_shrink_to_line(y, var, x, slope_max=0.0) + assert b == 0.0 + _, _, _, _, b2 = dl_shrink_to_line(-y, var, x, slope_min=0.0) + assert b2 == 0.0 + + +@pytest.mark.unit +def test_shrink_replaces_undetermined_shells_by_the_line(): + x = torch.linspace(0.05, 0.35, 8, dtype=torch.float64) + y = 2.0 - 3.0 * x + var = torch.full((8,), 0.01, dtype=torch.float64) + y[2] = float("nan") + var[5] = float("inf") + out, w, _, a, b = dl_shrink_to_line(y, var, x) + assert float(w[2]) == 1.0 and float(w[5]) == 1.0 + assert torch.isfinite(out).all() + assert torch.allclose(out[[2, 5]], a + b * x[[2, 5]]) diff --git a/tests/unit/refinement/test_sigma_d.py b/tests/unit/refinement/test_sigma_d.py new file mode 100644 index 00000000..39ec8f72 --- /dev/null +++ b/tests/unit/refinement/test_sigma_d.py @@ -0,0 +1,358 @@ +"""Properties of the sigma_D difference-power estimator. + +Pinned on seeded synthetic differences with a KNOWN power law: the per-shell power is +recovered from a single dataset through ``mean(dF**2) - mean(sigma**2)``, the dark- +amplitude exponent is recovered and can be fixed, the moment identity with a difference +model holds exactly, clamps are counted and an all-noise input is flagged rather than +weighted, degenerate inputs stay finite, per-reflection weights lie in ``[0, 1)`` with the +shell mean of the power preserved, the fit is deterministic and device-independent, and +the cached estimator resets on demand. +""" + +import pytest +import torch + +from torchref.refinement.model_error_estimation._shells import interp_in_dss +from torchref.refinement.model_error_estimation.sigma_d import ( + GAMMA_DEFAULT, + SigmaDConfig, + SigmaDEstimator, + estimate_sigma_d, + sigma_d_per_reflection, +) + +#: Tolerance on the recovered shell power relative to the truth. A shell of 140 +#: reflections estimates ``B`` with a relative sd of ``sqrt(2/140) = 12%``; the line +#: shrinkage pools shells, and the decile means below average ~14 shells, so 10% is +#: ~3 sd of what remains. +POWER_RTOL = 0.10 +#: Tolerance on the fitted exponent. Its standard error at 30 000 reflections is ~0.04; +#: 0.15 is well above that and well below the difference between the pure shell model +#: (0) and the default (1). +GAMMA_ATOL = 0.15 + + +def synth_diff( + n=30000, + gamma=1.0, + sig_frac=1.0, + seed=7, + dtype=torch.float32, + device="cpu", + with_model=False, + alpha_true=0.8, +): + """Signed differences with power ``Sigma_N(d*^2) * (F / )**gamma``. + + ``F`` is Wilson-like (the modulus of a complex normal) so the amplitude classes are + populated realistically; ``Sigma_N`` falls with resolution. The measurement sigma is + ``sig_frac`` times the rms true difference, constant across reflections so the + inverse-variance and sigma_D weights differ only through ``S``. + """ + g = torch.Generator().manual_seed(seed) + dss = torch.linspace(0.02, 0.35, n, dtype=torch.float64) + f = ( + torch.randn(n, generator=g, dtype=torch.float64) ** 2 + + torch.randn(n, generator=g, dtype=torch.float64) ** 2 + ).sqrt() * 10.0 + sigma_n = 0.05 * torch.exp(-3.0 * dss) + s_true = sigma_n * (f / f.mean()) ** gamma + d_true = torch.randn(n, generator=g, dtype=torch.float64) * s_true.sqrt() + sig = torch.full( + (n,), float(sig_frac) * float(s_true.mean().sqrt()), dtype=torch.float64 + ) + d_obs = d_true + torch.randn(n, generator=g, dtype=torch.float64) * sig + out = { + "delta_obs": d_obs, + "sigma_diff": sig, + "d_star_sq": dss, + "f_dark": f, + "s_true": s_true, + "fit_mask": torch.ones(n, dtype=torch.bool), + } + if with_model: + beta_true = 0.3 * s_true + out["delta_calc"] = ( + d_true + torch.randn(n, generator=g, dtype=torch.float64) * beta_true.sqrt() + ) / alpha_true + return { + k: ( + v.to(device=device, dtype=dtype) + if v.dtype.is_floating_point + else v.to(device) + ) + for k, v in out.items() + } + + +def _decile_means(values, dss, n_dec=10): + order = torch.argsort(dss) + chunks = torch.chunk(values[order], n_dec) + return torch.stack([c.mean() for c in chunks]) + + +@pytest.mark.unit +def test_recovers_shell_power_from_one_dataset(any_device): + d = synth_diff(device=any_device) + sh = estimate_sigma_d( + d["delta_obs"], + d["sigma_diff"], + None, + d["d_star_sq"], + d["f_dark"], + d["fit_mask"], + ) + est = sigma_d_per_reflection(sh, d["d_star_sq"], None, d["f_dark"], d["sigma_diff"]) + assert not sh.degenerate and not sh.all_zero + assert (sh.Sigma_N > 0).all() + got = _decile_means(est.S, d["d_star_sq"]) + want = _decile_means(d["s_true"], d["d_star_sq"]) + assert torch.allclose(got, want, rtol=POWER_RTOL) + + +@pytest.mark.unit +@pytest.mark.parametrize("gamma", [1.0, 0.5]) +def test_recovers_the_amplitude_exponent(gamma): + d = synth_diff(gamma=gamma, dtype=torch.float64) + sh = estimate_sigma_d( + d["delta_obs"], + d["sigma_diff"], + None, + d["d_star_sq"], + d["f_dark"], + d["fit_mask"], + ) + assert sh.gamma_fitted and sh.diagnostics["gamma_reason"] == "fitted" + assert abs(sh.gamma - gamma) < GAMMA_ATOL + assert sh.diagnostics["gamma_se"] < GAMMA_ATOL + + +@pytest.mark.unit +def test_fixed_exponent_is_honoured(): + d = synth_diff(n=5000) + sh = estimate_sigma_d( + d["delta_obs"], + d["sigma_diff"], + None, + d["d_star_sq"], + d["f_dark"], + d["fit_mask"], + gamma=0.7, + ) + assert sh.gamma == 0.7 and not sh.gamma_fitted + assert sh.diagnostics["gamma_reason"] == "fixed" + with pytest.raises(ValueError): + SigmaDConfig(gamma=3.0) + + +@pytest.mark.unit +def test_without_dark_amplitude_the_power_is_flat_within_a_shell(): + d = synth_diff(n=5000) + sh = estimate_sigma_d( + d["delta_obs"], d["sigma_diff"], None, d["d_star_sq"], None, d["fit_mask"] + ) + assert sh.gamma == GAMMA_DEFAULT and not sh.gamma_fitted + assert sh.diagnostics["gamma_reason"] == "no_f_dark" + est = sigma_d_per_reflection(sh, d["d_star_sq"], None, None, d["sigma_diff"]) + # Reflections at the same resolution share the power regardless of amplitude. + order = torch.argsort(d["d_star_sq"]) + close = est.S[order][:200] + assert float(close.max() / close.min()) < 1.05 + + +@pytest.mark.unit +@pytest.mark.parametrize("dtype,rtol", [(torch.float32, 1e-4), (torch.float64, 1e-10)]) +def test_moment_identity_with_a_difference_model(dtype, rtol): + d = synth_diff(dtype=dtype, with_model=True) + sh = estimate_sigma_d( + d["delta_obs"], + d["sigma_diff"], + None, + d["d_star_sq"], + d["f_dark"], + d["fit_mask"], + delta_calc=d["delta_calc"], + shrink=False, + ) + assert sh.has_model + assert sh.diagnostics["n_s2_clamped"] == 0 + # Sampling noise can push alpha**2 Sigma_P above Sigma_N in a few shells; the clamp + # there is counted, and the identity is exact everywhere it did not fire. + unclamped = sh.alpha**2 * sh.Sigma_P <= sh.Sigma_N + assert ( + int(unclamped.sum()) + == sh.diagnostics["n_shell"] - sh.diagnostics["n_beta_clamped"] + ) + assert float(unclamped.float().mean()) > 0.8 + lhs = sh.alpha**2 * sh.Sigma_P + sh.beta_model + sh.S2 + assert torch.allclose(lhs[unclamped], sh.B[unclamped], rtol=rtol) + # alpha is the Gaussian coupling S / (S + beta_true) / alpha_true-scaled slope; it must + # be positive and below one for this generator. + assert (sh.alpha > 0).all() and (sh.alpha < 1).all() + + +@pytest.mark.unit +def test_clamps_are_counted_and_all_noise_is_flagged(): + d = synth_diff(n=5000, sig_frac=5.0) + sh = estimate_sigma_d( + d["delta_obs"], + d["sigma_diff"], + None, + d["d_star_sq"], + d["f_dark"], + d["fit_mask"], + ) + assert sh.diagnostics["n_s2_clamped"] > 0 + noise = synth_diff(n=5000, sig_frac=50.0) + sh2 = estimate_sigma_d( + noise["delta_obs"], + noise["sigma_diff"] * 1.2, + None, + noise["d_star_sq"], + noise["f_dark"], + noise["fit_mask"], + shrink=False, + ) + assert sh2.all_zero + est = sigma_d_per_reflection( + sh2, noise["d_star_sq"], None, noise["f_dark"], noise["sigma_diff"] + ) + assert torch.equal(est.w, torch.zeros_like(est.w)) + + +@pytest.mark.unit +def test_pure_noise_with_calibrated_sigma_gets_no_power(): + """Nothing tells the estimator whether a difference exists: on pure noise with + calibrated sigmas the shrinkage must not manufacture power from the positive half + of the noise in ``B - S2``, and the weights collapse onto inverse variance.""" + g = torch.Generator().manual_seed(3) + n = 30000 + dss = torch.linspace(0.02, 0.35, n, dtype=torch.float64) + f = torch.rand(n, generator=g, dtype=torch.float64) * 20.0 + 1.0 + sig = 0.2 + 0.8 * dss + d_obs = torch.randn(n, generator=g, dtype=torch.float64) * sig + mask = torch.ones(n, dtype=torch.bool) + sh = estimate_sigma_d(d_obs, sig, None, dss, f, mask) + assert not sh.degenerate + # Shell power is below a few per cent of the noise power in every shell. + assert (sh.Sigma_N <= 0.05 * sh.S2).all() + est = sigma_d_per_reflection(sh, dss, None, f, sig) + ivw = 1.0 / sig**2 + ivw = ivw / ivw.mean() + w_sd = est.w / est.w.mean().clamp(min=1e-30) + if not sh.all_zero: + assert torch.corrcoef(torch.stack([w_sd, ivw]))[0, 1] > 0.97 + + +@pytest.mark.unit +def test_degenerate_input_stays_finite(): + d = synth_diff(n=100) + mask = torch.zeros(100, dtype=torch.bool) + mask[0] = True + sh = estimate_sigma_d( + d["delta_obs"], d["sigma_diff"], None, d["d_star_sq"], d["f_dark"], mask + ) + assert sh.degenerate and not sh.all_zero + est = sigma_d_per_reflection(sh, d["d_star_sq"], None, d["f_dark"], d["sigma_diff"]) + assert torch.isfinite(est.S).all() and torch.isfinite(est.w).all() + assert est.S.shape == (100,) + + +@pytest.mark.unit +def test_per_reflection_weights_and_shell_mean(any_device): + d = synth_diff(device=any_device) + eps = torch.where( + torch.arange(d["delta_obs"].numel(), device=any_device) % 7 == 0, 2.0, 1.0 + ).to(d["delta_obs"].dtype) + sh = estimate_sigma_d( + d["delta_obs"], d["sigma_diff"], eps, d["d_star_sq"], d["f_dark"], d["fit_mask"] + ) + est = sigma_d_per_reflection(sh, d["d_star_sq"], eps, d["f_dark"], d["sigma_diff"]) + assert (est.w >= 0).all() and (est.w < 1).all() + assert est.S.device == d["delta_obs"].device + # The multiplier has shell mean one, so S / epsilon averages to Sigma_N over a shell. + counts = sh.counts.to(torch.long) # dtype-ok: split sizes; PyTorch requires int64 + order = torch.argsort(d["d_star_sq"]) + per_shell = torch.stack( + [c.mean() for c in torch.split((est.S / eps)[order], counts.tolist())] + ) + assert torch.allclose(per_shell, sh.Sigma_N, rtol=0.15) + # A missing dark amplitude means a multiplier of one. + f_missing = d["f_dark"].clone() + f_missing[:50] = float("nan") + est2 = sigma_d_per_reflection(sh, d["d_star_sq"], eps, f_missing, d["sigma_diff"]) + log_sn = interp_in_dss(d["d_star_sq"][:50], sh.bin_dss, torch.log(sh.Sigma_N)) + assert torch.allclose(est2.S[:50], eps[:50] * torch.exp(log_sn), rtol=1e-4) + + +@pytest.mark.unit +def test_deterministic_and_device_independent(any_device): + d_cpu = synth_diff() + a = estimate_sigma_d( + d_cpu["delta_obs"], + d_cpu["sigma_diff"], + None, + d_cpu["d_star_sq"], + d_cpu["f_dark"], + d_cpu["fit_mask"], + ) + b = estimate_sigma_d( + d_cpu["delta_obs"], + d_cpu["sigma_diff"], + None, + d_cpu["d_star_sq"], + d_cpu["f_dark"], + d_cpu["fit_mask"], + ) + assert torch.equal(a.Sigma_N, b.Sigma_N) and a.gamma == b.gamma + d_dev = synth_diff(device=any_device) + c = estimate_sigma_d( + d_dev["delta_obs"], + d_dev["sigma_diff"], + None, + d_dev["d_star_sq"], + d_dev["f_dark"], + d_dev["fit_mask"], + ) + assert torch.allclose(c.Sigma_N.cpu(), a.Sigma_N, rtol=1e-4) + assert abs(c.gamma - a.gamma) < 1e-3 + + +@pytest.mark.unit +def test_estimator_caches_until_reset_and_remaps(): + d = synth_diff(n=5000) + est = SigmaDEstimator(SigmaDConfig(gamma=1.0)) + first = est.get( + d["delta_obs"], + d["sigma_diff"], + None, + d["d_star_sq"], + d["f_dark"], + d["fit_mask"], + ) + assert ( + est.get( + d["delta_obs"], + d["sigma_diff"], + None, + d["d_star_sq"], + d["f_dark"], + d["fit_mask"], + ) + is first + ) + est.reset() + assert est._cache is None + target = d["d_star_sq"][:1000] + remapped = est.get( + d["delta_obs"], + d["sigma_diff"], + None, + d["d_star_sq"], + d["f_dark"], + d["fit_mask"], + target_dss=target, + out_f_dark=d["f_dark"][:1000], + out_sigma_diff=d["sigma_diff"][:1000], + ) + assert remapped.S.shape == (1000,) and est.shells is not None diff --git a/tests/unit/refinement/test_targets.py b/tests/unit/refinement/test_targets.py deleted file mode 100644 index fe2728e7..00000000 --- a/tests/unit/refinement/test_targets.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Unit tests for torchref.refinement.targets - -Tests target (loss) functions for crystallographic refinement. -Note: These are unit tests so we test the functions in isolation with mock data. -""" - -import pytest -import torch -import torch.nn as nn -import numpy as np - - -class TestTargetBase: - """Tests for base Target class.""" - - @pytest.mark.unit - def test_target_empty_initialization(self): - """Test Target can be initialized without arguments.""" - from torchref.refinement.targets import Target - - target = Target() - - assert target.verbose == 0 - - @pytest.mark.unit - def test_target_is_nn_module(self): - """Target should be a nn.Module.""" - from torchref.refinement.targets import Target - - target = Target() - - assert isinstance(target, nn.Module) - - -class TestGaussianNLL: - """Tests for Gaussian NLL calculation logic.""" - - @pytest.mark.unit - def test_gaussian_nll_identical_gives_small_loss(self, mock_F_obs, mock_F_sigma): - """When Fobs = Fcalc, NLL should be small (just the log sigma term).""" - from torchref.base.math_torch import nll_xray - - fobs = mock_F_obs(n_reflections=100) - sigma = mock_F_sigma(n_reflections=100) - fcalc = fobs.clone().to(torch.complex64) # |Fcalc| = Fobs - - # Calculate manually what Gaussian NLL should be - # NLL = 0.5*((fobs - |fcalc|)/sigma)^2 + log(sigma) + 0.5*log(2pi) - diff = fobs - torch.abs(fcalc) - expected_data_term = 0.5 * ((diff / sigma) ** 2) - - # Data term should be ~0 when fobs = |fcalc| - assert torch.allclose(expected_data_term, torch.zeros_like(expected_data_term), atol=1e-5) - - @pytest.mark.unit - def test_gaussian_nll_positive(self, mock_F_obs, mock_F_sigma): - """NLL should generally be positive or close to zero.""" - fobs = mock_F_obs(n_reflections=100) - sigma = mock_F_sigma(n_reflections=100) - fcalc = mock_F_obs(n_reflections=100, seed=123).to(torch.complex64) # Different - - # Simple Gaussian NLL - diff = fobs - torch.abs(fcalc) - eps = torch.median(sigma) * 0.1 - sigma_safe = torch.clamp(sigma, min=eps) - log_2pi = torch.log(torch.tensor(2.0 * np.pi)) - nll = 0.5 * (diff ** 2) / (sigma_safe ** 2) + torch.log(sigma_safe) + 0.5 * log_2pi - - # Mean NLL should be finite - assert torch.isfinite(nll.mean()) - - -class TestLeastSquaresTarget: - """Tests for Least Squares target calculation.""" - - @pytest.mark.unit - def test_least_squares_identical_zero(self, mock_F_obs): - """LS loss should be 0 when Fobs = Fcalc.""" - fobs = mock_F_obs(n_reflections=100) - fcalc = fobs.clone() - - # Simple LS: sum((fobs - fcalc)^2) - loss = torch.sum((fobs - fcalc) ** 2) - - assert torch.isclose(loss, torch.tensor(0.0, dtype=loss.dtype), atol=1e-10) - - @pytest.mark.unit - def test_least_squares_scaled(self, mock_F_obs): - """Test LS loss with scaled Fcalc.""" - fobs = mock_F_obs(n_reflections=100) - fcalc = fobs * 1.1 # 10% scaled - - loss = torch.mean((fobs - fcalc) ** 2) - - # Should be (0.1 * fobs)^2 on average - expected_loss = torch.mean((0.1 * fobs) ** 2) - assert torch.isclose(loss, expected_loss, rtol=1e-5) - - @pytest.mark.unit - def test_least_squares_weighted(self, mock_F_obs, mock_F_sigma): - """Test weighted LS with sigma weights.""" - fobs = mock_F_obs(n_reflections=100) - sigma = mock_F_sigma(n_reflections=100) - fcalc = mock_F_obs(n_reflections=100, seed=123) - - # Weighted LS: sum(w * (fobs - fcalc)^2) where w = 1/sigma^2 - weights = 1.0 / (sigma ** 2) - diff = fobs - fcalc - weighted_loss = torch.sum(weights * (diff ** 2)) - - assert torch.isfinite(weighted_loss) - assert weighted_loss >= 0 - - -class TestRiceNLL: - """Tests for Rice distribution NLL (used for acentric reflections).""" - - @pytest.mark.unit - def test_rice_nll_components(self, mock_F_obs, mock_F_sigma): - """Test components of Rice NLL calculation.""" - from torch.special import i0 - - fobs = mock_F_obs(n_reflections=50) - sigma = mock_F_sigma(n_reflections=50) - fcalc_amp = mock_F_obs(n_reflections=50, seed=123) - - # Rice NLL components - # NLL = (Fo^2 + Fc^2)/(2σ^2) - log(I0(Fo*Fc/σ^2)) - log(Fo/σ^2) - - # Check I0 calculation - x = fobs * fcalc_amp / (sigma ** 2) - bessel_i0 = i0(x) - - # I0 should be >= 1 for x >= 0 - assert torch.all(bessel_i0 >= 1.0) - - -class TestTargetDeviceHandling: - """Tests for proper device handling in targets.""" - - @pytest.mark.unit - def test_target_cpu_tensors(self, mock_F_obs, mock_F_sigma): - """Test calculations work on CPU.""" - fobs = mock_F_obs(n_reflections=100) - sigma = mock_F_sigma(n_reflections=100) - - # Simple calculation on CPU - loss = torch.mean((fobs / sigma) ** 2) - - assert loss.device.type == 'cpu' - assert torch.isfinite(loss) - - @pytest.mark.unit - @pytest.mark.gpu - def test_target_gpu_tensors(self, mock_F_obs, mock_F_sigma, gpu_device): - """Test calculations work on GPU.""" - fobs = mock_F_obs(n_reflections=100).to(gpu_device) - sigma = mock_F_sigma(n_reflections=100).to(gpu_device) - - loss = torch.mean((fobs / sigma) ** 2) - - assert loss.device.type == gpu_device.type - assert torch.isfinite(loss) - - -class TestNumericStability: - """Tests for numeric stability in target calculations.""" - - @pytest.mark.unit - def test_small_sigma_handling(self, mock_F_obs): - """Test handling of very small sigma values.""" - fobs = mock_F_obs(n_reflections=100) - sigma = torch.ones_like(fobs) * 1e-10 # Very small sigma - fcalc = mock_F_obs(n_reflections=100, seed=123) - - # Clamped sigma approach - eps = torch.median(sigma) * 0.1 - sigma_safe = torch.clamp(sigma, min=max(eps, 1e-6)) - - diff = fobs - fcalc - loss = torch.mean((diff / sigma_safe) ** 2) - - assert torch.isfinite(loss) - - @pytest.mark.unit - def test_zero_fcalc_handling(self, mock_F_obs, mock_F_sigma): - """Test handling of zero Fcalc values.""" - fobs = mock_F_obs(n_reflections=100) - sigma = mock_F_sigma(n_reflections=100) - fcalc = torch.zeros_like(fobs, dtype=torch.complex64) # All zero - - fcalc_amp = torch.abs(fcalc) # Will be zero - diff = fobs - fcalc_amp - - loss = torch.mean(diff ** 2) - - # Should just be mean of fobs^2 - expected = torch.mean(fobs ** 2) - assert torch.isclose(loss, expected, rtol=1e-5) diff --git a/tests/unit/refinement/test_targets_comprehensive.py b/tests/unit/refinement/test_targets_comprehensive.py index 21aa6c9e..0593a129 100644 --- a/tests/unit/refinement/test_targets_comprehensive.py +++ b/tests/unit/refinement/test_targets_comprehensive.py @@ -4,16 +4,16 @@ These tests focus on individual target classes with mock/minimal data to achieve higher coverage of the targets module. """ + +import numpy as np import pytest import torch -import numpy as np -from unittest.mock import MagicMock, PropertyMock - # ============================================================================= # Base Target Tests # ============================================================================= + @pytest.mark.unit class TestBaseTarget: """Test base Target class functionality.""" @@ -24,18 +24,19 @@ def test_target_initialization_empty(self): target = Target() assert target.verbose == 0 + assert isinstance(target, torch.nn.Module) def test_target_initialization_with_verbose(self): """Test initialization with verbose setting.""" from torchref.refinement.targets import Target - + target = Target(verbose=2) assert target.verbose == 2 def test_target_forward_not_implemented(self): """Test that forward raises NotImplementedError.""" from torchref.refinement.targets import Target - + target = Target() with pytest.raises(NotImplementedError): target.forward() @@ -45,6 +46,7 @@ def test_target_forward_not_implemented(self): # X-ray Target Tests # ============================================================================= + @pytest.mark.unit class TestXrayTargetBase: """Test XrayTarget base class.""" @@ -71,39 +73,9 @@ def test_gaussian_target_initialization(self): assert target._model is None assert target._data is None - def test_gaussian_nll_computation(self): - """Test Gaussian NLL computation with mock data.""" - from torchref.base.math_torch import nll_xray - - # Test the underlying function - fobs = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32) - fcalc = torch.tensor([1.1, 1.9, 3.2, 3.8], dtype=torch.float32) - sigma = torch.tensor([0.1, 0.1, 0.1, 0.1], dtype=torch.float32) - - loss = nll_xray(fobs, fcalc, sigma).mean() - assert torch.isfinite(loss) # NLL can be negative depending on normalization -@pytest.mark.unit -class TestLeastSquaresXrayTarget: - """Test LeastSquaresXrayTarget.""" - - def test_least_squares_computation(self): - """Test least squares computation with mock data.""" - # Least squares: sum of (fobs - fcalc)^2 / sigma^2 - fobs = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32) - fcalc = torch.tensor([1.1, 1.9, 3.2, 3.8], dtype=torch.float32) - sigma = torch.tensor([0.1, 0.1, 0.1, 0.1], dtype=torch.float32) - - diff = fobs - fcalc - weights = 1.0 / (sigma ** 2) - loss = 0.5 * torch.sum(weights * (diff ** 2)) - - assert torch.isfinite(loss) - assert loss >= 0 - - @pytest.mark.unit class TestRiceXrayTarget: """Test RiceXrayTarget.""" @@ -121,6 +93,7 @@ def test_rice_target_initialization(self): # Geometry Target Tests # ============================================================================= + @pytest.mark.unit class TestGeometryTargetBase: """Test GeometryTarget base class.""" @@ -144,30 +117,6 @@ def test_bond_target_initialization(self): target = BondTarget() assert target._model is None - def test_bond_deviation_calculation(self): - """Test bond deviation calculation with mock data.""" - # Create mock coordinates for a simple bond - xyz = torch.tensor([ - [0.0, 0.0, 0.0], - [1.5, 0.0, 0.0], # 1.5 Å bond - ], dtype=torch.float32) - - # Bond indices - i_atoms = torch.tensor([0]) - j_atoms = torch.tensor([1]) - - # Expected distance and sigma - d_expected = torch.tensor([1.54]) # Expected C-C bond - sigma = torch.tensor([0.02]) - - # Calculate actual distances - d_actual = torch.norm(xyz[i_atoms] - xyz[j_atoms], dim=1) - - # Calculate deviation - deviation = (d_actual - d_expected) / sigma - - assert torch.isfinite(deviation).all() - @pytest.mark.unit class TestAngleTarget: @@ -180,27 +129,6 @@ def test_angle_target_initialization(self): target = AngleTarget() assert target._model is None - def test_angle_calculation(self): - """Test angle calculation with mock data.""" - # Create mock coordinates for a 90-degree angle - xyz = torch.tensor([ - [1.0, 0.0, 0.0], # Atom 1 - [0.0, 0.0, 0.0], # Atom 2 (vertex) - [0.0, 1.0, 0.0], # Atom 3 - ], dtype=torch.float32) - - # Vectors - v1 = xyz[0] - xyz[1] - v2 = xyz[2] - xyz[1] - - # Calculate angle - cos_angle = torch.dot(v1, v2) / (torch.norm(v1) * torch.norm(v2)) - angle = torch.acos(cos_angle) - angle_deg = torch.rad2deg(angle) - - # Should be approximately 90 degrees - assert torch.isclose(angle_deg, torch.tensor(90.0), atol=0.1) - @pytest.mark.unit class TestTorsionTarget: @@ -213,34 +141,6 @@ def test_torsion_target_initialization(self): target = TorsionTarget() assert target._model is None - def test_torsion_angle_calculation(self): - """Test torsion angle calculation.""" - # Create mock coordinates for a torsion - # Atoms in a plane should give ~0 or ~180 degree torsion - xyz = torch.tensor([ - [0.0, 0.0, 0.0], - [1.0, 0.0, 0.0], - [2.0, 0.0, 0.0], - [3.0, 0.0, 0.0], - ], dtype=torch.float32) - - # Calculate torsion using standard formula - b1 = xyz[1] - xyz[0] - b2 = xyz[2] - xyz[1] - b3 = xyz[3] - xyz[2] - - # Normal vectors - n1 = torch.linalg.cross(b1, b2) - n2 = torch.linalg.cross(b2, b3) - - # Torsion angle - if torch.norm(n1) > 1e-6 and torch.norm(n2) > 1e-6: - cos_torsion = torch.dot(n1, n2) / (torch.norm(n1) * torch.norm(n2)) - # Clamp to valid range - cos_torsion = torch.clamp(cos_torsion, -1.0, 1.0) - torsion = torch.acos(cos_torsion) - assert torch.isfinite(torsion) - @pytest.mark.unit class TestPlanarityTarget: @@ -253,29 +153,6 @@ def test_planarity_target_initialization(self): target = PlanarityTarget() assert target._model is None - def test_planarity_calculation(self): - """Test planarity calculation for coplanar atoms.""" - # Atoms in the XY plane - xyz = torch.tensor([ - [0.0, 0.0, 0.0], - [1.0, 0.0, 0.0], - [1.0, 1.0, 0.0], - [0.0, 1.0, 0.0], - ], dtype=torch.float32) - - # Calculate centroid - centroid = xyz.mean(dim=0) - - # Center coordinates - centered = xyz - centroid - - # SVD to find plane - U, S, Vh = torch.linalg.svd(centered) - - # The smallest singular value indicates planarity - # For perfectly coplanar points, it should be ~0 - assert S[-1] < 0.1 - @pytest.mark.unit class TestChiralTarget: @@ -288,26 +165,6 @@ def test_chiral_target_initialization(self): target = ChiralTarget() assert target._model is None - def test_chiral_volume_calculation(self): - """Test chiral volume calculation.""" - # Create a tetrahedron - xyz = torch.tensor([ - [1.0, 0.0, -1.0/np.sqrt(2)], # Center - [0.0, 0.0, 1.0/np.sqrt(2)], # Atom 1 - [1.0, 1.0, 0.0], # Atom 2 - [1.0, -1.0, 0.0], # Atom 3 - ], dtype=torch.float32) - - # Vectors from center to other atoms - v1 = xyz[1] - xyz[0] - v2 = xyz[2] - xyz[0] - v3 = xyz[3] - xyz[0] - - # Chiral volume (scalar triple product) - chiral_vol = torch.dot(v1, torch.linalg.cross(v2, v3)) - - assert torch.isfinite(chiral_vol) - @pytest.mark.unit class TestNonBondedTarget: @@ -337,6 +194,7 @@ def test_total_geometry_target_initialization(self): # ADP Target Tests # ============================================================================= + @pytest.mark.unit class TestADPTargetBase: """Test ADPTarget base class.""" @@ -349,61 +207,10 @@ def test_adp_target_initialization(self): assert target._model is None -@pytest.mark.unit -class TestADPSimilarityTarget: - """Test ADPSimilarityTarget (SIMU restraint).""" - - def test_simu_calculation(self): - """Test SIMU calculation with mock B-factors.""" - # Create mock B-factors for nearby atoms - b_factors = torch.tensor([20.0, 21.0, 22.0, 50.0], dtype=torch.float32) - - # Pairs of similar atoms (indices) - i_atoms = torch.tensor([0, 1]) - j_atoms = torch.tensor([1, 2]) - - # Calculate difference - diff = b_factors[i_atoms] - b_factors[j_atoms] - - # SIMU restraint loss - sigma = 1.0 # B-factor sigma - simu_loss = (diff / sigma).pow(2).mean() - - assert torch.isfinite(simu_loss) - assert simu_loss >= 0 - - @pytest.mark.unit class TestRigidBondTarget: """Test RigidBondTarget (DELU restraint).""" - def test_delu_calculation(self): - """Test DELU calculation with mock U matrices.""" - # Create mock anisotropic U matrices (6 parameters each) - # U11, U22, U33, U12, U13, U23 - u1 = torch.tensor([0.05, 0.06, 0.04, 0.01, 0.005, -0.01], dtype=torch.float32) - u2 = torch.tensor([0.05, 0.06, 0.04, 0.01, 0.005, -0.01], dtype=torch.float32) - - # Bond vector (normalized) - bond_vec = torch.tensor([1.0, 0.0, 0.0], dtype=torch.float32) - - # Calculate U components along bond direction - # For Uij, the component along direction v is v^T U v - def u_along_direction(u_params, direction): - """Calculate U component along a direction.""" - U11, U22, U33, U12, U13, U23 = u_params - vx, vy, vz = direction - return (U11 * vx * vx + U22 * vy * vy + U33 * vz * vz + - 2 * U12 * vx * vy + 2 * U13 * vx * vz + 2 * U23 * vy * vz) - - u1_bond = u_along_direction(u1, bond_vec) - u2_bond = u_along_direction(u2, bond_vec) - - # DELU restraint: difference should be small - diff = u1_bond - u2_bond - - assert torch.isfinite(diff) - def test_aniso_path_runs_and_routes_grad_to_u(self, pdb_dir): """The anisotropic DELU path actually executes and feeds gradient to the U tensors. Regression for the dead ``hasattr(model, "u_aniso")`` gate, @@ -467,27 +274,11 @@ def test_matches_inverse_gamma_nll(self): beta = float(b.mean()) * (alpha - 1.0) mode = beta / (alpha + 1.0) - expected = ( - -sps.invgamma.logpdf(b.numpy(), alpha, scale=beta).sum() - + sps.invgamma.logpdf(mode, alpha, scale=beta) * len(b) - ) + expected = -sps.invgamma.logpdf( + b.numpy(), alpha, scale=beta + ).sum() + sps.invgamma.logpdf(mode, alpha, scale=beta) * len(b) assert float(adp_sigd_math(b, a, s0)) == pytest.approx(expected, rel=1e-10) - def test_alpha_sets_log_width(self): - """std(log B) = sqrt(trigamma(alpha)), the bridge the design rests on. - - This is what lets alpha play the role the log-normal's sigma played, and - is the basis for reporting ``implied_std_log_adp``. - """ - from scipy import stats as sps - from scipy.special import polygamma - - for alpha in (3.5, 7.4): - draws = sps.invgamma.rvs(alpha, scale=100.0, size=400000, random_state=1) - assert np.log(draws).std() == pytest.approx( - np.sqrt(polygamma(1, alpha)), rel=2e-2 - ) - def test_monotonically_increasing_in_spread(self): """The loss must never reward spreading the B distribution out. @@ -577,6 +368,7 @@ def test_gradient_pushes_toward_the_mode(self): # R-factor Tests # ============================================================================= + @pytest.mark.unit class TestRfactorCalculations: """Test R-factor calculation functions.""" @@ -584,15 +376,15 @@ class TestRfactorCalculations: def test_get_rfactors_basic(self): """Test basic R-factor calculation.""" from torchref.base.math_torch import get_rfactors - + fobs = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32) fcalc = torch.tensor([1.1, 2.1, 3.1, 4.1, 5.1], dtype=torch.float32) - + # Create rfree mask (1 reflection in test set) rfree_mask = torch.tensor([True, True, True, True, False], dtype=torch.bool) - + r_work, r_free = get_rfactors(fobs, fcalc, rfree_mask) - + # Both should be small since fcalc is close to fobs assert r_work < 0.2 # r_free only has one reflection @@ -600,33 +392,33 @@ def test_get_rfactors_basic(self): def test_get_rfactors_perfect_fit(self): """Test R-factor with perfect fit.""" from torchref.base.math_torch import get_rfactors - + fobs = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32) fcalc = fobs.clone() # Perfect fit - + rfree_mask = torch.tensor([True, True, True, True, False], dtype=torch.bool) - + r_work, r_free = get_rfactors(fobs, fcalc, rfree_mask) - + assert r_work < 0.001 # Should be ~0 def test_bin_wise_rfactors(self): """Test bin-wise R-factor calculation.""" from torchref.base.math_torch import bin_wise_rfactors - + n_refl = 100 n_bins = 5 - + fobs = torch.rand(n_refl) + 1.0 fcalc = fobs * (1 + 0.1 * torch.randn(n_refl)) # Note: rfree=True means work set (not free set) rfree_mask = torch.rand(n_refl) > 0.1 - + # Ensure all bins are represented bins = torch.arange(n_refl) % n_bins - + r_work_bins, r_free_bins = bin_wise_rfactors(fobs, fcalc, rfree_mask, bins) - + # Should have results for each bin assert len(r_work_bins) == n_bins assert len(r_free_bins) == n_bins @@ -636,88 +428,7 @@ def test_bin_wise_rfactors(self): # Loss Function Tests # ============================================================================= -@pytest.mark.unit -class TestLossFunctions: - """Test individual loss functions from math_torch.""" - - def test_nll_xray(self): - """Test NLL X-ray loss function.""" - from torchref.base.math_torch import nll_xray - - fobs = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) - fcalc = torch.tensor([1.1, 2.1, 3.1], dtype=torch.float32) - sigma = torch.tensor([0.1, 0.1, 0.1], dtype=torch.float32) - - loss = nll_xray(fobs, fcalc, sigma).mean() - - # NLL can be negative depending on normalization - assert torch.isfinite(loss) - - def test_least_squares_manual(self): - """Test least squares loss calculation.""" - # Manual least squares implementation - fobs = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) - fcalc = torch.tensor([1.1, 2.1, 3.1], dtype=torch.float32) - sigma = torch.tensor([0.1, 0.1, 0.1], dtype=torch.float32) - - diff = fobs - fcalc - weights = 1.0 / (sigma ** 2) - loss = 0.5 * torch.sum(weights * (diff ** 2)) / len(fobs) - - assert torch.isfinite(loss) - assert loss >= 0 - - def test_nll_xray_with_mask(self): - """Test NLL X-ray with masking.""" - from torchref.base.math_torch import nll_xray - - fobs = torch.tensor([1.0, 2.0, 3.0, float('nan')], dtype=torch.float32) - fcalc = torch.tensor([1.1, 2.1, 3.1, 0.0], dtype=torch.float32) - sigma = torch.tensor([0.1, 0.1, 0.1, 0.1], dtype=torch.float32) - - # Only use finite values - valid = torch.isfinite(fobs) - loss = nll_xray(fobs[valid], fcalc[valid], sigma[valid]).mean() - - assert torch.isfinite(loss) - # ============================================================================= # Helper Function Tests # ============================================================================= - -@pytest.mark.unit -class TestTargetHelpers: - """Test helper functions used in targets.""" - - def test_distance_calculation(self): - """Test distance calculation between atom pairs.""" - xyz = torch.tensor([ - [0.0, 0.0, 0.0], - [3.0, 4.0, 0.0], # Distance = 5.0 - ], dtype=torch.float32) - - distance = torch.norm(xyz[1] - xyz[0]) - - assert torch.isclose(distance, torch.tensor(5.0)) - - def test_angle_from_vectors(self): - """Test angle calculation from vectors.""" - v1 = torch.tensor([1.0, 0.0, 0.0], dtype=torch.float32) - v2 = torch.tensor([0.0, 1.0, 0.0], dtype=torch.float32) - - cos_angle = torch.dot(v1, v2) / (torch.norm(v1) * torch.norm(v2)) - angle = torch.acos(cos_angle) - angle_deg = torch.rad2deg(angle) - - assert torch.isclose(angle_deg, torch.tensor(90.0)) - - def test_cross_product(self): - """Test cross product for normal vectors.""" - v1 = torch.tensor([1.0, 0.0, 0.0], dtype=torch.float32) - v2 = torch.tensor([0.0, 1.0, 0.0], dtype=torch.float32) - - normal = torch.linalg.cross(v1, v2) - - # Should be [0, 0, 1] - assert torch.allclose(normal, torch.tensor([0.0, 0.0, 1.0])) diff --git a/tests/unit/refinement/test_targets_reach_the_loss_state.py b/tests/unit/refinement/test_targets_reach_the_loss_state.py new file mode 100644 index 00000000..170e5cfd --- /dev/null +++ b/tests/unit/refinement/test_targets_reach_the_loss_state.py @@ -0,0 +1,117 @@ +"""Every registered target must arrive in the LossState under a key its weight can reach. + +``LossState.register_targets`` keys each component off ``target.name``, falling back to +``Target.name`` -- which is the literal string ``"model_target"``. A target that forgets +to declare its own name therefore registers under that, collides with every other target +that forgot, and no hierarchical weight can address it. The term is constructed, callable +and correct in isolation; it simply never enters the loss. + +That is what happened to ``adp/node_load`` and ``adp/node_smoothness``: both shipped +unnamed, so the node-coverage barrier was never in any refinement's loss despite having a +weight of 10.0 in ``DEFAULT_GROUP_WEIGHTS`` and passing every test that called it +directly. These tests check the plumbing rather than the arithmetic. +""" + +import pytest + +from torchref.refinement.targets.base import ModelTarget + + +def _leaf_target_classes(): + """Every concrete ModelTarget subclass that is a loss component, not a container.""" + import inspect + import pkgutil + import importlib + + import torchref.refinement.targets as pkg + from torchref.refinement.targets.combined import CombinedModelTargets + + for mod in pkgutil.walk_packages(pkg.__path__, pkg.__name__ + "."): + try: + importlib.import_module(mod.name) + except Exception: + continue + seen = {} + stack = [ModelTarget] + while stack: + cls = stack.pop() + for sub in cls.__subclasses__(): + stack.append(sub) + # Skip group bases (ADPTarget, GeometryTarget, ...). They are not formally + # abstract -- nothing in them is an abstractmethod -- so the only reliable + # marker is that other targets derive from them. A base legitimately carries + # the placeholder name because it is never registered itself. + if sub.__subclasses__(): + continue + if inspect.isabstract(sub) or issubclass(sub, CombinedModelTargets): + continue + seen[sub.__qualname__] = sub + return seen + + +@pytest.mark.unit +def test_no_component_target_inherits_the_placeholder_name(): + """A component using the base placeholder cannot be addressed by any weight.""" + offenders = { + name: cls.name + for name, cls in _leaf_target_classes().items() + if getattr(cls, "name", None) in (None, "base_target", "model_target", + "data_target") + } + assert not offenders, ( + "these targets would register under the base placeholder name, colliding with " + "each other and unreachable by any hierarchical weight:\n " + + "\n ".join(f"{k}: name={v!r}" for k, v in sorted(offenders.items())) + ) + + +@pytest.mark.unit +def test_adp_component_names_are_hierarchical(): + """An ADP component must sit under the ``adp`` group or the group weight misses it.""" + import torchref.refinement.targets.adp as adp_pkg + + bad = {} + for attr in dir(adp_pkg): + cls = getattr(adp_pkg, attr) + if not isinstance(cls, type) or not issubclass(cls, ModelTarget): + continue + if cls.__subclasses__(): + continue # a group base, never registered itself + name = getattr(cls, "name", "") + if not isinstance(name, str) or not name.startswith("adp/"): + bad[attr] = name + assert not bad, f"ADP targets not under the adp group: {bad}" + + +@pytest.mark.unit +@pytest.mark.parametrize("mode_set", [None, "rigid_dilation"]) +def test_field_components_are_registered_and_weighted(pdb_dir, mtz_dir, mode_set): + """End to end: the components the representation declares must be in the loss. + + Compares the combined target's own component set against the LossState's keys, so a + component that exists but never registers is caught -- which is the failure mode that + a direct ``adp_target['node_load']()`` call cannot see. + """ + from torchref.refinement.lbfgs_refinement import LBFGSRefinement + + ref = LBFGSRefinement( + data_file=str(mtz_dir / "1DAW.mtz"), pdb=str(pdb_dir / "1DAW.pdb"), + verbose=0, adp_mode="field_aniso", adp_mode_set=mode_set, + ) + components = set(ref.adp_target.target_losses()) + state = ref.complete_loss_state() + + for component in components: + key = f"adp/{component}" + assert key in state.targets, ( + f"{component!r} is a component of TotalADPTarget but never reached the " + f"LossState. Registered adp keys: " + f"{sorted(k for k in state.targets if k.startswith('adp'))}" + ) + # And the weight has to be addressable, not merely present. + assert state.get_effective_weight(key) is not None + + assert "node_load" in components, "field mode must carry the coverage barrier" + assert state.get_effective_weight("adp/node_load") > 0.0, ( + "the coverage barrier registered but at zero effective weight, so it is inert" + ) diff --git a/tests/unit/refinement/test_wilson_prior.py b/tests/unit/refinement/test_wilson_prior.py index a3d008ba..75635189 100644 --- a/tests/unit/refinement/test_wilson_prior.py +++ b/tests/unit/refinement/test_wilson_prior.py @@ -32,7 +32,7 @@ def setup_target(): ) ens.cell = data.cell ens.spacegroup = data.spacegroup - ens.setup_grid(max_res=data.get_max_res()) + ens.max_res = data.get_max_res() scaler = Scaler(model=ens, data=data, nbins=10, verbose=0) fcalc0 = ens(data.hkl) diff --git a/tests/unit/restraints/__init__.py b/tests/unit/restraints/__init__.py deleted file mode 100644 index 571dfe85..00000000 --- a/tests/unit/restraints/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Unit tests for torchref.restraints module.""" diff --git a/tests/unit/scaling/test_f_sol_override_contract.py b/tests/unit/scaling/test_f_sol_override_contract.py new file mode 100644 index 00000000..0178db00 --- /dev/null +++ b/tests/unit/scaling/test_f_sol_override_contract.py @@ -0,0 +1,123 @@ +"""Solvent overrides preserve output shape and leave the default solvent cache intact.""" + +from types import SimpleNamespace + +import pytest +import torch + +from torchref.config import get_complex_dtype, get_float_dtype, get_int_dtype + + +class _StubSolvent: + """Minimal stand-in for :class:`SolventModel` on the k_sol/B_sol path.""" + + optimize_phase = False + + def __init__(self, device, value: float = 1.0): + self.device = device + self.value = value + self.n_reads = 0 + + def k_solvent(self): + return torch.tensor(0.35, device=self.device, dtype=get_float_dtype()) + + def damping(self, s_half_sq): + return torch.ones_like(s_half_sq) + + def get_rec_solvent(self, hkl): + self.n_reads += 1 + return torch.full( + (hkl.shape[0],), self.value, dtype=get_complex_dtype(), device=self.device + ) + + +@pytest.fixture +def scaler_with_stub(any_device): + """A bare ``Scaler`` with only the solvent branch live.""" + from torchref.scaling.scaler import Scaler + + n = 6 + scaler = Scaler(device=any_device) + dev = scaler.device + scaler.bins = torch.zeros(n, dtype=get_int_dtype(), device=dev) + scaler._s_half_sq = torch.zeros(n, device=dev, dtype=get_float_dtype()) + # The no-override path reads the solvent model at self.hkl, which is a read-only + # property over self._data. + scaler._data = SimpleNamespace( + hkl=torch.zeros((n, 3), dtype=get_int_dtype(), device=dev) + ) + scaler.solvent = _StubSolvent(dev) + scaler._f_sol_raw = None + return scaler, n, dev + + +class TestOverrideDoesNotMutateTheCache: + + @pytest.mark.unit + def test_a_later_call_without_an_override_sees_the_model_solvent( + self, scaler_with_stub + ): + """A solvent override applies only to the call that supplies it.""" + scaler, n, dev = scaler_with_stub + fcalc = torch.ones(n, dtype=get_complex_dtype(), device=dev) + + baseline = scaler.forward(fcalc).clone() # stub solvent, value 1.0 + scaler._f_sol_raw = None # as update_solvent() would leave it + + far_off = torch.full((n,), 99.0, dtype=get_complex_dtype(), device=dev) + scaler.forward(fcalc, f_sol_override=far_off) + after = scaler.forward(fcalc) + + assert torch.allclose(after, baseline), ( + "a plain forward() after an override call returned the override's " + "solvent contribution" + ) + + +class TestOverridePreservesRank: + + @pytest.mark.unit + def test_each_batch_row_matches_the_unbatched_call(self, scaler_with_stub): + """Rank alone is not enough -- the rows must also be the right ones.""" + scaler, n, dev = scaler_with_stub + t = 3 + fcalc = torch.stack( + [ + torch.full((n,), float(i + 1), dtype=get_complex_dtype(), device=dev) + for i in range(t) + ] + ) + override = torch.stack( + [ + torch.full( + (n,), float(10 * (i + 1)), dtype=get_complex_dtype(), device=dev + ) + for i in range(t) + ] + ) + + batched = scaler.forward(fcalc, f_sol_override=override) + assert batched.shape == (t, n) + for i in range(t): + scaler._f_sol_raw = None + single = scaler.forward(fcalc[i], f_sol_override=override[i]) + assert torch.allclose( + batched[i], single + ), f"batched row {i} does not match the equivalent unbatched call" + + @pytest.mark.unit + def test_unbatched_override_is_unchanged(self, scaler_with_stub): + """The ``(N,)`` solvent path must keep working; it is what every single-dataset + caller uses.""" + scaler, n, dev = scaler_with_stub + fcalc = torch.ones(n, dtype=get_complex_dtype(), device=dev) + override = torch.full((n,), 2.0, dtype=get_complex_dtype(), device=dev) + + out = scaler.forward(fcalc, f_sol_override=override) + + assert out.shape == (n,) + # fcalc + k_sol * f_sol, with damping == 1 and no aniso/Chebyshev/per-bin B. + expected = 1.0 + 0.35 * 2.0 + assert torch.allclose( + out.real, torch.full((n,), expected, device=dev, dtype=get_float_dtype()) + ) diff --git a/tests/unit/scaling/test_multiplicative_scale.py b/tests/unit/scaling/test_multiplicative_scale.py new file mode 100644 index 00000000..479e8ed4 --- /dev/null +++ b/tests/unit/scaling/test_multiplicative_scale.py @@ -0,0 +1,42 @@ +"""The scaler's observed-to-model factor without the bulk-solvent term. + +Pinned: an uninitialised scaler reports ones; after initialisation the factor +reproduces ``forward`` exactly once the additive solvent term is removed, so dividing +observed amplitudes by it returns them to the model's absolute scale. +""" + +import pytest +import torch + +from torchref.io import ReflectionData +from torchref.model.model_ft import ModelFT +from torchref.scaling.scaler import Scaler + + +@pytest.fixture +def scaler(sample_structure_pair): + model = ModelFT() + model.load_cif(str(sample_structure_pair["model"])) + data = ReflectionData(verbose=0) + data.load_mtz(str(sample_structure_pair["reflections"])) + return Scaler(model=model, data=data, nbins=10, verbose=0) + + +@pytest.mark.unit +def test_uninitialised_scaler_reports_ones(scaler): + factor = scaler.multiplicative_scale() + assert factor.shape == (int(scaler.bins.numel()),) + assert torch.equal(factor, torch.ones_like(factor)) + assert factor.device == scaler.device + + +@pytest.mark.integration +def test_factor_reproduces_forward_without_solvent(scaler): + scaler.initialize() + fcalc = scaler.compute_fcalc() + factor = scaler.multiplicative_scale() + assert (factor > 0).all() and torch.isfinite(factor).all() + assert not factor.requires_grad + with torch.no_grad(): + scaled = scaler(fcalc, f_sol_override=torch.zeros_like(fcalc)) + assert torch.allclose(scaled, factor.to(scaled.dtype) * fcalc, rtol=1e-5, atol=1e-6) diff --git a/tests/unit/scaling/test_weighting.py b/tests/unit/scaling/test_weighting.py new file mode 100644 index 00000000..6bd49cd2 --- /dev/null +++ b/tests/unit/scaling/test_weighting.py @@ -0,0 +1,152 @@ +"""The weighting half of the split: properties, not a preferred answer. + +Every assertion here is something that has to be true whatever the caps end up +being. The numbers themselves are screened, not asserted -- pinning a cap in a +unit test would make the screen unfalsifiable. +""" + +import pytest +import torch + +from torchref.scaling.weighting import ( + DEFAULT_SNR_CAP, empirical_sigma_a, information_weight, + inverse_variance_weight, normalise_weight, snr_from_amplitude, +) + +pytestmark = pytest.mark.unit + + +def test_information_weight_saturates_at_one(): + """Past the cap, better measurement buys nothing: the model is the limit.""" + snr = torch.tensor([0.0, 1.0, 5.0, 50.0, 1e4], dtype=torch.float64) + w = information_weight(snr, cap=5.0) + assert float(w[0]) == 0.0 + assert float(w[2]) == pytest.approx(0.5), "the cap is where it reaches a half" + assert float(w[-1]) == pytest.approx(1.0, abs=1e-6) + assert bool((w[1:] > w[:-1]).all()), "must be monotone in signal-to-noise" + + +def test_information_weight_is_the_variance_ratio_in_disguise(): + """``w = 1/(1 + sigma_meas^2/sigma_model^2)`` -- not a sigmoid picked by eye. + + With the cap standing for the signal-to-noise at which the two errors are + equal, the weight must equal that expression exactly. + """ + snr = torch.linspace(0.1, 40.0, 200, dtype=torch.float64) + cap = 5.0 + expected = 1.0 / (1.0 + (cap / snr) ** 2) + assert torch.allclose(information_weight(snr, cap=cap), expected) + + +def test_below_the_cap_the_weight_is_quadratic_in_snr(): + """Where measurement error dominates, information goes as snr^2.""" + snr = torch.tensor([0.01, 0.02, 0.04], dtype=torch.float64) + w = information_weight(snr, cap=DEFAULT_SNR_CAP) + assert float(w[1] / w[0]) == pytest.approx(4.0, rel=1e-3) + assert float(w[2] / w[1]) == pytest.approx(4.0, rel=1e-3) + + +def test_snr_uses_the_intensity_convention(): + """``I/sigma_I`` with ``I = F^2`` is ``F/(2 sigma_F)``, half the amplitude's. + + Only a factor of two, and only a rescaling of the cap -- but quoting a cap + against the wrong signal-to-noise silently doubles it. + """ + F = torch.tensor([10.0, 100.0], dtype=torch.float64) + sig = torch.tensor([1.0, 1.0], dtype=torch.float64) + assert torch.allclose(snr_from_amplitude(F, sig), + torch.tensor([5.0, 50.0], dtype=torch.float64)) + + +def test_the_weight_rises_as_either_error_falls(): + """Inverse variance: better data or a better model both mean more weight.""" + snr = torch.tensor([1.0, 1.0, 10.0], dtype=torch.float64) + sa = torch.tensor([0.1, 0.9, 0.1], dtype=torch.float64) + w = inverse_variance_weight(snr, sa, cap=1e9) + assert float(w[1]) > float(w[0]), "a more reliable model must weigh more" + assert float(w[2]) > float(w[0]), "a better measurement must weigh more" + + +def test_measurement_error_bounds_the_low_resolution_weight(): + """``sigma_A -> 1`` sends the model variance to zero; ``1/snr^2`` is what + stops the weight diverging, and it must, because that regime is where the + strongest reflections live and one of them could otherwise carry the run.""" + sa = torch.tensor([0.999999, 0.999999], dtype=torch.float64) + snr = torch.tensor([5.0, 50.0], dtype=torch.float64) + w = inverse_variance_weight(snr, sa, cap=1e9) + assert bool(torch.isfinite(w).all()) + # snr^2 is the ceiling, approached from below: the residual model variance + # is small but not zero, and it bites harder the better the measurement. + assert float(w[0]) <= 25.0 and float(w[0]) == pytest.approx(25.0, rel=1e-2) + assert float(w[1]) <= 2500.0 and float(w[1]) == pytest.approx(2500.0, rel=1e-2) + + +def test_the_weight_keeps_its_dependence_on_model_error(): + """The factorised form lost this, which is why it was abandoned. + + A product of a ``snr`` term and ``sigma_A/(eps - sigma_A^2)`` came out + identical for a 0.5 A and a 1.0 A coordinate error, because the singularity + set the shape and the model error dropped out. The coupled form must not. + """ + import math + s = torch.linspace(0.02, 0.5, 12, dtype=torch.float64) + snr = torch.full_like(s, 8.0) + curves = [] + for dv in (0.5, 1.0): + sa = torch.exp(-(2.0 / 3.0) * (math.pi ** 2) * s * s * dv * dv) + curves.append(normalise_weight( + inverse_variance_weight(snr, sa, cap=1e9))) + spread = float((curves[0] / curves[1] - 1).abs().max()) + assert spread > 0.1, ( + f"the weight barely moved ({spread:.3f}) between a 0.5 A and a 1.0 A " + f"model error; it has stopped carrying model information" + ) + + +def test_epsilon_enters_the_variance_not_the_signal(): + """``V = eps - sigma_A^2``, so higher multiplicity means more variance and + therefore less weight, at fixed model reliability and measurement error.""" + sa = torch.full((4,), 0.5, dtype=torch.float64) + snr = torch.full((4,), 10.0, dtype=torch.float64) + plain = inverse_variance_weight(snr, sa, cap=1e9) + axial = inverse_variance_weight( + snr, sa, eps=torch.full((4,), 2.0, dtype=torch.float64), cap=1e9) + assert bool((axial < plain).all()) + + +def test_normalise_weight_sets_the_mean(): + w = torch.rand(1000, dtype=torch.float64) * 7.0 + 0.5 + assert float(normalise_weight(w).mean()) == pytest.approx(1.0) + + +def test_a_constant_weight_survives_normalisation_as_ones(): + w = torch.full((100,), 3.7, dtype=torch.float64) + assert torch.allclose(normalise_weight(w), torch.ones_like(w)) + + +def test_empirical_sigma_a_ignores_the_absolute_scale(): + """The data's scale is arbitrary and the model's is electrons; neither is information. + + Before this held, the ratio's level set the answer: a flat sigma_A of 0.2 on + one structure and 0.33 on another, with no resolution dependence at all. + """ + s = torch.linspace(0.07, 0.25, 40, dtype=torch.float64) + obs = torch.exp(-30.0 * s * s) + calc = torch.exp(-30.0 * s * s) + base = empirical_sigma_a(obs, calc) + torch.testing.assert_close(empirical_sigma_a(1000.0 * obs, calc), base) + torch.testing.assert_close(empirical_sigma_a(obs, 1e-3 * calc), base) + + +def test_empirical_sigma_a_reads_the_shape(): + """Identical shapes: full trust everywhere. A low-resolution deficit: less trust there.""" + s = torch.linspace(0.07, 0.25, 40, dtype=torch.float64) + calc = torch.exp(-30.0 * s * s) + same = empirical_sigma_a(7.0 * calc, calc) + assert float(same.min()) > 0.999 + # The model predicts more scattering at low resolution than the data have + # -- the bulk solvent it lacks -- and matches at high resolution. + deficit = torch.where(s < 0.12, torch.full_like(s, 0.4), torch.ones_like(s)) + sa = empirical_sigma_a(calc * deficit, calc) + assert float(sa[s < 0.12].max()) < float(sa[s > 0.15].min()) + assert bool((sa <= 1.0).all()) and bool((sa > 0.0).all()) diff --git a/tests/unit/scaling/test_wilson_normaliser.py b/tests/unit/scaling/test_wilson_normaliser.py new file mode 100644 index 00000000..0ff47375 --- /dev/null +++ b/tests/unit/scaling/test_wilson_normaliser.py @@ -0,0 +1,214 @@ +"""Invariants of the absolute Wilson normaliser. + +The load-bearing one is the first: `` = 1`` is the *stationarity condition* +of the Gamma GLM's constant term, not a normalisation applied afterwards. Every +consumer that centres an intensity as ``E^2 - 1`` depends on it being exact +rather than approximate -- the previous convention centred against a measured +`` = 0.954`` and nothing noticed. + +The rest pin the properties that make two fits comparable, which is what the +weighting design will be built on: a shared abscissa, epsilon reaching both +sides of the ratio, and invariance to the units the data arrive in. +""" + +import pytest +import torch + +from torchref.scaling.wilson import WilsonNormaliser +from torchref.symmetry import Cell, SpaceGroup + +pytestmark = pytest.mark.unit + + +def _wilson_data(n=20000, seed=0, centric_frac=0.1, eps_value=None): + """Intensities actually drawn from the distribution the fit assumes.""" + g = torch.Generator().manual_seed(seed) + s = torch.rand(n, generator=g, dtype=torch.float64) * 0.45 + 0.05 + # A real curve to recover: Wilson falloff plus a low-resolution deficit of + # the shape a missing bulk solvent produces. + sigma = (torch.exp(-120.0 * (s / 2) ** 2) * 3000.0 + * (1 - 0.9 * torch.exp(-300.0 * (s / 2) ** 2))) + centric = torch.zeros(n, dtype=torch.bool) + centric[: int(n * centric_frac)] = True + k = torch.where(centric, 0.5, 1.0).to(torch.float64) + eps = (torch.ones(n, dtype=torch.float64) if eps_value is None + else torch.full((n,), float(eps_value), dtype=torch.float64)) + # `_standard_gamma` takes no generator, so seed the global RNG too -- + # otherwise the draw depends on whatever ran before it and the test + # passes alone and fails in a suite. + torch.manual_seed(seed) + I = torch._standard_gamma(k.clone()) / k * (eps * sigma) + return I, s, eps, centric, sigma + + +def _k_weighted_mean(v, centric): + k = torch.where(centric, 0.5, 1.0).to(torch.float64) + return float((k * v.to(torch.float64)).sum() / k.sum()) + + +def test_unit_mean_is_an_identity_of_the_fit(): + """The constant column's score equation IS `` = 1``. + + ``sum_h k_h (I_h/mu_h - 1) = 0`` at the optimum, and the fit puts the + intercept on it in closed form, so this does NOT degrade as the convergence + tolerance is loosened. Measured 1e-9 to 2e-7 over five draws in float32; the + bar is the package's usual 1e-4 relative, so a failure here means the + intercept solve is broken rather than that the fit stopped early. + """ + I, s, eps, centric, _ = _wilson_data() + w = WilsonNormaliser(I, s, eps=eps, centric=centric, n_coeff=6) + assert _k_weighted_mean(w.E_squared, centric) == pytest.approx(1.0, rel=1e-4) + + +@pytest.mark.parametrize("n_coeff", [1, 2, 6, 12]) +def test_unit_mean_holds_at_every_order(n_coeff): + """It is the intercept that pins the mean, so the order must not matter.""" + I, s, eps, centric, _ = _wilson_data() + w = WilsonNormaliser(I, s, eps=eps, centric=centric, n_coeff=n_coeff) + assert _k_weighted_mean(w.E_squared, centric) == pytest.approx(1.0, abs=1e-6) + + +def test_a_uniform_epsilon_cancels_out_of_the_ratio(): + """``E^2 = (I/eps) / ``, so a constant eps must change nothing. + + This is the defect the previous convention carried: it divided the shell + mean by eps without dividing the intensity by it, leaving `` = `` + -- 1 on a primitive lattice and 2 on a centred one, so a normaliser whose + absolute scale depended on the space group. + """ + I, s, _, centric, _ = _wilson_data() + plain = WilsonNormaliser(I, s, centric=centric, n_coeff=6) + doubled = WilsonNormaliser( + I, s, eps=torch.full_like(s, 2.0), centric=centric, n_coeff=6, + ) + # eps=2 halves the intensity going in AND halves Sigma, so E is unchanged. + assert torch.allclose(doubled.E, plain.E, rtol=1e-4, atol=1e-6) + + +def test_invariant_to_the_units_the_data_arrive_in(): + I, s, eps, centric, _ = _wilson_data() + base = WilsonNormaliser(I, s, eps=eps, centric=centric, n_coeff=6).E + for c in (1e-6, 1e6): + scaled = WilsonNormaliser( + I * c, s, eps=eps, centric=centric, n_coeff=6, + ).E + assert torch.allclose(scaled, base, rtol=1e-4, atol=1e-6), ( + f"scaling I by {c:g} moved E by " + f"{float((scaled - base).abs().max()):.3e}" + ) + + +def test_the_resolution_trend_is_removed(): + I, s, eps, centric, _ = _wilson_data() + w = WilsonNormaliser(I, s, eps=eps, centric=centric, n_coeff=6) + order = torch.argsort(s) + means = [float(w.E_squared[order[i::10]].to(torch.float64).mean()) + for i in range(10)] + assert max(means) / min(means) < 1.15, f"residual trend: {means}" + + +def test_the_fitted_curve_recovers_the_true_one(): + I, s, eps, centric, sigma_true = _wilson_data(n=60000) + w = WilsonNormaliser(I, s, eps=eps, centric=centric, n_coeff=6) + rel = (w.sigma_wilson.to(torch.float64) / sigma_true - 1).abs() + assert float(rel.median()) < 0.05 + assert float(rel.max()) < 0.30 + + +def test_one_coefficient_is_a_single_global_scale(): + I, s, eps, centric, _ = _wilson_data() + w = WilsonNormaliser(I, s, eps=eps, centric=centric, n_coeff=1) + sig = w.sigma_wilson.to(torch.float64) + assert float((sig / sig[0] - 1).abs().max()) < 1e-9 + + +def test_the_range_only_matters_outside_the_fitted_data(): + """An affine remap does not change what a polynomial basis spans. + + So two fits over different ranges recover the same *function* where their + data overlap -- the coefficients differ, the curve does not. What the range + controls is the other side of it: ``u`` saturates at the ends, so beyond the + fitted data the curve is frozen flat rather than extrapolated. That is the + whole reason to pass an explicit range, and it is why a fit made on one + reflection set can be used on another only if the range covers both. + """ + I, s, eps, centric, _ = _wilson_data(n=30000) + lo, hi = float(s.min()), float(s.max()) + sub = s < 0.3 + kw = dict(eps=eps[sub], centric=centric[sub], n_coeff=6) + shared = WilsonNormaliser(I[sub], s[sub], s_lo=lo, s_hi=hi, **kw) + own = WilsonNormaliser(I[sub], s[sub], **kw) + + # Looser than the package's usual 1e-4, and the reason is the point of the + # test rather than an excuse. These are two INDEPENDENT fits, each stopped + # when its own objective stops improving by 1e-4 of what it has gained. The + # valley is flat along the high-order coefficients, so equal objectives + # there do not mean equal coefficients, and the curves separate by more than + # the objective did. Measured 0.2-1.6% over five draws; 3% catches a real + # dependence on the parameterisation without chasing the stopping rule. + inside = torch.linspace(0.06, 0.29, 40, dtype=torch.float64) + assert torch.allclose(shared.evaluate(inside), own.evaluate(inside), + rtol=3e-2), "the fitted function must not depend on " \ + "how the basis was parameterised" + + # Outside its own data, the narrow fit is pinned at its endpoint; the one + # given the full range keeps varying because it is still inside its basis. + outside = torch.linspace(0.32, 0.49, 20, dtype=torch.float64) + own_out = own.evaluate(outside).to(torch.float64) + assert float((own_out / own_out[0] - 1).abs().max()) < 1e-9, \ + "beyond the fitted range the curve should be flat, not extrapolated" + shared_out = shared.evaluate(outside).to(torch.float64) + assert float((shared_out / shared_out[0] - 1).abs().max()) > 1e-3 + + +def test_evaluate_reproduces_the_fitted_curve(): + I, s, eps, centric, _ = _wilson_data() + w = WilsonNormaliser(I, s, eps=eps, centric=centric, n_coeff=6) + assert torch.allclose(w.evaluate(s), w.sigma_wilson, rtol=1e-10, atol=1e-12) + + +def test_negative_intensities_are_kept_but_do_not_inform_the_fit(): + """Negative measurements are meaningful and unbiased; they are not errors. + + The Gamma likelihood has no support there, so they are held out of the + estimate -- but they still get a Sigma and a signed ``E_squared``, because + excluding them from the fit is not the same as refusing to normalise them. + """ + I, s, eps, centric, _ = _wilson_data() + I = I.clone() + I[:500] = -torch.rand(500, dtype=torch.float64) * 10.0 + w = WilsonNormaliser(I, s, eps=eps, centric=centric, n_coeff=6) + assert w.n_fitted == I.numel() - 500 + assert bool((w.E_squared[:500] < 0).all()), "sign must survive" + assert bool(torch.isfinite(w.sigma_wilson).all()) + assert bool((w.E[:500] == 0).all()), "E clamps, E_squared does not" + + +def test_it_raises_rather_than_quietly_degrading(): + """No fallback. A normaliser that becomes a different normaliser on the + hard cases is two normalisers wearing one name.""" + I, s, eps, centric, _ = _wilson_data(n=20) + with pytest.raises(ValueError, match="usable reflections"): + WilsonNormaliser(I[:3], s[:3], eps=eps[:3], centric=centric[:3], + n_coeff=6) + + +def test_from_hkl_excludes_systematic_absences(): + """Absences are zero by symmetry, not by measurement, so they say nothing + about Sigma -- and a Gamma fit told otherwise is dragged toward zero.""" + sg = SpaceGroup("P 43 21 2") + cell = Cell([70.0, 70.0, 90.0, 90.0, 90.0, 90.0]) + g = torch.Generator().manual_seed(5) + hkl = torch.randint(-14, 15, (12000, 3), generator=g) + hkl = hkl[hkl.abs().sum(dim=-1) > 0] + absent = sg.is_absent(hkl).to(torch.bool) + assert int(absent.sum()) > 0, "test needs a group with real absences" + + I = torch.rand(hkl.shape[0], generator=g, dtype=torch.float64) * 100 + 1 + I = torch.where(absent, torch.zeros_like(I), I) # absences really are 0 + + w = WilsonNormaliser.from_hkl(I, hkl, sg, cell, n_coeff=6) + assert w.n_fitted == int((~absent).sum()) + assert bool(torch.isfinite(w.sigma_wilson).all()) + # The zeros must not have dragged the curve down. + assert float(w.sigma_wilson.min()) > 0.0 diff --git a/tests/unit/scattering/test_anomalous.py b/tests/unit/scattering/test_anomalous.py index 145f9b3a..6fb0951b 100644 --- a/tests/unit/scattering/test_anomalous.py +++ b/tests/unit/scattering/test_anomalous.py @@ -246,7 +246,9 @@ def test_modelft_disable_anomalous(self, test_pdb_file): assert model.wavelength is None # Create HKL reflections - hkl = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.int32) + hkl = torch.tensor( + [[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.int32, device=model.device + ) # Should compute structure factors without anomalous correction sf = model.get_structure_factor(hkl) @@ -260,7 +262,11 @@ def test_anomalous_correction_applied(self, test_pdb_file): model = ModelFT(wavelength=1.0, anomalous_threshold=0.5, verbose=0) model.load_pdb(test_pdb_file) - hkl = torch.tensor([[1, 0, 0], [2, 1, 0], [1, 1, 1]], dtype=torch.int32) + hkl = torch.tensor( + + [[1, 0, 0], [2, 1, 0], [1, 1, 1]], dtype=torch.int32, device=model.device + + ) # Compute with anomalous correction sf_with = model.get_structure_factor( @@ -288,7 +294,11 @@ def test_friedel_pair_asymmetry(self, test_pdb_file): model = ModelFT(wavelength=1.0, anomalous_threshold=0.5, verbose=0) model.load_pdb(test_pdb_file) - hkl = torch.tensor([[1, 2, 3], [2, 1, 0], [3, 3, 3]], dtype=torch.int32) + hkl = torch.tensor( + + [[1, 2, 3], [2, 1, 0], [3, 3, 3]], dtype=torch.int32, device=model.device + + ) sf_plus = model.get_structure_factor(hkl, apply_anomalous=True, recalc=True) sf_minus = model.get_structure_factor(-hkl, apply_anomalous=True, recalc=True) @@ -361,7 +371,11 @@ def test_gradient_flow(self, test_pdb_file): model.load_pdb(test_pdb_file) # xyz.refinable_params should already have requires_grad=True by default - hkl = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.int32) + hkl = torch.tensor( + + [[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.int32, device=model.device + + ) sf = model.get_structure_factor(hkl, apply_anomalous=True, recalc=True) diff --git a/tests/unit/structure_factor/conftest.py b/tests/unit/structure_factor/conftest.py index bcbf9be4..7ab42565 100644 --- a/tests/unit/structure_factor/conftest.py +++ b/tests/unit/structure_factor/conftest.py @@ -13,13 +13,11 @@ import torch import torchref -from torchref.config import device as device_cfg, dtypes - -from tests.conftest import _accelerator +from tests.fixtures.devices import _accelerator +from tests.fixtures.precision import cpu_double_precision from . import helpers as H - # --------------------------------------------------------------------------- # Device axis # --------------------------------------------------------------------------- @@ -86,7 +84,9 @@ def ds_device_dtype_kernels(): for name in H.ds_kernels_for(device, dtype): out.append( pytest.param( - device, dtype, name, + device, + dtype, + name, id=f"{device.type}-{str(dtype).replace('torch.float', 'f')}-{name}", marks=dev_param.marks, ) @@ -103,30 +103,9 @@ def ds_device_dtype_kernels(): # --------------------------------------------------------------------------- @pytest.fixture(scope="package", autouse=True) def _float64_cpu(): - """float64/complex128 on CPU for this package; restore afterwards. - - Required, not cosmetic: ``iso_structure_factor_torched`` casts ``hkl`` to the - *global* ``dtypes.float`` (``torchref/base/direct_summation/isotropic.py:121``), so - under the default float32 config a float64 leaf produces a dtype-mismatched matmul. - That is why the pre-existing tests wrapped every eager-SF call in a ``double_cpu`` - fixture. - - ``sigma_cutoff_ed`` is restored here too -- the three copies of ``double_cpu`` this - replaces did not, so a test that changed the cutoff leaked it into everything that - ran after it. - """ - f0, c0, d0 = dtypes.float, dtypes.complex, device_cfg.current - s0 = torchref.sigma_cutoff_ed.value - dtypes.float = torch.float64 - dtypes.complex = torch.complex128 - device_cfg.current = torch.device("cpu") - try: + """Scope the package's CPU double-precision reference configuration.""" + with cpu_double_precision(): yield - finally: - dtypes.float = f0 - dtypes.complex = c0 - device_cfg.current = d0 - torchref.sigma_cutoff_ed.value = s0 @pytest.fixture diff --git a/tests/unit/structure_factor/helpers.py b/tests/unit/structure_factor/helpers.py index a853925f..fcebf894 100644 --- a/tests/unit/structure_factor/helpers.py +++ b/tests/unit/structure_factor/helpers.py @@ -442,17 +442,18 @@ def sf_fft_for( Pass ``fineness=1.0`` for a deliberately under-sampled grid; see :data:`GRID_FINENESS` for why that is the sampling-limited regime. """ + from torchref.model.context import ModelContext from torchref.model.sf_fft import SfFFT - - sf = SfFFT( - cell=scene.cell, - spacegroup=spacegroup, - max_res=scene.d_min / fineness, - dtype_float=dtype, - device=torch.device("cpu"), + from torchref.symmetry import Cell, SpaceGroup + + cpu = torch.device("cpu") + # A private context: the engine reads the crystal live, so it must not share + # the module-scoped scene's cell with other tests. + ctx = ModelContext( + cell=Cell(scene.cell.data, dtype=dtype, device=cpu), + spacegroup=SpaceGroup(spacegroup, dtype=dtype, device=cpu), ) - sf.setup_grid() - return sf + return SfFFT(ctx, max_res=scene.d_min / fineness, dtype_float=dtype, device=cpu) # --------------------------------------------------------------------------- @@ -465,11 +466,11 @@ def sf_fft_for( # falls back to the portable splat (``main.py`` catches and falls through), so a # dispatch-driven test can pass while measuring a different kernel than the one it # names. Calling the kernel directly settles that by construction. -# 2. **No global-config coupling.** ``SfFFT`` builds its grid through ``get_real_grid``, -# which reads the *global* ``dtypes.float`` and takes no dtype argument -- so an MPS -# ``SfFFT`` under this package's float64 pin would try to allocate float64 on MPS and -# fail. ``ifft`` and ``extract_structure_factor_from_grid`` read no global config at -# all, so :func:`density_to_F` needs no config switching. +# 2. **No global-config coupling.** ``build_electron_density`` allocates its map at the +# *global* ``dtypes.float`` when no ``dtype`` is passed, so a dispatch-driven MPS test +# under this package's float64 pin would try to allocate float64 on MPS and fail. +# ``ifft`` and ``extract_structure_factor_from_grid`` read no global config at all, so +# :func:`density_to_F` needs no config switching. # # The dispatch ladder is a separate concern, tested in ``test_dispatch.py``. @@ -684,8 +685,8 @@ def fft_sf(scene: Scene, sf_fft, xyz=None, occ=None, third=None, *, aniso=False) ``apply_symmetry=False`` on both calls: P1 isolates the truncation and sampling budget, and a symmetric comparison would cancel the symmetry algebra anyway, since - both routes call the same ``compute_symmetry_equivalent_hkls`` / - ``compute_translation_phases``. Symmetry is validated against gemmi instead, in + both routes call the same ``Symmetry.expand_reciprocal`` / + ``Symmetry.phase_factors``. Symmetry is validated against gemmi instead, in ``test_forward.py``. """ xyz = scene.xyz if xyz is None else xyz diff --git a/tests/unit/structure_factor/test_dispatch.py b/tests/unit/structure_factor/test_dispatch.py index 4368dc49..c78e14f3 100644 --- a/tests/unit/structure_factor/test_dispatch.py +++ b/tests/unit/structure_factor/test_dispatch.py @@ -45,18 +45,14 @@ def _build(scene, device, dtype, force_portable=False, aniso=False): """ s = scene.to(device=device, dtype=dtype) dims = H._grid_dims(s) - grid = torch.zeros(*dims, 3, dtype=dtype, device=device) - voxel = torch.tensor( - [float(s.cell.data[i]) / dims[i] for i in range(3)], dtype=dtype, device=device - ) empty1 = s.xyz.new_zeros(0) empty3 = s.xyz.new_zeros(0, 3) empty5 = s.A.new_zeros(0, 5) kw = dict( - real_space_grid=grid, + grid_shape=dims, + device=device, inv_frac_matrix=s.inv_frac_matrix, frac_matrix=s.frac_matrix, - voxel_size=voxel, dtype=dtype, ) kw["force_portable"] = force_portable @@ -393,9 +389,16 @@ def test_sfds_backend_toggle_end_to_end(scene_fine): s = scene_fine.to(device=cuda, dtype=torch.float32) obs = H.synthetic_obs(H.ds_direct(scene_fine, "eager").detach()).to(cuda, torch.float32) + from torchref.model.context import ModelContext + from torchref.symmetry import SpaceGroup + def run(force_portable): + ctx = ModelContext( + cell=s.cell, + spacegroup=SpaceGroup("P212121", dtype=torch.float32, device=cuda), + ) sf = SfDS( - cell=s.cell, spacegroup="P212121", force_portable=force_portable, + ctx, force_portable=force_portable, dtype_float=torch.float32, device=cuda, max_memory_gb=2.0, ) leaves = tuple(t.clone().requires_grad_(True) for t in (s.xyz, s.adp, s.occ)) diff --git a/tests/unit/structure_factor/test_forward.py b/tests/unit/structure_factor/test_forward.py index 0d4caa00..ecd0b065 100644 --- a/tests/unit/structure_factor/test_forward.py +++ b/tests/unit/structure_factor/test_forward.py @@ -245,19 +245,23 @@ def test_sfds_matches_gemmi_with_symmetry(gemmi_iso_symmetry): This is also the only symmetric comparison in the package. A DS-vs-FFT check cannot validate symmetry, because both routes call the same - ``compute_symmetry_equivalent_hkls`` / ``compute_translation_phases`` and the shared + ``Symmetry.expand_reciprocal`` / ``Symmetry.phase_factors`` and the shared algebra cancels; gemmi does not share it. """ scene, structure = gemmi_iso_symmetry assert len(structure.cell.images) > 0, "structure was not set up with symmetry" F_gemmi = H.gemmi_sf(structure, scene.hkl_list) - ds = SfDS( + from torchref.model.context import ModelContext + from torchref.symmetry import SpaceGroup + + ctx = ModelContext( cell=scene.cell, - spacegroup=scene.spacegroup, - dtype_float=torch.float64, - device=torch.device("cpu"), + spacegroup=SpaceGroup( + scene.spacegroup, dtype=torch.float64, device=torch.device("cpu") + ), ) + ds = SfDS(ctx, dtype_float=torch.float64, device=torch.device("cpu")) with torch.no_grad(): F_sym, _ = ds.compute_structure_factors( scene.hkl, scene.xyz, scene.adp, scene.occ, scene.A, scene.B, diff --git a/tests/unit/symmetry/test_canonicalize_hkl.py b/tests/unit/symmetry/test_canonicalize_hkl.py index 60471c7f..1764baf6 100644 --- a/tests/unit/symmetry/test_canonicalize_hkl.py +++ b/tests/unit/symmetry/test_canonicalize_hkl.py @@ -4,7 +4,31 @@ import pytest import torch -from torchref.symmetry.reciprocal_symmetry import canonicalize_hkl +from torchref.symmetry import SpaceGroup + + +# The HKL verbs live on the space group now. These adapters keep the assertions +# below -- which pin the phase-sign contract -- expressed in terms of the space +# group specifications the cases are parametrised over. +def canonicalize_hkl(hkl, sg, include_friedel=True, device=None): + return SpaceGroup(sg).canonicalize_hkl( + hkl, include_friedel=include_friedel, device=device + ) + + +def expand_hkl(hkl, sg, include_friedel=True, remove_absences=True, device=None): + return SpaceGroup(sg).expand_hkl( + hkl, + include_friedel=include_friedel, + remove_absences=remove_absences, + device=device, + ) + + +def reduce_hkl(hkl, sg, include_friedel=True, device=None): + return SpaceGroup(sg).reduce_hkl( + hkl, include_friedel=include_friedel, device=device + ) # --------------------------------------------------------------------------- @@ -41,8 +65,6 @@ def test_idempotency(self, sg): ) def test_equivalents_converge(self, sg): """All symmetry equivalents of a reflection map to the same canonical HKL.""" - from torchref.symmetry.reciprocal_symmetry import expand_hkl - # Use (1,1,5) which satisfies centering conditions for C2 hkl_asu = torch.tensor([[1, 1, 5]], dtype=torch.int32) hkl_p1, _, _ = expand_hkl(hkl_asu, sg, include_friedel=True) @@ -70,7 +92,6 @@ def test_phase_roundtrip(self, sg): then verify canonicalization produces a single consistent SF. """ import gemmi - from torchref.symmetry.spacegroup import SpaceGroup # Reference SF at canonical h F_ref = 10.0 diff --git a/tests/unit/symmetry/test_cell_identity.py b/tests/unit/symmetry/test_cell_identity.py new file mode 100644 index 00000000..9132d751 --- /dev/null +++ b/tests/unit/symmetry/test_cell_identity.py @@ -0,0 +1,102 @@ +"""Value identity of :class:`~torchref.symmetry.cell.Cell`: ``key``, ``__eq__``, ``__hash__``.""" + +import pytest +import torch + +from torchref.symmetry import Cell + +PARAMS = [50.0, 60.0, 70.0, 90.0, 90.0, 90.0] + + +@pytest.mark.unit +def test_equal_parameters_compare_and_hash_equal(): + a = Cell(PARAMS, dtype=torch.float32, device="cpu") + b = Cell(PARAMS, dtype=torch.float32, device="cpu") + assert a is not b + assert a == b + assert hash(a) == hash(b) + assert a.key == tuple(PARAMS) + + +@pytest.mark.unit +def test_clone_compares_equal(): + a = Cell(PARAMS) + assert a.clone() == a + assert hash(a.clone()) == hash(a) + + +@pytest.mark.unit +def test_different_parameters_compare_unequal(): + a = Cell(PARAMS) + perturbed = list(PARAMS) + perturbed[0] += 1e-3 + b = Cell(perturbed) + assert a != b + assert a.key != b.key + + +@pytest.mark.unit +def test_comparison_with_other_types_is_false(): + a = Cell(PARAMS) + assert (a == "50 60 70 90 90 90") is False + assert (a == PARAMS) is False + assert a != None # noqa: E711 + + +@pytest.mark.unit +def test_usable_as_dict_key_and_in_set(): + a = Cell(PARAMS) + b = Cell(PARAMS) + cache = {a: "grid"} + assert cache[b] == "grid" + assert len({a, b}) == 1 + + +@pytest.mark.unit +def test_key_is_cached_and_reset_with_the_other_derived_quantities(): + cell = Cell(PARAMS, dtype=torch.float32, device="cpu") + assert "key" not in cell._cache + first = cell.key + assert cell._cache["key"] is first + + cell.to(dtype=torch.float64) + assert "key" not in cell._cache, "reset_cache must drop the key with the rest" + assert cell.key == first + + +@pytest.mark.unit +def test_in_place_edit_is_refused_at_the_next_derived_read(): + cell = Cell(PARAMS) + _ = cell.volume + cell.data[0] = 51.0 + with pytest.raises(RuntimeError, match="create a new one"): + cell.fractional_matrix + with pytest.raises(RuntimeError, match="Please don't edit Cell objects"): + cell.key + + +@pytest.mark.unit +def test_in_place_edit_before_any_read_is_refused_too(): + cell = Cell(PARAMS) + cell.data.mul_(2.0) + with pytest.raises(RuntimeError, match="edited in place"): + cell.volume + + +@pytest.mark.unit +def test_constructor_owns_its_tensor(): + t = torch.tensor(PARAMS, dtype=torch.float32) + cell = Cell(t, dtype=torch.float32, device="cpu") + t[0] = 99.0 + assert cell.key[0] == 50.0 + assert float(cell.volume) == pytest.approx(210000.0) + + +@pytest.mark.unit +def test_device_and_dtype_moves_are_not_edits(): + cell = Cell(PARAMS, dtype=torch.float32, device="cpu") + _ = cell.fractional_matrix + cell.to(dtype=torch.float64) + assert cell.fractional_matrix.dtype == torch.float64 + assert cell.key == tuple(PARAMS) + assert cell.clone() == cell diff --git a/tests/unit/symmetry/test_epsilon_conventions.py b/tests/unit/symmetry/test_epsilon_conventions.py new file mode 100644 index 00000000..78a84bf5 --- /dev/null +++ b/tests/unit/symmetry/test_epsilon_conventions.py @@ -0,0 +1,121 @@ +"""``Symmetry.epsilon`` carries two conventions, and they must not drift together. + +Folding Friedel mates into epsilon mixes two different effects. Operations that +map ``h -> h`` add coherently and set the **mean**, ``<|F|^2> = eps * Sigma`` -- +the conventional crystallographic epsilon. Operations that map ``h -> -h`` leave +the mean alone and make ``F`` real, which changes the **distribution**; that is +centricity, and ``is_centric`` already carries it. + +Both settings have a consumer: sigma_A estimation is calibrated against the +Friedel-folded default, and the molecular-replacement likelihood wants the +conventional count for its ``V = eps - sigma_A**2``. So the risk is not that one +is wrong -- it is that a later edit quietly makes them the same, or flips which +one is the default, and nothing notices. These pin the difference, its location, +and its size. + +Counts are the measured ones rather than recomputed, so a regression shows up as +a specific wrong number instead of a green test that changed convention. +""" + +import pytest +import torch + +from torchref.symmetry import SpaceGroup + +pytestmark = pytest.mark.unit + +#: ``(hm, centred)``. Spans primitive and centred lattices, and the point groups +#: where the alignment package's own epsilon was measured to disagree. +SPACEGROUPS = [ + ("C 1 2 1", True), + ("P 21 21 2", False), + ("P 31 2 1", False), + ("P 43 21 2", False), + ("P 4 3 2", False), + ("P 65 2 2", False), +] + + +def _hkl(n=4000, seed=11): + g = torch.Generator().manual_seed(seed) + hkl = torch.randint(-12, 13, (n, 3), generator=g) + return hkl[(hkl.abs().sum(dim=-1) > 0)] # drop (0,0,0) + + +@pytest.mark.parametrize("hm, centred", SPACEGROUPS) +def test_friedel_changes_centric_reflections_and_only_those(hm, centred): + """The switch must move exactly the centric reflections, never a general one. + + This is the property that makes the two conventions safe to hold at once: if + the difference ever spread beyond centrics, one of them would have stopped + meaning what its docstring says. + """ + sg = SpaceGroup(hm) + hkl = _hkl() + with_f = sg.epsilon(hkl, friedel=True) + without = sg.epsilon(hkl, friedel=False) + centric = sg.is_centric(hkl).to(torch.bool) + + differs = with_f != without + assert not bool((differs & ~centric).any()), ( + f"{hm}: epsilon differs on {int((differs & ~centric).sum())} ACENTRIC " + f"reflections; the Friedel term must only reach centrics" + ) + # And the difference is a doubling where it lands, not an arbitrary shift. + if bool(differs.any()): + ratio = (with_f[differs] / without[differs]) + assert torch.allclose(ratio, torch.full_like(ratio, 2.0)), ( + f"{hm}: Friedel folding is not a factor of two where it applies" + ) + + +@pytest.mark.parametrize("hm, centred", SPACEGROUPS) +def test_the_conventional_count_is_never_larger(hm, centred): + sg = SpaceGroup(hm) + hkl = _hkl() + assert bool((sg.epsilon(hkl, friedel=False) + <= sg.epsilon(hkl, friedel=True)).all()) + + +@pytest.mark.parametrize("hm, centred", SPACEGROUPS) +def test_centring_cosets_are_counted_by_both(hm, centred): + """A centred lattice gives every reflection the centring order as a factor. + + Not asserted as desirable -- it is a documented property with one consequence, + that epsilon is inflated wherever it is used as a *term* rather than a factor. + Pinned so the behaviour is deliberate rather than discovered again. + """ + sg = SpaceGroup(hm) + hkl = _hkl() + eps = sg.epsilon(hkl, friedel=False) + general_min = float(eps.min()) + if centred: + assert general_min >= 2.0, ( + f"{hm} is centred; every reflection should carry the centring order " + f"but the minimum epsilon is {general_min}" + ) + else: + assert general_min == 1.0, ( + f"{hm} is primitive; general reflections should have epsilon 1, got " + f"{general_min}" + ) + + +def test_the_default_is_the_calibrated_one(): + """sigma_A estimation is calibrated against Friedel-folded epsilon. + + Flipping this default would decalibrate the refinement path silently, so the + default is pinned separately from the behaviour of either branch. + """ + sg = SpaceGroup("P 21 21 2") + hkl = _hkl() + assert torch.equal(sg.epsilon(hkl), sg.epsilon(hkl, friedel=True)) + + +def test_epsilon_is_at_least_one_and_finite(): + for hm, _ in SPACEGROUPS: + sg = SpaceGroup(hm) + for friedel in (True, False): + eps = sg.epsilon(_hkl(), friedel=friedel) + assert bool(torch.isfinite(eps).all()) + assert float(eps.min()) >= 1.0 diff --git a/tests/unit/symmetry/test_hkl_symmetry_gemmi.py b/tests/unit/symmetry/test_hkl_symmetry_gemmi.py index 4536f9c5..1bd1e126 100644 --- a/tests/unit/symmetry/test_hkl_symmetry_gemmi.py +++ b/tests/unit/symmetry/test_hkl_symmetry_gemmi.py @@ -1,12 +1,12 @@ """Regression tests for reciprocal-space (Miller-index) symmetry against gemmi. These guard the ``h' = h·R = Rᵀ·h`` reciprocal-space convention used by -``SpaceGroup.apply_to_hkl``. A previous bug applied the real-space transform +``Symmetry.expand_reciprocal``. A previous bug applied the real-space transform ``R·h`` instead, which is only correct when the fractional rotation matrix is symmetric. For space groups whose rotation matrices are non-symmetric (trigonal, hexagonal, and permutation-type cubic operations) ``R·h`` produced the wrong set of symmetry equivalents, corrupting the centric flags -(``is_centric_from_hkl``) and epsilon multiplicities (``epsilon_from_hkl``) +(``Symmetry.is_centric``) and epsilon multiplicities (``Symmetry.epsilon``) that feed French-Wilson intensity conversion and ML sigma_A weighting. Ground truth comes from gemmi: @@ -14,7 +14,7 @@ - centric: ``GroupOps.is_reflection_centric`` - epsilon: ``GroupOps.epsilon_factor_without_centering`` (Friedel-doubled for centric reflections, matching the Friedel-aware count in - ``epsilon_from_hkl``) + ``Symmetry.epsilon``) """ import numpy as np @@ -24,7 +24,7 @@ from torchref.config import get_default_device, get_float_dtype, get_int_dtype # Space groups whose rotation matrices are non-symmetric — these are the ones -# that regress if apply_to_hkl uses R·h instead of Rᵀ·h. Plus orthorhombic / +# that regress if expand_reciprocal uses R·h instead of Rᵀ·h. Plus orthorhombic / # tetragonal controls (symmetric matrices) that must remain correct either way. _TRIGONAL_HEXAGONAL = ["P 32 2 1", "P 31 2 1", "P 61 2 2", "P 6 2 2", "P 3 1 2"] _CONTROLS = ["P 21 21 21", "P 43 21 2", "P 1"] @@ -53,8 +53,8 @@ def _hkl_tensor(): @pytest.mark.unit @pytest.mark.parametrize("sg_name", _SPACE_GROUPS) -def test_apply_to_hkl_matches_gemmi_transform(sg_name): - """apply_to_hkl must reproduce gemmi's per-operation reciprocal transform. +def test_expand_reciprocal_matches_gemmi_transform(sg_name): + """expand_reciprocal must reproduce gemmi's per-operation reciprocal transform. Compared as the *set* of equivalents per reflection so the test is insensitive to operation ordering between torchref and gemmi. @@ -66,33 +66,32 @@ def test_apply_to_hkl_matches_gemmi_transform(sg_name): sg = SpaceGroup(sg_name) hkl = _hkl_tensor() - # torchref: (N, 3, ops) -> per-reflection set of equivalents - out = sg.apply_to_hkl(hkl.to(get_float_dtype())) - out_int = torch.round(out).to(torch.int64).cpu().numpy() # (N, 3, ops) + # torchref: (ops, N, 3) -> per-reflection set of equivalents + out_int = sg.expand_reciprocal(hkl).cpu().numpy() # (ops, N, 3) gemmi_ops = list(gemmi.SpaceGroup(sg_name).operations().sym_ops) for n, h in enumerate(_HKLS): - got = {tuple(out_int[n, :, o]) for o in range(out_int.shape[2])} + got = {tuple(out_int[o, n, :]) for o in range(out_int.shape[0])} expected = {tuple(op.apply_to_hkl(h)) for op in gemmi_ops} assert got == expected, ( - f"{sg_name} hkl={h}: apply_to_hkl equivalents {sorted(got)} " + f"{sg_name} hkl={h}: expand_reciprocal equivalents {sorted(got)} " f"!= gemmi {sorted(expected)}" ) @pytest.mark.unit @pytest.mark.parametrize("sg_name", _SPACE_GROUPS) -def test_is_centric_from_hkl_matches_gemmi(sg_name): +def test_is_centric_matches_gemmi(sg_name): """Centric flags must match gemmi's is_reflection_centric.""" import gemmi - from torchref.base.french_wilson import is_centric_from_hkl + from torchref.symmetry import SpaceGroup ops = gemmi.SpaceGroup(sg_name).operations() hkl = _hkl_tensor() - got = is_centric_from_hkl(hkl, sg_name).cpu().numpy().astype(bool) + got = SpaceGroup(sg_name).is_centric(hkl).cpu().numpy().astype(bool) expected = np.array([ops.is_reflection_centric(h) for h in _HKLS], dtype=bool) assert np.array_equal(got, expected), ( @@ -129,7 +128,7 @@ def test_epsilon_from_hkl_matches_gemmi(sg_name): @pytest.mark.unit -def test_apply_to_hkl_is_transpose_not_plain_rotation(): +def test_expand_reciprocal_is_transpose_not_plain_rotation(): """Explicit guard on the exact bug: apply_to_hkl == Rᵀ·h, not R·h. Uses P3, whose 3-fold rotation matrix is non-symmetric, so R·h and Rᵀ·h @@ -143,14 +142,16 @@ def test_apply_to_hkl_is_transpose_not_plain_rotation(): mats = sg.matrices # (ops, 3, 3) hkl = torch.tensor([[1, 2, 3]], dtype=mats.dtype, device=mats.device) - out = sg.apply_to_hkl(hkl) # (1, 3, ops) + out = sg.expand_reciprocal(hkl) # (ops, 1, 3) # Correct reciprocal transform: Rᵀ·h - expected_t = torch.einsum("oji,nj->nio", mats, hkl) - # The old (buggy) real-space transform: R·h - plain_r = torch.einsum("oij,nj->nio", mats, hkl) + expected_t = torch.einsum("oji,nj->oni", mats, hkl) + # The buggy real-space transform: R·h + plain_r = torch.einsum("oij,nj->oni", mats, hkl) - assert torch.allclose(out, expected_t), "apply_to_hkl should compute Rᵀ·h" + assert torch.allclose(out.to(mats.dtype), expected_t), ( + "expand_reciprocal should compute Rᵀ·h" + ) # For P3 the two conventions must actually differ (non-symmetric matrix), # otherwise this test would not detect a regression. assert not torch.allclose(expected_t, plain_r), ( diff --git a/tests/unit/symmetry/test_phase_convention.py b/tests/unit/symmetry/test_phase_convention.py index 2d6436e4..6647eb63 100644 --- a/tests/unit/symmetry/test_phase_convention.py +++ b/tests/unit/symmetry/test_phase_convention.py @@ -35,12 +35,31 @@ import torch from torchref.config import get_float_dtype -from torchref.symmetry.reciprocal_symmetry import ( - canonicalize_hkl, - expand_hkl, - reduce_hkl, -) -from torchref.symmetry.spacegroup import SpaceGroup +from torchref.symmetry import SpaceGroup + + +# The HKL verbs live on the space group now. These adapters keep the assertions +# below -- which pin the phase-sign contract -- expressed in terms of the space +# group specifications the cases are parametrised over. +def canonicalize_hkl(hkl, sg, include_friedel=True, device=None): + return SpaceGroup(sg).canonicalize_hkl( + hkl, include_friedel=include_friedel, device=device + ) + + +def expand_hkl(hkl, sg, include_friedel=True, remove_absences=True, device=None): + return SpaceGroup(sg).expand_hkl( + hkl, + include_friedel=include_friedel, + remove_absences=remove_absences, + device=device, + ) + + +def reduce_hkl(hkl, sg, include_friedel=True, device=None): + return SpaceGroup(sg).reduce_hkl( + hkl, include_friedel=include_friedel, device=device + ) # Groups spanning the three regimes above. P1 is the degenerate control (no # translations at all); the screw-axis groups are the ones with real signal. diff --git a/tests/unit/symmetry/test_spacegroup_identity.py b/tests/unit/symmetry/test_spacegroup_identity.py new file mode 100644 index 00000000..aaa21eee --- /dev/null +++ b/tests/unit/symmetry/test_spacegroup_identity.py @@ -0,0 +1,52 @@ +"""Value identity of :class:`~torchref.symmetry.spacegroup.SpaceGroup` and the +setting-preserving round trip through gemmi.""" + +import gemmi +import pytest + +from torchref.symmetry import SpaceGroup + + +@pytest.mark.unit +def test_same_number_different_setting_are_unequal(): + a = SpaceGroup("P 1 21 1") + b = SpaceGroup("P 1 1 21") + assert a.number == b.number == 4 + assert a != b + assert hash(a) != hash(b) + # The reason the identity must be setting-aware: the screw axis moves. + assert a.grid_requirements() != b.grid_requirements() + + +@pytest.mark.unit +def test_same_setting_compares_and_hashes_equal(): + a = SpaceGroup("P 21 21 21") + b = SpaceGroup(19) + assert a == b + assert hash(a) == hash(b) + assert a.key == b.key == a.xhm + + +@pytest.mark.unit +def test_copy_is_equal(): + a = SpaceGroup("C 1 2 1") + assert a.copy() == a + assert hash(a.copy()) == hash(a) + + +@pytest.mark.unit +def test_equality_with_gemmi_spacegroup_is_setting_aware(): + a = SpaceGroup("P 1 21 1") + assert a == gemmi.find_spacegroup_by_name("P 1 21 1") + assert a != gemmi.find_spacegroup_by_name("P 1 1 21") + assert (a == "P 1 21 1") is False + + +@pytest.mark.unit +@pytest.mark.parametrize("xhm", ["R 3:R", "R 3:H", "P 4/n:1", "P 4/n:2"]) +def test_rewrapping_preserves_the_setting(xhm): + sg = SpaceGroup(xhm) + assert sg.xhm == xhm + assert SpaceGroup(sg).xhm == xhm + assert sg._gemmi.xhm() == xhm + assert SpaceGroup(sg).n_ops == sg.n_ops diff --git a/tests/unit/symmetry/test_symmetry.py b/tests/unit/symmetry/test_symmetry.py index 1ddab2bf..afca8c07 100644 --- a/tests/unit/symmetry/test_symmetry.py +++ b/tests/unit/symmetry/test_symmetry.py @@ -6,7 +6,6 @@ import pytest import torch -import torch.nn as nn class TestSpaceGroupInitialization: @@ -131,7 +130,9 @@ def test_rotation_matrices_determinant(self): for i in range(sg.matrices.shape[0]): det = torch.linalg.det(sg.matrices[i]) - assert torch.isclose(torch.abs(det), torch.tensor(1.0, dtype=det.dtype), atol=1e-5) + assert torch.isclose( + torch.abs(det), torch.tensor(1.0, dtype=det.dtype), atol=1e-5 + ) class TestSpaceGroupApplication: @@ -143,20 +144,15 @@ def test_apply_identity(self, random_fractional_coordinates): from torchref.symmetry import SpaceGroup sg = SpaceGroup("P1") - # SpaceGroup expects (N, 3) format coords = random_fractional_coordinates(n_atoms=10) # Shape: (N, 3) - # Apply symmetry (P1 only has identity) - # Output shape is (n_atoms, 3, n_operations) - transformed = sg(coords) + # Expansions are operation-major: (n_ops, N, 3) + transformed = sg.expand_positions(coords) - # Should have shape (n_atoms, 3, n_operations) - assert transformed.shape[0] == 10 # 10 atoms - assert transformed.shape[1] == 3 # 3D coordinates - assert transformed.shape[2] == 1 # 1 operation (identity) - # First (and only) symmetry mate should match original + assert transformed.shape == (1, 10, 3) + # The only mate is the identity assert torch.allclose( - transformed[:, :, 0], + transformed[0], coords.to(device=transformed.device, dtype=transformed.dtype), atol=1e-5, ) @@ -169,26 +165,21 @@ def test_spacegroup_generates_mates(self, random_fractional_coordinates): sg = SpaceGroup("P21") # 2 operations coords = random_fractional_coordinates(n_atoms=5) # (N, 3) format - # Output shape is (n_atoms, 3, n_operations) - transformed = sg(coords) + transformed = sg.expand_positions(coords) - # Should have 2 symmetry operations - assert transformed.shape[2] == 2 + assert transformed.shape == (2, 5, 3) @pytest.mark.unit - def test_spacegroup_callable(self, random_fractional_coordinates): - """SpaceGroup should be callable.""" + def test_expand_to_p1_flattens(self, random_fractional_coordinates): + """expand_to_P1 flattens the operation axis into one coordinate list.""" from torchref.symmetry import SpaceGroup sg = SpaceGroup("P212121") coords = random_fractional_coordinates(n_atoms=10) # (N, 3) format - # Should be callable - # Output shape is (n_atoms, 3, n_operations) - result = sg(coords) + result = sg.expand_to_P1(coords) - assert result is not None - assert result.shape[2] == 4 # 4 symmetry operations + assert result.shape == (40, 3) # 4 operations x 10 atoms class TestSpaceGroupDeviceHandling: @@ -199,10 +190,10 @@ def test_spacegroup_cpu(self): """Test SpaceGroup on CPU.""" from torchref.symmetry import SpaceGroup - sg = SpaceGroup("P21", device=torch.device('cpu')) + sg = SpaceGroup("P21", device=torch.device("cpu")) - assert sg.matrices.device.type == 'cpu' - assert sg.translations.device.type == 'cpu' + assert sg.matrices.device.type == "cpu" + assert sg.translations.device.type == "cpu" @pytest.mark.unit @pytest.mark.gpu @@ -230,7 +221,23 @@ class TestSpaceGroupMapping: """Tests for space group name mapping.""" @pytest.mark.unit - @pytest.mark.parametrize("sg_name", ["P1", "P21", "P212121", "C2", "P21212"]) + @pytest.mark.parametrize( + "sg_name", + [ + "P1", + "P21", + "P212121", + "C2", + "P21212", + "P 1", + "P 21", + "P 21 21 21", + "P 43 21 2", + "P 3 2 1", + "P 6 2 2", + "P 2 3", + ], + ) def test_common_spacegroups(self, sg_name): """Test common crystallographic space groups.""" from torchref.symmetry import SpaceGroup @@ -252,24 +259,31 @@ def test_case_insensitivity(self): pass # Some case variations may not be supported -class TestSymmetryBackwardCompat: - """Tests for backward compatibility with Symmetry alias.""" +class TestSymmetryBase: + """Tests for the crystallography-free Symmetry base class.""" @pytest.mark.unit - def test_symmetry_alias_exists(self): - """Test that Symmetry alias is available.""" - from torchref.symmetry import Symmetry, SpaceGroup + def test_spacegroup_is_a_symmetry(self): + """SpaceGroup specialises Symmetry, so ops-only code accepts either.""" + from torchref.symmetry import SpaceGroup, Symmetry - # Symmetry should be an alias for SpaceGroup - assert Symmetry is SpaceGroup + assert issubclass(SpaceGroup, Symmetry) + assert isinstance(SpaceGroup("P21"), Symmetry) @pytest.mark.unit - def test_symmetry_alias_works(self): - """Test that Symmetry alias works identically to SpaceGroup.""" + def test_symmetry_from_raw_operations(self): + """A Symmetry can be built from an operation list with no space group.""" + import torch + from torchref.symmetry import Symmetry - sym = Symmetry("P21") + matrices = torch.eye(3).unsqueeze(0).repeat(2, 1, 1) + matrices[1] = -matrices[1] + translations = torch.zeros(2, 3) + + sym = Symmetry(matrices=matrices, translations=translations) - assert sym.matrices is not None assert sym.n_ops == 2 - assert sym.name == "P21" + # An inversion pair makes every reflection centric. + hkl = torch.tensor([[1, 2, 3], [4, 0, 1]]) + assert bool(sym.is_centric(hkl).all()) diff --git a/tests/unit/test_dtype_conformance.py b/tests/unit/test_dtype_conformance.py new file mode 100644 index 00000000..8dbb8347 --- /dev/null +++ b/tests/unit/test_dtype_conformance.py @@ -0,0 +1,97 @@ +"""Dtype conformance: no unjustified hardcoded float dtype anywhere in the source. + +The package resolves one float dtype at import (``config.get_float_dtype()``) +and every allocation on a live path is meant to honour it. A literal +``torch.float32`` / ``torch.float64`` baked into an allocation is a latent bug: +MPS has no float64, so a float64 config silently downcasts and a float32 config +silently upcasts -- neither raises, and the corruption surfaces far from its +cause. + +This guard makes the rule enforceable. Outside a small set of inherently-exempt +modules (triton kernels, the config dtype maps, offline scripts), every +``torch.`` reference must carry a one-line justification:: + + x = torch.tensor(v, dtype=torch.float64) # dtype-ok: SVD needs f64 stability + +The marker may sit on the reference's own line or the line immediately above it. +The point is not to ban hardcoded dtypes -- some are correct (dtype validation, +deliberate high-precision accumulation, backend capability declarations) -- but +to force each one to say *why* it deviates, so a reviewer can tell a considered +choice from an oversight at a glance. +""" + +from pathlib import Path + +import pytest + +from tests.helpers.dtype_inventory import ( + EXEMPT_PREFIXES, + JUSTIFY_MARKER, + find_hardcoded_dtypes, + is_exempt, +) + +_PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "torchref" + +# The config getter each category should defer to, quoted in the failure message. +_GETTER = { + "float": "config.get_float_dtype()", + "int": "config.get_int_dtype()", + "complex": "config.get_complex_dtype()", +} + + +@pytest.mark.unit +def test_no_unjustified_hardcoded_dtype(): + """Every hardcoded float/int/complex dtype on a live path is justified.""" + uses = find_hardcoded_dtypes(_PACKAGE_ROOT) + offenders = [u for u in uses if not is_exempt(u.rel_path) and not u.justified] + + def fmt(u): + return f"{u.where} {u.dtype} -> {_GETTER[u.category]} | {u.line}" + + assert not offenders, ( + f"{len(offenders)} hardcoded dtype(s) with no justification. Either switch " + f"to the config default for that category, or if the literal is deliberate " + f"add a '{JUSTIFY_MARKER} ' marker on the line or the block above:\n " + + "\n ".join(fmt(u) for u in offenders) + ) + + +@pytest.mark.unit +def test_justifications_carry_a_reason(): + """A ``# dtype-ok:`` marker must be followed by an actual reason, not left blank.""" + blank = [] + for path in sorted(_PACKAGE_ROOT.rglob("*.py")): + for i, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + if JUSTIFY_MARKER in line: + reason = line.split(JUSTIFY_MARKER, 1)[1].strip() + if not reason: + rel = path.relative_to(_PACKAGE_ROOT.parent) + blank.append(f"{rel}:{i}") + + assert not blank, ( + f"'{JUSTIFY_MARKER}' markers with no reason after the colon:\n " + + "\n ".join(blank) + ) + + +@pytest.mark.unit +def test_exempt_prefixes_still_match_something(): + """Stop ``EXEMPT_PREFIXES`` accumulating entries for paths that are long gone. + + An exemption that no longer matches any source is either a typo or a stale + excuse -- both hide the fact that the module it was meant to cover is now + unguarded (or renamed and silently re-included). + """ + uses = find_hardcoded_dtypes(_PACKAGE_ROOT) + seen_paths = {u.rel_path for u in uses} + stale = [ + p + for p in EXEMPT_PREFIXES + if not any(rp.startswith(p) for rp in seen_paths) + ] + assert not stale, ( + "EXEMPT_PREFIXES entries that match no source file with a hardcoded " + f"float dtype (stale or mistyped): {stale}" + ) diff --git a/tests/unit/topology/__init__.py b/tests/unit/topology/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/topology/test_energy_types.py b/tests/unit/topology/test_energy_types.py new file mode 100644 index 00000000..d296a038 --- /dev/null +++ b/tests/unit/topology/test_energy_types.py @@ -0,0 +1,112 @@ +"""CCP4 energy types travel from the monomer template onto the atom graph. + +What is pinned: the type column survives the CIF reader, link modifications retype the +atoms they change (an in-chain backbone nitrogen is an amide ``NH1``, only the +N-terminus keeps the free-amine ``NT3``), the template hydrogen count lands on every +template atom, and the bundled per-type table covers every type the standard residues +use. +""" + +import csv + +import numpy as np +import pytest + +from torchref import PATH_TORCHREF_DATA +from torchref.model.model import Model +from torchref.topology.hydrogens import template_atom_types + + +@pytest.fixture(scope="module") +def heavy_1daw(pdb_dir): + model = Model(verbose=0, strip_H=True, add_hydrogens=False) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + return model + + +@pytest.mark.unit +def test_reader_keeps_type_energy(heavy_1daw): + atoms = heavy_1daw.restraints.cif_dict["ASN"]["atoms"] + types = dict(zip(atoms["atom_id"].str.strip(), atoms["type_energy"])) + assert types["N"] == "NT3" + assert types["ND2"] == "NH2" + assert types["OD1"] == "O" + assert types["CA"] == "CH1" + + +@pytest.mark.unit +def test_template_atom_types_counts_hydrogens(heavy_1daw): + types, h_count = template_atom_types(heavy_1daw.restraints.cif_dict["ASN"]) + assert types["OXT"] == "OC" + assert h_count["ND2"] == 2 + assert h_count["CB"] == 2 + assert h_count["N"] == 3 + assert "OD1" not in h_count + + +@pytest.mark.unit +def test_atom_graph_carries_types_with_link_modifications(heavy_1daw): + """Peptide-linked backbone N is retyped NH1; the chain start stays NT3.""" + atoms = heavy_1daw.restraints.topology.atoms + names = atoms.name.astype(str) + resseq = heavy_1daw.pdb["resseq"].values + chain = heavy_1daw.pdb["chainid"].astype(str).values + resnames = heavy_1daw.pdb["resname"].str.strip().values + is_n = names == "N" + first_res = min(resseq[chain == chain[0]]) + n_types = atoms.energy_type[is_n] + n_first = atoms.energy_type[is_n & (resseq == first_res) & (chain == chain[0])] + assert set(n_first.tolist()) == {"NT3"} + # In-chain amide N is NH1; proline's tertiary N is NH0; only chain starts stay NT3. + assert set(n_types.tolist()) <= {"NH1", "NH0", "NT3"} + assert (n_types == "NH1").sum() > 0.8 * is_n.sum() + assert (atoms.energy_type[is_n & (resnames == "PRO")] == "NH0").all() + assert (atoms.energy_type[(names == "O") & (resnames != "HOH")] == "O").all() + assert (atoms.energy_type[(names == "CA") & (resnames != "GLY")] == "CH1").all() + assert (atoms.energy_type[(names == "CA") & (resnames == "GLY")] == "CH2").all() + + +@pytest.mark.unit +def test_template_h_count_and_implicit_hydrogens(heavy_1daw): + """A heavy-only model is missing exactly the hydrogens its templates carry.""" + atoms = heavy_1daw.restraints.topology.atoms + names = atoms.name.astype(str) + polymer = heavy_1daw.pdb["ATOM"].astype(str).str.strip().values == "ATOM" + single_atom = np.isin(heavy_1daw.pdb["resname"].str.strip().values, ["HOH", "MG"]) + counts = atoms.template_h_count.cpu().numpy() + assert (counts[names == "CB"] >= 1).all() + assert (counts[(names == "O") & polymer] == 0).all() + missing = atoms.implicit_h_count().cpu().numpy() + np.testing.assert_array_equal(missing[counts >= 0], counts[counts >= 0]) + # Waters and ions have no template, so their count is unknown; every polymer + # atom's is known. + assert (counts[polymer] >= 0).all() + assert (counts[single_atom] == -1).all() + + +@pytest.mark.unit +def test_hydrogenated_model_completes_charged_amines(pdb_dir): + """Explicit hydrogen generation fills the polymer's template hydrogen counts.""" + model = Model(verbose=0, add_hydrogens=True) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + atoms = model.restraints.topology.atoms + missing = atoms.implicit_h_count().cpu().numpy() + polymer = model.pdb["ATOM"].astype(str).str.strip().values == "ATOM" + assert (missing[polymer] == 0).all() + ammonium = np.isin(atoms.energy_type, ["NT", "NT1", "NT2", "NT3", "NT4"]) + assert ammonium.any() + assert (missing[ammonium & polymer] == 0).all() + + +@pytest.mark.unit +def test_bundled_table_covers_standard_residue_types(heavy_1daw): + with open(f"{PATH_TORCHREF_DATA}/ener_lib_atoms.csv") as handle: + rows = [r for r in csv.DictReader(l for l in handle if not l.startswith("#"))] + table = {row["type"] for row in rows} + used = set(heavy_1daw.restraints.topology.atoms.energy_type.tolist()) - {""} + assert used and used <= table + by_type = {row["type"]: row for row in rows} + assert by_type["NH1"]["hb_type"] == "D" + assert by_type["O"]["hb_type"] == "A" + assert by_type["OH1"]["hb_type"] == "B" + assert float(by_type["CH3"]["vdwh_radius"]) > float(by_type["CH3"]["vdw_radius"]) diff --git a/tests/unit/topology/test_equivalence.py b/tests/unit/topology/test_equivalence.py new file mode 100644 index 00000000..0e276804 --- /dev/null +++ b/tests/unit/topology/test_equivalence.py @@ -0,0 +1,326 @@ +"""The topology graph must reproduce the restraint builders' edge sets exactly. + +Compares as **sets** of atom-index tuples, per edge type and per origin. Row order is +deliberately not compared: the graph lays edges out in a canonical order, whereas the +old storage concatenated origins in Python ``set`` iteration order and so varied +between processes. +""" + +import pytest +import torch + +from torchref.model.model import Model +from torchref.topology import build_topology + +#: Structures chosen to cover the cases the builders treat specially: alternative +#: conformations, pre-existing hydrogens, disulfides, multi-compound ligands, and a +#: metal plus a nucleotide analogue. +STRUCTURES = [ + "7L84", + "1AK5_with_H", + "3VRJ", + "3A5V", + "1DAW", +] + + +def _tuples(indices: torch.Tensor) -> set: + """Rows of an index tensor as a set of int tuples.""" + return {tuple(int(v) for v in row) for row in indices.cpu().numpy()} + + +def _current_edges(restraints) -> dict: + """``{edge type: {origin: set of tuples}}`` from the existing restraint storage.""" + out = {} + for rtype in ("bond", "angle", "torsion"): + out[rtype] = {} + for origin in restraints.restraints[rtype].keys(): + if origin == "all": + continue + group = restraints.restraints[rtype][origin] + if group is None or group.get("indices") is None: + continue + out[rtype][origin] = _tuples(group["indices"]) + + out["chiral"] = {} + chiral = restraints.restraints.get("chiral") + if chiral is not None and chiral.get("indices") is not None: + out["chiral"]["intra"] = _tuples(chiral["indices"]) + + out["plane"] = {} + for key in restraints.restraints["plane"].keys(): + group = restraints.restraints["plane"][key] + if group is None or group.get("indices") is None: + continue + out["plane"][key] = _tuples(group["indices"]) + return out + + +@pytest.fixture(scope="module") +def built(request, pdb_dir): + """``(topology, current edge sets)`` for one structure, built once per module.""" + cache = {} + + def _build(code): + if code not in cache: + path = pdb_dir / f"{code}.pdb" + if not path.exists(): + pytest.skip(f"{code}.pdb not bundled") + model = Model(verbose=0) + model.load_pdb(str(path)) + model.set_restraints_cif(None) + restraints = model.restraints + topology = build_topology( + model.pdb, + restraints.cif_dict, + link_dict=getattr(restraints, "link_dict", None), + link_list=getattr(restraints, "link_list", None), + links=restraints.links, + xyz=model.xyz().detach(), + verbose=0, + ) + cache[code] = (topology, _current_edges(restraints)) + return cache[code] + + return _build + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +@pytest.mark.parametrize("edge_type", ["bond", "angle", "torsion", "chiral"]) +def test_edge_sets_match_builders(built, code, edge_type): + """Every origin of every edge type holds exactly the builders' index tuples.""" + topology, current = built(code) + block = topology.edge_block(edge_type) + graph = {origin: block.tuple_set(origin) for origin in block.origins()} + expected = current.get(edge_type, {}) + + assert set(graph) == set(expected), ( + f"{code} {edge_type}: origins differ -- " + f"graph {sorted(graph)} vs builders {sorted(expected)}" + ) + for origin in sorted(expected): + missing = expected[origin] - graph[origin] + extra = graph[origin] - expected[origin] + assert not missing and not extra, ( + f"{code} {edge_type}/{origin}: {len(missing)} edges the builders " + f"produced are absent from the graph, {len(extra)} are only in the graph" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +def test_plane_sets_match_builders(built, code): + """Planes match per atom count, pooling the origins the graph splits them into.""" + topology, current = built(code) + graph = {} + for size, block in topology.atoms.planes.items(): + graph.setdefault(f"{size}_atoms", set()).update(block.tuple_set()) + expected = current["plane"] + + assert set(graph) == set(expected), ( + f"{code} planes: size groups differ -- " + f"graph {sorted(graph)} vs builders {sorted(expected)}" + ) + for key in sorted(expected): + assert graph[key] == expected[key], ( + f"{code} plane/{key}: " + f"{len(expected[key] - graph[key])} missing, " + f"{len(graph[key] - expected[key])} extra" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +def test_exclusions_reproduce_current_set(built, code): + """The restraint-edge exclusions equal what the non-bonded term is given today.""" + topology, _ = built(code) + from_edges = topology.atoms.exclusions_from_restraint_edges() + + expected = set() + for edge_type, cols in (("bond", (0, 1)), ("angle", (0, 2)), ("torsion", (0, 3))): + block = topology.edge_block(edge_type) + for row in block.indices.cpu().numpy(): + a, b = int(row[cols[0]]), int(row[cols[1]]) + if a != b: + expected.add((min(a, b), max(a, b))) + + assert from_edges == expected + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +def test_connectivity_exclusions_are_a_superset(built, code): + """Connectivity-derived exclusions cover the restraint-derived ones, and then some. + + The difference is the defect the connectivity path fixes: a pair that is 1-3 or 1-4 + bonded but whose angle or torsion the monomer library does not restrain is currently + not excluded, so the non-bonded term pushes it apart. + """ + topology, _ = built(code) + from_edges = topology.atoms.exclusions_from_restraint_edges() + from_bonds = topology.atoms.exclusions_12_13_14() + + assert from_edges <= from_bonds, ( + f"{code}: {len(from_edges - from_bonds)} restraint-derived exclusions are " + f"not reachable within three bonds, which should be impossible" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("code", ["7L84", "3VRJ"]) +def test_adjacency_matches_bond_block(built, code): + """Every bond appears in both atoms' neighbour lists, and nothing else does.""" + topology, _ = built(code) + atoms = topology.atoms + + from_adjacency = set() + for i in range(atoms.n_atoms): + for j in atoms.neighbors(i).cpu().tolist(): + from_adjacency.add((min(i, j), max(i, j))) + + from_block = set() + for a, b in atoms.bonds.indices.cpu().numpy(): + from_block.add((min(int(a), int(b)), max(int(a), int(b)))) + + assert from_adjacency == from_block + # Each distinct partner once: a bond row repeated per altloc conformer or per + # duplicated LINK record does not add to the degree. + assert int(atoms.degree().sum()) == 2 * len(from_block) + + +@pytest.mark.unit +def test_layout_is_reproducible(built, pdb_dir): + """Two builds of one structure lay the blocks out identically. + + The canonical order removes the process-to-process variation the old storage had, + where origins were concatenated in ``set`` iteration order. + """ + topology_a, _ = built("7L84") + + model = Model(verbose=0) + model.load_pdb(str(pdb_dir / "7L84.pdb")) + model.set_restraints_cif(None) + restraints = model.restraints + topology_b = build_topology( + model.pdb, + restraints.cif_dict, + link_dict=getattr(restraints, "link_dict", None), + link_list=getattr(restraints, "link_list", None), + links=restraints.links, + xyz=model.xyz().detach(), + verbose=0, + ) + + for edge_type in ("bond", "angle", "torsion", "chiral"): + a = topology_a.edge_block(edge_type) + b = topology_b.edge_block(edge_type) + assert a.origin_bounds == b.origin_bounds, edge_type + assert torch.equal(a.indices, b.indices), edge_type + + +@pytest.mark.unit +def test_origin_slices_share_storage(built): + """A per-origin view is a slice of the block, not a copy.""" + topology, _ = built("7L84") + block = topology.atoms.bonds + origin = block.origins()[0] + view = block.origin(origin) + + assert view.data_ptr() == block.indices.data_ptr() + saved = int(block.indices[0, 0]) + block.indices[0, 0] = saved + 1000 + assert int(view[0, 0]) == saved + 1000 + block.indices[0, 0] = saved + + +@pytest.mark.unit +def test_residue_identity_includes_insertion_code(built): + """Residue nodes are keyed on ``(chain, resseq, icode)``. + + The restraint builders group on ``(chain, resseq)`` alone, which merges residue 100 + with 100A and leaves the inserted residue without intra-residue restraints. + """ + topology, _ = built("7L84") + residues = topology.residues + keys = [residues.key(i) for i in range(residues.n_residues)] + + assert len(keys) == len(set(keys)), "residue identity is not unique" + assert all(len(k) == 3 for k in keys) + + +@pytest.mark.unit +def test_residue_join_recovers_per_atom_identity(built): + """Per-atom residue name is recovered via ``residue_of``, not stored per atom.""" + topology, _ = built("7L84") + + for atom in (0, topology.n_atoms // 2, topology.n_atoms - 1): + residue = topology.residue_of_atom(atom) + assert topology.resname_of_atom(atom) == topology.residues.resname[residue] + assert ( + topology.residues.atom_start[residue] + <= atom + < topology.residues.atom_end[residue] + ) + + +@pytest.mark.unit +def test_connectivity_exclusions_match_brute_force(): + """The vectorised path-walk agrees with a plain breadth-first reference. + + ``test_connectivity_exclusions_are_a_superset`` alone would also pass for an + implementation that returned every pair, so the walk is pinned against an + independent one on a graph small enough to enumerate: a six-ring with two + substituents, which exercises the ring closure and the 1-4 wrap-around. + """ + import numpy as np + + from torchref.topology import EdgeBlock + from torchref.topology.atom_graph import AtomGraph + + bonds = [ + (0, 1), + (1, 2), + (2, 3), + (3, 4), + (4, 5), + (5, 0), # six-ring + (0, 6), # substituent on 0 + (6, 7), # and one more out + ] + n = 8 + + def block(rows, arity, edge_type): + return EdgeBlock.from_origins( + {"intra": np.asarray(rows, dtype=np.int64).reshape(-1, arity)}, + arity, + edge_type, + ) + + graph = AtomGraph( + name=np.array([f"A{i}" for i in range(n)]), + element=np.array(["C"] * n), + altloc=np.array([" "] * n), + residue_of=torch.zeros(n, dtype=torch.int64), + bonds=block(bonds, 2, "bond"), + angles=EdgeBlock.empty(3), + torsions=EdgeBlock.empty(4), + chirals=EdgeBlock.empty(4), + ) + + adjacency = {i: set() for i in range(n)} + for a, b in bonds: + adjacency[a].add(b) + adjacency[b].add(a) + + expected = set() + for start in range(n): + frontier = {start} + seen = {start} + for _ in range(3): + frontier = {j for i in frontier for j in adjacency[i]} - seen + seen |= frontier + for other in frontier: + expected.add((min(start, other), max(start, other))) + + assert graph.exclusions_12_13_14() == expected diff --git a/tests/unit/topology/test_hydrogen_frames.py b/tests/unit/topology/test_hydrogen_frames.py new file mode 100644 index 00000000..268f6a52 --- /dev/null +++ b/tests/unit/topology/test_hydrogen_frames.py @@ -0,0 +1,154 @@ +"""Riding frames read off the bond graph, and their bookkeeping across table edits. + +Every hydrogen bonded to a heavy atom gets a frame; a parent with a single heavy +neighbour borrows a grandparent so the hydrogen turns with its torsion; and the frames +planned before an insertion agree, after remapping, with frames rebuilt on the +inserted table. +""" + +import numpy as np +import pytest +import torch + +from torchref.base.coordinates.local_frame import ( + frame_is_degenerate, + local_frame_coordinates, + place_local_frame, +) +from torchref.model.model import Model +from torchref.topology.hydrogens import ( + HydrogenFrames, + augment_atom_table_with_maps, + hydrogen_frames, + optimise_free_torsions, + plan_hydrogens, +) + + +@pytest.fixture(scope="module") +def heavy_and_plan(pdb_dir): + """Heavy-only 1DAW with its hydrogen plan.""" + model = Model(verbose=0, strip_H=True, add_hydrogens=False) + model.load_pdb(str(pdb_dir / "1DAW.pdb")) + restraints = model.restraints + xyz = model.xyz().detach() + plan = plan_hydrogens(restraints.topology, restraints.cif_dict, xyz) + optimise_free_torsions(plan, restraints.topology, xyz) + return model, restraints, plan + + +@pytest.fixture(scope="module") +def hydrogenated(heavy_and_plan): + """The same structure with the plan inserted, plus the maps the insertion made.""" + model, restraints, plan = heavy_and_plan + augmented, old_to_new, plan_to_new = augment_atom_table_with_maps( + model.pdb, plan, restraints.topology + ) + full = Model(verbose=0, strip_H=False, add_hydrogens=False) + cell, spacegroup = model.cell.data.cpu().numpy(), model.spacegroup + + def reader(): + return augmented, cell, spacegroup + + reader.links = model.ctx.links + full.load(reader) + return full, old_to_new, plan_to_new + + +@pytest.mark.unit +def test_every_bonded_hydrogen_gets_a_frame_or_orientation(hydrogenated): + """Hydrogens have a heavy-atom frame or an independently rotatable water group.""" + full, _, _ = hydrogenated + frames = hydrogen_frames(full.restraints.topology) + n_h = int((full.pdb["element"].str.strip() == "H").sum()) + assert frames.n_hydrogens == n_h + assert (frames.frame_valid | (frames.rotation_group >= 0)).all() + assert (frames.parent_row >= 0).all() + is_h = full.restraints.topology.atoms.is_hydrogen.cpu().numpy() + assert not is_h[frames.parent_row].any() + assert not is_h[frames.n1_row[frames.n1_row >= 0]].any() + assert not is_h[frames.n2_row[frames.n2_row >= 0]].any() + + +@pytest.mark.unit +def test_single_neighbour_parents_borrow_the_grandparent(hydrogenated): + """A hydroxyl or methyl hydrogen is framed on the bond it rotates about.""" + full, _, _ = hydrogenated + frames = hydrogen_frames(full.restraints.topology) + names = full.pdb["name"].str.strip().values + resnames = full.pdb["resname"].str.strip().values + seen = {} + for parent, n1, n2 in zip(frames.parent_row, frames.n1_row, frames.n2_row): + key = (resnames[parent], names[parent]) + seen.setdefault(key, (names[n1], names[n2])) + assert seen[("SER", "OG")] == ("CB", "CA") + assert seen[("LYS", "NZ")] == ("CE", "CD") + # A two-neighbour parent frames on its own neighbours. + assert set(seen[("ALA", "CA")]) <= {"N", "C", "CB"} + + +@pytest.mark.unit +def test_planned_frames_match_frames_rebuilt_on_the_augmented_table( + heavy_and_plan, hydrogenated +): + """Remapping the pre-insertion frames reproduces the post-insertion ones.""" + _, restraints, plan = heavy_and_plan + full, old_to_new, plan_to_new = hydrogenated + planned = hydrogen_frames(restraints.topology, plan) + assert planned.n_planned == plan.n_hydrogens + carried = planned.remap(old_to_new).fill_planned_rows(plan_to_new).sorted_by_row() + rebuilt = hydrogen_frames(full.restraints.topology).sorted_by_row() + for field in ("h_row", "parent_row", "n1_row", "n2_row"): + np.testing.assert_array_equal(getattr(carried, field), getattr(rebuilt, field)) + np.testing.assert_array_equal(carried.frame_valid, rebuilt.frame_valid) + + +@pytest.mark.unit +def test_positions_round_trip_through_their_frames(hydrogenated): + """Placed hydrogens are reproduced exactly from heavy atoms and local offsets.""" + full, _, _ = hydrogenated + frames = hydrogen_frames(full.restraints.topology) + xyz = full.xyz().detach().cpu() + t = frames.to_tensors() + p, n1, n2 = xyz[t["parent_row"]], xyz[t["n1_row"]], xyz[t["n2_row"]] + h = xyz[t["h_row"]] + assert not frame_is_degenerate(p, n1, n2)[t["frame_valid"]].any() + local = local_frame_coordinates(p, n1, n2, h) + back = place_local_frame(p, n1, n2, local, t["frame_valid"], h - p) + tolerance = 1e-4 if xyz.dtype == torch.float32 else 1e-9 + assert float((back - h).norm(dim=1).max()) < tolerance + + +@pytest.mark.unit +def test_remap_drops_orphans_and_degrades_lost_frames(): + """A hydrogen whose parent vanishes is dropped; a lost n2 leaves a rigid frame.""" + frames = HydrogenFrames( + h_row=np.array([5, 6, 7]), + parent_row=np.array([1, 2, 3]), + n1_row=np.array([0, 1, 2]), + n2_row=np.array([2, 3, 4]), + frame_valid=np.array([True, True, True]), + ) + # Drop atoms 2 and 7: the first hydrogen loses n2, the second loses its parent, + # the third is gone itself. + old_to_new = np.array([0, 1, -1, 2, 3, 4, 5, -1]) + out = frames.remap(old_to_new) + assert out.h_row.tolist() == [4] + assert out.parent_row.tolist() == [1] + assert out.n2_row.tolist() == [-1] + assert out.frame_valid.tolist() == [False] + + +@pytest.mark.unit +def test_tensor_round_trip_preserves_frames(): + """to_tensors / from_tensors is lossless.""" + frames = HydrogenFrames( + h_row=np.array([3, 4]), + parent_row=np.array([1, 1]), + n1_row=np.array([0, 0]), + n2_row=np.array([2, -1]), + frame_valid=np.array([True, False]), + ) + back = HydrogenFrames.from_tensors(**frames.to_tensors()) + for field in ("h_row", "parent_row", "n1_row", "n2_row", "frame_valid"): + np.testing.assert_array_equal(getattr(back, field), getattr(frames, field)) diff --git a/tests/unit/topology/test_hydrogens.py b/tests/unit/topology/test_hydrogens.py new file mode 100644 index 00000000..bcd2620d --- /dev/null +++ b/tests/unit/topology/test_hydrogens.py @@ -0,0 +1,401 @@ +"""Hydrogen generation from monomer templates, driven by the bond graph. + +The properties asserted here are the ones that make template instantiation +trustworthy: every hydrogen lands at its library bond length, none is placed in a +direction the geometry does not determine, the count per parent respects the valence +left over after the graph's real bonds, and the free-torsion set is exactly the centres +whose dihedral the template cannot know. +""" + +import numpy as np +import pytest + +from torchref.model.model import Model +from torchref.topology.hydrogens import ( + STANDARD_VALENCE, + _template, + augment_atom_table, + optimise_free_torsions, + plan_hydrogens, +) + +# 3E98 brings HETATM selenomethionines bonded through LINK records and split side chains. +STRUCTURES = ["7L84", "1DAW", "3E98"] + + +@pytest.fixture(scope="module") +def built(pdb_dir): + """``(model, restraints, plan)`` per structure, built once.""" + cache = {} + + def _build(code): + if code not in cache: + # add_hydrogens=False: these tests exercise generation itself, so the model + # has to arrive without the hydrogens the loader would otherwise add. + model = Model(verbose=0, add_hydrogens=False, strip_H=True) + model.load_pdb(str(pdb_dir / f"{code}.pdb")) + model.set_restraints_cif(None) + restraints = model.restraints + plan = plan_hydrogens( + restraints.topology, restraints.cif_dict, model.xyz().detach() + ) + cache[code] = (model, restraints, plan) + return cache[code] + + return _build + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +def test_hydrogens_sit_at_their_library_bond_length(built, code): + """Placement is exact, not approximate: the parent distance is the library value.""" + model, _, plan = built(code) + assert plan.n_hydrogens > 0 + + coords = model.xyz().detach().cpu().numpy() + distance = np.linalg.norm(plan.position - coords[plan.parent], axis=1) + assert np.abs(distance - plan.bond_length).max() < 1e-9 + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +def test_every_candidate_hydrogen_is_placed(built, code): + """No hydrogen is dropped for want of a determined direction. + + A hydrogen is only planned once a strategy has fixed its direction, so a shortfall + here means some centre fell through all three. The earlier two-shell alignment left + 12% of side-chain hydrogens beyond 1.5 A of their parent and discarded them. + """ + model, restraints, plan = built(code) + topology = restraints.topology + atoms = topology.atoms + residues = topology.residues + assert plan.n_hydrogens, "no hydrogens planned" + + is_h = atoms.is_hydrogen + altlocs = np.char.strip(atoms.altloc.astype(str)) + names = np.char.strip(atoms.name.astype(str)) + checked = 0 + for residue in range(residues.n_residues): + start, end = int(residues.atom_start[residue]), int(residues.atom_end[residue]) + template = _template( + restraints.cif_dict, str(residues.resname[residue]).strip() + ) + if template is None: + continue + # Residues with altlocs plan one hydrogen per conformer; the two-sided count + # below is for the plain case, where the graph degree is the whole story. + if (altlocs[start:end] != "").any(): + continue + for parent in range(start, end): + template_h = template["h_count"].get(names[parent], 0) + if template_h == 0: + continue + neighbours = atoms.neighbors(parent) + heavy = int((~is_h[neighbours]).sum()) + element = str(atoms.element[parent]).strip().upper() + template_heavy = len(template["heavy_adjacency"].get(names[parent], [])) + extra_bonds = max(0, heavy - template_heavy) + valence = STANDARD_VALENCE.get(element, 4) + if ( + element == "N" + and extra_bonds == 0 + and str(atoms.energy_type[parent]).startswith("NT") + ): + valence = 4 + expected = max(0, min(valence - heavy, template_h - extra_bonds)) + placed = int((plan.parent == parent).sum()) + assert placed == expected, ( + f"{code}: atom {parent} ({names[parent]} {element}) has {heavy} heavy " + f"bonds against {template_heavy} in the template and {template_h} " + f"template hydrogens, so {expected} expected, {placed} planned" + ) + checked += 1 + assert checked > 0 + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +def test_free_torsions_are_exactly_the_single_neighbour_centres(built, code): + """A dihedral is free when the parent has one heavy neighbour, and only then.""" + _, restraints, plan = built(code) + topology = restraints.topology + is_h = topology.atoms.is_hydrogen + + for i in range(plan.n_hydrogens): + parent = int(plan.parent[i]) + neighbours = topology.atoms.neighbors(parent) + heavy = int((~is_h[neighbours]).sum()) + residue = int(topology.atoms.residue_of[parent]) + water = str(topology.residues.resname[residue]).strip() == "HOH" + assert (plan.group[i] >= 0) == (heavy == 1 and not water), ( + f"{code}: hydrogen {plan.name[i]} on atom {parent} with {heavy} heavy " + f"neighbours has group {plan.group[i]}" + ) + + +@pytest.mark.unit +def test_hydroxyl_rotates_and_backbone_amide_does_not(built): + """The chemistry the graph criterion is meant to capture, spot-checked. + + A serine hydroxyl hangs off an oxygen bonded only to CB, so its dihedral is free. A + backbone amide nitrogen is bonded to CA and to the preceding residue's carbon, which + fixes its hydrogen entirely. + """ + _, restraints, plan = built("7L84") + topology = restraints.topology + names = topology.atoms.name.astype(str) + resnames = topology.residues.resname + + free_parents = {int(p) for p, g in zip(plan.parent, plan.group) if g >= 0} + fixed_parents = {int(p) for p, g in zip(plan.parent, plan.group) if g < 0} + + hydroxyl = [ + int(p) + for p in free_parents + if names[p] == "OG" + and str(resnames[topology.residue_of_atom(p)]).strip() == "SER" + ] + assert hydroxyl, "no serine hydroxyl was treated as a free torsion" + + amide = [p for p in fixed_parents if names[p] == "N"] + assert amide, "no backbone amide nitrogen was treated as determined" + + # A backbone nitrogen rotates exactly when nothing is bonded to it on the other + # side: an N-terminal ammonium does, an in-chain amide does not. Tied to the + # residue graph's link edges, so a missing peptide bond would show up here. + peptide_links = topology.residues.links_of_kind("TRANS") + accepts_link = set(peptide_links[:, 1].tolist()) if len(peptide_links) else set() + + for parent in free_parents: + if names[parent] != "N": + continue + residue = topology.residue_of_atom(parent) + assert residue not in accepts_link, ( + f"nitrogen {parent} in residue {topology.residues.key(residue)} was " + f"treated as rotatable even though a peptide bond reaches it" + ) + for parent in fixed_parents: + if names[parent] != "N": + continue + residue = topology.residue_of_atom(parent) + assert residue in accepts_link, ( + f"nitrogen {parent} in residue {topology.residues.key(residue)} was " + f"treated as determined but no peptide bond reaches it" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +def test_torsion_scan_preserves_bond_lengths(built, code): + """The scan rotates about a bond, so it cannot change any bond length.""" + model, restraints, plan = built(code) + coords = model.xyz().detach() + + scanned = plan_hydrogens(restraints.topology, restraints.cif_dict, coords) + before = scanned.position.copy() + optimise_free_torsions(scanned, restraints.topology, coords) + + numpy_coords = coords.cpu().numpy() + distance = np.linalg.norm(scanned.position - numpy_coords[scanned.parent], axis=1) + assert np.abs(distance - scanned.bond_length).max() < 1e-9 + + moved = np.linalg.norm(scanned.position - before, axis=1) > 1e-6 + assert moved.any(), "the scan changed nothing at all" + assert not moved[ + ~scanned.rotatable + ].any(), "the scan moved a hydrogen whose torsion is not free" + + +@pytest.mark.unit +def test_scan_reduces_clash(built): + """Scanned hydrogens end up no closer to heavy atoms than they started.""" + model, restraints, plan = built("7L84") + coords = model.xyz().detach() + topology = restraints.topology + + scanned = plan_hydrogens(topology, restraints.cif_dict, coords) + numpy_coords = coords.cpu().numpy() + heavy = numpy_coords[~topology.atoms.is_hydrogen.cpu().numpy()] + + def closest(positions): + gaps = np.linalg.norm(positions[:, None, :] - heavy[None, :, :], axis=-1) + # The parent itself is always the nearest heavy atom; take the next one. + return np.sort(gaps, axis=1)[:, 1] + + rotatable = scanned.rotatable + before = closest(scanned.position[rotatable]) + optimise_free_torsions(scanned, topology, coords) + after = closest(scanned.position[rotatable]) + + assert after.min() >= before.min() - 1e-9, "the scan made the worst clash worse" + + +@pytest.mark.unit +@pytest.mark.parametrize("code", STRUCTURES) +def test_augmented_table_keeps_residues_contiguous(built, code): + """Hydrogens are inserted into their residue, not appended after everything. + + The residue partition is built from contiguous runs of ``(chain, resseq, icode)``, + so appending hydrogens at the end would split every hydrogenated residue in two. + """ + model, restraints, plan = built(code) + augmented = augment_atom_table(model.pdb, plan, restraints.topology) + + assert len(augmented) == len(model.pdb) + plan.n_hydrogens + assert (augmented["index"].values == np.arange(len(augmented))).all() + + key = ( + augmented[["chainid", "resseq", "icode"]] + .astype(str) + .agg("|".join, axis=1) + .values + ) + runs = 1 + int((key[1:] != key[:-1]).sum()) + assert runs == len(set(key)), "a residue was split into non-adjacent runs" + + +@pytest.mark.unit +def test_waters_receive_two_hydrogens_with_initial_orientations(built): + """Explicit generation supplies two HOH hydrogens, including coordinated waters.""" + _, restraints, plan = built("7L84") + topology = restraints.topology + + waters = [ + i + for i in range(topology.n_residues) + if str(topology.residues.resname[i]).strip() == "HOH" + ] + assert waters, "7L84 has no waters, so this asserts nothing" + for residue in waters: + assert int((plan.residue == residue).sum()) == 2 + + +@pytest.mark.unit +def test_hydrogenate_returns_a_consistent_model(pdb_dir): + """The end-to-end path yields a model whose tensors, table and restraints agree.""" + model = Model(verbose=0, add_hydrogens=False, strip_H=True) + model.load_pdb(str(pdb_dir / "7L84.pdb")) + model.set_restraints_cif(None) + n_heavy = len(model.pdb) + + hydrogenated = model.hydrogenate(verbose=0) + + assert hydrogenated.ctx.strip_H is False + assert len(hydrogenated.pdb) > n_heavy + assert hydrogenated.xyz().shape[0] == len(hydrogenated.pdb) + assert hydrogenated.adp().shape[0] == len(hydrogenated.pdb) + assert len(model.pdb) == n_heavy, "the original model was modified" + + elements = hydrogenated.pdb["element"].astype(str).str.strip().values + n_h = int((elements == "H").sum()) + assert n_h == len(hydrogenated.pdb) - n_heavy + + # Every hydrogen carries exactly one bond restraint, at library geometry. + restraints = hydrogenated.restraints + bonds = restraints.restraints["bond"]["all"]["indices"].cpu().numpy() + references = restraints.restraints["bond"]["all"]["references"].cpu().numpy() + coords = hydrogenated.xyz().detach().cpu().numpy() + is_h = elements == "H" + involves_h = is_h[bonds[:, 0]] | is_h[bonds[:, 1]] + + assert int(involves_h.sum()) == n_h + lengths = np.linalg.norm(coords[bonds[:, 0]] - coords[bonds[:, 1]], axis=1) + deviation = np.sqrt(((lengths[involves_h] - references[involves_h]) ** 2).mean()) + assert deviation < 0.02, f"placed hydrogens deviate by {deviation:.4f} A RMS" + + +@pytest.mark.unit +def test_strip_H_removes_deposited_hydrogens(pdb_dir): + """The opt-out drops the hydrogens the file carries, as it always did.""" + model = Model(verbose=0, strip_H=True) + model.load_pdb(str(pdb_dir / "1AK5_with_H.pdb")) + elements = model.pdb["element"].astype(str).str.strip().values + assert not (elements == "H").any() + + +def _row(model, chain, resseq, name, altloc=""): + pdb = model.pdb + mask = ( + (pdb["chainid"].astype(str) == chain) + & (pdb["resseq"].astype(int) == resseq) + & (pdb["name"].astype(str).str.strip() == name) + & (pdb["altloc"].astype(str).str.strip() == altloc) + ) + (row,) = np.nonzero(mask.values)[0] + return int(row) + + +@pytest.mark.unit +def test_linked_nitrogen_keeps_one_hydrogen(built): + """A peptide bond supplied by a LINK record displaces two of the template's three. + + MSE is a HETATM residue, so its backbone bonds come only from LINK records. MSE65 + has a split side chain: its shared N gets one hydrogen per conformer, and its + carbonyl carbon, bonded to CA(A), CA(B), O and the next N, gets none. + """ + model, _, plan = built("3E98") + n73 = _row(model, "A", 73, "N") + assert plan.name[plan.parent == n73].tolist() == ["H"] + + n65 = _row(model, "A", 65, "N") + on_n65 = plan.parent == n65 + assert plan.name[on_n65].tolist() == ["H", "H"] + assert sorted(plan.altloc[on_n65].tolist()) == ["A", "B"] + assert not (plan.parent == _row(model, "A", 65, "C")).any() + + +@pytest.mark.unit +def test_split_side_chain_keeps_its_alpha_hydrogen(built): + """A CA bonded to two altloc copies of CB is not saturated: one HA per conformer.""" + model, _, plan = built("3E98") + ca_a = _row(model, "A", 65, "CA", "A") + ca_b = _row(model, "A", 65, "CA", "B") + assert plan.name[plan.parent == ca_a].tolist() == ["HA"] + assert plan.name[plan.parent == ca_b].tolist() == ["HA"] + + +# 1U19 chain A: an acetyl cap bonded to MET1 through a LINK record. The ACE template +# is acetaldehyde-like, with a hydrogen on the carbonyl carbon that the peptide link +# displaces; the element valence alone (degree 3 of 4) would still have generated it. +_ACE_MET = """\ +CRYST1 96.680 96.680 150.200 90.00 90.00 90.00 P 41 8 +LINK C ACE A 0 N MET A 1 1555 1555 1.33 +HETATM 1 C ACE A 0 53.553 -7.050 35.606 1.00 47.40 C +HETATM 2 O ACE A 0 52.916 -7.860 34.934 1.00 46.96 O +HETATM 3 CH3 ACE A 0 54.727 -7.523 36.434 1.00 47.42 C +ATOM 4 N MET A 1 53.284 -5.731 35.670 1.00 47.11 N +ATOM 5 CA MET A 1 52.214 -5.077 34.913 1.00 46.26 C +ATOM 6 C MET A 1 52.674 -4.891 33.485 1.00 46.68 C +ATOM 7 O MET A 1 53.849 -4.563 33.283 1.00 46.86 O +ATOM 8 CB MET A 1 51.887 -3.719 35.536 1.00 45.60 C +ATOM 9 CG MET A 1 51.426 -3.792 36.982 1.00 44.48 C +ATOM 10 SD MET A 1 49.945 -4.797 37.183 1.00 46.06 S +ATOM 11 CE MET A 1 48.647 -3.745 36.534 1.00 43.77 C +END +""" + + +@pytest.mark.unit +def test_acetyl_cap_carbon_gets_no_hydrogen(tmp_path): + """The template's own hydrogen count, minus the link, caps the carbonyl carbon.""" + from torchref.topology.monomer.cif import find_cif_file_in_library + + if find_cif_file_in_library("ACE") is None: + pytest.skip("ACE not in the monomer library") + path = tmp_path / "ace_met.pdb" + path.write_text(_ACE_MET) + model = Model(verbose=0, add_hydrogens=False, strip_H=True) + model.load_pdb(str(path)) + model.set_restraints_cif(None) + restraints = model.restraints + plan = plan_hydrogens( + restraints.topology, restraints.cif_dict, model.xyz().detach() + ) + + by_parent = {} + for parent, name in zip(plan.parent.tolist(), plan.name.tolist()): + by_parent.setdefault(parent, []).append(name) + assert _row(model, "A", 0, "C") not in by_parent + assert len(by_parent[_row(model, "A", 0, "CH3")]) == 3 + assert by_parent[_row(model, "A", 1, "N")] == ["H"] diff --git a/tests/unit/topology/test_insertion_codes.py b/tests/unit/topology/test_insertion_codes.py new file mode 100644 index 00000000..0e5aec1b --- /dev/null +++ b/tests/unit/topology/test_insertion_codes.py @@ -0,0 +1,283 @@ +"""Residues distinguished only by an insertion code. + +A deposited structure may number two residues 100 and 100A. They are different residues +with different chemistry, and the only thing separating them is the insertion code. The +topology keys residues on ``(chain, resseq, icode)`` for that reason; the restraint +builders key on ``(chain, resseq)`` alone, which merges them into one residue whose +atom names then collide, so the name-to-index map keeps the first of each and every +restraint belonging to the later residues is silently lost. + +No bundled structure has an insertion code, so the case is synthesised here rather +than shipped as another data file: the rewrite is then visible, and it is obvious that +nothing but the numbering changed. +""" + +import pytest + +from torchref.model.model import Model +from torchref.topology import build_topology + +#: Base structure: chain A, no altlocs, no insertion codes anywhere. +BASE = "3GR5" + +#: The three consecutive residues collapsed onto one sequence number. Their real +#: identities differ (SER, LEU, GLU), so the restraints of the second and third are +#: distinguishable from the first's rather than being duplicates of it. +STRETCH = (23, 24, 25) + + +def _rewrite_with_insertion_codes(source, destination): + """Copy a PDB, renumbering ``STRETCH`` as ``N``, ``NA``, ``NB``. + + Only columns 23-27 change -- the sequence number and the insertion code. Every + atom, coordinate and residue name is untouched, and the residues stay in file order, + so they remain contiguous exactly as a real insertion would be. + + Returns + ------- + tuple of tuple + The ``(resseq, icode)`` pairs written, in order. + """ + first = STRETCH[0] + codes = ["", "A", "B"] + mapping = {old: (first, codes[i]) for i, old in enumerate(STRETCH)} + + out = [] + for line in source.read_text().splitlines(keepends=True): + if line.startswith(("ATOM", "HETATM")): + resseq = int(line[22:26]) + if resseq in mapping: + new_seq, icode = mapping[resseq] + line = f"{line[:22]}{new_seq:>4d}{icode:1s}{line[27:]}" + out.append(line) + destination.write_text("".join(out)) + return tuple((first, code) for code in codes) + + +@pytest.fixture(scope="module") +def inserted(pdb_dir, tmp_path_factory): + """``(topology, restraints, expected keys, model)`` for the rewritten file.""" + path = tmp_path_factory.mktemp("icode") / f"{BASE}_icode.pdb" + expected = _rewrite_with_insertion_codes(pdb_dir / f"{BASE}.pdb", path) + + model = Model(verbose=0, strip_H=True, add_hydrogens=False) + model.load_pdb(str(path)) + model.set_restraints_cif(None) + restraints = model.restraints + + topology = build_topology( + model.pdb, + restraints.cif_dict, + link_dict=restraints.link_dict, + link_list=restraints.link_list, + links=restraints.links, + xyz=model.xyz().detach(), + verbose=0, + ) + return topology, restraints, expected, model + + +def _tuples(indices): + return {tuple(int(v) for v in row) for row in indices.cpu().numpy()} + + +@pytest.mark.unit +def test_the_rewrite_actually_produced_insertion_codes(inserted): + """Guard the fixture: if the rewrite silently failed the rest proves nothing.""" + _, _, _, model = inserted + icodes = model.pdb["icode"].astype(str).str.strip().values + assert set(icodes[icodes != ""]) == {"A", "B"} + + +@pytest.mark.unit +def test_the_graph_keeps_them_apart(inserted): + """Three residue nodes, one per insertion code.""" + topology, _, expected, _ = inserted + residues = topology.residues + + found = [ + residues.key(i) + for i in range(residues.n_residues) + if (int(residues.resseq[i]), str(residues.icode[i]).strip()) + in {(seq, code) for seq, code in expected} + ] + assert len(found) == 3, f"expected three inserted residues, found {found}" + assert len({key[2] for key in found}) == 3, "insertion codes were not distinguished" + + +@pytest.mark.unit +def test_the_builders_merge_them(inserted): + """The comparison only means something if the old grouping really does merge. + + ``PreprocessedPDB`` groups on ``(chain, resseq)``, so the three residues become one + with three sets of backbone atom names. + """ + from torchref.topology.builders import PreprocessedPDB + + _, _, expected, model = inserted + preprocessed = PreprocessedPDB(model.pdb) + + merged = [ + i + for i in range(preprocessed.n_residues) + if int(preprocessed.residue_resseqs[i]) == expected[0][0] + ] + assert len(merged) == 1, "the builders did not merge the inserted residues" + assert preprocessed.has_duplicate_atoms(merged[0]), ( + "the merged residue should carry duplicate atom names, which is what makes the " + "name-to-index map lose the later residues" + ) + + +def _legacy_intra_bonds(model, restraints): + """Intra-residue bonds as the ``(chain, resseq)``-keyed builder produces them. + + ``restraints.py`` now builds from the topology, so it cannot serve as the + baseline -- it *is* the graph. ``BondRestraintBuilder`` is the original path, still + keying residues on ``(chain, resseq)``, which is the behaviour under test. + """ + import torch + + from torchref.topology.builders import BondRestraintBuilder + + built = BondRestraintBuilder(verbose=0).build( + model.pdb, restraints.cif_dict, torch.device("cpu") + ) + if not built: + return set() + return {tuple(int(v) for v in row) for row in built["indices"].cpu().numpy()} + + +def _inserted_residue_indices(topology, expected): + wanted = {(seq, code) for seq, code in expected} + return [ + i + for i in range(topology.n_residues) + if (int(topology.residues.resseq[i]), str(topology.residues.icode[i]).strip()) + in wanted + ] + + +def _bonds_within(edges, topology, residue): + start = int(topology.residues.atom_start[residue]) + end = int(topology.residues.atom_end[residue]) + return {e for e in edges if all(start <= int(a) < end for a in e)} + + +@pytest.mark.unit +def test_the_legacy_grouping_loses_the_later_residues(inserted): + """The merged residue gets restraints for its first component only. + + This is the defect keying on ``(chain, resseq, icode)`` fixes. The three residues + become one, their backbone atom names collide, the name-to-index map keeps the first + of each, and the second and third end up with no intra-residue geometry at all. + """ + topology, restraints, expected, model = inserted + legacy = _legacy_intra_bonds(model, restraints) + residues = _inserted_residue_indices(topology, expected) + assert len(residues) == 3 + + counts = [len(_bonds_within(legacy, topology, r)) for r in residues] + assert counts[0] > 0, "even the first component lost its bonds; check the fixture" + assert counts[1:] == [0, 0], ( + f"the legacy grouping was expected to lose the second and third residues, " + f"but found {counts} bonds in them" + ) + + +@pytest.mark.unit +def test_the_graph_finds_what_the_legacy_grouping_lost(inserted): + """Every inserted residue gets its own bonds, and the graph is a strict superset. + + Localised, not merely larger: the bonds the graph adds all lie inside the inserted + residues, so this is the insertion-code fix rather than a general difference. + """ + topology, restraints, expected, model = inserted + legacy = _legacy_intra_bonds(model, restraints) + graph = topology.atoms.bonds.tuple_set("intra") + residues = _inserted_residue_indices(topology, expected) + + for residue in residues: + assert _bonds_within( + graph, topology, residue + ), f"residue {topology.residues.key(residue)} has no intra-residue bonds" + + gained = graph - legacy + assert gained, "the graph found nothing the legacy grouping missed" + + inserted_set = set(residues) + stray = [ + edge + for edge in gained + if not ({topology.residue_of_atom(a) for a in edge} & inserted_set) + ] + assert not stray, ( + f"{len(stray)} gained bonds lie outside the inserted residues, so the " + f"difference is not localised to the insertion codes: {stray[:3]}" + ) + + +@pytest.mark.unit +def test_the_inserted_residues_get_their_own_intra_restraints(inserted): + """Each of the three carries bonds of its own, not just the first. + + Under the merged grouping only the first residue's template matched, so the second + and third had no intra-residue geometry at all. + """ + topology, _, expected, _ = inserted + + inserted_residues = [ + i + for i in range(topology.n_residues) + if ( + int(topology.residues.resseq[i]), + str(topology.residues.icode[i]).strip(), + ) + in {(seq, code) for seq, code in expected} + ] + + intra = topology.atoms.bonds.origin("intra").cpu().numpy() + for residue in inserted_residues: + start = int(topology.residues.atom_start[residue]) + end = int(topology.residues.atom_end[residue]) + own = [ + row + for row in intra + if start <= int(row[0]) < end and start <= int(row[1]) < end + ] + assert own, ( + f"residue {topology.residues.key(residue)} " + f"({topology.residues.resname[residue]}) has no intra-residue bonds" + ) + + +@pytest.mark.unit +def test_the_inserted_stretch_is_peptide_linked(inserted): + """An insertion-code step is a sequence step, so the chain is not broken. + + ``find_peptide_links`` allows a ``resseq`` difference of 0 precisely for this: 100 + to 100A is consecutive. Without it the inserted residues would float free of the + chain. + """ + topology, _, expected, _ = inserted + + inserted_residues = { + i + for i in range(topology.n_residues) + if ( + int(topology.residues.resseq[i]), + str(topology.residues.icode[i]).strip(), + ) + in {(seq, code) for seq, code in expected} + } + + links = topology.residues.links_of_kind("TRANS") + internal = [ + pair + for pair in links + if int(pair[0]) in inserted_residues and int(pair[1]) in inserted_residues + ] + assert len(internal) == 2, ( + f"expected two peptide links inside the three inserted residues, got " + f"{len(internal)}" + ) diff --git a/tests/unit/topology/test_links.py b/tests/unit/topology/test_links.py new file mode 100644 index 00000000..0a9c2aec --- /dev/null +++ b/tests/unit/topology/test_links.py @@ -0,0 +1,110 @@ +"""Covalent links in the bond graph: counted once, and the same from PDB and mmCIF. + +A repeated LINK record, or a bond emitted once per altloc conformer between two shared +atoms, used to inflate an atom's graph degree. Hydrogen generation reads that degree as +the number of heavy partners, so a Schiff-base nitrogen listed in two identical LINK +records lost its hydrogen and a CA with a split side chain lost its HA. +""" + +import numpy as np +import pytest + +from torchref.model.model import Model + +# 3E98 chain A, LEU72 - MSE73: the MSE is a HETATM residue, so its peptide bonds come +# only from LINK records. +_LEU_MSE = """\ +CRYST1 53.841 88.114 60.963 90.00 107.92 90.00 P 1 21 1 4 +{links} +ATOM 206 N LEU A 72 25.385 1.315 55.882 1.00 51.26 N +ATOM 207 CA LEU A 72 24.279 2.045 56.497 1.00 52.28 C +ATOM 208 C LEU A 72 24.644 2.581 57.879 1.00 51.78 C +ATOM 209 O LEU A 72 24.190 3.667 58.275 1.00 52.62 O +ATOM 210 CB LEU A 72 23.061 1.126 56.613 1.00 52.98 C +ATOM 211 CG LEU A 72 21.688 1.747 56.752 1.00 56.56 C +ATOM 212 CD1 LEU A 72 21.410 2.644 55.536 1.00 55.74 C +ATOM 213 CD2 LEU A 72 20.665 0.619 56.876 1.00 54.03 C +HETATM 214 N MSE A 73 25.451 1.824 58.619 1.00 50.52 N +HETATM 215 CA MSE A 73 25.888 2.255 59.952 1.00 51.25 C +HETATM 216 C MSE A 73 26.955 3.368 59.879 1.00 51.87 C +HETATM 217 O MSE A 73 26.968 4.263 60.731 1.00 52.14 O +HETATM 218 CB MSE A 73 26.424 1.077 60.754 1.00 50.55 C +HETATM 219 CG MSE A 73 25.351 0.118 61.312 1.00 56.44 C +HETATM 220 SE MSE A 73 26.197 -1.503 62.054 0.75 51.93 SE +HETATM 221 CE MSE A 73 27.068 -0.667 63.543 1.00 59.19 C +END +""" +_LINK = "LINK C LEU A 72 N MSE A 73 1555 1555 1.33" + + +def _row(model, chain, resseq, name, altloc=""): + pdb = model.pdb + mask = ( + (pdb["chainid"].astype(str) == chain) + & (pdb["resseq"].astype(int) == resseq) + & (pdb["name"].astype(str).str.strip() == name) + & (pdb["altloc"].astype(str).str.strip() == altloc) + ) + (row,) = np.nonzero(mask.values)[0] + return int(row) + + +def _load(path): + model = Model(verbose=0, add_hydrogens=False, strip_H=True) + model.load_pdb(str(path)) if str(path).endswith(".pdb") else model.load_cif(str(path)) + model.set_restraints_cif(None) + return model + + +@pytest.mark.unit +def test_repeated_link_record_contributes_one_edge(tmp_path): + """The parser keeps both records; the graph carries one bond and one restraint.""" + path = tmp_path / "dup_link.pdb" + path.write_text(_LEU_MSE.format(links=_LINK + "\n" + _LINK)) + model = _load(path) + restraints = model.restraints + atoms = restraints.topology.atoms + + assert len(restraints.links) == 2 + link_rows = atoms.bonds.origin("link") + assert len(link_rows) == 1 + n_mse = _row(model, "A", 73, "N") + assert int(atoms.degree(n_mse)) == 2 # CA and the previous C + + +@pytest.mark.unit +def test_shared_atom_degree_counts_each_partner_once(pdb_dir): + """A bond between two blank-altloc atoms is one edge however many conformers emit it.""" + model = _load(pdb_dir / "3E98.pdb") + atoms = model.restraints.topology.atoms + # MSE65 has split CA/CB/CG/SE/CE; its C is bonded to CA(A), CA(B), O and ARG66 N. + assert int(atoms.degree(_row(model, "A", 65, "C"))) == 4 + + model = _load(pdb_dir / "3A5V.pdb") + atoms = model.restraints.topology.atoms + # CYS53 has a split side chain: N sees C(prev) and CA; CA sees N, C, CB(A), CB(B). + assert int(atoms.degree(_row(model, "A", 53, "N"))) == 2 + assert int(atoms.degree(_row(model, "A", 53, "CA"))) == 4 + + +def _link_identities(model): + atoms = model.restraints.topology.atoms + pdb = model.pdb + key = lambda i: ( + str(pdb["chainid"].iloc[i]), + int(pdb["resseq"].iloc[i]), + str(pdb["icode"].iloc[i]).strip(), + str(pdb["name"].iloc[i]).strip(), + str(pdb["altloc"].iloc[i]).strip(), + ) + return {frozenset((key(int(a)), key(int(b)))) for a, b in atoms.bonds.origin("link")} + + +@pytest.mark.unit +@pytest.mark.parametrize("code", ["1DAW", "2DQ6", "3A5V", "3E98", "5BOV", "6G9X"]) +def test_link_edges_agree_between_pdb_and_cif(pdb_dir, cif_dir, code): + """mmCIF ``_struct_conn`` yields the link edges the PDB LINK records do.""" + from_pdb = _load(pdb_dir / f"{code}.pdb") + from_cif = _load(cif_dir / f"{code}.cif") + assert from_cif.ctx.links is not None and len(from_cif.ctx.links) > 0 + assert _link_identities(from_cif) == _link_identities(from_pdb) diff --git a/tests/unit/topology/test_storage.py b/tests/unit/topology/test_storage.py new file mode 100644 index 00000000..d3423297 --- /dev/null +++ b/tests/unit/topology/test_storage.py @@ -0,0 +1,207 @@ +"""The restraint storage: plain dict, views into the edge blocks, no per-access work. + +The geometry targets read ``restraints[edge_type][origin][property]`` on every +iteration, so the properties tested here are load-bearing rather than cosmetic: the +mapping must be a plain dict of already-materialised tensors, the per-origin entries +must alias the contiguous blocks rather than copy them, and none of that may come +undone on a device move or a copy. +""" + +import pytest +import torch + +from torchref.model.model import Model +from torchref.utils.caching import ParameterFingerprint + +KEYED_TYPES = ("bond", "angle", "torsion") + +#: The surface the geometry targets rely on. Guards the contract from drifting: 'phi' +#: and 'psi' are conformationally free and must NOT acquire a reference value or sigma, +#: and 'omega' must keep the proline flag its own target reads. +EXPECTED_PROPERTIES = { + ("bond", "all"): {"indices", "references", "sigmas"}, + ("bond", "intra"): {"indices", "references", "sigmas"}, + ("angle", "all"): {"indices", "references", "sigmas"}, + ("torsion", "all"): {"indices", "references", "sigmas", "periods"}, + ("torsion", "phi"): {"indices", "periods"}, + ("torsion", "psi"): {"indices", "periods"}, + ("torsion", "omega"): { + "indices", + "references", + "sigmas", + "periods", + "is_proline", + }, +} + + +@pytest.fixture(scope="module") +def restraints(pdb_dir): + """Restraints for a structure with altlocs, disulfides and peptide links.""" + model = Model(verbose=0) + model.load_pdb(str(pdb_dir / "7L84.pdb")) + model.set_restraints_cif(None) + return model.restraints + + +@pytest.mark.unit +def test_restraints_is_a_plain_dict(restraints): + """Not an accessor object that has to be constructed per access.""" + assert type(restraints.restraints) is dict + assert restraints.restraints is restraints.restraints + + +@pytest.mark.unit +def test_access_allocates_nothing(restraints): + """Every level of the lookup returns the same object each time. + + This is what makes the read cheap: three dict lookups and no construction. The + previous storage built a fresh accessor, then a fresh per-type accessor, then a + fresh dict of six string-keyed buffer lookups, on every call. + """ + first = restraints.restraints["bond"]["all"] + second = restraints.restraints["bond"]["all"] + assert first is second + assert first["indices"] is second["indices"] + + +@pytest.mark.unit +@pytest.mark.parametrize("edge_type", KEYED_TYPES) +def test_origin_entries_alias_the_block(restraints, edge_type): + """Per-origin indices are slices of one block, not copies of it.""" + block = restraints.topology.edge_block(edge_type) + entries = restraints.restraints[edge_type] + + for origin, bounds in block.origin_bounds.items(): + indices = entries[origin]["indices"] + assert indices.data_ptr() == block.indices[bounds[0] : bounds[1]].data_ptr() + assert indices.shape[0] == bounds[1] - bounds[0] + + +@pytest.mark.unit +@pytest.mark.parametrize("edge_type", KEYED_TYPES) +def test_all_group_is_a_view(restraints, edge_type): + """The combined group the targets read is a span of the block, not a concatenation. + + The block layout deliberately keeps each type's ``all`` members adjacent so this + holds; ``torsion`` is the one that needs it, since its group is only ``intra`` plus + ``disulfide``. + """ + block = restraints.topology.edge_block(edge_type) + combined = restraints.restraints[edge_type]["all"]["indices"] + assert combined.data_ptr() == block.indices.data_ptr() + + +@pytest.mark.unit +def test_in_place_block_edit_is_visible_through_every_entry(restraints): + """Shared storage means a block edit needs no invalidation to be seen.""" + block = restraints.topology.atoms.bonds + entries = restraints.restraints["bond"] + origin = block.origins()[0] + + saved = int(block.indices[0, 0]) + try: + block.indices[0, 0] = saved + 7 + assert int(entries[origin]["indices"][0, 0]) == saved + 7 + assert int(entries["all"]["indices"][0, 0]) == saved + 7 + finally: + block.indices[0, 0] = saved + + +@pytest.mark.unit +@pytest.mark.parametrize("key", sorted(EXPECTED_PROPERTIES)) +def test_expected_properties_present(restraints, key): + """Each group carries exactly the properties its consumers expect.""" + edge_type, origin = key + group = restraints.restraints[edge_type][origin] + assert set(group) == EXPECTED_PROPERTIES[key] + + +@pytest.mark.unit +def test_cat_dict_is_idempotent(restraints): + """Repeated calls leave the restraint counts alone. + + The previous implementation registered ``'all'`` as an origin when it wrote the + combined group, so a second call concatenated the group into itself and doubled + every bond, angle and torsion -- a silent 2x on the geometry weight. Deriving the + group as a span of the block makes that unrepresentable. + """ + before = { + edge_type: restraints.restraints[edge_type]["all"]["indices"].shape[0] + for edge_type in KEYED_TYPES + } + restraints.cat_dict() + restraints.cat_dict() + after = { + edge_type: restraints.restraints[edge_type]["all"]["indices"].shape[0] + for edge_type in KEYED_TYPES + } + assert before == after + + +@pytest.mark.unit +def test_entries_survive_a_device_apply(restraints): + """A ``.to()`` re-slices the entries instead of leaving them stale or duplicated. + + ``DeviceMixin``'s walk recurses into dicts, so without the ``_apply`` override each + view would be moved on its own and become an independent tensor. + """ + before = restraints.restraints["bond"]["all"]["indices"].clone() + n_before = { + t: restraints.restraints[t]["all"]["indices"].shape[0] for t in KEYED_TYPES + } + + restraints.to(torch.device("cpu")) + + block = restraints.topology.atoms.bonds.indices + after = restraints.restraints["bond"]["all"]["indices"] + assert after.data_ptr() == block.data_ptr(), "entries no longer alias the block" + # ``before`` was cloned prior to the move, so it sits on device.current; the move + # target here is CPU, which is only a no-op when those already agree. + assert torch.equal(after, before.cpu()) + assert { + t: restraints.restraints[t]["all"]["indices"].shape[0] for t in KEYED_TYPES + } == n_before + assert restraints.restraints["vdw"].get("indices") is not None + + +@pytest.mark.unit +def test_blocks_are_untouched_by_a_refinement_step(restraints): + """The edge tensors are constants; nothing in a loss evaluation may mutate them.""" + blocks = [restraints.topology.edge_block(t).indices for t in KEYED_TYPES] + fingerprint = ParameterFingerprint(blocks) + + loss = restraints.nll_bonds().sum() + restraints.nll_angles().sum() + loss.backward() + + assert fingerprint.matches( + [restraints.topology.edge_block(t).indices for t in KEYED_TYPES] + ) + + +@pytest.mark.unit +def test_rebuilding_entries_reslices_onto_the_current_blocks(restraints): + """Re-deriving the entries produces fresh views of the same blocks. + + This is the operation ``_apply`` and ``copy`` both rely on, and the one that has to + stay cheap: it re-slices rather than recomputing anything. + + ``Restraints.copy`` is not exercised here because it cannot run at all -- it is + ``deepcopy``, which walks the *borrowed* ``_xyz_fn`` wrapper, whose cache holds a + graph-attached tensor once ``xyz()`` has been evaluated. Verified to fail + identically at the commit before this change, so it is pre-existing rather than a + regression, and it is reached only through ``Model.copy`` on a model whose lazy + restraints have already been built. + """ + block = restraints.topology.atoms.bonds.indices + before = restraints.restraints["bond"]["all"]["indices"].clone() + + restraints._rebuild_entries() + + after = restraints.restraints["bond"]["all"]["indices"] + assert after.data_ptr() == block.data_ptr() + assert torch.equal(after, before) + for origin, bounds in restraints.topology.atoms.bonds.origin_bounds.items(): + entry = restraints.restraints["bond"][origin]["indices"] + assert entry.shape[0] == bounds[1] - bounds[0] + assert restraints.restraints["vdw"].get("indices") is not None diff --git a/tests/unit/topology/test_subset.py b/tests/unit/topology/test_subset.py new file mode 100644 index 00000000..60b36126 --- /dev/null +++ b/tests/unit/topology/test_subset.py @@ -0,0 +1,305 @@ +"""Subsetting and copying a topology. + +``subset`` reindexes what survives instead of rebuilding, so the properties that matter +are the ones a plausible-but-wrong implementation would break: that every surviving edge +keeps the *same atoms* it had before, that an edge touching a removed atom is gone +entirely rather than left dangling, that the blocks stay canonically ordered without a +re-sort, and that the residue level stays consistent with the atom level. +""" + +import numpy as np +import pytest +import torch + +from torchref.model.model import Model +from torchref.topology import EdgeBlock + +KEYED_TYPES = ("bond", "angle", "torsion", "chiral") + + +@pytest.fixture(scope="module") +def topology(pdb_dir): + """A topology with altlocs, disulfides, peptide links and hydrogens.""" + model = Model(verbose=0, add_hydrogens=False, strip_H=True) + model.load_pdb(str(pdb_dir / "7L84.pdb")) + model.set_restraints_cif(None) + return model.restraints.topology + + +def _named_edges(topology, edge_type): + """Edges as tuples of ``(chain, resseq, icode, atom name)``, identity not index. + + Comparing on identity rather than index is the whole point: an index-based check + passes trivially after a remap, whereas this catches a remap that points an edge at + the wrong atom. + """ + residues = topology.residues + names = topology.atoms.name.astype(str) + out = set() + block = topology.edge_block(edge_type) + for row in block.indices.cpu().numpy(): + out.add( + tuple( + (*residues.key(topology.residue_of_atom(int(a))), names[int(a)]) + for a in row + ) + ) + return out + + +@pytest.mark.unit +def test_subset_of_everything_is_the_same_graph(topology): + """Keeping every atom changes nothing -- the identity case.""" + whole = topology.subset(torch.ones(topology.n_atoms, dtype=torch.bool)) + + assert whole.n_atoms == topology.n_atoms + assert whole.n_residues == topology.n_residues + for edge_type in KEYED_TYPES: + assert torch.equal( + whole.edge_block(edge_type).indices, + topology.edge_block(edge_type).indices, + ) + assert ( + whole.edge_block(edge_type).origin_bounds + == topology.edge_block(edge_type).origin_bounds + ) + + +@pytest.mark.unit +def test_surviving_edges_keep_the_same_atoms(topology): + """Every edge in the subset joins exactly the atoms it joined before. + + Checked on residue-and-name identity, so a remap that silently shifted an index + would fail here even though the shapes still looked right. + """ + residues = topology.residues + chain_a = np.array( + [ + str(residues.chain[topology.residue_of_atom(atom)]) == "A" + for atom in range(topology.n_atoms) + ] + ) + assert chain_a.any() and not chain_a.all() or chain_a.all() + + keep = torch.zeros(topology.n_atoms, dtype=torch.bool) + keep[: topology.n_atoms // 2] = True + reduced = topology.subset(keep) + + for edge_type in KEYED_TYPES: + after = _named_edges(reduced, edge_type) + before = _named_edges(topology, edge_type) + assert after <= before, ( + f"{edge_type}: the subset invented {len(after - before)} edges that were " + f"not in the original" + ) + + +@pytest.mark.unit +def test_an_edge_dies_with_any_of_its_atoms(topology): + """No edge survives that touches a removed atom. + + A bond to an atom that is gone is not a bond, and leaving it would point an index + outside the graph. + """ + keep = torch.ones(topology.n_atoms, dtype=torch.bool) + keep[5] = False + keep[100] = False + reduced = topology.subset(keep) + + for edge_type in KEYED_TYPES: + block = reduced.edge_block(edge_type) + if block.n_edges == 0: + continue + assert int(block.indices.max()) < reduced.n_atoms + assert int(block.indices.min()) >= 0 + + # Precisely: the edges lost are exactly those that used a dropped atom. + dropped = {5, 100} + for edge_type in KEYED_TYPES: + original = topology.edge_block(edge_type).indices.cpu().numpy() + touching = sum(1 for row in original if dropped & set(int(a) for a in row)) + assert ( + reduced.edge_block(edge_type).n_edges + == topology.edge_block(edge_type).n_edges - touching + ), f"{edge_type}: wrong number of edges dropped" + + +@pytest.mark.unit +def test_blocks_stay_canonically_ordered(topology): + """Subsetting needs no re-sort, because the remap is monotone on survivors. + + If this fails the block is no longer canonical, and the ``all`` group stops being a + contiguous span -- which is what makes it a view rather than a copy. + """ + keep = torch.zeros(topology.n_atoms, dtype=torch.bool) + keep[::2] = True + reduced = topology.subset(keep) + + for edge_type in KEYED_TYPES: + block = reduced.edge_block(edge_type) + for origin in block.origins(): + rows = block.origin(origin).cpu().numpy() + if len(rows) < 2: + continue + order = np.lexsort( + tuple(rows[:, c] for c in reversed(range(rows.shape[1]))) + ) + assert ( + order == np.arange(len(rows)) + ).all(), f"{edge_type}/{origin} is no longer lexicographically sorted" + # bounds must remain contiguous and cover the block + spans = sorted(block.origin_bounds.values()) + assert all(a[1] == b[0] for a, b in zip(spans, spans[1:])) + if spans: + assert spans[0][0] == 0 and spans[-1][1] == block.n_edges + + +@pytest.mark.unit +def test_residue_level_stays_consistent_with_the_atoms(topology): + """Atom ranges, residue count and links all agree after subsetting.""" + keep = torch.zeros(topology.n_atoms, dtype=torch.bool) + keep[: topology.n_atoms // 3] = True + reduced = topology.subset(keep) + + residues = reduced.residues + assert residues.n_residues == reduced.n_residues + # Ranges partition the atoms, in order, with no gaps. + assert int(residues.atom_start[0]) == 0 + assert int(residues.atom_end[-1]) == reduced.n_atoms + assert (residues.atom_start[1:] == residues.atom_end[:-1]).all() + + # residue_of agrees with the ranges it is supposed to index. + for residue in range(residues.n_residues): + rows = range(int(residues.atom_start[residue]), int(residues.atom_end[residue])) + for row in rows: + assert reduced.residue_of_atom(row) == residue + + # No link edge points outside the surviving residues. + if len(residues.link_pairs): + assert residues.link_pairs.max() < residues.n_residues + assert residues.link_pairs.min() >= 0 + + +@pytest.mark.unit +def test_dropping_a_whole_residue_drops_its_links(topology): + """A residue with no atoms left is gone, and so is any link that reached it.""" + residues = topology.residues + linked = None + for pair in residues.links_of_kind("TRANS"): + linked = int(pair[0]) + break + assert linked is not None, "7L84 has peptide links" + + keep = torch.ones(topology.n_atoms, dtype=torch.bool) + for row in range(int(residues.atom_start[linked]), int(residues.atom_end[linked])): + keep[row] = False + reduced = topology.subset(keep) + + assert reduced.n_residues == topology.n_residues - 1 + before = len(residues.links_of_kind("TRANS")) + after = len(reduced.residues.links_of_kind("TRANS")) + assert after < before, "a link to the removed residue survived" + + +@pytest.mark.unit +def test_subset_accepts_indices_as_well_as_a_mask(topology): + """Integer indices give the same result as the equivalent mask.""" + indices = torch.arange(0, topology.n_atoms, 3) + mask = torch.zeros(topology.n_atoms, dtype=torch.bool) + mask[indices] = True + + by_index = topology.subset(indices) + by_mask = topology.subset(mask) + + assert by_index.n_atoms == by_mask.n_atoms + for edge_type in KEYED_TYPES: + assert torch.equal( + by_index.edge_block(edge_type).indices, + by_mask.edge_block(edge_type).indices, + ) + + +@pytest.mark.unit +def test_subset_ignores_the_order_of_the_indices(topology): + """A shuffled index list yields the topology's own atom order. + + Deliberate: the edge blocks stay canonical only under a monotone relabelling, so + honouring a caller's order would silently leave them unsorted. + """ + indices = torch.arange(0, topology.n_atoms, 5) + shuffled = indices[torch.randperm(len(indices))] + + assert torch.equal( + topology.subset(indices).atoms.bonds.indices, + topology.subset(shuffled).atoms.bonds.indices, + ) + + +@pytest.mark.unit +def test_subset_of_nothing_is_an_error(topology): + """An empty selection is a mistake, not an empty topology.""" + with pytest.raises(ValueError, match="no atoms"): + topology.subset(torch.zeros(topology.n_atoms, dtype=torch.bool)) + + +@pytest.mark.unit +def test_copy_shares_nothing(topology): + """A copy has equal contents and independent storage.""" + duplicate = topology.copy() + + assert duplicate.n_atoms == topology.n_atoms + assert duplicate.n_residues == topology.n_residues + for edge_type in KEYED_TYPES: + original = topology.edge_block(edge_type) + copied = duplicate.edge_block(edge_type) + assert torch.equal(copied.indices, original.indices) + assert copied.indices.data_ptr() != original.indices.data_ptr() + + saved = int(topology.atoms.bonds.indices[0, 0]) + duplicate.atoms.bonds.indices[0, 0] = saved + 13 + assert int(topology.atoms.bonds.indices[0, 0]) == saved + + duplicate.residues.resname[0] = "XXX" + assert topology.residues.resname[0] != "XXX" + + +@pytest.mark.unit +def test_copy_rebuilds_its_own_adjacency(topology): + """``neighbors`` on a copy reads the copy's bonds, not the original's.""" + duplicate = topology.copy() + atom = int(topology.atoms.bonds.indices[0, 0]) + assert torch.equal(duplicate.neighbors(atom), topology.neighbors(atom)) + assert ( + duplicate.atoms._adj_indices.data_ptr() + != topology.atoms._adj_indices.data_ptr() + ) + + +@pytest.mark.unit +def test_subset_adjacency_matches_its_own_bonds(topology): + """The reduced graph's adjacency is rebuilt, not carried over stale.""" + keep = torch.zeros(topology.n_atoms, dtype=torch.bool) + keep[: topology.n_atoms // 2] = True + reduced = topology.subset(keep) + + from_adjacency = set() + for atom in range(reduced.n_atoms): + for other in reduced.neighbors(atom).cpu().tolist(): + from_adjacency.add((min(atom, other), max(atom, other))) + + from_block = { + (min(int(a), int(b)), max(int(a), int(b))) + for a, b in reduced.atoms.bonds.indices.cpu().numpy() + } + assert from_adjacency == from_block + assert int(reduced.atoms.degree().sum()) == 2 * reduced.atoms.bonds.n_edges + + +@pytest.mark.unit +def test_empty_block_subsets_to_empty(): + """An edge type with no edges survives subsetting without special-casing.""" + block = EdgeBlock.empty(3) + remap = torch.arange(10, dtype=torch.int64) + reduced = block.subset(remap) + assert reduced.n_edges == 0 + assert reduced.arity == 3 diff --git a/tests/unit/utils/test_device_resolution.py b/tests/unit/utils/test_device_resolution.py index 4dab9045..6567108c 100644 --- a/tests/unit/utils/test_device_resolution.py +++ b/tests/unit/utils/test_device_resolution.py @@ -205,9 +205,14 @@ def test_sfds_refuses_a_cell_recast_after_construction(): """ from torchref.model.sf_ds import SfDS + from torchref.model.context import ModelContext + from torchref.symmetry import SpaceGroup + cell = Cell([50.0, 60.0, 70.0, 90.0, 90.0, 90.0], dtype=torch.float32, device="cpu") - sf = SfDS(cell=cell, spacegroup="P 1", dtype_float=torch.float32, - device=torch.device("cpu")) + ctx = ModelContext( + cell=cell, spacegroup=SpaceGroup("P 1", dtype=torch.float32, device="cpu") + ) + sf = SfDS(ctx, dtype_float=torch.float32, device=torch.device("cpu")) xyz = torch.zeros(3, 3, dtype=torch.float32) sf._cartesian_to_fractional(xyz) # consistent: fine diff --git a/tests/unit/utils/test_gradnorm.py b/tests/unit/utils/test_gradnorm.py index d8e1a621..62378f84 100644 --- a/tests/unit/utils/test_gradnorm.py +++ b/tests/unit/utils/test_gradnorm.py @@ -1,75 +1,34 @@ -""" -Unit tests for torchref.utils.gradnorm +"""Pin the RMS gradient norm across one or several parameter tensors.""" -Tests gradient norm calculation utilities. -""" +import math import pytest import torch -import torch.nn as nn +from torchref.config import get_default_device, get_float_dtype +from torchref.utils.gradnorm import gradnorm -class TestGradNorm: - """Tests for gradient norm calculation.""" +pytestmark = pytest.mark.unit - @pytest.mark.unit - def test_gradnorm_basic(self): - """Test basic gradient norm calculation.""" - from torchref.utils.gradnorm import gradnorm - - # Simple linear model - model = nn.Linear(10, 1, bias=False) - x = torch.randn(5, 10) - y = torch.randn(5, 1) - - # Forward pass - pred = model(x) - loss = ((pred - y) ** 2).mean() - - # Calculate gradient norm - grad_norm = gradnorm(loss, model.parameters()) - - assert isinstance(grad_norm, torch.Tensor) - assert grad_norm.ndim == 0 # Scalar - assert grad_norm >= 0 # Non-negative - @pytest.mark.unit - def test_gradnorm_zero_gradient(self): - """Gradient norm should handle zero gradients.""" - from torchref.utils.gradnorm import gradnorm - - model = nn.Linear(10, 1, bias=False) - - # Create a loss that depends on the model but has zero gradient - x = torch.randn(3, 10) - pred = model(x) - loss = (pred * 0.0).sum() # Zero gradient - # DON'T call backward before gradnorm - it calls backward internally - - grad_norm = gradnorm(loss, model.parameters()) - - # Should be 0 (zero gradients) - assert torch.isclose(grad_norm, torch.tensor(0.0, dtype=grad_norm.dtype), atol=1e-10) +@pytest.mark.parametrize("split", [False, True], ids=["single", "multiple"]) +def test_gradnorm_rms(split: bool) -> None: + """The norm weights individual gradient elements, not parameter tensors.""" + values = torch.tensor( + [1.0, 2.0, 3.0], dtype=get_float_dtype(), device=get_default_device() + ) + chunks = (values[:1], values[1:]) if split else (values,) + params = [chunk.clone().requires_grad_() for chunk in chunks] + loss = sum((param.square().sum() for param in params)) + expected = values.new_tensor(math.sqrt(56.0 / 3.0)) + torch.testing.assert_close(gradnorm(loss, iter(params)), expected) - @pytest.mark.unit - def test_gradnorm_multiple_params(self): - """Test gradient norm with multiple parameter groups.""" - from torchref.utils.gradnorm import gradnorm - - # Model with multiple layers - model = nn.Sequential( - nn.Linear(10, 5), - nn.ReLU(), - nn.Linear(5, 1) - ) - - x = torch.randn(3, 10) - y = torch.randn(3, 1) - - pred = model(x) - loss = ((pred - y) ** 2).mean() - - grad_norm = gradnorm(loss, model.parameters()) - - assert isinstance(grad_norm, torch.Tensor) - assert grad_norm >= 0 + +def test_gradnorm_zero_gradient() -> None: + """A connected loss with zero derivative has zero RMS gradient.""" + param = torch.ones( + 3, dtype=get_float_dtype(), device=get_default_device(), requires_grad=True + ) + torch.testing.assert_close( + gradnorm((param * 0).sum(), [param]), param.new_zeros(()) + ) diff --git a/torchref/__init__.py b/torchref/__init__.py index 5f580fd4..9f63eb6e 100644 --- a/torchref/__init__.py +++ b/torchref/__init__.py @@ -40,7 +40,7 @@ General utilities and debugging tools. """ -__version__ = "0.6.4" +__version__ = "0.7.0" import os @@ -92,13 +92,17 @@ # Data I/O from torchref.io import ( DatasetCollection, - ReflectionData, FcalcDataset, - read_mtz, + ReflectionData, + ScaledDataset, read_cif, + read_mtz, read_pdb, ) +# Maps +from torchref.maps import DifferenceMap, Map + # Model from torchref.model import Model, ModelFT from torchref.model.rigid_xyz import RigidXYZTensor @@ -106,20 +110,19 @@ # Refinement from torchref.refinement import LBFGSRefinement, Refinement from torchref.refinement.rigid_body_refinement import RigidBodyRefinementStep -from torchref.symmetry import Cell, SpaceGroup - -# Restraints -# from torchref.restraints import Restraints # Initialized lazily due to monomer library download requirement # Scaling -from torchref.scaling import Scaler, SolventModel, ScalerBase - -# Maps -from torchref.maps import DifferenceMap, Map +from torchref.scaling import Scaler, ScalerBase, SolventModel +from torchref.symmetry import Cell, SpaceGroup, Symmetry # Device movement mixin (public API for extension code) from torchref.utils.device_mixin import DeviceMixin +# Restraints +# torchref.topology.restraints.Restraints is not imported here: constructing it can +# trigger a monomer-library download, so it stays lazy. + + __all__ = [ # Version and paths "__version__", @@ -132,6 +135,7 @@ "sigma_cutoff_ed", # Data I/O "ReflectionData", + "ScaledDataset", "DatasetCollection", "read_mtz", "read_cif", @@ -151,6 +155,7 @@ # Symmetry "Cell", "SpaceGroup", + "Symmetry", # Maps "Map", "DifferenceMap", diff --git a/torchref/base/__init__.py b/torchref/base/__init__.py index bfbfe8c7..e62d4747 100644 --- a/torchref/base/__init__.py +++ b/torchref/base/__init__.py @@ -119,9 +119,6 @@ interpolate_for_rotation, smooth_reciprocal_grid, # Symmetry - compute_symmetry_equivalent_hkls, - compute_translation_phases, - extract_structure_factors_with_symmetry, ReciprocalSymmetryExtractor, ) @@ -283,9 +280,6 @@ "interpolate_structure_factor_from_grid", "interpolate_complex_from_grid", "trilinear_interpolate_patterson", - "compute_symmetry_equivalent_hkls", - "compute_translation_phases", - "extract_structure_factors_with_symmetry", "interpolate_for_rotation", "smooth_reciprocal_grid", # Structure factors diff --git a/torchref/base/coordinates/__init__.py b/torchref/base/coordinates/__init__.py index 3fd18aa4..5e5dd1cf 100644 --- a/torchref/base/coordinates/__init__.py +++ b/torchref/base/coordinates/__init__.py @@ -30,6 +30,13 @@ smallest_diff_aniso, ) +from .local_frame import ( + frame_is_degenerate, + local_frame_axes, + local_frame_coordinates, + place_local_frame, +) + __all__ = [ # PyTorch implementations "cartesian_to_fractional_torch", @@ -45,4 +52,9 @@ # Periodic boundary "smallest_diff", "smallest_diff_aniso", + # Local frames (riding hydrogens) + "local_frame_axes", + "place_local_frame", + "local_frame_coordinates", + "frame_is_degenerate", ] diff --git a/torchref/base/coordinates/local_frame.py b/torchref/base/coordinates/local_frame.py new file mode 100644 index 00000000..a482af3b --- /dev/null +++ b/torchref/base/coordinates/local_frame.py @@ -0,0 +1,207 @@ +"""Local orthonormal frames anchored on three atoms, and points expressed in them. + +A hydrogen that rides on its parent is stored as a constant offset in a frame built +from the parent ``p`` and two reference heavy atoms ``n1``, ``n2``:: + + e1 = unit(n1 - p) + e2 = unit((n2 - p) orthogonal to e1) + e3 = e1 x e2 + h = p + lx*e1 + ly*e2 + lz*e3 + +Every function here is a pure tensor op on already-gathered positions, so autograd +carries a force on ``h`` back onto ``p``, ``n1`` and ``n2`` through the exact frame +Jacobian. Coordinates are Cartesian Angstroms unless the caller chooses otherwise; the +only scale-dependent constant is ``eps``, which floors norms before division. +""" + +from typing import Tuple + +import torch + +#: Norm floor in the same units as the coordinates (Angstroms here). +DEFAULT_EPS = 1e-8 + +#: A frame whose reference bonds are shorter than this, or whose ``n1-p-n2`` angle has +#: a sine below ``MIN_FRAME_SINE``, is treated as degenerate and placed rigidly. +MIN_FRAME_NORM = 1e-3 +MIN_FRAME_SINE = 0.1 + + +def rotate_vectors(vectors: torch.Tensor, rotation: torch.Tensor) -> torch.Tensor: + """Rotate Cartesian vectors by axis-angle rotation vectors. + + Parameters + ---------- + vectors : torch.Tensor + Cartesian vectors, shape ``(..., 3)``, in Å. + rotation : torch.Tensor + Rotation vectors with the same shape as ``vectors``, in radians. + Direction gives the axis and length gives the right-handed angle. + + Returns + ------- + torch.Tensor + Rotated vectors in Å. First and second derivatives are finite at zero. + """ + angle2 = rotation.square().sum(-1, keepdim=True) + # Both torch.where branches must be safe at zero. The series also avoids + # cancellation in (1 - cos(angle)) / angle**2 in single precision. + angle = angle2.clamp_min(1e-4).sqrt() + small = angle2 < 1e-4 + a = torch.where( + small, 1 - angle2 / 6 + angle2.square() / 120, torch.sin(angle) / angle + ) + b = torch.where( + small, + 0.5 - angle2 / 24 + angle2.square() / 720, + 0.5 * torch.sinc(angle / (2 * torch.pi)).square(), + ) + cross = torch.cross(rotation, vectors, dim=-1) + return vectors + a * cross + b * torch.cross(rotation, cross, dim=-1) + + +def local_frame_axes( + p: torch.Tensor, + n1: torch.Tensor, + n2: torch.Tensor, + eps: float = DEFAULT_EPS, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Right-handed orthonormal axes of the frame anchored at ``p``. + + Parameters + ---------- + p, n1, n2 : torch.Tensor + Cartesian positions, shape ``(H, 3)``, same dtype and device. + eps : float + Norm floor guarding a collapsed reference bond. + + Returns + ------- + e1, e2, e3 : torch.Tensor + Unit vectors, each ``(H, 3)``. ``e1`` points along ``n1 - p``, ``e2`` lies in + the ``p, n1, n2`` plane, ``e3 = e1 x e2``. + """ + a = n1 - p + e1 = a / a.norm(dim=-1, keepdim=True).clamp(min=eps) + b = n2 - p + b_perp = b - (b * e1).sum(-1, keepdim=True) * e1 + e2 = b_perp / b_perp.norm(dim=-1, keepdim=True).clamp(min=eps) + e3 = torch.cross(e1, e2, dim=-1) + return e1, e2, e3 + + +def place_local_frame( + p: torch.Tensor, + n1: torch.Tensor, + n2: torch.Tensor, + local_offset: torch.Tensor, + frame_valid: torch.Tensor, + rigid_offset: torch.Tensor, + eps: float = DEFAULT_EPS, +) -> torch.Tensor: + """Positions of points stored as local-frame offsets. + + Parameters + ---------- + p, n1, n2 : torch.Tensor + Frame atoms, each ``(H, 3)``. Rows flagged invalid may hold any in-bounds + position; their frame is not used. + local_offset : torch.Tensor + Coordinates in the ``(e1, e2, e3)`` frame, shape ``(H, 3)``. + frame_valid : torch.Tensor + Boolean ``(H,)``; where False the point is placed rigidly at + ``p + rigid_offset``. + rigid_offset : torch.Tensor + Cartesian ``p -> point`` vector for the rigid fallback, shape ``(H, 3)``. + eps : float + Norm floor guarding degenerate frames. + + Returns + ------- + torch.Tensor + Cartesian positions, shape ``(H, 3)``, differentiable in ``p``, ``n1``, ``n2``. + """ + e1, e2, e3 = local_frame_axes(p, n1, n2, eps) + h_frame = ( + p + + local_offset[:, 0:1] * e1 + + local_offset[:, 1:2] * e2 + + local_offset[:, 2:3] * e3 + ) + h_rigid = p + rigid_offset + return torch.where(frame_valid.unsqueeze(-1), h_frame, h_rigid) + + +def local_frame_coordinates( + p: torch.Tensor, + n1: torch.Tensor, + n2: torch.Tensor, + point: torch.Tensor, + eps: float = DEFAULT_EPS, +) -> torch.Tensor: + """Inverse of :func:`place_local_frame`: express ``point`` in the frame at ``p``. + + Parameters + ---------- + p, n1, n2, point : torch.Tensor + Cartesian positions, each ``(H, 3)``. + eps : float + Norm floor guarding degenerate frames. + + Returns + ------- + torch.Tensor + Local coordinates ``(lx, ly, lz)``, shape ``(H, 3)``, such that + ``place_local_frame(p, n1, n2, result, True, ...)`` returns ``point``. + """ + e1, e2, e3 = local_frame_axes(p, n1, n2, eps) + d = point - p + return torch.stack([(d * e1).sum(-1), (d * e2).sum(-1), (d * e3).sum(-1)], dim=-1) + + +def frame_is_degenerate( + p: torch.Tensor, + n1: torch.Tensor, + n2: torch.Tensor, + min_norm: float = MIN_FRAME_NORM, + min_sine: float = MIN_FRAME_SINE, +) -> torch.Tensor: + """Frames too ill-conditioned to carry an offset. + + A frame is degenerate when either reference bond is shorter than ``min_norm`` or + the two reference bonds are within ``asin(min_sine)`` of collinear, in which case + ``e2`` is set by numerical noise and a riding point would swing with it. + + Parameters + ---------- + p, n1, n2 : torch.Tensor + Cartesian positions, each ``(H, 3)``. + min_norm : float + Shortest acceptable reference bond, same units as the coordinates. + min_sine : float + Smallest acceptable ``|sin(angle(n1 - p, n2 - p))|``. + + Returns + ------- + torch.Tensor + Boolean ``(H,)``, True where the frame must not be used. + """ + a = n1 - p + b = n2 - p + na = a.norm(dim=-1) + nb = b.norm(dim=-1) + cross = torch.cross(a, b, dim=-1).norm(dim=-1) + sine = cross / (na * nb).clamp(min=min_norm * min_norm) + return (na < min_norm) | (nb < min_norm) | (sine < min_sine) + + +__all__ = [ + "DEFAULT_EPS", + "MIN_FRAME_NORM", + "MIN_FRAME_SINE", + "rotate_vectors", + "local_frame_axes", + "place_local_frame", + "local_frame_coordinates", + "frame_is_degenerate", +] diff --git a/torchref/base/direct_summation/_backends.py b/torchref/base/direct_summation/_backends.py index 9b581f32..7011f6bb 100644 --- a/torchref/base/direct_summation/_backends.py +++ b/torchref/base/direct_summation/_backends.py @@ -74,7 +74,7 @@ def _ds_aniso_triton(hkl, s_vec, xyz_frac, occ, U, A, B, max_memory_gb): name="ds_triton", kernel=(_THIS, "_ds_iso_triton", "_ds_aniso_triton"), device="cuda", - dtypes=(torch.float32,), + dtypes=(torch.float32,), # dtype-ok: backend capability declaration, not an allocation # Every argument except ``hkl`` (position 0), whose dtype provably costs # nothing -- see the module docstring. probes=(1, 2, 3, 4, 5, 6), diff --git a/torchref/base/direct_summation/dispatch.py b/torchref/base/direct_summation/dispatch.py index d3005af8..4fe3fd64 100644 --- a/torchref/base/direct_summation/dispatch.py +++ b/torchref/base/direct_summation/dispatch.py @@ -20,6 +20,7 @@ import torch +from torchref.config import get_complex_dtype from torchref.base.direct_summation.isotropic import ( _estimate_batch_size, iso_structure_factor_torched, @@ -241,7 +242,7 @@ def ds_iso(hkl, s, xyz_frac, occ, adp, A, B, *, force_portable=None, max_memory_ reference path regardless. """ if xyz_frac.shape[0] == 0: - return torch.zeros(hkl.shape[0], dtype=torch.complex64, device=hkl.device) + return torch.zeros(hkl.shape[0], dtype=get_complex_dtype(), device=hkl.device) return _dispatch( False, hkl, s, xyz_frac, occ, adp, A, B, force_portable, max_memory_gb ) @@ -253,7 +254,7 @@ def ds_aniso(hkl, s_vec, xyz_frac, occ, U, A, B, *, force_portable=None, max_mem See :func:`ds_iso` on ``force_portable=None``. """ if xyz_frac.shape[0] == 0: - return torch.zeros(hkl.shape[0], dtype=torch.complex64, device=hkl.device) + return torch.zeros(hkl.shape[0], dtype=get_complex_dtype(), device=hkl.device) return _dispatch( True, hkl, s_vec, xyz_frac, occ, U, A, B, force_portable, max_memory_gb ) diff --git a/torchref/base/electron_density/_backends.py b/torchref/base/electron_density/_backends.py index b43ca3aa..ef83ab1c 100644 --- a/torchref/base/electron_density/_backends.py +++ b/torchref/base/electron_density/_backends.py @@ -39,7 +39,7 @@ name="cuda_triton", kernel=(_CUDA, "add_isotropic_cuda_var", "add_anisotropic_cuda_var"), device="cuda", - dtypes=(torch.float32,), + dtypes=(torch.float32,), # dtype-ok: backend capability declaration, not an allocation probes=_ATOM_ARGS, probe=(_CUDA, "why_unavailable"), expect_available="cuda", @@ -53,7 +53,7 @@ name="mps_metal", kernel=(_MPS, "add_isotropic_mps_var", "add_anisotropic_mps_var"), device="mps", - dtypes=(torch.float32,), + dtypes=(torch.float32,), # dtype-ok: backend capability declaration, not an allocation probes=_ATOM_ARGS, probe=("torchref.base.electron_density.kernels.mps.compile", "why_unavailable"), @@ -66,7 +66,7 @@ kernel=(_SPHERE, "add_isotropic_cpu_sphere_var", "add_anisotropic_cpu_sphere_var"), device="cpu", - dtypes=(torch.float32, torch.float64), + dtypes=(torch.float32, torch.float64), # dtype-ok: backend capability declaration, not an allocation # Uniformity, not membership: the kernel picks one ``scalar_t`` from the output # map and then reads every other tensor through a raw pointer of that type, so a # float64 map beside float32 atoms would be a 2x out-of-bounds read. diff --git a/torchref/base/electron_density/kernels/cpu/jit_reference.py b/torchref/base/electron_density/kernels/cpu/jit_reference.py index 5394ee5d..257fa6d0 100644 --- a/torchref/base/electron_density/kernels/cpu/jit_reference.py +++ b/torchref/base/electron_density/kernels/cpu/jit_reference.py @@ -136,9 +136,9 @@ def forward( ny: int = density_map.shape[1] nz: int = density_map.shape[2] strides = torch.tensor( - [ny * nz, nz, 1], device=voxel_indices.device, dtype=torch.long + [ny * nz, nz, 1], device=voxel_indices.device, dtype=torch.long # dtype-ok: CPU-kernel strides for flat voxel index arithmetic; indexing requires long ) - index_flat = torch.sum(voxel_indices.to(torch.long) * strides, dim=-1).view(-1) + index_flat = torch.sum(voxel_indices.to(torch.long) * strides, dim=-1).view(-1) # dtype-ok: voxel indices flattened for scatter; indexing requires long density_map.view(-1).scatter_add_(0, index_flat, density.reshape(-1)) return density_map @@ -231,9 +231,9 @@ def forward( ny: int = density_map.shape[1] nz: int = density_map.shape[2] index_flat = ( - voxel_indices[:, :, 0].to(torch.int64) * (ny * nz) - + voxel_indices[:, :, 1].to(torch.int64) * nz - + voxel_indices[:, :, 2].to(torch.int64) + voxel_indices[:, :, 0].to(torch.int64) * (ny * nz) # dtype-ok: voxel-index flat-arithmetic term for scatter; requires int64 + + voxel_indices[:, :, 1].to(torch.int64) * nz # dtype-ok: voxel-index flat-arithmetic term for scatter; requires int64 + + voxel_indices[:, :, 2].to(torch.int64) # dtype-ok: voxel-index flat-arithmetic term for scatter; requires int64 ).flatten() density_map.view(-1).scatter_add_(0, index_flat, density.flatten()) @@ -307,9 +307,9 @@ def _add_to_map_gpu_simple( ny, nz = density_map.shape[1], density_map.shape[2] index_flat = ( - voxel_indices[:, :, 0].to(torch.int64) * (ny * nz) - + voxel_indices[:, :, 1].to(torch.int64) * nz - + voxel_indices[:, :, 2].to(torch.int64) + voxel_indices[:, :, 0].to(torch.int64) * (ny * nz) # dtype-ok: voxel-index flat-arithmetic term for scatter; requires int64 + + voxel_indices[:, :, 1].to(torch.int64) * nz # dtype-ok: voxel-index flat-arithmetic term for scatter; requires int64 + + voxel_indices[:, :, 2].to(torch.int64) # dtype-ok: voxel-index flat-arithmetic term for scatter; requires int64 ).flatten() density_map.view(-1).scatter_add_(0, index_flat, density.flatten()) diff --git a/torchref/base/electron_density/kernels/cpu/sphere_splat.py b/torchref/base/electron_density/kernels/cpu/sphere_splat.py index 9310769f..8d7dde39 100644 --- a/torchref/base/electron_density/kernels/cpu/sphere_splat.py +++ b/torchref/base/electron_density/kernels/cpu/sphere_splat.py @@ -608,7 +608,7 @@ def _prep(density_map, xyz, radius_per_atom, *tensors): f"sphere_splat is a CPU kernel; got device {density_map.device}" ) dtype = density_map.dtype - if dtype not in (torch.float32, torch.float64): + if dtype not in (torch.float32, torch.float64): # dtype-ok: validation guard, not an allocation raise ValueError(f"sphere_splat supports float32/float64, got {dtype}") for t in (xyz, radius_per_atom) + tensors: if t.dtype != dtype: diff --git a/torchref/base/electron_density/kernels/cpu/variable_radius.py b/torchref/base/electron_density/kernels/cpu/variable_radius.py index ba3393c0..bb19ae93 100644 --- a/torchref/base/electron_density/kernels/cpu/variable_radius.py +++ b/torchref/base/electron_density/kernels/cpu/variable_radius.py @@ -52,7 +52,7 @@ def _bucket_by_radius(radius: torch.Tensor, center_1d: torch.Tensor): spans.append((float(r), cursor, cursor + idx.numel())) cursor += idx.numel() order = (torch.cat(order_parts) if order_parts - else torch.zeros(0, dtype=torch.long, device=radius.device)) + else torch.zeros(0, dtype=torch.long, device=radius.device)) # dtype-ok: empty voxel-index fallback; must stay long for indexing return order, spans @@ -88,7 +88,7 @@ def _canonical_setup(xyz, inv_frac, frac, grid_dims, radius_per_atom, dtype): nx, ny, nz = grid_dims grid_f = torch.tensor(grid_dims, device=device, dtype=dtype) xyz_frac = (xyz @ inv_frac.T) % 1.0 - center_idx = torch.round(xyz_frac * grid_f).to(torch.long) + center_idx = torch.round(xyz_frac * grid_f).to(torch.long) # dtype-ok: rounded voxel center indices; torch indexing requires long # w0: atom position relative to its anchor node, in Cartesian. This is what # centres the sphere on the atom rather than on the node. w0 = (xyz_frac - center_idx.to(dtype) / grid_f) @ frac.T @@ -111,8 +111,8 @@ def add_isotropic_plain_var(density_map, xyz, adp, occ, A, B, device, dtype = xyz.device, density_map.dtype nx, ny, nz = (int(s) for s in density_map.shape) grid_dims = (nx, ny, nz) - strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long) - grid_shape = torch.tensor(grid_dims, device=device, dtype=torch.long) + strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long) # dtype-ok: strides for flat voxel-index arithmetic; indexing requires long + grid_shape = torch.tensor(grid_dims, device=device, dtype=torch.long) # dtype-ok: grid_shape for flat voxel-index arithmetic; indexing requires long order, spans, center_idx, w0 = _canonical_setup( xyz, inv_frac_matrix, frac_matrix, grid_dims, radius_per_atom, dtype) @@ -151,8 +151,8 @@ def add_anisotropic_plain_var(density_map, xyz, u, occ, A, B, device, dtype = xyz.device, density_map.dtype nx, ny, nz = (int(s) for s in density_map.shape) grid_dims = (nx, ny, nz) - strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long) - grid_shape = torch.tensor(grid_dims, device=device, dtype=torch.long) + strides = torch.tensor([ny * nz, nz, 1], device=device, dtype=torch.long) # dtype-ok: strides for flat voxel-index arithmetic; indexing requires long + grid_shape = torch.tensor(grid_dims, device=device, dtype=torch.long) # dtype-ok: grid_shape for flat voxel-index arithmetic; indexing requires long order, spans, center_idx, w0 = _canonical_setup( xyz, inv_frac_matrix, frac_matrix, grid_dims, radius_per_atom, dtype) diff --git a/torchref/base/electron_density/kernels/cuda/variable_radius.py b/torchref/base/electron_density/kernels/cuda/variable_radius.py index a787afa0..2784916b 100644 --- a/torchref/base/electron_density/kernels/cuda/variable_radius.py +++ b/torchref/base/electron_density/kernels/cuda/variable_radius.py @@ -58,7 +58,7 @@ def _sym3_inv(a, b, c, d, e, f): @triton.jit def _wq_grid_fwd_kernel( n_items, - grid_ptr, density_map_ptr, + density_map_ptr, xyz_ptr, b_ptr, A_ptr, B_ptr, occ_ptr, r2cut_ptr, mask_ptr, inv_frac_ptr, frac_ptr, @@ -177,7 +177,7 @@ def _wq_grid_fwd_kernel( @triton.jit def _wq_grid_bwd_kernel( n_items, - grid_ptr, grad_density_map_ptr, + grad_density_map_ptr, xyz_ptr, b_ptr, A_ptr, B_ptr, occ_ptr, r2cut_ptr, mask_ptr, inv_frac_ptr, frac_ptr, @@ -327,7 +327,7 @@ def _wq_grid_bwd_kernel( @triton.jit def _wq_grid_aniso_fwd_kernel( n_items, - grid_ptr, density_map_ptr, + density_map_ptr, xyz_ptr, u_ptr, A_ptr, B_ptr, occ_ptr, r2cut_ptr, mask_ptr, inv_frac_ptr, frac_ptr, @@ -461,7 +461,7 @@ def _wq_grid_aniso_fwd_kernel( @triton.jit def _wq_grid_aniso_bwd_kernel( n_items, - grid_ptr, grad_density_map_ptr, + grad_density_map_ptr, xyz_ptr, u_ptr, A_ptr, B_ptr, occ_ptr, r2cut_ptr, mask_ptr, inv_frac_ptr, frac_ptr, @@ -666,12 +666,12 @@ def _wq_grid_aniso_bwd_kernel( def _launch_grid_fwd(out_flat, r2cut, mask, scene_buffers, dims): """Isotropic grid=(n_atoms,) forward (fixed FWD_BLOCK_V/FWD_NUM_WARPS).""" - (grid_flat, xyz, b, A, B, occ, inv_frac, frac) = scene_buffers + (xyz, b, A, B, occ, inv_frac, frac) = scene_buffers nx, ny, nz = dims n_atoms = r2cut.shape[0] _wq_grid_fwd_kernel[(n_atoms,)]( n_atoms, - grid_flat, out_flat, + out_flat, xyz, b, A, B, occ, r2cut, mask, inv_frac, frac, @@ -682,12 +682,12 @@ def _launch_grid_fwd(out_flat, r2cut, mask, scene_buffers, dims): def _launch_grid_aniso_fwd(out_flat, r2cut, mask, scene_buffers, dims): """Anisotropic grid=(n_atoms,) forward (fixed FWD_BLOCK_V/FWD_NUM_WARPS).""" - (grid_flat, xyz, u, A, B, occ, inv_frac, frac) = scene_buffers + (xyz, u, A, B, occ, inv_frac, frac) = scene_buffers nx, ny, nz = dims n_atoms = r2cut.shape[0] _wq_grid_aniso_fwd_kernel[(n_atoms,)]( n_atoms, - grid_flat, out_flat, + out_flat, xyz, u, A, B, occ, r2cut, mask, inv_frac, frac, @@ -705,14 +705,13 @@ class WorkQueueGridDensity(torch.autograd.Function): """ @staticmethod - def forward(ctx, density_map, real_space_grid, xyz, b, occ, A, B, + def forward(ctx, density_map, xyz, b, occ, A, B, r2cut, mask, inv_frac, frac): # Accumulate the splat into a copy of the running density_map (out = # density_map + splat) so the dispatch needs no separate zeros buffer + add. # A clone (not in-place) keeps this autograd-trivial AND safe for the AUTO # fallthrough: density_map is untouched if the kernel raises. - nx, ny, nz = real_space_grid.shape[:3] - grid_flat = real_space_grid.contiguous().view(-1) + nx, ny, nz = density_map.shape[:3] xyz = xyz.contiguous(); b = b.contiguous(); occ = occ.contiguous() A = A.contiguous(); B = B.contiguous() inv_frac_flat = inv_frac.contiguous().view(-1) @@ -720,19 +719,19 @@ def forward(ctx, density_map, real_space_grid, xyz, b, occ, A, B, out = density_map.contiguous().clone().view(-1) _launch_grid_fwd( out, r2cut, mask, - (grid_flat, xyz, b, A, B, occ, inv_frac_flat, frac_flat), + (xyz, b, A, B, occ, inv_frac_flat, frac_flat), (nx, ny, nz), ) - ctx.save_for_backward(real_space_grid, xyz, b, occ, A, B, + ctx.dims = (nx, ny, nz) + ctx.save_for_backward(xyz, b, occ, A, B, r2cut, mask, inv_frac, frac) return out.view(nx, ny, nz) @staticmethod def backward(ctx, grad_density_map): - (real_space_grid, xyz, b, occ, A, B, + (xyz, b, occ, A, B, r2cut, mask, inv_frac, frac) = ctx.saved_tensors - nx, ny, nz = real_space_grid.shape[:3] - grid_flat = real_space_grid.contiguous().view(-1) + nx, ny, nz = ctx.dims grad_dm = grad_density_map.contiguous().view(-1) inv_frac_flat = inv_frac.contiguous().view(-1) frac_flat = frac.contiguous().view(-1) @@ -741,7 +740,7 @@ def backward(ctx, grad_density_map): grad_occ = torch.zeros_like(occ) _wq_grid_bwd_kernel[(r2cut.shape[0],)]( r2cut.shape[0], - grid_flat, grad_dm, + grad_dm, xyz.contiguous(), b.contiguous(), A.contiguous(), B.contiguous(), occ.contiguous(), r2cut, mask, inv_frac_flat, frac_flat, @@ -750,8 +749,8 @@ def backward(ctx, grad_density_map): num_warps=BWD_NUM_WARPS, ) # out = density_map + splat -> grad wrt density_map is identity. - # grads for: density_map, real_space_grid, xyz, b, occ, A, B, r2cut, mask, inv_frac, frac - return (grad_density_map, None, grad_xyz, grad_b, grad_occ, None, None, + # grads for: density_map, xyz, b, occ, A, B, r2cut, mask, inv_frac, frac + return (grad_density_map, grad_xyz, grad_b, grad_occ, None, None, None, None, None, None) @@ -762,11 +761,10 @@ class WorkQueueGridDensityAniso(torch.autograd.Function): ``grad_u`` (6 components) in place of ``grad_b``.""" @staticmethod - def forward(ctx, density_map, real_space_grid, xyz, u, occ, A, B, + def forward(ctx, density_map, xyz, u, occ, A, B, r2cut, mask, inv_frac, frac): # Accumulate into a copy of the running density_map (see the iso forward). - nx, ny, nz = real_space_grid.shape[:3] - grid_flat = real_space_grid.contiguous().view(-1) + nx, ny, nz = density_map.shape[:3] xyz = xyz.contiguous(); u = u.contiguous(); occ = occ.contiguous() A = A.contiguous(); B = B.contiguous() inv_frac_flat = inv_frac.contiguous().view(-1) @@ -774,19 +772,19 @@ def forward(ctx, density_map, real_space_grid, xyz, u, occ, A, B, out = density_map.contiguous().clone().view(-1) _launch_grid_aniso_fwd( out, r2cut, mask, - (grid_flat, xyz, u, A, B, occ, inv_frac_flat, frac_flat), + (xyz, u, A, B, occ, inv_frac_flat, frac_flat), (nx, ny, nz), ) - ctx.save_for_backward(real_space_grid, xyz, u, occ, A, B, + ctx.dims = (nx, ny, nz) + ctx.save_for_backward(xyz, u, occ, A, B, r2cut, mask, inv_frac, frac) return out.view(nx, ny, nz) @staticmethod def backward(ctx, grad_density_map): - (real_space_grid, xyz, u, occ, A, B, + (xyz, u, occ, A, B, r2cut, mask, inv_frac, frac) = ctx.saved_tensors - nx, ny, nz = real_space_grid.shape[:3] - grid_flat = real_space_grid.contiguous().view(-1) + nx, ny, nz = ctx.dims grad_dm = grad_density_map.contiguous().view(-1) inv_frac_flat = inv_frac.contiguous().view(-1) frac_flat = frac.contiguous().view(-1) @@ -795,7 +793,7 @@ def backward(ctx, grad_density_map): grad_occ = torch.zeros_like(occ) _wq_grid_aniso_bwd_kernel[(r2cut.shape[0],)]( r2cut.shape[0], - grid_flat, grad_dm, + grad_dm, xyz.contiguous(), u.contiguous(), A.contiguous(), B.contiguous(), occ.contiguous(), r2cut, mask, inv_frac_flat, frac_flat, @@ -804,7 +802,7 @@ def backward(ctx, grad_density_map): num_warps=BWD_NUM_WARPS, ) # out = density_map + splat -> grad wrt density_map is identity. - return (grad_density_map, None, grad_xyz, grad_u, grad_occ, None, None, + return (grad_density_map, grad_xyz, grad_u, grad_occ, None, None, None, None, None, None) @@ -821,15 +819,9 @@ def backward(ctx, grad_density_map): # wrappers do for CUDA what ``add_*_mps_var`` already did for Metal: square the # radius and build the coefficient mask, rather than leaving that to the caller. # -# ``density_map`` is passed where the ``autograd.Function`` wants -# ``real_space_grid``. That is exact, not a convenience: ``forward``/``backward`` -# use that argument only for ``.shape[:3]`` and for a ``grid_flat`` pointer handed -# to the kernel as ``grid_ptr`` -- which none of the four Triton kernels ever -# loads, because voxel coordinates are derived arithmetically from ``frac`` and the -# grid dims. ``density_map`` has the same ``(nx, ny, nz)`` shape, so both uses are -# satisfied. If a kernel is ever changed to actually read ``grid_ptr``, this breaks -# and the right fix is to delete that dead parameter, not to thread a real grid -# back through here. +# No coordinate grid is threaded anywhere: every voxel's Cartesian position is +# derived arithmetically in-kernel from ``frac`` and the grid dims, so ``density_map`` +# alone carries the ``(nx, ny, nz)`` shape the launch needs. def why_unavailable(): @@ -869,7 +861,6 @@ def add_isotropic_cuda_var( """ return WorkQueueGridDensity.apply( density_map, - density_map, # stands in for real_space_grid; see the note above xyz, adp, occ, @@ -893,7 +884,6 @@ def add_anisotropic_cuda_var( """ return WorkQueueGridDensityAniso.apply( density_map, - density_map, # stands in for real_space_grid; see the note above xyz, u, occ, diff --git a/torchref/base/electron_density/main.py b/torchref/base/electron_density/main.py index e619ffe9..7a80492d 100644 --- a/torchref/base/electron_density/main.py +++ b/torchref/base/electron_density/main.py @@ -37,7 +37,8 @@ def build_electron_density( - real_space_grid: torch.Tensor, + grid_shape, + device: torch.device, xyz_iso: torch.Tensor, adp_iso: torch.Tensor, occ_iso: torch.Tensor, @@ -45,7 +46,6 @@ def build_electron_density( B_iso: torch.Tensor, inv_frac_matrix: torch.Tensor, frac_matrix: torch.Tensor, - voxel_size: torch.Tensor, xyz_aniso: Optional[torch.Tensor] = None, u_aniso: Optional[torch.Tensor] = None, occ_aniso: Optional[torch.Tensor] = None, @@ -61,18 +61,18 @@ def build_electron_density( Parameters ---------- - real_space_grid : torch.Tensor - Coordinate grid, shape ``(nx, ny, nz, 3)``. + grid_shape : tuple of int + Map dimensions ``(nx, ny, nz)``. No coordinate grid is needed or built: every + splat derives a voxel's Cartesian position arithmetically from its index and + ``inv_frac_matrix``. + device : torch.device + Device to allocate the map on. xyz_iso, adp_iso, occ_iso : torch.Tensor Isotropic positions ``(n_iso, 3)``, B-factors and occupancies ``(n_iso,)``. A_iso, B_iso : torch.Tensor ITC92 coefficients, shape ``(n_iso, 5)``. inv_frac_matrix, frac_matrix : torch.Tensor Cartesian-to-fractional and fractional-to-Cartesian, shape ``(3, 3)``. - voxel_size : torch.Tensor - **Unused by every splat** -- the truncation radius comes from each atom's B/U and - ``torchref.sigma_cutoff_ed``, and the enumeration box from ``inv_frac_matrix``. - Retained because ``SfFFT`` passes it positionally. xyz_aniso, u_aniso, occ_aniso : torch.Tensor, optional Anisotropic positions ``(n_aniso, 3)``, U ``(n_aniso, 6)``, occupancies. A_aniso, B_aniso : torch.Tensor, optional @@ -86,12 +86,7 @@ def build_electron_density( """ if dtype is None: dtype = get_float_dtype() - device = real_space_grid.device - density_map = torch.zeros( - real_space_grid.shape[:-1], - dtype=dtype, - device=device, - ) + density_map = torch.zeros(tuple(grid_shape), dtype=dtype, device=device) # --- isotropic atoms --- if len(xyz_iso) > 0: diff --git a/torchref/base/electron_density/map_building.py b/torchref/base/electron_density/map_building.py index 1d9b90fd..d00d7664 100644 --- a/torchref/base/electron_density/map_building.py +++ b/torchref/base/electron_density/map_building.py @@ -28,11 +28,11 @@ def scatter_add_nd(source, index, map): """Vectorized n-dimensional scatter-add: ``source`` ``(N,)`` into ``map`` ``(d1..dn)`` at ``index`` ``(N, ndim)``, returning the modified map. """ - map_shape = torch.tensor(map.shape, device=index.device, dtype=torch.int64) + map_shape = torch.tensor(map.shape, device=index.device, dtype=torch.int64) # dtype-ok: map_shape for stride/flat-index arithmetic feeding scatter_add; requires int64 # Convert n-dimensional indices to flat indices # For shape (d1, d2, d3, ..., dn), flat_index = i0 * (d1*d2*...*dn) + i1 * (d2*d3*...*dn) + ... + in - strides = torch.ones(len(map_shape), device=index.device, dtype=torch.int64) + strides = torch.ones(len(map_shape), device=index.device, dtype=torch.int64) # dtype-ok: strides for flat scatter_add index; requires int64 for i in range(len(map_shape) - 2, -1, -1): strides[i] = strides[i + 1] * map_shape[i + 1] diff --git a/torchref/base/electron_density/solvent_mask.py b/torchref/base/electron_density/solvent_mask.py index 08f1ceb7..2e36e709 100644 --- a/torchref/base/electron_density/solvent_mask.py +++ b/torchref/base/electron_density/solvent_mask.py @@ -113,7 +113,7 @@ def add_to_phenix_mask( ) # (N_atoms, N_voxels) # Flatten for scatter operations - voxel_indices_flat = voxel_indices.reshape(-1, 3).to(torch.long) + voxel_indices_flat = voxel_indices.reshape(-1, 3).to(torch.long) # dtype-ok: voxel indices for grid indexing; requires long # Create protein core mask using scatter_add int_dtype = dtypes.int diff --git a/torchref/base/electron_density/voxel_utils.py b/torchref/base/electron_density/voxel_utils.py index 88c2d037..3fc903d7 100644 --- a/torchref/base/electron_density/voxel_utils.py +++ b/torchref/base/electron_density/voxel_utils.py @@ -49,13 +49,13 @@ def find_relevant_voxels(real_space_grid, xyz, radius_angstrom=4, inv_frac_matri # This ensures atoms outside the unit cell are correctly wrapped xyz_frac = torch.matmul(inv_frac_matrix, xyz.T).T # (N, 3) xyz_frac = xyz_frac % 1.0 # Wrap to [0, 1] - center_idx = torch.round(xyz_frac * grid_shape.unsqueeze(0)).to(torch.int64) + center_idx = torch.round(xyz_frac * grid_shape.unsqueeze(0)).to(torch.int64) # dtype-ok: rounded voxel center indices; torch indexing requires int64 else: # Fallback for orthogonal cells (less accurate for non-orthogonal) voxelsize = real_space_grid[3, 3, 3] - real_space_grid[2, 2, 2] center_idx = torch.round( (xyz - grid_origin.unsqueeze(0)) / voxelsize.unsqueeze(0) - ).to(torch.int64) + ).to(torch.int64) # dtype-ok: voxel index cast; torch indexing requires int64 voxel_indices_wrapped = excise_angstrom_radius_around_coord( real_space_grid, center_idx, radius_angstrom diff --git a/torchref/base/french_wilson.py b/torchref/base/french_wilson.py index 60602f78..a03c6df2 100644 --- a/torchref/base/french_wilson.py +++ b/torchref/base/french_wilson.py @@ -116,7 +116,7 @@ 2.906, 3.004, ], - dtype=torch.float32, + dtype=get_float_dtype(), ) AC_ZJ_SD = torch.tensor( @@ -193,7 +193,7 @@ 0.994, 0.996, ], - dtype=torch.float32, + dtype=get_float_dtype(), ) AC_ZF = torch.tensor( @@ -270,7 +270,7 @@ 1.676, 1.706, ], - dtype=torch.float32, + dtype=get_float_dtype(), ) AC_ZF_SD = torch.tensor( @@ -347,7 +347,7 @@ 0.310, 0.304, ], - dtype=torch.float32, + dtype=get_float_dtype(), ) # Centric lookup tables from French-Wilson supplement (1978) @@ -435,7 +435,7 @@ 3.753, 3.962, ], - dtype=torch.float32, + dtype=get_float_dtype(), ) C_ZJ_SD = torch.tensor( @@ -522,7 +522,7 @@ 1.029, 1.028, ], - dtype=torch.float32, + dtype=get_float_dtype(), ) C_ZF = torch.tensor( @@ -609,7 +609,7 @@ 1.917, 1.945, ], - dtype=torch.float32, + dtype=get_float_dtype(), ) C_ZF_SD = torch.tensor( @@ -696,7 +696,7 @@ 0.278, 0.272, ], - dtype=torch.float32, + dtype=get_float_dtype(), ) @@ -1146,135 +1146,6 @@ def french_wilson( return F, sigma_F, valid_mask -def is_centric_from_hkl( - hkl: torch.Tensor, space_group: SpaceGroupLike = "P1" -) -> torch.Tensor: - """ - Determine if reflections are centric based on Miller indices and space group. - - Uses symmetry operations to check if reflections are invariant under - inversion through the origin (Friedel mates). A reflection is centric - if -h,-k,-l is symmetry equivalent to h,k,l. - - Parameters - ---------- - hkl : torch.Tensor - Miller indices of shape (..., 3). - space_group : str, int, or gemmi.SpaceGroup, optional - Space group specification. Default is "P1". - - Returns - ------- - torch.Tensor - Boolean mask of shape (...), True for centric reflections. - """ - original_shape = hkl.shape[:-1] - hkl_flat = hkl.reshape(-1, 3) - n_reflections = hkl_flat.shape[0] - - # Get symmetry operations from the SpaceGroup class - float_dtype = get_float_dtype() - spacegroup = SpaceGroup(space_group, dtype=float_dtype, device=hkl.device) - - # Convert HKL to the configured float dtype for symmetry operations - hkl_float = hkl_flat.to(float_dtype) # Shape: (n_reflections, 3) - - # Apply all spacegroup operations to all reflections at once - # For reciprocal space (Miller indices), only rotation applies, not translation - # hkl_float shape: (n_reflections, 3) - # spacegroup.apply_to_hkl returns shape: (n_reflections, 3, n_ops) - hkl_sym = spacegroup.apply_to_hkl(hkl_float) - - # Compute Friedel mates: -h, -k, -l - # Shape: (n_reflections, 3, 1) to broadcast against (n_reflections, 3, n_ops) - friedel_hkl = -hkl_float.unsqueeze(-1) # Shape: (n_reflections, 3, 1) - - # Check if any spacegroup operation produces the Friedel mate - # Round to nearest integer (Miller indices should be integers) - hkl_sym_rounded = torch.round(hkl_sym) - - # Compute difference for all reflections and all spacegroup operations - # Shape: (n_reflections, 3, n_ops) - diff = torch.abs(hkl_sym_rounded - friedel_hkl) - - # A reflection is centric if ANY spacegroup operation maps it to its Friedel mate - # Check if all 3 components (h,k,l) match (diff < 0.5) for any operation - # Shape: (n_reflections, n_ops) after checking all 3 components match - matches = torch.all(diff < 0.5, dim=1) # Check all 3 Miller indices match - - # A reflection is centric if it matches for ANY spacegroup operation - # Shape: (n_reflections,) - is_centric = torch.any(matches, dim=1) - - return is_centric.reshape(original_shape) - - -def epsilon_from_hkl(hkl: torch.Tensor, spacegroup) -> torch.Tensor: - """Per-reflection epsilon: number of rotation symops mapping h -> +/-h. - - Mirrors ``ReciprocalSymmetry.get_epsilon`` (Friedel-aware) but works directly - on the scattered HKL list. Returns ones if ``spacegroup`` is None or lacks - ``apply_to_hkl``. - - Unlike :func:`is_centric_from_hkl` this takes a constructed space group rather - than a specification, because its callers already hold one. - - Always returns on ``hkl.device``, whatever device the space group's symmetry - matrices live on: the caller multiplies this against per-reflection data - sitting beside ``hkl``. - """ - n = hkl.shape[0] - float_dtype = get_float_dtype() - if spacegroup is None or not hasattr(spacegroup, "apply_to_hkl"): - return torch.ones(n, device=hkl.device, dtype=float_dtype) - - with torch.no_grad(): - # Configured float dtype, not float64: MPS has no float64 and casting - # there raises. Symmetry arithmetic on Miller indices is exact in - # float32 (integer-valued rotation matrices, small indices), so the - # exact `==` comparisons below remain valid. - # - # ``apply_to_hkl`` moves its input onto the matrices' device, so build - # ``h`` there too -- otherwise ``Hs`` and ``h0`` land on different - # devices and the comparisons below raise. The space group wins for the - # arithmetic; the result is handed back on the caller's device. - sym_device = getattr(spacegroup, "matrices", hkl).device - h = hkl.to(device=sym_device, dtype=float_dtype) - Hs = spacegroup.apply_to_hkl(h) # (N,3,ops) - h0 = h.unsqueeze(-1) # (N,3,1) - same = (Hs == h0).all(dim=1) - friedel = (Hs == -h0).all(dim=1) - eps = (same | friedel).sum(dim=1).clamp(min=1).to(float_dtype) - return eps.to(hkl.device) - - -def get_centric_acentric_masks( - hkl: torch.Tensor, space_group: SpaceGroupLike = "P1" -) -> tuple[torch.Tensor, torch.Tensor]: - """ - Get both centric and acentric masks for reflections. - - Convenience function that returns both masks explicitly. - - Parameters - ---------- - hkl : torch.Tensor - Miller indices of shape (..., 3). - space_group : str, int, or gemmi.SpaceGroup, optional - Space group specification. Default is "P1". - - Returns - ------- - centric_mask : torch.Tensor - Boolean mask of shape (...), True for centric reflections. - acentric_mask : torch.Tensor - Boolean mask of shape (...), True for acentric reflections. - """ - centric_mask = is_centric_from_hkl(hkl, space_group) - acentric_mask = ~centric_mask - return centric_mask, acentric_mask - - def estimate_mean_intensity_by_resolution( I: torch.Tensor, d_spacings: torch.Tensor, n_bins: int = 60, min_per_bin: int = 40 ) -> torch.Tensor: @@ -1323,7 +1194,7 @@ def estimate_mean_intensity_by_resolution( # Use scatter_add to compute sum of intensities per bin bin_sums = torch.zeros(actual_n_bins, dtype=I.dtype, device=I.device) - bin_counts = torch.zeros(actual_n_bins, dtype=torch.long, device=I.device) + bin_counts = torch.zeros(actual_n_bins, dtype=torch.long, device=I.device) # dtype-ok: count accumulator; scatter_add source is long ones, dtype must match bin_sums.scatter_add_(0, bin_indices, I_sorted) bin_counts.scatter_add_(0, bin_indices, torch.ones_like(bin_indices)) @@ -1460,7 +1331,7 @@ def french_wilson_auto( ) # Step 2: Determine centric reflections from Miller indices - is_centric = is_centric_from_hkl(hkl, space_group=space_group) + is_centric = SpaceGroup(space_group, device=hkl.device).is_centric(hkl) # Step 3: Apply French-Wilson conversion F, sigma_F, valid_mask = french_wilson( @@ -1546,7 +1417,7 @@ def __init__( self.register_buffer("d_spacings", d_spacings) # Determine centric reflections - is_centric = is_centric_from_hkl(hkl, space_group) + is_centric = SpaceGroup(space_group, device=hkl.device).is_centric(hkl) self.register_buffer("is_centric", is_centric) # Set by forward(); None until the first conversion. Not a buffer -- it diff --git a/torchref/base/metrics/binwise_scale.py b/torchref/base/metrics/binwise_scale.py index 0ca6b213..2d803ef1 100644 --- a/torchref/base/metrics/binwise_scale.py +++ b/torchref/base/metrics/binwise_scale.py @@ -58,7 +58,7 @@ def binwise_scale( Fo = Fo.reshape(-1) device, dtype = Fc.device, Fc.dtype - bins = bins.reshape(-1).to(device=device, dtype=torch.int64) + bins = bins.reshape(-1).to(device=device, dtype=torch.int64) # dtype-ok: resolution-bin indices used as scatter_add index; requires int64 if nbins is None: nbins = int(bins.max().item()) + 1 if bins.numel() else 0 diff --git a/torchref/base/reciprocal/__init__.py b/torchref/base/reciprocal/__init__.py index 00003aac..cdaf6c20 100644 --- a/torchref/base/reciprocal/__init__.py +++ b/torchref/base/reciprocal/__init__.py @@ -33,12 +33,7 @@ smooth_reciprocal_grid ) -from .symmetry import ( - compute_symmetry_equivalent_hkls, - compute_translation_phases, - extract_structure_factors_with_symmetry, - ReciprocalSymmetryExtractor, -) +from .symmetry import ReciprocalSymmetryExtractor __all__ = [ # Basis functions @@ -62,8 +57,5 @@ "interpolate_for_rotation", "smooth_reciprocal_grid", # Symmetry - "compute_symmetry_equivalent_hkls", - "compute_translation_phases", - "extract_structure_factors_with_symmetry", "ReciprocalSymmetryExtractor", ] diff --git a/torchref/base/reciprocal/grid_operations.py b/torchref/base/reciprocal/grid_operations.py index 08fa138f..a631a640 100644 --- a/torchref/base/reciprocal/grid_operations.py +++ b/torchref/base/reciprocal/grid_operations.py @@ -47,14 +47,14 @@ def place_on_grid( dtype = structure_factor.dtype Nx, Ny, Nz = [int(x) for x in grid_size] hkls = hkls.to(device=device) - h = hkls[:, 0].to(torch.int64) - k = hkls[:, 1].to(torch.int64) - l = hkls[:, 2].to(torch.int64) + h = hkls[:, 0].to(torch.int64) # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long + k = hkls[:, 1].to(torch.int64) # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long + l = hkls[:, 2].to(torch.int64) # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long hi = torch.remainder(h, Nx) ki = torch.remainder(k, Ny) li = torch.remainder(l, Nz) - lin = (hi * (Ny * Nz) + ki * Nz + li).to(torch.int64) # (N,) + lin = (hi * (Ny * Nz) + ki * Nz + li).to(torch.int64) # (N,) # dtype-ok: flat grid index (lin) for scatter/gather; requires int64 grid = torch.zeros((B, Nx * Ny * Nz), dtype=dtype, device=device) grid = grid.index_add(1, lin, structure_factor) # (B, Nx*Ny*Nz) @@ -62,7 +62,7 @@ def place_on_grid( hi_sym = torch.remainder(-h, Nx) ki_sym = torch.remainder(-k, Ny) li_sym = torch.remainder(-l, Nz) - lin_sym = (hi_sym * (Ny * Nz) + ki_sym * Nz + li_sym).to(torch.int64) + lin_sym = (hi_sym * (Ny * Nz) + ki_sym * Nz + li_sym).to(torch.int64) # dtype-ok: symmetry flat grid index (lin_sym) for scatter/gather; requires int64 vals_conj = torch.conj(structure_factor) grid = grid.index_add(1, lin_sym, vals_conj) @@ -101,9 +101,9 @@ def extract_structure_factor_from_grid(reciprocal_grid, hkls) -> torch.Tensor: # Same wrapping convention as place_on_grid. hkls = hkls.to(device=device) - h = hkls[:, 0].to(torch.int64) - k = hkls[:, 1].to(torch.int64) - l = hkls[:, 2].to(torch.int64) + h = hkls[:, 0].to(torch.int64) # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long + k = hkls[:, 1].to(torch.int64) # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long + l = hkls[:, 2].to(torch.int64) # dtype-ok: hkl component cast to int64 for flat grid-index arithmetic; indexing requires long hi = torch.remainder(h, Nx) ki = torch.remainder(k, Ny) diff --git a/torchref/base/reciprocal/interpolation.py b/torchref/base/reciprocal/interpolation.py index e7db4f18..4b611634 100644 --- a/torchref/base/reciprocal/interpolation.py +++ b/torchref/base/reciprocal/interpolation.py @@ -39,6 +39,9 @@ def interpolate_structure_factor_from_grid( device = reciprocal_grid.device Nx, Ny, Nz = reciprocal_grid.shape + # dtype-ok: grid-index interpolation only. hkl are small integers, so floor + # and fractional weights are exact in float32; the weights are recast to the + # grid's dtype below, so nothing here mixes with config-dtype tensors. hkl_float = hkl_float.to(device=device, dtype=torch.float32) # Get the 8 corner indices for trilinear interpolation @@ -149,6 +152,9 @@ def interpolate_complex_from_grid( device = reciprocal_grid.device Nx, Ny, Nz = reciprocal_grid.shape + # dtype-ok: grid-index interpolation only. hkl are small integers, so floor + # and fractional weights are exact in float32; the weights are recast to the + # grid's dtype below, so nothing here mixes with config-dtype tensors. hkl_float = hkl_float.to(device=device, dtype=torch.float32) # Get the 8 corner indices for trilinear interpolation @@ -314,7 +320,9 @@ def interpolate_for_rotation(hkl, R, cell, reciprocal_space_grid): batched = False rotation_in_s = torch.einsum('ij, bjk -> bik', cell.reciprocal_basis_matrix, R.permute(0,2,1)) rotation_in_hkl = torch.einsum('bij, jk -> bik', rotation_in_s, cell.reciprocal_basis_matrix.inverse()) - reoriented_hkl = torch.einsum('aj, bji -> bai', hkl.to(torch.float32), rotation_in_hkl) + # Match the reciprocal-basis math's dtype (config float): a hardcoded float32 + # here mixes with a float64 rotation under a float64 config and raises. + reoriented_hkl = torch.einsum('aj, bji -> bai', hkl.to(rotation_in_hkl.dtype), rotation_in_hkl) shape = reoriented_hkl.shape reoriented_hkl = reoriented_hkl.reshape(-1, 3) interpolated = interpolate_structure_factor_from_grid(reciprocal_space_grid, reoriented_hkl).reshape(shape[0], shape[1]) diff --git a/torchref/base/reciprocal/symmetry.py b/torchref/base/reciprocal/symmetry.py index 72d4ed0d..390a13b4 100644 --- a/torchref/base/reciprocal/symmetry.py +++ b/torchref/base/reciprocal/symmetry.py @@ -1,215 +1,89 @@ """Reciprocal-space ("late") symmetry for structure factor calculation. -The alternative to symmetrizing the density map before the FFT ("early" symmetry, -:func:`~torchref.symmetry.MapSymmetry`): here symmetry is applied to the P1 -transform afterwards, avoiding the map symmetrization entirely. Per operation +The alternative to symmetrizing the density map before the FFT ("early" symmetry, via +:meth:`~torchref.symmetry.symmetry.Symmetry.symmetrize_map`): here symmetry is applied +to the P1 transform afterwards, avoiding the map symmetrization entirely. Per operation {R|t}, - F_sym(h) = Σ_ops exp(2πi h·t) · F_P1(Rᵀ·h) + F_sym(h) = sum_ops exp(2 pi i h.t) * F_P1(R^T h) -and because crystallographic R is integer-valued, Rᵀ·h lands exactly on grid -points and needs no interpolation. **Every grid or map argument here is the P1 -one** -- feeding in an already-symmetrized grid double-counts. +and because crystallographic R is integer-valued, ``R^T h`` lands exactly on grid +points and needs no interpolation. **Every grid or map argument here is the P1 one** -- +feeding in an already-symmetrized grid double-counts. + +Both halves of that sum come from :class:`~torchref.symmetry.symmetry.Symmetry`: +``R^T h`` from :meth:`~torchref.symmetry.symmetry.Symmetry.expand_reciprocal` and the +phases from :meth:`~torchref.symmetry.symmetry.Symmetry.phase_factors`. Reach this +class through +:meth:`~torchref.symmetry.symmetry.Symmetry.reciprocal_extractor`, which caches it. """ -from typing import Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Optional -import numpy as np import torch -from torchref.config import canonical_device, get_float_dtype +from torchref.config import canonical_device from torchref.utils.autograd_ops import gather_with_index_add - -from .grid_operations import extract_structure_factor_from_grid +from torchref.utils.device_mixin import DeviceMixin if TYPE_CHECKING: - from torchref.symmetry.spacegroup import SpaceGroup - - -def compute_symmetry_equivalent_hkls( - hkl: torch.Tensor, - rotation_matrices: torch.Tensor, -) -> torch.Tensor: - """ - Compute symmetry-equivalent HKLs for each operation. - - Row-vector convention: h' = h @ R, equivalently Rᵀ·h for column h. The - matrices are used as given -- do **not** pre-transpose them. - - Parameters - ---------- - hkl : torch.Tensor, shape (N, 3) - Miller indices. - rotation_matrices : torch.Tensor, shape (n_ops, 3, 3) - Real-space rotation matrices, applied directly as ``h @ R`` (no - transpose). - - Returns - ------- - torch.Tensor, shape (n_ops, N, 3) - Equivalent HKLs for each symmetry operation. - """ - device = hkl.device - dtype = get_float_dtype() - - hkl_float = hkl.to(dtype=dtype, device=device) # (N, 3) - rot_matrices = rotation_matrices.to(dtype=dtype, device=device) # (n_ops, 3, 3) - - n_ops = rot_matrices.shape[0] - - hkl_expanded = hkl_float.unsqueeze(0).expand(n_ops, -1, -1) # (n_ops, N, 3) - - equiv_hkl = torch.bmm(hkl_expanded, rot_matrices) - - # Exact for valid crystallographic ops; round only mops up float error. - equiv_hkl = torch.round(equiv_hkl).to(torch.int64) - - return equiv_hkl - - -def compute_translation_phases( - hkl: torch.Tensor, - translations: torch.Tensor, -) -> torch.Tensor: - """ - Compute the translation phase shifts exp(2πi h·t) for each operation. - - Parameters - ---------- - hkl : torch.Tensor, shape (N, 3) - Miller indices. - translations : torch.Tensor, shape (n_ops, 3) - Translation vectors in fractional coordinates. - - Returns - ------- - torch.Tensor, shape (n_ops, N) - Complex phase factors exp(2*pi*i * h.t). - """ - device = hkl.device - dtype = get_float_dtype() - - hkl_float = hkl.to(dtype=dtype, device=device) # (N, 3) - translations = translations.to(dtype=dtype, device=device) # (n_ops, 3) - - h_dot_t = torch.matmul(hkl_float, translations.T).T # (n_ops, N) + from torchref.symmetry.symmetry import Symmetry - # Leave ``phase`` at the configured float dtype; casting it to float32 would - # force complex64 output even under a float64 configuration. - phase = 2.0 * np.pi * h_dot_t - phase_factor = torch.exp(1j * phase) - return phase_factor # (n_ops, N) complex - - -def extract_structure_factors_with_symmetry( - reciprocal_grid: torch.Tensor, - hkl: torch.Tensor, - rotation_matrices: torch.Tensor, - translations: torch.Tensor, +def _equiv_hkls_to_flat_indices( + equiv_hkls: torch.Tensor, Nx: int, Ny: int, Nz: int ) -> torch.Tensor: - """ - Extract structure factors with symmetry applied in reciprocal space. - - Sums F over the symmetry-equivalent positions with their translation phases, - replacing the symmetrize-then-extract MapSymmetry route. For repeated calls - with the same hkl and symmetry, use :class:`ReciprocalSymmetryExtractor`. + """Flatten symmetry-equivalent Miller indices into linear grid indices. Parameters ---------- - reciprocal_grid : torch.Tensor, shape (Nx, Ny, Nz) - Complex reciprocal space grid from FFT of the **P1** density map. - hkl : torch.Tensor, shape (N, 3) - Target Miller indices. - rotation_matrices : torch.Tensor, shape (n_ops, 3, 3) - Real-space rotation matrices from symmetry operations. - translations : torch.Tensor, shape (n_ops, 3) - Translation vectors from symmetry operations. + equiv_hkls : torch.Tensor + Equivalent indices, shape ``(n_ops, N, 3)``. + Nx, Ny, Nz : int + Reciprocal grid dimensions. Returns ------- - torch.Tensor, shape (N,) - Complex structure factors with symmetry applied. + torch.Tensor + Flat indices, shape ``(n_ops * N,)``, dtype ``int64``, wrapped modulo the grid. """ - device = reciprocal_grid.device - Nx, Ny, Nz = reciprocal_grid.shape - - # Move everything to the same device - hkl = hkl.to(device=device) - rotation_matrices = rotation_matrices.to(device=device) - translations = translations.to(device=device) - - n_ops = rotation_matrices.shape[0] - N = hkl.shape[0] - - equiv_hkls = compute_symmetry_equivalent_hkls(hkl, rotation_matrices) - - # One flat gather for all symops. gather_with_index_add keeps the backward a - # single index_add_ instead of a radix-sort + dedup scatter. - flat_indices = _equiv_hkls_to_flat_indices(equiv_hkls, Nx, Ny, Nz) - f_all = gather_with_index_add( - reciprocal_grid.reshape(-1), flat_indices, - ) # (n_ops * N,) - f_p1 = f_all.view(n_ops, N) - - phases = compute_translation_phases(hkl, translations) - - f_sym = (f_p1 * phases).sum(dim=0) - - return f_sym - - -def _equiv_hkls_to_flat_indices( - equiv_hkls: torch.Tensor, Nx: int, Ny: int, Nz: int, -) -> torch.Tensor: - """Convert (n_ops, N, 3) equiv HKLs to flat linear grid indices.""" - all_hkl = equiv_hkls.reshape(-1, 3) # (n_ops*N, 3) + all_hkl = equiv_hkls.reshape(-1, 3) hi = torch.remainder(all_hkl[:, 0], Nx) ki = torch.remainder(all_hkl[:, 1], Ny) li = torch.remainder(all_hkl[:, 2], Nz) - return (hi * (Ny * Nz) + ki * Nz + li).to(torch.int64) - - -from torchref.utils.device_mixin import DeviceMixin + return (hi * (Ny * Nz) + ki * Nz + li).to(torch.int64) # dtype-ok: flat HKL grid index; int64 avoids overflow, used for indexing class ReciprocalSymmetryExtractor(DeviceMixin): - """ - Class-based interface for reciprocal space symmetry extraction. + """Precomputed symmetrized structure-factor extraction at fixed hkl and grid. - For repeated structure-factor evaluation at fixed hkl and symmetry, as in - refinement: the equivalent HKLs, phase factors and flat grid indices are - precomputed here, so each call is one gather, multiply and sum. The - precomputation binds ``grid_shape`` -- a differently shaped grid needs a new - extractor. + For repeated evaluation during refinement: the equivalent indices, phases and flat + gather indices are computed once here, so each call is one gather, multiply and + sum. The precomputation binds both ``hkl`` and ``grid_shape`` -- either changing + needs a new extractor, which is what + :meth:`~torchref.symmetry.symmetry.Symmetry.reciprocal_extractor` tracks. Parameters ---------- - hkl : torch.Tensor, shape (N, 3) - Target Miller indices. - symmetry : SpaceGroup - SpaceGroup object containing rotation matrices and translations. + hkl : torch.Tensor + Target Miller indices, shape ``(N, 3)``. + symmetry : Symmetry + The group supplying the operations. grid_shape : tuple of int - Reciprocal grid dimensions (Nx, Ny, Nz). + Reciprocal grid dimensions ``(Nx, Ny, Nz)``. device : torch.device, optional - Device for computation. - - Examples - -------- - >>> extractor = ReciprocalSymmetryExtractor(hkl, symmetry, grid_shape=(209, 86, 67)) - >>> f_calc = extractor.extract_from_grid(reciprocal_grid) + Device for the precomputed tensors. Defaults to ``hkl``'s. """ def __init__( self, hkl: torch.Tensor, - symmetry: "SpaceGroup", + symmetry: "Symmetry", grid_shape: tuple, device: Optional[torch.device] = None, ): - # ``is not None``, not ``or``: ``device=0`` means cuda:0/mps:0 and is - # falsy, so ``or`` silently discarded it. ``hkl`` is a bare tensor, so - # this reads its device rather than going through ``resolve_device``. + # ``is not None``, not ``or``: ``device=0`` means cuda:0/mps:0 and is falsy, so + # ``or`` silently discarded it. self.device = canonical_device( device if device is not None else hkl.device ) @@ -219,67 +93,56 @@ def __init__( self.N = len(hkl) self.grid_shape = grid_shape - # Precompute equivalent HKLs - self.equiv_hkls = compute_symmetry_equivalent_hkls( - self.hkl, - symmetry.matrices.to(device=self.device), - ) # (n_ops, N, 3) + self.equiv_hkls = symmetry.expand_reciprocal(self.hkl).to(device=self.device) + self.phases = symmetry.phase_factors(self.hkl).to(device=self.device) - # Precompute phase factors - self.phases = compute_translation_phases( - self.hkl, - symmetry.translations.to(device=self.device), - ) # (n_ops, N) complex - - # Precompute flat linear indices for single-gather extraction Nx, Ny, Nz = grid_shape self._flat_indices = _equiv_hkls_to_flat_indices( - self.equiv_hkls, Nx, Ny, Nz, - ) # (n_ops * N,) int64 + self.equiv_hkls, Nx, Ny, Nz + ) def __call__(self, density_map: torch.Tensor) -> torch.Tensor: """Alias for :meth:`extract`.""" return self.extract(density_map) def extract(self, density_map: torch.Tensor) -> torch.Tensor: - """ - Transform a P1 density map and extract symmetrized structure factors. + """Transform a P1 density map and extract symmetrized structure factors. Parameters ---------- - density_map : torch.Tensor, shape (Nx, Ny, Nz) - **P1** electron density map -- passing a symmetrized map double-counts. + density_map : torch.Tensor + **P1** electron density, shape ``(Nx, Ny, Nz)``. A symmetrized map + double-counts. Returns ------- - torch.Tensor, shape (N,) - Complex structure factors with symmetry applied. + torch.Tensor + Complex structure factors, shape ``(N,)``. """ from torchref.base.fourier.fft import ifft - reciprocal_grid = ifft(density_map) - return self.extract_from_grid(reciprocal_grid) + return self.extract_from_grid(ifft(density_map)) def extract_from_grid(self, reciprocal_grid: torch.Tensor) -> torch.Tensor: - """ - Extract structure factors from an already-transformed P1 grid. + """Extract structure factors from an already-transformed P1 grid. Parameters ---------- - reciprocal_grid : torch.Tensor, shape (Nx, Ny, Nz) - Complex reciprocal space grid from FFT of the **P1** map; its shape - must match the ``grid_shape`` this extractor was built for. + reciprocal_grid : torch.Tensor + Complex grid from the FFT of the **P1** map, shape ``(Nx, Ny, Nz)``; + its shape must match the ``grid_shape`` this extractor was built for. Returns ------- - torch.Tensor, shape (N,) - Complex structure factors with symmetry applied. + torch.Tensor + Complex structure factors, shape ``(N,)``. """ # gather_with_index_add keeps the backward a single ``index_add_`` # (atomic scatter, no radix sort + dedup). f_all = gather_with_index_add( - reciprocal_grid.reshape(-1), self._flat_indices, + reciprocal_grid.reshape(-1), self._flat_indices ) # (n_ops * N,) - f_sym = (f_all.view(self.n_ops, self.N) * self.phases).sum(dim=0) - return f_sym + return (f_all.view(self.n_ops, self.N) * self.phases).sum(dim=0) + +__all__ = ["ReciprocalSymmetryExtractor"] diff --git a/torchref/base/scattering/scattering_table.py b/torchref/base/scattering/scattering_table.py index 8faf3b5d..fbc26aab 100644 --- a/torchref/base/scattering/scattering_table.py +++ b/torchref/base/scattering/scattering_table.py @@ -168,7 +168,7 @@ def get_scattering_params_by_z( table = load_scattering_table(device=device, dtype=dtype) # Long, not the caller's int32: torch indexing requires it. - z_idx = z_tensor.to(device=device, dtype=torch.long) + z_idx = z_tensor.to(device=device, dtype=torch.long) # dtype-ok: z cast to long for scattering-table index lookup; indexing requires long A = table["A"][z_idx] B = table["B"][z_idx] @@ -254,4 +254,4 @@ def elements_to_z(elements: list, normalize: bool = True) -> torch.Tensor: z = element_to_z.get(elem, 0) z_values.append(z) - return torch.tensor(z_values, dtype=torch.int32) + return torch.tensor(z_values, dtype=torch.int32) # dtype-ok: atomic-number Z categorical codes; fixed int32 lookup keys diff --git a/torchref/base/targets/_dispatch.py b/torchref/base/targets/_dispatch.py index 40babcab..0b9562a0 100644 --- a/torchref/base/targets/_dispatch.py +++ b/torchref/base/targets/_dispatch.py @@ -54,7 +54,7 @@ def why_unavailable() -> Optional[str]: name="triton", kernel=None, # gate-only; see the module docstring device="cuda", - dtypes=(torch.float32,), + dtypes=(torch.float32,), # dtype-ok: backend capability declaration, not an allocation probe=(_THIS, "why_unavailable"), expect_available="cuda", # The probe handles availability, so this governs only a kernel that diff --git a/torchref/base/targets/dataset_scaling.py b/torchref/base/targets/dataset_scaling.py new file mode 100644 index 00000000..2a877dd7 --- /dev/null +++ b/torchref/base/targets/dataset_scaling.py @@ -0,0 +1,51 @@ +"""Profile a shared amplitude against independently measured datasets. + +The inputs use a common reflection axis with explicit presence masks. Uncertainties +remain in their measurement units; scaling and consensus estimation are differentiable. +""" + +import torch + + +def dataset_scaling_loss( + amplitudes: torch.Tensor, + sigmas: torch.Tensor, + log_corrections: torch.Tensor, + mask: torch.Tensor, +) -> torch.Tensor: + """Return the summed profiled Gaussian least-squares loss. + + Parameters + ---------- + amplitudes, sigmas : torch.Tensor + Measured amplitudes and positive uncertainties, shape (N, H), in each + dataset's input amplitude units. Masked entries may be non-finite. + log_corrections : torch.Tensor + Dimensionless log amplitude corrections, shape (N, H). + mask : torch.Tensor + Boolean usable-observation mask, shape (N, H). Only columns with at + least two observations contribute. + + Returns + ------- + torch.Tensor + Scalar dimensionless loss. Gradients include the consensus and the + scale dependence of the propagated uncertainties. + """ + active = mask & (mask.sum(dim=0, keepdim=True) >= 2) + obs = torch.where(active, amplitudes, torch.zeros_like(amplitudes)) + sigma = torch.where(active, sigmas, torch.ones_like(sigmas)) + log_k = torch.where(active, log_corrections, torch.zeros_like(log_corrections)) + # In measurement units the model is mu / k and sigma is fixed. Rescaling + # its whitened design column avoids overflow without changing the projection. + log_design = -log_k - sigma.log() + floor = torch.finfo(log_design.dtype).min + shift = torch.where(active, log_design, floor).amax(dim=0, keepdim=True) + shift = torch.where(active.any(dim=0, keepdim=True), shift, torch.zeros_like(shift)) + shifted = torch.where(active, log_design - shift, torch.zeros_like(obs)) + design = torch.where(active, shifted.exp(), torch.zeros_like(obs)) + whitened = obs / sigma + norm = design.square().sum(dim=0).clamp_min(torch.finfo(design.dtype).tiny) + consensus = (design * whitened).sum(dim=0) / norm + residual = torch.where(active, whitened - design * consensus, torch.zeros_like(obs)) + return 0.5 * residual.square().sum() diff --git a/torchref/base/targets/triton/place_hydrogens.py b/torchref/base/targets/triton/place_hydrogens.py index 31f74461..f88c13f0 100644 --- a/torchref/base/targets/triton/place_hydrogens.py +++ b/torchref/base/targets/triton/place_hydrogens.py @@ -1,7 +1,7 @@ """Triton forward + analytic backward for riding-hydrogen placement. One launch each way, where the eager helper (``_place_h_jit`` in -:mod:`torchref.restraints.hydrogen_topology`) fuses only the forward and leaves its +:mod:`torchref.topology.riding`) fuses only the forward and leaves its backward to run op-by-op through autograd -- ~100 launches at 3k hydrogens, which dominates the non-bonded backward. The math mirrors ``_place_h_jit`` exactly: diff --git a/torchref/base/targets/xray_likelihoods.py b/torchref/base/targets/xray_likelihoods.py index e385eeb7..e70c263d 100644 --- a/torchref/base/targets/xray_likelihoods.py +++ b/torchref/base/targets/xray_likelihoods.py @@ -22,6 +22,15 @@ indistinguishable from a change of x-ray weight. ``sigma_obs**2`` needs **no** conversion: it is already an amplitude variance, a 1-DOF error on a measured amplitude. +**The observable is a third axis, and it is only the variance and the mean that carry it.** +:func:`gaussian_per_refl` is a Gaussian on any real observable; :func:`nll_per_refl` is that +same function on amplitudes, and the intensity rows are it on intensities. Only the variance +builder has to know which: :func:`amplitude_var_from_sigma_obs` vs +:func:`intensity_var_from_sigma_obs`, and confusing them is wrong by ``2|F|`` -- which is +resolution-dependent, so it presents as a scale or B error rather than as a bug. Rice has no +intensity twin: Rice and the folded normal are distributions *of an amplitude*, and the +intensity analogue is the exponential / chi-square_1 Wilson distribution. + Do not pair ``sigma_obs`` with a Rice ``Sigma``: that asserts an isotropic *complex* error where ``sigma_obs`` carries no phase at all, and no regime makes it correct (it was tried, and was the worst of every target). Model error -- ``beta`` -- is what belongs in a Rice @@ -39,12 +48,59 @@ #: Floor on any variance before it reaches a division or a log. VAR_FLOOR = 1e-10 +#: Floor on a *measured* sigma, as a fraction of its median over the fitted subset. +#: Data-dependent rather than absolute, because the scale of a sigma is the scale of the +#: data. Merged intensities in particular are reported with ``sigma == 0`` rows. +SIGMA_FLOOR_FRAC = 1e-1 + +#: Backstop under :data:`SIGMA_FLOOR_FRAC` for the intensity builder, for the pathological +#: case of a median that is itself ~0. The amplitude builder deliberately has none -- see +#: :func:`floor_sigma_obs`. +SIGMA_FLOOR_ABS = 1e-12 + # ===================================================================== # Variance builders -- the axis that distinguishes the five targets # ===================================================================== +def floor_sigma_obs( + sigma: torch.Tensor, + mask: torch.Tensor = None, + abs_floor: float = 0.0, + floor=None, +) -> torch.Tensor: + """Clamp a measured sigma at :data:`SIGMA_FLOOR_FRAC` of its median. + + ``mask`` restricts the median to the fitted subset -- which matters when the unfitted + rows carry filler sigmas, as reindexed collection members do. ``mask=None`` takes the + median over everything. + + ``abs_floor`` is a backstop under the fractional floor. It defaults to **off** because + the amplitude builder shipped without one and has a Triton counterpart to stay + bit-identical to; the intensity builder passes :data:`SIGMA_FLOOR_ABS`. + + **Pass ``floor`` explicitly to make the result independent of which reflections are in + ``sigma``.** A median computed from the argument makes every per-reflection value + depend on the whole array, so the same reflection scores differently in a subset sum + than in a full-size residual -- measured at 0.09% on a work set and 1.8% on a free set + for intensities, whose sigmas span orders of magnitude. Callers that need the two to + agree (any target with both a ``forward`` and a ``residuals``) compute the floor once + from their own fitted subset and pass it here. + """ + if floor is None: + selected = sigma if mask is None else sigma[mask] + if selected.numel() == 0: + # No fitted reflections to take a median over. Any positive floor is arbitrary + # here; what matters is that it is finite, so a later log or division cannot + # produce a NaN that would poison the whole gradient. + return sigma.clamp(min=1e-6) + floor = torch.median(selected) * SIGMA_FLOOR_FRAC + if abs_floor > 0.0: + floor = torch.clamp(torch.as_tensor(floor), min=abs_floor) + return sigma.clamp(min=floor) + + def amplitude_var_from_sigma_obs(sigma: torch.Tensor) -> torch.Tensor: """Amplitude variance from the experimental sigma: ``clamp(sigma)**2``. @@ -53,8 +109,25 @@ def amplitude_var_from_sigma_obs(sigma: torch.Tensor) -> torch.Tensor: beta-derived builders below floor only at :data:`VAR_FLOOR`; the two conventions are deliberately not reconciled. """ - floor = torch.median(sigma) * 1e-1 - return torch.clamp(sigma, min=floor) ** 2 + return floor_sigma_obs(sigma) ** 2 + + +def intensity_var_from_sigma_obs( + sigma: torch.Tensor, mask: torch.Tensor = None, floor=None +) -> torch.Tensor: + """Intensity variance from the experimental ``sigma(I)``: ``clamp(sigma)**2``. + + The intensity twin of :func:`amplitude_var_from_sigma_obs`, and **not** interchangeable + with it: applying an amplitude sigma to an intensity residual is wrong by a factor of + ``2|F|``, which is resolution-dependent and so looks like a scale or B error rather + than like a mistake. + + Differs from the amplitude builder in taking a ``mask`` (merged collection members are + reindexed onto a common list, so the unfitted rows carry filler), an absolute backstop + at :data:`SIGMA_FLOOR_ABS`, and an explicit ``floor`` -- see :func:`floor_sigma_obs` on + why a caller with both a ``forward`` and a ``residuals`` must pass one. + """ + return floor_sigma_obs(sigma, mask, abs_floor=SIGMA_FLOOR_ABS, floor=floor) ** 2 def amplitude_var_from_complex( @@ -133,6 +206,36 @@ def _masked_sum(loss: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor: return (loss * mask).sum() +def gaussian_per_refl( + obs: torch.Tensor, + model: torch.Tensor, + var: torch.Tensor, + var_floor: float = VAR_FLOOR, +) -> torch.Tensor: + """Per-reflection Gaussian NLL on **any** real observable (NOT masked or summed). + + 0.5 * (obs - model)**2 / var + 0.5 * log(var) + 0.5 * log(2*pi) + + Observable-agnostic on purpose: ``var`` just has to be the variance of whatever + ``obs`` is. :func:`nll_per_refl` is this on amplitudes; the intensity rows are this on + intensities. Keeping one implementation is what stops the two drifting -- they were + separately written once, and the copies differed only in spelling ``log(sigma)`` + instead of ``0.5 * log(var)``. + + ``var_floor`` defaults to :data:`VAR_FLOOR` because the amplitude path has always + applied it. **Pass 0.0 when the variance builder has already floored the sigma**, as + :func:`intensity_var_from_sigma_obs` does. An absolute floor on a variance is + dimensionally arbitrary -- 1e-10 is a distortion, not a safeguard, on any dataset whose + sigmas are smaller than ~1e-5, and it silently reweights the whole objective by up to + the ratio of the two floors. Positivity is the builder's job; this is a backstop for + builders that do not do it. + """ + diff = obs - model + if var_floor > 0.0: + var = torch.clamp(var, min=var_floor) + return 0.5 * diff**2 / var + 0.5 * torch.log(var) + HALF_LOG_2PI + + def nll_per_refl( F_obs: torch.Tensor, F_calc: torch.Tensor, var: torch.Tensor ) -> torch.Tensor: @@ -143,10 +246,11 @@ def nll_per_refl( ``var`` is the **amplitude** variance. Build it with :func:`amplitude_var_from_sigma_obs` (``nll``) or :func:`amplitude_var_from_complex` (``nll_beta``) -- see the module docstring on why those are not interchangeable. + + The ``torch.abs`` is what makes this the amplitude entry point: callers pass a complex + or signed ``F_calc``. Everything else is :func:`gaussian_per_refl`. """ - diff = F_obs - torch.abs(F_calc) - var = torch.clamp(var, min=VAR_FLOOR) - return 0.5 * diff**2 / var + 0.5 * torch.log(var) + HALF_LOG_2PI + return gaussian_per_refl(F_obs, torch.abs(F_calc), var) def nll_math( diff --git a/torchref/base/targets/xray_ml_full.py b/torchref/base/targets/xray_ml_full.py index 381ed3e9..be37e6e8 100644 --- a/torchref/base/targets/xray_ml_full.py +++ b/torchref/base/targets/xray_ml_full.py @@ -220,6 +220,7 @@ def acentric_nll(F_obs, sigma, Fc, Sigma, n_quad=None, n_sigma=None, li0=log_i0) if ( li0 is log_i0 + # dtype-ok: validation guard (compile eligibility), not an allocation and F_obs.dtype is not torch.float64 and F_obs.numel() > 1 # 0/1-specialisation would force a 2nd compile and get_compile_targets() diff --git a/torchref/base/wilson_outliers.py b/torchref/base/wilson_outliers.py index a9283676..e1f0d6c7 100644 --- a/torchref/base/wilson_outliers.py +++ b/torchref/base/wilson_outliers.py @@ -520,6 +520,7 @@ def _normal_quantile(p: float) -> float: low, high = -40.0, 10.0 for _ in range(200): mid = 0.5 * (low + high) + # dtype-ok: deliberate float64 for a scalar CDF; extracted via float() value = float(log_normal_cdf(torch.tensor(mid, dtype=torch.float64))) if value < target: low = mid diff --git a/torchref/cli/__init__.py b/torchref/cli/__init__.py index dbcf9ebe..efba6bad 100644 --- a/torchref/cli/__init__.py +++ b/torchref/cli/__init__.py @@ -5,8 +5,8 @@ __all__ = [ "add_metadata", "collection_difference_refine", + "difference_map", "mtz2map", - "phased_difference_map", "refine", "strip_altlocs", "validate_ded", diff --git a/torchref/cli/_common.py b/torchref/cli/_common.py index 67a90602..135e62a8 100644 --- a/torchref/cli/_common.py +++ b/torchref/cli/_common.py @@ -144,12 +144,43 @@ def add_adp_mode_arg(parser: argparse.ArgumentParser) -> None: "--adp-mode", type=str, default="isotropic", - choices=["isotropic", "anisotropic"], + choices=["isotropic", "anisotropic", "field", "field_aniso", "preserve"], help="ADP parametrization: 'isotropic' (default) refines a per-atom " "B-factor; 'anisotropic' refines a 6-component U tensor for the atoms " "given by --anisotropic-selection. The model is converted between " "representations and the output PDB/mmCIF follows the convention " - "(ANISOU only for anisotropic atoms).", + "(ANISOU only for anisotropic atoms). 'field' and 'field_aniso' replace " + "the per-atom parameters with a node field, whose size is set from the " + "data rather than the atom count (--reflections-per-adp-parameter). " + "'preserve' leaves the input file's own ADPs untouched.", + ) + parser.add_argument( + "--adp-mode-set", + type=str, + default=None, + choices=["constant", "rigid", "rigid_dilation", "affine"], + help="Displacement-mode set for --adp-mode field_aniso. Each node carries " + "the covariance of these modes, so its ADP varies across the region it " + "serves: 'constant' is one U per node, 'rigid' is TLS, 'rigid_dilation' " + "adds uniform breathing, 'affine' adds shear and extension.", + ) + parser.add_argument( + "--reflections-per-adp-parameter", + type=float, + default=7.0, + metavar="R", + help="Work reflections per ADP parameter a node field is sized to hold " + "(--adp-mode field/field_aniso). Default 7. Node count follows from the " + "data rather than the atom count, and both directions from 7 measured " + "worse. Ignored by the per-atom modes.", + ) + parser.add_argument( + "--adp-nodes", + type=int, + default=None, + metavar="N", + help="Explicit node count for a field ADP mode, bypassing " + "--reflections-per-adp-parameter.", ) parser.add_argument( "--anisotropic-selection", @@ -282,6 +313,7 @@ def add_dual_model_args( parser: argparse.ArgumentParser, fraction_required: bool = True, fraction_default: Optional[float] = None, + light_model_required: bool = True, ) -> None: """Add the standard dual-model (dark/light) input arguments. @@ -289,6 +321,10 @@ def add_dual_model_args( ``-lm``/``--light-model``, ``-dsf``/``--dark-structure-factor``, ``-lsf``/``--light-structure-factor``, ``--fraction``, ``--cif`` and a *Column selection* group with per-side column flags. + + ``light_model_required=False`` makes ``-lm`` optional, for tools that can do + something useful with the dark model alone -- a weighted difference map needs only + the dark state's phases. Refinement cannot: it refines the light model. """ inp = parser.add_argument_group("Input files") inp.add_argument( @@ -298,12 +334,14 @@ def add_dual_model_args( type=str, help="Dark / reference state model file (PDB or CIF)", ) + light_help = "Light / triggered state model file (PDB or CIF)" + if not light_model_required: + light_help += ( + ". Optional: without it only the weighted difference map is written, " + "using the dark state's phases." + ) inp.add_argument( - "-lm", - "--light-model", - required=True, - type=str, - help="Light / triggered state model file (PDB or CIF)", + "-lm", "--light-model", required=light_model_required, type=str, help=light_help ) inp.add_argument( "-dsf", @@ -335,6 +373,59 @@ def add_dual_model_args( add_dual_column_args(col) +def add_all_columns_arg(parser: argparse.ArgumentParser) -> None: + """Add ``--all-columns`` for the difference MTZ writer. + + Off by default so the output file holds the map a reader wants and can identify. + The gated columns are alternative constructions of the same quantities -- a + model-phased difference, two more extrapolations, the intensity block -- which are + useful once you know which is which and misleading before then. + """ + parser.add_argument( + "--all-columns", action="store_true", default=False, + help="Write every alternative map coefficient and diagnostic column, not just " + "the default difference and extrapolated maps. Costs two further scale " + "fits for the extra extrapolations.", + ) + + +def add_ded_weight_args(parser: argparse.ArgumentParser) -> None: + """Add ``--ded-weight`` and ``--sigma-d-gamma`` for the difference-map writers. + + Every registered scheme's weight is written to the difference MTZ regardless; the + choice here decides which one the headline products (validate-ded correlations, + model-phased difference columns) carry. + """ + from torchref.maps.ded_weights import DEFAULT_SCHEME, SCHEMES + + parser.add_argument( + "--ded-weight", + choices=list(SCHEMES), + default=DEFAULT_SCHEME, + help="Per-reflection weight for difference coefficients: 'inverse_variance' " + "is 1/sigma^2, 'sigma_d' is the Wiener weight S/(S+sigma^2) from the " + "expected difference power (needs calibrated sigmas; check the reported " + f"clamped-shell count), 'none' is flat (default: {DEFAULT_SCHEME}). All " + "weights are written as columns.", + ) + parser.add_argument( + "--sigma-d-gamma", + type=float, + default=None, + metavar="GAMMA", + help="Fix the dark-amplitude exponent of the sigma_d power law in [0, 2] " + "instead of fitting it (default: fitted).", + ) + + +def sigma_d_config_from_args(args: argparse.Namespace): + """The :class:`~torchref.refinement.model_error_estimation.sigma_d.SigmaDConfig` + selected by ``--sigma-d-gamma``.""" + from torchref.refinement.model_error_estimation.sigma_d import SigmaDConfig + + return SigmaDConfig(gamma=getattr(args, "sigma_d_gamma", None)) + + def add_output_format_args(parser: argparse.ArgumentParser) -> None: """Add ``--output-format`` argument for coordinate file format.""" parser.add_argument( @@ -361,6 +452,14 @@ def add_metadata_args(parser: argparse.ArgumentParser) -> None: default=None, help="Author names for the output file header", ) + parser.add_argument( + "--output-remarks", + type=str, + default=None, + help="Free-text note for the output header (REMARK 3 OTHER REFINEMENT " + "REMARKS / _refine.details). Nothing is written here unless you ask " + "for it", + ) parser.add_argument( "--no-header", action="store_true", @@ -612,6 +711,8 @@ def load_model( device: Union[str, "torch.device", None] = None, verbose: int = 0, cif: Optional[Union[str, List[str]]] = None, + add_hydrogens: bool = False, + hydrogens_in_xray: bool = True, ) -> "ModelFT": """Load a model from PDB or CIF, auto-detected by file extension. @@ -626,7 +727,12 @@ def load_model( verbose : int Verbosity passed to ModelFT. cif : str or list of str, optional - CIF restraint file(s) to load after the model. + CIF restraint file(s), registered on the model before it loads so that hydrogen + generation and the restraints read the same dictionary. + add_hydrogens : bool, optional + Generate missing hydrogens on load. Default False. + hydrogens_in_xray : bool, optional + Whether hydrogens contribute to the structure factors. Default True. Returns ------- @@ -636,16 +742,19 @@ def load_model( from torchref.config import normalize_device device = normalize_device(device) - model = ModelFT(max_res=max_res, device=device, verbose=verbose) + model = ModelFT( + max_res=max_res, + device=device, + verbose=verbose, + cif_path=cif, + add_hydrogens=add_hydrogens, + hydrogens_in_xray=hydrogens_in_xray, + ) suffix = Path(path).suffix.lower() if suffix in (".cif", ".mmcif"): model.load_cif(path) else: model.load_pdb(path) - - if cif is not None: - model.set_restraints_cif(cif) - return model @@ -740,6 +849,45 @@ def write_refinement_outputs( metadata.title = args.title if getattr(args, "authors", None): metadata.authors = args.authors + if getattr(args, "output_remarks", None): + metadata.output_remarks = args.output_remarks + + # What was minimised and how. These live on the CLI namespace rather + # than on the refinement, which is why from_refinement cannot fill them + # and why the header carried no method line at all until now. + xray_mode = getattr(args, "xray_mode", None) + if xray_mode: + # Name the family as well as the registry key. "ML" alone is + # cryptic in a deposited header, and the family follows from the + # key's prefix, so there is no lookup table here to drift out of + # step with XRAY_TARGETS. + key = str(xray_mode).upper() + if key.startswith(("ML", "NLL")): + metadata.target_function = f"MAXIMUM LIKELIHOOD ({key})" + elif key.startswith("LS"): + metadata.target_function = f"LEAST SQUARES ({key})" + else: + metadata.target_function = key + optimizer_parts = [] + n_cycles = getattr(args, "n_cycles", None) + if n_cycles: + optimizer_parts.append( + f"{n_cycles} MACROCYCLE" + ("S" if n_cycles != 1 else "") + ) + mode = getattr(args, "mode", None) + if mode: + optimizer_parts.append(str(mode).upper()) + adp_mode = getattr(args, "adp_mode", None) + if adp_mode: + optimizer_parts.append(f"{str(adp_mode).upper()} ADP") + # The scale target changes the R-factors this very header reports, so a + # run cannot be attributed without it (see the note beside it in the + # refinement_history.json parameters block). + scale_target = getattr(args, "scale_target", None) + if scale_target: + optimizer_parts.append(f"SCALE TARGET {str(scale_target).upper()}") + if optimizer_parts: + metadata.optimizer = ", ".join(optimizer_parts) outputs = {"pdb": None, "cif": None} diff --git a/torchref/cli/add_metadata.py b/torchref/cli/add_metadata.py index 8bae256b..9aa86076 100644 --- a/torchref/cli/add_metadata.py +++ b/torchref/cli/add_metadata.py @@ -128,7 +128,12 @@ def main(): # Start with pass-through from input file input_suffix = input_path.suffix.lower() if input_suffix == ".pdb": - metadata = RefinementMetadata.from_pdb_file(str(input_path)) + # This tool annotates a file, it does not re-refine it -- so the input's + # REMARK 3 and AUTHOR records are not superseded by anything and are + # kept. Refinement output takes the default and drops them. + metadata = RefinementMetadata.from_pdb_file( + str(input_path), supersede_refinement=False + ) elif input_suffix in (".cif", ".mmcif"): metadata = RefinementMetadata.from_cif_file(str(input_path)) else: diff --git a/torchref/cli/collection_difference_refine.py b/torchref/cli/collection_difference_refine.py index 97a935be..a6b36909 100644 --- a/torchref/cli/collection_difference_refine.py +++ b/torchref/cli/collection_difference_refine.py @@ -7,25 +7,18 @@ parameters (overall scale, anisotropy, bulk solvent k_sol/B_sol) is shared across the dark and light datasets. -Writes refined dark and light models (PDB/CIF), a JSON summary, and a difference MTZ: - -Observed ``Fo_dark``, ``SIGFo_dark``, ``Fo_light``, ``SIGFo_light`` -Differences ``DF`` = F_light - F_dark, ``SIGDF`` (propagated), ``WDF`` (sigma-weighted) -Calculated ``Fc_dark``, ``Fc_light`` (mixed model), ``DFc`` (scalar), ``DFc_complex`` -DED map coeffs ``2mDFop-DFc``, ``mDFop-DFc`` (phase-aware, figure-of-merit weighted) -Phases ``PHIC_dark``, ``PHIC_mixed``, ``PHIC_diff``, ``PHIC_light`` -Extrapolations ``Fextp`` (phase-aware), ``Fextc`` (amplitude-only, no phases), - ``Fextb`` (empirical-Bayes shrinkage toward Fo_dark -- see - :func:`compute_bayes_extrapolated_amplitudes`), each with its sigma - and ``2F-Fc`` / ``F-Fc`` map coefficients +Writes refined dark and light models (PDB/CIF), a JSON summary, and a difference MTZ. +See :func:`write_results_mtz` for the columns and why each is where it is -- the table +used to live here, four hundred lines from the function that writes it, which is part of +how the output drifted from its own documentation. Examples -------- :: - torchref.difference-refine \\ - -dm dark.pdb -lm light.pdb \\ - -dsf dark.mtz -lsf light.mtz \\ + torchref.difference-refine \ + -dm dark.pdb -lm light.pdb \ + -dsf dark.mtz -lsf light.mtz \ --fraction 0.37 -o output/ """ @@ -38,8 +31,10 @@ import torch from torchref.cli._common import ( - add_dual_model_args, + add_all_columns_arg, + add_ded_weight_args, add_dmin_arg, + add_dual_model_args, add_general_args, add_metadata_args, add_outdir_arg, @@ -49,12 +44,19 @@ configure_unbuffered_output, load_model, load_reflection_data, + parse_device_str, parse_weights, register_timing, - parse_device_str, + sigma_d_config_from_args, validate_cif_files, validate_files, ) +from torchref.maps.ded_weights import ( + DEFAULT_SCHEME, + WEIGHT_COLUMNS, + all_ded_weights, + reflection_geometry, +) from torchref.utils.serialization import convert_to_serializable configure_unbuffered_output() @@ -65,7 +67,14 @@ DEFAULT_TARGET_WEIGHTS = { "xray/difference": 1.0, - "xray/rice": 0.0, + # Selected by --difference-target; the schedule drives whichever row is chosen. + "xray/difference_sd": 0.0, + # The absolute channel. Zero by default: the difference refinement fixes the + # dark model, so the overall level is already anchored and this term only adds + # the systematic errors the difference cancels. + "xray/ml": 0.0, + # Registered only under --two-moment; harmless in the dict either way. + "xray/two_moment": 1.0, # "geometry/bond": 1.0, # geometry restraint should never require tuning, so leave at 1.0 # "geometry/angle": 1.0, # "geometry/torsion": 1.0, @@ -111,8 +120,8 @@ def setup_model_collection(pdb_dark, pdb_light, fractions, cif, d_min, sys.stdout.flush() model_dark = model_dark.hydrogenate(verbose=max(0, verbose - 1)) model_light = model_light.hydrogenate(verbose=max(0, verbose - 1)) - model_dark.exclude_H_from_sf = True - model_light.exclude_H_from_sf = True + model_dark.hydrogens_in_xray = False + model_light.hydrogens_in_xray = False mc = ModelCollection([model_dark, model_light], dark_key="dark") mc.add_dark() @@ -158,35 +167,104 @@ def setup_scaler(dataset_collection, model_collection, device, verbose=1): return scaler +def setup_dark_only(pdb_dark, dc, cif, d_min, device, verbose, hydrogenate=False): + """Load the dark model alone and scale it against the dark data. + + This is everything a weighted difference map needs. The amplitude is + ``|Fo_light| - |Fo_dark|``, which :meth:`DatasetCollection.scale` has already put on + one scale without reference to any model, and the phase comes from the dark state. + So there is no mixed model, no occupancy fraction, and no joint model-to-data fit -- + the joint fit exists to share scale parameters between two models, and here there is + only one. + + A fitted scaler rather than a bare ``angle(model(hkl))`` because bulk solvent + contributes a phase at low resolution, which is exactly where difference density is + largest. + + Returns + ------- + tuple + ``(model_dark, scaler)`` -- ready to hand to :func:`write_results_mtz` with + ``mc=None``. + """ + from torchref import Scaler + + model_dark = load_model( + pdb_dark, max_res=d_min, device=device, verbose=verbose, cif=cif, + ) + if hydrogenate: + if verbose > 0: + print("Adding hydrogens...") + sys.stdout.flush() + model_dark = model_dark.hydrogenate(verbose=max(0, verbose - 1)) + model_dark.hydrogens_in_xray = False + + scaler = Scaler( + model_dark, dc["dark"], device=device, verbose=max(-1, verbose - 1), + ) + scaler.initialize().refine_lbfgs() + return model_dark, scaler + + def compute_rfactors(model, data, scaler): - """Compute R-work/R-free using forward_mixed for proper solvent. + """Compute R-work/R-free with the scaler's own solvent model applied. Routes through ``rfactor_work_free`` — the shared source of truth used by the refinement targets — so the validity mask is applied and the validation set is excluded from both work and free, matching every other reported R-factor. + + Takes either scaler kind; see the branch below. """ from torchref.base.metrics.rfactor import rfactor_work_free with torch.no_grad(): hkl = data.hkl fcalc = model(hkl) - fcalc_scaled = scaler.forward_mixed(fcalc, model.fractions) + # CollectionScaler needs fractions to mix components; a single-dataset + # Scaler consumes structure factors directly. + if hasattr(scaler, "forward_mixed"): + fcalc_scaled = scaler.forward_mixed(fcalc, model.fractions) + else: + fcalc_scaled = scaler(fcalc) return rfactor_work_free(data, torch.abs(fcalc_scaled)) -def setup_loss_state(dataset_collection, model_collection, scaler, - target_weights, device, similarity_alpha=2.0): +def setup_loss_state( + dataset_collection, + model_collection, + scaler, + target_weights, + device, + similarity_alpha=2.0, + two_moment=False, + difference_target="difference", + sigma_d_config=None, +): """Build LossState with collection-aware targets. Geometry and ADP restraints are applied only to the light base model (the dark model is a frozen reference). + + Parameters + ---------- + two_moment : bool, optional + Also register the two-moment intensity target, which fits merged intensities + under ``|F(alpha)|^2 + sigma_alpha^2 |dF|^2``. Requires I/SIGI on every + dataset. Default False. + difference_target : {"difference", "difference_sd"}, optional + Which difference row the weight schedule drives. Both are registered, as + ``xray/difference`` and ``xray/difference_sd``; the other keeps the weight in + ``target_weights`` (zero by default). + sigma_d_config : SigmaDConfig, optional + Exponent and shrinkage settings of the ``difference_sd`` row's estimator. """ from torchref.refinement import LossState - from torchref.experimental.kinetic.targets import ( + from torchref.refinement.targets import TotalADPTarget, TotalGeometryTarget + from torchref.refinement.targets.collection import ( + CollectionDifferenceSigmaDTarget, CollectionDifferenceTarget, - CollectionRiceTarget, + CollectionMLTarget, ) - from torchref.refinement.targets import TotalADPTarget, TotalGeometryTarget from torchref.refinement.targets.similarity import CoordinateSimilarityTarget state = LossState(device=device) @@ -197,7 +275,16 @@ def setup_loss_state(dataset_collection, model_collection, scaler, diff_target = CollectionDifferenceTarget( dataset_collection, model_collection, scaler=scaler, ) - rice_target = CollectionRiceTarget( + diff_sd_target = CollectionDifferenceSigmaDTarget( + dataset_collection, + model_collection, + scaler=scaler, + sigma_d_config=sigma_d_config, + ) + selected_diff = {"difference": diff_target, "difference_sd": diff_sd_target}[ + difference_target + ] + ml_target = CollectionMLTarget( dataset_collection, model_collection, scaler=scaler, ) geom_target = TotalGeometryTarget(model_light) @@ -208,19 +295,41 @@ def setup_loss_state(dataset_collection, model_collection, scaler, ) state.register_target("xray/difference", diff_target) - state.register_target("xray/rice", rice_target) + state.register_target("xray/difference_sd", diff_sd_target) + state.register_target("xray/ml", ml_target) state.register_target("geometry", geom_target) state.register_target("adp", adp_target) state.register_target("similarity", similarity_target) + if two_moment: + from torchref.refinement.targets import CollectionTwoMomentIntensityTarget + + two_moment_target = CollectionTwoMomentIntensityTarget( + dataset_collection, model_collection, scaler=scaler, verbose=1, + ) + # Match intensity and amplitude gradient norms to balance the targets + # against the geometry restraints despite their different units. + two_moment_target.calibrate_base_weight( + selected_diff, list(model_light.parameters()) + ) + state.register_target("xray/two_moment", two_moment_target) + state.set_weights(target_weights) return state def compute_bayes_extrapolated_amplitudes( - Fobs_dark, Fobs_light, sig_dark, sig_light, phi_dark, phi_mixed, f, - *, tau_sq_floor=1e-4, + Fobs_dark, + Fobs_light, + sig_ext, + phi_dark, + phi_mixed, + f, + *, + tau_sq_floor=1e-4, + epsilon=None, + d_star_sq=None, ): """Empirical Bayes shrinkage estimator for extrapolated SF amplitudes. @@ -229,280 +338,828 @@ def compute_bayes_extrapolated_amplitudes( Fo_dark, regularising noisy high-resolution and weakly-measured reflections:: F_ext = |F_dark*e^(iφ_d) + ΔF/f| (phase-aware amplitude) - σ_ext² = (σ_light² + σ_dark²) / f² - τ² = max(<(F_ext - Fo_dark)²> - <σ_ext²>, floor) - w(h) = τ² / (τ² + σ_ext²(h)) + S(h) = expected power of (F_ext - Fo_dark), per resolution shell + w(h) = S(h) / (S(h) + σ_ext²(h)) F_extb = w(h)·F_ext + (1-w(h))·Fo_dark (amplitude shrinkage) + With ``d_star_sq`` the signal power comes per resolution shell from + :func:`~torchref.refinement.model_error_estimation.sigma_d.estimate_sigma_d` + (``<(F_ext - Fo_dark)²> - <σ_ext²>`` per shell, shrunk toward a smooth curve); + without it the single global ``τ² = max(<(F_ext - Fo_dark)²> - <σ_ext²>, floor)`` + is used, which is the one-shell special case. + Parameters ---------- Fobs_dark, Fobs_light : Tensor (N,) Observed amplitudes. - sig_dark, sig_light : Tensor (N,) - Measurement uncertainties. + sig_ext : Tensor (N,) + Propagated uncertainty of the extrapolated amplitude. Taken from the caller + rather than rebuilt here: ``F_ext`` is linear in the observations with + ``dF_ext/dF_light = 1/f`` and ``dF_ext/dF_dark = 1 - 1/f = -(1-f)/f``, so the + dark term carries a ``(1-f)**2`` weight. phi_dark, phi_mixed : Tensor (N,) Calculated phases (radians) for the dark and mixed models. f : float or Tensor Excited-state population fraction. tau_sq_floor : float - Floor on the estimated signal variance τ². + Floor on the estimated signal variance. + epsilon, d_star_sq : Tensor (N,), optional + Reflection multiplicity and ``1/d**2`` in A^-2. Given ``d_star_sq`` the signal + power is estimated per resolution shell. Returns ------- tuple - ``(F_ext_bayes, var_ext_bayes, w_shrinkage, tau_sq)`` -- extrapolated - amplitudes **before** shrinkage, posterior variance and shrinkage weight per - reflection, and the global τ² as a float. + ``(F_ext_bayes, var_ext_bayes, w_shrinkage, tau_sq)`` -- the **shrunk** + extrapolated amplitude, its posterior variance and the shrinkage weight per + reflection, and the count-weighted mean signal variance as a float. """ F_dark_phased = Fobs_dark * torch.exp(1j * phi_dark) F_light_phased = Fobs_light * torch.exp(1j * phi_mixed) delta_F = F_light_phased - F_dark_phased - # Propagated variance - sig_sq_dF = sig_light**2 + sig_dark**2 - sig_sq_ext = sig_sq_dF / f**2 + sig_sq_ext = sig_ext**2 # Phase-aware extrapolated amplitude F_ext_complex = F_dark_phased + delta_F / f F_ext = torch.abs(F_ext_complex) - # Estimate signal variance τ² - residuals_sq = (F_ext - Fobs_dark) ** 2 - tau_sq = max((residuals_sq.mean() - sig_sq_ext.mean()).item(), tau_sq_floor) + residual = F_ext - Fobs_dark + if d_star_sq is None: + tau_sq = max( + (residual.square().mean() - sig_sq_ext.mean()).item(), tau_sq_floor + ) + S = torch.full_like(F_ext, tau_sq) + else: + from torchref.refinement.model_error_estimation.sigma_d import ( + estimate_sigma_d, + sigma_d_per_reflection, + ) + + fit = torch.isfinite(residual) & torch.isfinite(sig_ext) + shells = estimate_sigma_d( + residual, sig_ext, epsilon, d_star_sq, None, fit, gamma=0.0 + ) + est = sigma_d_per_reflection(shells, d_star_sq, epsilon, None, sig_ext) + S = est.S.clamp(min=tau_sq_floor) + weight = shells.counts.clamp(min=1.0) + tau_sq = max( + float((shells.Sigma_N * shells.counts).sum() / weight.sum()), tau_sq_floor + ) # Per-reflection shrinkage weight (in [0, 1]) - w = tau_sq / (tau_sq + sig_sq_ext) + w = S / (S + sig_sq_ext) # Posterior variance - var_ext_bayes = (tau_sq * sig_sq_ext) / (tau_sq + sig_sq_ext) + var_ext_bayes = (S * sig_sq_ext) / (S + sig_sq_ext) + + # Shrink the amplitude toward Fo_dark -- scalar, so no phase interference. + F_ext_bayes = w * F_ext + (1 - w) * Fobs_dark + + return F_ext_bayes, var_ext_bayes, w, tau_sq - return F_ext, var_ext_bayes, w, tau_sq +def _two_moment_columns(mc, dc, mask, fcalc_dark_full, fcalc_mixed_full, + *, weights, diff_Fobs, Fcalc_diff_amp, Fobs_dark, + sig_dark, phi_mixed, F_obs_dark_phased, all_columns=False): + """Two-moment diagnostic columns, or empty dicts when the model is off. -def write_results_mtz(dc, mc, scaler, filename): - """Write difference / extrapolated map coefficients to an MTZ file. + The observed light intensity carries a positive, phase-blind contamination + ``sigma_alpha^2 |dF|^2`` from the spread of activation across crystals. Subtracting + the model's estimate of it and converting back to an amplitude gives a difference + amplitude that is comparable across datasets, which the raw one is not. + + ``DDF`` is the diagnostic that matters: smooth and featureless against resolution + means the correction is collinear with a scale or overall-B error and should be + distrusted; structure in it is the signal. + + The decontaminated amplitude goes through the dataset's own French-Wilson estimator, + on the **full** reflection list, because subtracting the variance term pushes weak + reflections negative and that is exactly where a naive ``sqrt(clamp(I, 0))`` is worst. Parameters ---------- - dc : DatasetCollection - mc : ModelCollection - scaler : CollectionScaler - filename : str - Output MTZ path. + mask : torch.Tensor + The dark-and-light validity intersection the writer uses; the returned columns + are already reduced to it. + fcalc_dark_full, fcalc_mixed_full : torch.Tensor + Scaled complex structure factors on the **full** HKL list. + weights, diff_Fobs, Fcalc_diff_amp, Fobs_dark, sig_dark : numpy.ndarray + Masked quantities the writer has already computed, reused so the corrected + columns are constructed exactly like their uncorrected counterparts. + + Returns + ------- + tuple + ``(columns, types)`` -- the values, and the MTZ type letter for each. Carrying + the type beside the value preserves the crystallographic column type. """ import numpy as np - import reciprocalspaceship as rs - from torchref import ReflectionData, Scaler - data_dark = dc[mc.dark_key] - data_light = dc["light"] - dark_model = mc.dark_model - mixed_model = mc["light"] - model_light = mc.base_models[1] + empty = ({}, {}) + if float(mc.sigma_alpha_sq) == 0.0: + return empty - hkl_all = data_dark.hkl - Fobs_dark_full, sig_dark_full = data_dark.get_corrected_data() - Fobs_light_full, sig_light_full = data_light.get_corrected_data() + data_light = dc["light"] + if data_light.I is None: + return empty - mask = data_dark.masks().to(torch.bool) & data_light.masks().to(torch.bool) - hkl = hkl_all[mask] - Fobs_dark_vals = Fobs_dark_full[mask] - Fobs_light_vals = Fobs_light_full[mask] - sig_dark_vals = sig_dark_full[mask] - sig_light_vals = sig_light_full[mask] - rfree_flags_masked = ( - data_light.rfree_flags[mask] if data_light.rfree_flags is not None else None + with torch.no_grad(): + # Full-size, so French-Wilson sees the reflection list it was fitted on. + delta_F_full = fcalc_mixed_full - fcalc_dark_full + variance_full = mc.sigma_alpha_sq * delta_F_full.abs() ** 2 + + I_light_full, sig_I_full = data_light.get_corrected_intensities() + I_corrected_full = I_light_full - variance_full + + # The retained estimator is fitted on the dataset's HKL list *as loaded*; + # joining a collection expands the dataset onto the common grid, so it can be + # the wrong length by then. Rebuild against the current list when that happens. + fw = data_light._FrenchWilson + if fw is None or len(fw.d_spacings) != len(I_corrected_full): + from torchref.base.french_wilson import FrenchWilson + + fw = FrenchWilson( + data_light.hkl, data_light.cell, data_light.spacegroup, verbose=0 + ) + F_corr_full, sig_F_corr_full = fw(I_corrected_full, sig_I_full) + + def _np(t): + return t[mask].detach().cpu().numpy() + + variance = _np(variance_full) + I_light = _np(I_light_full) + sig_I_light = _np(sig_I_full) + I_coherent = _np(fcalc_mixed_full.abs() ** 2) + F_corr = _np(F_corr_full) + sig_F_corr = _np(sig_F_corr_full) + + I_two_moment = I_coherent + variance + DF_corr = F_corr - Fobs_dark + DDF = DF_corr - diff_Fobs + sig_DF_corr = np.sqrt(sig_F_corr**2 + sig_dark**2) + + # Use the modulus of the complex vector difference so the corrected + # coefficient retains the phase rotation between the dark and light states. + F_corr_phased = torch.as_tensor( + F_corr, dtype=F_obs_dark_phased.real.dtype, device=F_obs_dark_phased.device + ) * torch.exp(1j * phi_mixed) + Fobs_diff_phased_corr = ( + torch.abs(F_corr_phased - F_obs_dark_phased).detach().cpu().numpy() ) - fractions = mixed_model.fractions.detach() - w_dark = fractions[0] - w_light = fractions[1] + amp_2_corr = (2 * Fobs_diff_phased_corr - Fcalc_diff_amp) * weights + amp_1_corr = (Fobs_diff_phased_corr - Fcalc_diff_amp) * weights + + # The sigma_alpha^2-aware weight, on the same normalisation as the inverse-variance + # weight the existing DED coefficients carry, so the two are directly comparable. + w_two_moment = sig_I_light**2 / np.maximum(sig_I_light**2 + variance, 1e-12) + + columns = { + # The corrected difference map, on the same dark phases as DELFWT. + "DELFWT_corr": DF_corr * weights, + "Fo_light_corr": F_corr, + "SIGFo_light_corr": sig_F_corr, + "DF_corr": DF_corr, + "SIGDF_corr": sig_DF_corr, + "DDF": DDF, + } + types = { + "DELFWT_corr": "F", + "Fo_light_corr": "F", + "SIGFo_light_corr": "Q", + "DF_corr": "F", + "SIGDF_corr": "Q", + "DDF": "F", + } + if all_columns: + columns.update({ + "Io_light": I_light, + "SIGIo_light": sig_I_light, + "Ic_light_coh": I_coherent, + "Ic_light_2mom": I_two_moment, + "IVAR_ALPHA": variance, + "W_2MOM": w_two_moment, + "2mDFop-DFc_corr": amp_2_corr, + "mDFop-DFc_corr": amp_1_corr, + }) + types.update({ + "Io_light": "J", + "SIGIo_light": "Q", + "Ic_light_coh": "J", + "Ic_light_2mom": "J", + "IVAR_ALPHA": "J", + "W_2MOM": "W", + "2mDFop-DFc_corr": "F", + "mDFop-DFc_corr": "F", + }) + return columns, types + + +def _difference_columns( + data_dark, + data_light, + mask, + hkl_np, + *, + Fobs_dark, + sig_dark, + Fobs_light, + sig_light, + Fcalc_dark, + phases_dark, + diff_Fobs, + sig_diff, + weight_columns, + kscale, +): + """The difference map's amplitudes, phases and weights. + + ``DF``/``SIGDF`` is the signed amplitude difference ``|Fo_light| - |Fo_dark|`` with + its propagated uncertainty, ``PHDELWT`` the **dark** model's phase it is carried on: + the isomorphous difference Fourier, and the construction ``torchref.validate-ded`` + correlates against. One weight column per registered scheme (``W_SD``, ``W_IVW``; + MTZ type ``W``, mean one) sits beside it, so any weighting is ``DF`` times a column + and reproducible from the file: ``torchref.mtz2map -csf DF -cw W_SD -cphi PHDELWT``. + ``KSCALE`` (type ``R``) is the scaler's multiplicative factor from model to observed + scale, so ``DF / KSCALE`` is in electrons and ``mtz2map --units electrons`` gives + e/A^3. + + This layer needs no light-state model: the amplitude is ``|Fo_light| - |Fo_dark|`` + and the phase comes from the dark model. Keeping the light state's model out is the + point -- a phased construction puts its phases into the observed amplitude, biasing + the map toward the very model the experiment is testing. + """ + import numpy as np + + n = len(hkl_np) + + def _flags(data): + if data.rfree_flags is None: + return np.ones(n, dtype=int) + return data.rfree_flags[mask].cpu().numpy().astype(int) + + columns = { + "H": hkl_np[:, 0], + "K": hkl_np[:, 1], + "L": hkl_np[:, 2], + "Fo_dark": Fobs_dark, + "SIGFo_dark": sig_dark, + "Fo_light": Fobs_light, + "SIGFo_light": sig_light, + "DF": diff_Fobs, + "SIGDF": sig_diff, + "PHDELWT": phases_dark, + **weight_columns, + "KSCALE": kscale, + "Fc_dark": Fcalc_dark, + # 1 = work, 0 = free. Both are kept: the two datasets can disagree, and + # picking one would silently report an R-free against the wrong test set. + "FreeR_flag_dark": _flags(data_dark), + "FreeR_flag_light": _flags(data_light), + } + types = { + "H": "H", + "K": "H", + "L": "H", + "Fo_dark": "F", + "SIGFo_dark": "Q", + "Fo_light": "F", + "SIGFo_light": "Q", + "DF": "F", + "SIGDF": "Q", + "PHDELWT": "P", + **{name: "W" for name in weight_columns}, + "KSCALE": "R", + "Fc_dark": "F", + "FreeR_flag_dark": "I", + "FreeR_flag_light": "I", + } + return columns, types + + +def _phasing_columns(mc, scaler, hkl_all, mask, *, fcalc_dark, Fobs_dark_vals, + Fobs_light_vals, phi_dark, Fcalc_dark, weights, + all_columns=False): + """Mixed-model amplitude and phase, and the phased difference residuals. + + Everything here needs the light state's model. ``FC``/``PHIC`` are the mixed model's + scaled amplitude and phase. + + Under ``all_columns`` the phased difference residual coefficients come too -- + ``(|Fo_light e^{i phi_mixed} - Fo_dark e^{i phi_dark}| - |dFc|) * w`` on + ``PHIC_diff``. These are a *different object* from the plain difference Fourier in + ``DELFWT``, not a refinement of it: the light state's model phases enter the observed + amplitude, so they are model-biased where ``DELFWT`` is not. They are kept because + they are informative once that is understood, and gated because the name alone does + not say it. + + Returns ``(columns, types, ctx)``. ``ctx`` carries the intermediates the + extrapolation and two-moment layers need, so nothing is computed twice. + """ + mixed_model = mc["light"] - # Compute Fcalc on full HKL then mask (scalers fitted on full datasets) with torch.no_grad(): - fcalc_dark = scaler.forward_mixed(dark_model(hkl_all), dark_model.fractions)[mask] - fcalc_mixed = scaler.forward_mixed(mixed_model(hkl_all), mixed_model.fractions)[mask] + fcalc_mixed_full = scaler.forward_mixed( + mixed_model(hkl_all), mixed_model.fractions + ) + fcalc_mixed = fcalc_mixed_full[mask] fcalc_diff = fcalc_mixed - fcalc_dark - phi_dark = torch.angle(fcalc_dark) phi_mixed = torch.angle(fcalc_mixed) - F_obs_dark_phased = Fobs_dark_vals * torch.exp(1j * phi_dark) F_obs_light_phased = Fobs_light_vals * torch.exp(1j * phi_mixed) - # --- Phase-aware extrapolation --- - F_light_extra = (F_obs_light_phased - w_dark * F_obs_dark_phased) / w_light - amp_light_extra = torch.abs(F_light_extra) - sig_light_extra = torch.sqrt(sig_light_vals**2 + w_dark**2 * sig_dark_vals**2) / w_light - - data_light_extra = ReflectionData.from_tensors( - hkl=hkl, F=amp_light_extra, F_sigma=sig_light_extra, - cell=data_light.cell, spacegroup=data_light.spacegroup, - rfree_flags=rfree_flags_masked, device=str(hkl.device), verbose=0, - ) - scaler_extra = Scaler(model_light, data_light_extra, device=hkl.device, verbose=-1) - scaler_extra.initialize().refine_lbfgs() - F_calc_extra = scaler_extra(model_light(hkl)) - - amp_extra = torch.abs(F_light_extra) - amp_light_calc = torch.abs(F_calc_extra) - phi_light_calc = torch.angle(F_calc_extra) - amp_2fofc_light = 2 * amp_extra - amp_light_calc - amp_fextfc = amp_extra - amp_light_calc - - # --- Classic (amplitude-only) extrapolation --- - amp_extra_classic = (Fobs_light_vals - w_dark * Fobs_dark_vals) / w_light - sig_extra_classic = sig_light_extra - - data_extra_classic = ReflectionData.from_tensors( - hkl=hkl, F=amp_extra_classic, F_sigma=sig_extra_classic, - cell=data_light.cell, spacegroup=data_light.spacegroup, - rfree_flags=rfree_flags_masked, device=str(hkl.device), verbose=0, - ) - scaler_classic = Scaler(model_light, data_extra_classic, device=hkl.device, verbose=-1) - scaler_classic.initialize().refine_lbfgs() - F_calc_classic = scaler_classic(model_light(hkl)) + Fcalc_light = torch.abs(fcalc_mixed).detach().cpu().numpy() + Fcalc_diff_amp = torch.abs(fcalc_diff).detach().cpu().numpy() - amp_light_calc_classic = torch.abs(F_calc_classic) - amp_2fofc_classic = 2 * amp_extra_classic - amp_light_calc_classic - amp_fofc_classic = amp_extra_classic - amp_light_calc_classic + columns = { + "FC": Fcalc_light, + "PHIC": phi_mixed.detach().rad2deg().cpu().numpy(), + } + types = {"FC": "F", "PHIC": "P"} + + if all_columns: + Fobs_diff_phased = torch.abs( + F_obs_light_phased - F_obs_dark_phased + ).detach().cpu().numpy() + columns.update( + { + "2mDFop-DFc": (2 * Fobs_diff_phased - Fcalc_diff_amp) * weights, + "mDFop-DFc": (Fobs_diff_phased - Fcalc_diff_amp) * weights, + "PHIC_diff": torch.angle(fcalc_diff).detach().rad2deg().cpu().numpy(), + "DFc": Fcalc_light - Fcalc_dark, + # This column holds the real modulus of the complex vector difference. + "DFc_phased": Fcalc_diff_amp, + } + ) + types.update({ + "2mDFop-DFc": "F", "mDFop-DFc": "F", "PHIC_diff": "P", + "DFc": "F", "DFc_phased": "F", + }) + + ctx = { + "fcalc_mixed_full": fcalc_mixed_full, + "phi_mixed": phi_mixed, + "F_obs_dark_phased": F_obs_dark_phased, + "F_obs_light_phased": F_obs_light_phased, + "Fcalc_diff_amp": Fcalc_diff_amp, + } + return columns, types, ctx + + +def _extrapolation_columns( + mc, + dc, + hkl, + *, + Fobs_dark_vals, + Fobs_light_vals, + sig_dark_vals, + sig_light_vals, + phi_dark, + ctx, + rfree_flags_masked, + all_columns=False, + verbose=1, + geometry=None, +): + """Extrapolated light-state amplitudes and the map to refine against. + + Three constructions of the same quantity, all needing the light model: + + ``FEXT`` (default, Bayes-shrunk) + The phase-aware amplitude shrunk toward ``Fo_dark`` by a per-reflection weight + ``w(h) = tau^2 / (tau^2 + sigma_ext^2(h))``, which quiets the weak and + high-resolution reflections where the extrapolation is noisiest. + ``FEXT_PHASED`` (``all_columns``) + The unshrunk phase-aware amplitude. + ``FEXT_SCALAR`` (``all_columns``) + ``(Fo_light - w_dark * Fo_dark) / w_light`` on amplitudes only. + + Each needs ``F_calc`` rescaled against *its own* amplitudes -- the three sets differ + in overall scale -- so each costs one LBFGS scale fit. Only the default one runs + unless ``all_columns`` is set. + + ``FWT``/``PHWT`` is ``2 * FEXT - Fc`` with the phase from the Bayes fit. The phase + matters: the scaler contributes one through ``f_sol``, so the three fits do not agree + and pairing these coefficients with another fit's phase would be wrong. + """ + from torchref import ReflectionData, Scaler + from torchref.base.metrics.rfactor import rfactor_work_free - # --- Empirical Bayes extrapolation (amplitude-only shrinkage) --- - F_ext_bayes, var_ext_bayes, w_shrinkage, tau_sq = ( + data_light = dc["light"] + fractions = mc["light"].fractions.detach() + w_dark, w_light = fractions[0], fractions[1] + + def _fit(amp, sig): + """Rescale the light model against one set of extrapolated amplitudes.""" + data = ReflectionData.from_tensors( + hkl=hkl, F=amp, F_sigma=sig, + cell=data_light.cell, spacegroup=data_light.spacegroup, + rfree_flags=rfree_flags_masked, device=str(hkl.device), verbose=0, + ) + sc = Scaler(mc.base_models[1], data, device=hkl.device, verbose=-1) + sc.initialize().refine_lbfgs() + return data, sc(mc.base_models[1](hkl)) + + # The phase-aware amplitude, and the one propagated sigma for this extrapolation. + F_light_extra = ( + ctx["F_obs_light_phased"] - w_dark * ctx["F_obs_dark_phased"] + ) / w_light + sig_light_extra = torch.sqrt( + sig_light_vals**2 + w_dark**2 * sig_dark_vals**2 + ) / w_light + + eps, dss = geometry if geometry is not None else (None, None) + F_ext_bayes_amp, var_ext_bayes, w_shrinkage, tau_sq = ( compute_bayes_extrapolated_amplitudes( - Fobs_dark_vals, Fobs_light_vals, - sig_dark_vals, sig_light_vals, - phi_dark, phi_mixed, w_light, + Fobs_dark_vals, + Fobs_light_vals, + sig_light_extra, + phi_dark, + ctx["phi_mixed"], + w_light, + epsilon=eps, + d_star_sq=dss, ) ) sig_ext_bayes = torch.sqrt(var_ext_bayes) - # Shrink |F_ext| towards |Fo_dark| — scalar operation, no phase interference - F_ext_amp_only = torch.abs(F_light_extra) - F_ext_bayes_amp = w_shrinkage * F_ext_amp_only + (1 - w_shrinkage) * Fobs_dark_vals + data_bayes, F_calc_bayes = _fit(F_ext_bayes_amp, sig_ext_bayes) + amp_calc_bayes = torch.abs(F_calc_bayes) - data_extra_bayes = ReflectionData.from_tensors( - hkl=hkl, F=F_ext_bayes_amp, F_sigma=sig_ext_bayes, - cell=data_light.cell, spacegroup=data_light.spacegroup, - rfree_flags=rfree_flags_masked, device=str(hkl.device), verbose=0, - ) - scaler_bayes = Scaler(model_light, data_extra_bayes, device=hkl.device, verbose=-1) - scaler_bayes.initialize().refine_lbfgs() - F_calc_bayes = scaler_bayes(model_light(hkl)) + def _np(t): + return t.detach().cpu().numpy() - amp_calc_bayes = torch.abs(F_calc_bayes) - amp_2fofc_bayes = 2 * F_ext_bayes_amp - amp_calc_bayes - amp_fofc_bayes = F_ext_bayes_amp - amp_calc_bayes + columns = { + "FEXT": _np(F_ext_bayes_amp), + "SIGFEXT": _np(sig_ext_bayes), + "FWT": _np(2 * F_ext_bayes_amp - amp_calc_bayes), + "PHWT": _np(torch.angle(F_calc_bayes).rad2deg()), + } + types = {"FEXT": "F", "SIGFEXT": "Q", "FWT": "F", "PHWT": "P"} + + if verbose > 0: + print(" Bayes extrapolation rfactors:", + rfactor_work_free(data_bayes, amp_calc_bayes)) + print(f" Bayes: tau^2 = {tau_sq:.4f}, " + f"mean w(h) = {w_shrinkage.mean().item():.3f}") + + if all_columns: + amp_phased = torch.abs(F_light_extra) + data_phased, F_calc_phased = _fit(amp_phased, sig_light_extra) + amp_calc_phased = torch.abs(F_calc_phased) + + amp_scalar = (Fobs_light_vals - w_dark * Fobs_dark_vals) / w_light + data_scalar, F_calc_scalar = _fit(amp_scalar, sig_light_extra) + amp_calc_scalar = torch.abs(F_calc_scalar) + + columns.update({ + "FEXT_PHASED": _np(amp_phased), + "SIGFEXT_PHASED": _np(sig_light_extra), + "2FEXT_PHASED-Fc": _np(2 * amp_phased - amp_calc_phased), + "FEXT_PHASED-Fc": _np(amp_phased - amp_calc_phased), + "PHFEXT_PHASED": _np(torch.angle(F_calc_phased).rad2deg()), + "FEXT_SCALAR": _np(amp_scalar), + "SIGFEXT_SCALAR": _np(sig_light_extra), + "2FEXT_SCALAR-Fc": _np(2 * amp_scalar - amp_calc_scalar), + "FEXT_SCALAR-Fc": _np(amp_scalar - amp_calc_scalar), + "PHFEXT_SCALAR": _np(torch.angle(F_calc_scalar).rad2deg()), + }) + types.update({ + "FEXT_PHASED": "F", "SIGFEXT_PHASED": "Q", + "2FEXT_PHASED-Fc": "F", "FEXT_PHASED-Fc": "F", + "PHFEXT_PHASED": "P", + "FEXT_SCALAR": "F", "SIGFEXT_SCALAR": "Q", + "2FEXT_SCALAR-Fc": "F", "FEXT_SCALAR-Fc": "F", + "PHFEXT_SCALAR": "P", + }) + if verbose > 0: + print(" Phase-aware extrapolation rfactors:", + rfactor_work_free(data_phased, amp_calc_phased)) + print(" Scalar extrapolation rfactors:", + rfactor_work_free(data_scalar, amp_calc_scalar)) + + diagnostics = { + "tau_sq": float(tau_sq), + "w_shrinkage_mean": float(w_shrinkage.mean().item()), + } + return columns, types, diagnostics - from torchref.base.metrics.rfactor import rfactor_work_free - def _extrapolation_rfactors(data, fcalc_scaled): - return rfactor_work_free(data, torch.abs(fcalc_scaled)) +_DIFFERENCE_DATASET_COLUMNS = ( + "DF", "SIGDF", "PHDELWT", "KSCALE", *WEIGHT_COLUMNS.values() +) + +# One history line per MTZ dataset, in the order they are written. MTZ history lines +# are at most 80 characters. +_MTZ_DATASET_HISTORY = { + "observed": "observed: Fo_dark, Fo_light and flags on the shared scale; Fc_dark", + "difference": "difference: DF/SIGDF on dark phases PHDELWT; weights W_SD, W_IVW", + "light_model": ( + "light_model: FC/PHIC, amplitude and phase of the mixed dark+light model" + ), + "extrapolated_light": ( + "extrapolated_light: FWT/PHWT = 2*FEXT - Fc, the extrapolated light map" + ), + "two_moment": "two_moment: difference columns with the two-moment correction", +} + + +def _annotate_mtz(filename, datasets): + """Group the columns of a written MTZ into named datasets and describe them. + + The column labels are left alone -- Coot's auto-open looks for ``FWT``/``PHWT`` by + label and ignores the dataset -- so the grouping is what tells a reader which map a + standard label belongs to: the column chooser shows ``/torchref//FWT``, + and ``gemmi mtz`` or ``mtzdump`` print the history lines. + + Parameters + ---------- + filename : str + MTZ written by reciprocalspaceship, with all columns in one dataset. Rewritten + in place. + datasets : dict[str, str] + Column label to dataset name. Unlisted columns, Miller indices included, stay + in ``observed``. + """ + import gemmi + + mtz = gemmi.read_mtz_file(filename) + base = mtz.datasets[0] + base.project_name = "torchref" + base.crystal_name = "torchref" + base.dataset_name = "observed" + + for name in dict.fromkeys(datasets.values()): + mtz.add_dataset(name) + # Filled in afterwards: ``add_dataset`` returns a reference into a vector that the + # next call may reallocate, so writes through it can be lost. + base = mtz.datasets[0] + ids = {} + for ds in mtz.datasets: + ds.project_name = base.project_name + ds.crystal_name = base.crystal_name + ds.cell = base.cell + ds.wavelength = base.wavelength + ids[ds.dataset_name] = ds.id + for col in mtz.columns: + if col.label in datasets: + col.dataset_id = ids[datasets[col.label]] + + mtz.history = [ + "torchref difference-refine map coefficients, grouped by dataset:", + *(_MTZ_DATASET_HISTORY[name] for name in ids), + ] + mtz.write_to_file(filename) + + +def write_results_mtz( + dc, + dark_model, + scaler, + filename, + *, + mc=None, + all_columns=False, + verbose=1, + ded_weight=DEFAULT_SCHEME, + sigma_d_config=None, +): + """Write the difference map, and map coefficients when a light model is given. + + The default output is the **difference map**: ``DF``/``SIGDF`` on the dark model's + phases ``PHDELWT``, with one mean-one weight column per registered scheme + (``W_SD``, ``W_IVW``) and the observed-to-model scale ``KSCALE``; see + :func:`_difference_columns`. ``ded_weight`` selects the scheme the model-phased + difference columns and the two-moment columns are weighted with. That needs no + light-state model, which is why ``mc`` is optional -- with a dark model alone this + writes a difference map and nothing else, and no scale fit is run beyond the one that + produced ``scaler``. + + Given ``mc``, the layers that need the light state follow: its amplitude and phase, + the extrapolated amplitudes, and the two-moment correction. ``all_columns`` adds the + alternatives within each layer -- see :func:`_phasing_columns` and + :func:`_extrapolation_columns` for what each contains and why it is gated. + + Column labels are the standard CCP4 ones, so Coot auto-opens ``FWT``/``PHWT`` -- + here the *extrapolated light-state* map, not a ``2mFo-DFc``. What each label means + is recorded in the file: the columns are grouped into MTZ datasets (``observed``, + ``difference``, ``light_model``, ``extrapolated_light``, ``two_moment``) with one + history line describing each. + + Parameters + ---------- + dc : DatasetCollection + Dark and light data, already inter-scaled by :meth:`DatasetCollection.scale`. + dark_model : Model or _SharedMixedModel + Supplies ``Fc_dark`` and the phases the difference map is carried on. + scaler : Scaler or CollectionScaler + Scales ``dark_model`` against the dark data. A ``CollectionScaler`` when ``mc`` + is given, a single-dataset ``Scaler`` otherwise. + mc : ModelCollection, optional + The dark+light collection. Absent means difference map only. + filename : str + Output MTZ path. + ded_weight : str, optional + Weight scheme for the model-phased and two-moment difference columns; one of + :data:`torchref.maps.ded_weights.SCHEMES`. + sigma_d_config : SigmaDConfig, optional + Exponent and shrinkage settings of the ``sigma_d`` scheme. + + Returns + ------- + dict + Diagnostics worth recording outside the file -- currently the Bayes shrinkage's + ``tau_sq`` and mean ``w(h)``, which say whether the default extrapolated map is + over-shrunk. Empty when no light model was given. + """ + import reciprocalspaceship as rs + + data_dark = dc[mc.dark_key] if mc is not None else dc["dark"] + data_light = dc["light"] + + hkl_all = data_dark.hkl + Fobs_dark_full, sig_dark_full = data_dark.get_corrected_data() + Fobs_light_full, sig_light_full = data_light.get_corrected_data() + + mask = data_dark.masks().to(torch.bool) & data_light.masks().to(torch.bool) + hkl = hkl_all[mask] + Fobs_dark_vals = Fobs_dark_full[mask] + Fobs_light_vals = Fobs_light_full[mask] + sig_dark_vals = sig_dark_full[mask] + sig_light_vals = sig_light_full[mask] + rfree_flags_masked = ( + data_light.rfree_flags[mask] if data_light.rfree_flags is not None else None + ) + + # Fcalc on the full HKL list then masked, because the scalers were fitted on the + # full datasets. ``forward_mixed`` exists only on CollectionScaler; the dark-only + # path carries a single-dataset Scaler, whose ``forward`` is the plain call. + with torch.no_grad(): + if hasattr(scaler, "forward_mixed"): + fcalc_dark_full = scaler.forward_mixed( + dark_model(hkl_all), dark_model.fractions + ) + else: + fcalc_dark_full = scaler(dark_model(hkl_all)) + fcalc_dark = fcalc_dark_full[mask] - print("Phase-aware extrapolation rfactors:", - _extrapolation_rfactors(data_light_extra, F_calc_extra)) - print("Classic extrapolation rfactors:", - _extrapolation_rfactors(data_extra_classic, F_calc_classic)) - print("Bayes extrapolation rfactors:", - _extrapolation_rfactors(data_extra_bayes, F_calc_bayes)) - print(f" Bayes: tau^2 = {tau_sq:.4f}, mean w(h) = {w_shrinkage.mean().item():.3f}") + phi_dark = torch.angle(fcalc_dark) - # --- Build MTZ --- hkl_np = hkl.cpu().numpy() Fobs_dark = Fobs_dark_vals.cpu().numpy() Fobs_light = Fobs_light_vals.cpu().numpy() sig_dark = sig_dark_vals.cpu().numpy() sig_light = sig_light_vals.cpu().numpy() - Fcalc_dark = torch.abs(fcalc_dark).detach().cpu().numpy() - Fcalc_light = torch.abs(fcalc_mixed).detach().cpu().numpy() - phases_dark = torch.angle(fcalc_dark).detach().rad2deg().cpu().numpy() - phases_mixed = torch.angle(fcalc_mixed).detach().rad2deg().cpu().numpy() - - Fcalc_diff_amp = torch.abs(fcalc_diff).detach().cpu().numpy() - Fcalc_diff_scalar = Fcalc_light - Fcalc_dark - phases_diff = torch.angle(fcalc_diff).detach().rad2deg().cpu().numpy() - - Fobs_diff_phased = torch.abs( - F_obs_light_phased - F_obs_dark_phased - ).detach().cpu().numpy() - - diff_Fobs = Fobs_light - Fobs_dark - sig_diff = (sig_dark**2 + sig_light**2) ** 0.5 - weights = 1 / sig_diff**2 - weights = weights / weights.mean() - weighted_diff_Fobs = diff_Fobs * weights - - amp_2DFoDFc = (2 * Fobs_diff_phased - Fcalc_diff_amp) * weights - amp_DFoDFc = (Fobs_diff_phased - Fcalc_diff_amp) * weights + phases_dark = phi_dark.detach().rad2deg().cpu().numpy() + + diff_t = Fobs_light_vals - Fobs_dark_vals + sig_diff_t = torch.sqrt(sig_dark_vals**2 + sig_light_vals**2) + all_w = all_ded_weights( + delta_obs=diff_t, + sigma_diff=sig_diff_t, + hkl=hkl, + cell=data_dark.cell, + spacegroup=data_dark.spacegroup, + f_dark=Fobs_dark_vals, + sigma_d_config=sigma_d_config, + ) + selected = all_w[ded_weight] + weights = selected.weights.detach().cpu().numpy() + diff_Fobs = diff_t.detach().cpu().numpy() + sig_diff = sig_diff_t.detach().cpu().numpy() + weight_columns = { + WEIGHT_COLUMNS[name]: all_w[name].weights.detach().cpu().numpy() + for name in WEIGHT_COLUMNS + } + kscale = scaler.multiplicative_scale()[mask].detach().cpu().numpy() + geometry = reflection_geometry( + hkl, data_dark.cell, data_dark.spacegroup, diff_t.device, diff_t.dtype + ) + sd_diag = { + k: v + for k, v in all_w["sigma_d"].diagnostics.items() + if k != "weight_sigma_d_raw" + } + diagnostics = { + "ded_weights": { + "scheme": ded_weight, + "applied": selected.applied, + "sigma_d": sd_diag, + } + } + if verbose > 0: + print(f" Difference weights: {ded_weight} (applied: {selected.applied})") + print( + f" sigma_D: gamma = {sd_diag['gamma']:.3f} ({sd_diag['gamma_reason']}), " + f"tau = {sd_diag['tau']:.3f}, shells = {sd_diag['n_shell']}, " + f"shells without difference power = {sd_diag['n_s2_clamped']}" + ) + if "fallback_reason" in sd_diag: + print(f" sigma_D fallback: {sd_diag['fallback_reason']}") + if verbose > 1 and not sd_diag["degenerate"]: + table = sd_diag["shells"] + print(" sigma_D shells: d(A) n B S2 Sigma_N") + for dss, n, b, s2, sn in zip( + table["d_star_sq"], + table["counts"], + table["B"], + table["S2"], + table["Sigma_N"], + ): + print(f" {dss ** -0.5:6.2f} {int(n):5d} {b:9.4f} {s2:9.4f} {sn:9.4f}") + + columns, types = _difference_columns( + data_dark, + data_light, + mask, + hkl_np, + Fobs_dark=Fobs_dark, + sig_dark=sig_dark, + Fobs_light=Fobs_light, + sig_light=sig_light, + Fcalc_dark=Fcalc_dark, + phases_dark=phases_dark, + diff_Fobs=diff_Fobs, + sig_diff=sig_diff, + weight_columns=weight_columns, + kscale=kscale, + ) + datasets = {name: "difference" for name in _DIFFERENCE_DATASET_COLUMNS} + + if mc is not None: + phase_cols, phase_types, ctx = _phasing_columns( + mc, scaler, hkl_all, mask, + fcalc_dark=fcalc_dark, Fobs_dark_vals=Fobs_dark_vals, + Fobs_light_vals=Fobs_light_vals, phi_dark=phi_dark, + Fcalc_dark=Fcalc_dark, weights=weights, all_columns=all_columns, + ) + columns.update(phase_cols) + datasets.update(dict.fromkeys(phase_cols, "light_model")) + types.update(phase_types) + + ext_cols, ext_types, ext_diagnostics = _extrapolation_columns( + mc, + dc, + hkl, + Fobs_dark_vals=Fobs_dark_vals, + Fobs_light_vals=Fobs_light_vals, + sig_dark_vals=sig_dark_vals, + sig_light_vals=sig_light_vals, + phi_dark=phi_dark, + ctx=ctx, + rfree_flags_masked=rfree_flags_masked, + all_columns=all_columns, + verbose=verbose, + geometry=geometry, + ) + diagnostics.update(ext_diagnostics) + columns.update(ext_cols) + datasets.update(dict.fromkeys(ext_cols, "extrapolated_light")) + types.update(ext_types) + + tm_cols, tm_types = _two_moment_columns( + mc, dc, mask, fcalc_dark_full, ctx["fcalc_mixed_full"], + weights=weights, diff_Fobs=diff_Fobs, + Fcalc_diff_amp=ctx["Fcalc_diff_amp"], Fobs_dark=Fobs_dark, + sig_dark=sig_dark, phi_mixed=ctx["phi_mixed"], + F_obs_dark_phased=ctx["F_obs_dark_phased"], + all_columns=all_columns, + ) + columns.update(tm_cols) + datasets.update(dict.fromkeys(tm_cols, "two_moment")) + types.update(tm_types) df = rs.DataSet( - { - "H": hkl_np[:, 0], "K": hkl_np[:, 1], "L": hkl_np[:, 2], - # Observed - "Fo_dark": Fobs_dark, "SIGFo_dark": sig_dark, - "Fo_light": Fobs_light, "SIGFo_light": sig_light, - # Differences - "DF": diff_Fobs, "SIGDF": sig_diff, "WDF": weighted_diff_Fobs, - # Calculated - "Fc_dark": Fcalc_dark, "Fc_light": Fcalc_light, - "DFc": Fcalc_diff_scalar, "DFc_complex": Fcalc_diff_amp, - # DED map coefficients (phase-aware) - "2mDFop-DFc": amp_2DFoDFc, "mDFop-DFc": amp_DFoDFc, - # Phases - "PHIC_dark": phases_dark, "PHIC_mixed": phases_mixed, - "PHIC_diff": phases_diff, - "PHIC_light": phi_light_calc.detach().rad2deg().cpu().numpy(), - # Phase-aware extrapolation - "Fextp": amp_extra.detach().cpu().numpy(), - "2Fextp-Fc": amp_2fofc_light.detach().cpu().numpy(), - "Fextp-Fc": amp_fextfc.detach().cpu().numpy(), - # Classic extrapolation - "Fextc": amp_extra_classic.detach().cpu().numpy(), - "SIGFextc": sig_extra_classic.detach().cpu().numpy(), - "2Fextc-Fc": amp_2fofc_classic.detach().cpu().numpy(), - "Fextc-Fc": amp_fofc_classic.detach().cpu().numpy(), - # Bayes extrapolation (amplitude-only shrinkage) - "Fextb": F_ext_bayes_amp.detach().cpu().numpy(), - "SIGFextb": sig_ext_bayes.detach().cpu().numpy(), - "2Fextb-Fc": amp_2fofc_bayes.detach().cpu().numpy(), - "Fextb-Fc": amp_fofc_bayes.detach().cpu().numpy(), - # R-free flags (1=work, 0=free) - "FreeR_flag_dark": ( - data_dark.rfree_flags[mask].cpu().numpy().astype(int) - if data_dark.rfree_flags is not None - else np.ones(len(hkl_np), dtype=int) - ), - "FreeR_flag_light": ( - data_light.rfree_flags[mask].cpu().numpy().astype(int) - if data_light.rfree_flags is not None - else np.ones(len(hkl_np), dtype=int) - ), - }, + columns, cell=data_dark.cell.data.cpu().tolist(), spacegroup=data_dark.spacegroup.hm, ) - - df[["H", "K", "L"]] = df[["H", "K", "L"]].astype("H") - f_cols = [ - "Fo_dark", "Fo_light", "DF", "WDF", - "Fc_dark", "Fc_light", "DFc", "DFc_complex", - "2mDFop-DFc", "mDFop-DFc", - "Fextp", "2Fextp-Fc", "Fextp-Fc", - "Fextc", "2Fextc-Fc", "Fextc-Fc", - "Fextb", "2Fextb-Fc", "Fextb-Fc", - ] - df[f_cols] = df[f_cols].astype("F") - sig_cols = ["SIGFo_dark", "SIGFo_light", "SIGDF", "SIGFextc", "SIGFextb"] - df[sig_cols] = df[sig_cols].astype("Q") - phase_cols = ["PHIC_dark", "PHIC_mixed", "PHIC_diff", "PHIC_light"] - df[phase_cols] = df[phase_cols].astype("P") - df["FreeR_flag_dark"] = df["FreeR_flag_dark"].astype("I") - df["FreeR_flag_light"] = df["FreeR_flag_light"].astype("I") + # Carry MTZ types with the values; infer_mtz_dtypes also checks the + # result using the canonical writer's rules. + missing = set(columns) - set(types) + if missing: + raise AssertionError(f"columns with no declared MTZ type: {sorted(missing)}") + for name, letter in types.items(): + df[name] = df[name].astype(letter) + df = df.infer_mtz_dtypes() df.set_index(["H", "K", "L"], inplace=True) df.write_mtz(filename) - print(f" Results MTZ written to {filename}") - print(f" w_dark={w_dark.item():.3f}, w_light={w_light.item():.3f}") + _annotate_mtz(filename, datasets) + + if verbose > 0: + print(f" Results MTZ written to {filename} ({len(columns)} columns)") + for name in dict.fromkeys(["observed", *datasets.values()]): + print(f" {_MTZ_DATASET_HISTORY[name]}") + if mc is not None: + fractions = mc["light"].fractions.detach() + print(f" w_dark={fractions[0].item():.3f}, " + f"w_light={fractions[1].item():.3f}") + + return diagnostics def optimize_lbfgs(state, parameters, max_iter, nsteps, n_clean, verbose): @@ -563,8 +1220,19 @@ def main(): add_outdir_arg(output, help="Output directory for refined structures and maps") add_output_format_args(output) add_metadata_args(output) + add_all_columns_arg(output) + add_ded_weight_args(output) refine = parser.add_argument_group("Refinement") + refine.add_argument( + "--difference-target", + choices=("difference", "difference_sd"), + default="difference", + help="Difference row the weight schedule drives: 'difference' is the Gaussian " + "under the measurement variance, 'difference_sd' centres on " + "alpha*dF_calc with the sigma_D unexplained power added to the variance " + "(default: difference).", + ) refine.add_argument( "--weight-schedule", type=str, default="5,3,2", help="Comma-separated difference-target weights applied in " @@ -596,6 +1264,28 @@ def main(): help="Weight for dark/light coordinate similarity restraint " "(0 to disable, default: 1.0)", ) + two_moment = parser.add_argument_group("Activation heterogeneity (two-moment model)") + two_moment.add_argument( + "--two-moment", action="store_true", default=False, + help="Fit merged intensities with |F(alpha)|^2 + sigma_alpha^2 |dF|^2, " + "which accounts for crystal-to-crystal spread in activation. " + "Requires I/SIGI columns in both reflection files.", + ) + two_moment.add_argument( + "--lambda-twin", type=float, default=0.0, + help="Activation dispersion as a fraction of its maximum, in [0, 1]: " + "sigma_alpha^2 = alpha (1 - alpha) * lambda. 0 (default) is the " + "coherent model and reproduces the amplitude-only result. Needs " + "--two-moment: the dispersion belongs in the predicted intensity, not " + "in a weight.", + ) + two_moment.add_argument( + "--refine-lambda-twin", action="store_true", default=False, + help="Refine --lambda-twin instead of holding it fixed. Off by default: " + "the sigma_alpha^2 term is smooth and positive, so it is collinear " + "with a scale or overall-B error and can absorb one.", + ) + refine.add_argument( "--similarity-alpha", type=float, default=2.0, help="Log prior odds for spike-and-slab similarity restraint. " @@ -617,6 +1307,24 @@ def main(): return 1 fractions = [1.0 - args.fraction, args.fraction] + if not (0.0 <= args.lambda_twin <= 1.0): + print( + f"Error: --lambda-twin must lie in [0, 1] (got {args.lambda_twin})", + file=sys.stderr, + ) + return 1 + if (args.lambda_twin > 0.0 or args.refine_lambda_twin) and not args.two_moment: + # Activation heterogeneity changes the predicted mean intensity. Treating + # it as measurement variance downweights the reflections carrying the signal. + print( + "Error: --lambda-twin needs --two-moment. The dispersion enters the predicted " + "intensity, not a weight: as a variance it down-weights the reflections whose " + "difference signal is largest, which measurably worsens the recovered " + "displacement.", + file=sys.stderr, + ) + return 1 + # --- Parse weight schedule --- try: weight_schedule = [float(x) for x in args.weight_schedule.split(",")] @@ -629,7 +1337,10 @@ def main(): # --- Parse and merge target weights --- target_weights = dict(DEFAULT_TARGET_WEIGHTS) - target_weights["xray/difference"] = weight_schedule[0] + difference_key = f"xray/{args.difference_target}" + target_weights["xray/difference"] = 0.0 + target_weights["xray/difference_sd"] = 0.0 + target_weights[difference_key] = weight_schedule[0] target_weights["similarity"] = args.similarity_weight target_weights, err = parse_weights(args.weights, defaults=target_weights) if err: @@ -666,6 +1377,11 @@ def main(): print(f"Light data: {args.light_structure_factor}") frac_mode = "refinable" if args.refine_fractions else "frozen" print(f"Fractions: dark={fractions[0]}, light={fractions[1]} ({frac_mode})") + if args.two_moment: + lam_mode = "refinable" if args.refine_lambda_twin else "fixed" + print( + f"Activation spread: lambda_twin={args.lambda_twin} ({lam_mode})" + ) print(f"Output: {outdir}") print(f"Device: {device}") if args.dmin: @@ -747,8 +1463,32 @@ def main(): sys.stdout.flush() # --- Setup targets --- - state = setup_loss_state(dc, mc, scaler, target_weights, device, - similarity_alpha=args.similarity_alpha) + if args.two_moment: + missing = [k for k in dc.keys() if dc[k].I is None] + if missing: + print( + f"Error: --two-moment needs I/SIGI columns, but {missing} carry " + f"only amplitudes. Converting back with F**2 would reintroduce the " + f"French-Wilson distortion the intensity model exists to avoid.", + file=sys.stderr, + ) + return 1 + + # Set unconditionally: a non-zero dispersion also drives the difference target's + # per-reflection weighting, which needs no intensity data. + mc.set_lambda_twin(args.lambda_twin, refinable=args.refine_lambda_twin) + + state = setup_loss_state( + dc, + mc, + scaler, + target_weights, + device, + similarity_alpha=args.similarity_alpha, + two_moment=args.two_moment, + difference_target=args.difference_target, + sigma_d_config=sigma_d_config_from_args(args), + ) if args.verbose > 0: print("Initial loss breakdown:") @@ -761,10 +1501,14 @@ def main(): if args.refine_fractions: params = list(itertools.chain( - model_light.parameters(), [mixed.fraction_params] + model_light.parameters(), mc.fraction_parameters() )) else: params = list(model_light.parameters()) + if args.refine_lambda_twin: + # fraction_parameters() carries lambda once it is refinable; take only + # that, since the fractions themselves stay frozen here. + params.append(mc._lambda_logit) fraction_history = [] if args.refine_fractions: @@ -784,7 +1528,7 @@ def main(): ) sys.stdout.flush() - state.set_weights({"xray/difference": t_weight}) + state.set_weights({difference_key: t_weight}) optimize_lbfgs( state, params, max_iter=args.max_iter, @@ -816,14 +1560,17 @@ def main(): sys.stdout.flush() # --- Final statistics --- + # Computed unconditionally: the deposition metadata and the results MTZ both carry + # these, so they are not a reporting-only quantity. + r_work_d, r_free_d = compute_rfactors(dark, data_dark, scaler) + r_work_l, r_free_l = compute_rfactors(mixed, data_light, scaler) + r_work_dl, r_free_dl = compute_rfactors(dark, data_light, scaler) + if args.verbose > 0: print() print("=" * 72) print("Refinement complete") print("=" * 72) - r_work_d, r_free_d = compute_rfactors(dark, data_dark, scaler) - r_work_l, r_free_l = compute_rfactors(mixed, data_light, scaler) - r_work_dl, r_free_dl = compute_rfactors(dark, data_light, scaler) print(f" Final R-factor (dark vs dark data): R_work={r_work_d:.4f} R_free={r_free_d:.4f}") print(f" Final R-factor (mixed vs light data): R_work={r_work_l:.4f} R_free={r_free_l:.4f}") print(f" Final R-factor (dark vs light data): R_work={r_work_dl:.4f} R_free={r_free_dl:.4f}") @@ -892,7 +1639,7 @@ def _build_metadata(model, data, r_work, r_free): meta.n_atoms_solvent = int((pdb["ATOM"] == "HETATM").sum()) # Geometry deviations - if model.initialized and model._restraints is not None: + if model.ctx.initialized and model._restraints is not None: restraints = model.restraints with torch.no_grad(): if hasattr(restraints, "bond_deviations"): @@ -940,6 +1687,7 @@ def _build_metadata(model, data, r_work, r_free): if not has_altloc_dark and not has_altloc_light: import pandas as pd + from torchref import __version__ from torchref.io.metadata import RefinementMetadata @@ -1033,7 +1781,17 @@ def _mtz_to_cif(mtz_path, cif_path): print(f" Dark SF written to {dark_sf_mtz}, {dark_sf_cif}") print(f" Light SF written to {light_sf_mtz}, {light_sf_cif}") - write_results_mtz(dc, mc, scaler, diff_mtz_out) + map_diagnostics = write_results_mtz( + dc, + mc.dark_model, + scaler, + diff_mtz_out, + mc=mc, + all_columns=args.all_columns, + verbose=args.verbose, + ded_weight=args.ded_weight, + sigma_d_config=sigma_d_config_from_args(args), + ) # --- JSON summary --- summary = { @@ -1046,6 +1804,7 @@ def _mtz_to_cif(mtz_path, cif_path): "cif": args.cif, "dmin": args.dmin, }, + "dataset_scaling": dc.scaling_metrics, "parameters": { "weight_schedule": weight_schedule, "n_cycles": args.n_cycles, @@ -1054,15 +1813,23 @@ def _mtz_to_cif(mtz_path, cif_path): "weights": target_weights, }, "results": { - "r_factor_dark": dict(zip( - ["r_work", "r_free"], - compute_rfactors(dark, data_dark, scaler), - )), - "r_factor_light": dict(zip( - ["r_work", "r_free"], - compute_rfactors(mixed, data_light, scaler), - )), + "r_factor_dark": dict( + zip( + ["r_work", "r_free"], + compute_rfactors(dark, data_dark, scaler), + ) + ), + "r_factor_light": dict( + zip( + ["r_work", "r_free"], + compute_rfactors(mixed, data_light, scaler), + ) + ), "fractions": mixed.fractions.detach().cpu().tolist(), + "alpha_mean": float(mc.alpha_mean), + "lambda_twin": float(mc.lambda_twin), + "sigma_alpha_sq": float(mc.sigma_alpha_sq), + **map_diagnostics, }, "output_files": { "dark_pdb": dark_pdb_out, diff --git a/torchref/cli/difference_map.py b/torchref/cli/difference_map.py new file mode 100644 index 00000000..54f4bc6c --- /dev/null +++ b/torchref/cli/difference_map.py @@ -0,0 +1,244 @@ +#!/usr/bin/env python3 -u + +"""Difference and extrapolated map coefficients from dark/light data. + +Uses the ``torchref.difference-refine`` pipeline but performs **no refinement**: the input +models are used as-is. + +The default output is the difference map: the amplitude difference +``|Fo_light| - |Fo_dark|`` as ``DF``/``SIGDF`` carried on the **dark** model's phases +``PHDELWT``, with the per-reflection weights of every registered scheme beside it as +``W_IVW`` (inverse variance, the default) and ``W_SD`` (sigma_D Wiener weight), and the +observed-to-model scale ``KSCALE``. Build the map with +``torchref.mtz2map -csf DF -cw W_IVW -cphi PHDELWT`` (``--units electrons`` for e/A^3). +That needs no light-state model, so ``-lm`` is optional. It is also deliberately not a +*phased* difference map: putting the light state's model phases into the observed +amplitude biases the map toward the very model the experiment is testing. + +Given ``-lm``, the light state's amplitude and phase and the extrapolated map follow. +``--all-columns`` adds the alternative constructions of both. + +Examples +-------- +:: + + # difference map only -- no light model, no fraction needed + torchref.difference-map \\ + -dm dark.pdb -dsf dark.mtz -lsf light.mtz -o results.mtz + + torchref.difference-map \\ + -dm dark.pdb -lm light.pdb -dsf dark.mtz -lsf light.mtz \\ + --fraction 0.37 --dmin 1.7 --cif ligand.cif -o results.mtz +""" + +import argparse +import sys +from pathlib import Path + +import torch + +from torchref.cli._common import ( + add_all_columns_arg, + add_ded_weight_args, + add_dual_model_args, + add_dmin_arg, + add_general_args, + add_output_arg, + build_dual_column_names, + configure_unbuffered_output, + register_timing, + parse_device_str, + sigma_d_config_from_args, + validate_cif_files, + validate_files, +) + +configure_unbuffered_output() + + +def main(): + """Entry point for ``torchref.difference-map``; returns the exit code.""" + parser = argparse.ArgumentParser( + description="Compute difference and extrapolated map coefficients " + "(no refinement).", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # difference map only -- needs no light model and no fraction + torchref.difference-map \\ + -dm dark.pdb \\ + -dsf dark.mtz -lsf light.mtz -o results.mtz + + torchref.difference-map \\ + -dm dark.pdb -lm light.pdb \\ + -dsf dark.mtz -lsf light.mtz \\ + --fraction 0.37 -o results.mtz + + torchref.difference-map \\ + -dm dark.pdb -lm light.pdb \\ + -dsf dark.mtz -lsf light.mtz \\ + --fraction 0.37 --dmin 1.7 --all-columns -o results.mtz + """, + ) + + add_dual_model_args(parser, fraction_required=False, light_model_required=False) + + output = parser.add_argument_group("Output") + add_output_arg(output, help="Output MTZ file path (e.g. results.mtz)") + add_all_columns_arg(output) + add_ded_weight_args(output) + + res = parser.add_argument_group("Resolution") + add_dmin_arg(res) + + add_general_args(parser) + + args = parser.parse_args() + + register_timing() + + has_light_model = args.light_model is not None + if has_light_model: + if args.fraction is None: + parser.error( + "--fraction is required with -lm/--light-model: the extrapolation " + "divides by the light-state occupancy." + ) + fractions = [1.0 - args.fraction, args.fraction] + else: + fractions = None + if args.fraction is not None: + parser.error( + "--fraction needs -lm/--light-model. Without a light model only the " + "weighted difference map is written, and it carries no occupancy: the " + "amplitude is |Fo_light| - |Fo_dark| and the weight comes from the " + "sigmas." + ) + + to_check = [ + (args.dark_model, "dark model"), + (args.dark_structure_factor, "dark structure factor"), + (args.light_structure_factor, "light structure factor"), + ] + if has_light_model: + to_check.insert(1, (args.light_model, "light model")) + if validate_files(to_check): + return 1 + + if validate_cif_files(args.cif): + return 1 + + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + + device = parse_device_str(args.device) + + if args.verbose > 0: + print("=" * 72) + print("TorchRef Difference Map") + print("=" * 72) + print(f"Dark model: {args.dark_model}") + if has_light_model: + print(f"Light model: {args.light_model}") + print(f"Fraction: light={args.fraction} " + f"(dark={1.0 - args.fraction})") + else: + print("Light model: (none) -- difference map only") + print(f"Dark SF: {args.dark_structure_factor}") + print(f"Light SF: {args.light_structure_factor}") + print(f"Output: {args.output}") + print(f"Device: {device}") + if args.dmin: + print(f"Resolution cutoff: {args.dmin:.2f} A") + if args.cif: + print(f"CIF restraints: {', '.join(args.cif)}") + if args.all_columns: + print("Columns: all") + print("=" * 72) + print() + sys.stdout.flush() + + from torchref.cli.collection_difference_refine import ( + compute_rfactors, + setup_dark_only, + setup_model_collection, + setup_dataset_collection, + setup_scaler, + write_results_mtz, + ) + + d_min = args.dmin if args.dmin is not None else 1.0 + + if args.verbose > 0: + print("Loading reflection data...") + sys.stdout.flush() + + col_dark, col_light = build_dual_column_names(args) + + dc = setup_dataset_collection( + args.dark_structure_factor, args.light_structure_factor, args.dmin, device, + column_names_dark=col_dark, column_names_light=col_light, + ) + + if args.verbose > 0: + print("Setting up models...") + sys.stdout.flush() + + if has_light_model: + mc = setup_model_collection( + args.dark_model, args.light_model, fractions, + args.cif, d_min, device, args.verbose, + ) + mc["light"].freeze_fractions() + + if args.verbose > 0: + print("Setting up joint scaler...") + sys.stdout.flush() + scaler = setup_scaler(dc, mc, device, args.verbose) + dark_model = mc.dark_model + + if args.verbose > 0: + r_work_d, r_free_d = compute_rfactors(dark_model, dc["dark"], scaler) + r_work_l, r_free_l = compute_rfactors(mc["light"], dc["light"], scaler) + print(f" R-factor (dark): R_work={r_work_d:.4f} R_free={r_free_d:.4f}") + print(f" R-factor (mixed): R_work={r_work_l:.4f} R_free={r_free_l:.4f}") + print() + sys.stdout.flush() + else: + mc = None + dark_model, scaler = setup_dark_only( + args.dark_model, dc, args.cif, d_min, device, args.verbose, + ) + if args.verbose > 0: + r_work_d, r_free_d = compute_rfactors(dark_model, dc["dark"], scaler) + print(f" R-factor (dark): R_work={r_work_d:.4f} R_free={r_free_d:.4f}") + print() + sys.stdout.flush() + + if args.verbose > 0: + print("Computing map coefficients...") + sys.stdout.flush() + + with torch.no_grad(): + write_results_mtz( + dc, + dark_model, + scaler, + str(out_path), + mc=mc, + all_columns=args.all_columns, + verbose=args.verbose, + ded_weight=args.ded_weight, + sigma_d_config=sigma_d_config_from_args(args), + ) + + if args.verbose > 0: + print() + print("Done.") + sys.stdout.flush() + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/torchref/cli/mtz2map.py b/torchref/cli/mtz2map.py index b02f459f..e98fe12b 100644 --- a/torchref/cli/mtz2map.py +++ b/torchref/cli/mtz2map.py @@ -18,6 +18,7 @@ import numpy as np import torch +from torchref.config import get_float_dtype from torchref.cli._common import ( add_general_args, add_resolution_args, @@ -58,6 +59,25 @@ def main(): metavar="COL", help="Column name for phases in degrees (e.g. PHWT, PHDELWT, PH2FOFCWT).", ) + inp.add_argument( + "-cw", + "--column-weight", + default=None, + type=str, + metavar="COL", + help="Weight column multiplied into the amplitudes before the FFT " + "(e.g. W_SD, W_IVW from torchref.difference-map). Default: none.", + ) + inp.add_argument( + "-ck", + "--column-scale", + default=None, + type=str, + metavar="COL", + help="Per-reflection observed-to-model scale factor the amplitudes are divided " + "by for --units electrons (e.g. KSCALE from torchref.difference-map). " + "Default: KSCALE when the file has it.", + ) output = parser.add_argument_group("Output") output.add_argument( @@ -73,15 +93,24 @@ def main(): metavar=("NX", "NY", "NZ"), help="Override grid dimensions. Default: auto from cell and resolution.", ) + mapopts.add_argument( + "--units", + type=str, + choices=["sigma", "electrons", "raw"], + default=None, + help="Map units. 'sigma': zero mean and unit standard deviation (default). " + "'electrons': electrons per cubic Angstrom, sum_h F(h) exp(-2 pi i h.x) / V " + "with F divided by the --column-scale factor. 'raw': the plain FFT with the " + "1/N normalisation, no rescaling.", + ) mapopts.add_argument( "-n", "--normalize", type=str, - choices=['True', 'False'], - default='True', - help="Normalize amplitudes to unit variance. Accepts only the literal " - "strings 'True' or 'False' (case-sensitive); pass '-n False' to disable. " - "Default: True.", + choices=["True", "False"], + default=None, + help="Deprecated alias: '-n True' is '--units sigma', '-n False' is " + "'--units raw'.", ) res = parser.add_argument_group("Resolution") @@ -105,7 +134,24 @@ def main(): mtz = rs.read_mtz(args.structure_factor) available = list(mtz.columns) - normalize = args.normalize == 'True' + if args.units is not None and args.normalize is not None: + print("Error: --units and --normalize cannot both be given", file=sys.stderr) + sys.exit(1) + if args.units is not None: + units = args.units + elif args.normalize is not None: + units = "sigma" if args.normalize == "True" else "raw" + else: + units = "sigma" + scale_column = args.column_scale + if scale_column is None and units == "electrons" and "KSCALE" in available: + scale_column = "KSCALE" + if units == "electrons" and scale_column is None: + print( + "Error: --units electrons needs --column-scale (no KSCALE column found).", + file=sys.stderr, + ) + sys.exit(1) if args.column_structure_factor not in available: print( @@ -121,6 +167,14 @@ def main(): file=sys.stderr, ) sys.exit(1) + for label, col in (("weight", args.column_weight), ("scale", scale_column)): + if col is not None and col not in available: + print( + f"Error: {label} column '{col}' not found.\n" + f"Available columns: {available}", + file=sys.stderr, + ) + sys.exit(1) # Extract cell and spacegroup cell = np.array( @@ -140,9 +194,23 @@ def main(): hkl = df[["H", "K", "L"]].to_numpy().astype(np.int32) amplitudes = df[args.column_structure_factor].to_numpy().astype(np.float32) phases_deg = df[args.column_phase].to_numpy().astype(np.float32) + valid = np.isfinite(amplitudes) & np.isfinite(phases_deg) + if args.column_weight is not None: + weights = df[args.column_weight].to_numpy().astype(np.float32) + valid &= np.isfinite(weights) + amplitudes = amplitudes * weights + if args.verbose >= 1: + print(f" Weights: {args.column_weight} (mean {np.nanmean(weights):.3f})") + if units == "electrons": + kscale = df[scale_column].to_numpy().astype(np.float32) + valid &= np.isfinite(kscale) & (kscale > 0) + # Observed amplitudes carry the scaler's overall scale, B and anisotropy; + # dividing by that factor returns them to electrons. + amplitudes = amplitudes / np.where(valid, kscale, 1.0) + if args.verbose >= 1: + print(f" Absolute scale: dividing by {scale_column}") # Drop NaN reflections - valid = np.isfinite(amplitudes) & np.isfinite(phases_deg) if not valid.all(): n_drop = (~valid).sum() if args.verbose >= 1: @@ -178,15 +246,16 @@ def main(): f"{d_spacings.max():.2f} - {d_spacings.min():.2f} A") # --- Convert to torch --- - hkl_t = torch.tensor(hkl, dtype=torch.int32, device=device) - amp_t = torch.tensor(amplitudes, dtype=torch.float32, device=device) - phi_t = torch.tensor(phases_deg, dtype=torch.float32, device=device) * (np.pi / 180.0) + hkl_t = torch.tensor(hkl, dtype=torch.int32, device=device) # dtype-ok: hkl Miller indices fed to symmetry expand; fixed int32 crystallographic representation + amp_t = torch.tensor(amplitudes, dtype=get_float_dtype(), device=device) + phi_t = torch.tensor(phases_deg, dtype=get_float_dtype(), device=device) * (np.pi / 180.0) # --- Expand to P1 --- - from torchref.symmetry.reciprocal_symmetry import expand_hkl + from torchref.symmetry import Cell, SpaceGroup - hkl_p1, orig_idx, phase_shifts = expand_hkl( - hkl_t, spacegroup, include_friedel=False, remove_absences=True + sg = SpaceGroup(spacegroup) + hkl_p1, orig_idx, phase_shifts = sg.expand_hkl( + hkl_t, include_friedel=False, remove_absences=True ) amp_p1 = amp_t[orig_idx] @@ -199,13 +268,11 @@ def main(): coefficients = amp_p1 * torch.exp(1j * phi_p1) # --- Grid size --- - from torchref.symmetry.grid_utils import calculate_optimal_grid_size - if args.gridsize is not None: gridsize = tuple(args.gridsize) else: max_res = float(d_spacings.min()) - gridsize = calculate_optimal_grid_size(cell, max_res, spacegroup) + gridsize = sg.optimal_grid_size(Cell(cell), max_res) if args.verbose >= 1: print(f" Grid size: {gridsize[0]} x {gridsize[1]} x {gridsize[2]}") @@ -215,12 +282,17 @@ def main(): grid = place_on_grid(hkl_p1, coefficients, gridsize, enforce_hermitian=True) - # FFT to real space: rho(r) = sum_h F(h) * exp(-2*pi*i * h.r) + # FFT to real space with the 1/N normalisation: rho_raw(r) = (1/N) sum_h F(h) exp(-2 pi i h.r) real_map = torch.fft.fftn(grid, dim=(0, 1, 2), norm="forward").real - if normalize: + if units == "sigma": real_map = (real_map - real_map.mean()) / real_map.std() - + elif units == "electrons": + # rho(r) = (1/V) sum_h F(h) exp(-2 pi i h.r): undo the 1/N and divide by the + # cell volume, so the map is in electrons per cubic Angstrom. + volume = Cell(cell, device=device).volume.to(real_map.dtype) + real_map = real_map * (real_map.numel() / volume) + # --- Write output --- from torchref.io.cif import write_map @@ -229,6 +301,7 @@ def main(): if args.verbose >= 1: print(f" Written: {args.output}") sigma = float(real_map.std()) + print(f" Units: {units}") print(f" Map sigma: {sigma:.4f}") diff --git a/torchref/cli/phased_difference_map.py b/torchref/cli/phased_difference_map.py index 09648de0..3ad34cf7 100644 --- a/torchref/cli/phased_difference_map.py +++ b/torchref/cli/phased_difference_map.py @@ -1,184 +1,8 @@ -#!/usr/bin/env python3 -u +"""Run the difference-map CLI through its phased-difference-map alias.""" -"""Phased difference and extrapolated map coefficients from dark/light data. - -Uses the ``torchref.difference-refine`` pipeline but performs **no refinement**: the input -models are used as-is to compute phases, scale factors and every flavour of extrapolated -amplitude, and the result is one MTZ holding the observed, calculated, difference and -extrapolated columns. - -Examples --------- -:: - - torchref.phased-difference-map \\ - -dm dark.pdb -lm light.pdb -dsf dark.mtz -lsf light.mtz \\ - --fraction 0.37 --dmin 1.7 --cif ligand.cif -o results.mtz -""" - -import argparse -import sys -from pathlib import Path - -import torch - -from torchref.cli._common import ( - add_dual_model_args, - add_dmin_arg, - add_general_args, - add_output_arg, - build_dual_column_names, - configure_unbuffered_output, - register_timing, - parse_device_str, - validate_cif_files, - validate_files, -) - -configure_unbuffered_output() - - -def main(): - """Entry point for ``torchref.phased-difference-map``; returns the exit code.""" - parser = argparse.ArgumentParser( - description="Compute phased difference and extrapolated map " - "coefficients (no refinement).", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=""" -Examples: - torchref.phased-difference-map \\ - -dm dark.pdb -lm light.pdb \\ - -dsf dark.mtz -lsf light.mtz \\ - --fraction 0.37 -o results.mtz - - torchref.phased-difference-map \\ - -dm dark.pdb -lm light.pdb \\ - -dsf dark.mtz -lsf light.mtz \\ - --fraction 0.37 --dmin 1.7 -o results.mtz - """, - ) - - # --- Input files (creates "Input files" and "Column selection" groups) --- - add_dual_model_args(parser) - - output = parser.add_argument_group("Output") - add_output_arg(output, help="Output MTZ file path (e.g. results.mtz)") - - res = parser.add_argument_group("Resolution") - add_dmin_arg(res) - - add_general_args(parser) - - args = parser.parse_args() - - register_timing() - - # --- Parse fraction --- - fractions = [1.0 - args.fraction, args.fraction] - - # --- Validate input files --- - if validate_files([ - (args.dark_model, "dark model"), - (args.light_model, "light model"), - (args.dark_structure_factor, "dark structure factor"), - (args.light_structure_factor, "light structure factor"), - ]): - return 1 - - if validate_cif_files(args.cif): - return 1 - - # Ensure output directory exists - out_path = Path(args.output) - out_path.parent.mkdir(parents=True, exist_ok=True) - - # --- Device --- - device = parse_device_str(args.device) - - # --- Header --- - if args.verbose > 0: - print("=" * 72) - print("TorchRef Phased Difference Map") - print("=" * 72) - print(f"Dark model: {args.dark_model}") - print(f"Light model: {args.light_model}") - print(f"Dark SF: {args.dark_structure_factor}") - print(f"Light SF: {args.light_structure_factor}") - print(f"Fraction: light={args.fraction} (dark={1.0 - args.fraction})") - print(f"Output: {args.output}") - print(f"Device: {device}") - if args.dmin: - print(f"Resolution cutoff: {args.dmin:.2f} A") - if args.cif: - print(f"CIF restraints: {', '.join(args.cif)}") - print("=" * 72) - print() - sys.stdout.flush() - - from torchref.cli.collection_difference_refine import ( - compute_rfactors, - setup_model_collection, - setup_dataset_collection, - setup_scaler, - write_results_mtz, - ) - - # --- Resolution --- - d_min = args.dmin if args.dmin is not None else 1.0 - - # --- Setup models --- - if args.verbose > 0: - print("Setting up models...") - sys.stdout.flush() - - mc = setup_model_collection( - args.dark_model, args.light_model, fractions, - args.cif, d_min, device, args.verbose, - ) - mc["light"].freeze_fractions() - - # --- Load data --- - if args.verbose > 0: - print("Loading reflection data...") - sys.stdout.flush() - - col_dark, col_light = build_dual_column_names(args) - - dc = setup_dataset_collection( - args.dark_structure_factor, args.light_structure_factor, args.dmin, device, - column_names_dark=col_dark, column_names_light=col_light, - ) - - # --- Scale --- - if args.verbose > 0: - print("Setting up joint scaler...") - sys.stdout.flush() - - scaler = setup_scaler(dc, mc, device, args.verbose) - - if args.verbose > 0: - r_work_d, r_free_d = compute_rfactors(mc.dark_model, dc["dark"], scaler) - r_work_l, r_free_l = compute_rfactors(mc["light"], dc["light"], scaler) - print(f" R-factor (dark): R_work={r_work_d:.4f} R_free={r_free_d:.4f}") - print(f" R-factor (mixed): R_work={r_work_l:.4f} R_free={r_free_l:.4f}") - print() - sys.stdout.flush() - - # --- Write MTZ --- - if args.verbose > 0: - print("Computing map coefficients...") - sys.stdout.flush() - - with torch.no_grad(): - write_results_mtz(dc, mc, scaler, str(out_path)) - - if args.verbose > 0: - print() - print("Done.") - sys.stdout.flush() - - return 0 +from torchref.cli.difference_map import main +__all__ = ["main"] if __name__ == "__main__": - sys.exit(main()) + raise SystemExit(main()) diff --git a/torchref/cli/refine.py b/torchref/cli/refine.py index 238db84c..7aed867c 100644 --- a/torchref/cli/refine.py +++ b/torchref/cli/refine.py @@ -114,6 +114,20 @@ def main(): refine_group = parser.add_argument_group("Refinement") add_n_cycles_arg(refine_group) + refine_group.add_argument( + "--add-hydrogens", + action="store_true", + help="Generate missing hydrogens when loading the model (default: off). " + "Hydrogens already present in the input are retained either way.", + ) + refine_group.add_argument( + "--hydrogens-in-xray", + dest="hydrogens_in_xray", + action=argparse.BooleanOptionalAction, + default=True, + help="Include hydrogen atoms in the structure-factor calculation (default: on). " + "--no-hydrogens-in-xray keeps them in the restraints only.", + ) refine_group.add_argument( "--mode", type=str, @@ -248,12 +262,23 @@ def main(): print(f"Refinement mode: {args.mode}") print(f"X-ray target: {args.xray_mode}") print(f"Refinement cycles: {args.n_cycles}") + print(f"Add hydrogens: {'on' if args.add_hydrogens else 'off'}") + print(f"Hydrogens in Fcalc: {'on' if args.hydrogens_in_xray else 'off'}") if args.with_rigid_body: print(f"Rigid-body step: on (iterations/cutoff = {args.rigid_body_iter})") print(f"Device: {args.device}") if args.dmin: print(f"Resolution cutoff: {args.dmin:.2f} A") adp_line = f"ADP mode: {args.adp_mode}" + if args.adp_mode == "field_aniso" and args.adp_mode_set: + adp_line += f" ({args.adp_mode_set})" + if args.adp_mode in ("field", "field_aniso"): + adp_line += ( + f", {args.adp_nodes} nodes" + if args.adp_nodes + else f", sized at {args.reflections_per_adp_parameter:g} " + "work reflections per parameter" + ) if args.adp_mode == "anisotropic": adp_line += ( " (selection: " @@ -290,8 +315,13 @@ def main(): scale_target=args.scale_target, **_sigma_a_kwargs(args), adp_mode=args.adp_mode, + adp_mode_set=args.adp_mode_set, + n_nodes=args.adp_nodes, + reflections_per_adp_parameter=args.reflections_per_adp_parameter, aniso_selection=args.anisotropic_selection, wavelength=args.wavelength, + add_hydrogens=args.add_hydrogens, + hydrogens_in_xray=args.hydrogens_in_xray, ) # Merge onto DEFAULT_GROUP_WEIGHTS so unspecified groups keep their defaults; @@ -381,6 +411,9 @@ def main(): "n_cycles": args.n_cycles, "mode": args.mode, "adp_mode": args.adp_mode, + "adp_mode_set": args.adp_mode_set, + "adp_nodes": args.adp_nodes, + "reflections_per_adp_parameter": args.reflections_per_adp_parameter, "anisotropic_selection": ( args.anisotropic_selection if args.adp_mode == "anisotropic" else None ), diff --git a/torchref/cli/simulate_noisy_data.py b/torchref/cli/simulate_noisy_data.py new file mode 100644 index 00000000..7e37361d --- /dev/null +++ b/torchref/cli/simulate_noisy_data.py @@ -0,0 +1,282 @@ +#!/usr/bin/env python3 +"""Simulate a noisy reflection MTZ from a structure file. + +Two modes are supported, chosen by whether ``--reference-hkl`` is passed: + +**Reference-driven (preferred)** + The model is scaled to the reference dataset with torchref's + ``Scaler`` so that Fcalc and reference intensities share an absolute + scale. The simulation then inherits the reference's HKL list, and + noise for each reflection uses the reference's reported ``sigma(I)`` + directly — no parametric model, no fitting. Use this when you have a + CrystFEL ``.hkl`` (or equivalent) from a real experiment. + +**Parametric fallback** + When ``--reference-hkl`` is omitted, intensity sigma is built from a + three-term variance-additive model + ``sigma_I^2 = sigma_lin^2 * I + sigma_mul^2 * I^2 + sigma_abs_I^2``. + +In both modes, two independent noise realizations are drawn per +reflection; R-split and Pearson CC between the halves are printed, and +the mean is written as F-obs/SIGF-obs (default) or I-obs/SIGI-obs. + +Usage +----- +:: + + # Reference-driven + torchref.simulate-noisy-data input.pdb out.mtz \ + --reference-hkl td1.hkl --d-min 2.0 + + # Parametric (no reference) + torchref.simulate-noisy-data input.pdb out.mtz \ + --sigma-lin 5 --sigma-mul 0.05 --sigma-abs 0.33 + + # Intensity output + torchref.simulate-noisy-data input.pdb out.mtz \ + --reference-hkl td1.hkl --output-type intensities +""" + +import argparse +import sys +from pathlib import Path + +import pandas as pd +import torch + +from torchref.cli._common import ( + add_device_arg, + add_verbose_arg, + configure_unbuffered_output, + load_model, + parse_device_str, +) +from torchref.io import mtz +from torchref.io.datasets import FcalcDataset + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="torchref.simulate-noisy-data", + description=( + "Compute Fcalc from a structure, add Gaussian intensity noise " + "(optionally grafted from a CrystFEL reference hkl), and write " + "an MTZ as F-obs/SIGF-obs (default) or I-obs/SIGI-obs." + ), + ) + parser.add_argument("input", help="Input structure file (.pdb / .cif / .mmcif)") + parser.add_argument("output", help="Output MTZ file") + parser.add_argument( + "--d-min", + type=float, + default=2.0, + help="High-resolution limit in Angstroms (default: 2.0)", + ) + parser.add_argument( + "--d-max", + type=float, + default=None, + help="Low-resolution limit in Angstroms (default: no cutoff)", + ) + parser.add_argument( + "--output-type", + choices=("amplitudes", "intensities"), + default="amplitudes", + help="Write F-obs/SIGF-obs (amplitudes, default) or I-obs/SIGI-obs", + ) + parser.add_argument( + "--reference-hkl", + type=str, + default=None, + help=("CrystFEL partialator .hkl file. When given, the model is " + "scaled to this reference and per-reflection sigmas are " + "grafted from it (parametric sigma-* flags are ignored)."), + ) + parser.add_argument( + "--sigma-lin", + type=float, + default=0.0, + help=("Parametric only: Poisson coefficient — sigma_I^2 += " + "sigma_lin^2 * I (default: 0.0). Ignored if --reference-hkl."), + ) + parser.add_argument( + "--sigma-mul", + type=float, + default=0.05, + help=("Parametric only: multiplicative coefficient — sigma_I^2 += " + "(sigma_mul * I)^2 (default: 0.05). Ignored if --reference-hkl."), + ) + parser.add_argument( + "--sigma-abs", + type=float, + default=0.0, + help=("Parametric only: target 1/SNR at the resolution limit from " + "the absolute-noise term (default: 0.0). Ignored if " + "--reference-hkl."), + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="Seed for the random generator (default: non-reproducible)", + ) + add_device_arg(parser) + add_verbose_arg(parser) + return parser + + +def main(): + configure_unbuffered_output() + args = _build_parser().parse_args() + + input_path = Path(args.input) + if not input_path.is_file(): + print(f"Error: input file not found: {args.input}", file=sys.stderr) + return 1 + if args.reference_hkl is not None and not Path(args.reference_hkl).is_file(): + print( + f"Error: reference hkl not found: {args.reference_hkl}", + file=sys.stderr, + ) + return 1 + if args.sigma_lin < 0 or args.sigma_mul < 0 or args.sigma_abs < 0: + print("Error: all --sigma-* flags must be non-negative", file=sys.stderr) + return 1 + + device = parse_device_str(args.device) + + if args.verbose: + print(f"Loading structure: {args.input}") + model = load_model( + args.input, max_res=args.d_min, device=device, verbose=args.verbose + ) + + if args.reference_hkl is not None: + return _run_reference_mode(args, model, device) + return _run_parametric_mode(args, model, device) + + +def _run_reference_mode(args, model, device) -> int: + """Reference-driven: scale model to CrystFEL hkl and graft sigmas.""" + from torchref import ReflectionData + from torchref.base.reciprocal import get_d_spacing + from torchref.scaling import Scaler + + if args.verbose: + print(f"Loading CrystFEL reference: {args.reference_hkl}") + ref = ReflectionData(device=str(device), verbose=args.verbose).load_crystfel_hkl( + args.reference_hkl, cell=model.cell, spacegroup=model.spacegroup, + ) + # Prune reference tensors by resolution so the simulation output only + # covers the requested range. cut_res() only masks — we want the HKL + # list itself to be filtered so Scaler, model(), and add_noise all see + # the same reflection set. + if args.d_min is not None or args.d_max is not None: + res = ref.resolution + mask = torch.ones_like(res, dtype=torch.bool) + if args.d_min is not None: + mask &= res >= args.d_min + if args.d_max is not None: + mask &= res <= args.d_max + for field in ("hkl", "I", "I_sigma", "F", "F_sigma", "resolution", "rfree_flags"): + t = getattr(ref, field, None) + if t is not None: + setattr(ref, field, t[mask]) + # Replace the masks with a single all-True flagged_initial matching + # the new tensor size. An empty TensorMasks() returns None from + # __call__(), which downstream consumers (Scaler.get_bins) don't + # tolerate. + ref.masks = type(ref.masks)(device=ref.device) + ref.masks["flagged_initial"] = torch.ones( + len(ref.hkl), dtype=torch.bool, device=ref.device + ) + if args.verbose: + print(f"Reference: {len(ref.hkl)} reflections after resolution cuts") + + if args.verbose: + print("Scaling model to reference...") + scaler = Scaler(model, ref, device=device, verbose=args.verbose) + scaler.initialize().refine_lbfgs() + + if args.verbose: + print(f"Computing scaled Fcalc on {len(ref.hkl)} reference HKLs") + with torch.no_grad(): + fcalc_scaled = scaler(model(ref.hkl)) + + sim = FcalcDataset( + hkl=ref.hkl.clone(), + resolution=get_d_spacing(ref.hkl.float(), ref.cell.data), + cell=ref.cell, + spacegroup=ref.spacegroup, + device=device, + ) + sim.set_fcalc(fcalc_scaled) + + noisy = sim.add_noise(reference=ref, seed=args.seed, verbose=bool(args.verbose)) + _write_output(args, noisy) + return 0 + + +def _run_parametric_mode(args, model, device) -> int: + """No reference: build HKL from cell+resolution, use three-term model.""" + if args.verbose: + print( + f"Generating HKL to d_min={args.d_min} A" + + (f", d_max={args.d_max} A" if args.d_max is not None else "") + ) + dataset = FcalcDataset.from_cell_and_resolution( + cell=model.cell, + spacegroup=model.spacegroup, + d_min=args.d_min, + d_max=args.d_max, + device=device, + ) + if args.verbose: + print(f"Computing Fcalc for {len(dataset)} reflections") + hkl = dataset.hkl.to(device) + fcalc = model(hkl, recalc=True) + dataset.set_fcalc(fcalc) + + noisy = dataset.add_noise( + sigma_lin=args.sigma_lin, + sigma_mul=args.sigma_mul, + sigma_abs=args.sigma_abs, + seed=args.seed, + verbose=bool(args.verbose), + ) + _write_output(args, noisy) + return 0 + + +def _write_output(args, noisy: FcalcDataset) -> None: + """Write noisy amplitudes or intensities and their uncertainties to MTZ.""" + hkl_np = noisy.hkl.cpu().numpy() + columns = { + "H": hkl_np[:, 0], + "K": hkl_np[:, 1], + "L": hkl_np[:, 2], + } + + if args.output_type == "intensities": + # The intensity add_noise actually drew, not the square of the clamped + # amplitude. Squaring back would floor every negative reflection at zero, and + # that is a positive bias concentrated exactly where the noise dominates -- the + # same signature as a real positive perturbation of the merged intensity. + columns["I-obs"] = noisy.I.detach().cpu().numpy() + columns["SIGI-obs"] = noisy.I_sigma.detach().cpu().numpy() + else: + columns["F-obs"] = noisy.fcalc_amp.detach().cpu().numpy() + columns["SIGF-obs"] = noisy.fobs_sigma.detach().cpu().numpy() + + df = pd.DataFrame(columns) + mtz.write(df, noisy.cell.data, noisy.spacegroup, args.output) + + if args.verbose: + print( + f"Wrote {args.output} " + f"({args.output_type}, n={len(noisy)})" + ) + + +if __name__ == "__main__": + sys.exit(main() or 0) diff --git a/torchref/cli/validate_ded.py b/torchref/cli/validate_ded.py index 307d0c69..18b6d9a7 100644 --- a/torchref/cli/validate_ded.py +++ b/torchref/cli/validate_ded.py @@ -24,26 +24,33 @@ import argparse import json import sys +import warnings from pathlib import Path import numpy as np import torch from torchref.cli._common import ( - add_dual_model_args, + add_ded_weight_args, add_dmin_arg, + add_dual_model_args, add_general_args, add_outdir_arg, - build_dual_column_names, configure_unbuffered_output, load_model, load_reflection_data, - register_timing, parse_device_str, + register_timing, + sigma_d_config_from_args, validate_cif_files, validate_files, ) +from torchref.maps.ded_weights import ( + DEFAULT_SCHEME, + DedWeightFallbackWarning, + all_ded_weights, +) from torchref.utils.serialization import convert_to_serializable configure_unbuffered_output() @@ -63,6 +70,8 @@ def build_atom_mask(selection_xyz, real_space_grid, cell, mask_radius, device): """ from torchref.base.coordinates.transforms_torch import ( get_fractional_matrix, + ) + from torchref.base.coordinates.transforms_torch import ( get_inv_fractional_matrix_torch as get_inverse_fractional_matrix, ) from torchref.base.electron_density.solvent_mask import add_to_solvent_mask @@ -80,7 +89,7 @@ def build_atom_mask(selection_xyz, real_space_grid, cell, mask_radius, device): inv_frac_matrix=inv_frac, ) - mask = torch.zeros(grid_shape, dtype=torch.int32, device=device) + mask = torch.zeros(grid_shape, dtype=torch.int32, device=device) # dtype-ok: integer solvent-mask accumulator (mask>0); categorical count, not model-precision data mask = add_to_solvent_mask( surrounding_coords, voxel_indices, @@ -195,6 +204,8 @@ def setup_ded_context( col_light=None, n_bins=20, verbose=0, + ded_weight=DEFAULT_SCHEME, + sigma_d_config=None, ): """Load reflection data and prepare shared state for DED validation. @@ -227,10 +238,8 @@ def setup_ded_context( import gemmi from torchref import DatasetCollection - from torchref.symmetry.grid_utils import calculate_optimal_grid_size - from torchref.symmetry.reciprocal_symmetry import expand_hkl - - from torchref.config import normalize_device + from torchref.config import get_float_dtype, normalize_device + from torchref.symmetry import Cell, SpaceGroup device = normalize_device(device) @@ -249,15 +258,10 @@ def setup_ded_context( collection.add_dataset("dark", data_dark) collection.add_dataset("light", data_light) collection.scale() + data_dark, data_light = collection["dark"], collection["light"] if verbose >= 1: - print(f"Scale parameters after optimization:") - for name, ds in collection: - if hasattr(ds, "log_scale") and ds.log_scale is not None: - print( - f" {name}: log_scale={ds.log_scale.item():.6f} " - f"(scale={torch.exp(ds.log_scale).item():.6f})" - ) + print("Inter-dataset scaling:", collection.scaling_metrics) # Extract matched reflections hkl_all = data_dark.hkl @@ -286,11 +290,20 @@ def setup_ded_context( else: free_mask = work_mask = None - # Weighted difference Fo + # Difference Fo and the registered weights; the selected scheme is the headline. dfo = F_light - F_dark - sig_diff = (sig_dark**2 + sig_light**2) ** 0.5 - weights = 1 / sig_diff**2 - weights = weights / weights.mean() + sig_diff = torch.sqrt(sig_dark**2 + sig_light**2) + all_w = all_ded_weights( + delta_obs=dfo, + sigma_diff=sig_diff, + hkl=hkl, + cell=data_dark.cell, + spacegroup=data_dark.spacegroup, + f_dark=F_dark, + sigma_d_config=sigma_d_config, + ) + selected = all_w[ded_weight] + weights = selected.weights w_dfo = dfo * weights # Cell, spacegroup, d-spacings @@ -305,15 +318,19 @@ def setup_ded_context( ) if dmin is None: dmin = float(d_spacings.min()) - d_spacing = torch.tensor(d_spacings, dtype=torch.float32, device=device) + d_spacing = torch.tensor(d_spacings, dtype=get_float_dtype(), device=device) # P1 expansion and grid - gridsize = calculate_optimal_grid_size(cell_t, dmin, sg_name) - hkl_p1, orig_idx, phase_shifts = expand_hkl( - hkl, sg_name, include_friedel=False, remove_absences=True + sg = SpaceGroup(sg_name, device=device) + gridsize = sg.optimal_grid_size(Cell(cell_t, device=device), dmin) + hkl_p1, orig_idx, phase_shifts = sg.expand_hkl( + hkl, include_friedel=False, remove_absences=True ) w_dfo_p1 = w_dfo[orig_idx] weights_p1 = weights[orig_idx] + weights_by_scheme = { + name: (w.weights, w.weights[orig_idx]) for name, w in all_w.items() + } if verbose >= 1: print(f"Matched reflections: {len(hkl)}") @@ -331,6 +348,16 @@ def setup_ded_context( "refl_mask": refl_mask, "w_dfo": w_dfo, "weights": weights, + "dfo": dfo, + "dfo_p1": dfo[orig_idx], + "weights_by_scheme": weights_by_scheme, + "ded_weight": ded_weight, + "ded_weight_applied": selected.applied, + "ded_weight_diagnostics": { + k: v + for k, v in all_w["sigma_d"].diagnostics.items() + if k not in ("weight_sigma_d_raw", "shells") + }, "d_spacing": d_spacing, "cell_t": cell_t, "cell_np": cell_np, @@ -387,11 +414,11 @@ def compute_ded_maps( resolution_bins, reciprocal_cc_overall, reciprocal_cc_work, reciprocal_cc_free, w_delta_fcalc_asu. """ - from torchref.model.model_collection import ModelCollection + from torchref.base.fourier.grid import get_real_grid from torchref.cli.collection_difference_refine import ( setup_scaler as setup_collection_scaler, ) - from torchref.base.fourier.grid import get_real_grid + from torchref.model.model_collection import ModelCollection device = ctx["device"] @@ -531,10 +558,48 @@ def compute_ded_maps( if cc_work is not None: print(f" Work CC = {cc_work:.4f}, Free CC = {cc_free:.4f}") + # Every registered scheme on the same coefficients, for side-by-side reporting. + by_weight = {} + for name, (w_asu, w_p1) in ctx.get("weights_by_scheme", {}).items(): + with torch.no_grad(): + m_o = compute_map_from_coefficients( + ctx["dfo_p1"] * w_p1, phi_dark_p1, ctx["hkl_p1"], ctx["gridsize"] + ) + m_c = compute_map_from_coefficients( + delta_fcalc * w_p1, phi_dark_p1, ctx["hkl_p1"], ctx["gridsize"] + ) + entry = { + "realspace_correlation": { + mname: round(float(compute_correlation(m_o, m_c, mm)), 4) + for mname, mm in mask_dict.items() + } + } + wo, wc = ctx["dfo"] * w_asu, delta_fcalc_asu * w_asu + entry["reciprocal_cc_overall"] = round( + torch.corrcoef(torch.stack([wo, wc]))[0, 1].item(), 4 + ) + if free_mask is not None and free_mask.sum() > 10: + entry["reciprocal_cc_work"] = round( + torch.corrcoef(torch.stack([wo[work_mask], wc[work_mask]]))[ + 0, 1 + ].item(), + 4, + ) + entry["reciprocal_cc_free"] = round( + torch.corrcoef(torch.stack([wo[free_mask], wc[free_mask]]))[ + 0, 1 + ].item(), + 4, + ) + else: + entry["reciprocal_cc_work"] = entry["reciprocal_cc_free"] = None + by_weight[name] = entry + return { "map_dfo": map_dfo, "map_dfc": map_dfc, "mask_dict": mask_dict, + "by_weight": by_weight, "realspace_correlation": rs_corr, "resolution_bins": bin_results, "reciprocal_cc_overall": round(cc_overall, 4), @@ -551,11 +616,13 @@ def compute_ded_maps( def run_validation(args): """Run the DED validation pipeline.""" - from torchref.cli.collection_difference_refine import compute_rfactors - from torchref.model.model_collection import ModelCollection + from torchref.cli.collection_difference_refine import ( + compute_rfactors, + ) from torchref.cli.collection_difference_refine import ( setup_scaler as setup_collection_scaler, ) + from torchref.model.model_collection import ModelCollection device = parse_device_str(args.device) outdir = Path(args.outdir) @@ -570,16 +637,27 @@ def run_validation(args): print(f" Light SF: {args.light_structure_factor}") col_dark, col_light = build_dual_column_names(args) - ctx = setup_ded_context( - args.dark_structure_factor, - args.light_structure_factor, - dmin=args.dmin, - device=device, - col_dark=col_dark, - col_light=col_light, - n_bins=args.n_bins, - verbose=args.verbose, - ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DedWeightFallbackWarning) + ctx = setup_ded_context( + args.dark_structure_factor, + args.light_structure_factor, + dmin=args.dmin, + device=device, + col_dark=col_dark, + col_light=col_light, + n_bins=args.n_bins, + verbose=args.verbose, + ded_weight=args.ded_weight, + sigma_d_config=sigma_d_config_from_args(args), + ) + fallback_messages = [ + str(w.message) + for w in caught + if issubclass(w.category, DedWeightFallbackWarning) + ] + for message in fallback_messages: + print(f"WARNING: {message}") d_min = ctx["d_min"] # Load models @@ -638,9 +716,29 @@ def run_validation(args): "light_model": str(args.light_model), "fraction": args.fraction, "selection": args.selection, + # Mask choice changes which density enters the correlation, so record it + # alongside the score for reproducibility. + "mask_source": args.mask_source, "mask_radius": args.mask_radius, "dmin": d_min, }, + "weights": { + "requested": ctx["ded_weight"], + "applied": ctx["ded_weight_applied"], + **{ + k: ctx["ded_weight_diagnostics"].get(k) + for k in ( + "gamma", + "gamma_fitted", + "gamma_reason", + "tau", + "n_shell", + "n_s2_clamped", + "fallback_reason", + ) + }, + }, + "by_weight": result["by_weight"], "realspace_correlation": result["realspace_correlation"], "reciprocal_cc_overall": result["reciprocal_cc_overall"], "reciprocal_cc_work": result["reciprocal_cc_work"], @@ -689,10 +787,25 @@ def run_validation(args): # Summary if args.verbose >= 1: print(f"\n{'=' * 70}") - print("Summary:") + print( + f"Summary (headline weights: {ctx['ded_weight']}, " + f"applied: {ctx['ded_weight_applied']}):" + ) for name, corr in result["realspace_correlation"].items(): print(f" {name}: CC = {corr['cc']:.4f}") print(f" Reciprocal-space CC (overall): {result['reciprocal_cc_overall']}") + masks = list(result["realspace_correlation"]) + header = " {:<17s}".format("weights") + "".join( + f"{m[:12]:>13s}" for m in masks + ) + print(header + f"{'recip. CC':>13s}") + for name, entry in result["by_weight"].items(): + row = f" {name:<17s}" + "".join( + f"{entry['realspace_correlation'][m]:13.4f}" for m in masks + ) + print(row + f"{entry['reciprocal_cc_overall']:13.4f}") + for message in fallback_messages: + print(f" WARNING: {message}") print(f"{'=' * 70}") return 0 @@ -784,6 +897,7 @@ def main(): action="store_true", help="Write CCP4 map files for WDFo and WDFcalc", ) + add_ded_weight_args(analysis) res = parser.add_argument_group("Resolution") add_dmin_arg(res) diff --git a/torchref/config.py b/torchref/config.py index 617b7627..f528df65 100644 --- a/torchref/config.py +++ b/torchref/config.py @@ -568,3 +568,38 @@ def __repr__(self) -> str: def get_default_device() -> torch.device: """Get the current default device.""" return device.current + + +# --------------------------------------------------------------------------- +# Double-precision availability +# --------------------------------------------------------------------------- +#: Device types with no float64 at all. MPS is the live case and it *raises* +#: rather than quietly downcasting, so a float64 tensor there is an error and not +#: merely slow. +_NO_DOUBLE_DEVICE_TYPES = ("mps",) + + +def supports_double(dev=None) -> bool: + """Whether ``dev`` can hold float64 / complex128 at all.""" + return normalize_device(dev).type not in _NO_DOUBLE_DEVICE_TYPES + + +def widest_float_dtype(dev=None) -> torch.dtype: + """``float64`` where the device has it, else the configured working float. + + For computations whose *precision* is load-bearing rather than their storage: + accumulating single-precision data in double is the ordinary remedy, and the + dynamic range of an unnormalised recurrence is a hard requirement rather than + a preference. The right width for those is a property of the device, so it + belongs here and not in a constant at the call site. + + Where the device lacks float64 the caller gets the working dtype and whatever + accuracy that implies. That is the only option there, not a choice -- callers + that care should say what it costs in their own docstring. + """ + return torch.float64 if supports_double(dev) else get_float_dtype() + + +def widest_complex_dtype(dev=None) -> torch.dtype: + """``complex128`` where the device has it, else the configured working complex.""" + return torch.complex128 if supports_double(dev) else get_complex_dtype() diff --git a/torchref/data/ener_lib_atoms.csv b/torchref/data/ener_lib_atoms.csv new file mode 100644 index 00000000..1035e5ad --- /dev/null +++ b/torchref/data/ener_lib_atoms.csv @@ -0,0 +1,213 @@ +# Per-energy-type atom properties from the CCP4 monomer library ener_lib.cif (_lib_atom loop). +# hb_type: N neither, D donor, A acceptor, B both, H hydrogen able to hydrogen-bond. +# vdw_radius: contact radius in Angstrom; vdwh_radius: radius to use when the atom's own hydrogens are not modelled. +# Regenerate with python -m torchref.scripts.extract_ener_lib +type,element,hb_type,vdw_radius,vdwh_radius,ion_radius +CSP,C,N,1.700,1.700, +CSP1,C,N,1.700,1.700, +C,C,N,1.700,1.750, +C1,C,N,1.700,1.820, +C2,C,N,1.700,1.800, +CR1,C,N,1.700,1.800, +CR2,C,N,1.700,1.800, +CR1H,C,N,1.700,1.800, +CR15,C,N,1.700,1.740, +CR5,C,N,1.700,1.740, +CR56,C,N,1.700,1.740, +CR55,C,N,1.700,1.740, +CR16,C,N,1.700,1.820, +CR6,C,N,1.700,1.740, +CR66,C,N,1.700,1.740, +CH1,C,N,1.700,1.950, +CH2,C,N,1.700,1.920, +CH3,C,N,1.700,1.940, +CT,C,N,1.700,1.850, +NS,N,A,1.550,1.600,1.32 +NSP,N,A,1.550,1.600,1.32 +NS1,N,D,1.550,1.600,1.32 +NSP1,N,D,1.550,1.600,1.32 +N,N,N,1.550,1.600,1.32 +NC1,N,D,1.550,1.600,1.32 +NH0,N,N,1.550,1.600,1.32 +NH1,N,D,1.550,1.600,1.32 +NC2,N,D,1.550,1.600,1.32 +NH2,N,D,1.550,1.600,1.32 +NC3,N,D,1.550,1.600,1.32 +N20,N,A,1.550,1.600,1.32 +N21,N,B,1.550,1.600,1.32 +NT,N,N,1.550,1.600,1.32 +NT1,N,D,1.550,1.600,1.32 +NT2,N,D,1.550,1.600,1.32 +NT3,N,D,1.550,1.600,1.32 +NT4,N,D,1.550,1.600,1.32 +N30,N,A,1.550,1.600,1.32 +N31,N,B,1.550,1.600,1.32 +N32,N,B,1.550,1.600,1.32 +N33,N,B,1.550,1.600,1.32 +NPA,N,A,1.550,1.600,1.32 +NPB,N,A,1.550,1.600,1.32 +NR5,N,A,1.550,1.600,1.32 +NR15,N,D,1.550,1.600,1.32 +NRD5,N,A,1.550,1.600,1.32 +NR56,N,N,1.550,1.600,1.32 +NR55,N,N,1.550,1.600,1.32 +NR6,N,A,1.550,1.600,1.32 +NR66,N,N,1.550,1.600,1.32 +NR16,N,D,1.550,1.600,1.32 +NRD6,N,A,1.550,1.600,1.32 +OS,O,A,1.520,1.520,1.28 +O,O,A,1.520,1.520,1.28 +AF,AF,A,1.520,1.520,1.28 +O2,O,A,1.520,1.520,1.28 +OH1,O,B,1.520,1.520,1.28 +OH2,O,B,1.520,1.520,1.28 +OHA,O,B,1.520,1.680,1.28 +OHB,O,B,1.520,1.680,1.28 +OHC,O,B,1.520,1.680,1.28 +OC2,O,A,1.520,1.520,1.28 +OC,O,A,1.520,1.520,1.28 +OP,O,A,1.520,1.520,1.28 +OB,O,A,1.520,1.520,1.28 +P,P,N,1.800,1.88,1.79 +P1,P,N,1.800,1.88,1.79 +PS,P,N,1.800,1.90,1.79 +S,S,A,1.800,1.88,1.7 +S3,S,A,1.800,1.88,1.7 +S2,S,A,1.800,1.88,1.7 +S1,S,A,1.800,1.88,1.7 +ST,S,A,1.800,1.88,1.7 +SH1,S,B,1.800,1.950,1.7 +H,H,N,1.200,1.200, +HCH,H,N,1.200,1.200, +HCH1,H,N,1.200,1.200, +HCH2,H,N,1.200,1.200, +HCH3,H,N,1.200,1.200, +HCR1,H,N,1.200,1.200, +HC2,H,N,1.200,1.200, +HC1,H,N,1.200,1.200, +HC2,H,N,1.200,1.200, +HCR5,H,N,1.200,1.200, +HCR6,H,N,1.200,1.200, +HNC1,H,H,1.200,1.200, +HNC2,H,H,1.200,1.200, +HNH1,H,H,1.200,1.200, +HNH2,H,H,1.200,1.200, +HNR5,H,H,1.200,1.200, +HNR6,H,H,1.200,1.200, +HNT1,H,H,1.200,1.200, +HNT2,H,H,1.200,1.200, +HNT3,H,H,1.200,1.200, +HOH1,H,H,1.200,1.200, +HOH2,H,H,1.200,1.200, +HOHA,H,H,1.200,1.200, +HOHB,H,H,1.200,1.200, +HOHC,H,H,1.200,1.200, +HSH1,H,H,1.200,1.200, +SI,SI,N,2.10,2.10,0.40 +SI1,SI,N,2.10,2.10,0.40 +GE,GE,N,2.10,2.10,0.40 +GE1,GE,N,2.10,2.10,0.40 +SN,SN,N,2.17,2.17,0.69 +PB,PB,N,2.02,2.02,0.79 +LI,LI,N,1.82,1.82,0.73 +NA,NA,N,2.27,2.27,1.13 +K,K,N,2.75,2.75,1.51 +RB,RB,N,2.00,2.00,1.48 +CS,CS,N,2.98,2.98,1.81 +FR,FR,N,,,1.94 +BE,BE,N,1.12,1.12,0.41 +MG,MG,N,1.73,1.73,0.71 +CA,CA,N,1.94,1.94,1.14 +SR,SR,N,2.19,2.19,1.32 +BA,BA,N,2.53,2.53,1.49 +RA,RA,N,2.15,2.15,1.62 +SC,SC,N,1.60,1.60,0.885 +Y,Y,N,1.80,1.80,1.04 +LA,LA,N,1.95,1.95,1.172 +CE,CE,N,1.85,1.85,1.01 +PR,PR,N,1.85,1.85,0.99 +ND,ND,N,1.85,1.85,1.123 +PM,PM,N,,,1.11 +SM,SM,N,1.85,1.85,1.098 +EU,EU,N,1.85,1.85,1.087 +GD,GD,N,1.80,1.80,1.078 +TB,TB,N,1.75,1.75,0.90 +DY,DY,N,1.75,1.75,1.052 +HO,HO,N,1.75,1.75,1.041 +ER,ER,N,1.75,1.75,1.03 +TM,TM,N,1.75,1.75,1.02 +YB,YB,N,1.75,1.75,1.008 +LU,LU,N,1.75,1.75,1.001 +AC,AC,N,1.95,1.95,1.26 +TH,TH,N,1.80,1.80,1.08 +PA,PA,N,1.80,1.80,0.92 +U,U,N,1.86,1.86,0.66 +NP,NP,N,1.75,1.75,0.85 +PU,PU,N,1.75,1.75,0.85 +AM,AM,N,1.75,1.75,0.99 +CM,CM,N,,,0.99 +BK,BK,N,,,0.97 +CF,CF,N,,,0.961 +ES,ES,N,,, +FM,FM,N,,, +MD,MD,N,,, +NO,NO,N,,, +LR,LR,N,,, +TI,TI,N,1.40,1.40,0.56 +ZR,ZR,N,1.55,1.55,0.73 +HF,HF,N,1.55,1.55,0.72 +RF,RF,N,,, +V,V,N,1.35,1.35,0.68 +NB,NB,N,1.45,1.45,0.62 +TA,TA,N,1.45,1.45,0.78 +DB,DB,N,,, +CR,CR,N,1.40,1.40,0.53 +MO,MO,N,1.45,1.45,0.55 +W,W,N,1.35,1.35,0.56 +SG,SG,N,,, +MN,MN,N,1.40,1.40,0.46 +TC,TC,N,1.35,1.35,0.51 +RE,RE,N,1.35,1.35,0.52 +BH,BH,N,,, +FE,FE,N,1.40,1.40,0.68 +RU,RU,N,1.30,1.30,0.52 +OSE,OS,N,1.30,1.30,0.53 +HS,HS,N,,, +CO,CO,N,1.35,1.35,0.54 +RH,RH,N,1.35,1.35,0.69 +IR,IR,N,1.35,1.35,0.71 +MT,MT,N,,, +NI,NI,N,1.63,1.63,0.63 +PD,PD,N,1.63,1.63,0.78 +PT,PT,N,1.75,1.75,0.71 +CU,CU,N,1.40,1.40,0.71 +AG,AG,N,1.72,1.72,0.81 +AU,AU,N,1.66,1.66,0.71 +ZN,ZN,N,1.39,1.39,0.74 +CD,CD,N,1.58,1.58,0.92 +HG,HG,N,1.55,1.55,1.10 +B,B,N,0.85,0.85,0.25 +AL,AL,N,1.25,1.25,0.53 +GA,GA,N,1.87,1.87,0.61 +IN,IN,N,1.93,1.93,0.76 +TL,TL,N,1.96,1.96,0.89 +AS,AS,N,1.85,1.85,0.475 +AS1,AS,N,1.85,1.85,0.475 +SB,SB,N,1.8,1.8,0.90 +BI,BI,N,1.8,1.8,0.90 +SE,SE,N,1.90,1.90,0.42 +TE,TE,N,2.06,2.06,0.57 +PO,PO,N,2.0,2.0,0.81 +F,F,B,1.47,1.47,1.19 +CL,CL,A,1.75,1.75,1.67 +BR,BR,N,1.85,1.85,0.73 +I,I,N,1.98,1.98,0.56 +AT,AT,N,1.80,1.80,0.76 +HE,HE,N,1.40,1.40, +NE,NE,N,1.54,1.54,1.12 +AR,AR,N,1.88,1.88,1.54 +KR,KR,N,2.02,2.02,1.69 +XE,XE,N,2.16,2.16,1.90 +RN,RN,N,2.20,2.20,2.0 +DUM,O,N,0.6,0.6,0.6 +.,C,B,1.7,1.5,1.7 diff --git a/torchref/experimental/alignment/__init__.py b/torchref/experimental/alignment/__init__.py index 1934d2ac..797261f8 100644 --- a/torchref/experimental/alignment/__init__.py +++ b/torchref/experimental/alignment/__init__.py @@ -1,47 +1,35 @@ -""" -Experimental ball-harmonic molecular-replacement engine for TorchRef. - -Experimental / unstable API. This is the opt-in **ball-harmonic MR engine**; -the production / canonical MR entry point is ``torchref.alignment`` (the -consolidated FRF engine, default ``engine="frf_separate"``). Everything in this -package may change without notice (importing it emits a ``FutureWarning``). - -Provides molecular replacement functionality including: - -1. Fast Rotation Function: Ball harmonic transform for rotation search -2. Translation Search: FFT-based translation function -3. Rigid Body Refinement: Optimization of rotation and translation -4. Unified Pipeline: Complete MR workflow with early stopping - -Notes ------ -Optional dependency surface. The pipeline and ball rotation-search symbols -(``MolecularReplacementPipeline``, ``MRSolution``, ``ball_rotation_search``, -``ball_rotation_search_torch``, ``BallHarmonicCoefficients``, -``splat_evalues_to_ball``, ``compute_ball_harmonic_coefficients``, -``compute_ball_cross_correlation_coefficients``, ``evaluate_rotation_function``, -``find_rotation_peaks``, ``reduce_rotation_by_symmetry``, ``RotationCluster``, -``cluster_rotation_peaks`` and the other rotation/Euler helpers) require the -JAX/s2fft stack and are exported only when ``pip install torchref[alignment]`` -is present. Without that extra they fall back to stubs (or are absent from -``__all__``); only translation search, rigid-body refinement, transforms, clash -scoring, distributions, and the sampling utilities are unconditionally -available. - -The package-level ``cluster_rotation_peaks`` is the ``.ball_transform`` version -(signature ``(peaks, cluster_radius_deg=5.0, symmetry_matrices=None, -return_details=False)`` returning 6-tuples / ``RotationCluster`` objects). The -``.pipeline`` module keeps its own simpler variant for internal use, reachable -as ``torchref.experimental.alignment.pipeline.cluster_rotation_peaks``. - -A handful of public ball helpers (``compute_ball_harmonic_coefficients_analytical``, -``refine_peaks_analytical``, ``refine_peaks_subvoxel_wrapper``, -``evaluate_rotation_function_at_angles``, ``build_wigner_index_mapping``) are -documented in their modules but are not re-exported here and are not part of the -supported package API. - -Example - Full MR Pipeline --------------------------- +"""Molecular replacement: a rotation search feeding a translation search. + +Two stages and one normalisation between them. + +1. **Fast Rotation Function** (:func:`rotation_search`, over + :class:`~torchref.experimental.alignment.frf.FastRotationFunction`) -- + Phaser-faithful Bessel-radial x spherical-harmonic expansion against a dense + P1-box calc, with stable Wigner-d. It is a **shortlist generator**: over 30 + seeded cells it puts the true orientation at rank 0 six times, and inside the + top twenty essentially always. Only the second of those is required. +2. **Fast Translation Function** (:mod:`~torchref.experimental.alignment.translation`) + -- a Crowther-Blow correlation over the fractional cell, run per rotation + candidate, then an analytical-R local refine. On the same 30 cells it reaches + rank 0 in 24, and its likelihood in 27. Rotation ghosts are morphologically + identical to truth in a Patterson by construction; they stop being identical + once the crystal lattice is involved. + +There is deliberately nothing between the two. An ML rescore used to sit there +and was removed: it reordered a shortlist that already contained truth, and +end-to-end pose recovery was 18/30 with it against 24/30 without. + +Both stages normalise through :class:`torchref.scaling.WilsonNormaliser` and +weight through :mod:`torchref.scaling.weighting`, so ``E_obs`` means one thing +across the whole run. + +The pipeline returns a **placement** -- rotation and translation -- and stops. +Refining it is downstream refinement's job, and deleting the post-placement +polish that used to be here took pose recovery from 24/30 to 30/30, because on +the hard cases it walked a correct placement away from truth. + +Example +------- :: from torchref.experimental.alignment import MolecularReplacementPipeline @@ -51,37 +39,8 @@ data = ReflectionData().load_mtz('observed.mtz') model = ModelFT().load_pdb('search_model.pdb') - pipeline = MolecularReplacementPipeline(data, model) - solutions = pipeline.run(n_rotation_peaks=50, min_tries=3, max_tries=10) - print(f"Best R-factor: {solutions[0].r_factor:.3f}") - -Example - Individual Components -------------------------------- -:: - - from torchref.experimental.alignment import ( - ball_rotation_search_torch, - fft_translation_search_torch, - RigidBodyRefinement, - ) - - # E_obs, s_obs, E_calc, s_calc, F_obs, F_calc_rotated, hkl assumed - # computed beforehand from the data/model. - - # 1. Rotation search - rf, angles, peaks = ball_rotation_search_torch( - E_obs, s_obs, E_calc, s_calc, L=32, P=20 - ) - - # 2. Translation search for top rotation - alpha, beta, gamma, score, sigma = peaks[0] - corr_map, best_trans, trans_peaks = fft_translation_search_torch( - F_obs, F_calc_rotated, hkl - ) - - # 3. Rigid body refinement - rb = RigidBodyRefinement(model, data, initial_rotation=..., initial_translation=...) - result = rb.refine() + solutions = MolecularReplacementPipeline(data, model).run() + print(f"best R-work: {solutions[0].r_factor:.3f}") """ import warnings @@ -91,197 +50,59 @@ FutureWarning, ) -# ============================================================================= -# Pipeline & Rotation search require JAX + s2fft (dev dependencies) -# ============================================================================= -try: - from .pipeline import ( - MolecularReplacementPipeline, - MRSolution, - rotation_angular_distance, - euler_angular_distance, - ) - from .ball_transform import ( - ball_rotation_search, - ball_rotation_search_torch, - rotation_matrix_from_euler_zyz, - rotation_matrix_to_euler_zyz, - rotation_matrix_to_quaternion, - check_rotation_recovery, - BallHarmonicCoefficients, - splat_evalues_to_ball, - compute_ball_harmonic_coefficients, - compute_ball_cross_correlation_coefficients, - evaluate_rotation_function, - find_rotation_peaks, - reduce_rotation_by_symmetry, - reduce_peaks_by_symmetry, - reduce_peaks_by_symmetry_torch, - cluster_rotation_peaks, - cluster_rotation_peaks_torch, - RotationCluster, - ) - _HAS_BALL_TRANSFORM = True -except ImportError: - _HAS_BALL_TRANSFORM = False - - _BALL_TRANSFORM_MSG = ( - "The alignment pipeline and rotation search require jax, s2fft, " - "s2ball, spherical, and quaternionic. " - "Install with: pip install torchref[alignment]" - ) - - def _missing_dep_factory(name): - """Create a callable stub that raises ImportError with install hint.""" - def _stub(*args, **kwargs): - raise ImportError( - f"{name} is not available. {_BALL_TRANSFORM_MSG}" - ) - _stub.__name__ = name - _stub.__qualname__ = name - return _stub - - # Provide stubs so that attribute access works but calling raises - MolecularReplacementPipeline = _missing_dep_factory("MolecularReplacementPipeline") - MRSolution = _missing_dep_factory("MRSolution") - ball_rotation_search = _missing_dep_factory("ball_rotation_search") - ball_rotation_search_torch = _missing_dep_factory("ball_rotation_search_torch") - -# ============================================================================= -# Translation search -# ============================================================================= +from .frf import ( + FastRotationFunction, + RotationPeak, + dense_calc_via_box, + edmonds_euler_from_rotation_matrix, + phaser_lmax_resolution, + rotation_angular_distance_deg, + rotation_matrix_from_edmonds_euler, +) +from .rotation_search import ( + FRFInputs, + RotationSolutions, + prepare_frf_inputs, + rotation_search, +) from .translation import ( - fft_translation_search, - fft_translation_search_torch, + CandidateTransform, + TranslationObs, TranslationPeak, - find_translation_peaks, - apply_translation_to_fcalc, - apply_translation_to_fcalc_torch, + analytic_r_at, + fast_translation_function, + llg_at_translations, + prepare_candidate, ) - -# ============================================================================= -# Rigid body refinement -# ============================================================================= -from .rigid_body import RigidBodyRefinement, RigidBodyResult - -# ============================================================================= -# Rigid body transformations -# ============================================================================= -from .transform import ( - RigidTransform, - quaternion_normalize, - quaternion_conjugate, - quaternion_multiply, - quaternion_rotate, - quaternion_to_matrix, - matrix_to_quaternion, - axis_angle_to_quaternion, - quaternion_to_axis_angle, - quaternion_to_euler_zyz, - euler_zyz_to_quaternion, - rotation_matrix_from_euler, - sample_angles, +from .pipeline import ( + MolecularReplacementPipeline, + MRSolution, + align_model_to_data, ) -# ============================================================================= -# Clash scoring -# ============================================================================= -from .clashscore import ClashScoreCalculator, AtomSampler, compute_clash_score - -# ============================================================================= -# ML distributions -# ============================================================================= -from .distributions import ( - stable_log_bessel_i0, - rice_log_likelihood, - woolfson_log_likelihood, - combined_log_likelihood, - acentric_pdf, - centric_pdf, -) - -# ============================================================================= -# Sampling utilities -# ============================================================================= -from .sampling import VectorSampler, get_rotation_sampling_range - __all__ = [ - # ------------------------------------------------------------------------- + # Entry points + "align_model_to_data", + "MolecularReplacementPipeline", + "MRSolution", + # Rotation search + "rotation_search", + "RotationSolutions", + "FastRotationFunction", + "FRFInputs", + "prepare_frf_inputs", + "phaser_lmax_resolution", + "dense_calc_via_box", + "RotationPeak", + "rotation_matrix_from_edmonds_euler", + "edmonds_euler_from_rotation_matrix", + "rotation_angular_distance_deg", # Translation search - # ------------------------------------------------------------------------- - "fft_translation_search", - "fft_translation_search_torch", + "TranslationObs", "TranslationPeak", - "find_translation_peaks", - "apply_translation_to_fcalc", - "apply_translation_to_fcalc_torch", - # ------------------------------------------------------------------------- - # Rigid body refinement - # ------------------------------------------------------------------------- - "RigidBodyRefinement", - "RigidBodyResult", - # ------------------------------------------------------------------------- - # Transforms - # ------------------------------------------------------------------------- - "RigidTransform", - "quaternion_normalize", - "quaternion_conjugate", - "quaternion_multiply", - "quaternion_rotate", - "quaternion_to_matrix", - "matrix_to_quaternion", - "axis_angle_to_quaternion", - "quaternion_to_axis_angle", - "quaternion_to_euler_zyz", - "euler_zyz_to_quaternion", - "rotation_matrix_from_euler", - "sample_angles", - # ------------------------------------------------------------------------- - # Clash scoring - # ------------------------------------------------------------------------- - "ClashScoreCalculator", - "AtomSampler", - "compute_clash_score", - # ------------------------------------------------------------------------- - # Distributions - # ------------------------------------------------------------------------- - "stable_log_bessel_i0", - "rice_log_likelihood", - "woolfson_log_likelihood", - "combined_log_likelihood", - "acentric_pdf", - "centric_pdf", - # ------------------------------------------------------------------------- - # Utilities - # ------------------------------------------------------------------------- - "VectorSampler", - "get_rotation_sampling_range", + "fast_translation_function", + "CandidateTransform", + "analytic_r_at", + "prepare_candidate", + "llg_at_translations", ] - -if _HAS_BALL_TRANSFORM: - __all__ += [ - # Pipeline (main entry point) - "MolecularReplacementPipeline", - "MRSolution", - "rotation_angular_distance", - "euler_angular_distance", - # Rotation search - "ball_rotation_search", - "ball_rotation_search_torch", - "rotation_matrix_from_euler_zyz", - "rotation_matrix_to_euler_zyz", - "rotation_matrix_to_quaternion", - "check_rotation_recovery", - "BallHarmonicCoefficients", - "splat_evalues_to_ball", - "compute_ball_harmonic_coefficients", - "compute_ball_cross_correlation_coefficients", - "evaluate_rotation_function", - "find_rotation_peaks", - "reduce_rotation_by_symmetry", - "reduce_peaks_by_symmetry", - "reduce_peaks_by_symmetry_torch", - "cluster_rotation_peaks", - "cluster_rotation_peaks_torch", - "RotationCluster", - ] diff --git a/torchref/experimental/alignment/ball_transform.py b/torchref/experimental/alignment/ball_transform.py deleted file mode 100644 index f40463ce..00000000 --- a/torchref/experimental/alignment/ball_transform.py +++ /dev/null @@ -1,2042 +0,0 @@ -""" -Ball Harmonic Transform for Fast Rotation Function. - -Implements 3D ball transforms that preserve radial (resolution) information -for molecular replacement rotation searches. - -The ball function f(r, θ, φ) is expanded using: -- Uniform radial shells (resolution bins) -- Spherical harmonics for angular component (via s2fft) - -For rotation correlation of two ball functions f and g: - C(R) = ∫ f(x) g(R⁻¹x) dx - = Σ_{p,l,m,n} f*_{p,l,m} g_{p,l,n} D^l_{m,n}(R) - = Σ_{l,m,n} ξ_{l,m,n} D^l_{m,n}(R) - -where ξ_{l,m,n} = Σ_p w_p f*_{p,l,m} g_{p,l,n} sums over radial shells. - -Key property: Rotations only affect the angular part - radial indices are summed. -This preserves resolution information while reducing to a standard Wigner transform. - -Experimental / unstable API: part of ``torchref.experimental.alignment``, -the opt-in ball-harmonic MR engine. The production MR entry point is -``torchref.alignment`` (the consolidated FRF engine). Signatures and behavior -may change without notice. -""" - -import math -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple - -import numpy as np -import torch -from numba import njit - -# Enable JAX 64-bit precision before importing s2fft/s2ball -import jax -jax.config.update("jax_enable_x64", True) - -import s2fft -import s2ball.transform.wigner as wigner_transform -import spherical -import quaternionic - - -# ============================================================================= -# Numba-accelerated Analytical Spherical Harmonic Computation -# ============================================================================= - -@njit(cache=True) -def _compute_plm_recurrence(l_max: int, cos_theta: np.ndarray) -> np.ndarray: - """ - Compute associated Legendre polynomials P_l^m(cos(theta)) using recurrence. - - Parameters - ---------- - l_max : int - Maximum l value (exclusive), i.e., computes for l = 0, 1, ..., l_max-1. - cos_theta : np.ndarray - Cosine of colatitude angles, shape (N,). - - Returns - ------- - Plm : np.ndarray - Associated Legendre polynomials, shape (N, l_max, l_max). - Plm[i, l, m] = P_l^m(cos_theta[i]) for m >= 0. - """ - N = len(cos_theta) - sin_theta = np.sqrt(1 - cos_theta**2) - Plm = np.zeros((N, l_max, l_max)) - - # P_0^0 = 1 - Plm[:, 0, 0] = 1.0 - - if l_max > 1: - # P_1^0 = cos(theta) - Plm[:, 1, 0] = cos_theta - # P_1^1 = -sin(theta) - Plm[:, 1, 1] = -sin_theta - - # Recurrence for P_l^l (diagonal) - for l in range(2, l_max): - Plm[:, l, l] = -(2*l - 1) * sin_theta * Plm[:, l-1, l-1] - - # Recurrence for P_l^{l-1} (subdiagonal) - for l in range(2, l_max): - Plm[:, l, l-1] = cos_theta * (2*l - 1) * Plm[:, l-1, l-1] - - # Recurrence for P_l^m (general) - for l in range(2, l_max): - for m in range(0, l-1): - Plm[:, l, m] = ((2*l - 1) * cos_theta * Plm[:, l-1, m] - - (l + m - 1) * Plm[:, l-2, m]) / (l - m) - - return Plm - - -@njit(cache=True) -def _compute_sh_normalization_factors(l_max: int) -> np.ndarray: - """ - Compute normalization factors for spherical harmonics. - - K_l^m = sqrt((2l+1)/(4*pi) * (l-m)!/(l+m)!) - - Parameters - ---------- - l_max : int - Maximum l value (exclusive). - - Returns - ------- - K : np.ndarray - Normalization factors, shape (l_max, l_max). - K[l, m] for m >= 0. - """ - K = np.zeros((l_max, l_max)) - for l in range(l_max): - for m in range(l + 1): - # Compute log((l-m)!/(l+m)!) for numerical stability - log_factor = 0.0 - for k in range(l - m + 1, l + m + 1): - log_factor -= np.log(k) - K[l, m] = np.sqrt((2*l + 1) / (4 * np.pi) * np.exp(log_factor)) - return K - - -@njit(cache=True) -def _compute_sh_coeffs_analytical_numba( - E: np.ndarray, - cos_theta: np.ndarray, - phi: np.ndarray, - L: int, -) -> np.ndarray: - """ - Compute spherical harmonic coefficients analytically using numba. - - Computes a_lm = (4*pi/N) * sum_i(E_i * conj(Y_lm(theta_i, phi_i))) - - where Y_lm = K_l^m * P_l^m(cos(theta)) * exp(i*m*phi) - - Parameters - ---------- - E : np.ndarray - Function values at sample points, shape (N,). - cos_theta : np.ndarray - Cosine of colatitude angles, shape (N,). - phi : np.ndarray - Azimuthal angles in [0, 2*pi), shape (N,). - L : int - Angular bandlimit. - - Returns - ------- - flm : np.ndarray - Spherical harmonic coefficients, shape (L, 2*L-1). - flm[l, m + L - 1] = a_lm for m in [-l, l]. - """ - N = len(E) - - # Compute associated Legendre polynomials - Plm = _compute_plm_recurrence(L, cos_theta) - - # Compute normalization factors - K = _compute_sh_normalization_factors(L) - - # Precompute exp(i*m*phi) for all m values - exp_imphi = np.zeros((N, 2*L - 1), dtype=np.complex128) - for m in range(-(L-1), L): - exp_imphi[:, m + L - 1] = np.cos(m * phi) + 1j * np.sin(m * phi) - - # Compute coefficients - flm = np.zeros((L, 2*L - 1), dtype=np.complex128) - - for l in range(L): - for m in range(-l, l + 1): - # Compute Y_lm - if m >= 0: - Y_lm = K[l, m] * Plm[:, l, m] * exp_imphi[:, m + L - 1] - else: - # Y_l^{-|m|} = (-1)^|m| * conj(Y_l^|m|) - abs_m = -m - sign = 1.0 if abs_m % 2 == 0 else -1.0 - Y_lm = sign * K[l, abs_m] * Plm[:, l, abs_m] * np.conj(exp_imphi[:, abs_m + L - 1]) - - # a_lm = (4*pi/N) * sum(E * conj(Y_lm)) - a_lm = 0.0 + 0.0j - for i in range(N): - a_lm += E[i] * np.conj(Y_lm[i]) - flm[l, m + L - 1] = (4 * np.pi / N) * a_lm - - return flm - - - -# ============================================================================= -# Ball Harmonic Coefficient Container -# ============================================================================= - -@dataclass -class BallHarmonicCoefficients: - """ - Container for ball harmonic coefficients. - - Attributes - ---------- - flmp : np.ndarray - Spherical harmonic coefficients for each radial shell, shape (P, L, 2L-1). - L : int - Angular bandlimit. - P : int - Number of radial shells. - shell_edges : np.ndarray - Radial shell boundaries in Å⁻¹, shape (P+1,). - shell_centers : np.ndarray - Radial shell centers in Å⁻¹, shape (P,). - shell_counts : np.ndarray - Number of reflections in each shell, shape (P,). - """ - flmp: np.ndarray - L: int - P: int - shell_edges: np.ndarray - shell_centers: np.ndarray - shell_counts: np.ndarray - - -def _compute_radial_shells_np( - d_min: float, - d_max: float, - P: int, -) -> Tuple[np.ndarray, np.ndarray]: - """ - Compute uniform radial shell boundaries in reciprocal space. - - Parameters - ---------- - d_min : float - High resolution limit in Angstrom. - d_max : float - Low resolution limit in Angstrom. - P : int - Number of radial shells. - - Returns - ------- - shell_edges : np.ndarray - Shell boundaries in Angstrom^-1, shape (P+1,). - shell_centers : np.ndarray - Shell centers in Angstrom^-1, shape (P,). - """ - s_min = 1.0 / d_max # Low resolution end - s_max = 1.0 / d_min # High resolution end - - shell_edges = np.linspace(s_min, s_max, P + 1) - shell_centers = 0.5 * (shell_edges[:-1] + shell_edges[1:]) - - return shell_edges, shell_centers - - -def _compute_equal_count_shells_np( - s_mag: np.ndarray, - P: int, - d_min: float, - d_max: float, -) -> Tuple[np.ndarray, np.ndarray]: - """ - Compute radial shell boundaries such that each shell has equal reflection count. - - This addresses the issue that uniform spacing in s leads to highly imbalanced - shell counts (low resolution shells have few reflections, high resolution - shells have many). Equal-count binning ensures each shell contributes equally - to the spherical harmonic expansion. - - Parameters - ---------- - s_mag : np.ndarray - Magnitude of s-vectors (|s| = 1/d), shape (N,). - P : int - Number of radial shells. - d_min : float - High resolution limit in Angstrom. - d_max : float - Low resolution limit in Angstrom. - - Returns - ------- - shell_edges : np.ndarray - Shell boundaries in Angstrom^-1, shape (P+1,). - shell_centers : np.ndarray - Shell centers in Angstrom^-1, shape (P,). - """ - s_min = 1.0 / d_max # Low resolution end - s_max = 1.0 / d_min # High resolution end - - # Filter to resolution range - mask = (s_mag >= s_min) & (s_mag <= s_max) - s_in_range = s_mag[mask] - - if len(s_in_range) == 0: - # Fallback to uniform if no reflections in range - return _compute_radial_shells_np(d_min, d_max, P) - - # Compute quantiles for equal-count binning - # We want P bins, so we need P+1 edges at quantiles 0, 1/P, 2/P, ..., 1 - quantiles = np.linspace(0, 1, P + 1) - shell_edges = np.quantile(s_in_range, quantiles) - - # Ensure edges are strictly within bounds - shell_edges[0] = max(shell_edges[0], s_min) - shell_edges[-1] = min(shell_edges[-1], s_max) - - # Ensure strictly increasing (can happen with many identical values) - for i in range(1, len(shell_edges)): - if shell_edges[i] <= shell_edges[i-1]: - shell_edges[i] = shell_edges[i-1] + 1e-10 - - shell_centers = 0.5 * (shell_edges[:-1] + shell_edges[1:]) - - return shell_edges, shell_centers - - -def get_mw_grid(L: int) -> Tuple[np.ndarray, np.ndarray]: - """ - Get McEwen-Wiaux (MW) sampling grid positions. - - Internal helper: not exported in the package ``__all__`` and not part of the - public package API (treat as internal alongside ``splat_to_mw_grid``). - - Parameters - ---------- - L : int - Angular bandlimit. - - Returns - ------- - thetas : np.ndarray - Colatitude samples in [0, π], shape (L,). - phis : np.ndarray - Azimuth samples in [0, 2π), shape (2L-1,). - """ - thetas = np.array([(2*t + 1) * np.pi / (2*L) for t in range(L)]) - phis = np.array([2 * np.pi * p / (2*L - 1) for p in range(2*L - 1)]) - return thetas, phis - - -def splat_to_mw_grid( - theta: np.ndarray, - phi: np.ndarray, - values: np.ndarray, - L: int, - mean_center: bool = True, -) -> np.ndarray: - """ - Splat values onto MW sampling grid using bilinear interpolation. - - Internal helper: not exported in the package ``__all__`` and not part of the - public package API (treat as internal alongside ``get_mw_grid``). - - Parameters - ---------- - theta : np.ndarray - Colatitude angles in [0, π], shape (N,). - phi : np.ndarray - Azimuthal angles in [0, 2π), shape (N,). - values : np.ndarray - Values to splat, shape (N,). - L : int - Angular bandlimit. - mean_center : bool - If True, subtract mean from grid. - - Returns - ------- - grid : np.ndarray - Splatted grid of shape (L, 2L-1). - """ - n_theta = L - n_phi = 2 * L - 1 - - grid = np.zeros((n_theta, n_phi), dtype=np.float64) - weights = np.zeros((n_theta, n_phi), dtype=np.float64) - - # MW grid: theta_t = (2t + 1) * π / (2L) - theta_px = (theta * 2 * L / np.pi - 1) / 2 - phi_px = phi * (2 * L - 1) / (2 * np.pi) - - theta_px = np.clip(theta_px, 0, n_theta - 1 - 1e-6) - - theta_lo = np.floor(theta_px).astype(int) - phi_lo = np.floor(phi_px).astype(int) - - theta_frac = theta_px - theta_lo - phi_frac = phi_px - phi_lo - - theta_hi = np.clip(theta_lo + 1, 0, n_theta - 1) - theta_lo = np.clip(theta_lo, 0, n_theta - 1) - phi_hi = (phi_lo + 1) % n_phi - phi_lo = phi_lo % n_phi - - w00 = (1 - theta_frac) * (1 - phi_frac) - w01 = (1 - theta_frac) * phi_frac - w10 = theta_frac * (1 - phi_frac) - w11 = theta_frac * phi_frac - - np.add.at(grid, (theta_lo, phi_lo), w00 * values) - np.add.at(grid, (theta_lo, phi_hi), w01 * values) - np.add.at(grid, (theta_hi, phi_lo), w10 * values) - np.add.at(grid, (theta_hi, phi_hi), w11 * values) - - np.add.at(weights, (theta_lo, phi_lo), w00) - np.add.at(weights, (theta_lo, phi_hi), w01) - np.add.at(weights, (theta_hi, phi_lo), w10) - np.add.at(weights, (theta_hi, phi_hi), w11) - - mask = weights > 0 - grid[mask] /= weights[mask] - - if mean_center: - grid = grid - grid.mean() - - return grid - - -def _assign_to_shells_np( - s_mag: np.ndarray, - shell_edges: np.ndarray, -) -> np.ndarray: - """ - Internal NumPy version for ball-specific functions. - """ - shell_idx = np.digitize(s_mag, shell_edges) - 1 - P = len(shell_edges) - 1 - # Mark out-of-range as -1 - shell_idx[(shell_idx < 0) | (shell_idx >= P)] = -1 - return shell_idx - - -def splat_evalues_to_ball( - E_values: np.ndarray, - s_vectors: np.ndarray, - L: int, - P: int, - d_min: float, - d_max: float, - mean_center_shells: bool = True, -) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """ - Splat E-values onto a 3D ball grid using uniform radial shells. - - Parameters - ---------- - E_values : np.ndarray - E² values, shape (N,). - s_vectors : np.ndarray - Reciprocal space vectors in Å⁻¹, shape (N, 3). - L : int - Angular bandlimit. - P : int - Number of radial shells. - d_min : float - High resolution limit in Å. - d_max : float - Low resolution limit in Å. - mean_center_shells : bool - If True, mean-center each radial shell. - - Returns - ------- - ball_grid : np.ndarray - 3D ball grid of shape (P, L, 2L-1). - shell_edges : np.ndarray - Shell boundaries in Å⁻¹. - shell_centers : np.ndarray - Shell centers in Å⁻¹. - shell_counts : np.ndarray - Number of reflections per shell. - """ - # Compute uniform radial shells - shell_edges, shell_centers = _compute_radial_shells_np(d_min, d_max, P) - - # Compute |s| and angles - s_mag = np.linalg.norm(s_vectors, axis=1) - s_normed = s_vectors / np.maximum(s_mag[:, np.newaxis], 1e-10) - - # Spherical angles - theta = np.arccos(np.clip(s_normed[:, 2], -1, 1)) - phi = np.arctan2(s_normed[:, 1], s_normed[:, 0]) - phi = phi % (2 * np.pi) - - # Assign to shells - shell_idx = _assign_to_shells_np(s_mag, shell_edges) - - # Initialize ball grid - ball_grid = np.zeros((P, L, 2 * L - 1), dtype=np.float64) - shell_counts = np.zeros(P, dtype=np.int64) - - # Splat each shell - for p in range(P): - mask = shell_idx == p - count = mask.sum() - shell_counts[p] = count - - if count == 0: - continue - - ball_grid[p] = splat_to_mw_grid( - theta[mask], - phi[mask], - E_values[mask], - L, - mean_center=mean_center_shells, - ) - - return ball_grid, shell_edges, shell_centers, shell_counts - - -def compute_ball_harmonic_coefficients( - ball_grid: np.ndarray, - L: int, - shell_edges: np.ndarray, - shell_centers: np.ndarray, - shell_counts: np.ndarray, -) -> BallHarmonicCoefficients: - """ - Compute spherical harmonic coefficients for each radial shell using s2fft. - - Parameters - ---------- - ball_grid : np.ndarray - 3D ball grid of shape (P, L, 2L-1). - L : int - Angular bandlimit. - shell_edges : np.ndarray - Shell boundaries. - shell_centers : np.ndarray - Shell centers. - shell_counts : np.ndarray - Number of reflections per shell. - - Returns - ------- - coeffs : BallHarmonicCoefficients - Ball harmonic coefficients (SH coeffs for each shell). - """ - P = ball_grid.shape[0] - - # Compute SH coefficients for each shell using s2fft - flmp = np.zeros((P, L, 2*L - 1), dtype=np.complex128) - - for p in range(P): - if shell_counts[p] > 0: - # Use s2fft forward transform: grid -> SH coefficients - flmp[p] = s2fft.forward( - ball_grid[p], - L, - sampling="mw", - method="jax", - reality=False, - ) - - return BallHarmonicCoefficients( - flmp=flmp, - L=L, - P=P, - shell_edges=shell_edges, - shell_centers=shell_centers, - shell_counts=shell_counts, - ) - - -def compute_ball_harmonic_coefficients_analytical( - E_values: np.ndarray, - s_vectors: np.ndarray, - L: int, - P: int, - d_min: float, - d_max: float, - normalize_shells: bool = True, - equal_count_shells: bool = True, -) -> BallHarmonicCoefficients: - """ - Compute ball harmonic coefficients analytically without splatting. - - This method computes SH coefficients directly from the reflection positions - using the analytical formula: - a_lm = (4*pi/N) * sum_i(E_i * conj(Y_lm(theta_i, phi_i))) - - This avoids discretization errors from splatting onto a grid and is - significantly faster when using numba acceleration. - - Parameters - ---------- - E_values : np.ndarray - E-values (normalized structure factor amplitudes), shape (N,). - s_vectors : np.ndarray - Reciprocal space vectors in Angstrom^-1, shape (N, 3). - L : int - Angular bandlimit. - P : int - Number of radial shells. - d_min : float - High resolution limit in Angstrom. - d_max : float - Low resolution limit in Angstrom. - normalize_shells : bool - If True, mean-center and variance-normalize E-values in each shell. - This removes DC bias and makes correlations comparable across shells. - equal_count_shells : bool - If True (default), compute shell boundaries such that each shell has - approximately the same number of reflections. This ensures equal - contribution from each resolution shell to the spherical harmonic - expansion. If False, use uniform spacing in s (which leads to - imbalanced shell counts due to increasing reflection density at - higher resolution). - - Returns - ------- - coeffs : BallHarmonicCoefficients - Ball harmonic coefficients computed analytically. - """ - # Compute s-vector magnitudes (needed for shell assignment) - s_mag = np.linalg.norm(s_vectors, axis=1) - - # Compute radial shells - if equal_count_shells: - shell_edges, shell_centers = _compute_equal_count_shells_np(s_mag, P, d_min, d_max) - else: - shell_edges, shell_centers = _compute_radial_shells_np(d_min, d_max, P) - - # Compute normalized s-vectors for angular coordinates - s_normed = s_vectors / np.maximum(s_mag[:, np.newaxis], 1e-10) - - # Spherical angles - theta = np.arccos(np.clip(s_normed[:, 2], -1, 1)) - cos_theta = np.cos(theta) - phi = np.arctan2(s_normed[:, 1], s_normed[:, 0]) - phi = phi % (2 * np.pi) - - # Assign to shells - shell_idx = _assign_to_shells_np(s_mag, shell_edges) - - # Initialize coefficient arrays - flm = np.zeros((P, L, 2*L - 1), dtype=np.complex128) - shell_counts = np.zeros(P, dtype=np.int64) - - # Compute SH coefficients for each shell using numba-accelerated function - for p in range(P): - mask = shell_idx == p - count = mask.sum() - shell_counts[p] = count - - if count == 0: - continue - - E_shell = E_values[mask].copy() - cos_theta_shell = cos_theta[mask] - phi_shell = phi[mask] - - if normalize_shells: - # Mean-center (removes DC bias) - E_shell = E_shell - E_shell.mean() - # Variance-normalize (makes correlations comparable across shells) - E_std = E_shell.std() - if E_std > 1e-10: - E_shell = E_shell / E_std - - # Use numba-accelerated analytical SH computation - flm[p] = _compute_sh_coeffs_analytical_numba(E_shell, cos_theta_shell, phi_shell, L) - - return BallHarmonicCoefficients( - flmp=flm, - L=L, - P=P, - shell_edges=shell_edges, - shell_centers=shell_centers, - shell_counts=shell_counts, - ) - - -def compute_ball_cross_correlation_coefficients( - f_coeffs: BallHarmonicCoefficients, - g_coeffs: BallHarmonicCoefficients, - radial_weights: Optional[np.ndarray] = None, -) -> np.ndarray: - """ - Compute Wigner coefficients for ball cross-correlation. - - The cross-correlation is: - C(R) = Σ_{p,l,m,n} ξ_{l,m,n} D^l_{m,n}(R) - - To recover the rotation R itself (rather than its transpose) under the - s2ball Wigner D convention, g is conjugated and f/g are swapped between the - m and n indices (see the inline comment in the body). The coefficients - actually assembled are therefore - ξ_{l,m,n} = Σ_p w_p f_{p,l,n} conj(g_{p,l,m}) - summed over the P radial shells with weights w_p. - - Parameters - ---------- - f_coeffs : BallHarmonicCoefficients - Ball harmonic coefficients of function f (observed). - g_coeffs : BallHarmonicCoefficients - Ball harmonic coefficients of function g (calculated). - radial_weights : np.ndarray, optional - Weights for each radial shell, shape (P,). - Default: uniform weights based on shell counts. - - Returns - ------- - xi_nlm : np.ndarray - Wigner coefficients, shape (2N-1, L, 2L-1) where N=L. - """ - assert f_coeffs.L == g_coeffs.L, "Angular bandlimits must match" - assert f_coeffs.P == g_coeffs.P, "Radial bandlimits must match" - - L = f_coeffs.L - P = f_coeffs.P - N = L - - if radial_weights is None: - # Weight by number of reflections in each shell (normalized) - radial_weights = f_coeffs.shell_counts.astype(np.float64) - radial_weights = np.where(radial_weights > 0, radial_weights, 0) - - # Normalize weights - weight_sum = radial_weights.sum() - if weight_sum > 0: - radial_weights = radial_weights / weight_sum - else: - radial_weights = np.ones(P) / P - - # Initialize Wigner coefficients: (2N-1, L, 2L-1) = [n_idx, l, m_idx] - xi_nlm = np.zeros((2*N - 1, L, 2*L - 1), dtype=np.complex128) - - # f_coeffs.flmp has shape (P, L, 2L-1) - # Sum over radial index p - # - # For cross-correlation C(R) that finds rotation R such that g(R⁻¹x) ≈ f(x): - # C(R) = ∫ f(x) g(R⁻¹x) dx = Σ f*_{lm} g_{ln} D^l_{mn}(R) - # - # But D^l_{mn}(R) convention in s2ball gives R^T, so we swap f↔g: - # ξ[n_idx, l, m_idx] = Σ_p w_p * conj(g[l, m_idx]) * f[l, n_idx] - # - for p in range(P): - w = radial_weights[p] - f_lm = f_coeffs.flmp[p] # (L, 2L-1) - g_lm = g_coeffs.flmp[p] # (L, 2L-1) - - # Swap f and g to get R instead of R^T - g_conj = np.conj(g_lm) # (L, 2L-1) - - # For s2fft/s2ball wigner convention: - # coeffs[n_idx, l, m_idx] corresponds to D^l_{m,n} - # where m = m_idx - (L-1), n = n_idx - (N-1) - - # Build the product: for each l, compute outer product over m and n - for l in range(L): - # Valid m range: -l to l, i.e., m_idx from L-1-l to L-1+l - m_start = L - 1 - l - m_end = L - 1 + l + 1 - - # Valid n range: -l to l, i.e., n_idx from N-1-l to N-1+l - n_start = N - 1 - l - n_end = N - 1 + l + 1 - - # g_conj[l, m_start:m_end] shape: (2l+1,) - # f_lm[l, n_start:n_end] shape: (2l+1,) - g_l = g_conj[l, m_start:m_end] # (2l+1,) - conjugated g for m index - f_l = f_lm[l, n_start:n_end] # (2l+1,) - f for n index - - # Outer product: (2l+1, 2l+1) -> [n, m] - outer = np.outer(f_l, g_l) # (2l+1, 2l+1) = [n, m] - - xi_nlm[n_start:n_end, l, m_start:m_end] += w * outer - - return xi_nlm - - -def evaluate_rotation_function( - xi_nlm: np.ndarray, - L: int, -) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """ - Evaluate rotation function from Wigner coefficients using s2ball. - - Parameters - ---------- - xi_nlm : np.ndarray - Wigner coefficients, shape (2N-1, L, 2L-1). - L : int - Angular bandlimit. - - Returns - ------- - rotation_function : np.ndarray - Rotation function, shape (2N-1, L, 2L-1) = (gamma, beta, alpha). - alphas : np.ndarray - Alpha angle grid. - betas : np.ndarray - Beta angle grid. - gammas : np.ndarray - Gamma angle grid. - """ - N = L - - # Inverse Wigner transform - rotation_function = wigner_transform.inverse( - xi_nlm, - L=L, - N=N, - method="jax", - ) - - # MW sampling grid positions - betas = np.array([(2*t + 1) * np.pi / (2*L) for t in range(L)]) - alphas = np.array([2 * np.pi * p / (2*L - 1) for p in range(2*L - 1)]) - gammas = np.array([2 * np.pi * p / (2*N - 1) for p in range(2*N - 1)]) - - return np.asarray(rotation_function), alphas, betas, gammas - - -def ball_rotation_search( - E_obs: np.ndarray, - s_obs: np.ndarray, - E_calc: np.ndarray, - s_calc: np.ndarray, - L: int = 32, - P: int = 20, - d_min: float = 4.0, - d_max: float = 50.0, - n_peaks: int = 100, - radial_weights: Optional[np.ndarray] = None, - refine_subvoxel: bool = True, - refine_analytical: bool = False, - analytical_embedding: bool = True, - equal_count_shells: bool = True, - return_coefficients: bool = False, - verbose: bool = True, -) -> Tuple[np.ndarray, tuple, list]: - """ - Perform ball harmonic rotation function search. - - Main entry point for the ball-based fast rotation function. - - Parameters - ---------- - E_obs : np.ndarray - Observed E² values, shape (N_obs,). - s_obs : np.ndarray - Observed s-vectors in Å⁻¹, shape (N_obs, 3). - E_calc : np.ndarray - Calculated E² values, shape (N_calc,). - s_calc : np.ndarray - Calculated s-vectors in Å⁻¹, shape (N_calc, 3). - L : int - Angular bandlimit. - P : int - Number of radial shells. - d_min : float - High resolution limit in Å. - d_max : float - Low resolution limit in Å. - n_peaks : int - Number of peaks to extract. - radial_weights : np.ndarray, optional - Weights for radial shells. - refine_subvoxel : bool - If True, refine peak positions using fast quadratic fitting. - Ignored if refine_analytical=True. - refine_analytical : bool - If True, refine peak positions using analytical Wigner D-matrix - evaluation. This is more accurate but slower than subvoxel refinement. - Provides exact sub-grid positions for the band-limited rotation function. - analytical_embedding : bool - If True (default), compute spherical harmonic coefficients analytically - from reflection positions using numba-accelerated computation. This - avoids discretization errors from splatting and is significantly faster. - If False, use the original splatting approach. - equal_count_shells : bool - If True (default), compute radial shell boundaries such that each shell - has approximately the same number of reflections. This ensures equal - contribution from each resolution shell. Only applies when - analytical_embedding=True. - return_coefficients : bool - If True, also return the Wigner coefficients xi_nlm for later use - (e.g., for rescoring or additional analytical refinement). - verbose : bool - Print progress. - - Returns - ------- - rotation_function : np.ndarray - Full rotation function, shape (2L-1, L, 2L-1) = (gamma, beta, alpha). - angles_grid : tuple - (alphas, betas, gammas) angle grids. - peaks : list - List of (alpha, beta, gamma, score, sigma) tuples. - xi_nlm : np.ndarray, optional - Wigner coefficients, shape (2N-1, L, 2L-1). Only returned if - return_coefficients=True. - """ - import time - - start_time = time.time() - - if verbose: - print(f"Ball rotation search: L={L}, P={P}") - print(f"Resolution range: {d_min:.2f} - {d_max:.2f} Å") - print(f"Embedding method: {'analytical' if analytical_embedding else 'splatting'}") - - if analytical_embedding: - # Analytical approach: compute SH coefficients directly from reflection positions - # This avoids discretization errors from splatting and is much faster - if verbose: - print("Computing analytical spherical harmonic coefficients...") - if equal_count_shells: - print(" Using equal-count shell binning") - - coeffs_obs = compute_ball_harmonic_coefficients_analytical( - E_obs, s_obs, L, P, d_min, d_max, - normalize_shells=True, equal_count_shells=equal_count_shells - ) - coeffs_calc = compute_ball_harmonic_coefficients_analytical( - E_calc, s_calc, L, P, d_min, d_max, - normalize_shells=True, equal_count_shells=equal_count_shells - ) - - if verbose: - print(f" Coefficients shape: {coeffs_obs.flmp.shape}") - print(f" Reflections per shell (obs): min={coeffs_obs.shell_counts.min()}, max={coeffs_obs.shell_counts.max()}") - - else: - # Original splatting approach - # Step 1: Splat E-values onto ball grids (uniform radial shells) - if verbose: - print("Splatting E-values onto ball grid...") - - ball_obs, shell_edges, shell_centers, shell_counts_obs = splat_evalues_to_ball( - E_obs, s_obs, L, P, d_min, d_max - ) - ball_calc, _, _, shell_counts_calc = splat_evalues_to_ball( - E_calc, s_calc, L, P, d_min, d_max - ) - - if verbose: - print(f" Ball grid shape: {ball_obs.shape}") - print(f" Shell range: [{shell_edges[0]:.4f}, {shell_edges[-1]:.4f}] Å⁻¹") - print(f" Reflections per shell (obs): min={shell_counts_obs.min()}, max={shell_counts_obs.max()}") - - # Step 2: Compute spherical harmonic coefficients for each shell - if verbose: - print("Computing spherical harmonic coefficients per shell...") - - coeffs_obs = compute_ball_harmonic_coefficients( - ball_obs, L, shell_edges, shell_centers, shell_counts_obs - ) - coeffs_calc = compute_ball_harmonic_coefficients( - ball_calc, L, shell_edges, shell_centers, shell_counts_calc - ) - - if verbose: - print(f" Coefficients shape: {coeffs_obs.flmp.shape}") - - # Step 3: Compute cross-correlation Wigner coefficients - if verbose: - print("Computing cross-correlation coefficients...") - - xi_nlm = compute_ball_cross_correlation_coefficients( - coeffs_obs, coeffs_calc, radial_weights - ) - - if verbose: - print(f" Wigner coefficients shape: {xi_nlm.shape}") - print(f" Max |ξ|: {np.abs(xi_nlm).max():.6e}") - - # Step 4: Evaluate rotation function - if verbose: - print("Evaluating rotation function via inverse Wigner transform...") - - rotation_function, alphas, betas, gammas = evaluate_rotation_function(xi_nlm, L) - - rf_real = np.real(rotation_function) - - if verbose: - print(f" Rotation function shape: {rf_real.shape}") - print(f" RF range: [{rf_real.min():.4f}, {rf_real.max():.4f}]") - print(f" RF mean: {rf_real.mean():.4f}, std: {rf_real.std():.4f}") - - # Step 5: Find peaks - if verbose: - print("Finding peaks...") - - peaks = find_rotation_peaks(rf_real, alphas, betas, gammas, n_peaks=n_peaks) - - # Step 6: Optionally refine peaks - if peaks: - if refine_analytical: - # Analytical refinement using Wigner D-matrices (more accurate) - if verbose: - print("Refining peaks using analytical Wigner D-matrices...") - - rf_mean = rf_real.mean() - rf_std = rf_real.std() - peaks = refine_peaks_analytical( - peaks, xi_nlm, L, rf_mean=rf_mean, rf_std=rf_std, verbose=verbose - ) - elif refine_subvoxel: - # Fast quadratic subvoxel refinement - if verbose: - print("Refining peaks to sub-voxel accuracy...") - - peaks = refine_peaks_subvoxel_wrapper( - peaks, rf_real, alphas, betas, gammas - ) - - if verbose: - print(f" Refined {len(peaks)} peaks") - - elapsed = time.time() - start_time - if verbose: - print(f"Ball rotation search completed in {elapsed:.2f}s") - if peaks: - print(f"Top peak: alpha={np.degrees(peaks[0][0]):.2f}°, " - f"beta={np.degrees(peaks[0][1]):.2f}°, " - f"gamma={np.degrees(peaks[0][2]):.2f}°, " - f"sigma={peaks[0][4]:.2f}") - - if return_coefficients: - return rf_real, (alphas, betas, gammas), peaks, xi_nlm - return rf_real, (alphas, betas, gammas), peaks - - -def refine_peaks_subvoxel_wrapper( - peaks: list, - rotation_function: np.ndarray, - alphas: np.ndarray, - betas: np.ndarray, - gammas: np.ndarray, -) -> list: - """ - Refine peak positions to sub-voxel accuracy using quadratic fitting. - - Parameters - ---------- - peaks : list - List of (alpha, beta, gamma, score, sigma) tuples. - rotation_function : np.ndarray - Rotation function grid, shape (n_gamma, n_beta, n_alpha). - alphas, betas, gammas : np.ndarray - Angle grids. - - Returns - ------- - refined_peaks : list - List of (alpha, beta, gamma, score, sigma) tuples with refined positions. - """ - import jax.numpy as jnp - from torchref.experimental.alignment.jax_subpixel_peaks import refine_peaks_subvoxel - - if not peaks: - return peaks - - n_gamma, n_beta, n_alpha = rotation_function.shape - - # Compute angle spacings - d_alpha = alphas[1] - alphas[0] if len(alphas) > 1 else 2 * np.pi / n_alpha - d_beta = betas[1] - betas[0] if len(betas) > 1 else np.pi / n_beta - d_gamma = gammas[1] - gammas[0] if len(gammas) > 1 else 2 * np.pi / n_gamma - - # Convert peaks to grid indices - peak_indices = [] - for alpha, beta, gamma, score, sigma in peaks: - # Find nearest grid indices - a_idx = int(round((alpha - alphas[0]) / d_alpha)) % n_alpha - b_idx = int(round((beta - betas[0]) / d_beta)) - b_idx = max(0, min(b_idx, n_beta - 1)) - g_idx = int(round((gamma - gammas[0]) / d_gamma)) % n_gamma - - peak_indices.append([g_idx, b_idx, a_idx]) # shape matches rf: (gamma, beta, alpha) - - peak_indices = jnp.array(peak_indices, dtype=jnp.int32) - grid_jax = jnp.array(rotation_function) - - # Call JAX subvoxel refinement - refined_coords, refined_values = refine_peaks_subvoxel(grid_jax, peak_indices) - - # Convert back to numpy - refined_coords = np.array(refined_coords) - refined_values = np.array(refined_values) - - # Compute mean and std for sigma calculation - rf_mean = rotation_function.mean() - rf_std = rotation_function.std() - - # Convert refined grid indices back to angles - refined_peaks = [] - for i, ((alpha_orig, beta_orig, gamma_orig, score_orig, sigma_orig), refined_val) in enumerate( - zip(peaks, refined_values) - ): - g_refined, b_refined, a_refined = refined_coords[i] - - # Convert to angles (with periodic wrapping for alpha and gamma) - alpha_new = alphas[0] + a_refined * d_alpha - alpha_new = alpha_new % (2 * np.pi) - - beta_new = betas[0] + b_refined * d_beta - beta_new = np.clip(beta_new, 0, np.pi) - - gamma_new = gammas[0] + g_refined * d_gamma - gamma_new = gamma_new % (2 * np.pi) - - # Compute refined sigma - sigma_new = (refined_val - rf_mean) / rf_std if rf_std > 1e-10 else 0.0 - - refined_peaks.append((alpha_new, beta_new, gamma_new, float(refined_val), float(sigma_new))) - - # Note: We do NOT re-sort by refined score because the quadratic interpolation - # gives an approximate value, not the true score. Keeping original order - # preserves the ranking from the grid-based search. - return refined_peaks - - -# ============================================================================= -# Analytical Peak Refinement using Wigner D-matrices -# ============================================================================= - -def build_wigner_index_mapping(L: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """ - Build index mapping from xi_nlm array to Wigner D-matrix indices. - - Pre-computes the mapping between the xi_nlm coefficient array - (shape: 2N-1, L, 2L-1) and the flattened Wigner D-matrix from spherical. - - Parameters - ---------- - L : int - Angular bandlimit. - - Returns - ------- - xi_indices : np.ndarray - Array of (n_idx, l, m_idx) indices, shape (n_terms, 3). - D_indices : np.ndarray - Corresponding indices into spherical's D array, shape (n_terms,). - l_values : np.ndarray - The l value for each term, shape (n_terms,). - """ - N = L - xi_indices = [] - D_indices = [] - l_values = [] - - for ell in range(L): - for m in range(-ell, ell + 1): - m_idx = m + (L - 1) - for n in range(-ell, ell + 1): - n_idx = n + (N - 1) - D_idx = spherical.WignerDindex(ell, m, n) - xi_indices.append((n_idx, ell, m_idx)) - D_indices.append(D_idx) - l_values.append(ell) - - return (np.array(xi_indices, dtype=np.int32), - np.array(D_indices, dtype=np.int32), - np.array(l_values, dtype=np.int32)) - - -def evaluate_rotation_function_at_angles( - xi_nlm: np.ndarray, - alpha: float, - beta: float, - gamma: float, - wigner: spherical.Wigner, - xi_indices: np.ndarray, - D_indices: np.ndarray, - l_values: np.ndarray, -) -> float: - """ - Evaluate rotation function at arbitrary Euler angles using Wigner D-matrices. - - The s2ball inverse Wigner transform uses the normalization: - RF(α,β,γ) = Σ_{l,m,n} ξ_{l,m,n} * [(2l+1)/(8π²)] * D^l_{m,n}(α,β,γ) - - This allows exact evaluation of the band-limited rotation function at - any point, not just grid points. - - Parameters - ---------- - xi_nlm : np.ndarray - Wigner coefficients from compute_ball_cross_correlation_coefficients(), - shape (2N-1, L, 2L-1). - alpha, beta, gamma : float - ZYZ Euler angles in radians. - wigner : spherical.Wigner - Pre-initialized Wigner calculator. - xi_indices, D_indices : np.ndarray - Pre-computed index mappings from build_wigner_index_mapping(). - l_values : np.ndarray - The l value for each term, shape (n_terms,). - - Returns - ------- - float - Rotation function value (real part). - """ - R = quaternionic.array.from_euler_angles(alpha, beta, gamma) - D_all = wigner.D(R) - - xi_flat = xi_nlm[xi_indices[:, 0], xi_indices[:, 1], xi_indices[:, 2]] - D_flat = D_all[D_indices] - - # Apply the s2ball normalization factor (2l+1)/(8π²) for each term - norm_factors = (2 * l_values + 1) / (8 * np.pi**2) - - return np.sum(xi_flat * D_flat * norm_factors).real - - -def refine_peaks_analytical( - peaks: list, - xi_nlm: np.ndarray, - L: int, - rf_mean: float = None, - rf_std: float = None, - verbose: bool = False, -) -> list: - """ - Refine peak positions using analytical Wigner D-matrix evaluation. - - This provides exact sub-grid peak positions by evaluating the rotation - function directly from its Wigner coefficient expansion, rather than - using grid interpolation. - - Parameters - ---------- - peaks : list - List of (alpha, beta, gamma, score, sigma) tuples from grid search. - xi_nlm : np.ndarray - Wigner coefficients from compute_ball_cross_correlation_coefficients(), - shape (2N-1, L, 2L-1). - L : int - Angular bandlimit. - rf_mean, rf_std : float, optional - Mean and std of rotation function for sigma calculation. - If None, sigma values are preserved from input. - verbose : bool - Print progress. - - Returns - ------- - refined_peaks : list - List of (alpha, beta, gamma, score, sigma) tuples with refined positions. - """ - from scipy.optimize import minimize - - if not peaks: - return peaks - - # Pre-compute index mappings (done once) - if verbose: - print(" Building Wigner index mappings...") - xi_indices, D_indices, l_values = build_wigner_index_mapping(L) - wigner = spherical.Wigner(L - 1) - - # Grid spacing for search bounds - d_angle = 2 * np.pi / (2 * L - 1) - d_beta = np.pi / L - - refined_peaks = [] - - for i, (alpha, beta, gamma, score, sigma) in enumerate(peaks): - def neg_rf(angles): - return -evaluate_rotation_function_at_angles( - xi_nlm, angles[0], angles[1], angles[2], - wigner, xi_indices, D_indices, l_values - ) - - # Local search within ~2 grid spacings - bounds = [ - (alpha - 2*d_angle, alpha + 2*d_angle), - (max(0.01, beta - 2*d_beta), min(np.pi - 0.01, beta + 2*d_beta)), - (gamma - 2*d_angle, gamma + 2*d_angle) - ] - - result = minimize(neg_rf, [alpha, beta, gamma], - method='L-BFGS-B', bounds=bounds, - options={'ftol': 1e-10, 'gtol': 1e-10}) - - alpha_ref = result.x[0] % (2 * np.pi) - beta_ref = np.clip(result.x[1], 0, np.pi) - gamma_ref = result.x[2] % (2 * np.pi) - score_ref = -result.fun - - # Compute sigma if stats provided - if rf_mean is not None and rf_std is not None and rf_std > 1e-10: - sigma_ref = (score_ref - rf_mean) / rf_std - else: - sigma_ref = sigma # Keep original - - refined_peaks.append((alpha_ref, beta_ref, gamma_ref, score_ref, sigma_ref)) - - if verbose and (i + 1) % 100 == 0: - print(f" Refined {i + 1}/{len(peaks)} peaks") - - if verbose: - print(f" Refined {len(peaks)} peaks") - - return refined_peaks - - -def find_rotation_peaks( - rotation_function: np.ndarray, - alphas: np.ndarray, - betas: np.ndarray, - gammas: np.ndarray, - n_peaks: int = 100, - sigma_cutoff: float = 2.0, - cluster_radius_deg: float = 5.0, -) -> list: - """ - Extract and cluster peaks from rotation function. - - Parameters - ---------- - rotation_function : np.ndarray - Rotation function, shape (n_gamma, n_beta, n_alpha). - alphas, betas, gammas : np.ndarray - Angle grids. - n_peaks : int - Maximum number of peaks. - sigma_cutoff : float - Minimum sigma above mean. - cluster_radius_deg : float - Clustering radius in degrees. - - Returns - ------- - peaks : list - List of (alpha, beta, gamma, score, sigma) tuples. - """ - rf_mean = rotation_function.mean() - rf_std = rotation_function.std() - - if rf_std < 1e-10: - return [] - - threshold = rf_mean + sigma_cutoff * rf_std - - # Get sorted indices (descending) - flat_rf = rotation_function.flatten() - sorted_idx = np.argsort(flat_rf)[::-1] - - n_gamma, n_beta, n_alpha = rotation_function.shape - cluster_rad = np.radians(cluster_radius_deg) - - peaks = [] - used_angles = [] - - for flat_i in sorted_idx: - if len(peaks) >= n_peaks: - break - - g_idx, b_idx, a_idx = np.unravel_index(flat_i, rotation_function.shape) - score = rotation_function[g_idx, b_idx, a_idx] - - if score < threshold: - break - - alpha = alphas[a_idx] - beta = betas[b_idx] - gamma = gammas[g_idx] - - # Check if too close to existing peak - is_new = True - for prev_alpha, prev_beta, prev_gamma in used_angles: - da = min(abs(alpha - prev_alpha), 2*np.pi - abs(alpha - prev_alpha)) - db = abs(beta - prev_beta) - dg = min(abs(gamma - prev_gamma), 2*np.pi - abs(gamma - prev_gamma)) - if np.sqrt(da**2 + db**2 + dg**2) < cluster_rad: - is_new = False - break - - if is_new: - sigma = (score - rf_mean) / rf_std - peaks.append((alpha, beta, gamma, score, sigma)) - used_angles.append((alpha, beta, gamma)) - - return peaks - - -def rotation_matrix_from_euler_zyz( - alpha: float, - beta: float, - gamma: float, -) -> np.ndarray: - """ - Create rotation matrix from ZYZ Euler angles. - - R = Rz(alpha) @ Ry(beta) @ Rz(gamma) - """ - ca, sa = np.cos(alpha), np.sin(alpha) - cb, sb = np.cos(beta), np.sin(beta) - cg, sg = np.cos(gamma), np.sin(gamma) - - R = np.array([ - [ca*cb*cg - sa*sg, -ca*cb*sg - sa*cg, ca*sb], - [sa*cb*cg + ca*sg, -sa*cb*sg + ca*cg, sa*sb], - [-sb*cg, sb*sg, cb] - ]) - - return R - - -def check_rotation_recovery( - peaks: list, - true_alpha: float, - true_beta: float, - true_gamma: float, - symmetry_matrices: Optional[np.ndarray] = None, - tolerance_deg: float = 10.0, -) -> Tuple[bool, int, float]: - """ - Check if true rotation was recovered among top peaks. - - Parameters - ---------- - peaks : list - List of (alpha, beta, gamma, score, sigma) peaks. - true_alpha, true_beta, true_gamma : float - True rotation angles in radians. - symmetry_matrices : np.ndarray, optional - Symmetry operations for equivalence checking. - tolerance_deg : float - Angular tolerance in degrees. - - Returns - ------- - found : bool - Whether the rotation was found. - rank : int - Rank of matching peak (0-indexed), or -1. - min_error : float - Minimum angular error in degrees. - """ - R_true = rotation_matrix_from_euler_zyz(true_alpha, true_beta, true_gamma) - tolerance_rad = np.radians(tolerance_deg) - - min_error = float('inf') - best_rank = -1 - - for rank, (alpha, beta, gamma, score, sigma) in enumerate(peaks): - R_peak = rotation_matrix_from_euler_zyz(alpha, beta, gamma) - - error = _rotation_matrix_error(R_peak, R_true) - - if error < min_error: - min_error = error - best_rank = rank - - if symmetry_matrices is not None: - for S in symmetry_matrices: - S_rot = S[:3, :3] if S.shape[0] > 3 else S - for R_combined in [S_rot @ R_true, R_true @ S_rot]: - err = _rotation_matrix_error(R_peak, R_combined) - if err < min_error: - min_error = err - best_rank = rank - - found = min_error < tolerance_rad - return found, best_rank, np.degrees(min_error) - - -def _rotation_matrix_error(R1: np.ndarray, R2: np.ndarray) -> float: - """Compute angular error between rotation matrices.""" - R_diff = R1 @ R2.T - trace = np.clip(np.trace(R_diff), -1, 3) - angle = np.arccos((trace - 1) / 2) - return abs(angle) - - -@dataclass -class RotationCluster: - """ - Container for a cluster of rotation peaks. - - Attributes - ---------- - alpha : float - Alpha angle of best peak (radians). - beta : float - Beta angle of best peak (radians). - gamma : float - Gamma angle of best peak (radians). - best_score : float - Score of the best peak in cluster. - best_sigma : float - Sigma of the best peak. - size : int - Number of peaks in cluster. - sum_score : float - Sum of all scores in cluster. - mean_score : float - Mean score of peaks in cluster. - sum_sigma : float - Sum of all sigma values in cluster. - mean_sigma : float - Mean sigma of peaks in cluster. - score_std : float - Standard deviation of scores in cluster. - avg_alpha : float - Score-weighted average alpha (radians). - avg_beta : float - Score-weighted average beta (radians). - avg_gamma : float - Score-weighted average gamma (radians). - angular_spread : float - Average angular distance of peaks from cluster center (degrees). - """ - alpha: float - beta: float - gamma: float - best_score: float - best_sigma: float - size: int - sum_score: float - mean_score: float - sum_sigma: float - mean_sigma: float - score_std: float - avg_alpha: float - avg_beta: float - avg_gamma: float - angular_spread: float - - def to_tuple(self) -> tuple: - """Return basic tuple (alpha, beta, gamma, score, sigma, size).""" - return (self.alpha, self.beta, self.gamma, self.best_score, self.best_sigma, self.size) - - def __repr__(self) -> str: - return ( - f"RotationCluster(α={np.degrees(self.alpha):.1f}°, " - f"β={np.degrees(self.beta):.1f}°, γ={np.degrees(self.gamma):.1f}°, " - f"score={self.best_score:.4f}, size={self.size}, " - f"sum_score={self.sum_score:.4f}, spread={self.angular_spread:.2f}°)" - ) - - -def _average_quaternions_weighted(quaternions: np.ndarray, weights: np.ndarray) -> np.ndarray: - """ - Compute weighted average of quaternions. - - Uses the eigenvector method for quaternion averaging. - - Parameters - ---------- - quaternions : np.ndarray - Array of quaternions [w, x, y, z], shape (N, 4). - weights : np.ndarray - Weights for each quaternion, shape (N,). - - Returns - ------- - avg_q : np.ndarray - Averaged quaternion [w, x, y, z], shape (4,). - """ - # Normalize weights - weights = weights / weights.sum() - - # Build weighted outer product sum - M = np.zeros((4, 4)) - for q, w in zip(quaternions, weights): - M += w * np.outer(q, q) - - # Eigenvector with largest eigenvalue is the average - eigenvalues, eigenvectors = np.linalg.eigh(M) - avg_q = eigenvectors[:, -1] # Largest eigenvalue - - # Ensure w >= 0 - if avg_q[0] < 0: - avg_q = -avg_q - - return avg_q - - -def _quaternion_to_euler_zyz(q: np.ndarray) -> Tuple[float, float, float]: - """Convert quaternion [w, x, y, z] to ZYZ Euler angles.""" - w, x, y, z = q - - # Convert to rotation matrix - R = np.array([ - [1 - 2*(y*y + z*z), 2*(x*y - w*z), 2*(x*z + w*y)], - [2*(x*y + w*z), 1 - 2*(x*x + z*z), 2*(y*z - w*x)], - [2*(x*z - w*y), 2*(y*z + w*x), 1 - 2*(x*x + y*y)] - ]) - - return rotation_matrix_to_euler_zyz(R) - - -def cluster_rotation_peaks( - peaks: list, - cluster_radius_deg: float = 5.0, - symmetry_matrices: Optional[np.ndarray] = None, - return_details: bool = False, -) -> list: - """ - Cluster rotation peaks based on rotation matrix distance. - - Uses greedy clustering: iterates through peaks (sorted by score), and - assigns each peak to an existing cluster if within the angular threshold, - otherwise creates a new cluster. Returns the best peak from each cluster. - - Parameters - ---------- - peaks : list - List of (alpha, beta, gamma, score, sigma) tuples. - Angles are in radians. - cluster_radius_deg : float - Maximum angular distance (in degrees) for peaks to be in the same cluster. - symmetry_matrices : np.ndarray, optional - If provided, considers symmetry-equivalent rotations when clustering. - Shape (N, 3, 3) or (N, 4, 4). - return_details : bool - If True, return list of RotationCluster objects with full statistics. - If False, return simple tuples (alpha, beta, gamma, score, sigma, size). - - Returns - ------- - clustered_peaks : list - If return_details=False: List of (alpha, beta, gamma, score, sigma, cluster_size) tuples. - If return_details=True: List of RotationCluster objects with aggregate statistics. - Sorted by best score (descending). - """ - if not peaks: - return [] - - cluster_radius_rad = np.radians(cluster_radius_deg) - - # Sort peaks by score (descending) - sorted_peaks = sorted(peaks, key=lambda p: p[3], reverse=True) - - # Each cluster stores: (best_peak, rotation_matrix, list_of_all_peaks) - clusters = [] - - for peak in sorted_peaks: - alpha, beta, gamma = peak[0], peak[1], peak[2] - R_peak = rotation_matrix_from_euler_zyz(alpha, beta, gamma) - - # Find closest cluster - best_cluster_idx = -1 - best_distance = float('inf') - - for i, (_, R_cluster, _) in enumerate(clusters): - # Direct distance - dist = _rotation_matrix_error(R_peak, R_cluster) - - # Check symmetry-equivalent rotations if provided - if symmetry_matrices is not None: - for S in symmetry_matrices: - S_rot = S[:3, :3] if S.shape[0] > 3 else S - # Check both S @ R_peak and R_peak @ S - for R_equiv in [S_rot @ R_peak, R_peak @ S_rot]: - d = _rotation_matrix_error(R_equiv, R_cluster) - dist = min(dist, d) - - if dist < best_distance: - best_distance = dist - best_cluster_idx = i - - if best_distance < cluster_radius_rad and best_cluster_idx >= 0: - # Add to existing cluster - best_peak, R_cluster, peak_list = clusters[best_cluster_idx] - peak_list.append(peak) - clusters[best_cluster_idx] = (best_peak, R_cluster, peak_list) - else: - # Create new cluster - clusters.append((peak, R_peak, [peak])) - - # Compute cluster statistics - clustered_peaks = [] - for best_peak, R_center, peak_list in clusters: - alpha, beta, gamma, score, sigma = best_peak[:5] - size = len(peak_list) - - # Basic statistics - scores = np.array([p[3] for p in peak_list]) - sigmas = np.array([p[4] for p in peak_list]) - - sum_score = scores.sum() - mean_score = scores.mean() - sum_sigma = sigmas.sum() - mean_sigma = sigmas.mean() - score_std = scores.std() if size > 1 else 0.0 - - # Compute weighted average rotation using quaternions - if size > 1: - quaternions = [] - weights = [] - for p in peak_list: - R_p = rotation_matrix_from_euler_zyz(p[0], p[1], p[2]) - q = rotation_matrix_to_quaternion(R_p) - quaternions.append(q) - weights.append(max(p[3], 0.0)) # Use score as weight (non-negative) - - quaternions = np.array(quaternions) - weights = np.array(weights) - - if weights.sum() > 0: - avg_q = _average_quaternions_weighted(quaternions, weights) - avg_alpha, avg_beta, avg_gamma = _quaternion_to_euler_zyz(avg_q) - else: - avg_alpha, avg_beta, avg_gamma = alpha, beta, gamma - else: - avg_alpha, avg_beta, avg_gamma = alpha, beta, gamma - - # Compute angular spread (average distance from center) - if size > 1: - distances = [] - for p in peak_list: - R_p = rotation_matrix_from_euler_zyz(p[0], p[1], p[2]) - dist = _rotation_matrix_error(R_p, R_center) - distances.append(dist) - angular_spread = np.degrees(np.mean(distances)) - else: - angular_spread = 0.0 - - if return_details: - cluster = RotationCluster( - alpha=alpha, - beta=beta, - gamma=gamma, - best_score=score, - best_sigma=sigma, - size=size, - sum_score=sum_score, - mean_score=mean_score, - sum_sigma=sum_sigma, - mean_sigma=mean_sigma, - score_std=score_std, - avg_alpha=avg_alpha, - avg_beta=avg_beta, - avg_gamma=avg_gamma, - angular_spread=angular_spread, - ) - clustered_peaks.append(cluster) - else: - clustered_peaks.append((alpha, beta, gamma, score, sigma, size)) - - # Sort by best score (descending) - if return_details: - clustered_peaks.sort(key=lambda c: c.best_score, reverse=True) - else: - clustered_peaks.sort(key=lambda p: p[3], reverse=True) - - return clustered_peaks - - -def cluster_rotation_peaks_torch( - peaks: list, - cluster_radius_deg: float = 5.0, - symmetry_matrices: Optional[torch.Tensor] = None, - return_details: bool = False, -) -> list: - """ - Torch wrapper for cluster_rotation_peaks. - - Parameters - ---------- - peaks : list - List of (alpha, beta, gamma, score, sigma) tuples. - cluster_radius_deg : float - Maximum angular distance for clustering. - symmetry_matrices : torch.Tensor, optional - Symmetry matrices, shape (N, 3, 3) or (N, 4, 4). - return_details : bool - If True, return RotationCluster objects with full statistics. - - Returns - ------- - clustered_peaks : list - If return_details=False: List of (alpha, beta, gamma, score, sigma, cluster_size) tuples. - If return_details=True: List of RotationCluster objects. - """ - sym_np = None - if symmetry_matrices is not None: - sym_np = symmetry_matrices.detach().cpu().numpy() - return cluster_rotation_peaks(peaks, cluster_radius_deg, sym_np, return_details) - - -# Torch convenience wrappers -def ball_rotation_search_torch( - E_obs: torch.Tensor, - s_obs: torch.Tensor, - E_calc: torch.Tensor, - s_calc: torch.Tensor, - **kwargs, -) -> Tuple[np.ndarray, tuple, list]: - """Torch wrapper for ball_rotation_search.""" - return ball_rotation_search( - E_obs.detach().cpu().numpy(), - s_obs.detach().cpu().numpy(), - E_calc.detach().cpu().numpy(), - s_calc.detach().cpu().numpy(), - **kwargs, - ) - - -# ============================================================================= -# Symmetry Reduction Functions -# ============================================================================= - -def rotation_matrix_to_quaternion(R: np.ndarray) -> np.ndarray: - """ - Convert a 3x3 rotation matrix to a unit quaternion [w, x, y, z]. - - Uses Shepperd's method for numerical stability. - - Parameters - ---------- - R : np.ndarray - Rotation matrix, shape (3, 3). - - Returns - ------- - q : np.ndarray - Unit quaternion [w, x, y, z], shape (4,). - """ - trace = np.trace(R) - - if trace > 0: - s = 0.5 / np.sqrt(trace + 1.0) - w = 0.25 / s - x = (R[2, 1] - R[1, 2]) * s - y = (R[0, 2] - R[2, 0]) * s - z = (R[1, 0] - R[0, 1]) * s - elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]: - s = 2.0 * np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) - w = (R[2, 1] - R[1, 2]) / s - x = 0.25 * s - y = (R[0, 1] + R[1, 0]) / s - z = (R[0, 2] + R[2, 0]) / s - elif R[1, 1] > R[2, 2]: - s = 2.0 * np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) - w = (R[0, 2] - R[2, 0]) / s - x = (R[0, 1] + R[1, 0]) / s - y = 0.25 * s - z = (R[1, 2] + R[2, 1]) / s - else: - s = 2.0 * np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) - w = (R[1, 0] - R[0, 1]) / s - x = (R[0, 2] + R[2, 0]) / s - y = (R[1, 2] + R[2, 1]) / s - z = 0.25 * s - - q = np.array([w, x, y, z]) - # Normalize - q = q / np.linalg.norm(q) - # Ensure w >= 0 for canonical form (q and -q represent same rotation) - if q[0] < 0: - q = -q - return q - - -def quaternion_to_euler_zyz(q: np.ndarray) -> Tuple[float, float, float]: - """ - Convert unit quaternion [w, x, y, z] to ZYZ Euler angles. - - Parameters - ---------- - q : np.ndarray - Unit quaternion [w, x, y, z], shape (4,). - - Returns - ------- - alpha, beta, gamma : float - ZYZ Euler angles in radians. - """ - w, x, y, z = q - - # First convert to rotation matrix - R = np.array([ - [1 - 2*(y*y + z*z), 2*(x*y - w*z), 2*(x*z + w*y)], - [2*(x*y + w*z), 1 - 2*(x*x + z*z), 2*(y*z - w*x)], - [2*(x*z - w*y), 2*(y*z + w*x), 1 - 2*(x*x + y*y)] - ]) - - return rotation_matrix_to_euler_zyz(R) - - -def rotation_matrix_to_euler_zyz(R: np.ndarray) -> Tuple[float, float, float]: - """ - Extract ZYZ Euler angles from rotation matrix. - - R = Rz(alpha) @ Ry(beta) @ Rz(gamma) - - Parameters - ---------- - R : np.ndarray - Rotation matrix, shape (3, 3). - - Returns - ------- - alpha, beta, gamma : float - ZYZ Euler angles in radians, with: - - alpha in [0, 2π) - - beta in [0, π] - - gamma in [0, 2π) - """ - # beta from R[2,2] = cos(beta) - cos_beta = np.clip(R[2, 2], -1.0, 1.0) - beta = np.arccos(cos_beta) - - sin_beta = np.sin(beta) - - if np.abs(sin_beta) > 1e-10: - # General case: sin(beta) != 0 - # alpha from R[0,2] = cos(alpha)*sin(beta), R[1,2] = sin(alpha)*sin(beta) - alpha = np.arctan2(R[1, 2], R[0, 2]) - # gamma from R[2,0] = -sin(beta)*cos(gamma), R[2,1] = sin(beta)*sin(gamma) - gamma = np.arctan2(R[2, 1], -R[2, 0]) - else: - # Gimbal lock: beta ≈ 0 or beta ≈ π - # Only (alpha + gamma) or (alpha - gamma) is determined - gamma = 0.0 - if cos_beta > 0: # beta ≈ 0 - alpha = np.arctan2(R[1, 0], R[0, 0]) - else: # beta ≈ π - alpha = np.arctan2(-R[1, 0], -R[0, 0]) - - # Normalize to [0, 2π) - alpha = alpha % (2 * np.pi) - gamma = gamma % (2 * np.pi) - - return alpha, beta, gamma - - -def reduce_rotation_by_symmetry( - alpha: float, - beta: float, - gamma: float, - symmetry_matrices: np.ndarray, -) -> Tuple[float, float, float, int]: - """ - Map rotation angles to a canonical representative using crystal symmetry. - - Given a rotation R and symmetry operations {S_i}, finds the canonical - representative among all symmetry-equivalent rotations {S_i @ R}. - The canonical form is chosen as the one with the smallest Euler angle - norm (closest to identity rotation), with all angles positive. - - Parameters - ---------- - alpha, beta, gamma : float - ZYZ Euler angles in radians. - symmetry_matrices : np.ndarray - Crystallographic symmetry rotation matrices, shape (N, 3, 3). - These should be the rotation parts of the space group operations. - - Returns - ------- - alpha_red, beta_red, gamma_red : float - Reduced ZYZ Euler angles in radians (all positive, smallest norm). - sym_index : int - Index of the symmetry operation that gave the canonical form. - """ - R = rotation_matrix_from_euler_zyz(alpha, beta, gamma) - - best_norm = float('inf') - best_angles = (alpha % (2 * np.pi), beta, gamma % (2 * np.pi)) - best_sym_idx = 0 - - for i, S in enumerate(symmetry_matrices): - # Extract 3x3 rotation part if needed - S_rot = S[:3, :3] if S.shape[0] > 3 else S - - # Apply symmetry: S @ R - R_equiv = S_rot @ R - - # Extract Euler angles (already positive from rotation_matrix_to_euler_zyz) - a, b, g = rotation_matrix_to_euler_zyz(R_equiv) - - # Compute norm of the angle vector - norm = np.sqrt(a**2 + b**2 + g**2) - - if norm < best_norm: - best_norm = norm - best_angles = (a, b, g) - best_sym_idx = i - - return (*best_angles, best_sym_idx) - - -def reduce_peaks_by_symmetry( - peaks: list, - symmetry_matrices: np.ndarray, -) -> list: - """ - Reduce a list of rotation peaks to canonical symmetry representatives. - - Parameters - ---------- - peaks : list - List of (alpha, beta, gamma, score, sigma) tuples. - Angles are in radians. - symmetry_matrices : np.ndarray - Crystallographic symmetry rotation matrices, shape (N, 3, 3). - - Returns - ------- - reduced_peaks : list - List of (alpha, beta, gamma, score, sigma, sym_index) tuples. - Angles are the canonical representatives in radians. - """ - reduced_peaks = [] - - for alpha, beta, gamma, score, sigma in peaks: - alpha_red, beta_red, gamma_red, sym_idx = reduce_rotation_by_symmetry( - alpha, beta, gamma, symmetry_matrices - ) - reduced_peaks.append((alpha_red, beta_red, gamma_red, score, sigma, sym_idx)) - - return reduced_peaks - - -def reduce_peaks_by_symmetry_torch( - peaks: list, - symmetry_matrices: torch.Tensor, -) -> list: - """ - Torch wrapper for reduce_peaks_by_symmetry. - - Parameters - ---------- - peaks : list - List of (alpha, beta, gamma, score, sigma) tuples. - symmetry_matrices : torch.Tensor - Symmetry matrices, shape (N, 3, 3) or (N, 4, 4). - - Returns - ------- - reduced_peaks : list - List of (alpha, beta, gamma, score, sigma, sym_index) tuples. - """ - sym_np = symmetry_matrices.detach().cpu().numpy() - return reduce_peaks_by_symmetry(peaks, sym_np) diff --git a/torchref/experimental/alignment/clashscore.py b/torchref/experimental/alignment/clashscore.py deleted file mode 100644 index 2dae5a7c..00000000 --- a/torchref/experimental/alignment/clashscore.py +++ /dev/null @@ -1,403 +0,0 @@ -""" -Clash score calculator for crystallographic alignment. - -Provides clash-based scoring to complement Patterson alignment by detecting -steric clashes between symmetry-related molecules. - -Experimental / unstable API: part of ``torchref.experimental.alignment``, -the opt-in ball-harmonic MR engine. The production MR entry point is -``torchref.alignment`` (the consolidated FRF engine). Signatures and behavior -may change without notice. -""" - -from typing import TYPE_CHECKING, List, Optional - -import torch -import torch.nn as nn - -from torchref.base.coordinates import ( - cartesian_to_fractional_torch, - fractional_to_cartesian_torch, -) -from torchref.config import get_default_device, get_float_dtype -from torchref.symmetry import Cell, SpaceGroup -from torchref.symmetry.spacegroup import SpaceGroupLike -from torchref.utils.device_mixin import DeviceMixin - -from .transform import RigidTransform - -if TYPE_CHECKING: - from torchref.model.model import Model - - -class AtomSampler: - """ - Select representative atoms for clash checking. - - Provides methods to create atom selection masks based on the type of - molecular structure (protein vs. ligand-containing). - """ - - @staticmethod - def from_model( - model: "Model", - mode: str = "auto", - ) -> torch.Tensor: - """ - Create atom selection mask from a Model. - - Parameters - ---------- - model : Model - The crystallographic model containing atomic data. - mode : str, default 'auto' - Selection mode: - - 'auto': Use CA atoms if protein present (ATOM records), - else all atoms (for small molecules with only HETATM) - - 'ca_only': Only CA atoms (alpha carbons) - - 'all_atoms': All atoms in the structure - - Returns - ------- - torch.Tensor - Boolean mask of shape (n_atoms,) indicating which atoms to use. - - Examples - -------- - :: - - from torchref.model import Model - model = Model().load_pdb('protein.pdb') - mask = AtomSampler.from_model(model, mode='auto') - print(f"Selected {mask.sum()} atoms out of {len(mask)}") - """ - pdb = model.pdb - - if mode == "auto": - # Check if structure has normal ATOM records (protein/nucleic acid) - has_atom = (pdb["ATOM"] == "ATOM").any() - if has_atom: - # Has protein/nucleic acid: use CA atoms for efficiency - return torch.tensor((pdb["name"] == "CA").values, dtype=torch.bool) - else: - # Only HETATM (small molecule): use all atoms - return torch.ones(len(pdb), dtype=torch.bool) - elif mode == "ca_only": - return torch.tensor((pdb["name"] == "CA").values, dtype=torch.bool) - elif mode == "all_atoms": - return torch.ones(len(pdb), dtype=torch.bool) - else: - raise ValueError( - f"Unknown mode '{mode}'. Use 'auto', 'ca_only', or 'all_atoms'." - ) - - -class ClashScoreCalculator(DeviceMixin, nn.Module): - """ - Calculate clash scores between symmetry-related molecules. - - Computes steric clash violations between an asymmetric unit (ASU) and - its symmetry-related copies. Automatically filters symmetry mates based - on the actual input coordinates to only consider those that can potentially - clash. - - Uses a steep **4 penalty: (radius² - dist²)² which rises quickly as - atoms get closer than the clash radius. - - Parameters - ---------- - symmetry : str, int, gemmi.SpaceGroup, or SpaceGroup - Space group specification for symmetry expansion. - default_clash_radius : float, default 5.0 - Minimum allowed distance between atoms stored on the instance. Note: - ``forward`` does NOT read this attribute; it uses its own - ``clash_radius`` argument (also defaulting to 5.0), so this stored value - is currently unused and ``forward(clash_radius=...)`` is authoritative. - dtype : torch.dtype, default torch.float32 - Data type for computations. - device : torch.device, default 'cpu' - Device for computations. - - Examples - -------- - :: - - from torchref.experimental.alignment.clashscore import ClashScoreCalculator, AtomSampler - from torchref.model import Model - - model = Model().load_pdb('structure.pdb') - calc = ClashScoreCalculator(symmetry=model.spacegroup) - mask = AtomSampler.from_model(model) - score = calc(xyz=model.xyz(), cell=model.cell, atom_mask=mask) - print(f"Clash score: {score.item():.4f}") - """ - - # Cell offsets for neighboring cells (7 cells: central + 6 face neighbors) - _cell_offsets = [ - (0, 0, 0), # Central - (-1, 0, 0), - (1, 0, 0), # x neighbors - (0, -1, 0), - (0, 1, 0), # y neighbors - (0, 0, -1), - (0, 0, 1), # z neighbors - ] - - def __init__( - self, - symmetry: SpaceGroupLike, - default_clash_radius: float = 5.0, - dtype: torch.dtype = None, - device: torch.device = None, - ): - super().__init__() - if dtype is None: - dtype = get_float_dtype() - if device is None: - device = get_default_device() - self.default_clash_radius = default_clash_radius - self.dtype = dtype - self._device = device - - # Initialize symmetry handler. Use the user-configured dtype so the - # SpaceGroup stays MPS-compatible; ``_get_valid_transforms`` casts to - # CPU+float64 internally where high-precision symmetry math is needed. - if isinstance(symmetry, SpaceGroup): - self.symmetry = symmetry - else: - self.symmetry = SpaceGroup(symmetry, dtype=self.dtype, device=device) - - def _get_valid_transforms( - self, - cell: torch.Tensor, - centroid_frac: torch.Tensor, - molecule_radius: float, - clash_radius: float, - ) -> List[RigidTransform]: - """ - Compute which symmetry mates can potentially clash with the ASU. - - Parameters - ---------- - cell : torch.Tensor - Unit cell parameters [a, b, c, alpha, beta, gamma]. - centroid_frac : torch.Tensor - Fractional coordinates of molecule centroid. - molecule_radius : float - Approximate radius of the molecule in Angstroms. - clash_radius : float - Minimum allowed distance between atoms. - - Returns - ------- - List[RigidTransform] - List of symmetry transforms that could produce clashes. - - Notes - ----- - The 7 cell offsets in ``_cell_offsets`` (central + 6 face neighbours) - are crossed with every symmetry operation, so up to ``7 * n_ops`` - candidate transforms are examined, not just 7 neighbours. - """ - # Compute fractionalization matrix using Cell. Pin to CPU because the - # rest of this method does CPU-only float64 symmetry math. - cell_obj = Cell(cell, dtype=torch.float64, device="cpu") - B = cell_obj.fractional_matrix - - centroid_frac = centroid_frac.to(device="cpu", dtype=torch.float64) - - # Threshold distance for filtering - # Two molecules can clash if centroid distance < 2*radius + clash_radius - threshold = 2 * molecule_radius + clash_radius - - # Identity matrix for comparison with rotation matrices - I = torch.eye(3, dtype=torch.float64) - - valid_transforms = [] - n_ops = self.symmetry.n_ops - - for op_idx in range(n_ops): - R = self.symmetry.matrices[op_idx].cpu().to(torch.float64) - t = self.symmetry.translations[op_idx].cpu().to(torch.float64) - - for offset in self._cell_offsets: - # Skip identity operation in central cell (self-interaction) - if op_idx == 0 and offset == (0, 0, 0): - continue - - offset_tensor = torch.tensor(offset, dtype=torch.float64) - - # Displacement between ASU centroid and this symmetry mate's centroid - # Symmetry mate position: R @ x + t + offset - # Displacement from ASU (identity, no offset): (R - I) @ centroid + t + offset - d_frac = (R - I) @ centroid_frac + t + offset_tensor - - # Convert to Cartesian distance - d_cart = B @ d_frac - dist = d_cart.norm().item() - - if dist < threshold: - # Create RigidTransform for this symmetry operation in fractional space - t_total = t + offset_tensor - transform = RigidTransform.from_matrix(R, t_total) - valid_transforms.append(transform) - - return valid_transforms - - def forward( - self, - xyz: torch.Tensor, - cell: torch.Tensor, - atom_mask: Optional[torch.Tensor] = None, - clash_radius: float = 5.0, - ) -> torch.Tensor: - """ - Compute clash score for given coordinates. - - Automatically determines which symmetry mates could clash based on - the actual input coordinates. Uses squared distances for efficiency - and a steep **4 penalty for clashes. - - Parameters - ---------- - xyz : torch.Tensor - Cartesian coordinates of shape (N, 3). - cell : torch.Tensor - Unit cell parameters [a, b, c, alpha, beta, gamma]. - atom_mask : torch.Tensor, optional - Boolean mask of shape (N,) selecting atoms to use. - If None, all atoms are used. - clash_radius : float, default 5.0 - Minimum allowed distance between atoms in Angstroms. - Atoms closer than this will contribute to the clash score. - - Returns - ------- - torch.Tensor - Scalar clash score. Lower values indicate fewer clashes. - Zero indicates no clashes within the clash radius. - """ - device = xyz.device - dtype = xyz.dtype - - # Apply atom mask - if atom_mask is not None: - atom_mask = atom_mask.to(device) - xyz_selected = xyz[atom_mask] - else: - xyz_selected = xyz - - n_atoms = xyz_selected.shape[0] - - if n_atoms == 0: - return torch.tensor(0.0, device=device, dtype=dtype) - - # Compute molecule properties from actual coordinates - # Use Cell object to get transformation matrices - cell_obj = Cell(cell, dtype=dtype, device=device) - B = cell_obj.fractional_matrix - B_inv = cell_obj.inv_fractional_matrix - - centroid = xyz_selected.mean(dim=0) - molecule_radius = (xyz_selected - centroid).norm(dim=1).max().item() - centroid_frac = cartesian_to_fractional_torch( - centroid.unsqueeze(0), cell_obj.data, B_inv - ).squeeze(0) - - # Get valid transforms for these coordinates - valid_transforms = self._get_valid_transforms( - cell=cell, - centroid_frac=centroid_frac, - molecule_radius=molecule_radius, - clash_radius=clash_radius, - ) - - # If no valid mates after filtering, no clashes possible - if len(valid_transforms) == 0: - return torch.tensor(0.0, device=device, dtype=dtype) - - # Convert ASU to fractional coordinates - xyz_frac = cartesian_to_fractional_torch(xyz_selected, cell_obj.data, B_inv) - - # Precompute squared clash radius threshold - clash_radius_sq = clash_radius**2 - - # Accumulate score over valid symmetry mates - total_score = torch.tensor(0.0, device=device, dtype=dtype) - n_clashes = 0 - - for transform in valid_transforms: - # Apply symmetry operation in fractional space using RigidTransform - xyz_mate_frac = transform.apply(xyz_frac) - - # Convert back to Cartesian - xyz_mate_cart = fractional_to_cartesian_torch( - xyz_mate_frac, cell_obj.data, B - ) - - # Compute squared pairwise distances (more efficient, no sqrt) - diff = xyz_selected.unsqueeze(1) - xyz_mate_cart.unsqueeze(0) # (N, N, 3) - dists_sq = (diff**2).sum(dim=-1) # (N, N) - - # Compute violations: (radius² - dist²), clipped to 0 - # This gives a steep penalty that increases rapidly as atoms get closer - violations_sq = torch.clamp(clash_radius_sq - dists_sq, min=0.0) - - # Apply **2 to squared violations = **4 penalty on distance violation - # (radius² - dist²)² rises steeply as dist approaches 0 - total_score = total_score + (violations_sq**2).sum() - n_clashes += (violations_sq > 0).sum().item() - - # Normalize by number of atom pairs checked - n_pairs = n_atoms * n_atoms * len(valid_transforms) - if n_pairs > 0: - total_score = total_score / n_pairs - - return total_score - - -def compute_clash_score( - model: "Model", - mode: str = "auto", - clash_radius: float = 5.0, -) -> torch.Tensor: - """ - Convenience function to compute clash score for a model. - - Parameters - ---------- - model : Model - Crystallographic model with coordinates and symmetry. - mode : str, default 'auto' - Atom selection mode ('auto', 'ca_only', 'all_atoms'). - clash_radius : float, default 5.0 - Minimum allowed distance between atoms in Angstroms. - - Returns - ------- - torch.Tensor - Scalar clash score. - - Examples - -------- - :: - - from torchref.model import Model - from torchref.experimental.alignment.clashscore import compute_clash_score - model = Model().load_pdb('structure.pdb') - score = compute_clash_score(model) - print(f"Clash score: {score.item():.4f}") - """ - calc = ClashScoreCalculator( - symmetry=model.spacegroup, - device=model.device, - ) - - atom_mask = AtomSampler.from_model(model, mode=mode) - - return calc( - xyz=model.xyz(), - cell=model.cell, - atom_mask=atom_mask, - clash_radius=clash_radius, - ) diff --git a/torchref/experimental/alignment/distributions.py b/torchref/experimental/alignment/distributions.py deleted file mode 100644 index 47acd06d..00000000 --- a/torchref/experimental/alignment/distributions.py +++ /dev/null @@ -1,321 +0,0 @@ -""" -Statistical distributions for Maximum Likelihood molecular replacement. - -This module provides numerically stable implementations of the probability -distributions used in crystallographic ML target functions: - -- Rice distribution for acentric reflections -- Woolfson (folded normal) distribution for centric reflections -- Stable log-Bessel I_0 computation for large arguments - -The key numerical challenge is computing log(I_0(x)) for large x (up to ~10000), -which requires an asymptotic expansion to avoid overflow. - -Experimental / unstable API: part of ``torchref.experimental.alignment``, -the opt-in ball-harmonic MR engine. The production MR entry point is -``torchref.alignment`` (the consolidated FRF engine). Signatures and behavior -may change without notice. - -References ----------- -- Read, R.J. (2001). Pushing the boundaries of molecular replacement with - maximum likelihood. Acta Cryst. D57, 1373-1382. -- McCoy et al. (2007). Phaser crystallographic software. J. Appl. Cryst. 40, 658-674. -""" - -import math - -import torch - - -def stable_log_bessel_i0(x: torch.Tensor) -> torch.Tensor: - """ - Compute log(I_0(x)) with numerical stability for large x. - - The modified Bessel function I_0(x) grows exponentially, so direct - computation overflows for x > ~700. This implementation uses: - - For small x (< 50): log(I_0e(x)) + |x| using torch.special.i0e - - For large x (>= 50): asymptotic expansion x - 0.5*log(2*pi*x) - - Parameters - ---------- - x : torch.Tensor - Input tensor of non-negative values. - - Returns - ------- - torch.Tensor - log(I_0(x)) for each element, same shape as input. - - Notes - ----- - The asymptotic expansion is: - log(I_0(x)) ~ x - 0.5*log(2*pi*x) + 1/(8x) - 1/(128x^2) + ... - - For x >= 50, the leading two terms give excellent accuracy. - For very large x (> 10000), higher order terms can be added if needed. - - Examples - -------- - :: - - x = torch.tensor([0.1, 10.0, 100.0, 1000.0]) - log_i0 = stable_log_bessel_i0(x) - # torch.all(torch.isfinite(log_i0)) is True - """ - # Ensure non-negative input - x_abs = torch.abs(x) - - # Initialize output - result = torch.zeros_like(x) - - # Small x regime: use i0e which is I_0(x) * exp(-|x|) - # So log(I_0(x)) = log(I_0e(x)) + |x| - small_mask = x_abs < 50.0 - if small_mask.any(): - x_small = x_abs[small_mask] - i0e_val = torch.special.i0e(x_small) - # Guard against i0e returning 0 for very small x - i0e_val = torch.clamp(i0e_val, min=1e-300) - result[small_mask] = torch.log(i0e_val) + x_small - - # Large x regime: asymptotic expansion - # log(I_0(x)) ~ x - 0.5*log(2*pi*x) + 1/(8x) - 1/(128x^2) + ... - large_mask = ~small_mask - if large_mask.any(): - x_large = x_abs[large_mask] - # Leading terms of asymptotic expansion - log_i0_asymp = ( - x_large - - 0.5 * torch.log(2.0 * math.pi * x_large) - + 1.0 / (8.0 * x_large) - - 1.0 / (128.0 * x_large * x_large) - ) - result[large_mask] = log_i0_asymp - - return result - - -def rice_log_likelihood( - F_obs: torch.Tensor, - F_mean: torch.Tensor, - variance: torch.Tensor, -) -> torch.Tensor: - """ - Log-likelihood for acentric reflections (Rice distribution). - - The Rice distribution describes the distribution of |F_obs| given - the expected |F_calc| and variance for acentric reflections: - - p(F_obs | F_mean, sigma^2) = (F_obs / sigma^2) * - exp(-(F_obs^2 + F_mean^2) / (2*sigma^2)) * I_0(F_obs*F_mean / sigma^2) - - Parameters - ---------- - F_obs : torch.Tensor - Observed structure factor amplitudes |F_obs|. - F_mean : torch.Tensor - Expected structure factor amplitudes D * |F_calc|. - variance : torch.Tensor - Variance parameter sigma^2 = epsilon * (Sigma_N - D^2 * <|F_calc|^2>). - - Returns - ------- - torch.Tensor - Log-likelihood for each reflection. - - Notes - ----- - The log-likelihood is: - log p = log(F_obs) - log(variance) - (F_obs^2 + F_mean^2)/(2*variance) - + log(I_0(F_obs * F_mean / variance)) - - Numerical stability is ensured by using stable_log_bessel_i0. - """ - # Numerical guards - variance_safe = torch.clamp(variance, min=1e-8) - F_obs_safe = torch.clamp(F_obs, min=1e-10) - F_mean_safe = torch.clamp(F_mean, min=0.0) - - # Compute Bessel argument - bessel_arg = F_obs_safe * F_mean_safe / variance_safe - - # Log-likelihood components - log_likelihood = ( - torch.log(F_obs_safe) - - torch.log(variance_safe) - - (F_obs_safe**2 + F_mean_safe**2) / (2.0 * variance_safe) - + stable_log_bessel_i0(bessel_arg) - ) - - return log_likelihood - - -def woolfson_log_likelihood( - F_obs: torch.Tensor, - F_mean: torch.Tensor, - variance: torch.Tensor, -) -> torch.Tensor: - """ - Log-likelihood for centric reflections (folded normal / Woolfson distribution). - - For centric reflections, the phase is restricted to 0 or pi, so the - distribution becomes a folded normal (Woolfson distribution): - - p(F_obs | F_mean, sigma^2) = (2/sigma) * (2*pi)^(-0.5) * - cosh(F_obs * F_mean / sigma^2) * exp(-(F_obs^2 + F_mean^2)/(2*sigma^2)) - - Parameters - ---------- - F_obs : torch.Tensor - Observed structure factor amplitudes |F_obs|. - F_mean : torch.Tensor - Expected structure factor amplitudes D * |F_calc|. - variance : torch.Tensor - Variance parameter sigma^2. Used directly as supplied; this function - applies no factor-of-2 rescaling, so any centric/acentric scaling - convention must be handled by the caller. - - Returns - ------- - torch.Tensor - Log-likelihood for each centric reflection. - - Notes - ----- - The log-likelihood is: - log p = log(2) - 0.5*log(2*pi*variance) - (F_obs^2 + F_mean^2)/(2*variance) - + log(cosh(F_obs * F_mean / variance)) - - For numerical stability with large arguments: - log(cosh(x)) ~ |x| - log(2) for |x| > ~20 - """ - # Numerical guards - variance_safe = torch.clamp(variance, min=1e-8) - F_obs_safe = torch.clamp(F_obs, min=1e-10) - F_mean_safe = torch.clamp(F_mean, min=0.0) - - sigma = torch.sqrt(variance_safe) - - # Compute argument for cosh - cosh_arg = F_obs_safe * F_mean_safe / variance_safe - - # Stable log(cosh(x)): use |x| - log(2) for large x - log_cosh = torch.where( - cosh_arg < 20.0, - torch.log(torch.cosh(cosh_arg)), - torch.abs(cosh_arg) - math.log(2.0), - ) - - # Log-likelihood components - log_likelihood = ( - math.log(2.0) - - 0.5 * torch.log(2.0 * math.pi * variance_safe) - - (F_obs_safe**2 + F_mean_safe**2) / (2.0 * variance_safe) - + log_cosh - ) - - return log_likelihood - - -def combined_log_likelihood( - F_obs: torch.Tensor, - F_mean: torch.Tensor, - variance: torch.Tensor, - centric_flags: torch.Tensor, -) -> torch.Tensor: - """ - Combined log-likelihood dispatching to Rice/Woolfson based on centric flags. - - Parameters - ---------- - F_obs : torch.Tensor - Observed structure factor amplitudes |F_obs|. - F_mean : torch.Tensor - Expected structure factor amplitudes D * |F_calc|. - variance : torch.Tensor - Variance parameters for each reflection. - centric_flags : torch.Tensor - Boolean mask, True for centric reflections. - - Returns - ------- - torch.Tensor - Log-likelihood for each reflection using the appropriate distribution. - - Examples - -------- - :: - - F_obs = torch.tensor([10.0, 20.0, 15.0]) - F_mean = torch.tensor([9.5, 18.0, 14.0]) - variance = torch.tensor([5.0, 8.0, 6.0]) - centric = torch.tensor([False, True, False]) - ll = combined_log_likelihood(F_obs, F_mean, variance, centric) - """ - # Initialize with Rice likelihood (acentric default) - log_likelihood = rice_log_likelihood(F_obs, F_mean, variance) - - # Replace centric reflections with Woolfson likelihood - if centric_flags.any(): - centric_ll = woolfson_log_likelihood( - F_obs[centric_flags], - F_mean[centric_flags], - variance[centric_flags], - ) - log_likelihood[centric_flags] = centric_ll - - return log_likelihood - - -def acentric_pdf( - F_obs: torch.Tensor, - F_mean: torch.Tensor, - variance: torch.Tensor, -) -> torch.Tensor: - """ - Probability density for acentric reflections (Rice distribution). - - This is the PDF (not log) for cases where the actual probability is needed. - For numerical reasons, prefer rice_log_likelihood when possible. - - Parameters - ---------- - F_obs : torch.Tensor - Observed structure factor amplitudes. - F_mean : torch.Tensor - Expected structure factor amplitudes. - variance : torch.Tensor - Variance parameter. - - Returns - ------- - torch.Tensor - Probability density for each reflection. - """ - return torch.exp(rice_log_likelihood(F_obs, F_mean, variance)) - - -def centric_pdf( - F_obs: torch.Tensor, - F_mean: torch.Tensor, - variance: torch.Tensor, -) -> torch.Tensor: - """ - Probability density for centric reflections (Woolfson distribution). - - Parameters - ---------- - F_obs : torch.Tensor - Observed structure factor amplitudes. - F_mean : torch.Tensor - Expected structure factor amplitudes. - variance : torch.Tensor - Variance parameter. - - Returns - ------- - torch.Tensor - Probability density for each reflection. - """ - return torch.exp(woolfson_log_likelihood(F_obs, F_mean, variance)) diff --git a/torchref/experimental/alignment/frf/__init__.py b/torchref/experimental/alignment/frf/__init__.py new file mode 100644 index 00000000..991e1896 --- /dev/null +++ b/torchref/experimental/alignment/frf/__init__.py @@ -0,0 +1,39 @@ +"""Fast Rotation Function — single, validated implementation. + +Phaser-faithful engine: chunked Bessel-SH expansion, stable Wigner-d +(``wigner_d.small_d_stable``), resolution↔bandwidth coupling +(``phaser_lmax_resolution``, default cap=48), dense P1-box calc, all under +``no_grad``. Solved the high-symmetry cases that broke the earlier ball +and Phaser-mimic engines (4BX9 342→4–7, 6G9X 77→1–4). + +Shared leaf math (``..sh``, ``..wigner``) lives in the parent ``alignment`` +package; this sub-package imports it "up". +""" +from .api import FastRotationFunction, phaser_lmax_resolution +from .dense_calc import dense_calc_via_box, model_sf_abs +from .rotation_utils import ( + edmonds_euler_from_rotation_matrix, + rotation_angular_distance_deg, + rotation_matrix_from_edmonds_euler, +) +from .types import ( + AdaptiveRotationFunction, + BesselSHCoefficients, + RotationPeak, +) + +__all__ = [ + # Engine + "FastRotationFunction", + "phaser_lmax_resolution", + "dense_calc_via_box", + "model_sf_abs", + # Rotation geometry helpers + "rotation_matrix_from_edmonds_euler", + "edmonds_euler_from_rotation_matrix", + "rotation_angular_distance_deg", + # Types + "AdaptiveRotationFunction", + "BesselSHCoefficients", + "RotationPeak", +] diff --git a/torchref/experimental/alignment/frf/_backends.py b/torchref/experimental/alignment/frf/_backends.py new file mode 100644 index 00000000..35df4ae0 --- /dev/null +++ b/torchref/experimental/alignment/frf/_backends.py @@ -0,0 +1,64 @@ +"""Dispatch policy for the Legendre-recurrence-and-shell-accumulation stage. + +One row per kernel, with every criterion for choosing it as a field. See +:mod:`torchref.utils.backends` for what each field means, and +:mod:`torchref.base.electron_density._backends` for the table this follows. + +Both kernels take the same arguments:: + + (Tr, Ti, rep_cos, rep_sin, Dr, Di, shell, a_coef, b_coef, sect) + +and both accumulate into ``Tr``/``Ti`` in place, which is what lets the dispatch +site be a single call. +""" + +from __future__ import annotations + +import torch + +from torchref.utils.backends import Backend, BackendTable + +_CPU = "torchref.experimental.alignment.frf.kernels.cpu.legendre_shell" +_PORTABLE = "torchref.experimental.alignment.frf.kernels.portable" + +#: Argument positions carrying the device/dtype contract: the two accumulators and +#: the four per-cluster float arrays. ``shell`` is int64 and the three coefficient +#: tables are built by the caller at the working dtype, so probing them would only +#: restate what the caller already chose. +_FLOAT_ARGS = (0, 1, 2, 3, 4, 5) + +LEGENDRE_BACKENDS = BackendTable( + name="FRF Legendre/shell accumulation", + backends=( + Backend( + name="cpu_fused", + kernel=(_CPU, "legendre_shell_accumulate", "legendre_shell_accumulate"), + device="cpu", + dtypes=(torch.float32,), # dtype-ok: backend capability declaration + # The kernel reads every array through a raw `float*`, so a mixed-dtype + # call would reinterpret the buffer rather than convert it. The gate + # keeps that from reaching the kernel; the kernel checks it too, since + # a table row is easier to widen by accident than a TORCH_CHECK. + require_uniform_dtype=True, + probes=_FLOAT_ARGS, + probe=(_CPU, "why_unavailable"), + # A compiler is not guaranteed on every host, and a missing one is a + # performance problem, not an outage. + expect_available="never", + # NOT "degrade": the kernel accumulates into Tr/Ti as it goes, so a + # mid-run failure leaves them partly written and the portable path + # would add its own contributions on top. There is nothing to fall + # back to once it has started. + on_failure="raise", + second_order=False, + ), + Backend( + name="portable", + kernel=(_PORTABLE, "legendre_shell_accumulate", + "legendre_shell_accumulate"), + second_order=False, + ), + ), +) + +__all__ = ["LEGENDRE_BACKENDS"] diff --git a/torchref/experimental/alignment/frf/api.py b/torchref/experimental/alignment/frf/api.py new file mode 100644 index 00000000..c77be474 --- /dev/null +++ b/torchref/experimental/alignment/frf/api.py @@ -0,0 +1,426 @@ +"""The fast rotation function engine: obs-side preprocessing, then scoring. + +Pipeline (mirrors Phaser ``run_FRF()``): + 1. Resolution mask (both sides). + 2. Wilson normalisation, optionally + French-Wilson + DFAC on obs. + 3. Build LERF1 obs intensity. + 4. Optional per-shell variance reweight on obs intensity. + 5. σA Eterm on calc intensity. + 6. Detect ZSYMM → m-symmetry filter on obs SH coefficients. + 7. Bessel-SH expand both sides. + 8. Cross-correlate on the radial axis → ξ_{lmn}. + 9. Per-β fixed-shape FFT + adaptive sample list bilinear interp → + adaptive rotation function. + 10. Greedy SO(3) NMS peak finding. +""" +from __future__ import annotations + +import math +import warnings +from typing import List, Optional, Tuple + +import torch + +from torchref.scaling.weighting import ( + DEFAULT_SNR_CAP, DEFAULT_TRUST_CAP, empirical_sigma_a, + information_weight, inverse_variance_weight, normalise_weight, + snr_from_amplitude, +) +from torchref.scaling import WilsonNormaliser +from .data_mr import bessel_sh_expand, cross_correlate_xi +from .peak_finder import find_rotation_peaks +from .preprocessing import ( + apply_shell_variance_weights, + build_lerf1_intensity, + detect_zsymm, + eterm_sigma_a, +) +from .sitelist_ang import evaluate_rotation_function +from .types import AdaptiveRotationFunction, RotationPeak + +__all__ = ["FastRotationFunction", "phaser_lmax_resolution"] + +#: Chebyshev order of the Wilson fit, both sides. Provisional -- inherited from +#: :mod:`torchref.scaling.wilson`, and never screened against a metric sensitive +#: to it. Screen it on the fit's own residual trend, not on truth rank: the +#: rotation function is a correlation, where any per-resolution scaling cancels. +WILSON_N_COEFF = 6 + + +def phaser_lmax_resolution( + model_radius_A: float, + d_min_data: float, + lmax_cap: int = 64, +): + """Couple the spherical-harmonic bandwidth to the rotation-function + resolution, as Phaser does (``runMR_FRF.cc:427-448``):: + + sphereOuter = 2 * mean_radius + LMAX = ceil(2 pi sphereOuter / HIRES) # HIRES = the data's d_min + if LMAX is odd: LMAX += 1 + LMAX = min(LMAX, lmax_cap) + LMAX_RESO = (LMAX hit the cap) ? 2 pi sphereOuter / LMAX : HIRES + + Data finer than the bandwidth can represent contributes aliasing rather than + signal -- the discrete ``Y_lm`` are not orthogonal over scattered + reflections -- and that background buries the symmetry-diluted true peak on + large or high-symmetry structures. So either the bandwidth rises to meet the + resolution, or, once it is capped, the resolution is coarsened to + ``LMAX_RESO`` and the finer reflections are dropped (``DataMR.cc:984``). + + For real protein search models ``ceil(2 pi 2r / d_min)`` is 100 to 170, so + the cap binds and this reduces to: use ``L = lmax_cap``, and keep only data + coarser than ``2 pi (2r) / lmax_cap`` -- the finest resolution that + bandwidth can carry for a molecule of that radius. Bigger molecule, coarser + cutoff. The variable-``L`` branch matters only for small models or + low-resolution data. + + Parameters + ---------- + model_radius_A : float + Mean distance of the model's atoms from its centroid (Angstrom). + ``sphereOuter = 2 * model_radius_A``. + d_min_data : float + High-resolution limit of the data (Angstrom). + lmax_cap : int + Bandwidth ceiling. The production value lives in + :data:`torchref.experimental.alignment.rotation_search.LMAX_CAP`, which + carries the measurement behind it; this default exists only for direct + callers. Cost grows about as ``l^3`` in time and ``l^2`` in memory. + + Returns + ------- + (L, d_min_eff) : Tuple[int, float] + ``L`` is the bandwidth in this package's convention (``lmax = L - 1``, + even); ``d_min_eff`` is the resolution to expand at. + """ + sphere_outer = 2.0 * float(model_radius_A) + lmax = int(math.ceil(2.0 * math.pi * sphere_outer / float(d_min_data))) + if lmax % 2 != 0: + lmax += 1 + lmax = min(lmax, int(lmax_cap)) + if lmax >= int(lmax_cap): + d_min_eff = 2.0 * math.pi * sphere_outer / lmax # coarsen to match cap + else: + d_min_eff = float(d_min_data) + if lmax >= 256: + warnings.warn( + f"phaser_lmax_resolution chose lmax={lmax} (>=256): the Wigner " + f"contraction is numerically fine here, but its cost grows about as " + f"l^3 in time and l^2 in memory. Consider a tighter lmax_cap.", + RuntimeWarning, stacklevel=2, + ) + return lmax + 1, d_min_eff # L = lmax + 1 (our bandwidth convention) + + +def _resolution_mask( + s_vec: torch.Tensor, + extra: Tuple[torch.Tensor, ...], + d_min: Optional[float], + d_max: Optional[float], +): + smag = s_vec.norm(dim=-1) + lo = 1.0 / d_max if d_max is not None else 0.0 + hi = 1.0 / d_min if d_min is not None else float("inf") + keep = (smag >= lo) & (smag <= hi) + return s_vec[keep], tuple(e[keep] for e in extra), smag[keep] + + +class FastRotationFunction: + """Reusable obs-side preprocessor + SH expansion. + + Instantiate once per (obs reflections, sym, config) tuple, then + call ``score_model(s_calc, F_calc)`` for each candidate model. + """ + + def __init__( + self, + s_obs: torch.Tensor, + F_obs: torch.Tensor, + centric_obs: torch.Tensor, + sym_mats: Optional[torch.Tensor], + *, + L: int = 24, + d_min: Optional[float] = None, + d_max: Optional[float] = None, + delta_vrms_A: float = 1.0, + n_wilson_shells: int = 20, + sig_F_obs: Optional[torch.Tensor] = None, + grid_sampling_deg: float = 2.0, + asu_idx: Optional[torch.Tensor] = None, + s_mag_asu: Optional[torch.Tensor] = None, + wilson_n_coeff: int = WILSON_N_COEFF, + obs_weight: str = "inverse_variance", + snr_cap: float = DEFAULT_SNR_CAP, + trust_cap: float = DEFAULT_TRUST_CAP, + shell_variance_weights: bool = False, + sym_cart: Optional[torch.Tensor] = None, + ): + self.device = s_obs.device + # The point group as Cartesian rotations, for the peak finder: with it, + # the returned peaks are distinct orientations rather than an + # orientation and its mates. `sym_mats` above is in the fractional + # basis and only detects the z-axis order; a direct caller without a + # cell cannot supply this and gets the plain suppression. + self.sym_cart = sym_cart + + # `L` and `d_min` arrive already coupled: the caller runs + # `phaser_lmax_resolution` because it needs the same pair to size the + # dense calc box, so re-deriving them here would only rediscover what it + # computed a line earlier. `d_min` is therefore the *coarsened* limit -- + # the resolution this bandwidth can represent -- not the data's own. + self.L = L + self.d_min = d_min + self.d_max = d_max + self.delta_vrms_A = delta_vrms_A + self.n_wilson_shells = n_wilson_shells + self.grid_sampling_deg = grid_sampling_deg + # How much each observation counts. Separate from the convention above, + # which decides only what is compared -- see torchref.scaling.weighting. + # "information" is the saturating I/sigma weight; "none" is unit weight, + # which is the control that says what the weighting is worth at all. + self.obs_weight = obs_weight + self.snr_cap = float(snr_cap) + self.trust_cap = float(trust_cap) + # Off by default. It is a PER-SHELL weight, and per-shell weights are + # absorbed: it renormalises whatever scale the convention produced, which + # is precisely why every E convention measured as gauge. With the + # resolution weighting now declared on the calc side, this is a second + # mechanism for a job that already has one. + self.shell_variance_weights = bool(shell_variance_weights) + + # Chebyshev order of the Wilson fit. Both sides are normalised by the + # same class at the same order over the same abscissa -- that is what + # puts them on a common footing, and it is the reason Sigma_obs/Sigma_calc + # is a meaningful ratio rather than two unrelated curves. + self.wilson_n_coeff = int(wilson_n_coeff) + + # 1. Resolution window. + # + # With `asu_idx` the caller has already masked, and `F_obs` / `sig_F_obs` + # / `centric_obs` are ONE ROW PER UNIQUE REFLECTION while `s_obs` carries + # the full symmetry-unrolled geometry. Everything from here to the + # expansion is a per-reflection function of (F, sigma_F, |s|, centric), + # all four of which are symmetry-invariant, so the chain runs on the + # unique set and is broadcast at the end. That is exact -- not an + # approximation -- and it is the difference between doing the + # French-Wilson posterior once and doing it n_ops times. + # + # Without `asu_idx` every array is per-row of `s_obs` and the window is + # applied here, which is the path direct callers and the synthetic tests + # take. + if asu_idx is None: + extras = (F_obs, centric_obs) + if sig_F_obs is not None: + extras = extras + (sig_F_obs,) + s_obs, extras, smag_src = _resolution_mask(s_obs, extras, d_min, d_max) + F_obs, centric_obs = extras[0], extras[1] + if sig_F_obs is not None: + sig_F_obs = extras[2] + else: + if s_mag_asu is None: + raise ValueError("asu_idx requires s_mag_asu (|s| per unique row)") + if int(asu_idx.shape[0]) != int(s_obs.shape[0]): + raise ValueError( + f"asu_idx has {int(asu_idx.shape[0])} entries for " + f"{int(s_obs.shape[0])} unrolled reflections" + ) + smag_src = s_mag_asu + + if F_obs.shape[0] < n_wilson_shells * 5: + raise ValueError( + f"Too few obs reflections ({F_obs.shape[0]}) for " + f"{n_wilson_shells} Wilson shells in [{d_min}, {d_max}] Å." + ) + + # 2. Bessel scaling — Phaser's lmax · d_min (DataMR.cc:1107). + # `bessel_h_scale = 2 pi R_patt` is the Patterson integration radius + # (the chi_Omega sphere): the Bessel argument is + # `h = bessel_h_scale |s|`, so the radial basis represents the + # Patterson out to `R_patt = bessel_h_scale / (2 pi)`. With the + # bandwidth-coupled `d_min` the caller passes, that comes to + # `R_patt ~ sphereOuter = 2 x mean radius`. + if d_min is None: + raise ValueError("d_min is required to set the Bessel scaling") + lmax = L - 1 + lmax_even = lmax if lmax % 2 == 0 else lmax - 1 + self.bessel_h_scale = float(lmax_even) * float(d_min) + + # 3. ONE shell assignment, shared by everything below. + # + # The French-Wilson posterior, the LERF1 build and the variance reweight + # all normalise per resolution shell, and each used to derive its own + # equal-count edges from the same |s| -- one in numpy, one in torch, with + # different quantile-rank rounding. That put a handful of boundary + # reflections in different shells depending on which consumer asked, + # which is a difference of ~2e-4 relative on their normalisation for no + # reason. Assign once, pass it down. + from ..sh import assign_shells, equal_count_shell_edges + + shell_edges, _ = equal_count_shell_edges(smag_src, n_wilson_shells) + obs_shell_idx = assign_shells(smag_src, shell_edges) + + # 3b. Wilson normalisation, so that = 1 on the observations. + # One resolution window for both sides, taken from the bandwidth + # coupling rather than from whichever reflections each side happens to + # contain. Two fits on their own extremes span the same polynomial + # space, so they agree where their data overlap -- but the basis + # saturates at the ends, so each is frozen flat outside its own range + # and the two curves stop being usable against each other. Everything + # that reads Sigma_obs/Sigma_calc needs them on one abscissa. + self._s_lo = 1.0 / float(d_max) if d_max else float(smag_src.min()) + self._s_hi = 1.0 / float(d_min) + # No epsilon here: the observations reach this point symmetry-unrolled, + # which puts each reflection into the sum once per operation that maps + # to it, so multiplicity is already carried by the geometry. Centricity + # is separate and does enter -- it is the Gamma shape. + conv_obs = WilsonNormaliser( + F_obs * F_obs, smag_src, centric=centric_obs, + n_coeff=self.wilson_n_coeff, s_lo=self._s_lo, s_hi=self._s_hi, + ) + self._conv_obs = conv_obs + + # 4. LERF1 obs intensity, with the measurement-information weight. + if obs_weight == "none" or sig_F_obs is None: + w_obs = torch.ones_like(conv_obs.E) + elif obs_weight == "information": + # Measurement error only. The control that says what folding the + # model error in is actually worth. + w_obs = normalise_weight(information_weight( + snr_from_amplitude(F_obs, sig_F_obs), cap=self.snr_cap, + )) + elif obs_weight == "inverse_variance": + # Both error sources in one denominator. sigma_A is the same + # Luzzati falloff the calculated side uses for its expected signal; + # evaluated here at the OBSERVED reflections, because that is the + # side carrying sigmas and the weight has to be per reflection to + # be worth anything -- a weight constant within a shell is absorbed. + sigma_a_obs = eterm_sigma_a(smag_src, self.delta_vrms_A) + w_obs = normalise_weight(inverse_variance_weight( + snr_from_amplitude(F_obs, sig_F_obs), sigma_a_obs, + cap=self.trust_cap, + )) + else: + raise ValueError( + f"obs_weight={obs_weight!r}; expected 'inverse_variance', " + f"'information' or 'none'." + ) + self._w_obs = w_obs + intensity_obs = build_lerf1_intensity( + conv_obs.E, centric_obs, weight=w_obs, use_centric_weight=True, + ) + if self.shell_variance_weights: + intensity_obs = apply_shell_variance_weights( + intensity_obs, smag_src, n_var_shells=n_wilson_shells, + shell_idx=obs_shell_idx, + ) + if asu_idx is not None: + # One value per unique reflection -> one per unrolled reflection. + intensity_obs = intensity_obs[asu_idx] + + # 5. ZSYMM m-filter on the obs SH coefficients. The calc side is never + # filtered -- see score_model. + zsymm = detect_zsymm(sym_mats) + + # 6. Bessel-SH expand the obs side. + # `s_obs` stays at the caller's (wider) dtype so the expansion's + # clustering keys keep their resolution; the intensity does not need to, + # and the expansion casts it to the working precision anyway. + self._c_obs = bessel_sh_expand( + s_obs, intensity_obs, + L=L, bessel_h_scale=self.bessel_h_scale, + zsymm=zsymm, + ) + + + def score_model( + self, + s_calc: torch.Tensor, + F_calc: torch.Tensor, + *, + n_peaks: int = 500, + sigma_threshold: float = -5.0, + sigma_a_source: str = "empirical", + apply_bulk_solvent: bool = False, + solvent_fsol: float = 0.95, + solvent_bsol: float = 300.0, + ) -> Tuple[AdaptiveRotationFunction, List[RotationPeak]]: + """Score one model's transform against the prepared observations. + + ``s_calc`` / ``F_calc`` must already lie inside ``[d_max, d_min]`` -- + this method does not re-mask them. The window is the caller's because + the caller had to know it to build the calc set in the first place. + """ + # The caller owns the calc-side resolution window: `dense_calc_via_box` + # samples the box over `[d_max, d_min]` already, and re-masking here was + # measured to drop 0 of 339040 reflections on 3K7M and 0 of 271630 on + # 1DAW. So take `s_calc` as given and only derive |s| from it. + smag_calc = s_calc.norm(dim=-1) + # The same normaliser, at the same order, on the same abscissa. The calc + # side is a single molecular transform sampled in a P1 box, so there is + # no multiplicity and nothing is centric. + self._conv_calc = WilsonNormaliser( + F_calc * F_calc, smag_calc, + n_coeff=self.wilson_n_coeff, s_lo=self._s_lo, s_hi=self._s_hi, + ) + E_calc = self._conv_calc.E + if sigma_a_source == "empirical": + # Measured, not assumed. Both curves were fitted on one abscissa + # (see `self._s_lo`/`_s_hi`), which is what makes evaluating them at the same + # |s| meaningful. Subsumes the Babinet term: the low-resolution + # solvent deficit is simply what the ratio measures, per structure, + # instead of two universal constants. + eterm = empirical_sigma_a( + self._conv_obs.evaluate(smag_calc), + self._conv_calc.evaluate(smag_calc), + ).to(F_calc.dtype) + elif sigma_a_source == "luzzati": + eterm = eterm_sigma_a(smag_calc, self.delta_vrms_A) + else: + raise ValueError( + f"sigma_a_source={sigma_a_source!r}; expected 'luzzati' or " + f"'empirical'." + ) + # Optional Babinet bulk-solvent factor: Phaser folds it into σ_A as + # `σ_A_eff = solTerm(s²) · Luzzati(s², vrms)` (EnsemblePDB.cc:96-100). + # Default OFF; flip after v25 validates. + if apply_bulk_solvent: + from .preprocessing import bulk_solvent_factor + sol = bulk_solvent_factor( + smag_calc, fsol=solvent_fsol, bsol=solvent_bsol, + ).to(eterm.dtype) + eterm = eterm * sol + # sigma_A**2 here is the expected moving-model intensity, i.e. part of + # the SIGNAL, not a weight -- which is why it stays on this side and the + # inverse-variance weight goes on the observations. Weighting the model + # by its own reliability and then also weighting the data by it would + # count the same thing twice. + w_calc = eterm * eterm + intensity_calc = w_calc * (E_calc * E_calc - 1.0) + + # Bessel-SH expand calc. The calc is NEVER m-filtered (zsymm=1): the + # model carries no crystal symmetry, and projecting it onto the obs's + # invariant-m subspace destroys the orientation information that + # discriminates truth — a calc-side m-filter was tested and is strongly + # harmful on high-symmetry cases (3K7M rank 8→92), so the knob was removed. + c_calc = bessel_sh_expand( + s_calc, intensity_calc, + L=self.L, bessel_h_scale=self.bessel_h_scale, + zsymm=1, + ) + + # 8. Cross-correlate over the radial axis. + xi = cross_correlate_xi(self._c_obs, c_calc) + + # 9. FRF on the adaptive SO(3) grid. + arf = evaluate_rotation_function(xi, grid_sampling_deg=self.grid_sampling_deg) + + # 10. Peak finding (greedy SO(3) NMS). + peaks = find_rotation_peaks( + arf, + n_peaks=n_peaks, + sigma_threshold=sigma_threshold, + nms_radius_deg=max(2.0 * self.grid_sampling_deg, 6.0), + sym_cart=self.sym_cart, + ) + return arf, peaks diff --git a/torchref/experimental/alignment/frf/data_mr.py b/torchref/experimental/alignment/frf/data_mr.py new file mode 100644 index 00000000..3ddc7440 --- /dev/null +++ b/torchref/experimental/alignment/frf/data_mr.py @@ -0,0 +1,605 @@ +"""Bessel-radial × spherical-harmonic expansion of obs and calc Pattersons. + +Mirrors ``DataMR::dataMR_FRF`` (DataMR.cc) and the helper sums in +``Ensemble.cc``. Contains both the spherical-Bessel table (Miller +recurrence) and the chunked obs/calc Bessel×SH expansion fed into the +cross-correlation that ``SiteListAng::get_FRF`` consumes. + +Citations: + * Bessel-radial × SH expansion: DataMR.cc:993, 1107-1117 + * sqrt(2u+1) · j_u(h)/h radial weight: DataMR.cc:993 + * Even-l only (Patterson centrosymmetry): implicit in DataMR.cc + * m-symmetry filter: DataMR.cc:863-870, 1117 +""" +from __future__ import annotations + +import math +import os +import time + +import torch + +_PROFILE = bool(os.environ.get("FRF_PROFILE")) + +#: Byte budget for the per-chunk transients in :func:`bessel_sh_expand`. +#: +#: The chunk holds the Legendre recurrence's rolling rows, so this is really a +#: cache-residency knob, and it has an interior optimum. Measured on an EPYC +#: 9335 at four threads, seconds for the whole rotation function at cap 100 on +#: 3K7M: 24.9 at 2 MB, 12.1 at 8, **9.6 at 32**, 9.9 at 128, 11.0 at 256, 12.3 +#: at 1024. Same shape at cap 64. Below the optimum the 100-iteration loop over +#: l is re-run for too many chunks and Python and dispatch overhead dominate; +#: above it the rolling rows stop fitting in cache and the recurrence becomes +#: memory-bound. The truth rank was identical at every setting. +CLUSTER_CHUNK_BYTES = 32_000_000 + + +#: Grouping resolution for |s|, i.e. for the RADIAL factor. Reflections whose |s| +#: agrees to 1/this share one Bessel evaluation. +#: +#: This one has to be fine. The Bessel argument is `bessel_h_scale * |s|`, of +#: order 250 for a protein at L=64, and j_u oscillates on a scale of 2*pi in its +#: argument -- so an error in |s| is amplified by ~250 before it reaches j_u. +#: Against an ungrouped reference the expansion is bitwise exact at 1e9 and +#: 1.8e-8 to 7.5e-8 relative at 1e7, where this used to sit: a systematic error +#: at or above the engine's own run-to-run spread. Costs nothing where the +#: grouping pays most: the dense P1 calc box is exactly degenerate, so it groups +#: identically at 1e7, 1e9 and 1e11 alike. +_GROUP_SCALE_S = 10_000_000 + +#: Grouping resolution for cos(theta), i.e. for the ANGULAR factor. +#: +#: This one can be coarse, and that is where the speed is. The Legendre factor +#: varies smoothly in cos(theta) with no amplification, and Phaser -- the +#: reference this engine is validated against -- buckets cos(theta) at 1e-3 +#: (`lib/sphericalY.h:43`, COSTHETA_LIMIT), evaluating the Legendre polynomials +#: once per bucket. Setting this finer than Phaser buys accuracy the rest of the +#: chain does not have; setting it coarser than Phaser would be a new +#: approximation and needs its own evidence. +_GROUP_SCALE_COS = 10_000_000 + +from ....config import get_complex_dtype, get_float_dtype +from ....utils.backends import run_or_degrade, select +from ..sh import legendre_recurrence_coefficients +from ._backends import LEGENDRE_BACKENDS + +from .types import BesselSHCoefficients + +__all__ = [ + "bessel_sh_expand", + "cross_correlate_xi", + "spherical_bessel_table", +] + + +#: Exponent of the running-magnitude rescale in :func:`spherical_bessel_table`. +#: +#: The unnormalised downward ladder is enormous before it is renormalised: at the +#: FRF's low-resolution end (``x = bessel_h_scale / d_max``, about 1.3) the +#: intermediates reach 1e157, and they overflow float32 for every ``x`` below +#: ~35 -- most of the resolution range. Rescaling by a fixed factor whenever the +#: running value crosses it keeps the ladder in range at ANY working precision. +#: +#: It has to be a power of two. Dividing by one only decrements the exponent, so +#: the mantissas of every stored value and of the closing renormalisation are +#: untouched and the rescale introduces no rounding at all -- the table comes out +#: bit-identical to the un-rescaled version. Measured event counts over the FRF's +#: range: 5 rescales at x=1.26, 4 at 1.9, 3 at 5, 2 at 10, 1 at 20, 0 at 64. +_BESSEL_RESCALE_EXP = 100 + + +def spherical_bessel_table( + x: torch.Tensor, + u_max: int, + n_extra: int = 25, +) -> torch.Tensor: + """Tabulate spherical Bessel ``j_u(x)`` for ``u ∈ [0, u_max]``, batched over ``x``. + + Uses Miller's downward recurrence (the standard stable choice for + ``j_n(x)`` with ``n > x``): + + j_{u-1}(x) = (2u + 1) / x · j_u(x) − j_{u+1}(x) + + Seed: ``n_start = u_max + n_extra``, ``j_{n_start+1} = 0``, + ``j_{n_start} = 1`` (unnormalised), recur down to ``j_0``, then + renormalise using the exact ``j_0(x) = sin(x) / x``. + + The ladder is rescaled on the fly by ``2**-_BESSEL_RESCALE_EXP`` whenever it + grows past that magnitude -- see that constant for why the recurrence needs + it and why it costs no accuracy. Every rescale divides ``j_mid``, ``j_high`` + **and every row already written**, so the whole table stays in one common + frame and the closing renormalisation cancels it exactly. + + Note that the high-``u`` rows are genuinely negligible rather than merely + small: at ``x = 1.9`` every ``u >= 33`` is below float32's smallest normal, + which is 50% of the band, so those entries are already flushed to zero + wherever the caller works in single precision. + + Returns + ------- + j_table : torch.Tensor + Shape ``(*x.shape, u_max + 1)``, dtype = ``x.dtype``. + """ + real_dtype = x.dtype + device = x.device + # The ladder runs in the argument's own dtype. The rescaling above is what + # keeps the ~90-step downward recurrence in range -- without it the ladder + # overflows float32 for every x below ~35 -- and with it the single + # precision the expansion passes recovers every pose on the benchmark + # panel. A double argument gets a double ladder, which is what the + # bit-identity and scipy checks in the tests exercise. + work_dtype = x.dtype + x64 = x.to(work_dtype) + safe_x = x64.clamp(min=1e-30) + inv_x = 1.0 / safe_x + + n_start = max(u_max + n_extra, u_max + 2) + j_high = torch.zeros_like(x64) + j_mid = torch.ones_like(x64) + j_table = torch.zeros( + (u_max + 1, *x64.shape), dtype=work_dtype, device=device, + ) + threshold = float(2 ** _BESSEL_RESCALE_EXP) + inv_threshold = 1.0 / threshold + # Rescales applied so far, per element. Every element's ladder sits in the + # single frame 2**(-_BESSEL_RESCALE_EXP * n_rescales). + n_rescales = torch.zeros_like(x64, dtype=torch.int32) # dtype-ok: small integer counter + + for n in range(n_start, 0, -1): + j_low = (2.0 * n + 1.0) * inv_x * j_mid - j_high + if n - 1 <= u_max: + j_table[n - 1] = j_low + j_high = j_mid + j_mid = j_low + + # The `.any()` costs a host sync per step (~90 per call, a handful of + # calls per search) and buys skipping a pass over the written rows on + # every step that does not need one. The rows are the expensive part. + over = j_mid.abs() > threshold + if bool(over.any()): + factor = torch.where(over, inv_threshold, 1.0) + j_mid = j_mid * factor + j_high = j_high * factor + if n - 1 <= u_max: + j_table[n - 1:] = j_table[n - 1:] * factor + n_rescales = n_rescales + over.to(torch.int32) # dtype-ok: small integer counter + + true_j0 = torch.sin(x64) * inv_x + true_j0 = torch.where(x64 < 1e-30, torch.ones_like(x64), true_j0) + computed_j0 = j_table[0] + # The degeneracy guard's 1e-30 is an absolute bound on the UNSCALED j_0, so + # express it in the frame the ladder actually ended up in. `frame` is an + # exact power of two; it underflows to 0 past ~10 rescales, at which point + # the guard simply stops firing (it is unreachable for the FRF anyway, whose + # x >= bessel_h_scale / d_max keeps sin(x)/x well away from zero). + frame = torch.ldexp( + torch.ones_like(x64), -_BESSEL_RESCALE_EXP * n_rescales, + ) + safe_j0 = torch.where( + computed_j0.abs() < 1e-30 * frame, frame, computed_j0, + ) + scale = true_j0 / safe_j0 + j_table = j_table * scale.unsqueeze(0) + + perm = list(range(1, j_table.dim())) + [0] + j_table = j_table.permute(*perm).contiguous() + return j_table.to(real_dtype) + + +def _unit_power_ladder(z: torch.Tensor, L: int) -> torch.Tensor: + """``[z^0, z^1, ..., z^{L-1}]`` for unit-modulus ``z``, shape ``(n, L)``. + + Built by doubling -- the block of powers already computed, times the next + power -- rather than by ``torch.cumprod``, for portability: MPS has no + complex cumulative kernels at all (torch 2.9.1 raises "cumulative ops are + not yet supported for complex"), and this was the last thing in the rotation + search that could not run on Apple silicon. + + Not an accuracy change. Measured against the exact powers in double over + 5000 angles at L=101, the ladder gives 5.7e-6 where ``cumprod`` gives 4.1e-6 + in complex64, and 3.4e-14 against 3.3e-14 in complex128 -- the same, because + the error is dominated by ``z``'s own rounding amplified by ``p``, which no + grouping of the multiplies avoids. ``log2(L)`` wide multiplies in place of + one fused pass, so it is not a cost change either. + """ + out = torch.ones((z.shape[0], L), dtype=z.dtype, device=z.device) + width = 1 # out[:, :width] is filled + while width < L: + z_w = out[:, width - 1] * z # z^width + take = min(width, L - width) + out[:, width:width + take] = out[:, :take] * z_w.unsqueeze(1) + width += take + return out + + +def bessel_sh_expand( + s_vectors: torch.Tensor, + intensity: torch.Tensor, + *, + L: int, + bessel_h_scale: float, + zsymm: int = 1, +) -> BesselSHCoefficients: + """Phaser-style ``c_nlm = Σ_h Y*_lm(ŝ) · I · sqrt(2u+1) · j_u(h)/h``. + + Memory-bounded and chunked, verified element-wise by + ``tests/unit/frf_separate``. A direct implementation materialises the full + ``(M, L, N_radial)`` Bessel table and + ``(M, u_max+1)`` j-table for *all* reflections at once — at L≈100 with + a symmetry-unrolled obs set (≳10⁶ reflections) that is tens of GB and + OOMs. Here the j-table, Bessel weights and Y_lm are all computed + *inside* the reflection-chunk loop, so peak memory is set by one chunk. + + Citations: + * radial × SH expansion, sqrt(2u+1)·j_u(h)/h weight: DataMR.cc:993, 1107 + * even-l only (Patterson centrosymmetry) + m-filter: DataMR.cc:863-870, 1117 + + **No antipodal copy.** The Patterson's centrosymmetry is already encoded + twice here -- only even ``l`` are computed, and the negative-``m`` half is + mirrored rather than summed -- and both of those *save* work. Concatenating + ``-s`` onto the reflection set was a third encoding that *cost* work and + bought nothing: for even ``l``, ``Y_lm(-s_hat) = Y_lm(s_hat)``, and the + intensity, Bessel weight and Legendre factor are all unchanged under + negation, so it doubled ``c_nlm`` exactly. Both sides doubled scaled the + rotation function by 4, which the z-score normalisation removes. + + Measured before removal, over 10 benchmark structures x 10 seeded trials: + truth ranks 98/100 identical (1 better, 1 worse), the top score exactly + 0.2499974 to 0.2500036 of the doubled value, and the search 22.5% faster on + 3K7M / 16.4% on 1DAW. Note that Phaser does include the mate (cctbx's + ``conjugate_flag``), so our coefficients are now half of its -- which + matters only to the coefficient-level comparison in + ``alignment_lab/diagnostics/frf_encode_compare.py``. + + Two precisions are in play and they are deliberately different. + + The **clustering keys** are computed on the host in double, whatever dtype + ``s_vectors`` arrive in: ``_GROUP_SCALE_S`` keys ``|s|`` at 1e-7 and that is + exactly where float32's resolution runs out -- at ``|s| = 0.5`` a float32 + rounding is ~0.3 of a key step, so reflections that are mathematically + degenerate would sometimes land in adjacent keys and the degeneracy + collapse the cost model depends on would fray. The host always has double, + the key computation is O(N), and nothing double ever touches the device. + + Everything else -- the Legendre/Y_lm precompute, the radial weights, the + Bessel ladder (kept in range by rescaling), the contraction and the returned + coefficients -- runs at :func:`torchref.config.get_float_dtype`, this + codebase's working precision and the dtype the fused CPU kernel is built + for. + """ + assert s_vectors.dim() == 2 and s_vectors.shape[-1] == 3 + assert intensity.dim() == 1 and intensity.shape[0] == s_vectors.shape[0] + + device = s_vectors.device + # The working precision, and the dtype of everything returned. Not derived + # from the input: the input is deliberately wider (see the docstring). + comp_real = get_float_dtype() + complex_dtype = get_complex_dtype() + + lmax = L - 1 + lmax_even = lmax if (lmax % 2 == 0) else (lmax - 1) + if lmax_even < 2: + raise ValueError(f"L={L} too small; need lmax_even >= 2 (so L >= 3).") + N_radial = (lmax_even - 2) // 2 + 1 + u_max = lmax_even + 1 + + # (l, n) -> u = l + 2n + 1 and the sqrt(2u+1) weight, precomputed once as + # index tensors so the per-chunk Bessel fill is a single advanced-indexed + # assignment instead of a Python loop over ~N_radial·(lmax/2) (l, n) pairs. + even_ls = list(range(2, lmax_even + 1, 2)) + l_list, n_list, u_list, w_list = [], [], [], [] + for l in even_ls: + # Phaser's per-l radial band: nmax = (lmax - l + 2)/2 (DataMR.cc:894), + # narrowing from N_radial terms at l=2 to a single term at l=lmax, so + # the high-l bands cannot carry more radial detail than the reflection + # set supports. `N_radial` above is only the allocated width (Phaser's + # widest band); the populated support is this per-l count. + n_l = (lmax_even - l) // 2 + 1 + for n in range(n_l): + u = l + 2 * n + 1 + l_list.append(l) + n_list.append(n) + u_list.append(u) + w_list.append(math.sqrt(float(2 * u + 1))) + l_idx = torch.tensor(l_list, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + n_idx = torch.tensor(n_list, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + u_idx = torch.tensor(u_list, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + w_vec = torch.tensor(w_list, dtype=comp_real, device=device) + # Only even degrees l ∈ [2, lmax_even] carry signal (odd-l and l=0 are zeroed + # by Patterson centrosymmetry). Compute / contract Y_lm on these rows only — + # the assembly + einsum are the bottleneck, so this ~halves them. The full + # c_nlm keeps the (L, ...) shape with odd/zero rows left at zero. + even_l_idx = torch.tensor(even_ls, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + + M = s_vectors.shape[0] + einsum_dtype = complex_dtype + + prof = {"cluster": 0.0, "dbuild": 0.0, "bessel": 0.0, "legendre": 0.0, + "scatter": 0.0, "contract": 0.0} if _PROFILE else None + + def _tick(t0): + if device.type == "cuda": + torch.cuda.synchronize() + return time.perf_counter() - t0 + + if _PROFILE: + t0 = time.perf_counter() + + # ---- Cluster reflections by (|s|, cosθ) --------------------------------- + # Both the radial Bessel weight (a function of |s|) and the Legendre barP (a + # function of cosθ) are constant within a cluster, so the per-reflection SH + # expansion factorises — only the azimuthal phase e^{imφ} and the intensity + # vary inside a cluster. This generalises Phaser's cosθ clustering + # (DataMR.cc:918, HKL_clustered) by ALSO factoring the radial term, which + # collapses the dominant contraction from O(M·L³) to O(n_clusters·L³). On the + # dense P1 calc box (and on cubic/tetragonal obs lattices) reflections sharing + # (H²+K²+L², L_z) land in one cluster → n_clusters ≪ M (~16× fewer). For a + # non-degenerate reflection set it degrades gracefully to ≈ the per-reflection + # cost (clusters are singletons), with no change in the result. + s_mag_all = s_vectors.norm(dim=-1).clamp(min=1e-30) + cos_all = (s_vectors[..., 2] / s_mag_all).clamp(min=-1.0, max=1.0) + phi_all = torch.atan2(s_vectors[..., 1], s_vectors[..., 0]) + # Separate resolutions for the two factors: the radial term needs a fine + # |s| key, the angular term does not. One shared key forces the finer of the + # two on both, which costs merges the angular part never needed. The keys + # are formed on the host in double -- see the docstring -- and only the + # integer keys come back. + s_key = s_vectors.detach().cpu().to(torch.float64) # dtype-ok: exact clustering key on the host; the device never sees it + s_mag_key = s_key.norm(dim=-1).clamp(min=1e-30) + cos_key = (s_key[..., 2] / s_mag_key).clamp(min=-1.0, max=1.0) + k_s = (s_mag_key * _GROUP_SCALE_S).round().to(torch.int64) # dtype-ok: exact clustering key + k_c = (cos_key * _GROUP_SCALE_COS).round().to(torch.int64) + _GROUP_SCALE_COS # dtype-ok: exact clustering key + key = (k_s * (2 * _GROUP_SCALE_COS + 1) + k_c).to(s_vectors.device) + uniq_key, inverse = torch.unique(key, return_inverse=True) + n_clusters = int(uniq_key.shape[0]) + # Per-group geometry: the MEAN over the group's members, not an arbitrary + # one of them. Members agree to the key's resolution but not exactly, so a + # scatter-assignment leaves whichever member was written last -- an error up + # to the full bin width, and biased. The mean costs one extra reduction and + # centres it. + def _group_mean(values, index, n_groups): + tot = torch.zeros(n_groups, dtype=values.dtype, device=device) + cnt = torch.zeros(n_groups, dtype=values.dtype, device=device) + tot.index_add_(0, index, values) + cnt.index_add_(0, index, torch.ones_like(values)) + return tot / cnt.clamp(min=1.0) + + rep_cos = _group_mean(cos_all.to(comp_real), inverse, n_clusters) + rep_sin = torch.sqrt((1.0 - rep_cos * rep_cos).clamp(min=0.0)) + + # Resolution shells: the distinct |s| values, and which shell each cluster + # belongs to. The radial factor depends on |s| alone, so it is applied once + # per shell rather than once per cluster -- and there are far fewer shells + # than clusters, because many directions share a |s| on a lattice. Measured + # over the benchmark: 2.7 to 39 clusters per shell. + uniq_ks, inv_s = torch.unique(k_s, return_inverse=True) + # `k_s` is one of the host-side keys, so its inverse comes back on the host + # while everything it indexes -- `shell_of_cluster`, `s_mag_all` -- is on the + # compute device. Bring it across here, once, rather than leaving a host + # index to meet device values. + inv_s = inv_s.to(device) + n_shells = int(uniq_ks.shape[0]) + shell_of_cluster = torch.zeros(n_clusters, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + shell_of_cluster[inverse] = inv_s + shell_smag = _group_mean(s_mag_all.to(comp_real), inv_s, n_shells) + + # Reorder the clusters so each shell's members are adjacent. The angular + # accumulation below scatters every cluster into its shell's row of T; in + # cluster order those writes land all over T, in shell order they sweep it + # once. Same arithmetic, and the sort is one pass over n_clusters against a + # scatter of n_clusters x n_even x L. + order = torch.argsort(shell_of_cluster) + shell_of_cluster = shell_of_cluster[order] + rep_cos = rep_cos[order] + rep_sin = rep_sin[order] + if _PROFILE: + prof["cluster"] += _tick(t0); t0 = time.perf_counter() + + # ---- S[c, p] = Σ_{h∈c} I_h e^{-i p φ_h}, for p = 0 .. L-1 --------------- + # Only the non-negative half of m is built. The coefficients obey + # c[n, l, -p] = (-1)^p conj(c[n, l, +p]) + # exactly -- the intensity, the Bessel weight and the Legendre factor are all + # real, so m enters only through the azimuthal phase, and P_{l,|m|} does not + # distinguish +p from -p. Verified bit-exact against the full-range build. + # That halves both this sum and the contraction below. + # + # The Y_lm convention (sh.evaluate_ylm) carries C(m, φ) = (-1)^m e^{imφ} for + # m >= 0, so conj(C) contributes a (-1)^p factor. It is applied once per + # (cluster, p) after the sum rather than once per (reflection, p) -- the same + # number for a factor of M/n_clusters less work. + p_idx = torch.arange(L, device=device) # (L,) + # `inverse` maps a reflection to its cluster in the ORIGINAL cluster order; + # the clusters were just permuted into shell order, so compose the two. + rank_of_cluster = torch.empty_like(order) + rank_of_cluster[order] = torch.arange(n_clusters, device=device) + cluster_of_refl = rank_of_cluster[inverse] + Sp = torch.zeros((n_clusters, L), dtype=einsum_dtype, device=device) + dchunk = 262_144 + for start_i in range(0, M, dchunk): + stop = min(start_i + dchunk, M) + ph = phi_all[start_i:stop].to(comp_real) # (c,) + i_c = intensity[start_i:stop].to(comp_real) # (c,) + # e^{-i p phi} = z^p with z = e^{-i phi}, so one transcendental per + # reflection and a power ladder over p, rather than a transcendental + # per (reflection, p). At L=101 over 2.6e6 reflections that is 2.6e8 + # sincos calls replaced by 2.6e6 of them plus a few complex multiplies + # each. `_unit_power_ladder` builds it by doubling rather than with a + # cumulative product, which no complex MPS kernel implements. + z = torch.polar(torch.ones_like(ph), -ph) # (c,) + e_neg = _unit_power_ladder(z, L) # (c, L) = z^p + e_neg = (e_neg * i_c.unsqueeze(1)).to(einsum_dtype) + Sp.index_add_(0, cluster_of_refl[start_i:stop], e_neg) + sign_p = ((-1.0) ** p_idx.to(comp_real)).to(einsum_dtype) # (L,) + Dp = Sp * sign_p.unsqueeze(0) # (n_clusters, L) + if _PROFILE: + prof["dbuild"] += _tick(t0); t0 = time.perf_counter() + + # ---- contraction, in two steps ------------------------------------------ + # Direct: + # c[n,l,p] = Σ_c B[c,l,n] · P[c,l,p] · D[c,p] + # but B depends only on |s| and P only on cos(theta), while the cluster index + # carries both. Summing that way re-multiplies the radial factor once per + # distinct direction at the same resolution. Grouping the clusters by shell i: + # T[i,l,p] = Σ_{c in shell i} P[c,l,p] · D[c,p] (no radial axis) + # c[n,l,p] = Σ_i B[i,l,n] · T[i,l,p] (shells, not clusters) + # which trades n_clusters·N_radial for n_clusters + n_shells·N_radial. On the + # benchmark that is 2.5x to 18x fewer multiply-adds, and it shrinks the Bessel + # table by the same clusters-per-shell factor. Exact, not an approximation. + # + # The Legendre recurrence is run here rather than called, so each row can be + # accumulated into T the moment it exists and the (chunk, n_even, L) table is + # never built. That table was what bounded the chunk width, and the loop over + # l had to be repeated for every chunk -- 100 iterations of a handful of + # small kernels, 71 times over, which cost more in launch overhead than the + # arithmetic did. Without it the chunks are wide enough that the loop runs + # once or twice in total. + # + # `a_coef` and `b_coef` are zero for m >= l, so the vertical recurrence runs + # at full width; slicing to [:l] instead makes every iteration a differently + # shaped, mostly tiny kernel. + n_even = len(even_ls) + le_idx = (l_idx - 2) // 2 # l value -> even-l row index + a_coef, b_coef, sect = legendre_recurrence_coefficients(L, comp_real, device) + + # The whole per-shell sum T and the whole radial table B used to be built at + # full size, and both are large: at L=101 with 35k shells they are 2.8 GB and + # 0.7 GB. They are also touched once each, so that is pure memory traffic -- + # and the scatter into a 2.8 GB target misses cache on essentially every + # write. + # + # The clusters are sorted by shell, so a chunk of clusters spans a + # *contiguous* range of shells. That lets both be per-chunk, and lets the + # radial contraction be folded into the same loop: once a chunk's shells are + # complete, contract them and drop them. The arithmetic is identical -- the + # contraction still costs n_shells x n_even x N_radial x L in total -- but the + # scatter target is now tens of MB rather than gigabytes, and the running + # answer c_pos is a few MB, so both stay in cache. + c_pos = torch.zeros((N_radial, n_even, L), dtype=einsum_dtype, device=device) + + rbytes = 4 if comp_real == torch.float32 else 8 # dtype-ok: byte-size lookup for a memory estimate + per_cluster = rbytes * 6 * L + cstep = max(1, min(n_clusters, CLUSTER_CHUNK_BYTES // max(1, per_cluster))) + for cs in range(0, n_clusters, cstep): + ce = min(cs + cstep, n_clusters) + if _PROFILE: + t0 = time.perf_counter() + sh_abs = shell_of_cluster[cs:ce] + s0 = int(sh_abs[0]) + s1 = int(sh_abs[-1]) + 1 # sorted, so this is the range + nb = s1 - s0 + sh = sh_abs - s0 # shell index within the chunk + + # Radial weights for this chunk's shells only. + x_s = (bessel_h_scale * shell_smag[s0:s1]).clamp(min=1e-30) + j_all = spherical_bessel_table(x_s, u_max) # (nb, u_max+1) + B = torch.zeros((nb, n_even, N_radial), dtype=comp_real, device=device) + B[:, le_idx, n_idx] = w_vec.unsqueeze(0) * j_all[:, u_idx] / x_s.unsqueeze(-1) + if _PROFILE: + prof["bessel"] += _tick(t0); t0 = time.perf_counter() + + # (n_even, nb, L) each, so Tr[pos] is contiguous for index_add_. Two real + # accumulators rather than one complex: the scatter stays a real + # operation and no complex temporary is allocated per (l, chunk). + Tr = torch.zeros((n_even, nb, L), dtype=comp_real, device=device) + Ti = torch.zeros((n_even, nb, L), dtype=comp_real, device=device) + + # Legendre recurrence and per-shell accumulation. Both stages are + # memory-bound in plain torch -- every row of the recurrence makes a + # round trip -- so this dispatches to a fused kernel where the row stays + # in cache, and falls back to the torch reference when none is built. + # `shell_of_cluster` is sorted, which the fused kernel requires: it + # partitions work by shell so that no two threads write the same + # accumulator row. + args = (Tr, Ti, rep_cos[cs:ce], rep_sin[cs:ce], + Dp[cs:ce].real.contiguous(), Dp[cs:ce].imag.contiguous(), + sh, a_coef, b_coef, sect) + backend = select(LEGENDRE_BACKENDS, args[:6]) + run_or_degrade(LEGENDRE_BACKENDS, backend, False, *args) + if _PROFILE: + prof["legendre"] += _tick(t0); t0 = time.perf_counter() + + c_pos += torch.complex( + torch.einsum("iln,lim->nlm", B, Tr), + torch.einsum("iln,lim->nlm", B, Ti), + ).to(einsum_dtype) + if _PROFILE: + prof["contract"] += _tick(t0) + + # Mirror onto m < 0: c[-p] = (-1)^p conj(c[+p]). + c_e = torch.zeros((N_radial, n_even, 2 * L - 1), dtype=einsum_dtype, + device=device) + c_e[:, :, (L - 1):] = c_pos + if L > 1: + mirror = c_pos[:, :, 1:].conj() * sign_p[1:].view(1, 1, -1) + c_e[:, :, :(L - 1)] = torch.flip(mirror, dims=(-1,)) + c_nlm = torch.zeros((N_radial, L, 2 * L - 1), dtype=complex_dtype, device=device) + c_nlm[:, even_l_idx, :] = c_e.to(complex_dtype) + if _PROFILE: + prof["contract"] += _tick(t0) + + if _PROFILE: + tot = sum(prof.values()) + 1e-30 + print(f"[FRF_PROFILE] M={M} n_clusters={n_clusters} ({M/max(1,n_clusters):.1f}x) " + f"n_shells={n_shells} ({n_clusters/max(1,n_shells):.1f} clu/shell) " + f"L={L} dtype={comp_real} | " + + " ".join(f"{k}={v*1000:.0f}ms({100*v/tot:.0f}%)" for k, v in prof.items()), + flush=True) + + # m-symmetry filter (observed side only; caller passes zsymm=1 for calc). + if zsymm > 1: + m_vals = torch.arange(-(L - 1), L, device=device) + c_nlm[:, :, (m_vals.abs() % zsymm) != 0] = 0.0 + + return BesselSHCoefficients( + coeffs=c_nlm, + L=L, + bessel_h_scale=float(bessel_h_scale), + ) + + +def cross_correlate_xi( + c_obs: BesselSHCoefficients, + c_calc: BesselSHCoefficients, +) -> torch.Tensor: + """Contract obs/calc Bessel-SH coefficients on the radial-Bessel axis. + + Phaser source: the radial sum ``Σ_n c_obs[n,l,m] · conj(c_calc[n,l,m'])`` + happens inside ``DataMR::dataMR_FRF`` before being fed into + ``SiteListAng::DoRfftStuff`` as the ``clmn`` tensor (FastRot.cc:39). + + Convention: + xi[l, m, n] = Σ_r c_obs[r, l, n] · conj(c_calc[r, l, m]) + so that the peak Euler triple satisfies ``s_calc = R · s_obs``. + + Accumulated at the configured complex dtype. The radial sum runs over + oscillating ``j_u``, so the terms alternate in sign and cancel, and this was + once accumulated one step wider for that reason. Measured, the width is not + what the result needs: from one set of complex64 coefficients, a complex64 + contraction lands within 1.5e-5 of the complex128 one whose peak magnitude + is 109. What it does need is for the conjugate below to be *materialised* -- + see the ``resolve_conj`` note. + + Returns + ------- + xi : torch.Tensor (complex), shape (L, 2L-1, 2L-1) + """ + if c_obs.L != c_calc.L: + raise ValueError(f"L mismatch: obs={c_obs.L} calc={c_calc.L}") + # The configured complex dtype. The oscillatory radial sum used to be + # accumulated one step wider than the coefficients; measured, that moved + # scores by 1e-4 relative and reordered the deep peak list without moving + # the top peak, and the placement search now consumes only the top few + # distinct orientations. Single precision recovers every pose on the panel. + acc = get_complex_dtype() + # `resolve_conj()` is load-bearing, not tidiness. `torch.conj` returns a + # lazy view carrying a conjugate BIT, and MPS's batched complex matmul -- + # which is what this einsum lowers to -- ignores that bit and contracts the + # unconjugated values. It is silent: the result is a plausible tensor that + # is simply wrong, here by 173% of |xi|max, which then reorders the whole + # peak list. Elementwise ops, `where`, `index_add` and 2-D matmul all honour + # the bit; only the batched matmul path does not. + return torch.einsum( + "rln,rlm->lmn", + c_obs.coeffs.to(acc), + torch.conj(c_calc.coeffs).resolve_conj().to(acc), + ) diff --git a/torchref/experimental/alignment/frf/dense_calc.py b/torchref/experimental/alignment/frf/dense_calc.py new file mode 100644 index 00000000..d95921c5 --- /dev/null +++ b/torchref/experimental/alignment/frf/dense_calc.py @@ -0,0 +1,105 @@ +"""Dense P1-box sampling of a model's molecular transform. + +The Fast Rotation Function correlates the obs +Patterson against the *model* transform, and sampling that transform at the +sparse crystal lattice under-determines the high-l spherical-harmonic modes for +large molecules. Phaser avoids this by computing the model transform on a dense, +oversampled P1-box FFT grid (``EnsemblePDB.cc:122-135``). This module does the +same: drop the (single, un-symmetry-expanded) model into a cubic P1 box and +reuse ``ModelFT``'s own structure-factor machinery (real ITC92 form factors + +per-atom B/occ) to sample ``|F_calc|`` on the box's dense reciprocal grid. + +Unlike the original benchmark helper this operates on ``model.copy()`` so the +caller's ``cell``/``spacegroup``/``max_res`` are never mutated, and the whole SF +build runs under ``torch.no_grad()`` (the load-bearing memory fix — a forward-only +SF build otherwise accumulates a backward graph and OOMs on big grids). +""" +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Tuple + +import torch + +from torchref.config import get_float_dtype + +if TYPE_CHECKING: + from torchref.model import ModelFT + + +def dense_calc_via_box( + model: "ModelFT", + d_max: float, + d_min: float, + *, + pad: float = 2.0, + verbose: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Sample the model transform on a dense P1-box grid in ``[d_max, d_min]``. + + Parameters + ---------- + model : ModelFT + Search model (in whatever orientation the rotation search should treat + as the reference). **Not mutated** — a ``model.copy()`` is used internally. + d_max, d_min : float + Low- and high-resolution limits (Å) for the returned reflections. + pad : float, optional + P1-box edge as a multiple of the molecular diameter (default 2.0, the + validated v19 value). A bigger box = finer reciprocal sampling. + verbose : bool, optional + If True, print a one-line ``[DENSE_FT]`` summary (box edge / grid size). + + Returns + ------- + (s_vec, F_calc) : Tuple[torch.Tensor, torch.Tensor] + ``s_vec`` is the Cartesian reciprocal grid (N, 3) and ``F_calc`` the + amplitudes (N,), in the configured float dtype on the model's device. + The expansion forms its exact clustering keys on the host itself. + """ + from torchref.symmetry.cell import Cell + + # Isolate the box mutation from the caller: the three setters below replace + # the FFT submodule, so anything the copy set up for the crystal cell is + # thrown away. That used to cost 104 ms on 3K7M building a 250**3 x 3 + # coordinate grid the box never looks at; `real_space_grid` is built on + # demand now rather than stored, so the copy is cheap. + m = model.copy() + with torch.no_grad(): + coords = m.xyz() + dev = coords.device + # Cubic P1 box sized to ``pad`` diameters; no symmetry keeps the grid small. + extent = (coords - coords.mean(0)).norm(dim=-1).max().item() + a = float(pad * 2.0 * extent) + # The grid is derived lazily from (cell, space group, max_res) on first + # use, so the order of these assignments no longer matters. + m.max_res = float(d_min) + m.spacegroup = "P 1" + m.cell = Cell([a, a, a, 90.0, 90.0, 90.0], device=dev) + + nmax = int(math.ceil(a / d_min)) + idx = torch.arange(-nmax, nmax + 1, device=dev) + H, K, Lg = torch.meshgrid(idx, idx, idx, indexing="ij") + hkl = torch.stack( + [H.reshape(-1), K.reshape(-1), Lg.reshape(-1)], dim=-1 + ).to(torch.long) # dtype-ok: Miller indices are integers + # Cubic box: |s| = |hkl| / a. + real = get_float_dtype() + smag = hkl.to(real).norm(dim=-1) / a + keep = (smag >= 1.0 / d_max) & (smag <= 1.0 / d_min) + hkl = hkl[keep].contiguous() + F = model_sf_abs(m, hkl) + s_vec = hkl.to(real) / a + + if verbose: + print( + f"[DENSE_FT] box={a:.0f}A n_grid={hkl.shape[0]} max_res={d_min:.2f}", + flush=True, + ) + return s_vec, F + + +def model_sf_abs(model: "ModelFT", hkl: torch.Tensor) -> torch.Tensor: + """``|F_calc|`` for ``hkl`` via the model's SF machinery (no grad), in the model's dtype.""" + with torch.no_grad(): + return model.get_structure_factor(hkl, recalc=True).abs() diff --git a/torchref/experimental/alignment/frf/kernels/__init__.py b/torchref/experimental/alignment/frf/kernels/__init__.py new file mode 100644 index 00000000..e92588d1 --- /dev/null +++ b/torchref/experimental/alignment/frf/kernels/__init__.py @@ -0,0 +1 @@ +"""Kernels for the fast rotation function's spherical-harmonic expansion.""" diff --git a/torchref/experimental/alignment/frf/kernels/cpu/__init__.py b/torchref/experimental/alignment/frf/kernels/cpu/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/torchref/experimental/alignment/frf/kernels/cpu/legendre_shell.py b/torchref/experimental/alignment/frf/kernels/cpu/legendre_shell.py new file mode 100644 index 00000000..e7d56746 --- /dev/null +++ b/torchref/experimental/alignment/frf/kernels/cpu/legendre_shell.py @@ -0,0 +1,295 @@ +"""Fused Legendre recurrence and shell accumulation, as one C++ kernel. + +The portable version runs the vertical recurrence as one torch operation per +``l`` and then scatters the row, so every row makes a round trip to memory. At +L=101 over 4.4e5 clusters that is ~108 GB for the recurrence and ~71 GB for the +scatter, both measured at ~50 GB/s -- the stages are bandwidth-bound, and the +arithmetic underneath is a small fraction of the time. + +float32 throughout, matching the rest of this codebase's kernels. The radial +Bessel recurrence is a separate stage and keeps its float64 internals, where the +downward recurrence's cancellation actually needs them. + +Fusing them removes the round trip: one cluster's three rows are 1.2 kB of stack, +so ``cur`` is produced, multiplied and accumulated without ever reaching memory. +Two further things fall out of writing it as a loop nest: + +* **Ragged widths are free.** ``bar_P[l, m]`` is zero for m > l, so step ``l`` + needs only columns 0..l -- ``for (m = 0; m <= l; ++m)`` and nothing more. In + torch the same saving needs narrowed views, and that was measured *slower*, + because a strided scatter target costs more than the zeros it skips. +* **No atomics.** The clusters arrive sorted by shell, so a thread that owns a + range of shells owns every write into those shells' rows. Parallelising over + clusters instead would race on the shared accumulator. + +The accumulator rows for one shell are ``n_even * L`` scalars -- 40 kB at L=101 -- +so they stay in cache across that shell's clusters, which is the point of +grouping by shell in the first place. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from torchref.base.electron_density.kernels.cpu._cpp_build import build_extension + +_CPP_SRC = r""" +#include +#include +#ifdef _OPENMP +#include +#else +#include +#endif + +// One shell's worth of work: every cluster in [c0, c1) contributes +// T[pos][s][m] += barP(l, m) * D[c][m] for even l >= 2, m <= l +// with barP built by the vertical recurrence in registers/stack. +template +static void shell_range( + int64_t s_begin, int64_t s_end, + const int64_t* __restrict off, + const int64_t* __restrict shell, + const scalar_t* __restrict rep_cos, + const scalar_t* __restrict rep_sin, + const scalar_t* __restrict Dr, + const scalar_t* __restrict Di, + const scalar_t* __restrict a_coef, + const scalar_t* __restrict b_coef, + const scalar_t* __restrict sect, + scalar_t* __restrict Tr, + scalar_t* __restrict Ti, + int64_t L, int64_t nb, int64_t n_even, scalar_t seed) { + + std::vector buf(3 * L, scalar_t(0)); + scalar_t* prev2 = buf.data(); + scalar_t* prev1 = buf.data() + L; + scalar_t* cur = buf.data() + 2 * L; + + for (int64_t s = s_begin; s < s_end; ++s) { + for (int64_t c = off[s]; c < off[s + 1]; ++c) { + const int64_t row = shell[c]; // == s, carried explicitly + const scalar_t co = rep_cos[c]; + const scalar_t si = rep_sin[c]; + const scalar_t* dr = Dr + c * L; + const scalar_t* di = Di + c * L; + + for (int64_t m = 0; m < L; ++m) { prev1[m] = scalar_t(0); prev2[m] = scalar_t(0); } + prev1[0] = seed; // bar_P_0^0 + + for (int64_t l = 1; l < L; ++l) { + const scalar_t* a = a_coef + l * L; + const scalar_t* b = b_coef + l * L; + // Vertical recurrence, only where the row can be non-zero. + for (int64_t m = 0; m < l; ++m) { + cur[m] = a[m] * co * prev1[m] - b[m] * prev2[m]; + } + // Sectoral m == l, which MUST be in place before the products below: + // it is this row's diagonal entry. + cur[l] = sect[l] * si * prev1[l - 1]; + + if (l >= 2 && (l % 2) == 0) { + const int64_t pos = (l - 2) / 2; + scalar_t* tr = Tr + (pos * nb + row) * L; + scalar_t* ti = Ti + (pos * nb + row) * L; + for (int64_t m = 0; m <= l; ++m) { + tr[m] += cur[m] * dr[m]; + ti[m] += cur[m] * di[m]; + } + } + // Rotate the three buffers; nothing is copied. + scalar_t* t = prev2; prev2 = prev1; prev1 = cur; cur = t; + } + } + } +} + +void legendre_shell_accumulate( + torch::Tensor Tr, torch::Tensor Ti, + torch::Tensor rep_cos, torch::Tensor rep_sin, + torch::Tensor Dr, torch::Tensor Di, + torch::Tensor shell, torch::Tensor offsets, + torch::Tensor a_coef, torch::Tensor b_coef, torch::Tensor sect, + double seed) { + + TORCH_CHECK(Tr.is_contiguous() && Ti.is_contiguous(), "T must be contiguous"); + TORCH_CHECK(rep_cos.scalar_type() == Tr.scalar_type() + && rep_sin.scalar_type() == Tr.scalar_type() + && Dr.scalar_type() == Tr.scalar_type() + && Di.scalar_type() == Tr.scalar_type() + && a_coef.scalar_type() == Tr.scalar_type() + && b_coef.scalar_type() == Tr.scalar_type() + && sect.scalar_type() == Tr.scalar_type(), + "every array must share the accumulator's dtype"); + TORCH_CHECK(Dr.is_contiguous() && Di.is_contiguous(), "D must be contiguous"); + TORCH_CHECK(shell.scalar_type() == torch::kLong, "shell must be int64"); + TORCH_CHECK(offsets.scalar_type() == torch::kLong, "offsets must be int64"); + + const int64_t n_even = Tr.size(0); + const int64_t nb = Tr.size(1); + const int64_t L = Tr.size(2); + TORCH_CHECK(offsets.numel() == nb + 1, "offsets must have n_shells + 1 entries"); + + // float32 only, by policy: this codebase has no float64 kernels. The caller + // is checked rather than dispatched on, so a float64 accumulator is a loud + // error instead of a silent reinterpretation of the buffer. + TORCH_CHECK(Tr.scalar_type() == torch::kFloat, + "legendre_shell_accumulate is float32 only, got ", Tr.scalar_type()); + { + using scalar_t = float; + const int64_t* off = offsets.data_ptr(); + const int64_t* sh = shell.data_ptr(); + const scalar_t* rc = rep_cos.data_ptr(); + const scalar_t* rs = rep_sin.data_ptr(); + const scalar_t* dr = Dr.data_ptr(); + const scalar_t* di = Di.data_ptr(); + const scalar_t* ac = a_coef.data_ptr(); + const scalar_t* bc = b_coef.data_ptr(); + const scalar_t* sc = sect.data_ptr(); + scalar_t* tr = Tr.data_ptr(); + scalar_t* ti = Ti.data_ptr(); + const scalar_t sd = static_cast(seed); + +#ifdef _OPENMP + // Dynamic, because clusters per shell varies (measured 2.7 to 39 across the + // benchmark) so equal shell counts are not equal work. +#pragma omp parallel for schedule(dynamic, 8) + for (int64_t s = 0; s < nb; ++s) { + shell_range(s, s + 1, off, sh, rc, rs, dr, di, ac, bc, sc, + tr, ti, L, nb, n_even, sd); + } +#else + // Apple Clang rejects -fopenmp, so carve the shells into contiguous blocks. + int nthreads = std::max(1u, std::thread::hardware_concurrency()); + if (nthreads > nb) nthreads = static_cast(std::max(nb, 1)); + std::vector pool; + const int64_t per = (nb + nthreads - 1) / std::max(nthreads, 1); + for (int t = 0; t < nthreads; ++t) { + const int64_t s0 = t * per; + const int64_t s1 = std::min(nb, s0 + per); + if (s0 >= s1) break; + pool.emplace_back([=] { + shell_range(s0, s1, off, sh, rc, rs, dr, di, ac, bc, sc, + tr, ti, L, nb, n_even, sd); + }); + } + for (auto& th : pool) th.join(); +#endif + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("legendre_shell_accumulate", &legendre_shell_accumulate, + "Fused Legendre recurrence and per-shell accumulation"); +} +""" + +_module = None +_module_failed = False +_module_error: Optional[Tuple[str, str]] = None + + +def _get_module(): + """The compiled extension, or None if it could not be built.""" + global _module, _module_failed, _module_error + if _module is not None: + return _module + if _module_failed: + return None + _module, _module_error = build_extension("frf_legendre_shell", _CPP_SRC) + if _module is None: + _module_failed = True + return _module + + +def why_unavailable() -> Optional[str]: + """``None`` if the fused kernel is usable, else why it is not. + + The single availability probe for this backend, read by + :mod:`torchref.utils.backends`. A missing compiler and a compile error are + different problems, and the captured diagnostic is what separates them. + """ + if _get_module() is not None: + return None + reason = _module_error[0] if _module_error else "unknown reason" + return ( + f"the fused CPU Legendre/shell kernel is not available ({reason}); see " + "torchref.experimental.alignment.frf.kernels.cpu.legendre_shell." + "last_error()" + ) + + +def available() -> bool: + """Whether the fused kernel compiled and is ready to dispatch.""" + return why_unavailable() is None + + +def last_error() -> Optional[Tuple[str, str]]: + """``(message, traceback)`` from the last failed build attempt, if any.""" + _get_module() + return _module_error + + +def clear_cache() -> None: + """Forget the build result, so the next call retries. For tests.""" + global _module, _module_failed, _module_error + _module, _module_failed, _module_error = None, False, None + + +def shell_offsets(shell: torch.Tensor, n_shells: int) -> torch.Tensor: + """Start index of each shell in a shell-sorted cluster array, plus the end. + + ``(n_shells + 1,)`` int64. The kernel needs the ranges rather than the + per-cluster labels so that a thread can own a set of shells outright and + write their accumulator rows without atomics. + """ + counts = torch.bincount(shell, minlength=n_shells) + offsets = torch.zeros(n_shells + 1, dtype=torch.long, device=shell.device) # dtype-ok: index tensor; index_add_/gather need int64 + torch.cumsum(counts, dim=0, out=offsets[1:]) + return offsets + + +def legendre_shell_accumulate( + Tr: torch.Tensor, + Ti: torch.Tensor, + rep_cos: torch.Tensor, + rep_sin: torch.Tensor, + Dr: torch.Tensor, + Di: torch.Tensor, + shell: torch.Tensor, + a_coef: torch.Tensor, + b_coef: torch.Tensor, + sect: torch.Tensor, +) -> None: + """Fused recurrence and accumulation, in place on ``Tr``/``Ti``. + + Same signature and same effect as + :func:`torchref.experimental.alignment.frf.kernels.portable.legendre_shell_accumulate`. + ``shell`` must be sorted non-decreasing -- the kernel partitions work by + shell to avoid atomics, and unsorted input would silently drop + contributions rather than merely run slowly. + """ + from ....sh import LEGENDRE_SEED + + module = _get_module() + if module is None: + raise RuntimeError(why_unavailable()) + offsets = shell_offsets(shell, Tr.shape[1]) + module.legendre_shell_accumulate( + Tr, Ti, rep_cos.contiguous(), rep_sin.contiguous(), + Dr.contiguous(), Di.contiguous(), shell.contiguous(), offsets, + a_coef.contiguous(), b_coef.contiguous(), sect.contiguous(), + float(LEGENDRE_SEED), + ) + + +__all__ = [ + "available", + "clear_cache", + "last_error", + "legendre_shell_accumulate", + "shell_offsets", + "why_unavailable", +] diff --git a/torchref/experimental/alignment/frf/kernels/portable.py b/torchref/experimental/alignment/frf/kernels/portable.py new file mode 100644 index 00000000..7bcc2db3 --- /dev/null +++ b/torchref/experimental/alignment/frf/kernels/portable.py @@ -0,0 +1,75 @@ +"""Portable Legendre-recurrence-and-shell-accumulation, in plain torch. + +The reference the fused kernel is checked against, and the fallback whenever that +kernel cannot be built. One step of the vertical recurrence per ``l``, then the +row is multiplied by the azimuthal sums and scattered into its cluster's shell. + +Both stages are memory-bound: at L=101 over 4.4e5 clusters the recurrence moves +about 108 GB and the scatter about 71 GB, and both measure ~50 GB/s, which is +roughly what four threads get from a server memory controller. The arithmetic is +a small fraction of that -- which is the whole reason for a fused kernel, where +the row never leaves cache. +""" + +from __future__ import annotations + +import torch + +from ...sh import LEGENDRE_SEED + + +def legendre_shell_accumulate( + Tr: torch.Tensor, + Ti: torch.Tensor, + rep_cos: torch.Tensor, + rep_sin: torch.Tensor, + Dr: torch.Tensor, + Di: torch.Tensor, + shell: torch.Tensor, + a_coef: torch.Tensor, + b_coef: torch.Tensor, + sect: torch.Tensor, +) -> None: + """Accumulate ``sum_c barP[c, l, m] * D[c, m]`` into ``Tr``/``Ti``, in place. + + Parameters + ---------- + Tr, Ti : torch.Tensor + ``(n_even, n_shells, L)`` real accumulators, added into. + rep_cos, rep_sin : torch.Tensor + ``(n_clusters,)`` cos and sin of the polar angle, per cluster. + Dr, Di : torch.Tensor + ``(n_clusters, L)`` real and imaginary azimuthal sums. + shell : torch.Tensor + ``(n_clusters,)`` int64 shell index of each cluster, into ``Tr``'s middle + axis. + a_coef, b_coef : torch.Tensor + ``(L, L)`` recurrence coefficients, zero for ``m >= l``. + sect : torch.Tensor + ``(L,)`` sectoral factors. + """ + L = Tr.shape[-1] + cos_e = rep_cos.unsqueeze(-1) + prev2 = torch.zeros_like(Dr) + prev1 = torch.zeros_like(Dr) + prev1[:, 0] = LEGENDRE_SEED # bar_P_0^0 + for l in range(1, L): + # `a_coef` and `b_coef` are zero for m >= l, so this runs at full width. + # Narrowing it to the l+1 columns that can be non-zero was measured and + # is slower: the saving is real (the summed width at L=101 falls from + # 10100 to 6481) but a strided scatter target costs more than it, and the + # recurrence did not speed up at all -- it is not arithmetic-bound. The + # fused kernel gets the ragged widths for free, as loop bounds. + cur = a_coef[l] * cos_e * prev1 - b_coef[l] * prev2 + # The sectoral term must land BEFORE the products below are formed: it is + # the m = l entry of this very row. Forming `cur * Dr` first silently + # drops that entry for every even l. + cur[:, l] = sect[l] * rep_sin * prev1[:, l - 1] + if l >= 2 and (l % 2 == 0): + pos = (l - 2) // 2 + Tr[pos].index_add_(0, shell, cur * Dr) + Ti[pos].index_add_(0, shell, cur * Di) + prev2, prev1 = prev1, cur + + +__all__ = ["legendre_shell_accumulate"] diff --git a/torchref/experimental/alignment/frf/peak_finder.py b/torchref/experimental/alignment/frf/peak_finder.py new file mode 100644 index 00000000..19692db1 --- /dev/null +++ b/torchref/experimental/alignment/frf/peak_finder.py @@ -0,0 +1,162 @@ +"""Peak finding on the adaptive SO(3) sample list. + +Phaser source: ``SiteListAng::findpeaks`` (referenced from FastRot.cc; +implementation in ``SiteListAng.cc`` ``findpeaks`` and the related NMS +routines). Phaser's strategy is essentially: + + 1. Compute mean + std of all samples → z-score per sample. + 2. Sort by descending value. + 3. Greedy non-max suppression on SO(3) by *angular distance* between + rotations (not by α, β, γ box distance — that would double-count + near the poles). + +We implement the same flow in PyTorch, vectorised where possible, and the +suppression is **modulo the crystal's point group** when the Cartesian +symmetry rotations are supplied: an orientation and its symmetry mates are one +answer, and without this the shortlist handed downstream is mostly copies. On +3K7M (P432) 187 of the 300 pairs among the top 25 peaks were mates of each +other; on 2DQ6 (P3(1)21) 15 of 25 candidates were one orientation. + +The group acts on the **right**: a peak ``R`` maps the search-model frame onto +the crystal frame, and its mates are ``R R_g``. Measured, not assumed -- +composing on the left finds zero coincident pairs on every structure tried, the +right side finds all of them. +""" +from __future__ import annotations + +import math +from typing import List, Optional + +import torch + +from ....base.alignment.rotation import rotation_matrix_euler_zyz +from .types import AdaptiveRotationFunction, RotationPeak + +__all__ = [ + "find_rotation_peaks", +] + + +def _so3_greedy_nms( + alphas: torch.Tensor, + betas: torch.Tensor, + gammas: torch.Tensor, + values: torch.Tensor, + nms_radius_deg: float, + keep_at_most: int, + sym_cart: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Return indices (into the input order) of kept peaks after SO(3) NMS. + + Greedy: walk the values in descending order; keep a candidate if its + angular distance from every already-kept rotation -- and, with + ``sym_cart``, from every point-group mate ``R R_g`` of it -- is + > nms_radius_deg. + """ + n = values.shape[0] + if n == 0: + return torch.empty(0, dtype=torch.int64, device=values.device) # dtype-ok: index tensor; index_add_/gather need int64 + # The greedy walk is inherently sequential and latency-bound; on GPU a + # per-iteration `.item()` sync would dominate. Move the (tiny) candidate + # rotations to CPU once and run the loop there with no device syncs, a + # preallocated kept-buffer (no repeated torch.stack), and a cosine threshold + # (no per-iteration arccos). Result is identical to the original distance test. + order = torch.argsort(values, descending=True).cpu().tolist() + # `rotation_matrix_euler_zyz` is the shared implementation, and it is now the + # only one -- the alignment package's own batch copy was deleted after being + # measured bit-identical to it over 180k elements in both float32 and + # float64. (The three-matrix product it used reduces to the same two-term + # sums, because the rotation factors carry exact zeros and ones.) Rounding + # matters here beyond tidiness: the NMS threshold below flips for pairs + # sitting exactly on it. + R_all = ( + rotation_matrix_euler_zyz(torch.stack([alphas, betas, gammas], dim=-1)) + .cpu().to(torch.float64) # dtype-ok: 3x3 rotation algebra in double on the host + ) # (n, 3, 3) + # angle > nms_radius ⇔ cos(angle) < cos(nms_radius); cos(angle) from trace. + cos_thresh = math.cos(math.radians(nms_radius_deg)) + # The orbit of each kept rotation, R R_g over the point group; the identity + # alone when no symmetry is supplied. + # `.cpu()` first in both: the widening happens on the host, which always + # has float64, and the peaks arrive on the compute device. + G = (torch.eye(3, dtype=torch.float64).unsqueeze(0) if sym_cart is None # dtype-ok: 3x3 rotation algebra in double on the host + else sym_cart.cpu().to(torch.float64)) # dtype-ok: 3x3 rotation algebra in double on the host + kept_idx: List[int] = [] + kept_orbit = torch.empty((keep_at_most, G.shape[0], 3, 3), dtype=torch.float64) # dtype-ok: 3x3 rotation algebra in double on the host + count = 0 + for i_t in order: + Ri = R_all[i_t] + if count > 0: + trace = torch.einsum("kgij,ij->kg", kept_orbit[:count], Ri) + cos_theta = ((trace - 1.0) * 0.5).clamp(min=-1.0, max=1.0) + # Some kept rotation, or a mate of one, within nms_radius → skip. + if bool((cos_theta > cos_thresh).any()): + continue + kept_orbit[count] = Ri.unsqueeze(0) @ G + kept_idx.append(i_t) + count += 1 + if count >= keep_at_most: + break + return torch.tensor(kept_idx, dtype=torch.int64, device=values.device) # dtype-ok: index tensor; index_add_/gather need int64 + + +def find_rotation_peaks( + arf: AdaptiveRotationFunction, + n_peaks: int = 500, + sigma_threshold: float = -5.0, + nms_radius_deg: float = 6.0, + sym_cart: Optional[torch.Tensor] = None, +) -> List[RotationPeak]: + """Greedy SO(3) NMS over the adaptive sample list. + + Returns peaks sorted by descending value, capped at ``n_peaks`` and + filtered by ``sigma >= sigma_threshold``. With ``sym_cart`` -- the crystal's + point-group rotations as Cartesian ``(n_ops, 3, 3)`` matrices -- symmetry + mates of a kept peak are suppressed too, so every returned peak is a + distinct orientation. + """ + values = arf.values + if values.numel() == 0: + return [] + + mean = values.mean() + std = values.std().clamp(min=1e-30) + sigma = (values - mean) / std + + # Pre-filter by sigma threshold to keep the NMS loop tractable. + keep_mask = sigma >= sigma_threshold + if not keep_mask.any(): + return [] + + idx_filtered = torch.nonzero(keep_mask, as_tuple=False).squeeze(-1) + # Optionally cap candidate set so the O(n_kept · n_cand) NMS stays small. + candidate_cap = max(n_peaks * 20, 2000) + if idx_filtered.numel() > candidate_cap: + top_vals = values[idx_filtered] + top_keep = torch.topk(top_vals, candidate_cap).indices + idx_filtered = idx_filtered[top_keep] + + a = arf.alphas[idx_filtered] + b = arf.betas[idx_filtered] + g = arf.gammas[idx_filtered] + v = values[idx_filtered] + + kept = _so3_greedy_nms( + a, b, g, v, + nms_radius_deg=nms_radius_deg, + keep_at_most=n_peaks, + sym_cart=sym_cart, + ) + + # Gather kept peaks and move to CPU once (avoids a per-peak device sync). + a_k = a[kept].cpu().tolist() + b_k = b[kept].cpu().tolist() + g_k = g[kept].cpu().tolist() + v_k = v[kept] + s_k = ((v_k - mean) / std).cpu().tolist() + v_k = v_k.cpu().tolist() + peaks: List[RotationPeak] = [ + RotationPeak(alpha=a_k[i], beta=b_k[i], gamma=g_k[i], score=v_k[i], sigma=s_k[i]) + for i in range(len(a_k)) + ] + return peaks diff --git a/torchref/experimental/alignment/frf/preprocessing.py b/torchref/experimental/alignment/frf/preprocessing.py new file mode 100644 index 00000000..63977817 --- /dev/null +++ b/torchref/experimental/alignment/frf/preprocessing.py @@ -0,0 +1,360 @@ +"""Observed-side preprocessing chain. + +Mirrors the chain in Phaser ``DataMR::dataMR_FRF`` (DataMR.cc:863-1133) and the +auxiliary helpers in ``lib/math_RiceLLG.cc``, and carries the Phaser source +citations for each piece. + +Normalisation is deliberately **not** here. Turning amplitudes into E values is +:class:`~torchref.scaling.WilsonNormaliser`'s job, shared with the translation +search and with everything else in the repo that asks what the mean intensity at +a resolution is. What lives here is the LERF1 intensity built from those E +values, the multiplicity handling, and the symmetry detection. +""" +from __future__ import annotations + +import math +from typing import Optional + +import torch + +from ....config import get_float_dtype + +from ..sh import ( + get_high_order_axis, # phaser's highOrderAxis() + compute_patterson_shell_variance, + equal_count_shell_edges, + assign_shells, +) + + +def eterm_sigma_a(s_mag: torch.Tensor, delta_vrms_A: float) -> torch.Tensor: + """Phaser's σA Eterm, literal port of ``Ensemble.cc:42``: + + Eterm(s) = exp(-(2π²/3) · s² · ΔVRMS_var) + + where ``ΔVRMS_var`` is the *coordinate variance* in Ų. We accept the + RMS coordinate error ``delta_vrms_A`` (Å) per the standard σA + convention and square it internally: ``ΔVRMS_var = delta_vrms_A²``. + """ + s2 = s_mag * s_mag + return torch.exp(-(2.0 / 3.0) * (math.pi ** 2) * s2 * (delta_vrms_A ** 2)) + +__all__ = [ + "eterm_sigma_a", + "get_high_order_axis", + "build_lerf1_intensity", + "apply_shell_variance_weights", + "detect_zsymm", + "bulk_solvent_factor", + "oeffner_vrms", + "fit_relative_wilson_b", +] + + +def build_lerf1_intensity( + eEobs: torch.Tensor, + centric_obs: torch.Tensor, + weight: Optional[torch.Tensor] = None, + use_centric_weight: bool = True, +) -> torch.Tensor: + """LERF1 observed intensity: ``cweight · (eEobs² − 1) · weight``. + + Phaser source: ``DataMR::m_LETF1`` (DataMR.cc:1326-1431) — the + intensity that gets fed into the Bessel-SH expansion. cweight is + ε(h) · (1 for centric, 2 for acentric); we use the centric/acentric + factor only (the ε(h) multiplicity is implicit in the symmetry + reduction of the input reflection set). + + ``weight`` is the per-reflection information weight that travels with + ``eEobs`` -- ``DFAC**2`` for the French-Wilson convention, ones for a + convention that does not model measurement error. It arrives already + squared because it is the E convention that decides what the weight *is*; + this function's job is to apply one, not to know it came from a D factor. + + Note the ``- 1``: the LERF1 intensity is CENTRED, which is what makes + `` = 1`` load-bearing rather than cosmetic. A convention whose + mean square is not one puts a constant offset into every shell. + """ + if use_centric_weight: + cw = torch.where( + centric_obs.bool(), + torch.ones_like(eEobs), + 2.0 * torch.ones_like(eEobs), + ) + else: + cw = torch.ones_like(eEobs) + if weight is None: + weight = torch.ones_like(eEobs) + return cw * (eEobs * eEobs - 1.0) * weight + + +def apply_shell_variance_weights( + intensity: torch.Tensor, + s_mag: torch.Tensor, + n_var_shells: int = 20, + shell_idx: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Per-shell empirical variance reweight. + + Downweights shells whose observed Patterson intensity is dominated + by noise. Mean-normalised so total scale doesn't shift. Closest + Phaser analog is per-shell BINS + ``best(r)`` in ``Ensemble.cc``. + + ``shell_idx`` reuses an assignment the caller already made. Worth passing: + binning here independently of the Wilson normalisation puts the two on + edges that disagree for the reflections sitting on a boundary. + """ + if shell_idx is None: + edges, _ = equal_count_shell_edges(s_mag, n_var_shells) + shell_idx = assign_shells(s_mag, edges) + valid = shell_idx >= 0 + var_p = compute_patterson_shell_variance( + intensity[valid], + shell_idx[valid], + P=n_var_shells, + ) + inv_sqrt_var = 1.0 / var_p.sqrt().clamp(min=1e-30) + inv_sqrt_var = inv_sqrt_var * ( + n_var_shells / inv_sqrt_var.sum().clamp(min=1e-30) + ) + weights = torch.ones_like(intensity) + weights[valid] = inv_sqrt_var[shell_idx[valid]].to(intensity.dtype) + return intensity * weights + + +def detect_zsymm(sym_mats: Optional[torch.Tensor]) -> int: + """Detect ``ZSYMM`` for the **z-axis** m-symmetry filter. + + Phaser source: ``highOrderAxis()`` in ``rotationgroup.h``. The m-filter + zeroes obs SH coefficients with ``|m| mod ZSYMM != 0`` — but ``m`` is the + azimuthal order **about the z axis of the SH basis**, so the filter is only + valid when the crystal's high-order rotation axis is actually along z + (cubic / tetragonal / hexagonal: principal axis = c ∥ z). + + For spacegroups whose high-order axis is x or y — e.g. monoclinic C2 / P2₁ + with the 2-fold along b ∥ y — applying a z-axis filter is WRONG (it filters + about the wrong axis and corrupts the obs coefficients). Phaser rotates the + high-order axis to z before the expansion (DataMR.cc:962-979); we do not, so + here we conservatively return ``ZSYMM=1`` (no filter) when the axis is not z, + rather than apply a wrong filter. Verified: for all benchmark cubic/tetra/hex + cases the axis is z (filter unchanged); only monoclinic 1DAW/3E98/3VRJ change + (they already rank 0-3, and a 2-fold m-filter is a weak constraint anyway). + """ + if sym_mats is None: + return 1 + # No cast and no copy: `get_axis_order` works at the operators' own width + # and batches its one readback. Widening them here was also a crash on a + # backend with no float64, since these arrive on the compute device. + axis, zsymm = get_high_order_axis(sym_mats) + if axis != 2: # high-order axis not along z → don't apply a wrong filter + return 1 + return int(zsymm) + + +def bulk_solvent_factor( + s_mag: torch.Tensor, + fsol: float = 0.95, + bsol: float = 300.0, + sigA_min: float = 0.01, +) -> torch.Tensor: + """Phaser's Babinet bulk-solvent term — ``solTerm.h:9``:: + + solTerm(s²) = max(SIGA_MIN, 1 − fsol · exp(−bsol · s²/4)) + + Models the bulk solvent's contribution to the structure factor via Babinet's + principle. At low resolution (s→0) the term → ``1 − fsol`` ≈ 0.05 (with the + default ``fsol=0.95``), aggressively suppressing the calc — physically, the + model represents only the macromolecule, but the diffraction data sees + macromolecule + bulk solvent, and at low resolution the solvent's flat + average density partially cancels the macromolecule's contribution. At high + resolution (s→∞) the term → 1 (no effect). + + Phaser folds this into the effective σ_A via + ``σ_A_eff(s) = solTerm(s²) · DLuzzati(s², vrms)`` (EnsemblePDB.cc:96-100). + For callers that work in σ_A space (the rescore, the FRF eterm), multiplying + by this factor reproduces that behaviour. + + Defaults match Phaser (``DEF_SOLPAR_BULK_FSOL=0.95``, + ``DEF_SOLPAR_BULK_BSOL=300``, ``DEF_SOLPAR_SIGA_MIN=0.01``). + + Parameters + ---------- + s_mag : tensor + Per-reflection reciprocal-space magnitude |s| (Å^-1). + fsol, bsol, sigA_min : float + Babinet parameters. Defaults match Phaser. + + Returns + ------- + torch.Tensor + Per-reflection solvent multiplier, same shape as ``s_mag``. Always in + ``[sigA_min, 1]``. + """ + s2 = s_mag * s_mag + babinet = 1.0 - float(fsol) * torch.exp(-float(bsol) * s2 / 4.0) + return babinet.clamp(min=float(sigA_min)) + + +def oeffner_vrms(n_residues: int, identity: float = 1.0) -> float: + """Phaser's Oeffner empirical vrms estimate — ``rms_estimate.cc:37``:: + + vrms = A · (B + clamp(n_residues, 125, 1500))^(1/3) · exp(C · (1 − ident)) + + with ``A = 0.0569``, ``B = 173``, ``C = 1.52``. The clamp avoids extrapolating + beyond the well-populated range of Oeffner et al.'s training set + (Acta Cryst. (2013) D69:2209-2215). For a perfect model (``identity=1``) and + a typical protein (~300 residues), this gives vrms ≈ 0.47 Å; large + assemblies (clamped at 1500) give vrms ≈ 0.67 Å. Phaser uses this as the + Luzzati ``vrms`` for the σ_A computation. + + Parameters + ---------- + n_residues : int + Sequence length of the search model. Internally clamped to [125, 1500]. + identity : float, optional + Sequence identity to expected target on [0, 1]. Default 1.0 (perfect). + + Returns + ------- + float + Coordinate RMS estimate in Å, suitable as ``delta_vrms_A`` for + :func:`compute_sigma_a_luzzati` / :func:`eterm_sigma_a`. + """ + A, B, C = 0.0569, 173.0, 1.52 + n_clamped = max(125, min(int(n_residues), 1500)) + return A * (B + n_clamped) ** (1.0 / 3.0) * math.exp(C * (1.0 - float(identity))) + + +def fit_relative_wilson_b( + F_obs: torch.Tensor, + F_calc: torch.Tensor, + s_mag: torch.Tensor, + n_shells: int = 20, + clamp_b: float = 50.0, + s_mag_calc: Optional[torch.Tensor] = None, +) -> float: + """Phaser's relative Wilson-B fit — ``EnsemblePDB.cc:793-851``. + + Estimates the per-model relative Wilson B-factor that brings the calc's + per-shell <|F_calc|²> into agreement with the data's per-shell <|F_obs|²>. + Implemented as a weighted linear regression of + ``log(_shell / _shell)`` against ``s²`` over equal-count + shells; returns ``WilsonB = -2 · slope``. + + Per-shell weighting mirrors Phaser (EnsemblePDB.cc:830-835): + - ``s² < 0.009`` (d > 10.5 Å): weight = 0 (no reliable BEST curve). + - ``s² < 0.04`` (d > 5 Å): weight = ``(s²/0.04)²`` (down-weighted). + - ``s² ≥ 0.04``: weight = 1. + + Apply at call sites as ``F_calc · exp(-WilsonB · s² / 4)``. + + Parameters + ---------- + F_obs : (N_obs,) tensor + F_calc : (N_calc,) tensor + Calc amplitudes. May be on a different reciprocal grid than obs (e.g. + the dense P1-box from ``dense_calc_via_box``); per-shell means handle + the binning independently. + s_mag : (N_obs,) tensor + Reciprocal-space magnitudes for OBS. Used to derive shell edges by + equal-count binning on the obs distribution. + s_mag_calc : (N_calc,) tensor, optional + Reciprocal-space magnitudes for CALC. Defaults to ``s_mag`` (when obs + and calc share a grid). When given, obs and calc are binned into the + SAME edges (derived from obs ``s_mag``) but with independent counts. + n_shells : int + Number of equal-count resolution shells (on obs). + clamp_b : float + Clamps the fitted B to ``[-clamp_b, +clamp_b]``. + + Returns + ------- + float + Relative Wilson B-factor (Ų). ``0`` if too few shells contribute. + """ + if s_mag_calc is None: + s_mag_calc = s_mag + if F_obs.shape[0] != s_mag.shape[0]: + raise ValueError( + f"F_obs / s_mag length mismatch: {F_obs.shape[0]} vs {s_mag.shape[0]}" + ) + if F_calc.shape[0] != s_mag_calc.shape[0]: + raise ValueError( + f"F_calc / s_mag_calc length mismatch: " + f"{F_calc.shape[0]} vs {s_mag_calc.shape[0]}" + ) + + edges, _ = equal_count_shell_edges(s_mag, n_shells) + # Bin obs and calc into the SAME edges (independently — different N's). + shell_idx_obs = assign_shells(s_mag, edges) + shell_idx_calc = assign_shells(s_mag_calc, edges) + valid_obs = shell_idx_obs >= 0 + valid_calc = shell_idx_calc >= 0 + if not valid_obs.any() or not valid_calc.any(): + return 0.0 + + real = get_float_dtype() + F2_obs = (F_obs * F_obs).to(real) + F2_calc = (F_calc * F_calc).to(real) + s2_obs = (s_mag * s_mag).to(real) + + counts_obs = torch.zeros(n_shells, dtype=torch.int64, device=s_mag.device) # dtype-ok: per-shell counts + counts_calc = torch.zeros(n_shells, dtype=torch.int64, device=s_mag.device) # dtype-ok: per-shell counts + sum_F2obs = torch.zeros(n_shells, dtype=real, device=s_mag.device) + sum_F2calc = torch.zeros(n_shells, dtype=real, device=s_mag.device) + sum_s2 = torch.zeros(n_shells, dtype=real, device=s_mag.device) + idx_v_obs = shell_idx_obs[valid_obs] + idx_v_calc = shell_idx_calc[valid_calc] + counts_obs.index_add_(0, idx_v_obs, torch.ones_like(idx_v_obs)) + counts_calc.index_add_(0, idx_v_calc, torch.ones_like(idx_v_calc)) + sum_F2obs.index_add_(0, idx_v_obs, F2_obs[valid_obs]) + sum_F2calc.index_add_(0, idx_v_calc, F2_calc[valid_calc]) + sum_s2.index_add_(0, idx_v_obs, s2_obs[valid_obs]) # obs-side s² for the regression abscissa + + # Drop shells empty on either side. + keep = (counts_obs > 0) & (counts_calc > 0) + mean_F2obs = sum_F2obs[keep] / counts_obs[keep].to(real) + mean_F2calc = sum_F2calc[keep] / counts_calc[keep].to(real) + mean_s2 = sum_s2[keep] / counts_obs[keep].to(real) + + # log(Σ_N / Σ_P) per shell. + eps = 1e-30 + log_ratio = (mean_F2obs.clamp(min=eps) / mean_F2calc.clamp(min=eps)).log() + + # Phaser per-shell weights (EnsemblePDB.cc:830-835). + weights = torch.ones_like(mean_s2) + low_mask = mean_s2 < 0.04 + weights[low_mask] = (mean_s2[low_mask] / 0.04) ** 2 + weights[mean_s2 < 0.009] = 0.0 + if (weights > 0).sum().item() < 2: + return 0.0 + + # Weighted linear fit y = slope · x (no intercept), per the Phaser source. + w = weights + x = mean_s2 + y = log_ratio + sw = w.sum() + swx = (w * x).sum() + swy = (w * y).sum() + swx2 = (w * x * x).sum() + swxy = (w * x * y).sum() + denom = (sw * swx2 - swx * swx).item() + if abs(denom) < 1e-30: + return 0.0 + slope = (sw * swxy - swx * swy).item() / denom + # Phaser: WilsonB_intensity = -4·slope, then halved → WilsonB = -2·slope. + wilson_b = -2.0 * slope + return float(max(-clamp_b, min(wilson_b, clamp_b))) + + +# Note: a naive OLS-on-log-F² fit of anisotropic Wilson U was attempted on +# 2026-05-28 and didn't work. The Wilson left tail (small F values produce huge +# negative log F²) dominates the regression, returning U components of order +# 10²–10³ Ų on real data — three orders of magnitude beyond physical, on both +# easy (1DAW) and hard (2DQ6) cases. Robustifying via |F|-weighting + ridge +# only made the fit saturate any sensible clamp. Phaser's ``scaleANIS`` +# (``DataB.cc``, ~500 LoC) is an iterative ML fit on ``logSigmaEsq``; that's +# the right approach if obs-side aniso ever becomes the next lever. The 2DQ6 +# benchmark failure we were chasing turned out to be tNCS +# (``<(E²−1)²>_acentric = 5.5`` vs Wilson = 1.0), not anisotropy, so this +# branch is not in the immediate critical path. diff --git a/torchref/experimental/alignment/frf/rotation_utils.py b/torchref/experimental/alignment/frf/rotation_utils.py new file mode 100644 index 00000000..51fab749 --- /dev/null +++ b/torchref/experimental/alignment/frf/rotation_utils.py @@ -0,0 +1,92 @@ +"""Pure-geometry helpers shared by the FRF, the rescore, and tests. + +Edmonds active ZYZ convention throughout: a rotation matrix is built as +``R = R_z(α) R_y(β) R_z(γ)``. ``α, γ ∈ [0, 2π)``, ``β ∈ [0, π]``. +""" +from __future__ import annotations + +import math +from typing import Tuple + +import torch + + +def rotation_matrix_from_edmonds_euler( + alpha: float, beta: float, gamma: float, dtype=torch.float64, # dtype-ok: 3x3 rotation algebra in double on the host +) -> torch.Tensor: + """Build ``R = R_z(α) R_y(β) R_z(γ)`` (Edmonds active ZYZ). + + Equivalent to passing ``[γ, β, α]`` to + ``torchref.experimental.alignment.transform.rotation_matrix_from_euler``. + """ + ca, sa = math.cos(alpha), math.sin(alpha) + cb, sb = math.cos(beta), math.sin(beta) + cg, sg = math.cos(gamma), math.sin(gamma) + Rz_a = torch.tensor([[ca, -sa, 0.0], [sa, ca, 0.0], [0.0, 0.0, 1.0]], dtype=dtype) + Ry_b = torch.tensor([[cb, 0.0, sb], [0.0, 1.0, 0.0], [-sb, 0.0, cb]], dtype=dtype) + Rz_c = torch.tensor([[cg, -sg, 0.0], [sg, cg, 0.0], [0.0, 0.0, 1.0]], dtype=dtype) + return Rz_a @ Ry_b @ Rz_c + + +def edmonds_euler_from_rotation_matrix(R: torch.Tensor) -> Tuple[float, float, float]: + """Recover ``(α, β, γ)`` such that ``R = R_z(α) R_y(β) R_z(γ)``. + + Returns angles in radians; ``α, γ ∈ [0, 2π)``, ``β ∈ [0, π]``. Singular + when ``β = 0`` or ``π`` (only ``α+γ`` is determined); in those cases + ``γ=0`` is returned. + """ + R = R.to(torch.float64) # dtype-ok: 3x3 rotation algebra in double on the host + cos_beta = R[2, 2].clamp(-1.0, 1.0).item() + beta = math.acos(cos_beta) + sin_beta = math.sin(beta) + if abs(sin_beta) < 1e-9: + alpha = math.atan2(R[1, 0].item(), R[0, 0].item()) + gamma = 0.0 + else: + alpha = math.atan2(R[1, 2].item(), R[0, 2].item()) + gamma = math.atan2(R[2, 1].item(), -R[2, 0].item()) + alpha = alpha % (2.0 * math.pi) + gamma = gamma % (2.0 * math.pi) + return alpha, beta, gamma + + +def axis_angle_to_matrix(omega: torch.Tensor) -> torch.Tensor: + """Rodrigues axis-angle → SO(3). ``omega = θ · axis`` (radians). + + Accepts ``(3,)`` for a single rotation or ``(..., 3)`` for a batched stack + and returns ``(3, 3)`` or ``(..., 3, 3)``. The small-θ limit is handled + implicitly (sin θ→0, (1−cos θ)→0 ⇒ R→I); ``clamp(min=1e-30)`` guards the + axis normalisation at θ=0. + + Preferred over ``base.alignment.rotation.axis_angle_to_rotation_matrix``, + which accepts only ``(3,)``/``(N, 3)`` and switches the axis to ``[0, 0, 1]`` + below θ = 1e-10 rather than letting the trigonometric factors vanish. Above + that threshold the two agree term for term. + """ + if omega.dtype not in (torch.float32, torch.float64): # dtype-ok: 3x3 rotation algebra in double on the host + omega = omega.to(torch.float64) # dtype-ok: 3x3 rotation algebra in double on the host + single = omega.dim() == 1 + if single: + omega = omega.unsqueeze(0) + th = omega.norm(dim=-1, keepdim=True) # (..., 1) + axis = omega / th.clamp(min=1e-30) # (..., 3) + zeros = torch.zeros_like(axis[..., 0]) + K = torch.stack([ + torch.stack([zeros, -axis[..., 2], axis[..., 1]], dim=-1), + torch.stack([axis[..., 2], zeros, -axis[..., 0]], dim=-1), + torch.stack([-axis[..., 1], axis[..., 0], zeros], dim=-1), + ], dim=-2) # (..., 3, 3) + th_b = th.unsqueeze(-1) # (..., 1, 1) + eye = torch.eye(3, dtype=omega.dtype, device=omega.device).expand( + *omega.shape[:-1], 3, 3 + ) + R = eye + torch.sin(th_b) * K + (1.0 - torch.cos(th_b)) * (K @ K) + return R.squeeze(0) if single else R + + +def rotation_angular_distance_deg(R1: torch.Tensor, R2: torch.Tensor) -> float: + """Geodesic distance on SO(3) in degrees: ``arccos((tr(R1 R2^T) − 1)/2)``.""" + R = R1.to(torch.float64) @ R2.to(torch.float64).T # dtype-ok: 3x3 rotation algebra in double on the host + tr = (R[0, 0] + R[1, 1] + R[2, 2]).clamp(-1.0, 3.0).item() + cos_a = max(-1.0, min(1.0, (tr - 1.0) / 2.0)) + return math.degrees(math.acos(cos_a)) diff --git a/torchref/experimental/alignment/frf/sitelist_ang.py b/torchref/experimental/alignment/frf/sitelist_ang.py new file mode 100644 index 00000000..c297c3c2 --- /dev/null +++ b/torchref/experimental/alignment/frf/sitelist_ang.py @@ -0,0 +1,343 @@ +"""Per-β rotation function evaluation on Phaser's adaptive SO(3) sample list. + +Mirrors ``SiteListAng`` from +``reverse_engineering/phenix/phenix-1.20-4459/modules/phaser/codebase/phaser/src/FastRot.cc``. + +The crucial point is that **the FFT itself is NOT per-β-adaptive**. +Phaser does: + +1. ``get_FRF`` (FastRot.cc:90-167) loops over a uniform β grid + ``β_b = b · Δ`` for ``b ∈ [0, bmax)``, ``bmax = ceil(180/Δ)``. +2. For each β: ``DoRfftStuff`` (FastRot.cc:19-88) builds the per-β + Fourier-mode amplitudes + ``S_{m1, m2}(β) = Σ_l ξ_{l, m1, m2} · d^l_{m1, m2}(β)`` + on the full ``(2L-1) × (2L-1)`` grid (asymmetric-unit storage only — + the Friedel mate is added by cctbx via ``conjugate_flag=true``). +3. The 2D inverse FFT runs at a **fixed shape** + ``amax = find_fft_friendly_size(2·max(bmax, lmax))`` for every β. + The result is a dense ``M_β(α, γ)`` map indexed in ``[0, 1)`` along + each axis. +4. The **adaptive sample list** is built once by ``allocate_memory`` + (FastRot.cc:169-262): for each β, + ``pmax(β) = 720/Δ · cos(β/2)`` + ``qmax(β) = 360/Δ · sin(β/2)`` + and the ``(p, q)`` lattice is mapped to ``(α, γ)`` via + ``α = (p/pmax + q/qmax) mod 1`` + ``γ = sign · (p/pmax − q/qmax) mod 1`` (FastRot.cc:216-219) + with the β=0 special case keeping only the ``p == p`` diagonal + (FastRot.cc:189-207) because only ``α + γ`` is meaningful at the pole. +5. ``M_β`` is **bilinearly interpolated** at each ``(α, γ)`` sample point + to give the RF value (FastRot.cc:146-152, ``four_point_interpolation``). + +Making the FFT shape itself ``(pmax(β), qmax(β))`` would collapse to +``(N, 1)`` at small β and lose all γ Fourier information. The FFT stays +dense; adaptivity is only in the sample list and the interpolation. +""" +from __future__ import annotations + +import math +from typing import List, Tuple + +import torch + +from ....config import canonical_device +from ....symmetry.symmetry import find_fft_friendly_size +from .types import AdaptiveRotationFunction +from .wigner_d import wigner_contraction_per_beta + +__all__ = [ + "build_dense_map_per_beta", + "build_adaptive_sample_list", + "evaluate_rotation_function", +] + + +def build_dense_map_per_beta( + xi_lmn: torch.Tensor, + betas: torch.Tensor, + fft_size: int, +) -> torch.Tensor: + """Return the dense FFT map ``M_β(α, γ)`` for every β. + + Phaser source: ``DoRfftStuff`` (FastRot.cc:19-88), but tensor-batched + over β and with a single 2D ``torch.fft.ifft2`` per β instead of a + cctbx ``real_to_complex_3d`` of shape ``(1, amax, amax)`` — they + produce equivalent dense (α, γ) grids. + + Parameters + ---------- + xi_lmn : torch.Tensor (complex), shape (L, 2L-1, 2L-1) + Cross-correlation coefficients with l ∈ [0, L), m, n ∈ [-(L-1), L-1]. + betas : torch.Tensor (real), shape (n_beta,) + β values in radians. + fft_size : int + Fixed FFT grid size N. The map ``M_β`` will be ``(N, N)`` for every β, + indexed as ``M[k', l'] = RF(2π k'/N, β, 2π l'/N) / N²``. + + Returns + ------- + M : torch.Tensor (complex), shape (n_beta, fft_size, fft_size) + The dense maps. Use bilinear interpolation in the (α, γ) plane to + evaluate at non-grid points. + """ + L = xi_lmn.shape[0] + n_beta = betas.shape[0] + device = xi_lmn.device + + # 1. Per-β Wigner-d contraction: S[k, m1+L-1, m2+L-1] = Σ_l ξ d^l_{m1,m2}(β_k). + S = wigner_contraction_per_beta(xi_lmn, betas) # (n_beta, 2L-1, 2L-1) + + # 2. Place S into the (fft_size, fft_size) Fourier grid by FFT-frequency + # mapping: m → (m mod N), n → (n mod N). For m, n ∈ [-(L-1), L-1] and + # N >> 2L-1, this puts negative frequencies at the high end of each axis. + if fft_size < 2 * L - 1: + raise ValueError( + f"fft_size={fft_size} must be >= 2L-1={2*L-1} to avoid aliasing" + ) + pad = torch.zeros( + (n_beta, fft_size, fft_size), dtype=S.dtype, device=device, + ) + m_vals = torch.arange(-(L - 1), L, device=device) + idx = (m_vals % fft_size).to(torch.int64) # dtype-ok: index tensor; index_add_/gather need int64 + pad[:, idx.unsqueeze(1), idx.unsqueeze(0)] = S + + # 3. Forward 2D FFT — torch convention: + # fft2(X)[k, l] = Σ_{m, n} X[m, n] · exp(-2πi (m·k/N + n·l/M)) + # which gives M[k, l] = RF(α = 2π·k/N, β, γ = 2π·l/N) directly, with + # Edmonds D^l_{m,n} = exp(-imα) d^l_{m,n}(β) exp(-inγ). Using ifft2 + # here (the cctbx default in Phaser's pipeline) would give RF at + # (-α, -γ) which then requires negating alpha_frac/gamma_frac when + # recording sample Euler angles — Phaser's FastRot.cc:153 does this + # explicit ``-360 * alpha`` flip. We do the equivalent by using fft2 + # so the Euler labels are already in the right sign. + M = torch.fft.fft2(pad, dim=(-2, -1)) + return M + + +# Module-level memo for the data-independent sample list, keyed on +# (grid_sampling_deg, device-str, dtype). The list depends only on geometry, so +# repeat FRF calls at the same grid reuse it; a single cold call still pays the +# (now vectorised, CPU-built) construction once. +_SAMPLE_LIST_CACHE: dict = {} + + +def build_adaptive_sample_list( + grid_sampling_deg: float, + dtype: torch.dtype = torch.float64, # dtype-ok: sample-list geometry follows the accumulator's width + device: torch.device = torch.device("cpu"), +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the per-β (α, γ) sample list. + + Phaser source: ``SiteListAng::allocate_memory`` (FastRot.cc:169-262). + Returns (in radians): + alphas : (N_samples,) α value per sample + betas : (N_samples,) β value per sample + gammas : (N_samples,) γ value per sample + beta_starts: (bmax + 1,) int64 slice [beta_starts[b]:beta_starts[b+1]] + is the samples at β = b · Δ + beta_grid : (bmax,) the β values in radians + + The construction is purely geometric (independent of the data / ξ), so it is + memoised on ``(grid_sampling_deg, device, dtype)``. It is also built on the + CPU — the per-β tensors are tiny and the work is launch-latency-bound, so a + single host-side build + one device transfer is far cheaper than thousands of + small CUDA kernels. The per-key dedup is vectorised (no ``.tolist()`` / + Python scan), so there is no host sync inside the loop. + """ + device = torch.device(device) if not isinstance(device, torch.device) else device + cache_key = (float(grid_sampling_deg), str(canonical_device(device)), dtype) + cached = _SAMPLE_LIST_CACHE.get(cache_key) + if cached is not None: + return cached + + bmax = int(math.ceil(180.0 / grid_sampling_deg)) + if bmax < 1: + raise ValueError(f"grid_sampling_deg={grid_sampling_deg} too coarse") + cpu = torch.device("cpu") + + alphas_list: List[torch.Tensor] = [] + gammas_list: List[torch.Tensor] = [] + betas_list: List[torch.Tensor] = [] + beta_starts: List[int] = [0] + + deg2rad = math.pi / 180.0 + for b in range(bmax): + beta_rad = b * grid_sampling_deg * deg2rad # plain math: no sync + cosb = math.cos(beta_rad / 2.0) + sinb = math.sin(beta_rad / 2.0) + # Truncation toward zero, with NO clamp to a minimum of 1 -- Phaser + # has none (FastRot.cc:214-215). Near beta = 180 deg, cos(beta/2) drives + # pmax to 0 and Phaser's `for (p=0; p 0 and qmax == 0): + beta_starts.append(beta_starts[-1]) + continue + + if b == 0: + # β=0: only α = γ = p/pmax for p < pmax/2 (FastRot.cc:189-207). + p_idx = torch.arange(pmax, device=cpu) + p_ratio = p_idx.to(torch.float64) / pmax # dtype-ok: sample-list geometry follows the accumulator's width + keep = p_ratio < 0.5 + p_ratio = p_ratio[keep] + alpha_frac = p_ratio + gamma_frac = p_ratio + else: + # (p, q) ∈ [0, pmax) × [0, qmax), mapped to (α, γ) via + # FastRot.cc:216-219 — with the negative branch for γ when + # p_ratio < q_ratio (gives γ ∈ [0, 1) without negative values). + p_idx = torch.arange(pmax, device=cpu) + q_idx = torch.arange(qmax, device=cpu) + p_ratio = (p_idx.to(torch.float64) / pmax).unsqueeze(1) # (pmax, 1) # dtype-ok: sample-list geometry follows the accumulator's width + q_ratio = (q_idx.to(torch.float64) / qmax).unsqueeze(0) # (1, qmax) # dtype-ok: sample-list geometry follows the accumulator's width + alpha_frac = torch.fmod(p_ratio + q_ratio, 1.0) + diff = p_ratio - q_ratio + gamma_frac = torch.where( + diff >= 0.0, + torch.fmod(diff, 1.0), + 1.0 - torch.fmod(-diff, 1.0), + ) + alpha_frac = alpha_frac.reshape(-1) + gamma_frac = gamma_frac.reshape(-1) + + # Dedup: when p ≥ pmax/2, the (p, q) → (α, γ) map can collide + # with (p − pmax/2, q') for some q' (FastRot.cc:222-241). Keep the + # first occurrence (in original order) of each rounded (α, γ) key. + # Vectorised first-occurrence: stable-sort the unique-group labels, + # mark group boundaries, scatter back — same kept set & order as the + # original dict scan, but no host sync / Python loop. + # Hash the two rounded fracs (each in [0, 1e6]) into one int64 so we + # can use the fast 1-D unique instead of a 2-D row lexsort. + a_round = (alpha_frac * 1_000_000).round().to(torch.int64) # dtype-ok: index tensor; index_add_/gather need int64 + g_round = (gamma_frac * 1_000_000).round().to(torch.int64) # dtype-ok: index tensor; index_add_/gather need int64 + key_hash = a_round * 1_000_001 + g_round + _, uniq_idx = torch.unique(key_hash, return_inverse=True) + n = uniq_idx.shape[0] + order = torch.argsort(uniq_idx, stable=True) + sorted_u = uniq_idx[order] + first_in_sorted = torch.ones(n, dtype=torch.bool, device=cpu) + first_in_sorted[1:] = sorted_u[1:] != sorted_u[:-1] + keep_mask = torch.zeros(n, dtype=torch.bool, device=cpu) + keep_mask[order[first_in_sorted]] = True + alpha_frac = alpha_frac[keep_mask] + gamma_frac = gamma_frac[keep_mask] + + n_this = alpha_frac.shape[0] + alphas_list.append((alpha_frac * (2.0 * math.pi)).to(dtype)) + gammas_list.append((gamma_frac * (2.0 * math.pi)).to(dtype)) + betas_list.append(torch.full((n_this,), beta_rad, dtype=dtype, device=cpu)) + beta_starts.append(beta_starts[-1] + n_this) + + # Concatenate on CPU, then move to the target device in one transfer each. + alphas = torch.cat(alphas_list).to(device) + gammas = torch.cat(gammas_list).to(device) + betas_flat = torch.cat(betas_list).to(device) + beta_starts_t = torch.tensor(beta_starts, dtype=torch.int64, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + b = torch.arange(bmax, dtype=torch.float64, device=cpu) # dtype-ok: sample-list geometry follows the accumulator's width + betas_rad = (b * grid_sampling_deg * deg2rad).to(device=device, dtype=dtype) + + result = (alphas, betas_flat, gammas, beta_starts_t, betas_rad) + _SAMPLE_LIST_CACHE[cache_key] = result + return result + + +def _bilinear_interp_periodic( + M: torch.Tensor, # (N, N) complex + alpha_frac: torch.Tensor, # (n,) in [0, 1) + gamma_frac: torch.Tensor, # (n,) in [0, 1) +) -> torch.Tensor: + """Periodic bilinear interpolation of M at (alpha_frac · N, gamma_frac · N). + + Mirrors ``four_point_interpolation`` (FastRot.cc:146) on a periodic + map: both axes wrap around modulo N. + """ + N = M.shape[-1] + af = (alpha_frac % 1.0) * N + gf = (gamma_frac % 1.0) * N + a0 = torch.floor(af).to(torch.int64) % N # dtype-ok: index tensor; index_add_/gather need int64 + g0 = torch.floor(gf).to(torch.int64) % N # dtype-ok: index tensor; index_add_/gather need int64 + a1 = (a0 + 1) % N + g1 = (g0 + 1) % N + da = (af - torch.floor(af)).to(M.real.dtype) + dg = (gf - torch.floor(gf)).to(M.real.dtype) + da_c = da.to(M.dtype) + dg_c = dg.to(M.dtype) + v00 = M[a0, g0] + v01 = M[a0, g1] + v10 = M[a1, g0] + v11 = M[a1, g1] + return ( + v00 * ((1 - da_c) * (1 - dg_c)) + + v01 * ((1 - da_c) * dg_c) + + v10 * (da_c * (1 - dg_c)) + + v11 * (da_c * dg_c) + ) + + +def evaluate_rotation_function( + xi_lmn: torch.Tensor, + grid_sampling_deg: float = 2.0, + fft_size: int = -1, +) -> AdaptiveRotationFunction: + """Compute the rotation function on Phaser's adaptive SO(3) sample list. + + Phaser source: composition of ``get_FRF`` + ``allocate_memory`` + + ``four_point_interpolation`` (FastRot.cc:90-262). Returns + real-valued samples (the rotation function is real). + + Parameters + ---------- + xi_lmn : torch.Tensor (complex), shape (L, 2L-1, 2L-1) + grid_sampling_deg : float + Phaser's ``grid_sampling`` keyword. β grid is uniform at this + spacing; (α, γ) sample density per β follows pmax/qmax. + fft_size : int, optional + Dense FFT shape. Default: ``find_fft_friendly_size(2·max(bmax, 2L-1))``. + """ + if xi_lmn.ndim != 3: + raise ValueError(f"xi_lmn must be 3-D (L, 2L-1, 2L-1), got {tuple(xi_lmn.shape)}") + L = xi_lmn.shape[0] + device = xi_lmn.device + real_dtype = ( + torch.float64 # dtype-ok: sample-list geometry follows the accumulator's width + if xi_lmn.dtype in (torch.complex128, torch.float64) # dtype-ok: sample-list geometry follows the accumulator's width + else torch.float32 # dtype-ok: sample-list geometry follows the accumulator's width + ) + + bmax = int(math.ceil(180.0 / grid_sampling_deg)) + if fft_size < 0: + fft_size = find_fft_friendly_size(2 * max(bmax, 2 * L - 1)) + + # 1. Build adaptive sample list (purely geometric — independent of xi). + alphas, betas_flat, gammas, beta_starts, beta_grid = build_adaptive_sample_list( + grid_sampling_deg, dtype=real_dtype, device=device, + ) + + # 2. Dense FFT map per β. + M = build_dense_map_per_beta(xi_lmn, beta_grid, fft_size) # (n_beta, N, N) + + # 3. Bilinear interp at each sample's (α, γ). + values = torch.zeros(alphas.shape[0], dtype=real_dtype, device=device) + n_beta = beta_grid.shape[0] + for b in range(n_beta): + i0, i1 = int(beta_starts[b].item()), int(beta_starts[b + 1].item()) + if i1 <= i0: + continue + af = alphas[i0:i1] / (2.0 * math.pi) + gf = gammas[i0:i1] / (2.0 * math.pi) + v_complex = _bilinear_interp_periodic(M[b], af, gf) + # The rotation function is real; the imaginary residue is at the + # numerical-noise level for a Hermitian-symmetric input ξ. Drop it. + values[i0:i1] = v_complex.real + + return AdaptiveRotationFunction( + alphas=alphas, + betas=betas_flat, + gammas=gammas, + values=values, + beta_starts=beta_starts, + beta_grid=beta_grid, + grid_sampling_deg=grid_sampling_deg, + ) diff --git a/torchref/experimental/alignment/frf/types.py b/torchref/experimental/alignment/frf/types.py new file mode 100644 index 00000000..14d85108 --- /dev/null +++ b/torchref/experimental/alignment/frf/types.py @@ -0,0 +1,66 @@ +"""Dataclasses shared across the fast rotation function. + +Mirrors the small "data carrier" structs in Phaser +(phenix-1.20-4459/modules/phaser/codebase/phaser/src/SiteListAng.h, + src/DataMR.h) without trying to keep the same names everywhere — +crystallographic intent first, C++ naming second. +""" +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass +class BesselSHCoefficients: + """Bessel-radial × spherical-harmonic coefficients ``c_{n, l, m}``. + + Phaser source: DataMR.cc constructs the equivalent ``c_lmn`` 3D complex + grid inside ``DataMR::dataMR_FRF`` (not shown here — see annotations). + + Shape: ``coeffs[n, l, m + L - 1]`` is complex, with l ∈ [0, L) even-only, + n ∈ [0, N_radial), m ∈ [-l, l] (zero-padded outside). + """ + + coeffs: torch.Tensor # (N_radial, L, 2L-1), complex + L: int # angular bandlimit (lmax = L - 1) + bessel_h_scale: float # h = bessel_h_scale * |s| + + +@dataclass +class AdaptiveRotationFunction: + """Rotation function values on Phaser's adaptive SO(3) sample list. + + The samples live on a *non-rectangular* point set per β — generated by + the ``allocate_memory`` routine in ``SiteListAng::allocate_memory`` + (FastRot.cc:169-262). Each β has a different number of samples; we + store them concatenated, with ``beta_starts`` indexing into the flat + arrays. + + The dense FFT map ``M_β(α, γ)`` (fixed shape across β) is *not* stored + here — only the interpolated samples on the adaptive list. + """ + + alphas: torch.Tensor # (N_samples,) real in [0, 2π) + betas: torch.Tensor # (N_samples,) real in [0, π] + gammas: torch.Tensor # (N_samples,) real in [0, 2π) + values: torch.Tensor # (N_samples,) real — RF values + beta_starts: torch.Tensor # (n_beta + 1,) int — slice indices per β + beta_grid: torch.Tensor # (n_beta,) real — the β values themselves + grid_sampling_deg: float + + +@dataclass +class RotationPeak: + """A single rotation function peak. + + Single source of truth across the FRF, the rescore, and the tests. + Convention: Edmonds ZYZ Euler in radians. + """ + + alpha: float + beta: float + gamma: float + score: float + sigma: float diff --git a/torchref/experimental/alignment/frf/wigner_d.py b/torchref/experimental/alignment/frf/wigner_d.py new file mode 100644 index 00000000..299133aa --- /dev/null +++ b/torchref/experimental/alignment/frf/wigner_d.py @@ -0,0 +1,171 @@ +"""Wigner small-d matrices and Wigner-D pointwise evaluation. + +Phaser source: ``phaser/lib/wigner.h`` (the C++ template +``djmn_recursive_table`` used in ``FastRot.cc:41`` per-l, per-β). + +Phaser uses the Sakurai recurrence convention; the equivalent Edmonds +(4.1.23) convention is used throughout this package. Its small-d identities are +guarded by ``tests/unit/frf_separate/test_invariants.py``. +``wigner_contraction_per_beta`` builds the small-d blocks it needs from the +``J_y`` eigendecomposition, which stays bounded to any ``l``. The blocks depend +only on the bandwidth and the β grid, so they are memoised for reuse across +searches -- see ``_WIGNER_D_CACHE`` for what that costs in memory. +""" +from __future__ import annotations + +import torch + +from ....config import canonical_device + +__all__ = ["clear_wigner_d_cache", "wigner_contraction_per_beta"] + +#: Memo for the per-l J_y eigendecomposition, keyed on L alone. +#: It depends only on the bandwidth, so repeat calls at the same L reuse it. Held +#: on the HOST in float64/complex128: it is an eigendecomposition, the most +#: precision-sensitive step here, and it is small (L-1 matrices, the largest +#: 129x129) and data-independent. Keeping it off the accelerator costs nothing +#: measurable and means no float64 is required there. +_WIGNER_EIG_CACHE: dict = {} + +#: Memo for the per-l small-d blocks, keyed on (L, betas, device-str). Holds at +#: most one entry, because the blocks are large: the per-l blocks together are +#: ``n_beta * sum_l (2l+1)^2`` float64 scalars, 176 MB at L=65 and 659 MB at +#: L=101 for the 60-value beta grid. One entry is all production wants -- +#: ``LMAX_CAP`` and ``GRID_SAMPLING_DEG`` in ``rotation_search`` are constants, +#: so every call arrives with the same key. A caller that alternates bandwidths +#: rebuilds each time, which is the uncached cost and not worse. +_WIGNER_D_CACHE: dict = {} + + +def _wigner_eig_table(L: int): + """Return [(w_l, V_l)] for l ∈ [1, L) — the J_y eigendecomposition per l. + + ``d^l(β) = Re(V_l · diag(e^{-iβ w_l}) · V_l^H)``. ``w_l ≈ [-l..l]`` and + ``V_l`` are independent of β and the data, so they are memoised. Built and + kept on the host; see ``_WIGNER_EIG_CACHE``. + """ + key = int(L) + cached = _WIGNER_EIG_CACHE.get(key) + if cached is not None: + return cached + table = [] + for l in range(1, L): + sz = 2 * l + 1 + p = torch.arange(sz - 1, dtype=torch.float64) # dtype-ok: eigendecomposition in double for the Wigner-d recursion + sup = 0.5 * torch.sqrt((2 * l - p) * (p + 1.0)) + A = torch.diag(sup, 1) - torch.diag(sup, -1) # A = -i J_y + w, V = torch.linalg.eigh(1j * A.to(torch.complex128)) # w∈[-l..l] # dtype-ok: eigendecomposition in double for the Wigner-d recursion + table.append((w, V)) + _WIGNER_EIG_CACHE[key] = table + return table + + +def clear_wigner_d_cache() -> None: + """Drop the memoised small-d blocks, releasing their memory.""" + _WIGNER_D_CACHE.clear() + + +def _wigner_d_blocks(L: int, betas: torch.Tensor, device: torch.device, + dtype: torch.dtype): + """Per-l real ``d^l(β)`` blocks for ``l ∈ [1, L)``, memoised. + + Each entry is ``(n_beta, 2l+1, 2l+1)`` in ``dtype``, on ``device``. They + depend only on the bandwidth and the β grid, not on the data, so a process + that runs more than one rotation search at the same bandwidth builds them + once. A single search asks for them exactly once and so pays the full build. + + Built on the host in float64 from the cached eigendecomposition and moved + once: the ``d^l`` entries are bounded in [-1, 1], so storing them at the + working precision loses nothing structural, and it halves the memo. + + The build is the dominant cost of :func:`wigner_contraction_per_beta`: it is + a batched ``(n_beta, sz, sz) @ (sz, sz)`` product per l, against the + contraction's elementwise ``(n_beta, sz, sz)``. See ``_WIGNER_D_CACHE`` for + the footprint that buys. + """ + # `canonical_device` fills in the default index: torch.device('cuda') and + # torch.device('cuda:0') name one physical device but stringify differently, + # and this memo holds exactly ONE entry -- so a mixed spelling would not add + # an entry, it would clear and rebuild the whole table on every call. + key = ( + int(L), + str(canonical_device(device)), + dtype, + tuple(betas.detach().cpu().to(torch.float64).tolist()), # dtype-ok: memo key: exact host-side doubles + ) + hit = _WIGNER_D_CACHE.get(key) + if hit is not None: + return hit + + eig_table = _wigner_eig_table(L) # host, cached + # `.cpu()` before the widening, not after: float64 cannot be materialised + # on every backend, and widening on the host is exact either way. + betas_host = betas.detach().cpu().to(torch.float64) # dtype-ok: memo key: exact host-side doubles + blocks = [] + for l in range(1, L): + w, V = eig_table[l - 1] # data-independent + phase = torch.exp(-1j * betas_host.unsqueeze(1) * w.unsqueeze(0)) # (n_beta, sz) + VP = V.unsqueeze(0) * phase.unsqueeze(1) # (n_beta, sz, sz) = (k,m,a) + blocks.append( + # `resolve_conj()`: a batched matmul on a lazy conjugate view drops + # the conjugation on MPS (see `cross_correlate_xi`). This block is + # built on the host today, where the bit is honoured, but the failure + # mode is silent and the copy is one small matrix. + (VP @ V.conj().resolve_conj().transpose(-1, -2)) + .real.to(device=device, dtype=dtype) + ) + + _WIGNER_D_CACHE.clear() # one entry only; see the footprint note + _WIGNER_D_CACHE[key] = blocks + return blocks + + +def wigner_contraction_per_beta( + xi_lmn: torch.Tensor, + betas: torch.Tensor, +) -> torch.Tensor: + """Compute ``S_{m1, m2}(β) = Σ_l ξ_{l, m1, m2} · d^l_{m1, m2}(β)``. + + Phaser source: ``SiteListAng::DoRfftStuff`` (FastRot.cc:39-59) — the + inner ``for (l_index, m1_index, m2_index)`` triple-loop. We do it + in one tensor contraction instead of a Python loop over l. + + Parameters + ---------- + xi_lmn : torch.Tensor (complex), shape (L, 2L-1, 2L-1) + SH-Bessel coefficients with l ∈ [0, L), |m|, |n| ≤ L-1, + zero-padded outside |m| > l or |n| > l. (Already n-summed over + the Bessel radial index by the caller.) + betas : torch.Tensor (real), shape (n_beta,) + β values to evaluate at, in radians. + + Returns + ------- + S : torch.Tensor (complex), shape (n_beta, 2L-1, 2L-1) + ``S[k, m1+L-1, m2+L-1] = Σ_l ξ_{l, m1, m2} · d^l_{m1, m2}(β_k)``. + """ + if xi_lmn.ndim != 3: + raise ValueError(f"xi_lmn must be 3-D, got shape {tuple(xi_lmn.shape)}") + L = xi_lmn.shape[0] + dim = 2 * L - 1 + device = xi_lmn.device + n_beta = betas.shape[0] + # Follow the input rather than forcing double: `xi` carries the expansion's + # working precision, so widening here would buy nothing and cost a 2x + # complex buffer in this stage and in the FFT it feeds. + xi = xi_lmn + real_dtype = torch.float64 if xi.dtype == torch.complex128 else torch.float32 # dtype-ok: follows the accumulator's width + + # Per-l loop over the small-d blocks, which come from the J_y + # eigendecomposition (small_d_stable's method, stable to any l). Contract + # each into S in turn: the full (n_beta, L, 2L-1, 2L-1) table is never + # materialised as one array, nor is a 4-D einsum intermediate. + S = torch.zeros((n_beta, dim, dim), dtype=xi.dtype, device=device) + c = L - 1 + S[:, c, c] += xi[0, c, c] # l=0: d^0 = 1 + blocks = _wigner_d_blocks(L, betas, device, real_dtype) + for l in range(1, L): + d_l = blocks[l - 1] # (n_beta, sz, sz) + lo, hi = c - l, c + l + 1 + S[:, lo:hi, lo:hi] += xi[l, lo:hi, lo:hi].unsqueeze(0) * d_l + return S diff --git a/torchref/experimental/alignment/jax_subpixel_peaks.py b/torchref/experimental/alignment/jax_subpixel_peaks.py deleted file mode 100644 index 9ee1fff9..00000000 --- a/torchref/experimental/alignment/jax_subpixel_peaks.py +++ /dev/null @@ -1,82 +0,0 @@ -""" -JAX sub-voxel peak refinement for the experimental ball-harmonic MR engine. - -Experimental / unstable API: part of ``torchref.experimental.alignment``, -the opt-in ball-harmonic MR engine. The production MR entry point is -``torchref.alignment`` (the consolidated FRF engine). Signatures and behavior -may change without notice. -""" - -import jax -import jax.numpy as jnp - - -@jax.jit -def refine_peaks_subvoxel( - grid: jax.Array, - peak_indices: jax.Array, -) -> tuple[jax.Array, jax.Array]: - """ - Refine discrete peaks to sub-voxel accuracy using separable quadratic fitting. - - Parameters - ---------- - grid : jax.Array - Array of values, shape (n0, n1, n2). - peak_indices : jax.Array - Integer indices of discrete maxima, shape (n_peaks, 3). - - Returns - ------- - refined_coords : jax.Array - Fractional coordinates, shape (n_peaks, 3). - refined_values : jax.Array - Interpolated peak heights, shape (n_peaks,). - """ - shape = jnp.array(grid.shape) - - def refine_single_peak(idx): - i0, i1, i2 = idx[0], idx[1], idx[2] - - def get_val(d0, d1, d2): - # Periodic boundary conditions - ii0 = (i0 + d0) % shape[0] - ii1 = (i1 + d1) % shape[1] - ii2 = (i2 + d2) % shape[2] - return grid[ii0, ii1, ii2] - - f0 = get_val(0, 0, 0) - - # Axis 0 - fm = get_val(-1, 0, 0) - fp = get_val(1, 0, 0) - c0 = (fp - 2*f0 + fm) / 2 - b0 = (fp - fm) / 2 - delta0 = jnp.where(jnp.abs(c0) > 1e-12, jnp.clip(-b0 / (2*c0), -1.0, 1.0), 0.0) - - # Axis 1 - fm = get_val(0, -1, 0) - fp = get_val(0, 1, 0) - c1 = (fp - 2*f0 + fm) / 2 - b1 = (fp - fm) / 2 - delta1 = jnp.where(jnp.abs(c1) > 1e-12, jnp.clip(-b1 / (2*c1), -1.0, 1.0), 0.0) - - # Axis 2 - fm = get_val(0, 0, -1) - fp = get_val(0, 0, 1) - c2 = (fp - 2*f0 + fm) / 2 - b2 = (fp - fm) / 2 - delta2 = jnp.where(jnp.abs(c2) > 1e-12, jnp.clip(-b2 / (2*c2), -1.0, 1.0), 0.0) - - # Refined coordinates - refined = jnp.array([i0 + delta0, i1 + delta1, i2 + delta2]) - - # Refined value: f(δ) ≈ f0 - b²/(4c) at peak - peak_val = f0 - peak_val -= jnp.where(jnp.abs(c0) > 1e-12, b0**2 / (4*c0), 0.0) - peak_val -= jnp.where(jnp.abs(c1) > 1e-12, b1**2 / (4*c1), 0.0) - peak_val -= jnp.where(jnp.abs(c2) > 1e-12, b2**2 / (4*c2), 0.0) - - return refined, peak_val - - return jax.vmap(refine_single_peak)(peak_indices) \ No newline at end of file diff --git a/torchref/experimental/alignment/pipeline.py b/torchref/experimental/alignment/pipeline.py index 664561f5..256df019 100644 --- a/torchref/experimental/alignment/pipeline.py +++ b/torchref/experimental/alignment/pipeline.py @@ -1,619 +1,744 @@ +"""Molecular replacement: the FRF hands a shortlist to the FTF. + +Two stages, and the division of labour between them is the design: + +1. **Fast Rotation Function** — Phaser-faithful Bessel-radial × SH expansion + against a dense P1-box calc. It is a *shortlist generator*. It does not have + to rank well, and it does not: over the panel its own ordering puts truth + first in a minority of cells. What it does reliably is put the true + orientation somewhere in the top twenty-five -- in every cell of every + panel run on record -- and its peaks are one per orientation, symmetry + mates suppressed. +2. **Fast Translation Function** — for *each* of the top-N orientations, one + Crowther-Blow FFT over the fractional cell on a resolution-sized grid, with + the rotation function's own normalised score equation as its coefficients, + then the Rice/Woolfson likelihood at the best few peaks. Rotation ghosts + are morphologically identical to truth in a rotation function by + construction; they are not identical once the crystal is involved. + +There is deliberately nothing between them. An ML re-ranking of the FRF peaks +used to sit there and was removed: it reorders a shortlist that already contains +truth, and rotation recovery was 18/30 with it against 24/30 without (McNemar +p = 0.031, 6-0 discordant). Those figures, like every figure on this pipeline +before September 2026, gated on the rotation alone; see ``rank_by`` for what +the pose-gated panel measures. + +Every candidate is placed and then the best is taken, with no early stopping. +Ten candidates by default: the rotation function's first distinct peak was +the true orientation in every pose-gated cell measured, so ten is a margin. +Stopping early made the pipeline's answer depend on the order the rotation +function happened to produce -- it walked the list until one placement beat an +R-factor threshold and returned that, so it could accept the third candidate +without ever scoring the tenth. On a structure where several orientations place +plausibly that is not a choice between them, and it made the selection rule +impossible to reason about or to measure against a ranking harness. + +The candidates are ranked by the translation search's likelihood -- see +``rank_by``, and the sort in :meth:`MolecularReplacementPipeline.run` for what +the three available scores measured against each other. The user-facing +solvent-aware R-work is computed once, on the winner. The pipeline returns a +*placement* -- refining it is the caller's job, and downstream refinement does +it better than a bolted-on polish did. + +``align_model_to_data`` delegates here; this class is the implementation of +record. The crystallographic stages live in +:mod:`torchref.experimental.alignment.align` and +:mod:`~torchref.experimental.alignment.translation`; this module owns the +control flow that wires them together. """ -Molecular replacement pipeline integrating rotation, translation, and refinement. -This module provides a unified pipeline for molecular replacement that chains: -1. Fast rotation function (ball transform) -2. FFT-based translation search -3. Clash filtering -4. Rigid body refinement - -The pipeline supports early stopping when a good solution is found. - -Experimental / unstable API: this is the opt-in ball-harmonic MR engine in -``torchref.experimental.alignment``. The production / canonical MR entry point -is ``torchref.alignment`` (the consolidated FRF engine). Signatures and -behavior may change without notice. -""" +from __future__ import annotations +import time +from contextlib import contextmanager from dataclasses import dataclass -from typing import List, Optional, Tuple, TYPE_CHECKING +from typing import List, Optional, TYPE_CHECKING + import numpy as np import torch from torchref.config import get_default_device, get_float_dtype - -from .ball_transform import ( - ball_rotation_search_torch, - rotation_matrix_from_euler_zyz, +from torchref.utils.device_mixin import DeviceMixin + +from .frf.rotation_utils import rotation_matrix_from_edmonds_euler +from .frf.types import RotationPeak +from .rotation_search import prepare_frf_inputs, search_peaks +from .translation import ( + TranslationObs, + analytic_r_at, + fast_translation_function, + llg_at_translations, + prepare_candidate, ) -from .translation import fft_translation_search_torch, TranslationPeak -from .rigid_body import RigidBodyRefinement, RigidBodyResult -from .clashscore import ClashScoreCalculator, AtomSampler if TYPE_CHECKING: - from torchref.model import ModelFT from torchref.io.datasets import ReflectionData + from torchref.model import ModelFT -def rotation_angular_distance(R1: np.ndarray, R2: np.ndarray) -> float: - """ - Compute angular distance between two rotation matrices in degrees. - - The angular distance is the angle of the rotation R2 @ R1.T. - - Parameters - ---------- - R1, R2 : np.ndarray - 3x3 rotation matrices. - - Returns - ------- - float - Angular distance in degrees. - """ - R_diff = R2 @ R1.T - # Clamp trace to valid range for arccos - trace = np.trace(R_diff) - trace = np.clip(trace, -1.0, 3.0) - angle_rad = np.arccos((trace - 1.0) / 2.0) - return np.degrees(angle_rad) - - -def euler_angular_distance( - euler1: Tuple[float, float, float], - euler2: Tuple[float, float, float], -) -> float: - """ - Compute angular distance between two ZYZ Euler angle sets. +# --------------------------------------------------------------------------- +# Stage timing +# --------------------------------------------------------------------------- - Parameters - ---------- - euler1, euler2 : tuple - (alpha, beta, gamma) Euler angles in radians. - Returns - ------- - float - Angular distance in degrees. - """ - R1 = rotation_matrix_from_euler_zyz(euler1[0], euler1[1], euler1[2]) - R2 = rotation_matrix_from_euler_zyz(euler2[0], euler2[1], euler2[2]) - return rotation_angular_distance(R1, R2) +class _StageTimer: + """Lightweight wall-clock accumulator. Gated by ``verbose >= 2``. + Two interleavable usages: + * ``with t.stage(name):`` block — records the block's wall time. + * ``t.start(name)`` / ``t.stop(name)`` — checkpoint pair, no indent. -def cluster_rotation_peaks( - peaks: list, - threshold_deg: float = 6.0, - symmetry_matrices: Optional[np.ndarray] = None, -) -> list: + The summary table prints stages aggregated by name; per-rotation loop + stages (translation search etc.) get aggregated counts. """ - Cluster rotation peaks by angular distance. - - Peaks within threshold_deg of each other are considered the same solution. - Only the highest-scoring peak from each cluster is kept. - Parameters - ---------- - peaks : list - List of rotation peaks as tuples (alpha, beta, gamma, score, sigma). - threshold_deg : float - Angular distance threshold for clustering in degrees. - symmetry_matrices : np.ndarray, optional - Point group symmetry matrices (N, 3, 3) to check symmetry equivalents. - - Returns - ------- - list - Clustered peaks with one representative per cluster. - """ - if not peaks: - return [] - - # Sort by sigma (descending) so we keep highest-scoring peaks - sorted_peaks = sorted(peaks, key=lambda p: p[4], reverse=True) - - clustered = [] - used_rotations = [] # Store rotation matrices of accepted peaks - - for peak in sorted_peaks: - alpha, beta, gamma, score, sigma = peak - R = rotation_matrix_from_euler_zyz(alpha, beta, gamma) - - # Check if this rotation is too close to any already accepted rotation - is_new = True - for R_used in used_rotations: - dist = rotation_angular_distance(R, R_used) - if dist < threshold_deg: - is_new = False - break - - # Also check symmetry equivalents if provided - if symmetry_matrices is not None and is_new: - for sym_op in symmetry_matrices: - R_sym = sym_op @ R - dist_sym = rotation_angular_distance(R_sym, R_used) - if dist_sym < threshold_deg: - is_new = False - break - if not is_new: - break - - if is_new: - clustered.append(peak) - used_rotations.append(R) - - return clustered + def __init__(self, enabled: bool): + self.enabled = enabled + self.records: list[tuple[str, float]] = [] + self._open: dict[str, float] = {} + + @contextmanager + def stage(self, name: str): + if not self.enabled: + yield + return + t0 = time.perf_counter() + try: + yield + finally: + self.records.append((name, time.perf_counter() - t0)) + + def start(self, name: str) -> None: + if self.enabled: + self._open[name] = time.perf_counter() + + def stop(self, name: str) -> None: + if not self.enabled: + return + t0 = self._open.pop(name, None) + if t0 is not None: + self.records.append((name, time.perf_counter() - t0)) + + def summary(self) -> str: + if not self.records: + return "" + # Aggregate repeated stage names (the per-rotation loop visits the + # translation stages once per candidate rotation). + agg: dict[str, list[float]] = {} + for name, dt in self.records: + agg.setdefault(name, []).append(dt) + total = sum(sum(v) for v in agg.values()) + lines = [ + f"{'stage':<32s} {'count':>5s} {'wall_s':>10s} {'%':>6s}", + "-" * 60, + ] + for name, vs in agg.items(): + wall = sum(vs) + lines.append( + f"{name:<32s} {len(vs):>5d} {wall:>10.3f} " + f"{100 * wall / total:>5.1f}%" + ) + lines.append("-" * 60) + lines.append(f"{'TOTAL':<32s} {'':>5s} {total:>10.3f} 100.0%") + return "\n".join(lines) @dataclass class MRSolution: - """ - Molecular replacement solution. + """A molecular-replacement placement. Attributes ---------- rotation : np.ndarray - ZYZ Euler angles (radians), shape (3,). - translation : np.ndarray - Fractional coordinates, shape (3,). + Recovered orientation as a 3×3 rotation matrix (``R_recovered`` — the + rotation that maps the *search-model* frame onto the *crystal* frame). + translation : np.ndarray or None + Fractional translation applied after rotation, shape (3,). ``None`` for + a rotation-only solution (``do_translation=False``). rotation_score : float - FRF sigma (Z-score). + The rotation function's score for this candidate. translation_score : float - Translation function correlation. - clash_score : float - Steric clash score (lower is better). + The fast translation function's score at the chosen peak, higher + better. Reported, not ranked -- see the sort in :meth:`run`. r_factor : float - R-factor after refinement. - refined_rotation : np.ndarray, optional - Refined Euler angles (radians). - refined_translation : np.ndarray, optional - Refined fractional translation. + The analytical-scale R at that placement, lower better. Reported, not + ranked. A single global scale, so it is not the number a full Scaler + would return -- build one on the returned model if that is wanted. + llg_score : float + **The ranking key**: the translation likelihood at that placement, + higher better. + candidate_index : int + Position of this orientation in the rotation function's own ordering, + so the depth of shortlist a solution came from can be read off. + model : ModelFT or None + The rotated and translated model. Built for the winner only -- copying + and moving a 20k-atom model 25 times was a quarter of the run on the + large structures, to produce 24 models nobody reads. + :meth:`MolecularReplacementPipeline.place` builds it for any other + solution on request. """ + rotation: np.ndarray - translation: np.ndarray + translation: Optional[np.ndarray] rotation_score: float translation_score: float - clash_score: float r_factor: float - refined_rotation: Optional[np.ndarray] = None - refined_translation: Optional[np.ndarray] = None - + model: Optional["ModelFT"] = None + llg_score: float = float("nan") + candidate_index: int = -1 -class MolecularReplacementPipeline: - """ - Unified MR pipeline: Rotation -> Translation -> Rigid Body Refinement. - This pipeline integrates the fast rotation function (ball transform), - FFT-based translation search, clash filtering, and rigid body refinement - into a single workflow with early stopping. +class MolecularReplacementPipeline(DeviceMixin): + """Canonical MR pipeline: FRF → FTF (per candidate) → post-refine. - Experimental: this is the opt-in ball-harmonic MR engine. The production - MR entry point is ``torchref.alignment`` (the consolidated FRF engine). - APIs here may change without notice. + Parameters mirror :func:`align_model_to_data` (which delegates here), so a + caller can either use ``align_model_to_data`` for the common case or drive this + class directly for finer control / access to the ranked candidate list. Parameters ---------- data : ReflectionData Observed reflection data. model : ModelFT - Search model with atomic coordinates. + Initialised search model. device : torch.device, optional - Computation device. Default is CPU. + Compute device (defaults to the model's device). verbose : int - Verbosity level (0=silent, 1=summary, 2=detailed). + How much the run says about itself. Each level is a superset of the one + below, and the boundaries are chosen so that a level is useful on its + own rather than being "a bit more of the same": + + 0 + Silent. + 1 + What happened: the search settings, one line per stage, and the + winner. Enough to see that a run did the expected work. + 2 + **Why it chose what it chose.** One ``CAND`` line per rotation + candidate carrying every score the selection could have used, plus + the per-stage wall-clock table. This is the level that makes the + pipeline diagnosable without a second implementation of its own + scoring -- see :meth:`_log_candidate`. + 3 + Per-translation-peak detail inside each candidate. Examples -------- :: from torchref.experimental.alignment import MolecularReplacementPipeline - from torchref.model import ModelFT - from torchref.io.datasets import ReflectionData - - data = ReflectionData().load_mtz('observed.mtz') - model = ModelFT().load_pdb('search_model.pdb') - pipeline = MolecularReplacementPipeline(data, model) - solutions = pipeline.run(n_rotation_peaks=50, min_tries=3, max_tries=10) - print(f'Best R: {solutions[0].r_factor:.3f}') + + pipe = MolecularReplacementPipeline(data, model) + solutions = pipe.run() + print(f"best R-work: {solutions[0].r_factor:.3f}") """ def __init__( self, data: "ReflectionData", model: "ModelFT", + *, device: Optional[torch.device] = None, - verbose: int = 1, + verbose: int = 0, + # --- data prep / FRF --- + d_min: float = 4.0, + d_max: float = 15.0, + n_shells: int = 20, + n_rotation_peaks: int = 500, + model_error_A: Optional[float] = None, + # --- candidate tree --- + # Distinct orientations carried into the translation search. A safety + # margin, not a requirement: with symmetry mates suppressed the + # rotation function's FIRST peak is the true orientation in 50 of 50 + # pose-gated cells (10 structures x 5 seeds), and the panel is 30/30 at + # 10 as at 25. Measured on the deposited models as search models; raise + # it for poorer models, since every candidate costs a structure-factor + # evaluation and a translation FFT. + n_rotation_candidates: int = 10, + # Peaks of the fast translation function re-scored by the likelihood + # for each orientation. The fast map only has to get the true peak + # into this many; the likelihood picks. + n_translation_candidates: int = 3, + # Which score picks the winner among placed candidates. "llg" is the + # translation likelihood; "r" the analytical-scale R-factor; "corr" the + # fast translation function's own score. Not a tuning knob -- it exists + # because the three can disagree and a rank-level proxy once got the + # ordering wrong, so the comparison has to be made end to end on poses. + # See the sort in `run` for what that measured. + rank_by: str = "llg", + # Resolution window for the translation set. None means the rotation + # search's own [d_max, d_min], so one window and one normalisation + # serve both stages. Pass 0.0 / inf to remove a cut -- and see + # `_prepare_translation_arrays` for what the uncut set does. + tf_d_min: Optional[float] = None, + tf_d_max: Optional[float] = None, ): self.data = data self.model = model self.device = device or get_default_device() self.verbose = verbose - # Lazy caches - self._clash_calc = None - self._e_obs = None - self._e_calc = None - self._s_vectors = None - self._mask = None - - def run( - self, - n_rotation_peaks: int = 100, - n_translation_peaks: int = 5, - min_tries: int = 3, - max_tries: int = 10, - rfactor_converged: float = 0.45, - max_clash_score: float = 100.0, - d_min: float = 4.0, - d_max: float = 50.0, - L: int = 32, - P: int = 20, - cluster_threshold_deg: float = 6.0, - ) -> List[MRSolution]: + self.d_min = d_min + self.d_max = d_max + self.n_shells = n_shells + self.n_rotation_peaks = n_rotation_peaks + # Expected r.m.s. coordinate error of the search model, in Angstrom: + # it sets the sigma_A fall-off in the rotation function. When the caller + # does not know it, estimate it from the model's length the way Phaser + # does (Oeffner et al. 2013), assuming the sequence is the target's -- + # roughly 8 heavy atoms per residue. + if model_error_A is None: + from .frf.preprocessing import oeffner_vrms + n_residues = max(1, int(model.xyz().shape[0] / 8)) + model_error_A = oeffner_vrms(n_residues, 1.0) + self.model_error_A = float(model_error_A) + + self.n_rotation_candidates = n_rotation_candidates + self.n_translation_candidates = n_translation_candidates + if rank_by not in ("r", "corr", "llg"): + raise ValueError( + f"rank_by={rank_by!r}; expected 'r', 'corr' or 'llg'.") + self.rank_by = rank_by + self.tf_d_min = float(d_min if tf_d_min is None else tf_d_min) + self.tf_d_max = float(d_max if tf_d_max is None else tf_d_max) + + self._timer = _StageTimer(enabled=verbose >= 2) + # Filled in by run(). + self._frf = None + self._obs = None + self._tmask = None + # One P1 copy of the search model, re-oriented in place per candidate + # -- see `_prepare_translation_arrays`. + self._p1 = None + self._p1_xyz0 = None + self._p1_center = None + + #: Levels are documented on the class. They are a contract, not a dial: + #: level 2 is specifically "one machine-readable line per candidate", and + #: anything added at that level should preserve that. + def _log(self, level: int, msg: str) -> None: + """Emit ``msg`` if the run is at least this verbose. + + One emitter rather than ``if self.verbose > 0: print(...)`` at every + site. The scattered form is how levels drift -- the same stage ends up + reporting at 1 in one place and 2 in another, and nothing enforces that + a level means the same thing twice. """ - Run full MR pipeline with early stopping. + if self.verbose >= level: + print(msg, flush=True) + + def _log_candidate(self, k: int, peak, r_analytic, t_frac, + tf_score=None, llg_score=None) -> None: + """One line per rotation candidate, with every score behind the choice. + + Machine-readable on purpose. Diagnosing a wrong placement means asking + which candidate won and on what, and the only alternative to emitting it + here is a harness that re-implements the placement loop -- which drifts + from the pipeline and then disagrees with it about which candidate the + pipeline picked. A caller that knows the true orientation (a benchmark) + can join these lines against it; the pipeline cannot, and does not try. + + Fields are ``key=value`` so a reader does not depend on column order: + ``k`` candidate index in rotation-function order, ``rf``/``rfz`` its + score and z, ``tf`` the fast translation function's score, ``llg`` the + translation likelihood, ``r`` the analytical-scale R, ``t`` the + fractional translation. All three placement scores are reported whichever one + ranks, because which of them a wrong placement disagreed on is the + question, and they do disagree. + """ + tf = "nan" if tf_score is None else f"{float(tf_score):.5f}" + llg = "nan" if llg_score is None else f"{float(llg_score):.1f}" + t = ",".join(f"{float(x):.4f}" for x in t_frac) + self._log(2, f"CAND k={k} rf={float(peak.score):.4f} " + f"rfz={float(peak.sigma):.3f} tf={tf} llg={llg} " + f"r={float(r_analytic):.5f} t={t}") + + + # ------------------------------------------------------------------ + # Public entry point + # ------------------------------------------------------------------ + def run(self, do_translation: bool = True) -> List[MRSolution]: + """Run the MR pipeline and return solutions, best first. + + Ranked by ``rank_by`` -- the translation likelihood by default. Parameters ---------- - n_rotation_peaks : int - Number of rotation peaks to try. - n_translation_peaks : int - Number of translation peaks per rotation. - min_tries : int - Floor on the number of refinements performed before early stopping - (convergence-based break) is allowed. It does NOT force additional - refinements beyond the available candidates: candidates are capped - at ``min(len(candidates), max_tries)``, so ``min_tries`` is not a - guaranteed minimum when fewer candidates exist. - max_tries : int - Maximum number of candidates to refine. - rfactor_converged : float - R-factor threshold for convergence (early stopping). - max_clash_score : float - Maximum clash score to accept a candidate. - d_min : float - High resolution limit for rotation search (Angstroms). - d_max : float - Low resolution limit for rotation search (Angstroms). - L : int - Angular bandlimit for rotation search. - P : int - Radial bandlimit for rotation search. - cluster_threshold_deg : float - Angular distance threshold for clustering rotation peaks (degrees). - Peaks within this angular distance are considered the same solution. - Default 6.0 degrees matches rigid body refinement convergence radius. + do_translation : bool + If ``False``, stop after rotation rescoring and return a single + rotation-only solution (the model rotated onto the best + orientation, no translation or refinement). Returns ------- - List[MRSolution] - Solutions sorted by R-factor. Early stops if converged. + list of MRSolution + Sorted by ``r_factor`` (ascending). The first element is the best + placement; its ``r_factor`` is the solvent-aware Scaler R-work. """ - # Step 1: Rotation search - if self.verbose: - print("Step 1: Rotation search...") - rotation_peaks = self._rotation_search(n_rotation_peaks, d_min, d_max, L, P) - - if not rotation_peaks: - if self.verbose: - print(" No rotation peaks found!") - return [] - - # Step 1b: Cluster rotation peaks - if self.verbose: - print(f" Found {len(rotation_peaks)} raw peaks, clustering with {cluster_threshold_deg}° threshold...") - - # Get symmetry matrices for clustering - sym_matrices = None - if hasattr(self.model, 'symmetry') and self.model.symmetry is not None: - sym_matrices = np.array([s.numpy() for s in self.model.symmetry.matrices]) - - rotation_peaks = cluster_rotation_peaks( - rotation_peaks, - threshold_deg=cluster_threshold_deg, - symmetry_matrices=sym_matrices, - ) + if not self.model.ctx.initialized: + raise RuntimeError( + "Cannot fit an uninitialized ModelFT. Load PDB data first." + ) - if self.verbose: - print(f" {len(rotation_peaks)} unique rotation clusters") - - if not rotation_peaks: - if self.verbose: - print(" No rotation peaks after clustering!") - return [] - - # Step 2: Translation search for each rotation - if self.verbose: - print(f"Step 2: Translation search for {len(rotation_peaks)} rotations...") - candidates = [] - for i, rot in enumerate(rotation_peaks): - trans_peaks = self._translation_search(rot, n_translation_peaks) - for trans in trans_peaks: - candidates.append((rot, trans)) - if self.verbose > 1 and (i + 1) % 10 == 0: - print(f" Processed {i+1}/{len(rotation_peaks)} rotations...") - - if self.verbose: - print(f" Generated {len(candidates)} rotation+translation candidates") - - # Step 3: Score and filter by clash - if self.verbose: - print("Step 3: Clash filtering...") - candidates = self._score_and_filter(candidates, max_clash_score) + timer = self._timer + timer.start("0_data_prep") + frf = prepare_frf_inputs( + self.model, self.data, + d_min=self.d_min, d_max=self.d_max, n_shells=self.n_shells, + verbose=self.verbose, + ) + timer.stop("0_data_prep") + self._frf = frf + # --- Stage 1: FRF rotation search --- + candidates = self._rotation_candidates(frf) if not candidates: - if self.verbose: - print(" All candidates rejected by clash filter!") - return [] - - if self.verbose: - print(f" {len(candidates)} candidates passed clash filter") - - # Step 4: Rigid body refinement with early stopping - if self.verbose: - print(f"Step 4: Refining candidates (min={min_tries}, max={max_tries})...") - - solutions = [] - best_r_factor = float('inf') - converged = False - - n_to_refine = min(len(candidates), max_tries) - for i, (rot, trans, clash) in enumerate(candidates[:n_to_refine]): - if self.verbose > 1: - print(f" Refining candidate {i+1}/{n_to_refine}...") - - try: - result = self._rigid_body_refine(rot, trans) - - solution = MRSolution( - rotation=np.array([rot[0], rot[1], rot[2]]), - translation=trans.translation, - rotation_score=rot[4], # sigma - translation_score=trans.score, - clash_score=clash, - r_factor=result.final_r_factor, - refined_rotation=np.array(result.final_rotation), - refined_translation=np.array(result.final_translation_frac), + raise RuntimeError("Rotation search produced no peaks.") + + if not do_translation: + rotated, R_rec = self._make_rotated(candidates[0]) + top = candidates[0] + self._log(1, f"mr: top peak RF = {top.score:.2f} " + f"(σ_Z = {top.sigma:.2f}); applying R⁻¹ to coords.") + self._log(2, "\n" + timer.summary()) + return [ + MRSolution( + rotation=R_rec.detach().cpu().numpy(), + translation=None, + rotation_score=float(top.score), + translation_score=float("nan"), + r_factor=float("nan"), + model=rotated, ) - solutions.append(solution) - - # Track best - if result.final_r_factor < best_r_factor: - best_r_factor = result.final_r_factor - - # Check early stopping - n_refined = i + 1 - if n_refined >= min_tries and best_r_factor < rfactor_converged: - converged = True - if self.verbose: - print(f" Converged! R-factor {best_r_factor:.4f} < {rfactor_converged}") - break - - except Exception as e: - if self.verbose > 1: - print(f" Refinement failed: {e}") + ] + + # --- Stage 2: per-candidate translation search --- + self._prepare_translation_arrays() + n_rot = min(self.n_rotation_candidates, len(candidates)) + if n_rot > 1: + self._log(1, f"mr: placing all {n_rot} rotation candidates…") + + solutions: List[MRSolution] = [] + for k in range(n_rot): + peak_k = candidates[k] + R_rec_k = rotation_matrix_from_edmonds_euler( + peak_k.alpha, peak_k.beta, peak_k.gamma) + self._orient_template(R_rec_k) + self._log(3, f"\nfit_to_data: rot{k} " + f"(RF={peak_k.score:.2f}, σ_Z={peak_k.sigma:.2f})") + placement = self._placement_for_candidate() + if placement is None: + self._log(2, f"CAND k={k} rf={float(peak_k.score):.4f} " + f"rfz={float(peak_k.sigma):.3f} tf=nan r=nan " + f"t=none # no translation peaks") continue + r_analytic, t_refined, tf_score, llg_score = placement + self._log_candidate(k, peak_k, r_analytic, t_refined, tf_score, + llg_score) + solutions.append( + MRSolution( + rotation=R_rec_k.detach().cpu().numpy(), + translation=t_refined.detach().cpu().numpy(), + rotation_score=float(peak_k.score), + translation_score=float(tf_score), + r_factor=float(r_analytic), + llg_score=float(llg_score), + candidate_index=k, + ) + ) - if self.verbose: - status = "converged" if converged else f"completed {len(solutions)}/{n_to_refine}" - print(f" Refinement {status}") - if solutions: - print(f" Best R-factor: {best_r_factor:.4f}") - - solutions.sort(key=lambda s: s.r_factor) + if not solutions: + raise RuntimeError("Translation + joint refine produced no candidates.") + + # Highest translation likelihood. The three scores are measured end to + # end on POSES -- rotation and translation, against Cartesian symmetry + # mates -- over six structures x ten seeds (the four the translation + # search used to mis-place, plus two controls; job 544953): + # + # llg 60/60 + # r 60/60 + # corr 60/60 + # + # and they do not merely tie: in every one of the 60 cells the three + # pick the SAME candidate, so the residual distributions are identical + # arm for arm. Once the translation objective was normalised there was + # nothing left for the selection rule to decide on this panel. + # + # The likelihood stays the default because it is the right object for + # the question -- an R-factor on a partial model at this resolution has + # little to distinguish with, and the fast score is an expansion of the + # likelihood rather than the likelihood -- and because the arm is + # selectable if a structure ever separates them. Every earlier figure + # for these arms (37/40, 36/40, 32/40) gated on the rotation alone, with + # a metric that miscounted trigonal mates; none of them stands. + if self.rank_by == "r": + solutions.sort(key=lambda s: s.r_factor) + elif self.rank_by == "corr": + solutions.sort(key=lambda s: -s.translation_score) + else: + solutions.sort(key=lambda s: -s.llg_score) + winner = solutions[0] + + # No solvent-aware Scaler refit. It used to run here on the winner to + # report an R-work, and cost about a third of the whole alignment -- 8.6 + # of 28 seconds on 2DQ6 -- to fit sixteen scaling parameters that change + # nothing about which placement is returned. The pipeline's contract is + # a placement; downstream refinement fits its own scaler properly, and + # doing a worse version of that here to print a number is not worth a + # third of the runtime. A caller that wants an R-work can build a + # `Scaler` on the returned model. + winner.model = self.place(winner) + self._log(1, f"mr: winner ({self.rank_by}) " + f"LLG={winner.llg_score:.1f} " + f"TF={winner.translation_score:.5f} " + f"analytic R={winner.r_factor:.4f}") + self._log(2, "\n" + timer.summary()) return solutions - def _rotation_search( - self, - n_peaks: int, - d_min: float, - d_max: float, - L: int, - P: int, - ) -> list: - """Run ball rotation search.""" - # Prepare E-values - E_obs, s_obs = self._get_e_values_obs(d_min, d_max) - E_calc, s_calc = self._get_e_values_calc(d_min, d_max) - - _, _, peaks = ball_rotation_search_torch( - E_obs, s_obs, E_calc, s_calc, - L=L, P=P, d_min=d_min, d_max=d_max, n_peaks=n_peaks, - verbose=self.verbose > 1, + # ------------------------------------------------------------------ + # Stage 1: rotation search + # ------------------------------------------------------------------ + def _rotation_candidates(self, frf) -> list: + """FRF rotation search; the peaks it returns, ranked by its own score.""" + timer = self._timer + + timer.start("3_rotation_search") + self._log(1, f"mr: rotation search (n_peaks={self.n_rotation_peaks}, " + f"model error {self.model_error_A:.2f} A)…") + peaks, _lmax, _d_min = search_peaks( + self.model, self.data, self.model_error_A, + U_aniso=frf.U_aniso, n_peaks=self.n_rotation_peaks, + verbose=self.verbose, ) - return peaks - - def _translation_search( - self, - rotation_peak: tuple, - n_peaks: int, - ) -> List[TranslationPeak]: - """Run translation search for a rotation.""" - alpha, beta, gamma, _, _ = rotation_peak - - # Apply rotation to model coordinates - R = torch.tensor( - rotation_matrix_from_euler_zyz(alpha, beta, gamma), - dtype=get_float_dtype(), - device=self.device, + timer.stop("3_rotation_search") + + # Rank by the FRF's own score and hand the shortlist to the + # translation search. There is no rescore here by design: an ML + # re-ranking of these peaks was measured to lower end-to-end pose + # recovery from 24/30 to 18/30 (McNemar p = 0.031), because it reorders + # a shortlist that already contains truth and sometimes pushes truth + # out of it. The translation function does the discrimination. + return sorted(peaks, key=lambda p: p.score, reverse=True) + + def place(self, solution: MRSolution) -> "ModelFT": + """Build the placed model for ``solution``: a copy of the search model, + rotated and translated, carrying the alignment provenance attributes.""" + R_rec = torch.as_tensor(solution.rotation, dtype=torch.float64) # dtype-ok: 3x3 rotation algebra in double on the host + placed = self.model.copy().rotate( + R_rec.T.contiguous().to(device=self.model.device, + dtype=self.model.dtype_float), ) - xyz = self.model.xyz() - xyz_centered = xyz - xyz.mean(dim=0) - xyz_rotated = xyz_centered @ R.T - - # Temporarily update model coordinates and compute F_calc - original_xyz = self.model.xyz().clone() - self.model.xyz[:] = xyz_rotated - - try: - hkl = self.data.hkl - F_obs = self.data.F - mask = self.data.get_valid_mask() - - with torch.no_grad(): - F_calc = self.model(hkl) - - # Apply mask - F_obs_masked = F_obs[mask] - F_calc_masked = F_calc[mask] - hkl_masked = hkl[mask] - - _, _, peaks = fft_translation_search_torch( - F_obs_masked, F_calc_masked, hkl_masked, n_peaks=n_peaks - ) - finally: - # Restore original coordinates - self.model.xyz[:] = original_xyz + if str(placed.spacegroup) != str(self.data.spacegroup): + placed.spacegroup = self.data.spacegroup.hm + if solution.translation is not None: + t = torch.as_tensor(solution.translation, dtype=self.model.dtype_float) + placed.translate(t, fractional=True) + placed.last_alignment_translation = t + placed.last_alignment_rotation = R_rec + placed.last_alignment_rfactor = solution.r_factor + return placed + + def _orient_template(self, R_rec: torch.Tensor) -> None: + """Write the candidate orientation into the shared P1 copy. + + ``xyz = R_rec^T (xyz0 - c) + c`` about the search model's centroid, the + same rotation ``Model.rotate`` would apply. The forward cache + fingerprints parameters by pointer and version, so the next + structure-factor call recomputes. + """ + p1 = self._p1 + R_app = R_rec.T.to(device=self._p1_xyz0.device, dtype=self._p1_xyz0.dtype) + p1.xyz[:] = (self._p1_xyz0 - self._p1_center) @ R_app.T + self._p1_center - return peaks + def _make_rotated(self, peak: "RotationPeak"): + """Rotate the search model onto a candidate orientation. - def _score_and_filter( - self, - candidates: list, - max_clash: float, - ) -> list: - """Score candidates and filter by clash.""" - if self._clash_calc is None: - self._clash_calc = ClashScoreCalculator( - symmetry=self.data.spacegroup, - default_clash_radius=4.0, - device=self.device, + Returns ``(rotated_model, R_recovered)`` where ``R_recovered`` maps the + search-model frame onto the crystal frame; the applied coordinate + rotation is ``R_recovered.T``. + """ + R_rec = rotation_matrix_from_edmonds_euler(peak.alpha, peak.beta, peak.gamma) + R_app = R_rec.T.contiguous() + # .copy() first: Model.rotate mutates in place and returns self, so + # rotating self.model directly would compound candidate k+1 onto k. + rot = self.model.copy().rotate( + R_app.to(device=self.model.device, dtype=self.model.dtype_float), + ) + rot.last_alignment_rotation = R_rec + return rot, R_rec + + # ------------------------------------------------------------------ + # Stage 2: per-candidate translation search + local refine + # ------------------------------------------------------------------ + def _prepare_translation_arrays(self) -> None: + """Mask the observations for the translation search and normalise them once. + + The window is ``[tf_d_max, tf_d_min]`` on top of the dataset's own + validity mask, and by default it is the rotation search's ``[d_max, + d_min]``: one resolution window, one Wilson normalisation, both stages. + + The uncut set is not a safe default. With all data -- 228k reflections + to 1.5 A on 2DQ6 -- the translation search places the four largest panel + structures (2DQ6, 3VRJ, 4BX9, 6G9X) at the right orientation and 20-56 A + from the true position, on every trial, while its own score is HIGHER + at the wrong place than at the deposited pose (0.665 against 0.350 on + 2DQ6, where the likelihood is 1616 against 157865). At 15-4 A the same + search recovers all thirty poses to within 0.32 A. The objective's + calc side is raw ``|F_calc|^2``, so at high resolution it is dominated + by whatever reflections happen to carry the largest calculated + intensity rather than by the fit; the window is the first line of + defence and the normalisation of that objective is the second. + """ + data = self.data + device = self.device + hkl_full = data.hkl + F_obs_full = data.F + if hasattr(data, "get_valid_mask"): + tmask = data.get_valid_mask() + else: + tmask = torch.ones( + F_obs_full.shape[0], dtype=torch.bool, device=F_obs_full.device, ) - - scored = [] - for rot, trans in candidates: - clash = self._compute_clash(rot, trans) - if clash <= max_clash: - scored.append((rot, trans, clash)) - - # Sort by combined score (rotation_sigma + trans_sigma - clash_penalty) - def combined_score(x): - rot, trans, clash = x - return rot[4] + trans.sigma - clash / 100.0 - - scored.sort(key=combined_score, reverse=True) - return scored - - def _compute_clash( - self, - rotation_peak: tuple, - trans_peak: TranslationPeak, - ) -> float: - """Compute clash score for a solution.""" - alpha, beta, gamma, _, _ = rotation_peak - R = torch.tensor( - rotation_matrix_from_euler_zyz(alpha, beta, gamma), - dtype=get_float_dtype(), - device=self.device, + real = get_float_dtype() + rec_basis = data.cell.reciprocal_basis_matrix.to(real) + s_all = (hkl_full.to(real) @ rec_basis.to(hkl_full.device)).norm(dim=-1) + if self.tf_d_min > 0.0: + tmask = tmask & (s_all <= 1.0 / self.tf_d_min) + if np.isfinite(self.tf_d_max): + tmask = tmask & (s_all >= 1.0 / self.tf_d_max) + self._tmask = tmask + + sig_F_full = getattr(data, "F_sigma", None) + self._obs = TranslationObs.build( + F_obs_full[tmask], hkl_full[tmask], + data.spacegroup, data.cell, + sig_F=None if sig_F_full is None else sig_F_full[tmask], + delta_vrms_A=self.model_error_A, + device=device, ) + if self.verbose >= 1: + d_hi = 1.0 / float(self._obs.s_mag.max()) + d_lo = 1.0 / float(self._obs.s_mag.min().clamp(min=1e-9)) + self._log(1, f"mr: translation set {self._obs.F_obs.numel()} " + f"reflections, {d_lo:.1f}-{d_hi:.2f} A" + + ("" if sig_F_full is not None + else " (no sigmas: unit weight)")) + + # One P1 copy of the search model for the whole run, re-oriented in + # place per candidate. Two copies per candidate -- one to rotate, one to + # set P1 on -- were a quarter of the run on the large structures. + # + # Its FFT grid is sized to the translation set, not to the model's + # default 1.0 A: |s| is invariant under the symmetry rotations, so every + # rotated index the evaluator is asked for lies inside 1/tf_d_min. Two + # thirds of the window's resolution, not the resolution itself: at + # max_res = tf_d_min the transform's coherence with the 1.0 A grid over + # the 15-4 A set is 0.987 on 2DQ6 (0.9987-0.9999 on 1DAW, 3K7M, 4BX9); + # at tf_d_min/1.5 it is 0.9995-1.0000 everywhere, at 10-38 ms against + # 200-860 ms. max_res first -- the space-group setter rebuilds the FFT + # and reads it. + p1 = self.model.copy() + if self.tf_d_min > 0.0: + p1.max_res = self.tf_d_min / 1.5 + p1.spacegroup = "P 1" + self._p1 = p1 + self._p1_xyz0 = p1.xyz().detach().clone() + self._p1_center = self._p1_xyz0.mean(dim=0) + + def _placement_for_candidate(self) -> Optional[tuple]: + """Translation search for the orientation currently in the P1 template. + + Returns ``(r_analytic, t, tf_score, llg)`` for the translation the + likelihood prefers among the fast search's top peaks, or ``None`` if the + map had no peaks. All three scores are at the same ``t``, so the + reported numbers belong to the placement that was actually chosen. + """ + data = self.data + timer = self._timer + obs = self._obs + + timer.start("5_candidate_transform") + cand = prepare_candidate(self._p1, obs, data.spacegroup, data.cell) + timer.stop("5_candidate_transform") + + # One FFT on a grid a third of the set's resolution apart: dense enough + # that the parabolic peak refinement lands within a fraction of a step, + # and no coarse-then-refine pair whose coarse half could miss the peak. + d_min_set = 1.0 / float(obs.s_mag.max()) + timer.start("6_translation_function") + _, t_peaks = fast_translation_function( + obs, cand, data.cell, + grid_spacing_A=d_min_set / 3.0, + n_peaks=self.n_translation_candidates, + cluster_radius_A=d_min_set, + ) + timer.stop("6_translation_function") + if not t_peaks: + return None - xyz = self.model.xyz() - xyz_centered = xyz - xyz.mean(dim=0) - xyz_rotated = xyz_centered @ R.T - - # Apply translation (fractional -> Cartesian) - trans_frac = torch.tensor( - trans_peak.translation, - dtype=get_float_dtype(), - device=self.device, + timer.start("7_translation_llg") + t_cands = torch.as_tensor( + np.stack([p.translation for p in t_peaks]), dtype=get_float_dtype(), ) - # cart = frac @ B.T (B = fractional_matrix); the transpose is required - # for non-orthogonal cells. - trans_cart = trans_frac @ self.data.cell.fractional_matrix.T.to(self.device) - xyz_final = xyz_rotated + trans_cart - - atom_mask = AtomSampler.from_model(self.model, mode='ca_only') - with torch.no_grad(): - clash = self._clash_calc( - xyz=xyz_final, - cell=self.data.cell.data, - atom_mask=atom_mask, - ).item() - return clash - - def _rigid_body_refine( - self, - rotation_peak: tuple, - trans_peak: TranslationPeak, - ) -> RigidBodyResult: - """Run rigid body refinement.""" - alpha, beta, gamma, _, _ = rotation_peak - - rb = RigidBodyRefinement( - model=self.model, - data=self.data, - initial_rotation=torch.tensor([alpha, beta, gamma], dtype=get_float_dtype()), - initial_translation=torch.tensor( - trans_peak.translation, dtype=get_float_dtype() - ), - device=self.device, - verbose=max(0, self.verbose - 1), + llg = llg_at_translations(obs, cand, t_cands) + k_best = int(llg.argmax()) + t_best = t_cands[k_best] + r_analytic = analytic_r_at(obs, cand, t_best) + timer.stop("7_translation_llg") + for k_t, tp in enumerate(t_peaks): + self._log(3, f" trans{k_t}: tf={tp.score:.4f} z={tp.sigma:.2f} " + f"llg={float(llg[k_t]):.1f} " + f"t={[round(float(x), 3) for x in tp.translation]}") + return (r_analytic, t_best, float(t_peaks[k_best].score), float(llg[k_best])) + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + + +def align_model_to_data( + model: "ModelFT", + data: "ReflectionData", + *, + d_min: float = 4.0, + d_max: float = 15.0, + n_shells: int = 20, + n_rotation_peaks: int = 500, + verbose: int = 0, + do_translation: bool = True, + n_translation_candidates: int = 3, + n_rotation_candidates: int = 10, + rank_by: str = "llg", + tf_d_min: Optional[float] = None, + tf_d_max: Optional[float] = None, + model_error_A: Optional[float] = None, +) -> "ModelFT": + """Place ``model`` in ``data``'s crystal: rotation search, then translation. + + Returns a new rotated+translated ``ModelFT`` carrying + ``last_alignment_rotation``, ``last_alignment_translation`` and + ``last_alignment_rfactor`` provenance attributes. It is a *placement*, not a + refined structure -- refine it downstream. + + `MolecularReplacementPipeline` is the implementation of record; this + function returns its single best solution. + """ + if not model.ctx.initialized: + raise RuntimeError( + "Cannot fit an uninitialized ModelFT. Load PDB data first." ) - return rb.refine() - - def _get_e_values_obs( - self, - d_min: float, - d_max: float, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Get E-values for observed data (cached).""" - if self._e_obs is None: - from torchref.base.alignment import F_squared_to_E_values - - F_obs = self.data.F - mask = self.data.get_valid_mask() - F_obs_masked = F_obs[mask] - - F2 = (F_obs_masked ** 2).to(torch.float64) - s = self._get_s_vectors()[mask] - - E_values, E_squared, _ = F_squared_to_E_values( - F2, s, n_shells=20, d_min=d_min, d_max=d_max - ) - self._e_obs = E_squared # Use E² for correlation - self._s_obs = s - self._mask = mask # Cache mask for later use - return self._e_obs, self._s_obs - - def _get_e_values_calc( - self, - d_min: float, - d_max: float, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Get E-values for calculated data (cached).""" - if self._e_calc is None: - from torchref.base.alignment import F_squared_to_E_values - - hkl = self.data.hkl - mask = self.data.get_valid_mask() - - with torch.no_grad(): - F_calc = self.model(hkl).abs() - F_calc_masked = F_calc[mask] - F2 = (F_calc_masked ** 2).to(torch.float64) - s = self._get_s_vectors()[mask] - - E_values, E_squared, _ = F_squared_to_E_values( - F2, s, n_shells=20, d_min=d_min, d_max=d_max - ) - self._e_calc = E_squared - self._s_calc = s - return self._e_calc, self._s_calc - - def _get_s_vectors(self) -> torch.Tensor: - """Get scattering vectors (cached).""" - if self._s_vectors is None: - from torchref.base import reciprocal_basis_matrix - rec_basis = reciprocal_basis_matrix(self.model.cell) - self._s_vectors = self.data.hkl.to(torch.float64) @ rec_basis.to(torch.float64) - return self._s_vectors - - def clear_cache(self): - """Clear cached E-values and s-vectors.""" - self._e_obs = None - self._e_calc = None - self._s_vectors = None - self._s_obs = None - self._s_calc = None - self._mask = None + pipeline = MolecularReplacementPipeline( + data, model, + verbose=verbose, + d_min=d_min, d_max=d_max, n_shells=n_shells, + n_rotation_peaks=n_rotation_peaks, + model_error_A=model_error_A, + n_rotation_candidates=n_rotation_candidates, + n_translation_candidates=n_translation_candidates, + rank_by=rank_by, + tf_d_min=tf_d_min, tf_d_max=tf_d_max, + ) + solutions = pipeline.run(do_translation=do_translation) + return solutions[0].model diff --git a/torchref/experimental/alignment/rigid_body.py b/torchref/experimental/alignment/rigid_body.py deleted file mode 100644 index e914dc13..00000000 --- a/torchref/experimental/alignment/rigid_body.py +++ /dev/null @@ -1,502 +0,0 @@ -""" -Rigid Body Refinement for Molecular Replacement. - -Implements rigid body refinement where rotation and translation parameters -are optimized to minimize the difference between F_calc and F_obs. -This follows the rotation search (FRF) and translation search stages. - -The refinement optimizes 6 parameters: -- 3 rotation angles (alpha, beta, gamma) as small perturbations -- 3 translation components (fractional coordinates) - -Key design: Bypasses Model/MixedTensor to maintain gradient flow. -Stores all required tensors and uses FFT.compute_structure_factors() directly. - -Gradient flow: - d_alpha → rotation_matrix → xyz_transformed → FFT.compute_structure_factors() → loss - -Uses ScalerBase for proper crystallographic scaling during optimization. - -Experimental / unstable API: part of ``torchref.experimental.alignment``, -the opt-in ball-harmonic MR engine. The production MR entry point is -``torchref.alignment`` (the consolidated FRF engine). Signatures and behavior -may change without notice. -""" - -from dataclasses import dataclass -from typing import Optional, Tuple - -import numpy as np -import torch -import torch.nn as nn - -from torchref.config import get_default_device -from torchref.scaling import ScalerBase -from torchref.model import SfFFT -from torchref.symmetry import spacegroup -from torchref.refinement.targets import RiceXrayTarget -from torchref.base import rotation_matrix_euler_zyz -from torchref.config import get_default_device -from torchref.model import SfFFT -from torchref.scaling import ScalerBase -from torchref.symmetry import spacegroup -from torchref.utils.device_mixin import DeviceMixin - - -@dataclass -class RigidBodyResult: - """ - Results from rigid body refinement. - - Attributes - ---------- - final_rotation : torch.Tensor - Final Euler angles (alpha, beta, gamma) in radians. - final_translation_frac : torch.Tensor - Final translation in fractional coordinates. - initial_r_factor : float - R-factor before refinement. - final_r_factor : float - R-factor after refinement. - final_ml_loss : float - Final ML loss value. - n_steps : int - Number of optimization steps performed. - LBFGS_iterations : int - Number of LBFGS iterations taken by the optimizer. - LBFGS_function_evaluations : int - Number of objective/closure evaluations performed by the LBFGS optimizer. - converged : bool - Whether the refinement converged. - """ - - final_rotation: torch.Tensor - final_translation_frac: torch.Tensor - initial_r_factor: float - final_r_factor: float - final_ml_loss: float - n_steps: int - LBFGS_iterations: int - LBFGS_function_evaluations: int - converged: bool - - -class RigidBodyRefinement(DeviceMixin, nn.Module): - """ - Rigid body refinement using FFT directly (bypasses Model/MixedTensor). - - Optimizes 6 parameters (3 rotation + 3 translation) to maximize - agreement between calculated and observed structure factors using a - Rice maximum-likelihood X-ray target (``RiceXrayTarget`` -- a PRIVATE target, not a - selectable ``--xray-mode``; see its docstring for why it still exists and why this - caller should probably move to ``nll`` once it has a test). - - Key design: Extracts all tensors from Model once at init, then uses - FFT.compute_structure_factors() directly. This maintains gradient flow: - d_alpha → rotation_matrix → xyz_transformed → FFT → loss - - Parameters - ---------- - model : ModelFT - Model with atomic coordinates (tensors extracted, not stored). - data : ReflectionData - Observed reflection data. - expected_rotational_error : float, optional - Half-width (in radians) of the rotation perturbation; the refinable - angles are bounded to +/- this value. Default is 0.1. - initial_rotation : torch.Tensor, optional - Initial Euler angles (alpha, beta, gamma) in radians. - Default is (0, 0, 0). - initial_translation : torch.Tensor, optional - Initial fractional translation vector (3,). - Default is [0, 0, 0]. - device : torch.device, optional - Computation device. Default is the configured default device. - rfactor_converged_threshold : float, optional - R-work value below which refinement is treated as converged. - Default is 0.45. - max_res : float, optional - High-resolution limit (Angstrom) for the FFT structure factors. - Default is 4.0. - verbose : int, optional - Verbosity level for progress printing. Default is 1. - - Attributes - ---------- - rotation_parameters : nn.Parameter - Unconstrained rotation perturbation parameters (3,); mapped to bounded - Euler-angle perturbations via the ``rotation`` property. - translation_frac : nn.Parameter - Refinable fractional translation. - scaler : ScalerBase - Scaler for crystallographic scaling (jointly optimized). - """ - - def __init__( - self, - model, # ModelFT - data, # ReflectionData - expected_rotational_error: float = 0.1, - initial_rotation: torch.Tensor = torch.tensor( - [0.0, 0.0, 0.0], dtype=torch.float32 - ), - initial_translation: Optional[torch.Tensor] = None, - device: torch.device = None, - rfactor_converged_threshold: float = 0.45, - max_res: float = 4.0, - verbose: int = 1, - ): - super().__init__() - if device is None: - device = get_default_device() - self.device = device - self.data = data - - xyz_iso, adp_iso, occ_iso, A_iso, B_iso = model.get_iso() - - self.register_buffer("xyz_initial", xyz_iso.detach().clone().to(device)) - self.register_buffer("adp_iso", adp_iso.detach().clone().to(device)) - self.register_buffer("occ_iso", occ_iso.detach().clone().to(device)) - self.register_buffer("A_iso", A_iso.detach().clone().to(device)) - self.register_buffer("B_iso", B_iso.detach().clone().to(device)) - centroid = torch.mean(self.xyz_initial, dim=0) - self.register_buffer("centroid", centroid) - - self.cell = data.cell - self.spacegroup = data.spacegroup - - self.fft = SfFFT(self.cell, self.spacegroup, max_res=max_res) - - self.verbose = verbose - self.rfactor_converged_threshold = rfactor_converged_threshold - # Get anisotropic atoms if any - xyz_aniso, u_aniso, occ_aniso, A_aniso, B_aniso = model.get_aniso() - self.has_aniso = len(xyz_aniso) > 0 - if self.has_aniso: - self.xyz_aniso_original = xyz_aniso.detach().clone().to(device) - self.u_aniso = u_aniso.detach().clone().to(device) - self.occ_aniso = occ_aniso.detach().clone().to(device) - self.A_aniso = A_aniso.detach().clone().to(device) - self.B_aniso = B_aniso.detach().clone().to(device) - else: - self.xyz_aniso_original = None - self.u_aniso = None - self.occ_aniso = None - self.A_aniso = None - self.B_aniso = None - - # Store initial rotation - self.register_buffer( - "initial_rotation", initial_rotation.to(device=device).clone() - ) - - self.rotation_parameters = nn.Parameter(torch.zeros(3, device=device)) - self.expected_rotational_error = expected_rotational_error - - # Refinable translation (fractional coordinates) - if initial_translation is None: - initial_translation = torch.zeros(3, device=device) - else: - initial_translation = initial_translation.to(device=device).clone() - self.translation_frac = nn.Parameter(initial_translation) - - self.scaler = ScalerBase(data=data, nbins=20, verbose=0, device=device) - fcalc_initial = self() - - self.scaler.initialize(fcalc_initial) - self.scaler.refine_lbfgs(fcalc=fcalc_initial) - - self.xray_target = RiceXrayTarget(data=self.data, scaler=self.scaler) - - def get_rotation_matrix(self) -> torch.Tensor: - """ - Compute current rotation matrix from Euler angles (differentiable). - - Returns - ------- - torch.Tensor - 3x3 rotation matrix combining initial and perturbation rotations. - """ - - angles = self.initial_rotation + self.rotation - return rotation_matrix_euler_zyz(angles) - - @property - def rotation(self) -> torch.Tensor: - """ - Get current rotation perturbation angles. - - Returns - ------- - torch.Tensor - Current (d_alpha, d_beta, d_gamma) in radians. - """ - return ( - 2 * torch.sigmoid(self.rotation_parameters) - 1 - ) * self.expected_rotational_error - - def get_current_rotation_angles(self) -> torch.Tensor: - """ - Get current rotation angles (initial + perturbation). - - Returns - ------- - torch.Tensor - Current (alpha, beta, gamma) in radians. - """ - angles = self.initial_rotation + self.rotation - return angles - - def get_transformed_xyz(self) -> torch.Tensor: - """ - Transform coordinates - maintains gradient flow. - - Applies rotation around centroid, then translation. - - Returns - ------- - torch.Tensor - Transformed atomic coordinates with shape (n_atoms, 3). - """ - R = self.get_rotation_matrix() - - # Rotate around centroid - xyz_centered = self.xyz_initial - self.centroid - xyz_rotated = xyz_centered @ R.T + self.centroid - - # Apply translation (fractional -> Cartesian). Use the Cell helper - # (cart = frac @ B.T) so non-orthogonal cells are handled correctly. - t_cart = self.cell.fractional_to_cartesian(self.translation_frac) - return xyz_rotated + t_cart - - def get_scale(self) -> float: - """Get current scale factor from scaler.""" - return self.scaler.get_scale() - - def forward(self, debug: bool = False) -> torch.Tensor: - """ - Compute unscaled structure factors using FFT directly. - - Gradient flows: rotation_parameters → R → xyz → density → SF. - Miller indices are read internally from ``self.data.hkl``. - - Parameters - ---------- - debug : bool, optional - If True, print gradient tracking info. Default is False. - - Returns - ------- - torch.Tensor - Unscaled calculated structure factors (scaling done by scaler). - """ - # Get transformed coordinates (has gradient to rotation params) - xyz_transformed = self.get_transformed_xyz() - - hkl = self.data.hkl - - if debug: - print( - f" xyz_transformed.requires_grad: {xyz_transformed.requires_grad}" - ) - print(f" xyz_transformed.grad_fn: {xyz_transformed.grad_fn}") - - # Transform anisotropic atoms if present - xyz_aniso = None - if self.has_aniso: - R = self.get_rotation_matrix() - xyz_aniso_centered = self.xyz_aniso_original - self.centroid - xyz_aniso_rotated = xyz_aniso_centered @ R.T + self.centroid - t_cart = self.cell.fractional_to_cartesian(self.translation_frac) - xyz_aniso = xyz_aniso_rotated + t_cart - - # Compute structure factors via FFT (bypasses MixedTensor!) - # Note: fractional matrices are now obtained from FFT's internal Cell object - sf, _ = self.fft.compute_structure_factors( - hkl=hkl, - xyz_iso=xyz_transformed, - adp_iso=self.adp_iso, - occ_iso=self.occ_iso, - A_iso=self.A_iso, - B_iso=self.B_iso, - xyz_aniso=xyz_aniso, - u_aniso=self.u_aniso if self.has_aniso else None, - occ_aniso=self.occ_aniso if self.has_aniso else None, - A_aniso=self.A_aniso if self.has_aniso else None, - B_aniso=self.B_aniso if self.has_aniso else None, - ) - - if debug: - print(f" sf.requires_grad: {sf.requires_grad}") - print(f" sf.grad_fn: {sf.grad_fn}") - - return sf - - def refine( - self, - n_tries: int = 1, - n_iter: int = 100, - ) -> RigidBodyResult: - """ - Run rigid body refinement with an LBFGS optimizer and ML target. - - Uses a strong-Wolfe line-search LBFGS optimizer and a Rice maximum- - likelihood X-ray target, jointly optimizing the rotation, translation, - and scaler parameters. If R-work has not dropped below - ``rfactor_converged_threshold``, the optimizer is restarted with added - gradient noise, up to ``n_tries`` times. - - Parameters - ---------- - n_tries : int, optional - Maximum number of optimizer restarts (with gradient noise on - non-convergence). Default is 1. - n_iter : int, optional - Currently a no-op: this argument is not read anywhere in the body. - The underlying LBFGS uses a hardcoded ``max_iter=100`` per restart - regardless of this value. Default is 100. - - Returns - ------- - RigidBodyResult - Refinement results including final parameters and R-factors. - """ - import sys - - if self.verbose > 2: - print( - f" Setting up LBFGS optimizer niter = {n_iter} and max tries = {n_tries}" - ) - sys.stdout.flush() - parameters = [ - self.rotation_parameters, - self.translation_frac, - *self.scaler.parameters(), - ] - - self.optimizer = torch.optim.LBFGS( - parameters, lr=1, max_iter=100, line_search_fn="strong_wolfe" - ) - - def loss(): - fcalc = self() - return self.xray_target(fcalc) - - noise = 0 - - def closure(): - self.optimizer.zero_grad() - current_loss = loss() - current_loss.backward() - gradnorm = self.optimizer.param_groups[0]["params"][0].norm().item() - if noise > 0: - self.optimizer.param_groups[0]["params"][0].grad += ( - noise - * torch.randn_like(self.optimizer.param_groups[0]["params"][0].grad) - * gradnorm - ) - return current_loss - - rwork_initial, rfree_initial = self.xray_target.get_rfactor(self()) - - initial_loss = closure().item() - if self.verbose > 0: - print( - f"Initial R-work: {rwork_initial:.4f}, R-free: {rfree_initial:.4f}, ML loss: {initial_loss:.4f}" - ) - - from time import time - - start_time = time() - - tries_needed = 0 - - while True: - tries_needed += 1 - - self.optimizer.step(closure) - with torch.no_grad(): - current_loss = loss().item() - if self.verbose > 1: - print(f"Iter {tries_needed} Current ML loss: {current_loss:.4f}") - final_loss = closure().item() - final_rwork, final_rfree = self.xray_target.get_rfactor(self()) - converged = final_rwork < self.rfactor_converged_threshold - - if converged or tries_needed >= n_tries: - if noise > 0: - noise = 0 - self.optimizer.step(closure) - if self.verbose > 1: - print( - f"Converged at iteration {tries_needed} with R-work: {final_rwork:.4f}" - ) - break - - else: - noise += 1e-2 - self.optimizer = torch.optim.LBFGS( - parameters, lr=1, max_iter=100, line_search_fn="strong_wolfe" - ) - - end_time = time() - - if self.verbose > 0: - print( - f"\nRefinement complete after {tries_needed} steps in {end_time - start_time:.2f} seconds." - ) - print( - f" Final R-work: {final_rwork:.4f} (improved by {rwork_initial - final_rwork:.4f})" - ) - print( - f" Final R-free: {final_rfree:.4f} (improved by {rfree_initial - final_rfree:.4f})" - ) - print( - f" Final rotation angles (deg):", - self.get_current_rotation_angles().rad2deg().detach().cpu().numpy(), - ) - print( - f" Final translation: {self.translation_frac.detach().cpu().numpy()}" - ) - - return RigidBodyResult( - final_rotation=self.get_current_rotation_angles() - .detach() - .cpu() - .numpy() - .tolist(), - final_translation_frac=self.translation_frac.detach() - .cpu() - .numpy() - .tolist(), - initial_r_factor=rwork_initial, - final_r_factor=final_rwork, - final_ml_loss=final_loss, - LBFGS_iterations=self.optimizer.state["n_iter"], - LBFGS_function_evaluations=self.optimizer.state["func_evals"], - n_steps=tries_needed, - converged=converged, - ) - - def get_final_parameters(self) -> dict: - """ - Get final refined parameters. - - Returns - ------- - dict - Dictionary with rotation angles (degrees), translation (fractional), - and scale factor. - """ - angles = self.get_current_rotation_angles().detach().cpu() - rotation_perturbation = self.rotation.detach().cpu() - return { - "alpha_deg": np.degrees(angles[0].item()), - "beta_deg": np.degrees(angles[1].item()), - "gamma_deg": np.degrees(angles[2].item()), - "translation_frac": self.translation_frac.detach().cpu().numpy(), - "scale": self.get_scale(), - "d_alpha_deg": np.degrees(rotation_perturbation[0].item()), - "d_beta_deg": np.degrees(rotation_perturbation[1].item()), - "d_gamma_deg": np.degrees(rotation_perturbation[2].item()), - } diff --git a/torchref/experimental/alignment/rotation_search.py b/torchref/experimental/alignment/rotation_search.py new file mode 100644 index 00000000..245c049b --- /dev/null +++ b/torchref/experimental/alignment/rotation_search.py @@ -0,0 +1,570 @@ +"""Fast rotation function: find the orientations of a search model in a crystal. + +One call, three inputs:: + + from torchref.experimental.alignment import rotation_search + + solutions = rotation_search(model, data, model_error_A=0.8) + placed = model.copy().rotate(solutions.rotations[0].T) + +Everything else is derived from the model, the data and that error, following +Phaser's own chain (``runMR_FRF.cc:419-448``): the spherical-harmonic bandwidth +from the model's mean radius and the data's resolution, the sigma_A fall-off +from the coordinate error, the Wilson normalisation and French-Wilson posterior +from the observations and their sigmas. + +The constants below are engine settings, not tuning knobs. They are scored on +whether the true orientation lands inside the candidate window the downstream +placement search carries forward, over the benchmark structures at seeded +orientations -- not on the median rank, which hides the cases that matter. Each +one's provenance is in its own comment. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional, Tuple + +import torch + +from torchref.config import get_default_device, get_float_dtype +from torchref.scaling.weighting import (DEFAULT_SNR_CAP, + DEFAULT_TRUST_CAP) +from .sh import ( + apply_overall_anisotropy, + assign_shells, + equal_count_shell_edges, + fit_overall_anisotropy, +) + +if TYPE_CHECKING: # pragma: no cover - typing only + from ...io.datasets.reflection_data import ReflectionData + from ...model.model_ft import ModelFT + from .frf.types import RotationPeak + +__all__ = ["FRFInputs", "RotationSolutions", "prepare_frf_inputs", + "rotation_search"] + +# Note for anyone reaching for the constants below programmatically: the package +# re-exports `rotation_search` (the function) under this module's own name, so +# `from torchref.experimental.alignment import rotation_search` binds the +# function, not the module. Use `importlib.import_module` for the module object. + + +#: Spherical-harmonic bandwidth ceiling. Where the model and the resolution ask +#: for more, the resolution is coarsened to match instead -- see +#: ``phaser_lmax_resolution``. +#: +#: Chosen by measurement, and the optimum is interior: over ten structures at +#: ten seeded orientations, truth lands in the top twenty on 95/100 cells at 48, +#: 98/100 at 64 and 98/100 at 100, but the binding case is 1AK5 (P 4 3 2), which +#: manages 6/10, 9/10 and 8/10. Only 64 clears nine of ten on every structure. +#: Phaser's own ceiling is 100 (``DEF_CLMN_LMAX``); here that is both worse on +#: 1AK5 and six to ten times slower, and it needs more than 32 GB on three of +#: the ten. +LMAX_CAP = 64 + +#: SO(3) sample spacing in degrees for the rotation-function grid. Also sets the +#: peak-suppression radius, as ``max(2 * this, 6)`` degrees. Inherited from the +#: configuration every measurement on this engine was made with. +GRID_SAMPLING_DEG = 3.0 + +#: Edge of the P1 box the model's transform is sampled in, as a multiple of the +#: molecular diameter. The box has to be large enough that the periodic images +#: do not overlap the molecule's own Patterson. +DENSE_CALC_PAD = 2.0 + +#: Equal-count resolution shells for the Wilson normalisation, the shell +#: variance weights, the relative Wilson-B fit and the anisotropy fit. +N_WILSON_SHELLS = 20 + +#: Discard rotation-function samples this far below the mean, in standard +#: deviations, before peak finding. Generous: it exists to bound the candidate +#: set, not to select. +SIGMA_THRESHOLD = -5.0 + +#: Babinet bulk-solvent parameters folded into sigma_A on the model side +#: (``EnsemblePDB.cc:96-100``). Phaser's defaults. +SOLVENT_FSOL = 0.95 +SOLVENT_BSOL = 300.0 + +#: Low-resolution cutoff in Angstrom. Effectively none: the rotation function +#: wants the low-resolution terms, which carry the molecular envelope. +LOW_RESOLUTION_CUTOFF_A = 100.0 + +# Two things deliberately absent, both measured and rejected on the same panel: +# +# * **Orbit-deduplicated obs unroll.** Keeping only the distinct positions in +# each reflection's orbit, as Phaser does, rather than all n_ops copies. It +# moves 28 of 100 cells and in both directions -- 26 better, 13 worse against +# the shipped configuration -- with the binding structure unchanged at 9/10. A +# quarter of the results churned for no net gain. +# * **Two-radius Patterson union.** Running the search at two integration radii +# and merging the peak lists by z-score. Exactly double the cost (8.7 s +# against 4.4 s median) and it changes 1 cell in 100, which is the engine's own +# run-to-run spread. An earlier measurement had favoured it; that result does +# not survive the anisotropy fix. + +#: Resolution window ``(d_max, d_min)`` the overall anisotropy is fitted in. +#: The tensor is then applied across the full range. Inherited from the range +#: every measurement on this engine was made with, and not itself measured -- +#: unlike the constants above, this pair has no evidence behind it beyond being +#: the one in use. +ANISO_FIT_WINDOW_A = (15.0, 4.0) + + +@dataclass +class RotationSolutions: + """Candidate orientations for a search model, best first. + + Attributes + ---------- + rotations : torch.Tensor + ``(n, 3, 3)`` float64. ``rotations[i]`` maps the search-model frame onto + the crystal frame, so the coordinate rotation that places the model is + its transpose: ``model.copy().rotate(rotations[i].T)``. Each is + determined only up to the crystal's rotational symmetry -- its mates are + ``rotations[i] @ R_g`` -- and the list carries one representative per + orbit, so consecutive entries are distinct orientations. + scores : torch.Tensor + ``(n,)`` rotation-function value at each orientation. + z_scores : torch.Tensor + ``(n,)`` standard deviations above the mean over the whole SO(3) sample + list. This is the scale to judge a solution on; the raw score is not + comparable between runs. + euler_zyz : torch.Tensor + ``(n, 3)`` Edmonds active ZYZ angles in radians, the engine's native + parametrisation: ``R = R_z(alpha) R_y(beta) R_z(gamma)``. + lmax : int + Spherical-harmonic bandwidth used. + d_min : float + High-resolution limit actually used (Angstrom), after the + bandwidth-resolution coupling. + model_error_A : float + The coordinate error the sigma_A fall-off was built from. + """ + + rotations: torch.Tensor + scores: torch.Tensor + z_scores: torch.Tensor + euler_zyz: torch.Tensor + lmax: int + d_min: float + model_error_A: float + + def __len__(self) -> int: + return int(self.rotations.shape[0]) + + +def fit_anisotropy( + data: "ReflectionData", + *, + d_min: float, + d_max: float, + n_shells: int = N_WILSON_SHELLS, + device=None, +) -> torch.Tensor: + """Fit the overall anisotropy tensor and project it onto the point group. + + Returns ``U`` in Angstrom squared as a ``(3, 3)`` at the configured float + dtype, in the convention ``F_corrected = F * exp(+pi^2 s.U.s)``. It stays on + the data's own device unless ``device`` says otherwise, so nothing crosses a + device boundary to be fitted and come back. + + It used to be pinned to the host in double. Neither is needed. Measured over + the 16 datasets in ``tests/files/mtz``, the same fit in float32 reproduces + the double one to 3.3e-5 relative in ``U`` and 4.5e-6 in the correction + factor it exists to produce; end to end, four of five panel cases return a + bit-identical peak list and the fifth (2DQ6, P3121, the most nearly + isotropic ``U`` of the panel) keeps its top orientation and reshuffles two + near-tied deep ranks. + + The projection matters: an unconstrained six-component fit can return a + tensor the lattice forbids, and applying it then modulates the observations + by a direction-dependent factor the crystal cannot have. Cubic lattices + admit one degree of freedom, tetragonal/trigonal/hexagonal two, + orthorhombic three. + """ + from .sh import hkl_symops_to_cartesian, symmetrize_anisotropy + + real = get_float_dtype() + dev = data.hkl.device if device is None else device + rec_basis = data.cell.reciprocal_basis_matrix.detach().to(device=dev, dtype=real) + hkl = data.hkl.detach().to(device=dev, dtype=real) + s_vec_all = hkl @ rec_basis + s_mag_all = s_vec_all.norm(dim=-1) + keep = (s_mag_all >= 1.0 / d_max) & (s_mag_all <= 1.0 / d_min) + if int(keep.sum()) < n_shells * 5: + raise ValueError( + f"Only {int(keep.sum())} reflections in [{d_min}, {d_max}] A, too " + f"few for {n_shells} shells." + ) + F_obs = data.F.detach().to(device=dev, dtype=real).abs()[keep] + s_vec = s_vec_all[keep] + s_mag = s_mag_all[keep] + centric = ( + data.centric.detach().to(dev)[keep].to(torch.bool) + if hasattr(data, "centric") + else torch.zeros_like(F_obs, dtype=torch.bool) + ) + + edges, _ = equal_count_shell_edges(s_mag, n_shells) + shell_idx = assign_shells(s_mag, edges) + U = fit_overall_anisotropy( + F_obs, s_vec, shell_idx, centric, P=n_shells, min_count=20, + ) + sym_cart = hkl_symops_to_cartesian( + data.spacegroup.matrices.detach().to(device=dev, dtype=real), rec_basis, + ) + return symmetrize_anisotropy(U, sym_cart) + + + +@dataclass +class FRFInputs: + """The observations the rotation search runs on, masked and corrected. + + ``F_obs`` is anisotropy-corrected. ``sig_F`` carries the same correction, + which is a multiplicative factor, so ``F/sigma`` survives it unchanged -- it + is here because the engine builds its measurement weight from the sigmas and + the earlier code discarded them immediately after the Wilson step. ``None`` + when the data carry no sigmas. + """ + + F_obs: torch.Tensor # (N,) anisotropy-corrected amplitudes + sig_F: Optional[torch.Tensor] # (N,) their sigmas, same correction + hkl: torch.Tensor # (N, 3) integer Miller indices + s_vec: torch.Tensor # (N, 3) reciprocal-space Cartesian + s_mag: torch.Tensor # (N,) inverse Angstrom + centric: torch.Tensor # (N,) bool + U_aniso: torch.Tensor # (3, 3) Popov-Bourenkov U + device: torch.device + + +def prepare_frf_inputs( + model: "ModelFT", + data: "ReflectionData", + *, + d_min: float, + d_max: float, + n_shells: int, + verbose: int = 0, +) -> FRFInputs: + """Mask the observations to ``[d_min, d_max]`` and correct their anisotropy. + + The anisotropy tensor comes from :func:`fit_anisotropy`, which is also what + the public :func:`rotation_search` uses. It used to be refitted here by a + second copy of the same six lines over the same window -- two paths to one + number is how they drift apart. + + Everything lands on the configured default device, not on whichever device + ``model`` happens to sit on. + """ + device = get_default_device() + real = get_float_dtype() + + F_obs = data.F.to(real).abs() + hkl_all = data.hkl + rec_basis = data.cell.reciprocal_basis_matrix.to(real) + s_vec_all = hkl_all.to(real) @ rec_basis + s_mag_all = s_vec_all.norm(dim=-1) + keep = (s_mag_all >= 1.0 / d_max) & (s_mag_all <= 1.0 / d_min) + if keep.sum().item() < n_shells * 5: + raise ValueError( + f"Too few reflections ({keep.sum().item()}) in [{d_min},{d_max}] A " + f"for {n_shells} shells; widen the resolution range." + ) + F_obs = F_obs[keep].to(device) + sig_F = getattr(data, "F_sigma", None) + if sig_F is not None: + sig_F = sig_F.to(real)[keep].to(device) + hkl = hkl_all[keep].to(device) + s_vec = s_vec_all[keep].to(device) + s_mag = s_mag_all[keep].to(device) + centric = ( + data.centric[keep].to(torch.bool).to(device) + if hasattr(data, "centric") + else torch.zeros_like(F_obs, dtype=torch.bool) + ) + + U_aniso = fit_anisotropy( + data, d_min=d_min, d_max=d_max, n_shells=n_shells, device=device, + ) + F_obs_aniso = apply_overall_anisotropy(F_obs, s_vec, U_aniso) + # Same multiplicative factor, so F/sigma survives the correction intact. + sig_F_aniso = (None if sig_F is None + else apply_overall_anisotropy(sig_F, s_vec, U_aniso)) + + return FRFInputs( + F_obs=F_obs_aniso, sig_F=sig_F_aniso, hkl=hkl, s_vec=s_vec, + s_mag=s_mag, centric=centric, U_aniso=U_aniso, device=device, + ) + + +def search_peaks( + model: "ModelFT", + data: "ReflectionData", + model_error_A: float, + *, + U_aniso: torch.Tensor, + n_peaks: int, + verbose: int = 0, + device: Optional[torch.device] = None, + obs_weight: str = "inverse_variance", + sigma_a_source: str = "empirical", + apply_bulk_solvent: bool = False, + shell_variance_weights: bool = False, + snr_cap: float = DEFAULT_SNR_CAP, + trust_cap: float = DEFAULT_TRUST_CAP, +) -> Tuple[List["RotationPeak"], int, float]: + """Run the rotation function, returning the engine's own peak list. + + Returns ``(peaks, lmax, d_min)``, where ``peaks`` is a list of + :class:`~torchref.experimental.alignment.frf.types.RotationPeak` in Edmonds + ZYZ. For the placement pipeline, which consumes peaks directly and has + already fitted ``U_aniso`` for its rescore stage; :func:`rotation_search` is + the entry point for everything else. + """ + from ...utils import resolve_device + from .frf.api import FastRotationFunction, phaser_lmax_resolution + from .frf.dense_calc import dense_calc_via_box + + # One device for both inputs, rather than whichever one this function + # happened to read first: `resolve_device` moves them into agreement (with a + # warning) and falls back to the configured default. Data first, matching + # the rest of the codebase. + device = resolve_device(data, model, device=device) + real = get_float_dtype() + with torch.no_grad(): + rec_basis = data.cell.reciprocal_basis_matrix.to(real).to(device) + hkl_all = data.hkl.to(device) + s_vec_all = hkl_all.to(real) @ rec_basis + s_mag_all = s_vec_all.norm(dim=-1) + + d_min_data = float(1.0 / s_mag_all.max().item()) + d_max = float(LOW_RESOLUTION_CUTOFF_A) + + # Couple the bandwidth to the resolution FIRST, because the mask below is + # the only place the coarsened limit is applied -- the engine takes the + # window as given. Data finer than L can represent contributes aliasing + # rather than signal and buries the symmetry-diluted true peak, so + # dropping this cut is not a small error: on 3K7M it moved truth from + # rank 18 to rank 238 and doubled the runtime. + model_radius_A = float( + (model.xyz() - model.xyz().mean(0)).norm(dim=-1).mean().item() + ) + L, d_min = phaser_lmax_resolution(model_radius_A, d_min_data, LMAX_CAP) + + # Masking before the unroll rather than after: |s| is symmetry-invariant + # (symmetry operations are isometries), so the two commute -- and this way + # the discarded high-resolution tail is not first replicated n_ops times. + # The low-resolution half is live, not decorative: 3K7M carries two + # reflections beyond 100 A that it removes. + keep = (s_mag_all >= 1.0 / d_max) & (s_mag_all <= 1.0 / d_min) + + s_asu = s_vec_all[keep] + s_mag_asu = s_mag_all[keep] + F_obs = apply_overall_anisotropy( + data.F.to(real).abs().to(device)[keep], s_asu, U_aniso, + ) + sigF = ( + data.F_sigma.to(real).to(device)[keep] + if getattr(data, "F_sigma", None) is not None + else None + ) + centric = ( + data.centric[keep].to(torch.bool).to(device) + if hasattr(data, "centric") + else torch.zeros_like(F_obs, dtype=torch.bool) + ) + # Unmerged Bijvoet data puts two rows on one canonical index, so the + # unroll below would weight those reflections twice. Detectable, so say + # so rather than quietly double-counting. + if getattr(data, "friedel_merged", True) is False: + warnings.warn( + "data are Bijvoet-unmerged: both members of a pair share a " + "canonical index, so the symmetry unroll weights those " + "reflections twice in the Patterson. Merge first for a clean " + "rotation function.", + RuntimeWarning, stacklevel=2, + ) + + # Expand the observations over the space group's rotations to fill + # reciprocal space. |F(hS)| = |F(h)|, and the harmonics need the full + # sphere: sampling only the asymmetric unit under-determines the + # invariant subspace, which is what breaks the high-symmetry cases. + # + # h' = h.S, i.e. the transpose contraction. It agrees with S.h only for + # orthogonal symmetry matrices, so using S.h works everywhere except + # trigonal and hexagonal, where it mixes non-equivalent reflections into + # one orbit. + sg_mats = data.spacegroup.matrices.to(real).to(device) + n_ops = int(sg_mats.shape[0]) + # `expand_reciprocal` is the package's one implementation of this + # contraction, and it carries the h.S convention so no call site has to + # re-decide the real/reciprocal transpose. It returns (n_ops, N, 3), i.e. + # already op-major, which is the flattening the accumulations downstream + # were measured with -- a different row order changes the summation order + # in the later index_add_/unique and the last bits with it. + # + # It rounds to int64 internally, so the products are exact and the cast + # below loses nothing. It also returns on the SPACE GROUP's device rather + # than the caller's, so the move is load-bearing whenever they differ. + hkl_unrolled = ( + data.spacegroup.expand_reciprocal(hkl_all[keep]) + .reshape(-1, 3) + .to(device=device, dtype=rec_basis.dtype) + ) + s_obs = hkl_unrolled @ rec_basis + # Only the GEOMETRY is unrolled. The amplitudes, sigmas and centric + # flags stay one row per unique reflection and the engine broadcasts + # them after its per-reflection chain, which is symmetry-invariant. The + # op-major flattening above means unrolled row `k * N + n` came from + # unique row `n`, so the map is `arange(N)` tiled n_ops times. + n_unique = int(F_obs.shape[0]) + asu_idx = ( + torch.arange(n_unique, device=device) + .unsqueeze(0) + .expand(n_ops, -1) + .reshape(-1) + ) + + # The model's transform on a dense P1 grid rather than at the crystal's + # own reflections: the crystal lattice is too sparse to determine the + # high-l harmonics for a large molecule. `L` / `d_min` came from the + # bandwidth coupling above, which the obs mask also used. + s_calc, F_calc = dense_calc_via_box( + model, d_max, d_min, pad=DENSE_CALC_PAD, verbose=verbose > 0, + ) + s_calc = s_calc.to(device) + F_calc = F_calc.to(device) + + # No relative Wilson-B match here any more. It multiplied `F_calc` by + # exp(-B s^2/4) -- a smooth function of |s| -- and the engine's very next + # step divides out exactly such a function when it normalises. Measured: + # a relative B of +-30 A^2 moves E by at most 1.3e-7, the fit's own + # convergence tolerance. It was computing a number and having it undone. + # + # Not the same as the earlier finding that knocking it out was + # rank-neutral; that was a measurement about whether it mattered, this is + # that it is arithmetically cancelled. `fit_relative_wilson_b` survives in + # `frf/preprocessing` with no production caller at all -- it was kept for + # the ML rescore, and that was deleted. + + # Point-group rotations in the Cartesian frame, so the peak finder can + # treat an orientation and its symmetry mates as one peak. As a set + # these equal B S B^-1; `hkl_symops_to_cartesian` returns the same + # rotations in a different order. Built on the host in double, where + # the peak finder's 3x3 algebra runs anyway. + from .sh import hkl_symops_to_cartesian + sym_cart = hkl_symops_to_cartesian( + data.spacegroup.matrices.detach().cpu().to(torch.float64), # dtype-ok: 3x3 rotation algebra in double on the host + data.cell.reciprocal_basis_matrix.detach().cpu().to(torch.float64), # dtype-ok: 3x3 rotation algebra in double on the host + ) + + engine = FastRotationFunction( + s_obs, F_obs, centric, sg_mats, + L=L, d_min=d_min, d_max=d_max, + delta_vrms_A=float(model_error_A), + n_wilson_shells=N_WILSON_SHELLS, + sig_F_obs=sigF, + grid_sampling_deg=GRID_SAMPLING_DEG, + asu_idx=asu_idx, + s_mag_asu=s_mag_asu, + obs_weight=obs_weight, snr_cap=snr_cap, trust_cap=trust_cap, + shell_variance_weights=shell_variance_weights, + sym_cart=sym_cart, + ) + _arf, peaks = engine.score_model( + s_calc, F_calc, n_peaks=n_peaks, + sigma_threshold=SIGMA_THRESHOLD, + sigma_a_source=sigma_a_source, + apply_bulk_solvent=apply_bulk_solvent, + solvent_fsol=SOLVENT_FSOL, solvent_bsol=SOLVENT_BSOL, + ) + + return peaks, int(L - 1), float(d_min) + + +def _solutions(peaks: List["RotationPeak"], lmax: int, d_min: float, + model_error_A: float) -> RotationSolutions: + """Package a peak list as the public return type.""" + from torchref.base.alignment.rotation import rotation_matrix_euler_zyz + + euler = torch.tensor( + [[p.alpha, p.beta, p.gamma] for p in peaks], dtype=torch.float64, # dtype-ok: RotationSolutions are documented as float64 + ).reshape(-1, 3) + rotations = ( + rotation_matrix_euler_zyz(euler) + if euler.numel() + else torch.zeros((0, 3, 3), dtype=torch.float64) # dtype-ok: RotationSolutions are documented as float64 + ) + return RotationSolutions( + rotations=rotations, + scores=torch.tensor([p.score for p in peaks], dtype=torch.float64), # dtype-ok: RotationSolutions are documented as float64 + z_scores=torch.tensor([p.sigma for p in peaks], dtype=torch.float64), # dtype-ok: RotationSolutions are documented as float64 + euler_zyz=euler, + lmax=lmax, + d_min=d_min, + model_error_A=float(model_error_A), + ) + + +def rotation_search( + model: "ModelFT", + data: "ReflectionData", + model_error_A: float, + *, + n_peaks: int = 500, + verbose: int = 0, + device: Optional[torch.device] = None, +) -> RotationSolutions: + """Find the orientations of ``model`` consistent with ``data``. + + Parameters + ---------- + model : ModelFT + Search model. Its orientation in the file is the frame the returned + rotations are relative to; its position is irrelevant, since the + rotation function works on the Patterson. + data : ReflectionData + Observed amplitudes. ``F_sigma`` is used for the French-Wilson posterior + when present. + model_error_A : float + Expected r.m.s. coordinate error of the model against the target, in + Angstrom. This sets the sigma_A fall-off, and so how much weight the + high-resolution terms carry. Use + :func:`~torchref.experimental.alignment.frf.preprocessing.oeffner_vrms` + to estimate it from the model's length and sequence identity if it is + not otherwise known. + n_peaks : int, optional + How many candidate orientations to return, best first. Default 500. This + bounds the answer, not the search. + verbose : int, optional + Progress reporting. Default 0, silent. + device : torch.device, optional + Where to run. Default ``None`` takes ``data``'s device, moving ``model`` + to match; an explicit value moves both. With neither carrying one, the + configured default applies. + Returns + ------- + RotationSolutions + Ranked orientations. See that class for the rotation convention. + + Raises + ------ + RuntimeError + If ``model`` has no coordinates loaded. + ValueError + If the data carry too few reflections to bin. + """ + if not model.ctx.initialized: + raise RuntimeError("model has no coordinates; load a PDB first.") + d_max_fit, d_min_fit = ANISO_FIT_WINDOW_A + U_aniso = fit_anisotropy(data, d_min=d_min_fit, d_max=d_max_fit) + peaks, lmax, d_min = search_peaks( + model, data, model_error_A, + U_aniso=U_aniso, n_peaks=n_peaks, verbose=verbose, device=device, + ) + return _solutions(peaks, lmax, d_min, model_error_A) diff --git a/torchref/experimental/alignment/sampling.py b/torchref/experimental/alignment/sampling.py deleted file mode 100644 index 4ced91cb..00000000 --- a/torchref/experimental/alignment/sampling.py +++ /dev/null @@ -1,348 +0,0 @@ -""" -Atom-pair sampling for Patterson-style vector matching. - -Supports weighted sampling to prioritize informative pairs -(heavy atoms, close distances). This is a general weighted atom-pair sampler; -it is consumed by Patterson-style vector-matching code. - -Experimental / unstable API: part of ``torchref.experimental.alignment``, -the opt-in ball-harmonic MR engine. The production MR entry point is -``torchref.alignment`` (the consolidated FRF engine). Signatures and behavior -may change without notice. -""" - -from typing import Optional, Tuple - -import numpy as np -import torch - -from torchref.config import get_float_dtype - - -class VectorSampler: - """ - Samples atom pairs for Patterson vector matching. - - Supports weighted sampling to prioritize informative pairs - (heavy atoms via Z-weighting). Per-atom weights also fold in a 1/B-factor - term (``Z**2 / B`` for ``'Z2'`` weighting, ``1 / B`` for ``'uniform'``), - so more-ordered (low-B) atoms are favoured; see :meth:`_compute_weights`. - Samples pairs from the asymmetric unit (ASU) only - symmetry is already - encoded in the Patterson map. - - Parameters - ---------- - model : Model - TorchRef Model object. The caller is responsible for filtering - atoms (e.g., excluding waters) before passing to this class. - weighting : str, optional - Weighting scheme: 'uniform' or 'Z2' (weight by atomic number squared). - Default is 'Z2'. - seed : int, optional - Random seed for reproducibility. Default is None. - - Attributes - ---------- - model : Model - The model used for sampling. - n_atoms : int - Number of atoms in the model. - weighting : str - Weighting scheme used. - weights : torch.Tensor - Sampling weights for each atom (n_atoms,), combining the Z-weighting - scheme with a 1/B-factor term (see :meth:`_compute_weights`). - rng : torch.Generator - Random number generator. - """ - - def __init__(self, model, weighting: str = "Z2", seed: int = None): - """ - Initialize the VectorSampler. - - Parameters - ---------- - model : Model - TorchRef Model object. The caller is responsible for filtering - atoms (e.g., excluding waters) before passing to this class. - weighting : str - Weighting scheme for sampling. - seed : int, optional - Random seed for reproducibility. - """ - self.model = model - self.n_atoms = len(model.pdb) - self.weighting = weighting - self.rng = ( - torch.Generator().manual_seed(seed) - if seed is not None - else torch.Generator() - ) - self.weights = self._compute_weights() - - def _compute_weights( - self, - ) -> torch.Tensor: - """ - Compute sampling probability for each atom based on atomic number and B-factor. - - Weights are computed as: Z^2 / B (for Z2 weighting) or 1/B (for uniform). - Atoms with lower B-factors (more ordered) get higher weights since they - contribute more signal to the Patterson map. - - Returns - ------- - torch.Tensor - Weight for each atom with shape (n_atoms,). - """ - from torchref.utils.pse import PERIODIC_TABLE - - elements = self.model.pdb.element.values - - Zs = torch.tensor( - [PERIODIC_TABLE[el]["number"] for el in elements], - dtype=get_float_dtype(), - device=self.model.device, - ) - - # Get B-factors and compute reciprocal weights - # Use 1/B so atoms with lower B-factors get higher weights - B_factors = torch.tensor( - self.model.pdb["tempfactor"].values, - dtype=get_float_dtype(), - device=self.model.device, - ) - # Clamp B-factors to avoid division by zero or very small values - B_factors = torch.clamp(B_factors, min=1.0) - B_weights = 1.0 / B_factors - - if self.weighting == "Z2": - weights = Zs**2 * B_weights - else: # uniform - weights = B_weights - - weights = weights / weights.sum() # Normalize to probabilities - return weights - - def sample( - self, n_vectors: int, weights: Optional[torch.Tensor] = None - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Sample atom pairs according to weighting scheme. - - Parameters - ---------- - n_vectors : int - Number of atom pairs to sample. - weights : torch.Tensor, optional - Override weights for sampling. If None, uses self.weights. - - Returns - ------- - tuple[torch.Tensor, torch.Tensor] - Two tensors of shape (n_vectors,) containing - the indices of the sampled atom pairs. - """ - w = weights if weights is not None else self.weights - - # Sample first indices according to weights - idx1 = torch.multinomial(w, n_vectors, replacement=True, generator=self.rng) - - # Sample second indices according to weights - idx2 = torch.multinomial(w, n_vectors, replacement=True, generator=self.rng) - - # Redraw idx2 where it equals idx1 - same_mask = idx1 == idx2 - max_attempts = 100 # Prevent infinite loop - attempt = 0 - while same_mask.any() and attempt < max_attempts: - n_resample = same_mask.sum().item() - idx2[same_mask] = torch.multinomial( - w, n_resample, replacement=True, generator=self.rng - ) - same_mask = idx1 == idx2 - attempt += 1 - - return idx1, idx2 - - -def get_rotation_sampling_range( - rotation_matrices: torch.Tensor, -) -> Tuple[float, float, float]: - """ - Determine rotation angle sampling ranges given point group symmetry operations. - - Given the rotation matrices from a spacegroup's point group, this function - computes the asymmetric unit in rotation space (SO(3)) and returns the - maximum Euler angles (alpha, beta, gamma) needed to cover the asymmetric unit. - - The function uses ZYZ Euler angle convention where: - - alpha: rotation about Z axis, range [0, 2*pi) - - beta: rotation about Y axis, range [0, pi] - - gamma: rotation about Z axis, range [0, 2*pi) - - For crystals, the point group symmetry reduces the search space: - - Triclinic (1): Full SO(3) - (2*pi, pi, 2*pi) - - Monoclinic (2): Half of SO(3) - (2*pi, pi, pi) - - Orthorhombic (222): 1/4 of SO(3) - (pi, pi, pi) - - Tetragonal (4, 422): 1/8 or 1/16 of SO(3) - - Trigonal (3, 32): 1/6 or 1/12 of SO(3) - - Hexagonal (6, 622): 1/12 or 1/24 of SO(3) - - Cubic (23, 432): 1/12 or 1/24 of SO(3) - - Parameters - ---------- - rotation_matrices : torch.Tensor - Point group rotation matrices with shape (N, 3, 3), where N is the - number of symmetry operations. These should be the pure rotation - parts of the spacegroup operations (no translations). - - Returns - ------- - Tuple[float, float, float] - Maximum values for (alpha, beta, gamma) Euler angles in radians. - These define the asymmetric unit in rotation space that needs to - be sampled during molecular replacement searches. - - Examples - -------- - :: - - from torchref.symmetry import SpaceGroup - sg = SpaceGroup('P212121') # Orthorhombic - ranges = get_rotation_sampling_range(sg.matrices) - print(f"alpha: {ranges[0]:.4f}, beta: {ranges[1]:.4f}, gamma: {ranges[2]:.4f}") - # alpha: 3.1416, beta: 3.1416, gamma: 3.1416 - - sg = SpaceGroup('P1') # Triclinic - need full SO(3) - ranges = get_rotation_sampling_range(sg.matrices) - print(f"alpha: {ranges[0]:.4f}, beta: {ranges[1]:.4f}, gamma: {ranges[2]:.4f}") - alpha: 6.2832, beta: 3.1416, gamma: 6.2832 - """ - n_ops = rotation_matrices.shape[0] - - # Convert to numpy for analysis - if isinstance(rotation_matrices, torch.Tensor): - R_ops = rotation_matrices.detach().cpu().numpy() - else: - R_ops = np.array(rotation_matrices) - - # Analyze the point group to determine fold symmetries - # We look for rotation axes and their orders - - # Default: full SO(3) coverage - alpha_max = 2 * np.pi - beta_max = np.pi - gamma_max = 2 * np.pi - - # Identity only (P1) - need full search - if n_ops == 1: - return (alpha_max, beta_max, gamma_max) - - # Analyze rotation axes and angles - axes_and_angles = [] - for R in R_ops: - # Skip identity - trace = np.trace(R) - if np.abs(trace - 3.0) < 1e-6: - continue - - # Rotation angle from trace: trace = 1 + 2*cos(theta) - cos_theta = (trace - 1.0) / 2.0 - cos_theta = np.clip(cos_theta, -1.0, 1.0) - angle = np.arccos(cos_theta) - - if angle < 1e-6: - continue - - # Get rotation axis from antisymmetric part of R - # axis is proportional to (R - R^T) - axis = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]]) - norm = np.linalg.norm(axis) - if norm > 1e-6: - axis = axis / norm - else: - # 180-degree rotation - get axis from R + I - # For 180° rotation, axis is eigenvector with eigenvalue 1 - eigvals, eigvecs = np.linalg.eig(R) - idx = np.argmin(np.abs(eigvals - 1.0)) - axis = np.real(eigvecs[:, idx]) - axis = axis / np.linalg.norm(axis) - - axes_and_angles.append((axis, angle)) - - # Determine fold along principal axes - z_axis = np.array([0, 0, 1]) - y_axis = np.array([0, 1, 0]) - x_axis = np.array([1, 0, 0]) - - z_fold = 1 - y_fold = 1 - x_fold = 1 - - for axis, angle in axes_and_angles: - # Check if axis is along z - if np.abs(np.abs(np.dot(axis, z_axis)) - 1.0) < 0.1: - fold = int(round(2 * np.pi / angle)) - z_fold = max(z_fold, fold) - # Check if axis is along y - elif np.abs(np.abs(np.dot(axis, y_axis)) - 1.0) < 0.1: - fold = int(round(2 * np.pi / angle)) - y_fold = max(y_fold, fold) - # Check if axis is along x - elif np.abs(np.abs(np.dot(axis, x_axis)) - 1.0) < 0.1: - fold = int(round(2 * np.pi / angle)) - x_fold = max(x_fold, fold) - - # Also check for 2-fold along diagonal (orthorhombic has 3 perpendicular 2-folds) - has_three_twofolds = False - twofold_count = 0 - for axis, angle in axes_and_angles: - fold = int(round(2 * np.pi / angle)) - if fold == 2: - twofold_count += 1 - if twofold_count >= 3: - has_three_twofolds = True - - # Determine sampling ranges based on symmetry analysis - # For ZYZ Euler angles: - # - z_fold reduces alpha range - # - 2-fold perpendicular to z reduces beta range to [0, pi/2] in some cases - # - Combined symmetry reduces gamma range - - # Alpha reduction based on z-axis fold - if z_fold > 1: - alpha_max = 2 * np.pi / z_fold - - # Check for 2-fold perpendicular to z-axis (reduces gamma) - has_perp_twofold = False - for axis, angle in axes_and_angles: - fold = int(round(2 * np.pi / angle)) - if fold == 2: - # Check if axis is perpendicular to z - if np.abs(np.dot(axis, z_axis)) < 0.1: - has_perp_twofold = True - break - - if has_perp_twofold: - gamma_max = np.pi - - # For point groups with higher symmetry, apply additional reductions - # Based on number of operations (proxy for point group order) - if n_ops >= 24: - # Cubic or high-symmetry hexagonal - alpha_max = min(alpha_max, np.pi / 2) - gamma_max = min(gamma_max, np.pi / 2) - elif n_ops >= 12: - # Hexagonal 622, Cubic 23, etc. - alpha_max = min(alpha_max, np.pi) - gamma_max = min(gamma_max, np.pi) - elif n_ops >= 8: - # Tetragonal 422 - gamma_max = min(gamma_max, np.pi) - elif has_three_twofolds: - # Orthorhombic 222 - alpha_max = np.pi - gamma_max = np.pi - - return (alpha_max, beta_max, gamma_max) diff --git a/torchref/experimental/alignment/sh.py b/torchref/experimental/alignment/sh.py new file mode 100644 index 00000000..0da057fe --- /dev/null +++ b/torchref/experimental/alignment/sh.py @@ -0,0 +1,579 @@ +"""Leaf mathematics the rotation function needs: Legendre seeds, shells, anisotropy. + +Three unrelated things share this module because they share one consumer. + +* **Legendre recurrence coefficients** and their seed, for the fully-normalised + associated Legendre ``bar_P_l^m(cos theta)``, which the Bessel-SH expansion in + :mod:`~torchref.experimental.alignment.frf.data_mr` and its compiled kernels + build on. Normalised so ``(2l)!`` is never formed explicitly. +* **Equal-count resolution shells** (:func:`equal_count_shell_edges`, + :func:`assign_shells`). Assigned once and passed down: two consumers deriving + their own edges from the same ``|s|`` disagree about the reflections sitting on + a boundary. +* **Overall anisotropy** -- fit in intensity space, projected onto the point + group, applied to amplitudes with the half exponent. The projection is + load-bearing: an unconstrained six-component fit can return a tensor the + lattice forbids. + +This module used to also carry a full spherical-harmonic expansion +(``evaluate_ylm``, ``sh_expand_ball``). Nothing called it -- the FRF's own +expansion superseded it -- so it went. ``_bar_legendre_recurrence`` survived it: +production does not call that either, but the FRF expansion's only *independent* +test reference is built on it. +""" + +from __future__ import annotations + +import math +from typing import Optional, Tuple + +import torch + +from ...config import get_float_dtype + + +def legendre_recurrence_coefficients(L: int, dtype, device): + """Coefficient tables for the fully-normalised Legendre recurrence. + + Returns ``(a, b, sect)`` with ``a``, ``b`` of shape ``(L, L)`` indexed + ``[l, m]`` and ``sect`` of shape ``(L,)``:: + + a_l^m = sqrt((2l-1)(2l+1) / ((l-m)(l+m))) + b_l^m = sqrt((2l+1)(l+m-1)(l-m-1) / ((l-m)(l+m)(2l-3))) + sect_m = sqrt((2m+1) / (2m)) + + ``a`` and ``b`` are **zero wherever m >= l**, which lets a caller run the + vertical recurrence at full width instead of slicing to ``[:l]``: the + out-of-support entries come out zero on their own. ``b`` also vanishes at + l = m+1 through its ``(l-m-1)`` factor, so that case needs no special + handling. + + Split out of :func:`_bar_legendre_recurrence` because the rotation function + runs this recurrence itself, fused with its own accumulation, and two copies + of these formulae would be two chances to get them subtly different. + """ + # Small integers, exact in any float dtype; the results are cast to `dtype`. + ll = torch.arange(L, dtype=dtype, device=device).view(L, 1) + mm = torch.arange(L, dtype=dtype, device=device).view(1, L) + valid = ll > mm + denom = (ll - mm) * (ll + mm) + denom_safe = torch.where(valid, denom, torch.ones_like(denom)) + a = torch.sqrt((2.0 * ll - 1.0) * (2.0 * ll + 1.0) / denom_safe) + b_num = (2.0 * ll + 1.0) * (ll + mm - 1.0) * (ll - mm - 1.0) + b_den = denom_safe * (2.0 * ll - 3.0) + b = torch.sqrt(torch.clamp( + b_num / torch.where(b_den == 0, torch.ones_like(b_den), b_den), min=0.0)) + a = torch.where(valid, a, torch.zeros_like(a)).to(dtype) + b = torch.where(valid, b, torch.zeros_like(b)).to(dtype) + m_arange = torch.arange(L, dtype=dtype, device=device) + sect = torch.sqrt( + (2.0 * m_arange + 1.0) / (2.0 * m_arange).clamp(min=1.0)).to(dtype) + return a, b, sect + + +#: bar_P_0^0. +LEGENDRE_SEED = 1.0 / math.sqrt(4.0 * math.pi) + + +def _bar_legendre_recurrence( + cos_theta: torch.Tensor, + sin_theta: torch.Tensor, + L: int, + keep_l: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Fully-normalised associated Legendre ``bar_P_l^m(cos theta)``, all l < L, m <= l. + + **Kept for its test, deliberately.** Production does not call this: the + Bessel-SH expansion runs the same recurrence inside its compiled kernels, + from :func:`legendre_recurrence_coefficients` and :data:`LEGENDRE_SEED`. + That is exactly why this stays -- it is a second, independent, pure-torch + implementation, and ``tests/unit/frf_separate/test_bessel_sh_grouping.py`` + builds a slow reference expansion on it to check the fused one. The other + tests there compare ``bessel_sh_expand`` against *itself* at a different + grouping, so a dropped term cancels; one did, and only this reference caught + it. Deleting it would leave the expansion checked only against itself. + + Definition: + bar_P_l^m(x) = √[(2l+1)/(4π) · (l-m)!/(l+m)!] · P_l^m(x) + where P_l^m is the *unsigned* associated Legendre (no Condon-Shortley phase). + + The recurrence needs only levels ``l-1`` and ``l-2``, so it carries two + rolling rows rather than reading back out of the full table. That matters at + the sizes the rotation function uses: an all-l table is ``(batch, L, L)``, + which for 1e5 batch entries at L=65 is several GB touched once. + + Parameters + ---------- + cos_theta, sin_theta : torch.Tensor + Matching real tensors of any batch shape. + L : int + Bandwidth; l runs over [0, L). + keep_l : torch.Tensor, optional + Which ``l`` rows to return, as an increasing index tensor. ``None`` + returns all of them. Passing only the rows the caller needs -- the even + ones, for a centrosymmetric Patterson -- halves the output. + + Returns + ------- + bar_P : torch.Tensor, real + Shape ``(..., L, L)``, or ``(..., len(keep_l), L)`` when ``keep_l`` is + given. Entries with m > l are zero. + """ + batch_shape = cos_theta.shape + dtype = cos_theta.dtype + device = cos_theta.device + + inv_sqrt_4pi = LEGENDRE_SEED + + a_coef, b_coef, sect = legendre_recurrence_coefficients( + L, dtype, device) + + cos_e = cos_theta.unsqueeze(-1) # (..., 1) + sin_e = sin_theta.unsqueeze(-1) + + if keep_l is None: + rows = torch.arange(L, device=device) + else: + rows = keep_l.to(device=device, dtype=torch.long) # dtype-ok: index tensor; index_add_/gather need int64 + # l -> its position in the output, or -1 when it is not kept. + where = torch.full((L,), -1, dtype=torch.long, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + where[rows] = torch.arange(rows.numel(), device=device) + where_list = where.tolist() + + out = torch.zeros((*batch_shape, rows.numel(), L), dtype=dtype, device=device) + prev2 = torch.zeros((*batch_shape, L), dtype=dtype, device=device) + prev1 = torch.zeros((*batch_shape, L), dtype=dtype, device=device) + prev1[..., 0] = inv_sqrt_4pi # bar_P_0^0 + if where_list[0] >= 0: + out[..., where_list[0], :] = prev1 + + # One loop over l, updating every m in [0, l] at once. The vertical + # recurrence (m < l) and the sectoral diagonal (m = l) read only levels l-1 + # and l-2, which is what the two rolling rows hold. + for l in range(1, L): + cur = torch.zeros_like(prev1) + cur[..., :l] = ( + a_coef[l, :l] * cos_e * prev1[..., :l] - b_coef[l, :l] * prev2[..., :l] + ) + cur[..., l] = sect[l] * sin_e[..., 0] * prev1[..., l - 1] + pos = where_list[l] + if pos >= 0: + out[..., pos, :] = cur + prev2, prev1 = prev1, cur + + return out + + +def get_axis_order(sym_mats: torch.Tensor, axis: int) -> int: + """ + Order of the highest-multiplicity proper rotation about a principal axis. + + `sym_mats` is the spacegroup rotation matrices (n_ops, 3, 3). `axis` is + 0/1/2 for x/y/z. Returns the largest n such that some R in `sym_mats` is a + rotation by 2π/n around that axis. For non-rotational operations (or + rotations not aligned with the axis), the spacegroup element is skipped. + Returns 1 if no proper rotation around the axis exists. + + Used by the Phaser-style m-symmetry filter on the spherical-harmonic + coefficients: the Patterson is invariant under the spacegroup rotations, + so m-values that violate the highest-order axis symmetry are pure noise. + """ + dtype = sym_mats.dtype if sym_mats.is_floating_point() else get_float_dtype() + R = sym_mats.to(dtype) + a = torch.zeros(3, dtype=dtype, device=R.device) + a[axis] = 1.0 + # The working precision is enough here and the operators can stay wherever + # they are: the entries of a Miller-index symop are small integers, exact in + # float32, and the Cartesian form of one is accurate to ~1e-7 -- an order of + # magnitude inside the 1e-3 and 1e-6 tolerances this test uses. Batched, so + # the answer costs one transfer of two short vectors rather than a device + # sync per operation. + # + # Axis must be invariant under R (proper or improper rotation about it), and + # the trace of a rotation by angle θ about the preserved axis is 1+2cosθ. + keeps = (R @ a - a).norm(dim=-1) <= 1e-3 + cos_a = ((R.diagonal(dim1=-2, dim2=-1).sum(dim=-1) - 1.0) * 0.5).clamp(-1.0, 1.0) + keeps = keeps & (cos_a < 1.0 - 1e-6) # identity (angle ~0) → order 1 + orders = [ + round(2 * math.pi / math.acos(c)) + for c in cos_a[keeps].detach().cpu().tolist() + ] + return max(orders) if orders else 1 + + +def get_high_order_axis(sym_mats: torch.Tensor) -> Tuple[int, int]: + """ + Return (axis, zsymm) where `axis` ∈ {0, 1, 2} (x/y/z) maximises + `get_axis_order`, with z preferred on ties (matches Phaser's + `highOrderAxis()` in SpaceGroup.cc). + """ + orders = [get_axis_order(sym_mats, a) for a in (0, 1, 2)] + # Phaser: axis=3 (z); axis=2 if orderY > orderZ; axis=1 if orderX > both. + if orders[0] > orders[1] and orders[0] > orders[2]: + axis = 0 + elif orders[1] > orders[2]: + axis = 1 + else: + axis = 2 + return axis, orders[axis] + + +def equal_count_shell_edges( + s_magnitudes: torch.Tensor, + P: int, + s_min: Optional[float] = None, + s_max: Optional[float] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Compute equal-count radial shell edges for a list of |s| magnitudes. + + Each of the P shells receives (approximately) the same number of reflections. + Reflections outside [s_min, s_max] (if given) are excluded. + + Returns + ------- + edges : torch.Tensor, real, shape (P+1,) + Shell boundaries, increasing. + centers : torch.Tensor, real, shape (P,) + Mid-points of each shell. + """ + s = s_magnitudes + if s_min is not None: + s = s[s >= s_min] + if s_max is not None: + s = s[s <= s_max] + s_sorted, _ = torch.sort(s) + N = s_sorted.numel() + # quantile-based partition + idx = torch.linspace(0, N - 1, P + 1, dtype=s.dtype, device=s.device).round().long() + edges = s_sorted[idx] + # nudge endpoints so the data is fully covered (avoid floating-point miss) + if s_min is not None: + edges[0] = min(edges[0].item(), s_min) + else: + edges[0] = edges[0] - 1e-6 + if s_max is not None: + edges[-1] = max(edges[-1].item(), s_max) + else: + edges[-1] = edges[-1] + 1e-6 + centers = 0.5 * (edges[:-1] + edges[1:]) + return edges, centers + + +def fit_overall_anisotropy( + F_obs: torch.Tensor, + s_vectors: torch.Tensor, + shell_idx: torch.Tensor, + centric: torch.Tensor, + P: int, + min_count: int = 20, + n_iter: int = 12, +) -> torch.Tensor: + """ + Fit the overall anisotropy tensor U from F_obs alone (no model needed). + + The Popov-Bourenkov correction models the observed intensities as a + per-shell isotropic Wilson piece modulated by an overall anisotropic + Debye-Waller term:: + + E[ I(h) / _shell ] = c * exp(-2 pi^2 s.U.s) + + That expectation is exact in **intensity** space, which is where this fits + it: a free constant ``c`` absorbs the overall scale, the weights come from + ``Var(I/)`` -- 1 for acentric reflections and 2 for centric ones -- and + non-positive or non-finite amplitudes are dropped. Gauss-Newton from + ``U = 0``. + + Fitting the same relation in log space instead is what the earlier version + did, and it is biased: ``E[ln(I/)]`` is ``-gamma = -0.577`` for acentric + and ``-gamma - ln 2`` for centric reflections, not zero. Without a constant + term that offset can only be absorbed by the quadratic form, so U comes back + with a large spurious component -- and because centric reflections lie on + the zones perpendicular to the symmetry axes, the bias is + direction-dependent rather than a harmless overall scale. + + The returned U is the correction to *apply* in the form:: + + F_obs_corrected(h) = F_obs(h) * exp(+pi^2 s.U.s) + + so the corrected amplitudes have the same mean square in every direction. + Project it onto the point group with :func:`symmetrize_anisotropy` before + applying it: an unconstrained six-component fit can return a tensor the + lattice forbids. + + Parameters + ---------- + F_obs : (N,) real + s_vectors : (N, 3) real, reciprocal-space Cartesian (1/Angstrom) + shell_idx : (N,) int64 -- shell of each reflection, in [0, P); negative + entries are excluded + centric : (N,) bool + P : int, number of shells + min_count : int, optional + Shells with fewer reflections than this are dropped, since their mean + intensity is too noisy to normalise against. + n_iter : int, optional + Gauss-Newton iterations. + + Returns + ------- + U : (3, 3) symmetric real tensor (Angstrom squared). Zero if too few + reflections survive to constrain seven parameters. + """ + valid = shell_idx >= 0 + # The fit runs at the amplitudes' own width, wherever they are. It used to + # force double on the host, which was measured against this: over the 16 + # datasets in ``tests/files/mtz``, float32 reproduces U to 3.3e-5 relative + # and the correction it exists to apply, exp(+pi^2 s.U.s), to 4.5e-6. The + # design matrix is well scaled by construction -- a constant column beside + # 2 pi^2 s.s terms of order 0.1-1 over the fitting window -- so there is no + # precision cliff for seven parameters to fall off. + work = F_obs.dtype if F_obs.is_floating_point() else get_float_dtype() + F = F_obs[valid].to(work) + s = s_vectors[valid].to(work) + idx = shell_idx[valid] + cen = centric[valid].bool() + + ok = torch.isfinite(F) & (F > 0) + F, s, idx, cen = F[ok], s[ok], idx[ok], cen[ok] + + I = F * F + count = torch.zeros(P, dtype=torch.int64, device=F.device) # dtype-ok: index tensor; index_add_/gather need int64 + total = torch.zeros(P, dtype=work, device=F.device) + count.index_add_(0, idx, torch.ones_like(idx)) + total.index_add_(0, idx, I) + mean_I = (total / count.clamp(min=1).to(work)).clamp(min=1e-30) + + keep = (count >= min_count)[idx] + if int(keep.sum()) < 50: + return torch.zeros((3, 3), dtype=F_obs.dtype, device=F_obs.device) + ratio = I[keep] / mean_I[idx[keep]] + sk, cenk = s[keep], cen[keep] + + x, y, z = sk[:, 0], sk[:, 1], sk[:, 2] + # s.U.s = Uxx sx^2 + Uyy sy^2 + Uzz sz^2 + # + 2 Uxy sx sy + 2 Uxz sx sz + 2 Uyz sy sz + quad = torch.stack([x * x, y * y, z * z, + 2 * x * y, 2 * x * z, 2 * y * z], dim=1) + # Column 0 is the free constant ln(c); the rest carry -2 pi^2 s.U.s. + A = torch.cat([torch.ones_like(x).unsqueeze(1), + -2.0 * (torch.pi ** 2) * quad], dim=1) + w = torch.where(cenk, torch.full_like(ratio, 0.5), torch.ones_like(ratio)) + + theta = torch.zeros(7, dtype=work, device=F.device) + for _ in range(n_iter): + model = torch.exp((A @ theta).clamp(min=-20.0, max=20.0)) + J = model.unsqueeze(1) * A + Jw = J * w.unsqueeze(1) + H = J.transpose(0, 1) @ Jw + grad = Jw.transpose(0, 1) @ (ratio - model) + H = H + torch.eye(7, dtype=H.dtype, device=H.device) * 1e-12 * float( + torch.diagonal(H).abs().max().clamp(min=1e-30)) + theta = theta + torch.linalg.solve(H, grad) + + u = theta[1:] + return torch.tensor( + [[u[0], u[3], u[4]], [u[3], u[1], u[5]], [u[4], u[5], u[2]]], + dtype=F_obs.dtype, device=F_obs.device, + ) + + +def hkl_symops_to_cartesian( + sg_mats: torch.Tensor, + rec_basis: torch.Tensor, +) -> torch.Tensor: + """ + Convert spacegroup symmetry operators that act on integer Miller indices + into the equivalent rotation operators acting on Cartesian reciprocal- + space vectors (`s = h @ rec_basis`, column-vector form: `s = M @ h` with + `M = rec_basis^T`). + + For column vectors: `s' = M · S · M⁻¹ · s` so `P_cart = M · S · M⁻¹`. + + For orthogonal cells (orthorhombic+) M is diagonal and `P_cart == S` + exactly. For non-orthogonal cells (monoclinic, hex/trig with γ=120°, + triclinic), the Cartesian form differs and matters for any operation + that mixes the axes (e.g. averaging tensors over the point group). + + Parameters + ---------- + sg_mats : torch.Tensor, shape (n_ops, 3, 3) + Integer Miller-index symops (data.spacegroup.matrices). + rec_basis : torch.Tensor, shape (3, 3) + Reciprocal basis matrix such that `s = h @ rec_basis` (row-vector + convention used throughout torchref). + + Returns + ------- + sym_mats_cart : torch.Tensor, shape (n_ops, 3, 3), real + """ + # Whatever width the caller brought. Double when it hands us double -- the + # peak finder's host-side 3x3 algebra does -- and the working precision + # otherwise, so this is usable on a backend without float64. The integer + # symops carry no width of their own, hence the fallback. + dtype = rec_basis.dtype if rec_basis.is_floating_point() else get_float_dtype() + M = rec_basis.to(dtype).transpose(-1, -2).contiguous() # (3, 3) + # `.contiguous()` is not cosmetic on a 3x3: MPS's `linalg.inv` trips an + # internal contiguity assert on the transposed view (torch 2.9.1), and the + # copy costs nine elements. + M_inv = torch.linalg.inv(M) + S = sg_mats.to(device=M.device, dtype=dtype) # (n_ops, 3, 3) + # S^T, not S: reciprocal space transforms as h' = h.S, so the operator + # acting on Cartesian s as a column vector is (B^-1 S B)^T = M S^T M^-1 + # with M = B^T. Using S here returns matrices that are not rotations at all + # in a non-orthogonal basis -- measured orthogonality error 5.33 for + # P 3_1 2 1 and P 6_5 2 2, versus 2e-7 with the transpose. The two agree + # whenever the symmetry matrices are orthogonal, i.e. everywhere except + # trigonal/hexagonal, which is why this survived. + return torch.einsum("ij,klj,lm->kim", M, S, M_inv) + + +def symmetrize_anisotropy( + U: torch.Tensor, + sym_mats_cart: torch.Tensor, +) -> torch.Tensor: + """ + Project a symmetric 3×3 tensor `U` onto the point-group-invariant + subspace of the spacegroup by averaging over its Cartesian rotation + operators: + + U_sym = (1/n) Σ_k P_k · U · P_k^T + + This mirrors Phaser's `site_symmetry.average_u_star()` and the + `RefineANO.cc:116-142` constraint construction. After symmetrisation the + tensor automatically satisfies the crystal's point-group symmetry: + + - cubic: U_sym = (trace U / 3) · I (1 DOF) + - tetragonal: U_sym = diag(λ, λ, μ) (2 DOF) + - orthorhombic: U_sym = diag(λ, μ, ν) (3 DOF) + - monoclinic: diagonal + one off-diagonal (4 DOF) + - triclinic: unchanged (6 DOF) + + This is the structural fix that prevents an unconstrained 6-component + regression from producing physically impossible anisotropy in + high-symmetry cells (e.g. fitting a 70 Ų eigenvalue on a cubic dataset + where every eigenvalue must be equal by symmetry). + + Parameters + ---------- + U : torch.Tensor, shape (3, 3), symmetric + sym_mats_cart : torch.Tensor, shape (n_ops, 3, 3) + Output of `hkl_symops_to_cartesian` — Cartesian-space rotation + operators of the spacegroup. + + Returns + ------- + U_sym : torch.Tensor, shape (3, 3), symmetric, same dtype/device as U + """ + dtype = U.dtype + device = U.device + sym = sym_mats_cart.to(device=device, dtype=dtype) + # Vectorised: U_avg = mean_k P_k · U · P_k^T + U_avg = torch.einsum("kij,jl,knl->in", sym, U, sym) / sym.shape[0] + # Enforce symmetry (numerical hygiene; pure rotation averaging preserves + # symmetry exactly, but FP drift can leave 1e-15 asymmetry). + return 0.5 * (U_avg + U_avg.transpose(-1, -2)) + + +def apply_overall_anisotropy( + F: torch.Tensor, + s_vectors: torch.Tensor, + U: torch.Tensor, +) -> torch.Tensor: + """ + Apply the inverse of the anisotropy tensor to amplitudes: + F_corrected(h) = F(h) · exp(+π²·s·U·s) + (Inverse direction = +π² to undo the Debye-Waller effect.) + """ + device = F.device + dtype = F.dtype + # One `.to` per tensor, not two: `.to(device).to(dtype)` materialises the + # source width on the target device first, which throws on a backend that + # has no float64 -- and a host-side double U is a thing this is handed. + s = s_vectors.to(device=device, dtype=dtype) + U_t = U.to(device=device, dtype=dtype) + s_dot_U = s @ U_t # (N, 3) + arg = (torch.pi ** 2) * (s_dot_U * s).sum(dim=-1) + return F * torch.exp(arg.clamp(min=-10.0, max=10.0)) + + +def compute_patterson_shell_variance( + patt: torch.Tensor, + shell_idx: torch.Tensor, + P: int, + min_count: int = 8, + eps: float = 1e-3, +) -> torch.Tensor: + """ + Empirical per-shell variance of a Patterson coefficient (typically E²−1). + + For each shell p, returns Var_p = _p - _p². + + Shells with fewer than `min_count` reflections inherit the variance of the + nearest shell that does meet the count threshold (search outward, prefer + higher-resolution / larger-index shells first). The result is clamped at + `eps` so the inverse-sqrt weight `1/√Var_p` stays bounded. + + Parameters + ---------- + patt : torch.Tensor, shape (N,) + Per-reflection Patterson coefficient (e.g. E² − 1). + shell_idx : torch.Tensor, shape (N,), int64 + Shell assignment in [0, P). Reflections with index -1 are ignored. + P : int + Number of shells. + min_count : int, default 8 + Shells with fewer reflections than this borrow from a neighbour. + eps : float, default 1e-3 + Lower bound on returned variance. + + Returns + ------- + var : torch.Tensor, shape (P,), same dtype as `patt` + """ + device = patt.device + dtype = patt.dtype + valid = shell_idx >= 0 + patt_v = patt[valid] + idx_v = shell_idx[valid] + count = torch.zeros(P, dtype=torch.int64, device=device) # dtype-ok: index tensor; index_add_/gather need int64 + count.index_add_(0, idx_v, torch.ones_like(idx_v)) + sum1 = torch.zeros(P, dtype=dtype, device=device) + sum2 = torch.zeros(P, dtype=dtype, device=device) + sum1.index_add_(0, idx_v, patt_v) + sum2.index_add_(0, idx_v, patt_v * patt_v) + safe_count = count.clamp(min=1).to(dtype) + mean = sum1 / safe_count + var = sum2 / safe_count - mean * mean + var = var.clamp(min=eps) + + counts_l = count.tolist() + var_l = var.tolist() + good = [i for i, c in enumerate(counts_l) if c >= min_count] + if not good: + return torch.full_like(var, max(eps, 1.0)) + for i in range(P): + if counts_l[i] >= min_count: + continue + nearest = min(good, key=lambda g, ii=i: (abs(g - ii), -g)) + var_l[i] = var_l[nearest] + return torch.tensor(var_l, dtype=dtype, device=device).clamp(min=eps) + + +def assign_shells( + s_magnitudes: torch.Tensor, + edges: torch.Tensor, +) -> torch.Tensor: + """ + Assign each reflection to a shell index in [0, P). + + Reflections strictly outside the edges get index -1 (caller may filter). + """ + # bucketize returns indices in [0, P+1]; we want shell indices in [0, P). + # Reflections with s == edges[0] go to bucket 0 (use right=True trick). + idx = torch.bucketize(s_magnitudes.contiguous(), edges.contiguous(), right=True) - 1 + P = edges.shape[0] - 1 + invalid = (idx < 0) | (idx >= P) + idx = idx.clamp(min=0, max=P - 1) + idx = torch.where(invalid, torch.full_like(idx, -1), idx) + return idx diff --git a/torchref/experimental/alignment/transform.py b/torchref/experimental/alignment/transform.py deleted file mode 100644 index 876e0e17..00000000 --- a/torchref/experimental/alignment/transform.py +++ /dev/null @@ -1,947 +0,0 @@ -""" -Rigid body transformations for crystallographic alignment. - -Provides unified handling of rotations and translations with quaternion-based -internal storage and multiple representation formats. - -Experimental / unstable API: part of ``torchref.experimental.alignment``, -the opt-in ball-harmonic MR engine. The production MR entry point is -``torchref.alignment`` (the consolidated FRF engine). Signatures and behavior -may change without notice. -""" - -from typing import Optional, Union - -import torch -import torch.nn as nn - -from torchref.config import get_default_device, get_float_dtype -from torchref.utils.device_mixin import DeviceMixin - -# ============================================================================= -# Quaternion Helper Functions -# ============================================================================= - - -def get_inverse_rotation_matrix(R: torch.Tensor) -> torch.Tensor: - """ - Compute inverse of rotation matrix (transpose for orthogonal matrices). - - Parameters - ---------- - R : torch.Tensor - Rotation matrix of shape (3, 3) or (N, 3, 3). - - Returns - ------- - torch.Tensor - Inverse rotation matrix of same shape. - """ - return R.transpose(-2, -1) - - -def sample_angles(sampling_pitch_rad, max_angles_rad): - """ - Sample a regular grid of Euler angles (in radians). - - Builds the Cartesian product of per-axis ranges from 0 to each maximum - angle (inclusive) with the given step. - - Parameters - ---------- - sampling_pitch_rad : float - Sampling pitch (step size) in radians, used for all three axes. - max_angles_rad : tuple of float - Maximum angles (alpha, beta, gamma) in radians. - - Returns - ------- - torch.Tensor - Sampled angles of shape (N, 3), where N is the number of grid points. - - See Also - -------- - torchref.experimental.alignment.sampling : Related atom-pair sampling - utilities (this Euler-angle grid sampler is thematically a sampling - helper but lives here in ``transform``). - """ - - angles = [] - max_alpha, max_beta, max_gamma = max_angles_rad - alpha = torch.arange(0, max_alpha + 1e-6, sampling_pitch_rad, dtype=torch.float32) - beta = torch.arange(0, max_beta + 1e-6, sampling_pitch_rad, dtype=torch.float32) - gamma = torch.arange(0, max_gamma + 1e-6, sampling_pitch_rad, dtype=torch.float32) - alpha, beta, gamma = torch.meshgrid(alpha, beta, gamma, indexing="ij") - - return torch.stack([alpha.flatten(), beta.flatten(), gamma.flatten()], dim=-1) - - -def rotation_matrix_from_euler(angles): - """ - Compute rotation matrices from Euler angles (in radians). - - Each row of ``angles`` is (alpha, beta, gamma) and the matrix is built as - ``Rz(gamma) @ Ry(beta) @ Rz(alpha)`` (an intrinsic ZYZ composition applied - in the order alpha, then beta, then gamma). - - Parameters - ---------- - angles : torch.Tensor - Euler angles of shape (N, 3), columns [alpha, beta, gamma]. - - Returns - ------- - torch.Tensor - Rotation matrices of shape (N, 3, 3). - """ - alpha = angles[:, 0] - beta = angles[:, 1] - gamma = angles[:, 2] - - R_alpha = torch.stack( - [ - torch.cos(alpha), - -torch.sin(alpha), - torch.zeros_like(alpha), - torch.sin(alpha), - torch.cos(alpha), - torch.zeros_like(alpha), - torch.zeros_like(alpha), - torch.zeros_like(alpha), - torch.ones_like(alpha), - ], - dim=-1, - ).reshape(-1, 3, 3) - - R_beta = torch.stack( - [ - torch.cos(beta), - torch.zeros_like(beta), - torch.sin(beta), - torch.zeros_like(beta), - torch.ones_like(beta), - torch.zeros_like(beta), - -torch.sin(beta), - torch.zeros_like(beta), - torch.cos(beta), - ], - dim=-1, - ).reshape(-1, 3, 3) - - R_gamma = torch.stack( - [ - torch.ones_like(gamma), - torch.zeros_like(gamma), - torch.zeros_like(gamma), - torch.zeros_like(gamma), - torch.cos(gamma), - -torch.sin(gamma), - torch.zeros_like(gamma), - torch.sin(gamma), - torch.cos(gamma), - ], - dim=-1, - ).reshape(-1, 3, 3) - - R = torch.einsum("rij,rjk,rkl->ril", R_gamma, R_beta, R_alpha) - - return R - - -def quaternion_normalize(q: torch.Tensor) -> torch.Tensor: - """ - Normalize quaternion to unit length. - - Parameters - ---------- - q : torch.Tensor - Quaternion(s) of shape (4,) or (N, 4). - - Returns - ------- - torch.Tensor - Normalized quaternion(s) of same shape. - """ - return q / q.norm(dim=-1, keepdim=True).clamp(min=1e-12) - - -def quaternion_conjugate(q: torch.Tensor) -> torch.Tensor: - """ - Compute quaternion conjugate (inverse for unit quaternions). - - For q = [w, x, y, z], conjugate is [w, -x, -y, -z]. - - Parameters - ---------- - q : torch.Tensor - Quaternion(s) of shape (4,) or (N, 4). - - Returns - ------- - torch.Tensor - Conjugate quaternion(s) of same shape. - """ - conj = q.clone() - conj[..., 1:] = -conj[..., 1:] - return conj - - -def quaternion_multiply(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: - """ - Compute Hamilton product of two quaternions. - - Parameters - ---------- - q1, q2 : torch.Tensor - Quaternions of shape (4,) or (N, 4). Format: [w, x, y, z]. - - Returns - ------- - torch.Tensor - Product quaternion of same shape. - """ - w1, x1, y1, z1 = q1[..., 0], q1[..., 1], q1[..., 2], q1[..., 3] - w2, x2, y2, z2 = q2[..., 0], q2[..., 1], q2[..., 2], q2[..., 3] - - w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 - x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 - y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 - z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 - - return torch.stack([w, x, y, z], dim=-1) - - -def quaternion_rotate(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - """ - Rotate vector(s) by quaternion using q * v * q^*. - - Parameters - ---------- - q : torch.Tensor - Unit quaternion of shape (4,). - v : torch.Tensor - Vector(s) of shape (3,) or (N, 3). - - Returns - ------- - torch.Tensor - Rotated vector(s) of same shape as v. - """ - # Convert vector to pure quaternion [0, x, y, z] - v_shape = v.shape - v_flat = v.reshape(-1, 3) - - v_quat = torch.zeros(v_flat.shape[0], 4, dtype=q.dtype, device=q.device) - v_quat[:, 1:] = v_flat - - # q * v * q^* - q_conj = quaternion_conjugate(q) - result = quaternion_multiply( - quaternion_multiply(q.unsqueeze(0), v_quat), q_conj.unsqueeze(0) - ) - - # Extract vector part - return result[:, 1:].reshape(v_shape) - - -def quaternion_to_matrix(q: torch.Tensor) -> torch.Tensor: - """ - Convert unit quaternion to 3x3 rotation matrix. - - Parameters - ---------- - q : torch.Tensor - Unit quaternion of shape (4,) or (N, 4). Format: [w, x, y, z]. - - Returns - ------- - torch.Tensor - Rotation matrix of shape (3, 3) or (N, 3, 3). - """ - q = quaternion_normalize(q) - - batched = q.dim() == 2 - if not batched: - q = q.unsqueeze(0) - - w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] - - # Rotation matrix from quaternion - R = torch.stack( - [ - torch.stack( - [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)], - dim=-1, - ), - torch.stack( - [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)], - dim=-1, - ), - torch.stack( - [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], - dim=-1, - ), - ], - dim=-2, - ) - - if not batched: - R = R.squeeze(0) - - return R - - -def matrix_to_quaternion(R: torch.Tensor) -> torch.Tensor: - """ - Convert 3x3 rotation matrix to unit quaternion. - - Uses Shepperd's method for numerical stability. - - Parameters - ---------- - R : torch.Tensor - Rotation matrix of shape (3, 3) or (N, 3, 3). - - Returns - ------- - torch.Tensor - Unit quaternion of shape (4,) or (N, 4). Format: [w, x, y, z]. - """ - batched = R.dim() == 3 - if not batched: - R = R.unsqueeze(0) - - batch_size = R.shape[0] - dtype, device = R.dtype, R.device - - # Shepperd's method - choose largest diagonal element - trace = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2] - - q = torch.zeros(batch_size, 4, dtype=dtype, device=device) - - # Case 1: trace > 0 - mask1 = trace > 0 - if mask1.any(): - s = torch.sqrt(trace[mask1] + 1.0) * 2 # s = 4 * w - q[mask1, 0] = 0.25 * s - q[mask1, 1] = (R[mask1, 2, 1] - R[mask1, 1, 2]) / s - q[mask1, 2] = (R[mask1, 0, 2] - R[mask1, 2, 0]) / s - q[mask1, 3] = (R[mask1, 1, 0] - R[mask1, 0, 1]) / s - - # Case 2: R[0,0] is largest diagonal - mask2 = ~mask1 & (R[:, 0, 0] > R[:, 1, 1]) & (R[:, 0, 0] > R[:, 2, 2]) - if mask2.any(): - s = torch.sqrt(1.0 + R[mask2, 0, 0] - R[mask2, 1, 1] - R[mask2, 2, 2]) * 2 - q[mask2, 0] = (R[mask2, 2, 1] - R[mask2, 1, 2]) / s - q[mask2, 1] = 0.25 * s - q[mask2, 2] = (R[mask2, 0, 1] + R[mask2, 1, 0]) / s - q[mask2, 3] = (R[mask2, 0, 2] + R[mask2, 2, 0]) / s - - # Case 3: R[1,1] is largest diagonal - mask3 = ~mask1 & ~mask2 & (R[:, 1, 1] > R[:, 2, 2]) - if mask3.any(): - s = torch.sqrt(1.0 + R[mask3, 1, 1] - R[mask3, 0, 0] - R[mask3, 2, 2]) * 2 - q[mask3, 0] = (R[mask3, 0, 2] - R[mask3, 2, 0]) / s - q[mask3, 1] = (R[mask3, 0, 1] + R[mask3, 1, 0]) / s - q[mask3, 2] = 0.25 * s - q[mask3, 3] = (R[mask3, 1, 2] + R[mask3, 2, 1]) / s - - # Case 4: R[2,2] is largest diagonal - mask4 = ~mask1 & ~mask2 & ~mask3 - if mask4.any(): - s = torch.sqrt(1.0 + R[mask4, 2, 2] - R[mask4, 0, 0] - R[mask4, 1, 1]) * 2 - q[mask4, 0] = (R[mask4, 1, 0] - R[mask4, 0, 1]) / s - q[mask4, 1] = (R[mask4, 0, 2] + R[mask4, 2, 0]) / s - q[mask4, 2] = (R[mask4, 1, 2] + R[mask4, 2, 1]) / s - q[mask4, 3] = 0.25 * s - - # Ensure positive w (canonical form) - q = torch.where(q[:, 0:1] < 0, -q, q) - - if not batched: - q = q.squeeze(0) - - return quaternion_normalize(q) - - -def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor: - """ - Convert axis-angle representation to quaternion. - - Parameters - ---------- - axis_angle : torch.Tensor - Axis-angle vector of shape (3,) or (N, 3). - Direction is rotation axis, magnitude is angle in radians. - - Returns - ------- - torch.Tensor - Unit quaternion of shape (4,) or (N, 4). - """ - batched = axis_angle.dim() == 2 - if not batched: - axis_angle = axis_angle.unsqueeze(0) - - angle = axis_angle.norm(dim=-1, keepdim=True).clamp(min=1e-12) - axis = axis_angle / angle - - half_angle = angle / 2 - w = torch.cos(half_angle) - xyz = axis * torch.sin(half_angle) - - q = torch.cat([w, xyz], dim=-1) - - if not batched: - q = q.squeeze(0) - - return q - - -def quaternion_to_axis_angle(q: torch.Tensor) -> torch.Tensor: - """ - Convert quaternion to axis-angle representation. - - Parameters - ---------- - q : torch.Tensor - Unit quaternion of shape (4,) or (N, 4). - - Returns - ------- - torch.Tensor - Axis-angle vector of shape (3,) or (N, 3). - """ - q = quaternion_normalize(q) - - batched = q.dim() == 2 - if not batched: - q = q.unsqueeze(0) - - # Ensure positive w for numerical stability - q = torch.where(q[:, 0:1] < 0, -q, q) - - w = q[:, 0].clamp(-1 + 1e-7, 1 - 1e-7) - xyz = q[:, 1:] - - angle = 2 * torch.acos(w) - sin_half = torch.sqrt(1 - w * w).clamp(min=1e-12) - - axis = xyz / sin_half.unsqueeze(-1) - axis_angle = axis * angle.unsqueeze(-1) - - # Handle near-identity case - near_identity = angle < 1e-6 - axis_angle = torch.where( - near_identity.unsqueeze(-1), 2 * xyz, axis_angle # Small angle approximation - ) - - if not batched: - axis_angle = axis_angle.squeeze(0) - - return axis_angle - - -def quaternion_to_euler_zyz(q: torch.Tensor) -> torch.Tensor: - """ - Convert quaternion to Euler angles (ZYZ convention). - - Parameters - ---------- - q : torch.Tensor - Unit quaternion of shape (4,). - - Returns - ------- - torch.Tensor - Euler angles [alpha, beta, gamma] of shape (3,). - Ranges: alpha in [0, 2pi), beta in [0, pi], gamma in [0, 2pi). - """ - # Convert to matrix first, then extract ZYZ angles - R = quaternion_to_matrix(q) - - # ZYZ convention: R = Rz(alpha) @ Ry(beta) @ Rz(gamma) - # beta = acos(R[2,2]) - # alpha = atan2(R[1,2], R[0,2]) - # gamma = atan2(R[2,1], -R[2,0]) - - beta = torch.acos(R[2, 2].clamp(-1, 1)) - - # Handle gimbal lock cases - if torch.abs(torch.sin(beta)) < 1e-6: - # beta ≈ 0 or pi, gimbal lock - alpha = torch.atan2(R[1, 0], R[0, 0]) - gamma = torch.zeros_like(alpha) - else: - alpha = torch.atan2(R[1, 2], R[0, 2]) - gamma = torch.atan2(R[2, 1], -R[2, 0]) - - # Normalize to [0, 2pi) for alpha and gamma - alpha = torch.remainder(alpha, 2 * torch.pi) - gamma = torch.remainder(gamma, 2 * torch.pi) - - return torch.stack([alpha, beta, gamma]) - - -def euler_zyz_to_quaternion(euler: torch.Tensor) -> torch.Tensor: - """ - Convert Euler angles (ZYZ convention) to quaternion. - - Parameters - ---------- - euler : torch.Tensor - Euler angles [alpha, beta, gamma] of shape (3,). - - Returns - ------- - torch.Tensor - Unit quaternion of shape (4,). - """ - alpha, beta, gamma = euler[0], euler[1], euler[2] - - # ZYZ: q = qz(alpha) * qy(beta) * qz(gamma) - ca, sa = torch.cos(alpha / 2), torch.sin(alpha / 2) - cb, sb = torch.cos(beta / 2), torch.sin(beta / 2) - cg, sg = torch.cos(gamma / 2), torch.sin(gamma / 2) - - # qz(alpha) = [cos(a/2), 0, 0, sin(a/2)] - # qy(beta) = [cos(b/2), 0, sin(b/2), 0] - # qz(gamma) = [cos(g/2), 0, 0, sin(g/2)] - - q_alpha = torch.stack([ca, torch.zeros_like(ca), torch.zeros_like(ca), sa]) - q_beta = torch.stack([cb, torch.zeros_like(cb), sb, torch.zeros_like(cb)]) - q_gamma = torch.stack([cg, torch.zeros_like(cg), torch.zeros_like(cg), sg]) - - return quaternion_multiply(quaternion_multiply(q_alpha, q_beta), q_gamma) - - -# ============================================================================= -# RigidTransform Class -# ============================================================================= - - -class RigidTransform(DeviceMixin, nn.Module): - """ - Rigid body transformation with quaternion-based rotation storage. - - Stores rotation internally as unit quaternion [w, x, y, z] and - translation as 3D vector. Provides methods for various representations - and transformation operations. - - Parameters - ---------- - quaternion : torch.Tensor, optional - Rotation as quaternion [w, x, y, z] of shape (4,). - translation : torch.Tensor, optional - Translation vector of shape (3,). Defaults to zeros. - rotation_matrix : torch.Tensor, optional - Alternative: initialize from rotation matrix of shape (3, 3). - axis_angle : torch.Tensor, optional - Alternative: initialize from axis-angle vector of shape (3,). - - Examples - -------- - :: - - T = RigidTransform.random() - coords = torch.randn(100, 3) - coords_transformed = T.apply(coords) - coords_back = T.inverse().apply(coords_transformed) - """ - - def __init__( - self, - quaternion: Optional[torch.Tensor] = None, - translation: Optional[torch.Tensor] = None, - rotation_matrix: Optional[torch.Tensor] = None, - axis_angle: Optional[torch.Tensor] = None, - ): - super().__init__() - - # Determine dtype and device from inputs (fallback to package defaults) - dtype = get_float_dtype() - device = get_default_device() - - if quaternion is not None: - dtype, device = quaternion.dtype, quaternion.device - elif rotation_matrix is not None: - dtype, device = rotation_matrix.dtype, rotation_matrix.device - elif axis_angle is not None: - dtype, device = axis_angle.dtype, axis_angle.device - elif translation is not None: - dtype, device = translation.dtype, translation.device - - # Convert to quaternion from alternative representations - if quaternion is not None: - q = quaternion_normalize(quaternion) - elif rotation_matrix is not None: - q = matrix_to_quaternion(rotation_matrix) - elif axis_angle is not None: - q = axis_angle_to_quaternion(axis_angle) - else: - # Identity rotation - q = torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=dtype, device=device) - - # Set translation - if translation is not None: - t = translation.to(dtype=dtype, device=device) - else: - t = torch.zeros(3, dtype=dtype, device=device) - - # Register as buffers (not parameters by default) - self.register_buffer("_quaternion", q) - self.register_buffer("_translation", t) - - # ========================================================================= - # Properties for different representations - # ========================================================================= - - @property - def quaternion(self) -> torch.Tensor: - """Get rotation as quaternion [w, x, y, z].""" - return self._quaternion - - @property - def rotation_matrix(self) -> torch.Tensor: - """Get rotation as 3x3 matrix.""" - return quaternion_to_matrix(self._quaternion) - - @property - def axis_angle(self) -> torch.Tensor: - """Get rotation as axis-angle vector.""" - return quaternion_to_axis_angle(self._quaternion) - - @property - def euler_zyz(self) -> torch.Tensor: - """Get rotation as Euler angles (ZYZ convention).""" - return quaternion_to_euler_zyz(self._quaternion) - - @property - def translation(self) -> torch.Tensor: - """Get translation vector.""" - return self._translation - - @property - def dtype(self) -> torch.dtype: - """Get data type.""" - return self._quaternion.dtype - - @property - def device(self) -> torch.device: - """Get device.""" - return self._quaternion.device - - # ========================================================================= - # Transformation methods - # ========================================================================= - - def apply(self, coords: torch.Tensor) -> torch.Tensor: - """ - Apply transformation to coordinates: x' = R @ x + t. - - Parameters - ---------- - coords : torch.Tensor - Coordinates of shape (N, 3) or (3,). - - Returns - ------- - torch.Tensor - Transformed coordinates of same shape. - """ - R = self.rotation_matrix - coords_dtype = coords.dtype - coords = coords.to(dtype=self.dtype) - - if coords.dim() == 1: - result = R @ coords + self._translation - else: - result = coords @ R.T + self._translation - - return result.to(dtype=coords_dtype) - - def apply_rotation_only(self, coords: torch.Tensor) -> torch.Tensor: - """ - Apply rotation only: x' = R @ x. - - Parameters - ---------- - coords : torch.Tensor - Coordinates of shape (N, 3) or (3,). - - Returns - ------- - torch.Tensor - Rotated coordinates of same shape. - """ - R = self.rotation_matrix - coords_dtype = coords.dtype - coords = coords.to(dtype=self.dtype) - - if coords.dim() == 1: - result = R @ coords - else: - result = coords @ R.T - - return result.to(dtype=coords_dtype) - - def forward(self, coords: torch.Tensor) -> torch.Tensor: - """nn.Module forward = apply().""" - return self.apply(coords) - - # ========================================================================= - # Composition and inversion - # ========================================================================= - - def inverse(self) -> "RigidTransform": - """ - Compute inverse transformation. - - For T(x) = R @ x + t, inverse is T^{-1}(x) = R^T @ (x - t). - - Returns - ------- - RigidTransform - Inverse transformation. - """ - q_inv = quaternion_conjugate(self._quaternion) - # t_inv = -R^T @ t = -R_inv @ t - t_inv = -quaternion_rotate(q_inv, self._translation) - return RigidTransform(quaternion=q_inv, translation=t_inv) - - def compose(self, other: "RigidTransform") -> "RigidTransform": - """ - Compose with another transformation: (self @ other)(x) = self(other(x)). - - Parameters - ---------- - other : RigidTransform - Transformation to compose with. - - Returns - ------- - RigidTransform - Composed transformation. - """ - q_new = quaternion_multiply(self._quaternion, other._quaternion) - # t_new = R_self @ t_other + t_self - t_new = ( - quaternion_rotate(self._quaternion, other._translation) + self._translation - ) - return RigidTransform(quaternion=q_new, translation=t_new) - - def __matmul__(self, other: Union["RigidTransform", torch.Tensor]): - """ - Support @ operator for composition or application. - - Parameters - ---------- - other : RigidTransform or torch.Tensor - If RigidTransform, composes transformations. - If tensor, applies transformation to coordinates. - - Returns - ------- - RigidTransform or torch.Tensor - Composed transformation or transformed coordinates. - """ - if isinstance(other, RigidTransform): - return self.compose(other) - return self.apply(other) - - # ========================================================================= - # Factory methods - # ========================================================================= - - @classmethod - def identity( - cls, - device: Union[str, torch.device] = None, - dtype: torch.dtype = None, - ) -> "RigidTransform": - """ - Create identity transformation. - - Parameters - ---------- - device : str or torch.device - Device for tensors. - dtype : torch.dtype - Data type for tensors. - - Returns - ------- - RigidTransform - Identity transformation. - """ - if device is None: - device = get_default_device() - if dtype is None: - dtype = get_float_dtype() - q = torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=dtype, device=device) - t = torch.zeros(3, dtype=dtype, device=device) - return cls(quaternion=q, translation=t) - - @classmethod - def from_matrix( - cls, - R: torch.Tensor, - t: Optional[torch.Tensor] = None, - ) -> "RigidTransform": - """ - Create from rotation matrix and translation. - - Parameters - ---------- - R : torch.Tensor - Rotation matrix of shape (3, 3). - t : torch.Tensor, optional - Translation vector of shape (3,). - - Returns - ------- - RigidTransform - Transformation from matrix representation. - """ - return cls(rotation_matrix=R, translation=t) - - @classmethod - def from_axis_angle( - cls, - axis_angle: torch.Tensor, - t: Optional[torch.Tensor] = None, - ) -> "RigidTransform": - """ - Create from axis-angle and translation. - - Parameters - ---------- - axis_angle : torch.Tensor - Axis-angle vector of shape (3,). - t : torch.Tensor, optional - Translation vector of shape (3,). - - Returns - ------- - RigidTransform - Transformation from axis-angle representation. - """ - return cls(axis_angle=axis_angle, translation=t) - - @classmethod - def from_euler_zyz( - cls, - euler: torch.Tensor, - t: Optional[torch.Tensor] = None, - ) -> "RigidTransform": - """ - Create from Euler angles (ZYZ convention) and translation. - - Parameters - ---------- - euler : torch.Tensor - Euler angles [alpha, beta, gamma] of shape (3,). - t : torch.Tensor, optional - Translation vector of shape (3,). - - Returns - ------- - RigidTransform - Transformation from Euler angle representation. - """ - q = euler_zyz_to_quaternion(euler) - return cls(quaternion=q, translation=t) - - @classmethod - def random( - cls, - device: Union[str, torch.device] = None, - dtype: torch.dtype = None, - translation_scale: float = 0.0, - ) -> "RigidTransform": - """ - Create random transformation with uniform rotation over SO(3). - - Uses Shoemake's quaternion-based uniform sampling. - - Parameters - ---------- - device : str or torch.device - Device for tensors. - dtype : torch.dtype - Data type for tensors. - translation_scale : float - Scale for random translation (0 for no translation). - - Returns - ------- - RigidTransform - Random transformation. - """ - if device is None: - device = get_default_device() - if dtype is None: - dtype = get_float_dtype() - # Shoemake's uniform random quaternion - u1, u2, u3 = torch.rand(3, dtype=dtype, device=device) - - q = torch.stack( - [ - torch.sqrt(1 - u1) * torch.sin(2 * torch.pi * u2), - torch.sqrt(1 - u1) * torch.cos(2 * torch.pi * u2), - torch.sqrt(u1) * torch.sin(2 * torch.pi * u3), - torch.sqrt(u1) * torch.cos(2 * torch.pi * u3), - ] - ) - - # Reorder to [w, x, y, z] format - q = torch.stack([q[3], q[0], q[1], q[2]]) - - if translation_scale > 0: - t = torch.randn(3, dtype=dtype, device=device) * translation_scale - else: - t = torch.zeros(3, dtype=dtype, device=device) - - return cls(quaternion=q, translation=t) - - # ========================================================================= - # Utility methods - # ========================================================================= - - def detach(self) -> "RigidTransform": - """ - Return detached copy (no gradient tracking). - - Returns - ------- - RigidTransform - Detached transformation. - """ - return RigidTransform( - quaternion=self._quaternion.detach(), - translation=self._translation.detach(), - ) - - def clone(self) -> "RigidTransform": - """ - Return deep copy. - - Returns - ------- - RigidTransform - Cloned transformation. - """ - return RigidTransform( - quaternion=self._quaternion.clone(), - translation=self._translation.clone(), - ) - - def __repr__(self) -> str: - q = self._quaternion - t = self._translation - return ( - f"RigidTransform(\n" - f" quaternion=[{q[0]:.4f}, {q[1]:.4f}, {q[2]:.4f}, {q[3]:.4f}],\n" - f" translation=[{t[0]:.4f}, {t[1]:.4f}, {t[2]:.4f}]\n" - f")" - ) diff --git a/torchref/experimental/alignment/translation.py b/torchref/experimental/alignment/translation.py index 4ab1c7e8..97913973 100644 --- a/torchref/experimental/alignment/translation.py +++ b/torchref/experimental/alignment/translation.py @@ -1,277 +1,469 @@ +"""Fast translation search: where in the cell does an oriented model sit? + +A translation shifts phase, ``F(h, t) = F(h) exp(2 pi i h.t)``, so scoring every +``t`` on a grid is a Fourier transform rather than a scan. The Crowther-Blow +form used here accumulates the pair coefficients +``sum_h c(h) G_i*(h) G_j(h)`` onto a reciprocal grid at +``(h R_j - h R_i) mod G`` and takes one inverse FFT, which replaces ``G^3`` grid +evaluations with a single transform. + +Both sides of that sum are **normalised**. The observed side is the rotation +search's own LERF1 intensity, ``cw (E_obs^2 - 1) w sigma_A^2``, built from the +run's one Wilson fit; the calculated side is the oriented model's transform +divided by its own Wilson curve, so ``<|E_calc(h, t)|^2> = 1`` per shell for +every candidate. The score is then a covariance of two normalised intensities +and every resolution shell carries the weight the model error gives it. The +previous form divided raw ``|F_calc|^2`` by its own sum, which is not a +correlation: on 2DQ6 it was 0.665 at a position 41 A from the deposited pose +and 0.350 at the pose itself, and the search followed it there. + +The grid is sized to the resolution of the translation set, one FFT per +candidate, and the best few peaks are re-scored with the full Rice/Woolfson +likelihood at fixed ``sigma_A``. That likelihood is also what ranks the +candidates against each other. + +The observed side is prepared **once** per run, by :class:`TranslationObs`, +and reused for every orientation. Normalisation, weighting and model error are +properties of the observations and the search model, which do not change when +the model moves. """ -Fast FFT-based translation search for molecular replacement. -Translation t shifts phase: F(hkl, t) = F(hkl) * exp(2*pi*i * hkl.t) -Correlation: C(t) = IFFT{ conj(F_obs) * F_calc } +from __future__ import annotations -This module provides efficient FFT-based translation search that finds the -optimal translation to position a model after rotation has been determined. - -Experimental / unstable API: part of ``torchref.experimental.alignment``, -the opt-in ball-harmonic MR engine. The production MR entry point is -``torchref.alignment`` (the consolidated FRF engine). Signatures and behavior -may change without notice. -""" +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple, TYPE_CHECKING import numpy as np import torch -from dataclasses import dataclass -from typing import List, Optional, Tuple - -@dataclass -class TranslationPeak: - """ - Translation search peak. +from torchref.base.targets.xray_likelihoods import rice_per_refl +from torchref.config import get_complex_dtype, get_default_device, get_float_dtype +from torchref.scaling import WilsonNormaliser +from torchref.scaling.weighting import (inverse_variance_weight, + normalise_weight, snr_from_amplitude) +from torchref.symmetry.symmetry import find_fft_friendly_size - Attributes - ---------- - translation : np.ndarray - Fractional coordinates (3,). - score : float - Correlation score. - sigma : float - Z-score above mean. - """ - translation: np.ndarray - score: float - sigma: float +from .frf.preprocessing import build_lerf1_intensity, eterm_sigma_a +if TYPE_CHECKING: # pragma: no cover - typing only + from ...model.model_ft import ModelFT -def fft_translation_search( - F_obs: np.ndarray, - F_calc: np.ndarray, - hkl: np.ndarray, - grid_shape: Optional[Tuple[int, int, int]] = None, - n_peaks: int = 10, - cluster_radius: float = 0.05, -) -> Tuple[np.ndarray, np.ndarray, List[TranslationPeak]]: - """ - FFT-based translation search (vectorized). - The translation function is: - TF(t) = Re{ IFFT{ conj(F_obs) * F_calc } } +#: Chebyshev order of the Wilson fit. Matches the rotation function's +#: ``frf.api.WILSON_N_COEFF``: the two stages score the same observations and a +#: different order on each would be two normalisations again. +WILSON_N_COEFF = 6 - This finds translation t such that F_calc shifted by t best matches F_obs. +#: Largest FFT grid per axis. 256^3 complex64 is 134 MB, which bounds the +#: translation map for an uncut high-resolution set on a long cell; at the +#: default window the grid never reaches it. +MAX_GRID_PER_AXIS = 256 - Parameters - ---------- - F_obs : np.ndarray - Observed structure factor amplitudes (or complex), shape (N,). Real - (amplitude-only) input is treated as zero-phase, so the ``conj()`` in - the translation function is a no-op on it. - F_calc : np.ndarray - Calculated structure factors (complex), shape (N,). - hkl : np.ndarray - Miller indices, shape (N, 3). - grid_shape : tuple, optional - (Nx, Ny, Nz) grid size for FFT. If None, auto-computed from HKL range. - n_peaks : int - Number of peaks to return. - cluster_radius : float - Minimum fractional distance between peaks for clustering. - Returns - ------- - correlation_map : np.ndarray - Full translation function, shape grid_shape. - best_translation : np.ndarray - Best translation in fractional coordinates, shape (3,). - peaks : list - Top peaks as TranslationPeak objects. - - Examples - -------- - :: - - import numpy as np - from torchref.experimental.alignment.translation import fft_translation_search - - # Known translation test - hkl = np.array([[1,0,0], [0,1,0], [1,1,0], [0,0,1]]) - F_obs = np.array([1.0, 1.0, 1.0, 1.0]) - F_calc = np.exp(2j * np.pi * hkl @ [0.25, 0.0, 0.0]) - _, best, peaks = fft_translation_search(F_obs, F_calc, hkl) - print(f'Recovered: {best}') # Should be ~[0.25, 0, 0] - """ - # Auto grid shape from HKL range - if grid_shape is None: - hkl_abs = np.abs(hkl).astype(int) - grid_shape = tuple(2 * (hkl_abs[:, i].max() + 1) for i in range(3)) +@dataclass +class TranslationObs: + """The observed side of a translation search, normalised and weighted once. - Nx, Ny, Nz = grid_shape - product_grid = np.zeros((Nx, Ny, Nz), dtype=np.complex128) + Everything here is a property of the observations and of the search model's + expected error, so none of it changes when the model rotates or moves. - # Standard translation function: TF(t) = Re{ IFFT{ conj(F_obs) * F_calc } } - # This finds t such that F_calc(t) = F_calc * exp(2*pi*i*hkl.t) matches F_obs - product = np.conj(F_obs) * F_calc + Attributes + ---------- + F_obs, hkl, s_mag, centric, eps + The masked observations and their crystallographic bookkeeping. + E_obs : torch.Tensor + ``F / sqrt(eps Sigma(s))``, with ``Sigma`` the shared Wilson fit, so + `` = 1`` as an identity of that fit. + weight : torch.Tensor + Mean-1 inverse-variance weight, from measurement error and model error + in one denominator. Uniform when the data carry no sigmas. + sigma_a : torch.Tensor + The Luzzati fall-off ``exp(-(2 pi^2 / 3) s^2 vrms^2)`` for the search + model's expected coordinate error -- the same term the rotation + function weights with. It is the ``D`` of the likelihood and the + calc-side weight of the fast search. A prior, not a fit: nothing can be + fitted before the model is placed. + coeff : torch.Tensor + The fast search's per-reflection coefficient, + ``cw (E_obs^2 - 1) weight sigma_A^2`` -- the rotation function's LERF1 + intensity with its calc-side ``sigma_A^2`` folded in. Centred, so a + placement that puts calculated intensity everywhere gains nothing. + fit : WilsonNormaliser + Kept, not discarded. Anything comparing an observed curve against a + calculated one needs the curve itself. + """ - # Place at HKL positions using vectorized add.at for accumulation - hkl_int = hkl.astype(int) - h_idx = hkl_int[:, 0] % Nx - k_idx = hkl_int[:, 1] % Ny - l_idx = hkl_int[:, 2] % Nz + F_obs: torch.Tensor + hkl: torch.Tensor + s_mag: torch.Tensor + centric: torch.Tensor + eps: torch.Tensor + E_obs: torch.Tensor + weight: torch.Tensor + sigma_a: torch.Tensor + coeff: torch.Tensor + fit: "WilsonNormaliser" + + @classmethod + def build( + cls, + F_obs: torch.Tensor, + hkl: torch.Tensor, + spacegroup, + real_cell, + *, + sig_F: Optional[torch.Tensor] = None, + delta_vrms_A: float = 1.0, + n_coeff: int = WILSON_N_COEFF, + device=None, + ) -> "TranslationObs": + """Normalise and weight one set of observations. + + Parameters + ---------- + F_obs : torch.Tensor + ``(N,)`` observed amplitudes; complex input is coerced to ``|.|``. + hkl : torch.Tensor + ``(N, 3)`` integer Miller indices, matching ``F_obs`` row for row. + spacegroup, real_cell + Supply multiplicity, centricity and the reciprocal basis. + sig_F : torch.Tensor, optional + ``(N,)`` measurement errors. Without them the weight is uniform, + which is the honest fallback: the varying part of the weight *is* + the measurement term, and inventing one would be worse than not + having it. + delta_vrms_A : float + R.m.s. coordinate error of the search model, which sets the model + half of the variance budget and the likelihood's ``sigma_A``. + """ + dev = get_default_device() if device is None else device + real = get_float_dtype() + # Cast and move in one `.to`, and take `abs()` before either. Chaining + # them the other way round -- `.to(dev)` and then `.to(real)` -- puts + # the caller's width on the device first, which throws for a float64 + # input on a backend that has none, and double observations are a + # perfectly ordinary thing to be handed. + F = F_obs.detach() + F = (F.abs() if F.is_complex() else F).to(device=dev, dtype=real) + hkl_i = hkl.detach().to(dev) + + rec_basis = real_cell.reciprocal_basis_matrix.to(device=dev, dtype=real) + s_mag = (hkl_i.to(real) @ rec_basis).norm(dim=-1) + + hkl_l = hkl_i.round().to(torch.int64) # dtype-ok: Miller indices are integers + # friedel=False: Wilson's = eps*Sigma counts the operations mapping + # h to itself, which add coherently and set the mean. The Friedel-folded + # branch changes the distribution instead, and that is centricity -- + # which enters separately, as the Gamma shape. + eps = spacegroup.epsilon(hkl_l, friedel=False).to(real).clamp(min=1.0) + centric = spacegroup.is_centric(hkl_l).to(torch.bool) + + fit = WilsonNormaliser( + F * F, s_mag, eps=eps, centric=centric, n_coeff=n_coeff, + ) + sigma_a = eterm_sigma_a(s_mag, float(delta_vrms_A)).to(real) + + if sig_F is None: + weight = torch.ones_like(F) + else: + sig = sig_F.detach().to(device=dev, dtype=real).abs() + weight = normalise_weight(inverse_variance_weight( + snr_from_amplitude(F, sig), sigma_a, eps=eps, + )) + + E_obs = fit.E.to(real) + coeff = build_lerf1_intensity(E_obs, centric, weight=weight) * sigma_a ** 2 + + return cls( + F_obs=F, hkl=hkl_i, s_mag=s_mag, centric=centric, eps=eps, + E_obs=E_obs, weight=weight, sigma_a=sigma_a, coeff=coeff, fit=fit, + ) - np.add.at(product_grid, (h_idx, k_idx, l_idx), product) - # IFFT gives correlation at all translations - correlation_map = np.fft.ifftn(product_grid).real +@dataclass +class CandidateTransform: + """One oriented model's transform at the symmetry-rotated indices. - # Find peaks - peaks = find_translation_peaks(correlation_map, n_peaks, cluster_radius) - best = peaks[0].translation if peaks else np.zeros(3) + Attributes + ---------- + G : torch.Tensor + ``(S, N)`` complex. ``G_i(h) = F_p1(h R_i) exp(2 pi i h.t_i) / norm(h)``, + so ``E_calc(h, t) = |sum_i G_i(h) exp(2 pi i (h R_i).t)|`` is the + **normalised** calculated amplitude: `` = 1`` per shell. + h_R : torch.Tensor + ``(S, N, 3)`` the rotated indices ``h R_i``. + norm : torch.Tensor + ``(N,)`` ``sqrt(eps n_ops Sigma_P(s))``: the raw amplitude is + ``E_calc * norm``. + """ - return correlation_map, best, peaks + G: torch.Tensor + h_R: torch.Tensor + norm: torch.Tensor + + def e_calc(self, t: torch.Tensor) -> torch.Tensor: + """``E_calc(h, t)`` for ``t`` of shape ``(3,)`` or ``(K, 3)``: ``(N,)`` or ``(K, N)``.""" + single = t.ndim == 1 + # One `.to`: a translation handed in as double -- a fractional vector + # from host-side algebra usually is -- must not be put on the device at + # its own width first. + tt = t.reshape(-1, 3).to(device=self.h_R.device, dtype=self.h_R.dtype) + phase_arg = torch.einsum("ind,kd->kin", self.h_R, tt) + phase = torch.exp((2j * math.pi) * phase_arg.to(self.G.dtype)) + E = (self.G.unsqueeze(0) * phase).sum(dim=1).abs() + return E[0] if single else E + + +def prepare_candidate( + model_p1: "ModelFT", + obs: TranslationObs, + spacegroup, + real_cell, +) -> CandidateTransform: + """Evaluate one orientation's transform and normalise it. + + ``model_p1`` is an ordinary :class:`~torchref.model.ModelFT` in P1 with the + crystal's cell, already in the candidate orientation; its grid is whatever + its ``max_res`` implies. The only per-candidate model evaluation: ``F_p1`` + at all ``S x N`` rotated indices in one call. The result is moved to the + configured default device, whatever device the model sits on, so the + translation stage never straddles two devices. The normalising curve ``Sigma_P(s)`` is the same + Wilson fit the observed side uses, on the same abscissa, fitted to the + transform's mean intensity over the ``S`` copies -- which is what the crystal + sum averages to over a shell, since the cross terms between symmetry copies + have zero mean over ``h``. The crystal's ``<|F_calc|^2>`` is then + ``eps n_ops Sigma_P``, and dividing by it is what puts every candidate's + ``E_calc`` on one footing with ``E_obs`` and with each other. + """ + device = get_default_device() + real = get_float_dtype() + cplx = get_complex_dtype() + + hkl = obs.hkl.to(device=device, dtype=real) + sym_R = spacegroup.matrices.detach().to(device=device, dtype=real) + sym_t = spacegroup.translations.detach().to(device=device, dtype=real) + S = int(sym_R.shape[0]) + N = int(hkl.shape[0]) + + # h_R[i, n, d] = sum_e hkl[n, e] sym_R[i, e, d]: the h.S convention. + h_R = torch.einsum("ne,ied->ind", hkl, sym_R) + phase = torch.exp((2j * math.pi) * torch.einsum("ne,ie->in", hkl, sym_t).to(cplx)) + hkl_SN = h_R.reshape(-1, 3).round().to(torch.int64).to(model_p1.xyz().device) # dtype-ok: Miller indices are integers + with torch.no_grad(): + F_all = model_p1(hkl_SN).to(device).reshape(S, N).to(cplx) + G_raw = F_all * phase + + I_P = (G_raw.abs() ** 2).mean(dim=0).to(real) + s_mag = obs.s_mag.to(device=device, dtype=real) + fit_P = WilsonNormaliser( + I_P, s_mag, n_coeff=WILSON_N_COEFF, + s_lo=float(s_mag.min()), s_hi=float(s_mag.max()), + ) + Sigma_c = S * fit_P.evaluate(s_mag).to(real) + norm = (obs.eps.to(device=device, dtype=real) * Sigma_c).clamp(min=1e-30).sqrt() + return CandidateTransform(G=G_raw / norm.to(cplx), h_R=h_R, norm=norm) -def find_translation_peaks( - correlation_map: np.ndarray, - n_peaks: int = 10, - cluster_radius: float = 0.05, -) -> List[TranslationPeak]: - """ - Extract and cluster peaks from translation function. +@dataclass +class TranslationPeak: + """A peak of the fast translation function. - Parameters + Attributes ---------- - correlation_map : np.ndarray - Translation function values, shape (Nx, Ny, Nz). - n_peaks : int - Maximum number of peaks to return. - cluster_radius : float - Minimum fractional distance between peaks (periodic). - - Returns - ------- - peaks : list - List of TranslationPeak objects sorted by score. + translation : np.ndarray + Fractional coordinates (3,), refined to sub-grid precision. + score : float + The fast search's score at the grid maximum. + sigma : float + Standard deviations above the map mean. """ - Nx, Ny, Nz = correlation_map.shape - mean_val = correlation_map.mean() - std_val = correlation_map.std() + translation: np.ndarray + score: float + sigma: float - if std_val < 1e-10: - return [] - flat = correlation_map.flatten() - sorted_idx = np.argsort(flat)[::-1] +def _grid_sizes(real_cell, grid_spacing_A: float) -> Tuple[int, int, int]: + """FFT-friendly grid, at most ``grid_spacing_A`` apart along each axis.""" + sizes = [] + for length in (real_cell.a, real_cell.b, real_cell.c): + n = int(math.ceil(float(length) / float(grid_spacing_A))) + n = find_fft_friendly_size(max(n, 4)) + sizes.append(min(n, MAX_GRID_PER_AXIS)) + return sizes[0], sizes[1], sizes[2] - peaks = [] - used = [] - for idx in sorted_idx: - if len(peaks) >= n_peaks: - break +def _parabolic_offset(fm: float, f0: float, fp: float) -> float: + """Sub-grid offset of a maximum from its three samples, in grid units.""" + denom = fm - 2.0 * f0 + fp + if denom >= 0.0: + return 0.0 + return float(min(0.5, max(-0.5, 0.5 * (fm - fp) / denom))) - pos_3d = np.unravel_index(idx, correlation_map.shape) - trans = np.array([pos_3d[0] / Nx, pos_3d[1] / Ny, pos_3d[2] / Nz]) - score = flat[idx] - sigma = (score - mean_val) / std_val - # Check clustering - skip if too close to existing peak +def _find_peaks( + score: torch.Tensor, + n_peaks: int, + radii_frac: Tuple[float, float, float], +) -> List[TranslationPeak]: + """Greedy non-maximum suppression on the periodic map, then sub-grid refinement.""" + nx, ny, nz = score.shape + flat = score.reshape(-1) + mean = float(flat.mean()) + std = float(flat.std().clamp(min=1e-30)) + n_take = min(flat.numel(), max(50, 20 * n_peaks)) + vals, idx = torch.topk(flat, n_take) + vals = vals.cpu().numpy() + idx = idx.cpu().numpy() + grid = np.array([nx, ny, nz], dtype=np.float64) + radii = np.asarray(radii_frac, dtype=np.float64) + score_np = score.cpu().numpy() + + kept: List[TranslationPeak] = [] + kept_t: List[np.ndarray] = [] + for v, i in zip(vals, idx): + ijk = np.array(np.unravel_index(int(i), (nx, ny, nz)), dtype=np.int64) + t_grid = ijk / grid is_new = True - for prev in used: - diff = np.abs(trans - prev) - diff = np.minimum(diff, 1 - diff) # Periodic boundary - if np.linalg.norm(diff) < cluster_radius: + for prev in kept_t: + d = np.abs(t_grid - prev) + d = np.minimum(d, 1.0 - d) + if np.all(d < radii): is_new = False break - - if is_new: - peaks.append(TranslationPeak(trans, score, sigma)) - used.append(trans) - - return peaks - - -def fft_translation_search_torch( - F_obs: torch.Tensor, - F_calc: torch.Tensor, - hkl: torch.Tensor, - **kwargs, -) -> Tuple[np.ndarray, np.ndarray, List[TranslationPeak]]: - """ - Torch wrapper for fft_translation_search. - - Parameters - ---------- - F_obs : torch.Tensor - Observed structure factor amplitudes (or complex). - F_calc : torch.Tensor - Calculated structure factors (complex). - hkl : torch.Tensor - Miller indices. - **kwargs - Additional arguments passed to fft_translation_search. - - Returns - ------- - correlation_map : np.ndarray - Full translation function. - best_translation : np.ndarray - Best translation in fractional coordinates. - peaks : list - Top peaks as TranslationPeak objects. - """ - return fft_translation_search( - F_obs.detach().cpu().numpy(), - F_calc.detach().cpu().numpy(), - hkl.detach().cpu().numpy(), - **kwargs, - ) - - -def apply_translation_to_fcalc( - F_calc: np.ndarray, - hkl: np.ndarray, - translation_frac: np.ndarray, -) -> np.ndarray: - """ - Apply translation phase shift to calculated structure factors. - - F(hkl, t) = F(hkl) * exp(2*pi*i * hkl.t) + if not is_new: + continue + # Parabolic refinement along each axis from the periodic neighbours. + offs = np.zeros(3) + for d, n in enumerate((nx, ny, nz)): + lo = ijk.copy(); lo[d] = (ijk[d] - 1) % n + hi = ijk.copy(); hi[d] = (ijk[d] + 1) % n + offs[d] = _parabolic_offset( + float(score_np[tuple(lo)]), float(v), float(score_np[tuple(hi)]), + ) + kept.append(TranslationPeak( + translation=(ijk + offs) / grid, score=float(v), + sigma=(float(v) - mean) / std, + )) + kept_t.append(t_grid) + if len(kept) >= n_peaks: + break + return kept + + +def fast_translation_function( + obs: TranslationObs, + cand: CandidateTransform, + real_cell, + *, + grid_spacing_A: float, + n_peaks: int = 3, + cluster_radius_A: float = 4.0, +) -> Tuple[torch.Tensor, List[TranslationPeak]]: + """The Crowther-Blow map of ``sum_h coeff(h) |E_calc(h, t)|^2`` and its peaks. + + ``coeff`` is :attr:`TranslationObs.coeff` and ``E_calc`` is normalised per + candidate by :func:`prepare_candidate`, so the map is the covariance of + two unit-mean intensities weighted by the model's expected reliability -- + the rotation function's own score equation, for translations. Expanding + ``|sum_i G_i exp(2 pi i (h R_i).t)|^2`` gives pair terms at frequency + ``h R_j - h R_i``; accumulating them onto a reciprocal grid and inverting + evaluates every grid translation in one FFT. Parameters ---------- - F_calc : np.ndarray - Calculated structure factors (complex), shape (N,). - hkl : np.ndarray - Miller indices, shape (N, 3). - translation_frac : np.ndarray - Translation in fractional coordinates, shape (3,). + grid_spacing_A : float + Target spacing of the translation grid along each axis. A third of the + translation set's resolution samples the peak densely enough for the + parabolic refinement to land within a fraction of a grid step. + n_peaks : int + How many distinct peaks to return, best first. + cluster_radius_A : float + Peaks closer than this (per axis, periodic) are one peak. Returns ------- - F_calc_shifted : np.ndarray - Phase-shifted structure factors, shape (N,). + score : torch.Tensor + The ``(nx, ny, nz)`` map, fractional grid ``t = (i/nx, j/ny, k/nz)``. + peaks : list of TranslationPeak """ - phase_shift = 2 * np.pi * (hkl @ translation_frac) - return F_calc * np.exp(1j * phase_shift) - - -def apply_translation_to_fcalc_torch( - F_calc: torch.Tensor, - hkl: torch.Tensor, - translation_frac: torch.Tensor, + device = get_default_device() + real = get_float_dtype() + cplx = get_complex_dtype() + + nx, ny, nz = _grid_sizes(real_cell, grid_spacing_A) + G = cand.G.to(device=device, dtype=cplx) + S, N = G.shape + coeff = obs.coeff.to(device=device, dtype=cplx) + h_R_int = cand.h_R.round().to(torch.int64) # dtype-ok: Miller indices are integers + + # The pair (j, i) is the conjugate of (i, j) at -dh, so the map is twice + # the real part of the upper triangle's transform plus the diagonal, which + # carries no t and is a constant. Half the scatter, which is the cost here. + W = torch.zeros(nx * ny * nz, dtype=cplx, device=device) + for i in range(S - 1): + pair = G[i].conj().view(1, -1) * G[i + 1:] # (S-i-1, N) + dh = h_R_int[i + 1:] - h_R_int[i:i + 1] # (S-i-1, N, 3) + flat = ((dh[..., 0] % nx) * ny + (dh[..., 1] % ny)) * nz + (dh[..., 2] % nz) + W.index_add_(0, flat.reshape(-1), (coeff.view(1, -1) * pair).reshape(-1)) + diag = (obs.coeff.to(device=device, dtype=real) * (G.abs() ** 2).sum(dim=0).to(real)).sum() + score = (2.0 * torch.fft.ifftn(W.view(nx, ny, nz), dim=(0, 1, 2)).real + * float(nx * ny * nz)).to(real) + diag + + radii = tuple(float(cluster_radius_A) / float(L) + for L in (real_cell.a, real_cell.b, real_cell.c)) + peaks = _find_peaks(score, n_peaks, radii) + return score, peaks + + +def translation_score_at(obs: TranslationObs, cand: CandidateTransform, + t: torch.Tensor) -> float: + """The fast search's score at one translation, without the FFT.""" + E2 = cand.e_calc(t) ** 2 + return float((obs.coeff.to(E2.device).to(E2.dtype) * E2).sum()) + + +def llg_at_translations( + obs: TranslationObs, + cand: CandidateTransform, + t_candidates: torch.Tensor, ) -> torch.Tensor: - """ - Apply translation phase shift to calculated structure factors (PyTorch). - - F(hkl, t) = F(hkl) * exp(2*pi*i * hkl.t) + """Rice/Woolfson log-likelihood gain at each of ``K`` translations. - Parameters - ---------- - F_calc : torch.Tensor - Calculated structure factors (complex). - hkl : torch.Tensor - Miller indices. - translation_frac : torch.Tensor - Translation in fractional coordinates. + ``LLG(t) = sum_h [LL(E_obs; sigma_A E_calc(h, t), 1 - sigma_A^2) + - LL(E_obs; 0, 1)]`` with the complex-variance convention of + :func:`~torchref.base.targets.xray_likelihoods.rice_per_refl`, which + derives the centric case from the same ``Sigma``. ``sigma_A`` is the + Luzzati prior carried by ``obs`` -- the same for every candidate, so the + values are comparable across orientations as well as across translations, + and no candidate is scored against a likelihood tuned to itself. - Returns - ------- - F_calc_shifted : torch.Tensor - Phase-shifted structure factors. + Returns ``(K,)``. + """ + E_calc = cand.e_calc(t_candidates) # (K, N) + K, N = E_calc.shape + dev, real = E_calc.device, E_calc.dtype + E_obs = obs.E_obs.to(device=dev, dtype=real).view(1, N).expand(K, N) + D = obs.sigma_a.to(device=dev, dtype=real).view(1, N) + Sigma = (1.0 - D * D).clamp(min=1e-3).expand(K, N) + cent = obs.centric.to(dev).view(1, N).expand(K, N) + ll = -rice_per_refl(E_obs, D * E_calc, Sigma, cent) # (K, N) + ll_wil = -rice_per_refl( + E_obs[0], torch.zeros(N, dtype=real, device=dev), + torch.ones(N, dtype=real, device=dev), cent[0], + ).sum() + return ll.sum(dim=1) - ll_wil + + +def analytic_r_at(obs: TranslationObs, cand: CandidateTransform, + t: torch.Tensor) -> float: + """``R = sum ||F_obs| - k |F_calc(t)|| / sum |F_obs|`` with one global scale. + + On raw amplitudes, because that is what a crystallographer reads; not the + number a full Scaler would return, since there is no bulk solvent and no + B-factor scaling behind ``k``. """ - phase_shift = 2 * torch.pi * (hkl.to(translation_frac.dtype) @ translation_frac) - return F_calc * torch.exp(1j * phase_shift) + F_c = cand.e_calc(t) * cand.norm.to(cand.G.device) + F_o = obs.F_obs.to(F_c.device).to(F_c.dtype) + k = (F_o * F_c).sum() / (F_c * F_c).sum().clamp(min=1e-30) + return float((F_o - k * F_c).abs().sum() / F_o.sum().clamp(min=1e-30)) diff --git a/torchref/experimental/ensemble/ensemble_amber_kl.py b/torchref/experimental/ensemble/ensemble_amber_kl.py index 56a5aca6..ebb5cfc3 100644 --- a/torchref/experimental/ensemble/ensemble_amber_kl.py +++ b/torchref/experimental/ensemble/ensemble_amber_kl.py @@ -136,7 +136,7 @@ def __init__( self.register_buffer( "_member_atom_idx", torch.as_tensor( - atom_idx_np, dtype=torch.long, device=self._model.device + atom_idx_np, dtype=torch.long, device=self._model.device # dtype-ok: atom index tensor for indexing; PyTorch requires int64 ), ) else: diff --git a/torchref/experimental/ensemble/ensemble_model.py b/torchref/experimental/ensemble/ensemble_model.py index 07896f9e..de52ce11 100644 --- a/torchref/experimental/ensemble/ensemble_model.py +++ b/torchref/experimental/ensemble/ensemble_model.py @@ -38,7 +38,7 @@ from __future__ import annotations -from typing import Optional +from typing import Optional, Tuple import numpy as np import pandas as pd @@ -340,10 +340,16 @@ def __init__( verbose: int = 1, device=None, strip_H: bool = True, + # An ensemble's atom set is the replicated single copy its factories build, and + # _finalize_ensemble reshapes by n_atoms_per_member, so generating hydrogens on + # load would invalidate that. Off by default here, unlike on the base class. + add_hydrogens: bool = False, max_res: float = 1.0, - gridsize: Optional[int] = None, + gridsize: Optional[Tuple[int, int, int]] = None, wavelength: float = 1.0, anomalous_threshold: float = 0.5, + cif_path=None, + hydrogens_in_xray: bool = True, ): if dtype_float is None: dtype_float = get_float_dtype() @@ -354,10 +360,13 @@ def __init__( verbose=verbose, device=device, strip_H=strip_H, + add_hydrogens=add_hydrogens, max_res=max_res, gridsize=gridsize, wavelength=wavelength, anomalous_threshold=anomalous_threshold, + cif_path=cif_path, + hydrogens_in_xray=hydrogens_in_xray, ) # Filled in by ``_finalize_ensemble`` after ``load`` returns. self.n_members: int = 0 @@ -452,6 +461,9 @@ def from_single( model = cls( verbose=verbose, device=device, strip_H=False, # already stripped + # The replicated table is the atom set; _finalize_ensemble reshapes by + # n_atoms_per_member, so generating hydrogens here would invalidate it. + add_hydrogens=False, max_res=max_res, **modelft_kwargs, ) @@ -538,6 +550,9 @@ def from_multimodel_pdb( model = cls( verbose=verbose, device=device, strip_H=False, + # See from_single: the replicated table is the atom set, and + # _finalize_ensemble reshapes by n_atoms_per_member. + add_hydrogens=False, max_res=max_res, **modelft_kwargs, ) @@ -752,6 +767,8 @@ def enable_low_rank(self, K: int) -> float: with torch.no_grad(): flat = self.xyz().detach() # (N*n_atoms, 3) + # dtype-ok: SVD seeding in float64 for numerical stability. Caveat: no + # .cpu() first, so this errors on MPS. X = flat.reshape(N, n_atoms * 3).to(torch.float64) mu = X.mean(dim=0) # (D,) Xc = X - mu.unsqueeze(0) diff --git a/torchref/experimental/ensemble/ensemble_refinement.py b/torchref/experimental/ensemble/ensemble_refinement.py index eee55188..a36563e7 100644 --- a/torchref/experimental/ensemble/ensemble_refinement.py +++ b/torchref/experimental/ensemble/ensemble_refinement.py @@ -488,7 +488,7 @@ def __init__( ) self.model.cell = self.reflection_data.cell self.model.spacegroup = self.reflection_data.spacegroup - self.model.setup_grid(max_res=self.max_res) + self.model.max_res = self.max_res # Rebuild the scaler against the ensemble model. self.scaler = Scaler( diff --git a/torchref/experimental/ensemble/pca_model.py b/torchref/experimental/ensemble/pca_model.py index 4b092e47..00896cb0 100644 --- a/torchref/experimental/ensemble/pca_model.py +++ b/torchref/experimental/ensemble/pca_model.py @@ -92,6 +92,8 @@ def from_ensemble( member matrix. ``K`` defaults to ``N-1`` (complete reparameterization).""" N = int(n_members) with torch.no_grad(): + # dtype-ok: SVD seeding in float64 for numerical stability; recast to + # xyz_flat.dtype below (line 107). Caveat: no .cpu(), so errors on MPS. X = xyz_flat.detach().reshape(N, n_atoms * 3).to(torch.float64) mu = X.mean(dim=0) Xc = X - mu.unsqueeze(0) diff --git a/torchref/experimental/ensemble/quasi_crystal_amber.py b/torchref/experimental/ensemble/quasi_crystal_amber.py index eddf1423..ee9986ca 100644 --- a/torchref/experimental/ensemble/quasi_crystal_amber.py +++ b/torchref/experimental/ensemble/quasi_crystal_amber.py @@ -31,8 +31,7 @@ - the antechamber pipeline + GAFF2 setup for non-standard residues; - the template OpenMM ``System`` (single-molecule, AMBER14 / GAFF2); -- the H virtual-site frame tables (``_build_h_attachment``) and the shared - local-frame placement (``_place_hydrogens_local_frame``); +- a complete atom map including the model's hydrogens; - the autograd Function ``_OpenMMAMBERFunction``. This target then replicates the template ``System`` into the symmetry-expanded @@ -44,13 +43,14 @@ 1. replicate the System into a supercell with :func:`_replicate_to_supercell_system`; 2. build a new ``Context`` on the supercell (CUDA > OpenCL > CPU); -3. tile the template's atom map + H-attachment indices per member. +3. tile the template's complete atom map per member. Forward reads ``ensemble.xyz_per_member``, applies the supercell layout's sym+tile -transform, scatters into the unified OpenMM position tensor (heavy via -``_compose_full_omm_xyz``-style scatter; H via the tiled local-frame -placement), and calls the same ``_OpenMMAMBERFunction.apply``. +transform, scatters all model atoms into the unified OpenMM position tensor, +and calls ``_OpenMMAMBERFunction.apply``. Hydrogen positions come from TorchRef. +Construction leaves coordinates unchanged unless ``relax_on_init=True`` is +explicitly requested. """ from __future__ import annotations @@ -63,7 +63,6 @@ from torchref.experimental.targets.amber_target import ( AmberTarget, _OpenMMAMBERFunction, - _place_hydrogens_local_frame, ) from .ensemble_model import build_single_copy_model from .supercell import SupercellLayout, _replicate_to_supercell_system @@ -176,6 +175,9 @@ class QuasiCrystalAmberTarget(AmberTarget): charge_method : str antechamber charge method ('gas' or 'bcc'). Default 'gas' (fast, no QM); matches the ensemble setup. + relax_on_init : bool, default False + If True, explicitly minimize with OpenMM and write relaxed positions + back to the ensemble. Leave False to keep TorchRef's initial coordinates. verbose : int Verbosity (0 = silent, 1 = setup messages). @@ -204,7 +206,7 @@ def __init__( gaff2_files: Optional[Dict[str, Tuple[str, str]]] = None, charge_method: str = "gas", drop_special_position_threshold_ang: float = 0.0, - relax_on_init: bool = True, + relax_on_init: bool = False, relax_max_iterations: int = 200, force_clamp: float = 10000.0, verbose: int = 0, @@ -291,7 +293,7 @@ def __init__( # As an AmberTarget subclass we run the full antechamber + ForceField # pipeline against a genuine single-conformation Model (the ensemble's # ``_pdb_single`` restricted to non-special-position atoms). This - # populates self._system / _pos_buf / _model_to_omm / _h_* for ONE + # populates self._system / _pos_buf / _model_to_omm for ONE # member; we replicate them into the supercell below. ``_model`` stays # the ensemble (its per-member coords drive forward()); the # single-molecule context the base builds is replaced by the supercell @@ -322,9 +324,7 @@ def __init__( template_map = np.asarray(self._model_to_omm, dtype=np.int64) self._template_model_to_omm = template_map # Index pairs: model atom `src_model_idx[k]` lives in OMM slot - # `dst_omm_idx[k]` (single-member, in [0, n_omm_per_member)). Atoms - # with template_map == -1 (waters, OXT, ligands tleap regenerated) - # are excluded — their OMM slot keeps the construction-time position. + # `dst_omm_idx[k]` (single-member, in [0, n_omm_per_member)). valid_mask = template_map >= 0 self._src_model_idx_np = np.where(valid_mask)[0].astype(np.int64) self._dst_omm_idx_np = template_map[valid_mask].astype(np.int64) @@ -349,29 +349,6 @@ def __init__( self._n_omm_total = int(supercell_pos_nm.shape[0]) assert self._n_omm_total == N * self._n_omm_per_member - # --- H attachment, template arrays (numpy, into [0, n_omm_per_member)). - # The tiling per member is deferred to the forward path: each H index - # gets ``+ m · n_omm_per_member`` added per member m. - # - # AmberTarget marks rigid-fallback Hs (no valid local frame) with - # sentinel ``-1`` in ``_h_n1_idx`` / ``_h_n2_idx``. The corresponding - # ``_h_frame_valid`` row is False, so the local-frame branch never - # uses these indices. But ``index_select`` still evaluates the lookup - # and errors on negative indices, so clamp the sentinels to 0 — a - # safe in-bounds dummy whose result is then masked away by the - # ``frame_valid`` ``torch.where`` in :meth:`_place_hydrogens`. - self._h_idx_template = np.asarray(self._h_idx, dtype=np.int64).copy() - self._h_parent_idx_template = np.asarray(self._h_parent_idx, dtype=np.int64).copy() - h_n1 = np.asarray(self._h_n1_idx, dtype=np.int64).copy() - h_n2 = np.asarray(self._h_n2_idx, dtype=np.int64).copy() - h_n1[h_n1 < 0] = 0 - h_n2[h_n2 < 0] = 0 - self._h_n1_idx_template = h_n1 - self._h_n2_idx_template = h_n2 - self._h_local_pos_template = np.asarray(self._h_local_pos, dtype=np.float64).copy() - self._h_frame_valid_template = np.asarray(self._h_frame_valid, dtype=bool).copy() - self._h_offset_template = np.asarray(self._h_offset, dtype=np.float64).copy() - # Build the supercell System (replicate + PME + PBC). self._system = _replicate_to_supercell_system( template_system, @@ -575,91 +552,33 @@ def _relax_against_amber(self, max_iterations: int) -> None: # Lazy device buffers # ------------------------------------------------------------------ - def _ensure_torch_buffers( - self, device: torch.device, dtype: torch.dtype - ) -> None: - """Move/build the torch buffers for ``forward``: atom maps, tiled H - indices, and the constant init-positions tensor. Caches per + def _ensure_torch_buffers(self, device: torch.device, dtype: torch.dtype) -> None: + """Move/build atom maps and initial positions for ``forward``. Cache per (device, dtype). No work on repeat calls with the same key.""" - if ( - self._buffers_device == device - and self._buffers_dtype == dtype - ): + if self._buffers_device == device and self._buffers_dtype == dtype: return N = self._n_members n_omm = self._n_omm_per_member - # Initial sym-tiled positions for every OMM atom (nm). Used as the - # "fallback" position for slots that don't have a model atom mapped to - # them (waters, OXT, etc. tleap regenerated). - self._pos_buf_torch = torch.from_numpy(self._pos_buf).to( - device=device, dtype=dtype - ) - # Index pairs (long) for the scatter from model atoms into OMM slots. self._src_model_idx_torch = torch.from_numpy(self._src_model_idx_np).to( - device=device, dtype=torch.long + device=device, + dtype=torch.long, # dtype-ok: atom/copy index tensor for indexing; PyTorch requires int64 ) self._dst_omm_idx_torch = torch.from_numpy(self._dst_omm_idx_np).to( - device=device, dtype=torch.long + device=device, + dtype=torch.long, # dtype-ok: atom/copy index tensor for indexing; PyTorch requires int64 ) # Index of ensemble-model atoms (in the FULL EnsembleModel layout) # that survived the special-position filter — used in forward to # subset ``xyz_per_member`` before applying the layout transform. - self._keep_atom_idx_torch = torch.from_numpy( - self._keep_atom_idx_np - ).to(device=device, dtype=torch.long) - - # Boolean mask: True where the OMM slot has NO model atom mapped to - # it (so we keep the init position there). - unmapped = torch.ones(n_omm, dtype=torch.bool, device=device) - unmapped[self._dst_omm_idx_torch] = False - self._unmapped_mask_torch = unmapped # (n_omm,) - - # H-attachment indices tiled per member: template indices live in - # [0, n_omm); full-tensor indices live in [0, N · n_omm). - member_offset = ( - torch.arange(N, device=device, dtype=torch.long).unsqueeze(1) - * n_omm - ) # (N, 1) - h_idx_t = torch.from_numpy(self._h_idx_template).to( + self._keep_atom_idx_torch = torch.from_numpy(self._keep_atom_idx_np).to( device=device, dtype=torch.long - ) - h_parent_t = torch.from_numpy(self._h_parent_idx_template).to( - device=device, dtype=torch.long - ) - h_n1_t = torch.from_numpy(self._h_n1_idx_template).to( - device=device, dtype=torch.long - ) - h_n2_t = torch.from_numpy(self._h_n2_idx_template).to( - device=device, dtype=torch.long - ) - - self._h_idx_tiled = (member_offset + h_idx_t.unsqueeze(0)).reshape(-1) - self._h_parent_idx_tiled = ( - member_offset + h_parent_t.unsqueeze(0) - ).reshape(-1) - self._h_n1_idx_tiled = ( - member_offset + h_n1_t.unsqueeze(0) - ).reshape(-1) - self._h_n2_idx_tiled = ( - member_offset + h_n2_t.unsqueeze(0) - ).reshape(-1) - - # Per-H constants tiled by member (same value for each member's - # corresponding H). - self._h_local_pos_tiled = torch.from_numpy( - self._h_local_pos_template - ).to(device=device, dtype=dtype).repeat(N, 1) - self._h_frame_valid_tiled = torch.from_numpy( - self._h_frame_valid_template - ).to(device=device, dtype=torch.bool).repeat(N) - self._h_offset_tiled = torch.from_numpy( - self._h_offset_template - ).to(device=device, dtype=dtype).repeat(N, 1) + ) # dtype-ok: atom/copy index tensor for indexing; PyTorch requires int64 + self._omm_to_model = self._omm_to_model.to(device) self._buffers_device = device self._buffers_dtype = dtype @@ -667,16 +586,11 @@ def _ensure_torch_buffers( # Position composition # ------------------------------------------------------------------ - def _compose_full_omm_xyz( - self, supercell_xyz_nm: torch.Tensor - ) -> torch.Tensor: + def _compose_full_omm_xyz(self, supercell_xyz_nm: torch.Tensor) -> torch.Tensor: """Build the full ``(N · n_omm_per_member, 3)`` OpenMM xyz tensor. - Mapped (heavy) OMM slots get the current model coords (sym + tile - applied via the supercell layout); unmapped slots keep the construction- - time positions (tleap-regenerated atoms — waters, OXT, etc. that don't - move with the model). H atoms are then placed analytically from the - heavy positions via the tiled local-frame machinery. + Every OpenMM slot receives the current model coordinate after the + symmetry and tile transforms, including all hydrogen coordinates. Parameters ---------- @@ -690,50 +604,12 @@ def _compose_full_omm_xyz( Flat OpenMM-order positions in nm, differentiable in ``supercell_xyz_nm`` (and thus in ``model.xyz_per_member``). """ - N = self._n_members - n_omm = self._n_omm_per_member - device = supercell_xyz_nm.device - dtype = supercell_xyz_nm.dtype - - pos_init = self._pos_buf_torch.view(N, n_omm, 3) # (N, n_omm, 3) - # Scatter mapped model atoms into a zero tensor at the OMM slots. - src = supercell_xyz_nm.index_select(1, self._src_model_idx_torch) - # index_copy is autograd-friendly and returns a new tensor. - scattered = torch.zeros( - (N, n_omm, 3), device=device, dtype=dtype - ).index_copy(1, self._dst_omm_idx_torch, src) - # Where the slot is unmapped, use the init position; otherwise use - # scattered (the current model coord). - mask = self._unmapped_mask_torch.view(1, n_omm, 1) - heavy = torch.where(mask, pos_init, scattered) # (N, n_omm, 3) - - # Place hydrogens (operates on the flat (N·n_omm, 3) view). - full_flat = heavy.reshape(-1, 3) - return self._place_hydrogens(full_flat) - - def _place_hydrogens(self, heavy_xyz_nm: torch.Tensor) -> torch.Tensor: - """Vectorized H placement across all members via tiled local frames. - - Mirrors :meth:`AmberTarget._place_hydrogens` but operates on the - ``(N·n_omm_per_member, 3)`` supercell positions with tiled parent / - neighbour indices and per-H constants. Frame: parent + first heavy - neighbour for ``e1``, second heavy neighbour projected for ``e2``, - cross for ``e3``; H position is ``p + Σ local_pos[k] · e_k``. Rigid - fallback for the small fraction of Hs without two heavy neighbours. - """ - # Same local-frame physics as the single-molecule path — one shared - # implementation, here applied with member-tiled index tensors. - h_pos = _place_hydrogens_local_frame( - heavy_xyz_nm, - self._h_parent_idx_tiled, - self._h_n1_idx_tiled, - self._h_n2_idx_tiled, - self._h_local_pos_tiled, - self._h_frame_valid_tiled, - self._h_offset_tiled, - ) - # Write H positions into the heavy tensor via functional index_copy. - return heavy_xyz_nm.index_copy(0, self._h_idx_tiled, h_pos) + if supercell_xyz_nm.shape != (self._n_members, self._n_model_per_member, 3): + raise ValueError( + "[QuasiCrystalAmberTarget] Atom layout changed; rebuild the target." + ) + return supercell_xyz_nm.index_select(1, self._omm_to_model).reshape(-1, 3) + # ------------------------------------------------------------------ # Forward diff --git a/torchref/experimental/ensemble/rank_penalty.py b/torchref/experimental/ensemble/rank_penalty.py index 9254c17e..bf214098 100644 --- a/torchref/experimental/ensemble/rank_penalty.py +++ b/torchref/experimental/ensemble/rank_penalty.py @@ -232,6 +232,8 @@ def spectrum_diagnostics(self) -> dict: variance in the top mode. These show the "purification" as the penalty ramps up. """ + # dtype-ok: float64 svdvals for read-only diagnostics; results extracted + # via float(). Caveat: no .cpu() first, so this errors on MPS. Xc = self._centered().detach().to(torch.float64) s = torch.linalg.svdvals(Xc) # (min(N, D),) s2 = s ** 2 diff --git a/torchref/experimental/ensemble/wilson_prior.py b/torchref/experimental/ensemble/wilson_prior.py index e2fdf7e8..44490f9c 100644 --- a/torchref/experimental/ensemble/wilson_prior.py +++ b/torchref/experimental/ensemble/wilson_prior.py @@ -172,7 +172,7 @@ def _build_bin_assignment(self) -> None: order = torch.argsort(res) n = res.numel() nbins = min(self.nbins, max(1, n // 50)) - bin_assign = torch.empty(n, dtype=torch.long, device=res.device) + bin_assign = torch.empty(n, dtype=torch.long, device=res.device) # dtype-ok: bin-assignment tensor used as scatter_add index; PyTorch requires int64 edges = torch.linspace(0, n, nbins + 1, device=res.device).round().long() for b in range(nbins): start = int(edges[b].item()) diff --git a/torchref/experimental/kinetic/__init__.py b/torchref/experimental/kinetic/__init__.py index 168acf93..0c9bbf59 100644 --- a/torchref/experimental/kinetic/__init__.py +++ b/torchref/experimental/kinetic/__init__.py @@ -30,7 +30,7 @@ from torchref.experimental.kinetic.refinement import KineticRefinement from torchref.experimental.kinetic.targets import ( CollectionDifferenceTarget, - CollectionRiceTarget, + CollectionMLTarget, MultiModelGeometryTarget, MultiModelADPTarget, KineticPriorTarget, @@ -44,7 +44,7 @@ "ModelCollection", "KineticRefinement", "CollectionDifferenceTarget", - "CollectionRiceTarget", + "CollectionMLTarget", "MultiModelGeometryTarget", "MultiModelADPTarget", "KineticPriorTarget", diff --git a/torchref/experimental/kinetic/occupancies.py b/torchref/experimental/kinetic/occupancies.py index 8d4aa720..c045bd89 100644 --- a/torchref/experimental/kinetic/occupancies.py +++ b/torchref/experimental/kinetic/occupancies.py @@ -18,6 +18,7 @@ from typing import Dict, List, Optional, Union, Tuple import numpy as np +from torchref.config import get_float_dtype from torchref.utils.device_mixin import DeviceMixin @@ -163,7 +164,7 @@ def __init__( # Convert time to tensor if needed if not isinstance(time, torch.Tensor): - time = torch.tensor(time, dtype=torch.float32) + time = torch.tensor(time, dtype=get_float_dtype()) self.register_buffer('time', time) # Initialize the kinetic model diff --git a/torchref/experimental/kinetic/refinement.py b/torchref/experimental/kinetic/refinement.py index 94514c0c..2aebd182 100644 --- a/torchref/experimental/kinetic/refinement.py +++ b/torchref/experimental/kinetic/refinement.py @@ -13,8 +13,8 @@ from torchref.experimental.kinetic import ModelCollection, KineticRefinement # Base models - model_dark = ModelFT(max_res=1.5).load_pdb("dark.pdb") - model_light = ModelFT(max_res=1.5).load_pdb("light.pdb") + model_dark = ModelFT(max_res=1.5, cif_path=cif_paths).load_pdb("dark.pdb") + model_light = ModelFT(max_res=1.5, cif_path=cif_paths).load_pdb("light.pdb") # Collections models = ModelCollection([model_dark, model_light]) @@ -30,6 +30,7 @@ ref.refine(macro_cycles=5) """ +import warnings from typing import TYPE_CHECKING, Dict, List, Optional import torch @@ -38,7 +39,7 @@ from torchref.refinement.loss_state import LossState, create_loss_state from torchref.experimental.kinetic.targets import ( CollectionDifferenceTarget, - CollectionRiceTarget, + CollectionMLTarget, MultiModelGeometryTarget, MultiModelADPTarget, KineticPriorTarget, @@ -66,8 +67,10 @@ class KineticRefinement(DeviceMixin, nn.Module): Collection of mixed models keyed by timepoint name. xray_weight_difference : float Weight for the difference X-ray target. - xray_weight_rice : float - Weight for the Rice amplitude target. + xray_weight_ml : float + Weight for the absolute (Read MLF) amplitude target. With K free base models a + purely relative loss leaves the overall level unconstrained, so this is the + anchor. geometry_weight : float Weight for geometry restraints. adp_weight : float @@ -85,7 +88,7 @@ def __init__( dataset_collection: "DatasetCollection", model_collection: "ModelCollection", xray_weight_difference: float = 2.0, - xray_weight_rice: float = 1.0, + xray_weight_ml: float = 1.0, geometry_weight: float = 10.0, adp_weight: float = 3.0, kinetic_prior_weight: float = 0.0, @@ -105,7 +108,7 @@ def __init__( # Default weights self._weights = { "xray/difference": xray_weight_difference, - "xray/rice": xray_weight_rice, + "xray/ml": xray_weight_ml, "geometry": geometry_weight, "adp": adp_weight, "kinetic_prior": kinetic_prior_weight, @@ -116,7 +119,7 @@ def __init__( self.loss_state: Optional[LossState] = None self.kinetic_prior_target: Optional[KineticPriorTarget] = None self._diff_target: Optional[CollectionDifferenceTarget] = None - self._rice_target: Optional[CollectionRiceTarget] = None + self._ml_target: Optional[CollectionMLTarget] = None self._kinetic_model = None self._timepoints_map: Optional[Dict[str, int]] = None @@ -163,14 +166,14 @@ def setup( scaler=self.scaler, verbose=self.verbose, ) - rice_target = CollectionRiceTarget( + ml_target = CollectionMLTarget( dc, mc, scaler=self.scaler, verbose=self.verbose, ) # Store direct references for refine_kinetics() self._diff_target = diff_target - self._rice_target = rice_target + self._ml_target = ml_target geom_target = MultiModelGeometryTarget(mc, verbose=self.verbose) adp_target = MultiModelADPTarget(mc, verbose=self.verbose) @@ -178,7 +181,7 @@ def setup( self.loss_state = create_loss_state(device=device) self.loss_state.register_target("xray/difference", diff_target) - self.loss_state.register_target("xray/rice", rice_target) + self.loss_state.register_target("xray/ml", ml_target) self.loss_state.register_target("geometry", geom_target) self.loss_state.register_target("adp", adp_target) @@ -241,7 +244,7 @@ def set_weights(self, **kwargs): # Map short names to full paths mapping = { "difference": "xray/difference", - "rice": "xray/rice", + "ml": "xray/ml", "geometry": "geometry", "adp": "adp", "kinetic_prior": "kinetic_prior", @@ -321,10 +324,9 @@ def _collect_parameters(self, structures=True, fractions=True): ) if fractions: - for name in mc.timepoint_names: - p = mc[name].fraction_params - if p.requires_grad: - params.append(p) + # One shared activation plus a branching row per timepoint, owned by the + # collection rather than by any single timepoint. + params.extend(p for p in mc.fraction_parameters() if p.requires_grad) # Scaler parameters (if not frozen) if self.scaler is not None: @@ -528,7 +530,7 @@ def refine_kinetics(self, niter: int = 200, lr: float = 1e-2): optimizer = torch.optim.Adam(params, lr=lr) w_diff = self._weights.get("xray/difference", 1.0) - w_rice = self._weights.get("xray/rice", 1.0) + w_ml = self._weights.get("xray/ml", 1.0) # Collect all model keys that have kinetic indices (including dark) all_overrides = {} @@ -546,8 +548,8 @@ def refine_kinetics(self, niter: int = 200, lr: float = 1e-2): for tp_name, t_idx in all_overrides.items(): mc[tp_name].set_fraction_override(kinetic_occ[:, t_idx]) - # Compute X-ray loss only (difference + Rice) - loss = w_diff * self._diff_target() + w_rice * self._rice_target() + # Compute X-ray loss only (difference + the absolute anchor) + loss = w_diff * self._diff_target() + w_ml * self._ml_target() loss.backward() @@ -560,16 +562,40 @@ def refine_kinetics(self, niter: int = 200, lr: float = 1e-2): if self.verbose > 1 and (step + 1) % 10 == 0: print(f" Kinetic opt step {step+1}/{niter}: loss = {loss.item():.6f}") - # Update free fraction parameters to match final kinetic predictions + # Update the population parameters to match the final kinetic predictions. + # + # The collection stores one shared activation plus a per-timepoint branching, + # so a predicted population vector is decomposed into the two. A kinetic model + # whose reactive fraction is genuinely constant gives the same activation at + # every timepoint; if it does not, the shared value cannot represent all of + # them and the closest one is kept, with the spread reported. with torch.no_grad(): kinetic_occ = kinetic_model() + implied = {} for tp_name, t_idx in all_overrides.items(): if tp_name == mc.dark_key: - continue # dark fractions stay frozen at [1,0,...,0] + continue # the reference is the alpha = 0 evaluation predicted = kinetic_occ[:, t_idx] - mc[tp_name].fraction_params.data = torch.log( - predicted.clamp(min=1e-6) - ) + alpha = float(1.0 - predicted[0]) + if alpha <= 1e-6: + continue + implied[tp_name] = alpha + mc.set_branching(tp_name, predicted[1:]) + + if implied: + alphas = list(implied.values()) + spread = max(alphas) - min(alphas) + mean_alpha = sum(alphas) / len(alphas) + mc.set_activation(mean_alpha) + if spread > 1e-3: + warnings.warn( + f"Kinetic populations imply activations spanning {spread:.4f} " + f"across timepoints, but one activation is shared by all of " + f"them; using the mean {mean_alpha:.4f}. Drive the timepoints " + f"with set_fraction_override() to keep them independent.", + UserWarning, + stacklevel=2, + ) mc.unfreeze_structures() diff --git a/torchref/experimental/kinetic/targets.py b/torchref/experimental/kinetic/targets.py index 54201230..10037c6c 100644 --- a/torchref/experimental/kinetic/targets.py +++ b/torchref/experimental/kinetic/targets.py @@ -7,8 +7,8 @@ CollectionDifferenceTarget Multi-timepoint difference target (primary optimization driver). -CollectionRiceTarget - Multi-timepoint Rice maximum-likelihood amplitude target. +CollectionMLTarget + Read MLF at one shared Luzzati beta; the absolute channel. MultiModelGeometryTarget Geometry restraints applied to the shared base models. MultiModelADPTarget @@ -28,9 +28,9 @@ # Back-compat re-exports of the relocated generic collection targets. from torchref.refinement.targets.collection import ( # noqa: F401 + CollectionDifferenceIntensityTarget, CollectionDifferenceTarget, CollectionMLTarget, - CollectionRiceTarget, MultiModelADPTarget, MultiModelGeometryTarget, ) diff --git a/torchref/experimental/monolithic_refinement/density_scaler.py b/torchref/experimental/monolithic_refinement/density_scaler.py index b0322c3d..afe76f95 100644 --- a/torchref/experimental/monolithic_refinement/density_scaler.py +++ b/torchref/experimental/monolithic_refinement/density_scaler.py @@ -127,7 +127,7 @@ def get_rec_solvent(self, hkl): Not detached: ``F_sol`` follows the moving atoms so gradients reach ``xyz``/``adp``. The scaler applies the contrast and falloff on top. """ - return self.density(hkl.to(torch.long)) + return self.density(hkl.to(torch.long)) # dtype-ok: hkl cast to long for density lookup indexing; PyTorch requires int64 def update_solvent(self): """No-op: the density mask is rebuilt live on every scaler forward.""" diff --git a/torchref/experimental/monolithic_refinement/density_solvent.py b/torchref/experimental/monolithic_refinement/density_solvent.py index bac26532..220b33ac 100644 --- a/torchref/experimental/monolithic_refinement/density_solvent.py +++ b/torchref/experimental/monolithic_refinement/density_solvent.py @@ -32,6 +32,7 @@ ifft, ) from torchref.config import get_default_device, get_float_dtype +from torchref.model.context import ModelContext from torchref.model.sf_fft import SfFFT from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMixin @@ -195,16 +196,20 @@ def __init__( # because the nonlinear occupancy needs the full-cell density assembled # before the mask. The per-atom splat radius is governed by # torchref.sigma_cutoff_ed inside the density builder. + # Its own context: the cell is shared with the model, but the space group + # is copied because it memoises operators per grid shape and this engine's + # coarse grid must not evict the model's. + solvent_ctx = ModelContext( + cell=model.cell, spacegroup=model.spacegroup.copy() + ) self.solvent_fft = SfFFT( - cell=model.cell, - spacegroup=model.fft.spacegroup, + ctx=solvent_ctx, max_res=self.solvent_res, dtype_float=float_type, device=device, verbose=max(0, verbose - 1), use_late_symmetry=False, ) - self.solvent_fft.setup_grid() # ------------------------------------------------------------------ # Density -> occupancy -> structure factor @@ -280,10 +285,13 @@ def _smooth(self, field): At ``sigma=0`` the kernel is the identity (no smoothing). Differentiable w.r.t. ``sigma_shell`` and -- through ``field`` -- w.r.t. atomic xyz/B. """ - grid = self.solvent_fft.real_space_grid # (nx, ny, nz, 3) Cartesian - dx = float((grid[1, 0, 0] - grid[0, 0, 0]).norm()) - dy = float((grid[0, 1, 0] - grid[0, 0, 0]).norm()) - dz = float((grid[0, 0, 1] - grid[0, 0, 0]).norm()) + # Axis spacings: column j of the fractional (frac->cart) matrix is cell edge + # vector j, so its norm over the sampling count along that axis is the step + # between grid points one index apart -- what differencing the coordinate grid + # used to measure, without building the grid. + frac = self.solvent_fft.cell.fractional_matrix + n = self.solvent_fft.grid_shape + dx, dy, dz = (float(frac[:, j].norm()) / n[j] for j in range(3)) sigma = self.sigma_shell.clamp(min=0.0, max=5.0) nx, ny, nz = field.shape two_pi2 = 2.0 * torch.pi**2 @@ -316,7 +324,7 @@ def _get_tau(self, z): def _nyquist_mask(self, hkl): """True where |h|,|k|,|l| are within the coarse grid Nyquist limit.""" - nc = self.solvent_fft.real_space_grid.shape[:-1] # (ncx, ncy, ncz) + nc = self.solvent_fft.grid_shape # (ncx, ncy, ncz) nyq = torch.tensor( [n // 2 for n in nc], device=hkl.device, dtype=hkl.dtype ) diff --git a/torchref/experimental/targets/amber_target.py b/torchref/experimental/targets/amber_target.py index 414616da..ca6383f2 100644 --- a/torchref/experimental/targets/amber_target.py +++ b/torchref/experimental/targets/amber_target.py @@ -1,49 +1,10 @@ -""" -AMBER14/GAFF2 Force Field as a Differentiable Restraint. - -Uses OpenMM to evaluate the AMBER14 energy for current model coordinates. -Analytical forces from OpenMM are bridged into PyTorch autograd via a -custom Function, making the energy fully differentiable w.r.t. xyz. - -Non-standard residues (HETATM not in AMBER14_STANDARD) are parameterised -automatically via antechamber/GAFF2. Results are cached under -``PATH_TORCHREF_DATA / "amber_cache" / {resname}/``. - -Intended workflow:: - - # Canonical one-liner — strips altlocs, adds H, then build target: - mh = (Model(verbose=0, strip_H=True) - .load_pdb('structure.pdb') - .strip_altlocs() - .generate_hydrogens()) - target = AmberTarget(model=mh) # protein-only - target = AmberTarget(model=mh, residue_charges={'LIG': -1}) # with ligand - - loss = target() # kJ/mol per atom - loss.backward() - # xyz gradient is now populated with AMBER forces - -Performance note ----------------- -OpenMM's ``Modeller.addHydrogens()`` is faster when H atoms are already present -in the model (it refines positions rather than building from scratch). -Gradient and energy are identical either way (H are stripped from the atom map; -``n_model_atoms`` changes only the energy normalisation). - -Design notes ------------- -- Standard-residues path uses pdbfixer to add missing terminal/sidechain - heavy atoms before OpenMM's Modeller adds H; the GAFF2 path uses tleap - for H addition. -- Altloc atoms are filtered before building the OpenMM system: only the - primary conformation (altloc == '' or 'A') is used. -- OXT and H atoms are excluded from the PDB written to tleap; tleap - re-adds them via its C-terminal and H-addition templates. -- H positions in the OpenMM context are set once at construction and are - NOT updated during forward() — a good approximation for small refinement - steps (< 0.1 Å heavy-atom displacement). -- model_to_omm maps model-atom index → OpenMM atom index for HEAVY atoms - only. Model H atoms receive -1 and are skipped in forward(). +"""Evaluate AMBER energies and forces for TorchRef-owned atomic coordinates. + +The model must already contain the atoms and protonation state required by the +force field. Construction validates a one-to-one atom map; it does not add atoms +to the model. Each evaluation reorders all coordinates, including hydrogens, +converts Cartesian Å to nm, and returns OpenMM forces through PyTorch autograd. +Riding geometry and orientation parameters belong to the model. """ from __future__ import annotations @@ -63,6 +24,8 @@ import torch from torchref import PATH_TORCHREF_DATA +from torchref.config import get_float_dtype, get_int_dtype +from torchref.refinement.targets.base import ModelTarget from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -71,8 +34,6 @@ stat, ) -from torchref.refinement.targets.base import ModelTarget - if TYPE_CHECKING: from torchref.model.model import Model @@ -131,24 +92,8 @@ def _find_ambertools_binary(name: str) -> str: } ) -# Atom names that tleap adds itself via terminal / template logic; must be -# excluded from the PDB handed to tleap to avoid "does not have a type" errors. _TLEAP_SKIP_ATOMS: frozenset = frozenset({"OXT", "OT1", "OT2"}) -# Residues handled by amber14-all.xml + amber14/tip3pfb.xml in OpenMM Modeller -# (does NOT include Mg/Zn/Ca/Fe etc. — those lack templates in the default XML set) -_MODELLER_FF_RESIDUES: frozenset = frozenset( - { - "ALA", "ARG", "ASN", "ASP", "CYS", "CYX", "GLN", "GLU", "GLY", - "HID", "HIE", "HIP", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", - "PRO", "SER", "THR", "TRP", "TYR", "VAL", - "ACE", "NME", - "HOH", "WAT", - "NA", "K", "CL", # ions in amber14/tip3pfb.xml - "A", "G", "C", "U", "T", "DA", "DG", "DC", "DT", - } -) - # Residues to exclude from the protein PDB written to tleap (GAFF2 path). # Currently empty: all AMBER14_STANDARD residues (protein, ions, water) are # included so they participate in both LJ (steric) and Coulomb gradients. @@ -177,17 +122,9 @@ class _OpenMMAMBERFunction(torch.autograd.Function): forward : full_xyz_nm (nm, float, [n_omm_total, 3]) → energy (kJ/mol) - The input tensor must already contain positions for **every** OpenMM - atom — heavy and H — in OpenMM's native atom order. Building this - tensor (scattering model heavy atoms + computing H positions - analytically from heavy positions) happens in - :meth:`AmberTarget._compose_full_omm_xyz`, upstream of this Function. - - backward: ∂E/∂full_xyz = −F (full OpenMM force vector). The H - contributions in F propagate naturally through ``_compose_full_omm_xyz`` - and ``_place_hydrogens`` upstream via PyTorch autograd, delivering - correctly-distributed gradients to the heavy model atoms (parent + - local-frame neighbors). + The input contains every model atom in OpenMM order. Backward returns + minus the force in kJ/mol/nm; the upstream gather and Å-to-nm conversion + return each gradient to its TorchRef coordinate or riding parameter. """ @staticmethod @@ -235,83 +172,6 @@ def backward(ctx, grad_output): return -forces * grad_output, None, None -# --------------------------------------------------------------------------- -# Differentiable hydrogen placement (single source of truth) -# --------------------------------------------------------------------------- - - -def _place_hydrogens_local_frame( - heavy_xyz: torch.Tensor, - parent_idx: torch.Tensor, - n1_idx: torch.Tensor, - n2_idx: torch.Tensor, - local_pos: torch.Tensor, - frame_valid: torch.Tensor, - offset: torch.Tensor, - eps: float = 1e-12, -) -> torch.Tensor: - """Place hydrogens from heavy-atom positions via captured local frames. - - The one and only implementation of the H-placement physics, shared by the - single-molecule / per-member path (:meth:`AmberTarget._place_hydrogens`) - and the tiled supercell path (``QuasiCrystalAmberTarget._place_hydrogens``). - Differentiable in ``heavy_xyz``: autograd distributes each H force onto its - parent + the two frame-reference atoms via the exact local-frame Jacobian. - - For each H, an orthonormal frame is built from its parent ``p`` and two - heavy neighbours ``n1, n2``:: - - e1 = û(n1 − p) - e2 = û((n2 − p) ⊥ e1) - e3 = e1 × e2 - h = p + lx·e1 + ly·e2 + lz·e3 - - Hs flagged ``frame_valid == False`` (no two heavy neighbours) fall back to - the rigid translation ``p + offset``. - - Parameters - ---------- - heavy_xyz : torch.Tensor, ``(M, 3)`` - Positions (nm) with all heavy-atom slots populated. May be a single - topology (``M = n_omm``) or a tiled supercell (``M = N · n_omm``). - parent_idx, n1_idx, n2_idx : torch.Tensor, ``(H,)`` long - Indices into ``heavy_xyz``. Invalid-frame neighbour indices must be - pre-clamped to a safe in-bounds value (their result is masked out). - local_pos : torch.Tensor, ``(H, 3)`` - Captured local-frame coordinates of each H. - frame_valid : torch.Tensor, ``(H,)`` bool - Whether the local-frame placement is used (else the rigid fallback). - offset : torch.Tensor, ``(H, 3)`` - Rigid-fallback ``p → H`` vector. - eps : float - Norm floor guarding degenerate frames. - - Returns - ------- - torch.Tensor, ``(H, 3)`` - H positions (nm). The caller writes these into the H slots. - """ - p = heavy_xyz.index_select(0, parent_idx) - n1 = heavy_xyz.index_select(0, n1_idx) - n2 = heavy_xyz.index_select(0, n2_idx) - - a = n1 - p - e1 = a / a.norm(dim=-1, keepdim=True).clamp(min=eps) - b = n2 - p - b_perp = b - (b * e1).sum(-1, keepdim=True) * e1 - e2 = b_perp / b_perp.norm(dim=-1, keepdim=True).clamp(min=eps) - e3 = torch.cross(e1, e2, dim=-1) - - h_frame = ( - p - + local_pos[:, 0:1] * e1 - + local_pos[:, 1:2] * e2 - + local_pos[:, 2:3] * e3 - ) - h_rigid = p + offset - return torch.where(frame_valid.unsqueeze(-1), h_frame, h_rigid) - - # --------------------------------------------------------------------------- # AmberTarget # --------------------------------------------------------------------------- @@ -321,47 +181,17 @@ class AmberTarget(ModelTarget): """ Differentiable AMBER14/GAFF2 force-field energy restraint. - On construction the target: - - 1. Detects non-standard residues (HETATM not in :data:`AMBER14_STANDARD`). - 2. Runs antechamber + parmchk2 (parallel, cached) for each non-standard - residue. - 3. Builds an OpenMM system: - - * **Standard path** (no non-standard residues): filter model PDB to - primary conformation + heavy atoms, use ``openmm.app.Modeller`` to - re-add H with AMBER14-compatible names, create system with - ``ForceField('amber14-all.xml')``. - * **GAFF2 path** (with non-standard residues): same protein PDB - (additionally removing OXT) handed to tleap together with each - ligand's mol2 via ``combine{}``. Combined AMBER14+GAFF2 topology - is parameterised by parmed. - - 4. Creates an OpenMM Context on the platform that matches the model's - device: CUDA for ``model.device.type == 'cuda'``, CPU otherwise. - Falls back CUDA → OpenCL → CPU if the preferred platform is unavailable. - 5. Builds a model-atom → OpenMM-atom index map so that only heavy atoms - are transferred; H positions are kept from the initial OpenMM setup. + Build chemistry once, then supply current coordinates for every atom to + OpenMM. The loss never generates or independently places hydrogens. Parameters ---------- model : Model - TorchRef model. Heavy-atom-only models (``strip_H=True``) are - accepted. H atoms are added internally by OpenMM's Modeller or - tleap and are NOT included in the atom map or gradient. - - Passing a model that already has H atoms (via - ``model.generate_hydrogens()`` or loading a PDB with H) speeds up - initialisation because ``Modeller.addHydrogens()`` converges - faster from existing positions. - - **GAFF2 ligands**: antechamber's BCC charge scheme runs a - semiempirical QM step (sqm) that needs a fully protonated molecule. - Heavy-only ligands are auto-protonated from the monomer library - (``generate_hydrogens``) first; an error is raised only if no - monomer CIF resolves AND the heavy-atom electron count is odd. - Calling ``model.generate_hydrogens()`` or loading the PDB with - ``strip_H=False`` beforehand avoids relying on that fallback. + Fully prepared, single-conformation model, including hydrogens and + terminal atoms required by AMBER. Existing atoms and coordinates are + preserved. Prepare protonation before constructing this target and + enable riding mode on the model when hydrogen geometry is constrained. + Incomplete or incompatible chemistry raises ValueError during setup. cutoff : float Non-bonded cutoff in Angstroms. Default 5.0. normalize_by_atoms : bool @@ -389,6 +219,15 @@ class AmberTarget(ModelTarget): multi-member ensemble, ``chem_model`` supplies the one conformation used to build the chemistry/topology; defaults to ``model`` for the single-molecule case. + + Notes + ----- + Reconstruct the target after changing atom identities, atom order, or + connectivity. Cartesian coordinate and riding-parameter changes need no + rebuild. OpenMM evaluation transfers coordinates and forces through CPU + memory and supports first derivatives only. Forces above 10000 kJ/mol/nm + are clipped per atom; in that regime the returned gradient is clipped + rather than the exact energy derivative. """ name: str = "amber" @@ -403,7 +242,7 @@ def __init__( charge_method: str = "gas", verbose: int = 0, chem_model: "Model" = None, - ): + ) -> None: try: import openmm # noqa: F401, PLC0415 except ImportError: @@ -438,7 +277,10 @@ def __init__( self._residue_charges = dict(residue_charges) if residue_charges else {} self._gaff2_files = dict(gaff2_files) if gaff2_files else {} - self.register_buffer("_cutoff_buf", torch.tensor(float(cutoff))) + self.register_buffer( + "_cutoff_buf", + torch.tensor(float(cutoff), dtype=get_float_dtype(), device=self.device), + ) # Internal state (None until fully initialised) self._context = None @@ -448,13 +290,8 @@ def __init__( self._n_omm_atoms: int = 0 self._n_model_atoms: int = 0 self._n_nonstandard: int = 0 - # GAFF2 path: ordered residue map for atom matching (None = standard path) - self._tleap_residue_map: Optional[List[Dict[str, int]]] = None - # Cached protonated chemistry PDB (filled lazily by the first ligand - # parameterisation that needs H). None = not yet computed; False = - # generate_hydrogens failed (don't retry). - self._protonated_pdb_cache = None - + # tleap renumbers residues; identify their original model instances. + self._tleap_residue_map: Optional[bool] = None if self._chem_model is None: return # Allow empty init for state_dict loading @@ -494,17 +331,14 @@ def _build(self) -> None: self._build_atom_map() del self._tleap_pos_nm + xyz = self._chem_model.xyz().detach().cpu().numpy() + positions_nm = np.asarray( + xyz[np.argsort(self._model_to_omm)] * 0.1, dtype=np.float64 + ) self._build_context(positions_nm) - # Pre-allocate nm position buffer: H positions pre-filled from OpenMM init self._pos_buf = positions_nm.copy() self._n_model_atoms = len(self._chem_model.pdb) - # Build (H, parent, offset) table so we can rigidly re-attach H atoms - # to their parent heavy atom each forward. Without this, H positions - # stay frozen at construction time while heavy atoms move, blowing up - # bond-stretch terms by orders of magnitude (the dominant pathology - # for any model.xyz() that excludes H). - self._build_h_attachment(positions_nm) if self.verbose >= 1: print( @@ -594,55 +428,9 @@ def _write_residue_pdb(self, res_atoms, path: Path) -> None: ) f.write("END\n") - def _protonated_chem_pdb(self): - """Protonated chemistry-model PDB DataFrame (cached), or ``None``. - Uses :meth:`Model.generate_hydrogens` once on the whole chemistry model - (which has a unit cell + full residue context, so gemmi's topology engine - is well-posed). H come from the monomer-library CIF at ideal geometry via - TorchRef's auto-fetching monomer library — no full CCP4 install needed. - Cached so repeated ligand parameterisations don't re-run it. - """ - if self._protonated_pdb_cache is None: - try: - m_h = self._chem_model.generate_hydrogens() - self._protonated_pdb_cache = ( - m_h.update_pdb() if hasattr(m_h, "update_pdb") else m_h.pdb - ) - except Exception as exc: # missing CIF/lib, gemmi failure, etc. - if self.verbose >= 1: - print(f"[AmberTarget] generate_hydrogens failed: {exc}") - self._protonated_pdb_cache = False - if self._protonated_pdb_cache is False: - return None - return self._protonated_pdb_cache - - def _protonate_residue_pdb(self, resname: str, out_pdb: Path) -> bool: - """Write a protonated single-residue PDB for ``resname`` to ``out_pdb``. - - antechamber/GAFF2 needs a protonated, valence-satisfied molecule because - the model is heavy-atom-only. Only topologically-correct H are required - here — charges are Gasteiger (connectivity-based, no QM) and the running- - system H are re-placed analytically each step — so the monomer library's - ideal geometry (via :meth:`_protonated_chem_pdb`) is ample. - - Returns ``True`` iff H were added for ``resname`` (a monomer CIF - resolved); ``False`` lets the caller fall back. - """ - pdb_h = self._protonated_chem_pdb() - if pdb_h is None: - return False - res = pdb_h[pdb_h["resname"].astype(str).str.strip() == resname] - h_mask = res["element"].astype(str).str.strip().isin(["H", "D"]) - if not bool(h_mask.any()): - return False - self._write_residue_pdb(res, out_pdb) - if self.verbose >= 1: - print( - f"[AmberTarget] protonated '{resname}' via monomer library: " - f"+{int(h_mask.sum())} H" - ) - return True + + def _run_antechamber_one( self, resname: str, charge: int @@ -653,24 +441,14 @@ def _run_antechamber_one( Cache is checked first. On a miss, work happens in a temp dir and results are atomically moved to the cache (write-then-rename). """ - pdb = self._chem_model.pdb + pdb = self._chem_model.pdb.copy() + pdb[["x", "y", "z"]] = self._chem_model.xyz().detach().cpu().numpy() res_atoms = pdb[pdb["resname"].astype(str).str.strip() == resname] + first = res_atoms.iloc[0] + for column in ("chainid", "resseq", "icode"): + res_atoms = res_atoms[res_atoms[column] == first[column]] atom_names = res_atoms["name"].astype(str).str.strip().tolist() - # antechamber needs a fully protonated molecule (sqm — used for BCC - # charges — needs an even electron count, and GAFF2 atom typing needs - # satisfied valences). The model is heavy-atom-only, so a ligand with no - # H is protonated below from the monomer library before antechamber runs. - # Compute the heavy-atom electron parity here to sanity-check the result. - _Z = {"H":1,"He":2,"Li":3,"Be":4,"B":5,"C":6,"N":7,"O":8,"F":9,"Ne":10, - "Na":11,"Mg":12,"Al":13,"Si":14,"P":15,"S":16,"Cl":17,"Ar":18, - "K":19,"Ca":20,"Cr":24,"Mn":25,"Fe":26,"Co":27,"Ni":28,"Cu":29, - "Zn":30,"Br":35,"I":53,"Se":34,"Mo":42,"W":74,"Pt":78,"Au":79} - elems = res_atoms["element"].astype(str).str.strip().str.capitalize() - n_protons = sum(_Z.get(e, 0) for e in elems) - n_electrons = n_protons - charge - has_h = bool(elems.isin(["H", "D"]).any()) - key = self._cache_key(resname, atom_names, charge, self._charge_method) cache_dir = self._get_cache_dir(resname) @@ -693,26 +471,7 @@ def _run_antechamber_one( self._write_residue_pdb(res_atoms, lig_pdb) - # Heavy-atom-only ligand → protonate before antechamber so GAFF2 - # typing sees satisfied valences (and sqm, if BCC, gets a closed- - # shell molecule). Hydrogens come from TorchRef's monomer-library - # placement at ideal geometry. antechamber_input = lig_pdb - if not has_h: - lig_h_pdb = work_dir / "lig_h.pdb" - if self._protonate_residue_pdb(resname, lig_h_pdb): - antechamber_input = lig_h_pdb - elif n_electrons % 2 != 0: - raise RuntimeError( - f"[AmberTarget] Cannot parameterise '{resname}': odd " - f"electron count ({n_electrons}) for charge {charge:+d} " - f"and no hydrogens could be added (no monomer-library CIF " - f"resolved for '{resname}').\nFix: pass an explicit charge " - f"via residue_charges={{'{resname}': }}, supply " - f"gaff2_files for this residue, or make a monomer CIF " - f"resolvable for auto-protonation (TORCHREF_MONOMER_LIB, " - f"or CLIBD_MON as an optional override)." - ) # antechamber r = subprocess.run( @@ -820,54 +579,16 @@ def _run_antechamber_parallel( # Step 3 — Build OpenMM system # ------------------------------------------------------------------ - def _filter_pdb_for_omm(self, include_nonstandard: bool = False): - """ - Return a filtered copy of model.pdb suitable for OpenMM / tleap: - - Primary conformation only (altloc == '' or 'A') - - Heavy atoms only (element != H or D) - - Optionally exclude non-standard residues (standard path) - - The returned DataFrame keeps the original model.pdb integer index - so that ``df.index`` can be used as model row indices in the atom map. - """ - pdb = self._chem_model.update_pdb() - - mask = pdb["altloc"].astype(str).str.strip().isin(["", "A"]) - mask &= ~pdb["element"].astype(str).str.strip().isin(["H", "D"]) - - if not include_nonstandard: - ns_resnames = { - rn for rn in pdb["resname"].astype(str).str.strip().unique() - if rn not in _MODELLER_FF_RESIDUES - } - if ns_resnames: - mask &= ~pdb["resname"].astype(str).str.strip().isin(ns_resnames) - # Do NOT reset_index: keep original model.pdb row positions as index - return pdb[mask].copy() def _filter_pdb_for_tleap(self): + """Export standard heavy atoms for tleap template parameterisation. + + The resulting topology must map back to every model atom, including + hydrogens and terminal oxygens, before a context can be constructed. """ - Filter model.pdb for the tleap protein PDB (GAFF2 path): - - - Primary conformation only (altloc == '' or 'A') - - Heavy atoms only (element != H or D) - - Standard AMBER residues only (``AMBER14_STANDARD``) — non-standard - HETATM residues are handled via antechamber / mol2 separately - - Waters (HOH/WAT) ARE included — ``_TLEAP_EXCLUDE_RESIDUES`` is - empty, so all ``AMBER14_STANDARD`` residues participate in the - LJ/Coulomb gradients (atom matching is position-based, so tleap's - water ordering does not break the map) - - Monatomic ions (MG, ZN, CA, …) ARE included — covered by - ``leaprc.water.tip3p`` (Li/Merz 12-6 set), appear in fixed PDB - order, important for electrostatics near charged ligands - - Terminal atoms tleap regenerates (OXT …) excluded - - Note: uses ``AMBER14_STANDARD`` (not ``_MODELLER_FF_RESIDUES``) - so that ions absent from amber14-all.xml are still sent to tleap. - Index is preserved (original model.pdb row positions). - """ - pdb = self._chem_model.update_pdb() + pdb = self._chem_model.pdb.copy() + pdb[["x", "y", "z"]] = self._chem_model.xyz().detach().cpu().numpy() mask = pdb["altloc"].astype(str).str.strip().isin(["", "A"]) mask &= ~pdb["element"].astype(str).str.strip().isin(["H", "D"]) @@ -886,21 +607,7 @@ def _filter_pdb_for_tleap(self): def _build_omm_system( self, gaff2_params: Dict[str, Tuple[Path, Path]] ) -> Tuple: - """ - Build OpenMM system. Returns ``(system, omm_topology, pos_nm_array)``. - - Standard path (no non-standard residues) - ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - Filter model PDB → heavy atoms, primary conformation, standard residues. - Use ``openmm.app.Modeller.addHydrogens()`` to re-add H with AMBER names. - Create system with ``ForceField('amber14-all.xml')``. - - GAFF2 path (non-standard residues present) - ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - Write protein PDB (no OXT, no H) + mol2 per ligand. - Combine via tleap ``combine{}`` command → prmtop/inpcrd. - Load with parmed → ``AmberParm.createSystem()``. - """ + """Parameterise the model with AMBER14 or AMBER14/GAFF2.""" import openmm as mm # noqa: PLC0415 import openmm.app as app # noqa: PLC0415 import openmm.unit as unit # noqa: PLC0415 @@ -922,73 +629,57 @@ def _build_omm_system( return system, topology, pos_nm def _build_standard(self, cutoff_A: float, app, unit) -> Tuple: - """ - AMBER14 standard-residue path using gemmi + pdbfixer + OpenMM. - - gemmi writes proper chain termination / TER records so pdbfixer - can detect and fix missing terminal atoms (OXT). pdbfixer also - handles missing sidechain atoms and non-standard residue names. - """ - import gemmi # noqa: PLC0415 - from pdbfixer import PDBFixer # noqa: PLC0415 - from torchref.io import pdb as pdbio # noqa: PLC0415 + """Parameterise existing atoms, retaining PDB serials through name aliases.""" + from torchref.io import pdb as pdbio - # Standard path: Modeller preserves chain/resseq → use key-based mapping self._tleap_residue_map = None - - pdb_heavy = self._filter_pdb_for_omm(include_nonstandard=False) - - tmp = tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) - tmp2 = tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) - tmp.close() - tmp2.close() - try: - # Write via torchref, then re-read/write with gemmi to get - # proper chain breaks and TER records that pdbfixer needs. - pdbio.write(pdb_heavy, tmp.name) - st = gemmi.read_structure(tmp.name) - st.setup_entities() - st.assign_subchains() - st.write_pdb(tmp2.name) - - # pdbfixer: add missing terminal atoms and sidechain atoms - fixer = PDBFixer(filename=tmp2.name) - fixer.findMissingResidues() - fixer.missingResidues = {} # don't fill gaps - fixer.findMissingAtoms() - - if self.verbose >= 1: - n_missing = sum(len(v) for v in fixer.missingAtoms.values()) - n_terminals = sum( - 1 for v in fixer.missingTerminals.values() if v - ) - if n_missing or n_terminals: - print( - f"[AmberTarget] pdbfixer: {n_missing} missing atoms, " - f"{n_terminals} terminal fixes" - ) - - fixer.addMissingAtoms() - finally: - os.unlink(tmp.name) - os.unlink(tmp2.name) - + pdb = self._chem_model.pdb.copy() + xyz = self._chem_model.xyz().detach().cpu().numpy() + pdb[["x", "y", "z"]] = xyz + pdb["serial"] = np.arange(1, len(pdb) + 1) + with tempfile.TemporaryDirectory(prefix="torchref_amber_") as directory: + filename = str(Path(directory) / "model.pdb") + pdbio.write(pdb, filename) + parsed = app.PDBFile(filename) + + topology = parsed.topology + source_rows = np.array([int(a.id) - 1 for a in topology.atoms()]) + if len(source_rows) != len(pdb) or not np.array_equal( + np.sort(source_rows), np.arange(len(pdb)) + ): + raise ValueError( + "[AmberTarget] PDB atom identities are ambiguous or duplicated. " + "Every TorchRef atom must correspond to exactly one AMBER particle." + ) + self._source_model_rows = source_rows + # Preserve explicit covalent links that the PDB bond templates cannot + # infer. PDB parsing also resolves standard hydrogen-name aliases. + atoms = list(topology.atoms()) + inverse = np.argsort(source_rows) + existing_bonds = {tuple(sorted((a.index, b.index))) for a, b in topology.bonds()} + for i, j in self._chem_model.restraints.topology.atoms.bonds.indices.cpu().tolist(): + pair = tuple(sorted((int(inverse[i]), int(inverse[j])))) + if pair not in existing_bonds: + topology.addBond(atoms[pair[0]], atoms[pair[1]]) + existing_bonds.add(pair) ff = app.ForceField("amber14-all.xml", "amber14/tip3pfb.xml") - modeller = app.Modeller(fixer.topology, fixer.positions) - modeller.addHydrogens(ff) - - system = ff.createSystem( - modeller.topology, - nonbondedMethod=app.CutoffNonPeriodic, - nonbondedCutoff=cutoff_A * unit.angstrom, - constraints=None, - ) - - # positions in nm - pos_nm = np.array( - modeller.positions.value_in_unit(unit.nanometer), dtype=np.float64 - ) - return system, modeller.topology, pos_nm + try: + system = ff.createSystem( + topology, + nonbondedMethod=app.CutoffNonPeriodic, + nonbondedCutoff=cutoff_A * unit.angstrom, + constraints=None, + rigidWater=False, + ) + except ValueError as exc: + raise ValueError( + "[AmberTarget] TorchRef model is not AMBER-compatible. " + "Prepare missing atoms, terminal groups and protonation in the " + "model before constructing the loss; no atoms were added. " + f"OpenMM: {exc}" + ) from exc + xyz = self._chem_model.xyz().detach().cpu().numpy() + return system, topology, np.asarray(xyz[source_rows] * 0.1, dtype=np.float64) def _build_gaff2( self, @@ -997,19 +688,10 @@ def _build_gaff2( app, unit, ) -> Tuple: - """ - AMBER14 + GAFF2 path via tleap + parmed. - - All AMBER14-standard heavy atoms (protein, ions, waters — no OXT, no H, - no non-standard HETATM) plus each ligand mol2 are combined by tleap - ``combine{}``. parmed loads the resulting prmtop/inpcrd. - - Atom mapping uses position-based matching (see :meth:`_build_atom_map`): - tleap's initial coordinates are taken directly from the PDB we write, - so model and tleap positions agree to 3 decimal places (PDB precision), - making a KD-tree nearest-neighbour search unambiguous. This avoids - relying on tleap's residue-sequential numbering, which is fragile for - water molecules. + """Build AMBER14/GAFF2 templates with one copy per ligand instance. + + tleap may reorder or rename atoms; the complete map is validated before + its system is used. Runtime coordinates always come from TorchRef. """ import parmed as pmd # noqa: PLC0415 from torchref.io import pdb as pdbio # noqa: PLC0415 @@ -1019,12 +701,8 @@ def _build_gaff2( prot_pdb = work_dir / "protein.pdb" pdb_tleap = self._filter_pdb_for_tleap() - # Signal GAFF2 path to _build_atom_map (position-based + name fallback) - self._tleap_residue_map = True # type: ignore[assignment] - # Store GAFF2 resnames so _build_atom_map can do name-based fallback - # for ligand atoms (mol2 may have old coords if model was refined first) - self._gaff2_resnames: set = set(gaff2_params.keys()) - + # tleap does not preserve the original chain and residue identifiers. + self._tleap_residue_map = True pdbio.write(pdb_tleap.reset_index(drop=True), str(prot_pdb)) prmtop = work_dir / "complex.prmtop" @@ -1038,7 +716,18 @@ def _build_gaff2( lig_loads.append(f"{rn} = loadMol2 {mol2}") lig_names.append(rn) - combine_list = " ".join(["protein"] + lig_names) + ligand_copies = [] + ligand_keys = [] + pdb = self._chem_model.pdb + for rn in lig_names: + rows = pdb[pdb["resname"].astype(str).str.strip() == rn] + for key, _ in rows.groupby(["chainid", "resseq", "icode"], sort=False): + copy_name = f"ligand{len(ligand_copies)}" + lig_loads.append(f"{copy_name} = copy {rn}") + ligand_copies.append(copy_name) + ligand_keys.append(tuple(key)) + self._gaff2_residue_keys = ligand_keys + combine_list = " ".join(["protein"] + ligand_copies) tleap_script = "\n".join( [ "source leaprc.protein.ff14SB", @@ -1072,6 +761,7 @@ def _build_gaff2( nonbondedMethod=app.CutoffNonPeriodic, nonbondedCutoff=cutoff_A * unit.angstrom, constraints=None, + rigidWater=False, ) topology = combined.topology pos_nm = np.array( @@ -1093,478 +783,173 @@ def _is_hydrogen(omm_atom) -> bool: return omm_atom.name.startswith("H") # heuristic fallback def _build_atom_map(self) -> None: - """ - Build ``self._model_to_omm``: int32 array [n_model] where entry *i* - is the OpenMM atom index corresponding to model atom *i*, or -1 for - unmatched atoms (H atoms, altloc-B atoms, non-standard HETATM, …). - - Two strategies depending on how the system was built: - - **Standard path** (``_tleap_residue_map is None``): - OpenMM Modeller preserves chain IDs and residue numbers from the input - PDB, so matching uses the key ``(chain_id, resseq, icode, atom_name)``. - - **GAFF2 path** (``_tleap_residue_map is not None``): - tleap strips chain IDs and renumbers residues sequentially, making - name/number-based matching unreliable (especially for waters). - Instead, the tleap initial positions are taken from the exact - coordinates we wrote to the PDB (via ``update_pdb()``), so model and - tleap positions agree to within PDB precision (0.001 Å = 0.0001 nm). - A KD-tree nearest-neighbour search with a tight threshold (0.005 nm) - unambiguously identifies each tleap heavy atom's model counterpart. - """ - from scipy.spatial import cKDTree # noqa: PLC0415 - + """Require a bijection between model rows and all OpenMM particles.""" pdb = self._chem_model.pdb n_model = len(pdb) - model_to_omm = np.full(n_model, -1, dtype=np.int32) - + atoms = list(self._topology.atoms()) + mapping = np.full(n_model, -1, dtype=np.int32) if self._tleap_residue_map is None: - # ---- Standard path: match by (chain, resseq, icode, atom_name) ---- - # No altlocs at this point (checked in _build). - model_key_to_idx: Dict[Tuple, int] = {} - for i in range(n_model): - row = pdb.iloc[i] - key = ( - str(row["chainid"]).strip(), - int(row["resseq"]), - str(row.get("icode", "")).strip(), - str(row["name"]).strip(), - ) - model_key_to_idx[key] = i - - for omm_atom in self._topology.atoms(): - if self._is_hydrogen(omm_atom): - continue - chain_id = omm_atom.residue.chain.id.strip() - try: - resseq = int(omm_atom.residue.id) - except ValueError: - raw = omm_atom.residue.id.strip() - resseq = int(raw.rstrip("ABCDEFGHIJKLMNOPQRSTUVWXYZ") or "0") - icode = (omm_atom.residue.insertionCode or "").strip() - idx = model_key_to_idx.get( - (chain_id, resseq, icode, omm_atom.name.strip()) - ) - if idx is not None: - model_to_omm[idx] = omm_atom.index - + mapping[self._source_model_rows] = np.arange(len(atoms)) else: - # ---- GAFF2 path: position-based matching via KD-tree ---- - # Collect tleap heavy-atom positions (nm) and their indices. - tleap_pos_nm = self._tleap_pos_nm # set by _build() before this call - tleap_ha_omm_idx: List[int] = [] - tleap_ha_pos: List[np.ndarray] = [] - for omm_atom in self._topology.atoms(): - if not self._is_hydrogen(omm_atom): - tleap_ha_omm_idx.append(omm_atom.index) - tleap_ha_pos.append(tleap_pos_nm[omm_atom.index]) - - tleap_ha_pos_arr = np.array(tleap_ha_pos) # (N_tleap_heavy, 3) nm - tree = cKDTree(tleap_ha_pos_arr) - - # Collect model primary-altloc heavy-atom positions (nm) and indices. - # Use update_pdb() coords — same values that were written to tleap PDB. - fresh_pdb = self._chem_model.update_pdb() - altloc_ok = fresh_pdb["altloc"].astype(str).str.strip().isin(["", "A"]) - not_h = ~fresh_pdb["element"].astype(str).str.strip().isin(["H", "D"]) - primary_heavy = np.where((altloc_ok & not_h).values)[0] - - model_pos_nm = np.column_stack([ - fresh_pdb["x"].values[primary_heavy], - fresh_pdb["y"].values[primary_heavy], - fresh_pdb["z"].values[primary_heavy], - ]) * 0.1 # Å → nm - - # Match: threshold = 0.005 nm (50× PDB precision of 0.0001 nm) - dists, nn_idx = tree.query(model_pos_nm, k=1) - matched = dists < 0.005 - for local_i, (model_i, nn_i) in enumerate(zip(primary_heavy, nn_idx)): - if matched[local_i]: - model_to_omm[model_i] = tleap_ha_omm_idx[nn_i] - - # Name-based fallback for GAFF2 ligand residues whose mol2 positions - # differ from the current model (e.g. after refinement steps). - # The cached mol2 retains original antechamber coordinates, so a - # second AmberTarget init after LBFGS will have position shifts. - gaff2_resnames = getattr(self, "_gaff2_resnames", set()) - if gaff2_resnames: - # Build (resname, atom_name) → model positional index for primary heavy - lig_key_to_model: Dict[Tuple[str, str], int] = {} - for arr_pos in primary_heavy: - rn = str(fresh_pdb["resname"].values[arr_pos]).strip() - if rn not in gaff2_resnames: - continue - aname = str(fresh_pdb["name"].values[arr_pos]).strip() - lig_key_to_model[(rn, aname)] = int(arr_pos) - - for omm_atom in self._topology.atoms(): - if self._is_hydrogen(omm_atom): - continue - rn = omm_atom.residue.name - if rn not in gaff2_resnames: - continue - aname = omm_atom.name.strip() - model_arr_pos = lig_key_to_model.get((rn, aname)) - if model_arr_pos is not None and model_to_omm[model_arr_pos] < 0: - model_to_omm[model_arr_pos] = omm_atom.index - - # Warn about UNEXPECTED unmatched heavy atoms. - # Expected to be unmatched (silently skipped in gradient): - # - H / D atoms - # - Waters, ions excluded from tleap (_TLEAP_EXCLUDE_RESIDUES) - # - C-terminal OXT regenerated by tleap (_TLEAP_SKIP_ATOMS) - # - Alternate conformer atoms (altloc != '' and != 'A') - elem_col = pdb["element"].astype(str).str.strip() - altloc_col = pdb["altloc"].astype(str).str.strip() - resname_col = pdb["resname"].astype(str).str.strip() - name_col = pdb["name"].astype(str).str.strip() - heavy_mask = ~elem_col.isin(["H", "D"]) - # Residues in AMBER14_STANDARD but without an amber14-all.xml template: - # unmatched on the standard (Modeller) path; matched via tleap on GAFF2 path. - _no_modeller_template = AMBER14_STANDARD - _MODELLER_FF_RESIDUES - expected_mask = ( - # Waters always excluded from tleap; no AMBER gradient expected - resname_col.isin(_TLEAP_EXCLUDE_RESIDUES) | - # Ions that lack Modeller templates (matched in GAFF2 path, not standard) - resname_col.isin(_no_modeller_template) | - # tleap-regenerated terminal atoms (OXT etc.) - name_col.isin(_TLEAP_SKIP_ATOMS) | - # alternate conformers (altloc B, C, …) - (~altloc_col.isin(["", "A"])) - ) - unexpected_unmatched = np.where( - heavy_mask.values & ~expected_mask.values & (model_to_omm < 0) - )[0] - if len(unexpected_unmatched) > 0: - ex = [ - f"{pdb.iloc[i]['name'].strip()} " - f"({pdb.iloc[i]['resname'].strip()} {pdb.iloc[i]['resseq']})" - for i in unexpected_unmatched[:5] - ] - warnings.warn( - f"[AmberTarget] {len(unexpected_unmatched)} heavy model atom(s) " - f"could not be matched to OpenMM topology " - f"(e.g. {', '.join(ex)}). Their gradients will be zero.", - UserWarning, - stacklevel=3, - ) - elif self.verbose >= 2: - unmatched_heavy = int(heavy_mask.values.sum()) - int( - (heavy_mask.values & (model_to_omm >= 0)).sum() - ) - print( - f"[AmberTarget] {unmatched_heavy} heavy atoms have model_to_omm=-1 " - f"(expected: non-standard HETATM / altloc-B / OXT)" - ) - - self._model_to_omm = model_to_omm - self._n_omm_atoms = self._system.getNumParticles() - - if self.verbose >= 2: - matched = int((model_to_omm >= 0).sum()) - print( - f"[AmberTarget] atom map: {matched}/{n_model} model atoms matched " - f"({self._n_omm_atoms} total OpenMM atoms)" - ) - - # ------------------------------------------------------------------ - # Hydrogen re-attachment - # ------------------------------------------------------------------ - - def _build_h_attachment(self, pos_nm: np.ndarray) -> None: - """ - Build the local-frame placement table for every H atom. - - Each H is placed at construction-time according to OpenMM's - ``Modeller.addHydrogens`` output. We freeze that placement in a - local frame defined by the parent heavy atom and 2 reference - heavy atoms. At forward time, the H position is recomputed in - differentiable PyTorch from the current heavy positions: - - e1 = (n1 − p) / |n1 − p| - e2 = perp(n2 − p, e1) / |perp(n2 − p, e1)| - e3 = e1 × e2 - h = p + lx·e1 + ly·e2 + lz·e3 - - where (lx, ly, lz) = (h − p) · [e1, e2, e3] is captured once. - - Backward through this formula in PyTorch autograd produces the - exact local-frame Jacobian — so the force on H from OpenMM gets - correctly distributed across p, n1, n2, not just onto p. - - Reference-atom selection per H: - - parent ``p`` : the unique heavy atom bonded to H - - neighbor ``n1`` : any heavy atom bonded to ``p`` (≠ H) - - neighbor ``n2`` : another heavy atom bonded to ``p``; if - ``p`` has only one heavy neighbour, fall - back to a heavy atom bonded to ``n1`` - (i.e. walk one bond further out). - - Hs with no usable triple — extremely rare in real chemistry — - fall back to the legacy ``h = p + offset`` rigid translation - path, with ``h_frame_valid=False``. - """ - if not hasattr(self, "_topology") or self._topology is None: - self._h_idx = None - self._h_parent_idx = None - self._h_n1_idx = None - self._h_n2_idx = None - self._h_local_pos = None - self._h_frame_valid = None - self._h_offset = None - return - - # ---- Walk topology bonds. Build heavy-neighbor adjacency and - # parent map for H atoms in one pass. ------------------------------ - from collections import defaultdict - - def _is_h(atom) -> bool: - return atom.element is not None and atom.element.symbol == "H" - - parent_of_h: Dict[int, int] = {} - heavy_neighbors: Dict[int, list] = defaultdict(list) - for bond in self._topology.bonds(): - a, b = bond[0], bond[1] - a_is_h = _is_h(a) - b_is_h = _is_h(b) - if a_is_h and not b_is_h: - parent_of_h[a.index] = b.index - elif b_is_h and not a_is_h: - parent_of_h[b.index] = a.index - elif not a_is_h and not b_is_h: - heavy_neighbors[a.index].append(b.index) - heavy_neighbors[b.index].append(a.index) - # H-H bonds are nonsense; ignored. - - if not parent_of_h: - self._h_idx = None - self._h_parent_idx = None - self._h_n1_idx = None - self._h_n2_idx = None - self._h_local_pos = None - self._h_frame_valid = None - self._h_offset = None - return - - # ---- Resolve (parent, n1, n2) per H ---------------------------- - h_indices = sorted(parent_of_h.keys()) - n_h = len(h_indices) - p_arr = np.empty(n_h, dtype=np.int64) - n1_arr = np.empty(n_h, dtype=np.int64) - n2_arr = np.empty(n_h, dtype=np.int64) - valid = np.zeros(n_h, dtype=bool) - - for k, h in enumerate(h_indices): - p = parent_of_h[h] - p_arr[k] = p - neigh = heavy_neighbors.get(p, []) - if len(neigh) >= 2: - n1_arr[k] = neigh[0] - n2_arr[k] = neigh[1] - valid[k] = True - elif len(neigh) == 1: - n1 = neigh[0] - further = [ - j for j in heavy_neighbors.get(n1, []) if j != p - ] - if further: - n1_arr[k] = n1 - n2_arr[k] = further[0] - valid[k] = True - else: - n1_arr[k] = n1 - n2_arr[k] = -1 - valid[k] = False - else: - n1_arr[k] = -1 - n2_arr[k] = -1 - valid[k] = False - - # ---- Compute local-frame coordinates from initial positions ---- - h_pos = pos_nm[np.asarray(h_indices, dtype=np.int64)] - p_pos = pos_nm[p_arr] - local_pos = np.zeros((n_h, 3), dtype=np.float64) - eps = 1e-12 - - valid_idx = np.where(valid)[0] - if valid_idx.size > 0: - n1_pos = pos_nm[n1_arr[valid_idx]] - n2_pos = pos_nm[n2_arr[valid_idx]] - a = n1_pos - p_pos[valid_idx] - b = n2_pos - p_pos[valid_idx] - e1 = a / np.maximum( - np.linalg.norm(a, axis=-1, keepdims=True), eps, - ) - b_perp = b - (b * e1).sum(-1, keepdims=True) * e1 - e2 = b_perp / np.maximum( - np.linalg.norm(b_perp, axis=-1, keepdims=True), eps, - ) - e3 = np.cross(e1, e2) - offset_v = h_pos[valid_idx] - p_pos[valid_idx] - local_pos[valid_idx, 0] = (offset_v * e1).sum(-1) - local_pos[valid_idx, 1] = (offset_v * e2).sum(-1) - local_pos[valid_idx, 2] = (offset_v * e3).sum(-1) - - # Rigid fallback offset (used for !valid Hs only) - offset_all = h_pos - p_pos - - # Stash numpy arrays for the forward path (the autograd Function - # converts to torch tensors lazily on the model's device). - self._h_idx = np.asarray(h_indices, dtype=np.int64) - self._h_parent_idx = p_arr - self._h_n1_idx = n1_arr - self._h_n2_idx = n2_arr - self._h_local_pos = local_pos.astype(np.float64) - self._h_frame_valid = valid - self._h_offset = offset_all # legacy field, used only when !valid + self._map_gaff2_atoms(mapping, atoms) - if self.verbose >= 1: - n_frame = int(valid.sum()) - n_fallback = int((~valid).sum()) - print( - f"[AmberTarget] H-attachment: {n_h} H atoms — " - f"{n_frame} via local-frame placement, " - f"{n_fallback} via rigid fallback" - ) - - # ------------------------------------------------------------------ - # Differentiable PyTorch placement (called from AmberTarget.forward) - # ------------------------------------------------------------------ - - def _place_hydrogens(self, heavy_omm_xyz_nm: torch.Tensor) -> torch.Tensor: - """ - Compute H positions from the current heavy-atom OpenMM-order tensor. - - Uses the local-frame data captured in :meth:`_build_h_attachment`: - for each H, build an orthonormal frame from (parent, n1, n2) and - place the H at its captured local-frame coordinates. Hs with - ``h_frame_valid=False`` fall back to ``parent + h_offset``. - - Differentiable: backward through this function distributes the - H force across the parent + n1 + n2 reference atoms via the exact - local-frame Jacobian (handled by PyTorch autograd). - - Parameters - ---------- - heavy_omm_xyz_nm : (n_omm_total, 3) tensor in nm, OpenMM atom order. - H slot values are ignored — they will be overwritten in the - returned tensor. - - Returns - ------- - h_xyz_nm : (n_H, 3) tensor in nm. Empty if no Hs. - """ - if self._h_idx is None or self._h_idx.size == 0: - return torch.zeros( - (0, 3), - dtype=heavy_omm_xyz_nm.dtype, - device=heavy_omm_xyz_nm.device, - ) - - device = heavy_omm_xyz_nm.device - dtype = heavy_omm_xyz_nm.dtype - # Lazily cache tensor views on the right device/dtype. + mapped = mapping[mapping >= 0] + missing_model = np.flatnonzero(mapping < 0) + missing_omm = sorted(set(range(len(atoms))) - set(mapped.tolist())) + duplicate = len(np.unique(mapped)) != len(mapped) if ( - getattr(self, "_h_tensors_dev", None) != device - or getattr(self, "_h_tensors_dtype", None) != dtype + len(atoms) != self._system.getNumParticles() + or duplicate + or len(missing_model) + or missing_omm ): - self._h_parent_idx_t = torch.as_tensor( - self._h_parent_idx, dtype=torch.long, device=device, - ) - # For invalid frames clamp neighbor indices to 0 so the gather is - # safe; the value is masked out by `where` below. - n1 = np.where(self._h_n1_idx >= 0, self._h_n1_idx, 0) - n2 = np.where(self._h_n2_idx >= 0, self._h_n2_idx, 0) - self._h_n1_idx_t = torch.as_tensor(n1, dtype=torch.long, device=device) - self._h_n2_idx_t = torch.as_tensor(n2, dtype=torch.long, device=device) - self._h_local_pos_t = torch.as_tensor( - self._h_local_pos, dtype=dtype, device=device, - ) - self._h_offset_t = torch.as_tensor( - self._h_offset, dtype=dtype, device=device, - ) - self._h_frame_valid_t = torch.as_tensor( - self._h_frame_valid, dtype=torch.bool, device=device, + model_examples = [ + f"{pdb.iloc[i]['chainid']}:{pdb.iloc[i]['resseq']}:" + f"{pdb.iloc[i]['name']}" + for i in missing_model[:5] + ] + omm_examples = [ + f"{atoms[i].residue.name}:{atoms[i].residue.id}:{atoms[i].name}" + for i in missing_omm[:5] + ] + raise ValueError( + "[AmberTarget] AMBER atom mapping is not one-to-one: " + f"unmatched model atoms={model_examples}, " + f"unmatched AMBER atoms={omm_examples}, duplicate matches={duplicate}. " + "Prepare matching atoms and protonation in TorchRef before " + "constructing the loss." ) - self._h_tensors_dev = device - self._h_tensors_dtype = dtype - - return _place_hydrogens_local_frame( - heavy_omm_xyz_nm, - self._h_parent_idx_t, - self._h_n1_idx_t, - self._h_n2_idx_t, - self._h_local_pos_t, - self._h_frame_valid_t, - self._h_offset_t, + self._model_to_omm = mapping + self._n_omm_atoms = len(atoms) + inverse = np.argsort(mapping) + self.register_buffer( + "_omm_to_model", + torch.as_tensor(inverse, dtype=get_int_dtype(), device=self._chem_model.device), ) - def _compose_full_omm_xyz( - self, - heavy_model_xyz_ang: torch.Tensor, - ) -> torch.Tensor: - """ - Build the full OpenMM-order position tensor (heavy + H) in nm. - - Heavy model atoms are scattered into their OpenMM slots via - ``self._model_to_omm``. Unmatched OpenMM heavy slots (e.g. - non-standard residues without a model match) are filled from - the construction-time ``_pos_buf`` snapshot so they're at least - consistent. H slots are filled by :meth:`_place_hydrogens`. + def _map_gaff2_atoms(self, mapping: np.ndarray, atoms: list) -> None: + """Match residue instances using heavy anchors, then names and H parents.""" + from scipy.spatial import cKDTree - Differentiable through ``heavy_model_xyz_ang`` — autograd routes - gradients on H slots back to their parent/n1/n2 reference atoms. - """ - device = heavy_model_xyz_ang.device - dtype = heavy_model_xyz_ang.dtype - n_omm = self._n_omm_atoms - - # Lazy tensorize index maps. - if ( - getattr(self, "_omm_tensors_dev", None) != device - or getattr(self, "_omm_tensors_dtype", None) != dtype - ): - valid_np = self._model_to_omm >= 0 - self._model_valid_t = torch.as_tensor( - valid_np, dtype=torch.bool, device=device, - ) - self._model_valid_model_idx_t = torch.as_tensor( - np.where(valid_np)[0], dtype=torch.long, device=device, - ) - self._model_valid_omm_idx_t = torch.as_tensor( - self._model_to_omm[valid_np], dtype=torch.long, device=device, + pdb = self._chem_model.pdb + keys = [ + tuple(row) + for row in pdb[["chainid", "resseq", "icode"]].itertuples( + index=False, name=None ) - # Construction-time snapshot for unmatched heavy slots + initial Hs - self._pos_buf_t = torch.as_tensor( - self._pos_buf, dtype=dtype, device=device, + ] + groups = {} + for i, key in enumerate(keys): + groups.setdefault(key, []).append(i) + residues = list(self._topology.residues()) + ligand_keys = self._gaff2_residue_keys + ligand_residues = ( + residues[len(residues) - len(ligand_keys) :] if ligand_keys else [] + ) + residue_map = {res.index: key for res, key in zip(ligand_residues, ligand_keys)} + xyz_nm = self._chem_model.xyz().detach().cpu().numpy() * 0.1 + names = pdb["name"].astype(str).str.strip().to_numpy() + elements = pdb["element"].astype(str).str.strip().str.upper().to_numpy() + heavy_rows = np.flatnonzero(~np.isin(elements, ["H", "D"])) + tree = cKDTree(xyz_nm[heavy_rows]) + for residue in residues: + if residue.index in residue_map: + continue + candidates = set() + for atom in residue.atoms(): + if self._is_hydrogen(atom): + continue + for local in tree.query_ball_point(self._tleap_pos_nm[atom.index], 0.005): + row = heavy_rows[local] + if elements[row] == atom.element.symbol.upper(): + candidates.add(keys[row]) + if len(candidates) != 1: + raise ValueError( + f"[AmberTarget] Cannot uniquely identify AMBER residue " + f"{residue.name} {residue.id} in TorchRef." + ) + residue_map[residue.index] = candidates.pop() + if len(set(residue_map.values())) != len(residue_map): + raise ValueError( + "[AmberTarget] Multiple AMBER residues match one model residue." ) - self._omm_tensors_dev = device - self._omm_tensors_dtype = dtype - # Start from the construction-time snapshot (provides values for - # unmatched heavy atoms and any non-frame-placed atom). Heavy and - # H slots will be overwritten below. - full = self._pos_buf_t.clone() - - heavy_model_xyz_nm = heavy_model_xyz_ang * 0.1 - heavy_matched = heavy_model_xyz_nm.index_select( - 0, self._model_valid_model_idx_t, - ) - full = full.index_copy(0, self._model_valid_omm_idx_t, heavy_matched) - - # Now derive H positions from the fully-populated heavy tensor. - if self._h_idx is not None and self._h_idx.size > 0: - h_xyz = self._place_hydrogens(full) - if not hasattr(self, "_h_idx_t_for_omm"): - self._h_idx_t_for_omm = torch.as_tensor( - self._h_idx, dtype=torch.long, device=device, + model_parents = {} + graph = self._chem_model.restraints.topology.atoms + for i, j in graph.bonds.indices.cpu().tolist(): + if elements[i] in {"H", "D"} and elements[j] not in {"H", "D"}: + model_parents[i] = j + elif elements[j] in {"H", "D"} and elements[i] not in {"H", "D"}: + model_parents[j] = i + omm_parents = {} + for a, b in self._topology.bonds(): + if self._is_hydrogen(a) and not self._is_hydrogen(b): + omm_parents[a.index] = b.index + elif self._is_hydrogen(b) and not self._is_hydrogen(a): + omm_parents[b.index] = a.index + for residue in residues: + rows = groups.get(residue_map[residue.index], []) + by_name = {names[i]: i for i in rows} + if len(by_name) != len(rows): + raise ValueError( + "[AmberTarget] Duplicate atom names within a model residue." ) - elif self._h_idx_t_for_omm.device != device: - self._h_idx_t_for_omm = self._h_idx_t_for_omm.to(device) - full = full.index_copy(0, self._h_idx_t_for_omm, h_xyz) + for atom in residue.atoms(): + row = by_name.get(atom.name) + if row is None and atom.name in {"H", "H1"}: + row = by_name.get("H1" if atom.name == "H" else "H") + if row is None and not self._is_hydrogen(atom): + candidates = [ + i + for i in rows + if mapping[i] < 0 + and elements[i] == atom.element.symbol.upper() + and ( + len(rows) == 1 + or np.linalg.norm(xyz_nm[i] - self._tleap_pos_nm[atom.index]) + < 0.005 + ) + ] + if len(candidates) == 1: + row = candidates[0] + if row is not None: + symbol = "H" if elements[row] == "D" else elements[row] + if symbol == atom.element.symbol.upper() and mapping[row] < 0: + mapping[row] = atom.index + for row in rows: + if mapping[row] >= 0 and row in model_parents: + expected_parent = mapping[model_parents[row]] + if omm_parents.get(mapping[row]) != expected_parent: + raise ValueError( + "[AmberTarget] Hydrogen attachment differs between " + f"TorchRef and AMBER: {residue.name} {names[row]}." + ) + # Equivalent hydrogens may use different numbering conventions. Only + # pair remaining H atoms attached to the same already-mapped parent. + used = set(mapping[mapping >= 0].tolist()) + for row in rows: + if mapping[row] >= 0 or row not in model_parents: + continue + parent = mapping[model_parents[row]] + choices = sorted( + a.index + for a in residue.atoms() + if a.index not in used and omm_parents.get(a.index) == parent + ) + if choices: + mapping[row] = choices[0] + used.add(choices[0]) - return full + def _compose_full_omm_xyz(self, model_xyz_ang: torch.Tensor) -> torch.Tensor: + """Gather all Cartesian model coordinates into OpenMM order and nm.""" + if model_xyz_ang.shape != (self._n_model_atoms, 3): + raise ValueError( + "[AmberTarget] Atom count changed; rebuild the target after " + "changing model topology." + ) + if self._omm_to_model.device != model_xyz_ang.device: + self._omm_to_model = self._omm_to_model.to(model_xyz_ang.device) + return model_xyz_ang.index_select(0, self._omm_to_model) * 0.1 # ------------------------------------------------------------------ # Step 5 — OpenMM Context @@ -1614,27 +999,10 @@ def _build_context(self, pos_nm: np.ndarray) -> None: # ------------------------------------------------------------------ def _energy(self, xyz_ang: torch.Tensor) -> torch.Tensor: - """AMBER14 energy for one conformation's heavy-atom coords. - - Parameters - ---------- - xyz_ang : torch.Tensor - ``(n_model_atoms, 3)`` heavy-atom coordinates in Å, in the order - of ``self._chem_model.pdb`` (the topology the system was built on). - - Returns - ------- - torch.Tensor - Scalar energy in kJ/mol (or kJ/mol/atom if ``normalize_by_atoms``). - Gradient flows to ``xyz_ang`` via OpenMM analytical forces - (heavy atoms direct) and via :meth:`_place_hydrogens` / PyTorch - autograd (H positions, redistributed onto their parent + - local-frame neighbors). - - Notes - ----- - Subclasses feed per-member coordinates here; the single-molecule - :meth:`forward` passes ``self._model.xyz()``. + """Evaluate all-atom Cartesian coordinates in Å, in chemistry-model order. + + Return a scalar in kJ/mol, divided by the atom count when normalization + is enabled. Gradients flow through the model's own coordinate wrapper. """ if self._context is None: raise RuntimeError( diff --git a/torchref/experimental/targets/forcefield_target.py b/torchref/experimental/targets/forcefield_target.py index 27332457..09f0e1de 100644 --- a/torchref/experimental/targets/forcefield_target.py +++ b/torchref/experimental/targets/forcefield_target.py @@ -178,11 +178,11 @@ def forward(self) -> torch.Tensor: Z = self.model.Z # Shape: (n_atoms,) # Ensure Z is long tensor - if Z.dtype != torch.long: + if Z.dtype != torch.long: # dtype-ok: dtype guard comparison against torch.long, not an allocation Z = Z.long() # Create batch tensor (single structure = all zeros) - batch = torch.zeros(len(Z), dtype=torch.long, device=xyz.device) + batch = torch.zeros(len(Z), dtype=torch.long, device=xyz.device) # dtype-ok: batch index tensor for TorchMD-Net graph scatter; PyTorch requires int64 # Compute energy via TorchMD-Net # Returns (energy, forces) or just energy depending on model config diff --git a/torchref/experimental/targets/realspace.py b/torchref/experimental/targets/realspace.py index 0e2a85f6..2024b595 100644 --- a/torchref/experimental/targets/realspace.py +++ b/torchref/experimental/targets/realspace.py @@ -19,8 +19,7 @@ import torch from torchref.base.reciprocal.grid_operations import place_on_grid -from torchref.symmetry.grid_utils import calculate_optimal_grid_size -from torchref.symmetry.reciprocal_symmetry import expand_hkl +from torchref.symmetry import SpaceGroup from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -112,20 +111,12 @@ def __init__( # Caches (not registered as buffers since they're lazily computed) self._data_p1 = None self._molecular_mask = None - self._gridsize = None # P1 expansion cache (ASU → P1 mapping) self._hkl_p1 = None self._p1_indices = None self._p1_phase_shifts = None - def _ensure_grid(self): - """Ensure model's SfFFT grid is set up.""" - if self._model is None: - raise RuntimeError("No model set for RealSpaceTarget") - if self._model.real_space_grid is None: - self._model.setup_grid() - def _get_data_p1(self) -> "ReflectionData": """Return P1-expanded ReflectionData, cached after first call.""" if self._data_p1 is None: @@ -136,9 +127,9 @@ def _ensure_p1_expansion(self): """Compute and cache the ASU → P1 expansion mapping.""" if self._hkl_p1 is not None: return - hkl_p1, indices, phase_shifts = expand_hkl( + sg = self._data.spacegroup or SpaceGroup("P1") + hkl_p1, indices, phase_shifts = sg.expand_hkl( self._data.hkl, - self._data.spacegroup or "P1", include_friedel=True, remove_absences=True, device=self._data.hkl.device, @@ -154,19 +145,11 @@ def _expand_to_p1(self, fcalc: torch.Tensor) -> torch.Tensor: return fcalc_p1 * torch.exp(1j * self._p1_phase_shifts) def _get_gridsize(self) -> Tuple[int, int, int]: - """ - Get grid size for map computation. - - Uses the model's FFT grid size to ensure compatibility with - the molecular mask (which is built on the model's grid). - """ - if self._gridsize is not None: - return self._gridsize - - self._ensure_grid() - gs = self._model.fft.gridsize - self._gridsize = tuple(int(x) for x in gs) - return self._gridsize + """Grid size for map computation: the model's, so it matches the + molecular mask built on the model's grid.""" + if self._model is None: + raise RuntimeError("No model set for RealSpaceTarget") + return self._model.fft.grid_shape def _compute_observed_map(self) -> torch.Tensor: """ @@ -244,7 +227,6 @@ def _build_molecular_mask(self): """ from torchref.scaling.solvent import SolventModel - self._ensure_grid() with torch.no_grad(): solvent = SolventModel( @@ -624,9 +606,8 @@ def _ensure_p1_expansion(self): return spacegroup = self._data_light.spacegroup - hkl_p1, indices, phase_shifts = expand_hkl( + hkl_p1, indices, phase_shifts = spacegroup.expand_hkl( self._hkl, - spacegroup, include_friedel=True, remove_absences=True, device=self._hkl.device, diff --git a/torchref/experimental/targets/sampled_ml_phase_target.py b/torchref/experimental/targets/sampled_ml_phase_target.py index 59575c8d..e505240b 100644 --- a/torchref/experimental/targets/sampled_ml_phase_target.py +++ b/torchref/experimental/targets/sampled_ml_phase_target.py @@ -129,7 +129,7 @@ def __init__( self.name = "xray_sampled_ml_work" if use_work_set else "xray_sampled_ml_test" # Register tunable parameters as buffers for state_dict access - self.register_buffer("_n_samples", torch.tensor(n_samples, dtype=torch.int64)) + self.register_buffer("_n_samples", torch.tensor(n_samples, dtype=torch.int64)) # dtype-ok: scalar sample-count buffer; categorical count, not model-precision data self.register_buffer("_sigma_model_log", torch.tensor(sigma_model_log)) self.register_buffer("_use_analytical", torch.tensor(use_analytical)) self.register_buffer("_use_antithetic", torch.tensor(use_antithetic)) @@ -545,7 +545,7 @@ def __init__( self.add_module("_scaler_dark", scaler_dark) # Tunable parameters as buffers - self.register_buffer("_n_samples", torch.tensor(n_samples, dtype=torch.int64)) + self.register_buffer("_n_samples", torch.tensor(n_samples, dtype=torch.int64)) # dtype-ok: scalar sample-count buffer; categorical count, not model-precision data self.register_buffer("_sigma_model_log", torch.tensor(sigma_model_log)) self.use_work_set = use_work_set diff --git a/torchref/io/__init__.py b/torchref/io/__init__.py index a20af5b7..79c69f9c 100644 --- a/torchref/io/__init__.py +++ b/torchref/io/__init__.py @@ -21,31 +21,33 @@ RestraintCIFReader, ) -# Metadata -from .metadata import RefinementMetadata - -# Top-level object-creation readers -from .readers import read_cif, read_mtz, read_pdb - # Dataset classes (primary API) from .datasets import ( CrystalDataset, DatasetCollection, - ReflectionData, FcalcDataset, + ReflectionData, + ScaledDataset, ) +# IHM ensemble support (mapping always available; reader/writer need python-ihm) +from .ihm_mapping import IHMEnsembleMapping, IHMModelGroupInfo, IHMStateInfo + +# Metadata +from .metadata import RefinementMetadata + # Reader classes (from format modules) from .mtz import MTZReader from .pdb import PDBReader -# IHM ensemble support (mapping always available; reader/writer need python-ihm) -from .ihm_mapping import IHMEnsembleMapping, IHMModelGroupInfo, IHMStateInfo +# Top-level object-creation readers +from .readers import read_cif, read_mtz, read_pdb __all__ = [ # Primary API - Datasets "CrystalDataset", "ReflectionData", + "ScaledDataset", "DatasetCollection", "FcalcDataset", # Top-level readers diff --git a/torchref/io/cif.py b/torchref/io/cif.py index 96bccc10..d616abcb 100644 --- a/torchref/io/cif.py +++ b/torchref/io/cif.py @@ -257,12 +257,33 @@ def dataframe_to_gemmi_structure(df, cell, spacegroup): return st -def _add_refine_categories(doc, metadata): - """Inject a :class:`RefinementMetadata`'s categories into ``doc``, in place.""" +def _cif_value(val) -> str: + """Render one value as a CIF token, quoting it when it needs quoting. + + The unset markers ``?`` and ``.`` are passed through bare: ``gemmi.cif.quote`` + would turn them into the quoted one-character strings ``'?'`` and ``'.'``, + which are data rather than nulls. Everything else goes through ``quote`` -- + an unquoted value containing whitespace silently splits into extra loop + columns when the file is read back. + """ import gemmi + text = str(val) + if text in ("?", "."): + return text + return gemmi.cif.quote(text) + + +def _add_refine_categories(doc, metadata): + """Inject a :class:`RefinementMetadata`'s categories into ``doc``, in place. + + Returns the set of category prefixes written (e.g. ``{"_refine.", + "_software."}``) so the caller can avoid copying the same categories in + again from another block and either clobbering or duplicating them. + """ block = doc.sole_block() cats = metadata.render_cif_categories() + written = set() for cat_name, items in cats.items(): # List values mean a loop category rather than key-value pairs. @@ -274,6 +295,7 @@ def _add_refine_categories(doc, metadata): prefix = tags[0].rsplit(".", 1)[0] + "." suffixes = [t.split(".")[-1] for t in tags] loop = block.init_loop(prefix, suffixes) + written.add(prefix) # All list values should have same length n_rows = max(len(v) for v in items.values() if isinstance(v, list)) for i in range(n_rows): @@ -281,13 +303,17 @@ def _add_refine_categories(doc, metadata): for tag in tags: val = items[tag] if isinstance(val, list): - row.append(str(val[i]) if i < len(val) else "?") + cell = val[i] if i < len(val) else "?" else: - row.append(str(val)) + cell = val + row.append(_cif_value(cell)) loop.add_row(row) else: for key, val in items.items(): - block.set_pair(key, gemmi.cif.quote(str(val))) + block.set_pair(key, _cif_value(val)) + written.add(key.rsplit(".", 1)[0] + ".") + + return written def write_model(df, filepath: str, metadata=None) -> None: @@ -329,7 +355,12 @@ def write_model(df, filepath: str, metadata=None) -> None: "_symmetry.space_group_name_H-M", gemmi.cif.quote(str(spacegroup)) ) - _add_refine_categories(meta_doc, metadata) + written = _add_refine_categories(meta_doc, metadata) + # Categories we just wrote from metadata, plus the two written above. + # The structure block gemmi builds from the DataFrame carries its own + # version of some of these; copying those in would clobber a pair or + # append a second loop for the same category. + written |= {"_cell.", "_symmetry."} st = dataframe_to_gemmi_structure(df, cell, spacegroup) struct_doc = st.make_mmcif_document() @@ -342,14 +373,15 @@ def write_model(df, filepath: str, metadata=None) -> None: tags = list(loop.tags) suffixes = [t.split(".")[-1] for t in tags] prefix = tags[0].rsplit(".", 1)[0] + "." + if prefix in written: + continue new_loop = meta_block.init_loop(prefix, suffixes) for row_idx in range(loop.length()): row = [loop[row_idx, col] for col in range(loop.width())] new_loop.add_row(row) elif item.pair is not None: tag, val = item.pair - # Skip cell/symmetry - already added - if not tag.startswith(("_cell.", "_symmetry.")): + if tag.rsplit(".", 1)[0] + "." not in written: meta_block.set_pair(tag, val) meta_doc.write_file(filepath) diff --git a/torchref/io/cif_readers.py b/torchref/io/cif_readers.py index 44d1b9e5..221bb955 100644 --- a/torchref/io/cif_readers.py +++ b/torchref/io/cif_readers.py @@ -12,6 +12,8 @@ import gemmi import numpy as np +import warnings + import pandas as pd #: Column holding the ``data_`` block a loop row was read from. Added only when @@ -1265,7 +1267,11 @@ def _get_value(self, data, possible_keys: List[str], default: Any = None) -> Any class ModelCIFReader: """ Reader for model/structure CIF files (e.g. ``*.cif`` from the PDB): - coordinates, altlocs, ANISOU, cell and space group. + coordinates, altlocs, ANISOU, cell, space group and covalent links. + + Covalent and metal ``_struct_conn`` rows are exposed as ``.links`` in the same table + the PDB reader builds from LINK records, so :meth:`torchref.model.model.Model.load` + picks them up either way. Calling the instance gives the same unpack order as the PDB reader:: @@ -1305,6 +1311,7 @@ def _extract_data(self): self.cell = cell_params self.spacegroup = self.get_space_group() + self.links = self.get_link_records() # Store as DataFrame attributes (like legacy PDB reader) self.dataframe.attrs["cell"] = self.cell @@ -1316,6 +1323,81 @@ def _extract_data(self): print(f" Atoms: {len(self.dataframe)}") print(f" Cell: {self.cell}") print(f" Spacegroup: {self.spacegroup}") + print(f" Links: {len(self.links)}") + + #: ``_struct_conn.conn_type_id`` prefixes that describe a covalent bond the topology + #: should carry. Disulfides are detected from SG-SG distance instead (the PDB reader + #: ignores SSBOND the same way); hydrogen bonds, salt bridges and mismatches are not + #: bonds. + _LINK_CONN_TYPES = ("covale", "metalc") + + #: Symmetry operators under which a ``_struct_conn`` row joins atoms of the same + #: asymmetric unit copy; the PDB reader keeps LINK records with ``1555`` or blank. + _LINK_SYMMETRY_OK = frozenset({"1_555", "", "?", "."}) + + def get_link_records(self) -> pd.DataFrame: + """Covalent and metal links from ``_struct_conn``, in the LINK-record table. + + Same columns as :func:`torchref.io.pdb.extract_link_records`. Rows whose + connection type is not covalent or metal, that cross a symmetry operator, or + whose residue numbers are unreadable are dropped. Blank alternative locations and + insertion codes (``?`` or ``.``) become empty strings, which is what the atom + table carries and what the LINK lookup compares against. + + Returns + ------- + pandas.DataFrame + Empty, with the LINK columns, when the file has no ``_struct_conn`` loop. + """ + from torchref.io.pdb import LINK_COLUMNS + + empty = pd.DataFrame(columns=list(LINK_COLUMNS)) + conn = self.cif.data.get("struct_conn") + if conn is None or len(conn) == 0: + return empty + + def column(names, default=""): + for name in names: + if name in conn.columns: + values = conn[name].astype(str).str.strip() + return values.where(~values.isin(["?", "."]), default) + return pd.Series([default] * len(conn), index=conn.index, dtype=object) + + kind = column(["_struct_conn.conn_type_id"]).str.lower() + keep = kind.str.startswith(self._LINK_CONN_TYPES) + for side in ("1", "2"): + keep &= column([f"_struct_conn.ptnr{side}_symmetry"]).isin( + self._LINK_SYMMETRY_OK + ) + + out = pd.DataFrame(index=conn.index) + for side in ("1", "2"): + ptnr = f"_struct_conn.ptnr{side}_" + pdbx = f"_struct_conn.pdbx_ptnr{side}_" + # Same precedence as get_atom_data: label_* for atom and residue names, + # auth_* for chain and residue number, so the lookup matches the table. + out[f"name{side}"] = column([ptnr + "label_atom_id", ptnr + "auth_atom_id"]) + out[f"altloc{side}"] = column([pdbx + "label_alt_id"]) + out[f"resname{side}"] = column([ptnr + "label_comp_id", ptnr + "auth_comp_id"]) + out[f"chainid{side}"] = column([ptnr + "auth_asym_id", ptnr + "label_asym_id"]) + out[f"resseq{side}"] = pd.to_numeric( + column([ptnr + "auth_seq_id", ptnr + "label_seq_id"], default="nan"), + errors="coerce", + ) + out[f"icode{side}"] = column([pdbx + "PDB_ins_code"]) + out["length"] = pd.to_numeric( + column(["_struct_conn.pdbx_dist_value"], default="nan"), errors="coerce" + ) + + keep &= out["resseq1"].notna() & out["resseq2"].notna() + out = out.loc[keep].copy() + if len(out) == 0: + return empty + out["resseq1"] = out["resseq1"].astype(int) + out["resseq2"] = out["resseq2"].astype(int) + if self.verbose > 1: + print(f"_struct_conn: kept {len(out)} of {len(conn)} rows as links") + return out[list(LINK_COLUMNS)].reset_index(drop=True) def read(self, filepath: str = None): """Re-read ``filepath`` (default: the init path); returns ``self``.""" @@ -1445,7 +1527,59 @@ def get_atom_data(self) -> pd.DataFrame: "_atom_site.aniso_U[2][3]", ] - if all(col in atom_df.columns for col in aniso_cols): + # The standard mmCIF home for anisotropic ADPs is the SEPARATE + # ``_atom_site_anisotrop`` loop, keyed by ``.id`` against ``_atom_site.id``. + # Only the legacy in-line ``_atom_site.aniso_U[i][j]`` form was read here, so a + # standards-conforming file -- every PDB-REDO entry, and anything the PDB emits + # as mmCIF -- silently loaded with no anisotropy at all and every atom marked + # isotropic. + aniso_df = getattr(self.cif, "data", {}).get("atom_site_anisotrop") + std_cols = [f"_atom_site_anisotrop.U[{i}][{j}]" + for i, j in ((1, 1), (2, 2), (3, 3), (1, 2), (1, 3), (2, 3))] + key, atom_key = "_atom_site_anisotrop.id", "_atom_site.id" + joined = None + if ( + aniso_df is not None + and all(c in aniso_df.columns for c in std_cols) + and key in aniso_df.columns + and atom_key in atom_df.columns + ): + # Join on the id as a STRING. Coercing to a number first silently produces + # NaN keys for any non-integer id and then mis-pairs U tensors with atoms, + # which is far worse than having no anisotropy: the model is scrambled but + # still refines. + left = pd.DataFrame({"_k": atom_df[atom_key].astype(str).str.strip()}) + right = aniso_df[[key] + std_cols].copy() + right["_k"] = right[key].astype(str).str.strip() + right = right.drop_duplicates("_k") + merged = left.merge(right, on="_k", how="left") + if len(merged) == len(atom_df): + joined = merged + + if joined is not None: + for name, col in zip(("u11", "u22", "u33", "u12", "u13", "u23"), std_cols): + result[name] = pd.to_numeric(joined[col].to_numpy(), errors="coerce") + result["anisou_flag"] = ~pd.isna(result["u11"]) + n_hit = int(result["anisou_flag"].sum()) + frac = n_hit / max(len(aniso_df), 1) + if frac < 0.9: + # Partial coverage is legitimate in small amounts -- waters and + # hydrogens often carry no ANISOU -- but a large shortfall means the two + # loops are not labelled the same way, and then the rows that DID match + # cannot be trusted to have matched the right atoms. Drop the anisotropy + # rather than apply a possibly mis-paired subset: an isotropic model is + # merely less informative, a scrambled one still refines and is wrong. + warnings.warn( + f"{self.filepath}: matched only {n_hit} of {len(aniso_df)} " + "anisotropic records to atoms, so _atom_site.id and " + "_atom_site_anisotrop.id do not agree; discarding the anisotropy " + "and loading isotropically.", + RuntimeWarning, + ) + for name in ("u11", "u22", "u33", "u12", "u13", "u23"): + result[name] = np.nan + result["anisou_flag"] = False + elif all(col in atom_df.columns for col in aniso_cols): result["u11"] = pd.to_numeric( atom_df["_atom_site.aniso_U[1][1]"], errors="coerce" ) @@ -1675,6 +1809,12 @@ def has_anisotropic_data(self) -> bool: "_atom_site.aniso_U[2][2]", "_atom_site.aniso_U[3][3]", ] + aniso_df = self.cif.data.get("atom_site_anisotrop") + if aniso_df is not None and all( + f"_atom_site_anisotrop.U[{i}][{j}]" in aniso_df.columns + for i, j in ((1, 1), (2, 2), (3, 3), (1, 2), (1, 3), (2, 3)) + ): + return True return all(col in self.cif.data["atom_site"].columns for col in aniso_cols) def get_coordinates(self) -> Optional[np.ndarray]: @@ -2157,6 +2297,15 @@ def _standardize_atoms(self, df: pd.DataFrame) -> pd.DataFrame: ), errors="coerce", ) + # The CCP4 energy type (NH1, OC, CH3, ...) keys the contact radii and the + # hydrogen-bond donor/acceptor roles; absent from eLBOW/Grade dictionaries. + type_cols = ["type_energy", "_chem_comp_atom.type_energy"] + if any(col in df.columns for col in type_cols): + result["type_energy"] = ( + self._extract_col(df, type_cols).astype(str).str.strip() + ) + else: + result["type_energy"] = "" # Include x,y,z if present (for ideal coordinates) for coord in ["x", "y", "z"]: diff --git a/torchref/io/datasets/__init__.py b/torchref/io/datasets/__init__.py index 86fb8f0f..65e4d9da 100644 --- a/torchref/io/datasets/__init__.py +++ b/torchref/io/datasets/__init__.py @@ -2,7 +2,8 @@ Crystallographic dataset containers. - :class:`CrystalDataset` -- base dataclass: fields, device moves, save/load -- :class:`ReflectionData` -- one crystal's observed reflections +- :class:`ReflectionData` -- one crystal's raw observed reflections +- :class:`ScaledDataset` -- live observations corrected by a shared DatasetScaler - :class:`FcalcDataset` -- calculated structure factors on a generated HKL set - :class:`DatasetCollection` -- several ReflectionData on one common HKL grid """ @@ -11,10 +12,12 @@ from .collection import DatasetCollection from .fcalc_data import FcalcDataset from .reflection_data import ReflectionData +from .scaled_dataset import ScaledDataset __all__ = [ "CrystalDataset", "ReflectionData", + "ScaledDataset", "FcalcDataset", "DatasetCollection", ] diff --git a/torchref/io/datasets/base.py b/torchref/io/datasets/base.py index 668290f3..5935fac3 100644 --- a/torchref/io/datasets/base.py +++ b/torchref/io/datasets/base.py @@ -9,13 +9,13 @@ import warnings from dataclasses import dataclass, field, fields -from typing import TYPE_CHECKING, Any, Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Optional import gemmi import torch -from torchref.config import get_default_device, normalize_device -from torchref.symmetry import Cell +from torchref.config import get_default_device, get_float_dtype, normalize_device +from torchref.symmetry import Cell, SpaceGroup from torchref.utils.device_mixin import DeviceMovementMixin if TYPE_CHECKING: @@ -67,13 +67,6 @@ class CrystalDataset(DeviceMovementMixin): # explicit Bijvoet pairs (separate signed-HKL rows). Gates the model's f'' term. friedel_merged: bool = True - # === E-value and anisotropy correction fields === - E: Optional[torch.Tensor] = None # E-values (N,) - E_squared: Optional[torch.Tensor] = None # E² values (N,) - F_squared_corrected: Optional[torch.Tensor] = None # Anisotropy-corrected F² (N,) - U_aniso: Optional[torch.Tensor] = None # Fitted anisotropy parameters (6,) - radial_shell_indices: Optional[torch.Tensor] = None # Shell assignments (N,) - # === Unit cell and symmetry === cell: Optional[Cell] = None # Cell object with [a, b, c, alpha, beta, gamma] spacegroup: Optional[str] = None # Space group name string @@ -125,21 +118,27 @@ def _tensor_fields(self): # ========== SERIALIZATION ========== def _get_state(self) -> Dict[str, Any]: - """State dict of all fields, tensors on CPU, cell/device/spacegroup - flattened to tensor/str, plus a ``"masks"`` entry. + """Return observation fields and masks with tensors on CPU. + + Cell/device/space group are flattened to tensors or strings. Loading + provenance and French-Wilson conversion caches are omitted. """ state = {} for f in fields(self): + if f.name in {"source", "reader", "dataset", "_FrenchWilson"}: + # Loading provenance and conversion caches are not observation state. + state[f.name] = None + continue val = getattr(self, f.name) if isinstance(val, torch.Tensor): - state[f.name] = val.cpu() + state[f.name] = val.detach().cpu() elif f.name == "cell" and val is not None: state[f.name] = val.data.cpu() elif f.name == "device": state[f.name] = str(val) elif f.name == "spacegroup" and val is not None: - state[f.name] = val.xhm() # Extended Hermann-Mauguin + state[f.name] = val.xhm # Extended Hermann-Mauguin else: state[f.name] = val # Masks are not a dataclass field, so handle them separately. @@ -183,10 +182,16 @@ def _from_state(cls, state: Dict[str, Any], device=None) -> "CrystalDataset": if "device" in state: state["device"] = torch.device(state["device"]) - # Spacegroup stays a string here; subclasses that want an object rewrap. + if isinstance(state.get("spacegroup"), str): + state["spacegroup"] = SpaceGroup(state["spacegroup"], device=device) if "cell" in state and state["cell"] is not None: if isinstance(state["cell"], torch.Tensor): - state["cell"] = Cell(state["cell"], dtype=torch.float32, device=device) + # Conform the reloaded cell to the config float dtype rather than + # pinning float32: a dataset saved and reloaded under a float64 + # config otherwise carries a float32 cell into reciprocal-basis math. + state["cell"] = Cell( + state["cell"], dtype=get_float_dtype(), device=device + ) obj = cls(**state) diff --git a/torchref/io/datasets/collection.py b/torchref/io/datasets/collection.py index 3e85c923..9418272b 100644 --- a/torchref/io/datasets/collection.py +++ b/torchref/io/datasets/collection.py @@ -7,12 +7,16 @@ """ from dataclasses import dataclass, field -from typing import Dict, Iterator, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple import torch from .base import CrystalDataset from .reflection_data import ReflectionData +from .scaled_dataset import ScaledDataset + +if TYPE_CHECKING: + from torchref.scaling import DatasetScaler @dataclass @@ -20,9 +24,9 @@ class DatasetCollection(CrystalDataset): """ Container for multiple related crystal datasets on a common HKL set. - Members are expanded in place onto the reference dataset's HKL grid - (:meth:`ReflectionData.validate_hkl`) and moved to the collection's device, - so adding a dataset MUTATES it. Dict-like access via ``[]``, ``keys()``, + Members are copied onto the union HKL grid without changing input datasets. + The reference supplies cell and space-group metadata, not a fixed scale. + ``scale()`` installs ScaledDataset members backed by one shared scaler. Dict-like access via ``[]``, ``keys()``, ``values()``, ``items()``, ``get()``, and iteration yields ``(name, dataset)`` in insertion order. @@ -55,55 +59,73 @@ class DatasetCollection(CrystalDataset): _cell: Optional[torch.Tensor] = field(default=None, repr=False) _spacegroup: Optional[str] = field(default=None, repr=False) _resolution: Optional[torch.Tensor] = field(default=None, repr=False) - _scale_factors: Dict[str, torch.Tensor] = field(default_factory=dict, repr=False) + scaler: Optional["DatasetScaler"] = field(default=None, repr=False) + scaling_metrics: dict = field(default_factory=dict, repr=False) def add_dataset( self, name: str, dataset: ReflectionData, set_as_reference: bool = False ) -> "DatasetCollection": - """ - Add a dataset, expanding it onto the reference HKL grid **in place**. + """Add a copied dataset and rebuild the union reflection grid. Parameters ---------- name : str - Identifier for this dataset. + Unique member name. dataset : ReflectionData - The dataset to add. - set_as_reference : bool, optional - If True, this dataset's HKL becomes the reference. The first dataset - added becomes the reference regardless. + Raw or scaled observations; scaled inputs contribute their raw values. + set_as_reference : bool + Use this dataset's cell and symmetry as collection metadata. Returns ------- DatasetCollection - Self, for method chaining. - - Raises - ------ - ValueError - If a dataset with the same name already exists. + Self. Membership changes discard the fitted joint scaler; call scale() + again to fit all members. Existing raw inputs are never mutated. """ if name in self._datasets: raise ValueError(f"Dataset '{name}' already exists in collection") - - if len(self._datasets) == 0 or set_as_reference: + members = { + k: d.raw_data() if isinstance(d, ScaledDataset) else d + for k, d in self._datasets.items() + } + raw = dataset.raw_data() if isinstance(dataset, ScaledDataset) else dataset + members[name] = raw.__select__(torch.arange(len(raw), device=raw.device)) + members[name].source = None + members[name].spacegroup = raw.spacegroup.copy() + if ( + len({d.spacegroup.xhm for d in members.values()}) != 1 + or len({d.friedel_merged for d in members.values()}) != 1 + ): + raise ValueError( + "Datasets require compatible symmetry and Friedel conventions" + ) + if not self._dataset_order or set_as_reference: self._reference_dataset = name - self._common_hkl = dataset.hkl.clone() - if dataset.cell is not None: - self._cell = dataset.cell.clone() - self._spacegroup = dataset.spacegroup - - if self._common_hkl is not None and dataset.hkl is not None: - dataset.validate_hkl(self._common_hkl) - - dataset.to(self.device) - - self._datasets[name] = dataset + self._cell = raw.cell.clone() if raw.cell is not None else None + self._spacegroup = members[name].spacegroup self._dataset_order.append(name) - - if self.verbose > 0: - print(f"Added dataset '{name}' ({len(dataset)} reflections)") - + union_hkl = torch.unique( + torch.cat( + [ + (d.hkl if d.friedel_merged else d._hkl_for_sf()).to(self.device) + for d in members.values() + ] + ), + dim=0, + ) + identity_hkl = None + if raw.friedel_merged: + self._common_hkl = union_hkl + else: + canonical, _, _, order = raw.spacegroup.canonicalize_hkl(union_hkl) + self._common_hkl = canonical + identity_hkl = union_hkl[order] + for data in members.values(): + data.to(self.device) + data.validate_hkl(self._common_hkl, identity_hkl=identity_hkl) + self._datasets = members + self.scaler = None + self.scaling_metrics = {} return self @property @@ -245,74 +267,243 @@ def harmonize_partition( return self def __call__(self, mask: bool = True) -> Dict[str, Tuple]: - """ - Return all datasets' data scaled if scale factors are set. + """Return full observation arrays for every collection member. Parameters ---------- mask : bool, optional - Whether to apply masking. Default is True. + Wrap amplitudes and uncertainties in detached MaskedTensors carrying + validity masks. If False, return live observation tensors. Returns ------- dict - Name -> ``(hkl, F, F_sigma, rfree)``. Routed through the deprecated - ``ReflectionData.__call__``, so F/F_sigma are MaskedTensors and each - member emits a DeprecationWarning. + Name to (HKL, amplitudes, sigmas, work flags). HKL has shape (H, 3); + other arrays have shape (H,) in each dataset's observation units. """ - return {name: ds(mask=mask, scale=True) for name, ds in self} + from torch.masked import MaskedTensor + + result = {} + for name, data in self: + amplitudes, sigmas = data.F, data.F_sigma + if mask: + valid = data.masks() + if not bool(valid.any()): + raise ValueError(f"Dataset {name!r} has no valid observations") + amplitudes = MaskedTensor(amplitudes.detach().clone(), valid) + if sigmas is not None: + sigmas = MaskedTensor(sigmas.detach().clone(), valid) + result[name] = data.hkl, amplitudes, sigmas, data.rfree_flags + return result + + def scale(self, nsteps: int = 10, max_iter: int = 100) -> "DatasetCollection": + """Jointly scale observations and expose live ScaledDataset members. - def scale(self): + Parameters + ---------- + nsteps, max_iter : int + Outer steps and per-step iteration limit for the dedicated scaler. + + Returns + ------- + DatasetCollection + Self. Retrieve scaled observations from this collection; references + to original inputs remain raw. Repeated calls reuse the parameter owner. """ - Least-squares fit every non-reference dataset's scale and anisotropy - onto the reference, whose own parameters are left untouched. + from torchref.scaling.dataset_scaler import DatasetScaler + + raw = { + k: d.raw_data() if isinstance(d, ScaledDataset) else d + for k, d in self._datasets.items() + } + if self.scaler is None: + scaler = DatasetScaler(raw, device=self.device) + metrics = scaler.fit(nsteps=nsteps, max_iter=max_iter) + self._datasets = {k: ScaledDataset(d, scaler, k) for k, d in raw.items()} + self.scaler = scaler + else: + self.scaler.datasets = raw + metrics = self.scaler.fit(nsteps=nsteps, max_iter=max_iter) + self.scaling_metrics = metrics + return self - L-BFGS with strong-Wolfe line search, 10 outer steps of ``max_iter=100``. - Members' ``log_scale``/``U_aniso`` are mutated, and ``requires_grad`` is - turned on and back off around the fit. + def _get_state(self) -> dict: + raw = { + k: (d.raw_data() if isinstance(d, ScaledDataset) else d)._get_state() + for k, d in self._datasets.items() + } + scaler_state = None if self.scaler is None else self.scaler.get_state() + if scaler_state is not None: + scaler_state.pop("datasets") + return { + "datasets": raw, + "reference": self._reference_dataset, + "scaler": scaler_state, + "scaling_metrics": self.scaling_metrics, + } + + @classmethod + def _from_state(cls, state: dict, device=None) -> "DatasetCollection": + from torchref.scaling.dataset_scaler import DatasetScaler + + result = cls(device=device) if device is not None else cls() + for key, raw in state["datasets"].items(): + result.add_dataset( + key, + ReflectionData._from_state(dict(raw), device), + set_as_reference=key == state["reference"], + ) + if state["scaler"] is not None: + result.scaler = DatasetScaler.from_state( + {**state["scaler"], "datasets": state["datasets"]}, device + ) + result._datasets = { + k: ScaledDataset(d, result.scaler, k) + for k, d in result._datasets.items() + } + result.scaling_metrics = state.get("scaling_metrics", {}) + return result + + def _keys_or_all(self, keys: Optional[List[str]]) -> List[str]: + if keys is None: + return list(self._dataset_order) + missing = [k for k in keys if k not in self._datasets] + if missing: + raise KeyError(f"Unknown dataset keys: {missing}") + return list(keys) + + def stack_F_obs(self, keys: Optional[List[str]] = None) -> torch.Tensor: + """Scaled observed amplitudes, shape ``(n_datasets, n_reflections)``.""" + return torch.stack( + [self._datasets[k].F for k in self._keys_or_all(keys)], + dim=0, + ) + + def stack_F_sigma(self, keys: Optional[List[str]] = None) -> torch.Tensor: + """Scaled amplitude sigmas, shape ``(n_datasets, n_reflections)``.""" + return torch.stack( + [self._datasets[k].F_sigma for k in self._keys_or_all(keys)], + dim=0, + ) + + def stack_I_obs(self, keys: Optional[List[str]] = None) -> torch.Tensor: + """Scaled observed intensities, shape ``(n_datasets, n_reflections)``. Raises ------ ValueError - If no reference dataset is set, or there is nothing else to scale. + If any selected dataset carries no intensities. + """ + return torch.stack( + [self._require_intensities(k).I for k in self._keys_or_all(keys)], + dim=0, + ) + + def stack_I_sigma(self, keys: Optional[List[str]] = None) -> torch.Tensor: + """Scaled intensity sigmas, shape ``(n_datasets, n_reflections)``. + + Raises + ------ + ValueError + If any selected dataset carries no intensities. + """ + return torch.stack( + [self._require_intensities(k).I_sigma for k in self._keys_or_all(keys)], + dim=0, + ) + + def _require_intensities(self, key: str) -> ReflectionData: + """Return a dataset with intensities, naming it if the column is missing.""" + data = self._datasets[key] + if data.I_raw is None: + raise ValueError( + f"Dataset {key!r} carries no intensities; its reflection file had no " + f"I/SIGI columns. An intensity-space target needs them on every member." + ) + return data + + def stack_masks( + self, keys: Optional[List[str]] = None, use_set: str = "work" + ) -> torch.Tensor: + """Per-dataset boolean subset masks, shape ``(n_datasets, n_reflections)``. + + Uses the 3-way ``work``/``free``/``validation`` accessors, so validation + reflections are excluded from both work and free -- matching what the + collection targets fit. The 2-way ``rfree_flags`` cannot express that. + + Parameters + ---------- + keys : list of str, optional + Datasets to stack; all of them in insertion order by default. + use_set : {"work", "free", "val"}, optional + Which subset to select. Default ``"work"``. + """ + if use_set not in ("work", "free", "val"): + raise ValueError( + f"use_set must be 'work', 'free' or 'val'; got {use_set!r}" + ) + attr = {"work": "work", "free": "free", "val": "validation"}[use_set] + return torch.stack( + [getattr(self._datasets[k], attr).mask for k in self._keys_or_all(keys)], + dim=0, + ) + + def get_centric_flags(self) -> Optional[torch.Tensor]: + """Centric flags on the common HKL, from the reference dataset. + + A pure function of ``(hkl, spacegroup)``, so it is shared by every member and + needs no dataset axis. """ if self._reference_dataset is None: - raise ValueError("No reference dataset set for scaling") - - - ref_ds = self._datasets[self._reference_dataset] - to_scale = [ds for name, ds in self if name != self._reference_dataset] - - if not to_scale: - raise ValueError("No datasets to scale against reference") - - parameters = [p for data in to_scale for p in data.parameters()] - [p.requires_grad_(True) for p in parameters] - optimizer = torch.optim.LBFGS(parameters, max_iter=100, line_search_fn='strong_wolfe') - - # Get masks once (they don't change during optimization) - ref_mask = ref_ds.masks() - ds_masks = [ds.masks() for ds in to_scale] - - def closure(): - optimizer.zero_grad() - loss = 0.0 - # get_corrected_data, not __call__: MaskedTensor has no autograd. - ref_F_scaled, _ = ref_ds.get_corrected_data() - - for ds, ds_mask in zip(to_scale, ds_masks): - F_scaled, _ = ds.get_corrected_data() - combined_mask = ds_mask & ref_mask - F_data = F_scaled[combined_mask] - ref_F_data = ref_F_scaled[combined_mask] - loss = loss + torch.sum((F_data - ref_F_data) ** 2) - loss.backward() - return loss - - for i in range(10): - optimizer.step(closure) - [p.requires_grad_(False) for p in parameters] + return None + return self._datasets[self._reference_dataset].centric + + def component_structure_factors( + self, model_collection, recalc: bool = False + ) -> torch.Tensor: + """Per-base-model ``F_calc`` on the common HKL, in the canonical convention. + + The batched counterpart of :meth:`ReflectionData.structure_factors`: models are + evaluated at the **signed** indices so Bijvoet mates get distinct ``|F_calc|``, + and the result is returned on the canonical ASU index that :attr:`hkl` holds. + Use this rather than calling + :meth:`~torchref.model.model_collection.ModelCollection.compute_component_fcalcs` + on :attr:`hkl` directly, which would skip both halves of that convention. + + The convention is taken from the reference dataset. Members are all expanded onto + one HKL grid, but a member with different completeness can still carry different + ``friedel_flags`` (absent rows are filled ``False``); where they differ, the + returned **phases** follow the reference. Amplitudes are unaffected, so a target + working in moduli or intensities is insensitive to this. + + Parameters + ---------- + model_collection : ModelCollection + Supplies the shared base models. + recalc : bool, optional + Force recomputation rather than reusing each model's cached SF. + + Returns + ------- + torch.Tensor + Complex SFs of shape ``(n_base_models, n_reflections)``, row-aligned with + :attr:`hkl` on the reflection axis. + Raises + ------ + ValueError + If the collection has no reference dataset. + """ + if self._reference_dataset is None: + raise ValueError( + "No reference dataset set; add a dataset before computing " + "component structure factors." + ) + ref = self._datasets[self._reference_dataset] + stacked = model_collection.compute_component_fcalcs( + ref._hkl_for_sf(), recalc=recalc + ) + return ref.conjugate_friedel(stacked) def keys(self) -> List[str]: """Return list of dataset names.""" diff --git a/torchref/io/datasets/fcalc_data.py b/torchref/io/datasets/fcalc_data.py index 616c310d..d4a4a2aa 100644 --- a/torchref/io/datasets/fcalc_data.py +++ b/torchref/io/datasets/fcalc_data.py @@ -7,16 +7,19 @@ """ from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union import pandas as pd import torch -from torchref.config import get_default_device, get_float_dtype, normalize_device +from torchref.config import get_float_dtype, normalize_device from torchref.symmetry import Cell, SpaceGroup, SpaceGroupLike from .base import CrystalDataset +if TYPE_CHECKING: + from .reflection_data import ReflectionData + @dataclass class FcalcDataset(CrystalDataset): @@ -52,6 +55,7 @@ class FcalcDataset(CrystalDataset): fcalc: Optional[torch.Tensor] = None # Complex (N,) fcalc_amp: Optional[torch.Tensor] = None # |Fcalc| (N,) fcalc_phase: Optional[torch.Tensor] = None # Phase in radians (N,) + fobs_sigma: Optional[torch.Tensor] = None # Amp-space sigma (N,), set by add_noise @staticmethod def from_cell_and_resolution( @@ -129,7 +133,7 @@ def from_cell_and_resolution( # make_miller_array returns unique HKL for the asymmetric unit only. hkl_list = gemmi.make_miller_array(gemmi_cell, gemmi_sg, d_min) - hkl = torch.tensor(hkl_list, dtype=torch.int32, device=device) + hkl = torch.tensor(hkl_list, dtype=torch.int32, device=device) # dtype-ok: hkl Miller indices; fixed int32 crystallographic representation, not model-precision data resolution = get_d_spacing(hkl.float(), cell_tensor) @@ -174,6 +178,160 @@ def set_fcalc(self, fcalc: torch.Tensor) -> None: self.fcalc_amp = torch.abs(fcalc).to(device=self.device) self.fcalc_phase = torch.angle(fcalc).to(device=self.device) + def add_noise( + self, + reference: Optional["ReflectionData"] = None, + sigma_lin: float = 0.0, + sigma_mul: float = 0.05, + sigma_abs: float = 0.0, + seed: Optional[int] = None, + verbose: bool = True, + ) -> "FcalcDataset": + """ + Return a copy with Gaussian **intensity** noise, as the mean of two half-datasets. + + Two sigma sources, chosen by ``reference``: + + **Reference-driven (preferred).** Given a ``ReflectionData`` on the same HKL list, + its per-reflection ``I_sigma`` is grafted directly. Right when the Fcalc is already + on the reference's absolute scale. + + **Parametric.** Otherwise a three-term variance model:: + + sigma_I^2 = sigma_lin^2 * I + sigma_mul^2 * I^2 + (sigma_abs * _outer)^2 + + Either way two independent draws are made and averaged:: + + I_h{1,2} = I + N(0, sigma_I) + I_mean = (I_h1 + I_h2) / 2 + sigma_I_mean = sigma_I / sqrt(2) + + R-split and Pearson CC between the halves are reported, which makes the two draws + a ready-made split-half pair rather than only a noise model. + + **Negative intensities are kept.** ``I_mean`` and ``sigma_I_mean`` are stored + unclamped on the returned dataset, and only the *amplitude* is clamped, because an + amplitude cannot be negative. Clamping the intensity instead would put a positive + bias on exactly the weak reflections where noise dominates -- the same shape as a + genuine positive perturbation, and enough to swamp effects of order 1e-3. + + The amplitude sigma is propagated against the **true** amplitude, not the noisy + one, so it is not warped per draw. + + Parameters + ---------- + reference : ReflectionData, optional + Sigma donor. Requires ``torch.equal(self.hkl, reference.hkl)``. + sigma_lin, sigma_mul, sigma_abs : float, optional + Parametric coefficients, ignored when ``reference`` is given. + seed : int, optional + Seed for reproducibility. ``None`` uses the global RNG. + verbose : bool, optional + Print the R-split and CC between halves. Default True. + + Returns + ------- + FcalcDataset + New dataset carrying the noisy complex Fcalc, the unclamped ``I`` / + ``I_sigma``, and ``fobs_sigma``. + """ + if self.fcalc is None or self.fcalc_amp is None or self.fcalc_phase is None: + raise ValueError("No Fcalc values set. Call set_fcalc() first.") + + shape = self.fcalc_amp.shape + dtype = self.fcalc_amp.dtype + dev = self.device + if seed is not None: + g = torch.Generator(device=dev).manual_seed(int(seed)) + randn1 = torch.randn(shape, dtype=dtype, device=dev, generator=g) + randn2 = torch.randn(shape, dtype=dtype, device=dev, generator=g) + else: + randn1 = torch.randn(shape, dtype=dtype, device=dev) + randn2 = torch.randn(shape, dtype=dtype, device=dev) + + intensity = self.fcalc_amp**2 + + if reference is not None: + if reference.I_sigma is None: + raise ValueError( + "reference.I_sigma is None. Load the reference with intensity + " + "sigma columns (e.g. load_crystfel_hkl) before passing it here." + ) + ref_hkl = reference.hkl.to(device=self.hkl.device) + if ref_hkl.shape != self.hkl.shape or not torch.equal(ref_hkl, self.hkl): + raise ValueError( + "reference.hkl does not match self.hkl -- build the FcalcDataset " + "from the reference's HKL list to guarantee 1:1 sigma grafting." + ) + sigma_I = reference.I_sigma.to(device=dev, dtype=dtype) + if verbose: + print( + f"add_noise: grafting sigmas from reference ({len(sigma_I)} " + f"reflections, ={sigma_I.mean().item():.3g})" + ) + else: + if sigma_abs > 0: + if self.resolution is None: + raise ValueError( + "sigma_abs > 0 requires self.resolution (used to pick the " + "outer resolution shell)." + ) + n = len(self.resolution) + k = max(1, int(0.1 * n)) + outer_idx = torch.argsort(self.resolution)[:k] + sigma_abs_I = sigma_abs * intensity[outer_idx].mean() + else: + sigma_abs_I = torch.zeros((), dtype=dtype, device=dev) + + safe = intensity.clamp(min=0.0) + sigma_I = torch.sqrt( + (sigma_lin**2) * safe + + (sigma_mul**2) * safe * safe + + sigma_abs_I**2 + ) + + I_h1 = intensity + randn1 * sigma_I + I_h2 = intensity + randn2 * sigma_I + + diff_sum = (I_h1 - I_h2).abs().sum() + pair_sum = (I_h1 + I_h2).sum() + r_split = ((1.0 / (2.0**0.5)) * diff_sum / (0.5 * pair_sum)).item() + + x = I_h1 - I_h1.mean() + y = I_h2 - I_h2.mean() + cc = ( + (x * y).sum() + / torch.sqrt((x * x).sum() * (y * y).sum()).clamp(min=1e-30) + ).item() + if verbose: + print(f"add_noise: R-split = {r_split:.4f}, CC(half1, half2) = {cc:.4f}") + + I_mean = 0.5 * (I_h1 + I_h2) + sigma_I_mean = sigma_I / (2.0**0.5) + + # Only the amplitude is clamped; see the note in the docstring. + amp_noisy = torch.sqrt(I_mean.clamp(min=0.0)) + sigma_F = sigma_I_mean / (2.0 * self.fcalc_amp.clamp(min=1e-8)) + + fcalc_noisy = ( + amp_noisy * torch.exp(1j * self.fcalc_phase) + ).to(self.fcalc.dtype) + + new = FcalcDataset( + hkl=self.hkl.clone(), + resolution=( + self.resolution.clone() if self.resolution is not None else None + ), + cell=self.cell, + spacegroup=self.spacegroup, + device=self.device, + ) + new.set_fcalc(fcalc_noisy) + new.fobs_sigma = sigma_F.to(self.device) + new.I = I_mean.to(self.device) + new.I_sigma = sigma_I_mean.to(self.device) + return new + def write_mtz(self, filepath: str) -> None: """ Write Fcalc to MTZ as ``F-model`` / ``PH-model`` (phase in degrees). diff --git a/torchref/io/datasets/reflection_data.py b/torchref/io/datasets/reflection_data.py index fd62273b..58722edc 100644 --- a/torchref/io/datasets/reflection_data.py +++ b/torchref/io/datasets/reflection_data.py @@ -9,16 +9,15 @@ import warnings from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union import numpy as np import pandas as pd import torch -from torch.nn import Parameter from torchref.base import math_torch from torchref.base.french_wilson import FrenchWilson -from torchref.config import dtypes, get_default_device, normalize_device +from torchref.config import dtypes, normalize_device from torchref.io import cif, mtz from torchref.io.datasets.base import CrystalDataset from torchref.symmetry import Cell, SpaceGroup @@ -27,16 +26,6 @@ if TYPE_CHECKING: from torchref.model.model_ft import ModelFT -# Suppress PyTorch MaskedTensor prototype warnings globally -# MaskedTensor is stable enough for our use case (aggregations, element-wise ops) -warnings.filterwarnings( - "ignore", message=".*MaskedTensors is in prototype stage.*", category=UserWarning -) - -if TYPE_CHECKING: - from torch.masked import MaskedTensor - - class _ReflectionSubset: """ Lightweight view of one reflection subset (``work`` / ``free`` / @@ -95,25 +84,23 @@ def __len__(self) -> int: def n(self) -> int: return int(self.indices.numel()) - # -- amplitudes (scaled, matching the legacy ``data(scale=True)``) ----- + # Subset reads dispatch through the parent observation attributes. @property def F(self) -> torch.Tensor: - F_corr, _ = self._parent._corrected_or_raw() - return F_corr.index_select(0, self.indices) + return self._parent.F.index_select(0, self.indices) @property def sigF(self) -> torch.Tensor: - _, sig_corr = self._parent._corrected_or_raw() - return sig_corr.index_select(0, self.indices) + return self._parent.F_sigma.index_select(0, self.indices) # -- raw (uncorrected) amplitudes ------------------------------------- @property def F_raw(self) -> torch.Tensor: - return self._parent.F.index_select(0, self.indices) + return self._parent.F_raw.index_select(0, self.indices) @property def sigF_raw(self) -> torch.Tensor: - return self._parent.F_sigma.index_select(0, self.indices) + return self._parent.F_sigma_raw.index_select(0, self.indices) # -- common aliases --------------------------------------------------- @property @@ -124,9 +111,33 @@ def hkl(self) -> torch.Tensor: def rfree(self) -> torch.Tensor: return self._parent.rfree_flags.index_select(0, self.indices) + # -- intensities, corrected to match F/sigF above ----------------------- + @property + def I(self) -> torch.Tensor: # noqa: E743 - crystallographic name + """Scaled intensities, or None when this dataset carries no intensities. + + Corrected, like :attr:`F` -- both the anisotropy factor and the overall scale + enter squared. Use :attr:`I_raw` for the unscaled values. + """ + I_corr = self._parent.I + return I_corr.index_select(0, self.indices) if I_corr is not None else None + @property def sigI(self): - si = self._parent.I_sigma + """Scaled intensity sigmas, or None. See :attr:`I`.""" + sig_corr = self._parent.I_sigma + return sig_corr.index_select(0, self.indices) if sig_corr is not None else None + + @property + def I_raw(self): + """Unscaled intensities, or None.""" + i = self._parent.I_raw + return i.index_select(0, self.indices) if i is not None else None + + @property + def sigI_raw(self): + """Unscaled intensity sigmas, or None.""" + si = self._parent.I_sigma_raw return si.index_select(0, self.indices) if si is not None else None @property @@ -214,11 +225,7 @@ def __post_init__(self): """ # Call parent __post_init__ to initialize masks super().__post_init__() - self.setup_scale() - self.setup_anisotropy() - # Cached integer index maps for the work/free/validation subsets and - # a cache of the scaled (F, F_sigma). Both are invalidated by - # fingerprints (see _subset_indices / _corrected_or_raw). + # Subset membership is cached independently of observation values. self._subset_cache = { "work": None, "free": None, @@ -226,8 +233,6 @@ def __post_init__(self): "all": None, } self._subset_fp = None - self._corrected_cache = None - self._corrected_fp = None # ===================== work / free / validation ===================== @@ -301,7 +306,7 @@ def _subset_indices(self, kind: str) -> torch.Tensor: n = 0 if self.hkl is None else len(self.hkl) device = self.device if n == 0: - empty = torch.empty(0, dtype=torch.long, device=device) + empty = torch.empty(0, dtype=torch.long, device=device) # dtype-ok: empty index tensor; PyTorch requires int64 for indexing self._subset_cache = { "work": empty, "free": empty, @@ -334,26 +339,25 @@ def _subset_indices(self, kind: str) -> torch.Tensor: self._subset_fp = fp return self._subset_cache[kind] - def _corrected_or_raw(self) -> Tuple[torch.Tensor, torch.Tensor]: - """Return the scaled (F, F_sigma) (matching ``data(scale=True)``), - cached against the (log_scale, U_aniso) fingerprint. Falls back to the - raw (F, F_sigma) if scaling is not set up. - """ + @property + def F_raw(self) -> Optional[torch.Tensor]: + """Measured amplitudes, shape (N,), in the input amplitude units.""" + return self.F - def _tv(t): - return (t.data_ptr(), t._version) if isinstance(t, torch.Tensor) else None + @property + def F_sigma_raw(self) -> Optional[torch.Tensor]: + """Measured amplitude uncertainties, shape (N,), in amplitude units.""" + return self.F_sigma - fp = ( - _tv(getattr(self, "log_scale", None)), - _tv(getattr(self, "U_aniso", None)), - ) - if self._corrected_fp != fp or self._corrected_cache is None: - try: - self._corrected_cache = self.get_corrected_data() - except Exception: - self._corrected_cache = (self.F, self.F_sigma) - self._corrected_fp = fp - return self._corrected_cache + @property + def I_raw(self) -> Optional[torch.Tensor]: + """Measured intensities, shape (N,), in the input intensity units.""" + return self.I + + @property + def I_sigma_raw(self) -> Optional[torch.Tensor]: + """Measured intensity uncertainties, shape (N,), in intensity units.""" + return self.I_sigma # ===================== per-reflection field reindexing ===================== # @@ -373,10 +377,6 @@ def _tv(t): "I": 0.0, "phase": 0.0, "fom": 0.0, - "E": 0.0, - "E_squared": 0.0, - "F_squared_corrected": 0.0, - "radial_shell_indices": 0, "F_sigma": 1.0, "I_sigma": 1.0, "rfree_flags": 1, # missing reflections default to the work set @@ -390,10 +390,6 @@ def _tv(t): # lazily rebuilt by ``get_bins`` / the ``centric`` property). _REINDEX_DERIVED = ("resolution", "bin_indices", "_centric_flags") - # Tensor dataclass fields that are NOT per-reflection (exempt from the - # length invariant): overall anisotropy parameters have shape (6,). - _NON_PER_REFLECTION_TENSORS = frozenset({"U_aniso"}) - def _reindex_per_reflection( self, index_map: torch.Tensor, @@ -432,7 +428,7 @@ def _reindex_per_reflection( n_src = len(self.hkl) if self.hkl is not None else 0 new_hkl = new_hkl.to(dtype=dtypes.int, device=self.device) n_out = len(new_hkl) - index_map = index_map.to(device=self.device, dtype=torch.long) + index_map = index_map.to(device=self.device, dtype=torch.long) # dtype-ok: index map used for indexing/gather; PyTorch requires int64 present = index_map >= 0 src_idx = index_map[present] @@ -441,15 +437,6 @@ def _reindex_per_reflection( name = f.name if name == "hkl" or name in derived: continue - # Declared non-per-reflection fields are exempt by *name*, not by - # shape. The shape test below is a heuristic and collides whenever - # n_src equals the field's own length -- ``U_aniso`` is (6,), so a - # 6-reflection dataset would have it gathered as if it were - # per-reflection. ``_assert_per_reflection_consistent`` and - # ``reduce_to_spacegroup`` already exempt by name; this keeps all - # three routines consistent. - if name in self._NON_PER_REFLECTION_TENSORS: - continue val = getattr(self, name) if not isinstance(val, torch.Tensor): continue @@ -494,8 +481,6 @@ def _assert_per_reflection_consistent(self) -> None: n = len(self.hkl) if self.hkl is not None else 0 bad = [] for f in dc_fields(self): - if f.name in self._NON_PER_REFLECTION_TENSORS: - continue val = getattr(self, f.name) if isinstance(val, torch.Tensor) and val.ndim >= 1 and val.shape[0] != n: bad.append((f.name, tuple(val.shape))) @@ -508,13 +493,13 @@ def _canonicalize_in_place(self) -> None: """Remap HKL to canonical CCP4 ASU form and reorder all data in-place.""" from dataclasses import fields as dc_fields - from torchref.symmetry.reciprocal_symmetry import canonicalize_hkl - if self.hkl is None or self.spacegroup is None: return - canonical_hkl, phase_shifts, friedel_flags, sort_indices = canonicalize_hkl( - self.hkl, self.spacegroup, include_friedel=True, device=self.device + canonical_hkl, phase_shifts, friedel_flags, sort_indices = ( + self.spacegroup.canonicalize_hkl( + self.hkl, include_friedel=True, device=self.device + ) ) n_refl = len(self.hkl) @@ -709,8 +694,10 @@ def _group_any( Uses ``index_add_`` on float rather than ``scatter_reduce_(amax)``: the latter raises "not supported for torch.int64" on the MPS backend. """ + # dtype-ok: float32 count accumulator for the int64-scatter MPS workaround + # above; the result is reduced to bool (> 0), so precision is irrelevant. counts = torch.zeros(n_groups, dtype=torch.float32, device=mask.device) - counts.index_add_(0, group_id, mask.to(torch.float32)) + counts.index_add_(0, group_id, mask.to(torch.float32)) # dtype-ok: float32 counter for the MPS workaround above; reduced to bool return counts > 0 @staticmethod @@ -860,6 +847,13 @@ def load(self, reader, french_wilson: bool = True): rfree = rfree.clip(min=0, max=1).to(torch.bool) self.rfree_flags = rfree self.masks["flagged_initial"] = ~flagged + # Record the provenance for every file-sourced set, not only the + # ones that also carry a validation column: a header reporting + # R-free has to be able to say which test set produced it. Named + # after the reader rather than hardcoded "MTZ", since `load` also + # takes ReflectionCIFReader and any other compatible reader. + reader_name = type(reader).__name__ + self.rfree_source = f"{reader_name} FreeR" # A third (validation) column goes into the separate boolean # ``validation_flags``; ``rfree_flags`` stays binary work/free. if "Validation-flags" in data_dict: @@ -868,7 +862,7 @@ def load(self, reader, french_wilson: bool = True): device=self.device, requires_grad=False, ).to(torch.bool) - self.rfree_source = "MTZ FreeR+Validation" + self.rfree_source = f"{reader_name} FreeR+Validation" self._post_load_cleanup() @@ -1039,6 +1033,37 @@ def load_mtz( ).read(str(path)) return self.load(reader, french_wilson=french_wilson) + def load_crystfel_hkl( + self, path: str, cell, spacegroup, + ) -> "ReflectionData": + """ + Load a CrystFEL ``partialator`` ``.hkl`` reflection list. + + Unlike MTZ, the CrystFEL format carries no cell or space-group metadata, so both + must be supplied by the caller -- they usually live in a ``.cell`` file alongside. + + The format is intensity-native, so amplitudes are derived by French-Wilson on + load exactly as they are for an MTZ carrying I/SIGI columns. + + Parameters + ---------- + path : str + Path to the ``.hkl`` file. + cell : list | tuple | np.ndarray | Cell | torch.Tensor + Unit cell (a, b, c, alpha, beta, gamma). + spacegroup : str | gemmi.SpaceGroup | SpaceGroup + Space group identifier. + + Returns + ------- + ReflectionData + Self, for method chaining. + """ + from torchref.io import hkl as _hkl + + reader = _hkl.HKLReader(verbose=self.verbose).read(path, cell, spacegroup) + return self.load(reader) + def load_cif( self, path: Union[str, Path], @@ -1207,7 +1232,13 @@ def _generate_rfree_flags( flags[group_free[group_id]] = 0 self.rfree_flags = flags - self.rfree_source = "Generated (resolution-binned, ASU-grouped)" + # The seed belongs in the provenance string: without it "generated" + # names a draw nobody can reproduce. + self.rfree_source = ( + "Generated (resolution-binned, ASU-grouped" + + (f", seed {seed}" if seed is not None else "") + + ")" + ) n_free = (flags == 0).sum().item() n_work = (flags != 0).sum().item() @@ -1326,13 +1357,13 @@ def mean_res_per_bin(self) -> torch.Tensor: mean_resolutions = torch.scatter_add( mean_resolutions, 0, - self.bin_indices[mask].to(torch.int64), + self.bin_indices[mask].to(torch.int64), # dtype-ok: bin indices for scatter_add/index; PyTorch requires int64 self.resolution[mask], ) count_per_bin = torch.scatter_add( count_per_bin, 0, - self.bin_indices[mask].to(torch.int64), + self.bin_indices[mask].to(torch.int64), # dtype-ok: bin indices for scatter_add/index; PyTorch requires int64 torch.ones_like(self.resolution[mask], dtype=dtypes.int), ) mean_resolutions = mean_resolutions / count_per_bin.clamp(min=1).float() @@ -1361,12 +1392,12 @@ def mean_F_per_bin(self) -> torch.Tensor: count_per_bin = torch.zeros(self._n_bins, dtype=dtypes.int, device=self.device) mask = self.masks() mean_F = torch.scatter_add( - mean_F, 0, self.bin_indices[mask].to(torch.int64), self.F[mask] + mean_F, 0, self.bin_indices[mask].to(torch.int64), self.F[mask] # dtype-ok: bin indices for scatter_add index arg; PyTorch requires int64 ) count_per_bin = torch.scatter_add( count_per_bin, 0, - self.bin_indices[mask].to(torch.int64), + self.bin_indices[mask].to(torch.int64), # dtype-ok: bin indices for scatter_add index arg; PyTorch requires int64 torch.ones_like(self.F[mask], dtype=dtypes.int), ) mean_F = mean_F / count_per_bin.clamp(min=1).float() @@ -1395,12 +1426,12 @@ def mean_sigma_per_bin(self) -> Optional[torch.Tensor]: count_per_bin = torch.zeros(self._n_bins, dtype=dtypes.int, device=self.device) mask = self.masks() mean_sigma = torch.scatter_add( - mean_sigma, 0, self.bin_indices[mask].to(torch.int64), self.F_sigma[mask] + mean_sigma, 0, self.bin_indices[mask].to(torch.int64), self.F_sigma[mask] # dtype-ok: bin indices for scatter_add index arg; PyTorch requires int64 ) count_per_bin = torch.scatter_add( count_per_bin, 0, - self.bin_indices[mask].to(torch.int64), + self.bin_indices[mask].to(torch.int64), # dtype-ok: bin indices for scatter_add index arg; PyTorch requires int64 torch.ones_like(self.F_sigma[mask], dtype=dtypes.int), ) mean_sigma = mean_sigma / count_per_bin.clamp(min=1).float() @@ -1813,15 +1844,6 @@ def filter_by_resolution( return self - def get_mask(self): - """ - Placeholder for returning a combined mask from all active filters. - - Not implemented; the body is empty and this returns ``None``. Use - :meth:`masks` (the combined-validity callable) to obtain the boolean - mask combining all active filter conditions. - """ - def cut_res( self, highres: Optional[float] = None, lowres: Optional[float] = None ) -> "ReflectionData": @@ -1845,104 +1867,6 @@ def cut_res( """ return self.filter_by_resolution(d_min=highres, d_max=lowres) - def get_rfree_masks(self) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """ - Get boolean masks for work and test (free) sets. - - Returns - ------- - work_mask : torch.Tensor or None - Boolean tensor for work set (flag != 0). - test_mask : torch.Tensor or None - Boolean tensor for test/free set (flag == 0). - Both are None if no R-free flags are available. - - .. deprecated:: - Use ``data.work.mask`` / ``data.free.mask`` (which also apply the - validity masks), or ``data.work.indices`` / ``data.free.indices``. - """ - warnings.warn( - "ReflectionData.get_rfree_masks() is deprecated; use data.work.mask " - "/ data.free.mask (the work/free/validation accessor).", - DeprecationWarning, - stacklevel=2, - ) - if self.rfree_flags is None: - return None, None - - work_mask = self.rfree_flags != 0 - test_mask = self.rfree_flags == 0 - - return work_mask, test_mask - - def get_work_set(self) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """ - Get structure factors for the work set (R-free flag != 0). - - Returns - ------- - F_work : torch.Tensor - Structure factors for work set. - sigma_work : torch.Tensor or None - Uncertainties for work set, or None if not available. - - With no R-free flags this silently returns the *full* dataset (warning - printed), so a caller cannot tell work from all. - - .. deprecated:: - Use ``data.work.F`` / ``data.work.sigF``, which also apply the - validity masks and cache the subset indices. - """ - warnings.warn( - "ReflectionData.get_work_set() is deprecated; use data.work.F / " - "data.work.sigF (the work/free/validation accessor).", - DeprecationWarning, - stacklevel=2, - ) - if self.rfree_flags is None: - print("WARNING: No R-free flags available, returning full dataset") - return self.F, self.F_sigma - - work_mask = self.rfree_flags != 0 - F_work = self.F[work_mask] if self.F is not None else None - sigma_work = self.F_sigma[work_mask] if self.F_sigma is not None else None - - return F_work, sigma_work - - def get_test_set(self) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """ - Get structure factors for the test set (R-free flag == 0). - - Returns - ------- - F_test : torch.Tensor - Structure factors for test/free set. - sigma_test : torch.Tensor or None - Uncertainties for test set, or None if not available. - - Raises - ------ - ValueError - If no R-free flags are available. - - .. deprecated:: - Use ``data.free.F`` / ``data.free.sigF`` instead. - """ - warnings.warn( - "ReflectionData.get_test_set() is deprecated; use data.free.F / " - "data.free.sigF (the work/free/validation accessor).", - DeprecationWarning, - stacklevel=2, - ) - if self.rfree_flags is None: - raise ValueError("No R-free flags available in dataset") - - test_mask = self.rfree_flags == 0 - F_test = self.F[test_mask] if self.F is not None else None - sigma_test = self.F_sigma[test_mask] if self.F_sigma is not None else None - - return F_test, sigma_test - def get_max_res(self) -> Optional[float]: """Smallest d-spacing among valid reflections, in Ångströms.""" if self.resolution is None: @@ -2001,7 +1925,7 @@ def data_indexed( """ Return reflection data as compact (valid-only) tensors. - Note these are the RAW ``F``/``F_sigma``, not the scaled ones. + ScaledDataset returns its live corrected observations. Returns ------- @@ -2027,80 +1951,6 @@ def data_indexed( return hkl, F, F_sigma, rfree_flags - def __call__( - self, mask: bool = True, scale: bool = True - ) -> Tuple[torch.Tensor, "MaskedTensor", "MaskedTensor", torch.Tensor]: - """ - Return core reflection data with MaskedTensors for F and sigma. - - Everything is full size (N); invalid reflections are marked in the mask - rather than removed. With ``mask=True`` the returned ``F``/``F_sigma`` - are detached clones, so gradients do NOT flow through them (use - :meth:`get_corrected_data` when the graph is needed). - - Parameters - ---------- - mask : bool, optional - If True, wrap F and sigma as MaskedTensors. Default is True. - scale : bool, optional - If True, apply the current scale/anisotropy before returning. - - Returns - ------- - hkl : torch.Tensor - Miller indices of shape (N, 3), unfiltered. - F : MaskedTensor - Amplitudes of shape (N,) with invalid reflections masked. - F_sigma : MaskedTensor or None - Uncertainties of shape (N,) with invalid reflections masked. - rfree_flags : torch.Tensor or None - Flags of shape (N,), unfiltered. 1=work, 0=free. - - Raises - ------ - RuntimeError - If ``mask`` is True and every reflection is masked out. - - .. deprecated:: - Use the work/free/validation accessor (``data.work.F``, - ``data.free.F``, ``.sigF`` / ``.hkl`` / ``.select(...)``), or - ``data.get_corrected_data()`` for the full scaled (F, F_sigma). - """ - warnings.warn( - "Calling ReflectionData (data()) is deprecated; use the " - "data.work / data.free / data.validation accessor, or " - "data.get_corrected_data() for the full scaled arrays.", - DeprecationWarning, - stacklevel=2, - ) - return self._masked_unpack(mask=mask, scale=scale) - - def _masked_unpack( - self, mask: bool = True, scale: bool = True - ) -> Tuple[torch.Tensor, "MaskedTensor", "MaskedTensor", torch.Tensor]: - """Non-deprecated body of the legacy ``__call__``; see it for the contract. - - Internal only -- external callers should use the work/free/validation - accessor. - """ - from torch.masked import MaskedTensor - - hkl, F, F_sigma, rfree_flags = self.hkl, self.F, self.F_sigma, self.rfree_flags - - if scale: - F, F_sigma = self.get_corrected_data() - - if mask: - to_mask = self.masks() - if to_mask.sum() == 0: - raise RuntimeError( - "All reflections are masked! Check your filters/masks." - ) - F = MaskedTensor(F.detach().clone(), to_mask) - if F_sigma is not None: - F_sigma = MaskedTensor(F_sigma.detach().clone(), to_mask) - return hkl, F, F_sigma, rfree_flags - def data_fill_masked( self, mode="mean" ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: @@ -2127,23 +1977,31 @@ def data_fill_masked( R-free flags of shape (N,); filled-in reflections are assigned to the work set (True). """ - hkl, F, F_sigma, rfree = self._masked_unpack() + hkl, F, F_sigma = self.hkl, self.F, self.F_sigma + if F is None or F_sigma is None: + raise ValueError("Amplitude observations and uncertainties are required") + mask = self.masks() + if not bool(mask.any()): + raise ValueError("No valid reflections to fill from") + rfree = ( + self.rfree_flags.clone() + if self.rfree_flags is not None + else torch.ones_like(mask) + ) if mode == "mean": mean_F = self.mean_F_per_bin() mean_F_sigma = self.mean_sigma_per_bin() - F_data = F.get_data().clone() - F_sigma_data = F_sigma.get_data().clone() - mask = F.get_mask() + F_data = F.clone() + F_sigma_data = F_sigma.clone() F_data[~mask] = mean_F[self.bin_indices[~mask]] F_sigma_data[~mask] = mean_F_sigma[self.bin_indices[~mask]] rfree[~mask] = True # set missing to work set return hkl, F_data, F_sigma_data, rfree elif mode == "zero": - mask = F.get_mask() - F_data = F.get_data().clone() - F_sigma_data = F_sigma.get_data().clone() + F_data = F.clone() + F_sigma_data = F_sigma.clone() F_data[~mask] = 0.0 F_sigma_data[~mask] = 0.0 rfree[~mask] = True # set missing to work set @@ -2203,15 +2061,10 @@ def __select__(self, indices: torch.Tensor, op=None) -> "ReflectionData": if val is None: continue if isinstance(val, torch.Tensor): - # Exempt by name first: the shape test is a heuristic that - # collides when n_refl equals the field's own length (see - # ``_reindex_per_reflection``). - if f.name in self._NON_PER_REFLECTION_TENSORS: - setattr(selected, f.name, val.clone()) - elif val.shape and val.shape[0] == n_refl: + if val.shape and val.shape[0] == n_refl: setattr(selected, f.name, val[indices]) else: - # Non-matching tensor (e.g. U_aniso shape (6,)): copy as-is + # Preserve scalar tensor metadata. setattr(selected, f.name, val.clone()) elif isinstance(val, Cell): setattr(selected, f.name, val.clone()) @@ -2302,7 +2155,9 @@ def check_all_data_types(self): else: print(f"{key}: None") - def validate_hkl(self, hkl_ref: torch.Tensor) -> "ReflectionData": + def validate_hkl( + self, hkl_ref: torch.Tensor, *, identity_hkl: Optional[torch.Tensor] = None + ) -> "ReflectionData": """ Expand this dataset **in place** onto a reference HKL set. @@ -2317,6 +2172,10 @@ def validate_hkl(self, hkl_ref: torch.Tensor) -> "ReflectionData": hkl_ref : torch.Tensor Reference Miller indices of shape (N, 3), dtype int32; defines the canonical ordering for all aligned datasets. + identity_hkl : torch.Tensor, optional + Signed anomalous indices of shape (N, 3), distinguishing Bijvoet + observations that share a canonical HKL. When supplied, match these + against the dataset's signed indices and preserve their identities. Returns ------- @@ -2342,11 +2201,15 @@ def validate_hkl(self, hkl_ref: torch.Tensor) -> "ReflectionData": # Build lookup from data HKL to index # Use a dictionary with tuple keys for fast lookup - hkl_data_np = self.hkl.cpu().numpy() + source_hkl = self.hkl if identity_hkl is None else self._hkl_for_sf() + hkl_data_np = source_hkl.cpu().numpy() data_hkl_to_idx = {tuple(hkl): idx for idx, hkl in enumerate(hkl_data_np)} # For each reference HKL, find the corresponding data index (or -1 if missing) - hkl_ref_np = hkl_ref.cpu().numpy() + lookup_hkl = hkl_ref if identity_hkl is None else identity_hkl + if lookup_hkl.shape != hkl_ref.shape: + raise ValueError("identity_hkl must match the reference HKL shape") + hkl_ref_np = lookup_hkl.cpu().numpy() ref_to_data_idx = np.array( [data_hkl_to_idx.get(tuple(hkl), -1) for hkl in hkl_ref_np], dtype=np.int64 ) @@ -2357,6 +2220,9 @@ def validate_hkl(self, hkl_ref: torch.Tensor) -> "ReflectionData": # Reindex EVERY per-reflection field via the shared primitive. Masks are # handled separately below because they are not dataclass fields. presence_mask = self._reindex_per_reflection(valid_indices, hkl_ref) + if identity_hkl is not None: + self.hkl_anomalous = identity_hkl.to(self.hkl).clone() + self.friedel_flags = (self.hkl_anomalous != self.hkl).any(dim=-1) # Transfer existing masks to new indexing old_masks = dict(self.masks.items()) @@ -2486,11 +2352,11 @@ def flag_wilson_outliers( observations depart from Wilson statistics for reasons that have nothing to do with being outliers. """ - from torchref.base.french_wilson import ( + from torchref.base.french_wilson import intensities_from_amplitudes + from torchref.base.wilson_outliers import wilson_outlier_mask + from torchref.refinement.model_error_estimation.sigma_a import ( epsilon_from_hkl, - intensities_from_amplitudes, ) - from torchref.base.wilson_outliers import wilson_outlier_mask if self.F is None or self.F_sigma is None or self.resolution is None: return @@ -2642,8 +2508,8 @@ def _build_anomalous_dataframe( # The (+) member is the unconjugated row, (-) is the Friedel-flagged row. arange = torch.arange(N) - plus_idx = torch.full((M,), -1, dtype=torch.long) - minus_idx = torch.full((M,), -1, dtype=torch.long) + plus_idx = torch.full((M,), -1, dtype=torch.long) # dtype-ok: Friedel-mate index map (-1 sentinel) for indexing; PyTorch requires int64 + minus_idx = torch.full((M,), -1, dtype=torch.long) # dtype-ok: Friedel-mate index map (-1 sentinel) for indexing; PyTorch requires int64 # A Bijvoet mate only counts as present if it is a real, positive # observation. Stacked anomalous input (rs.stack_anomalous) carries a # row for every *absent* mate with a NaN intensity, which French-Wilson @@ -2982,11 +2848,9 @@ def centric(self): # Cached on the _centric_flags dataclass field, so it survives # serialization. if not hasattr(self, "_centric_flags") or self._centric_flags is None: - from torchref.base.french_wilson import is_centric_from_hkl + sg = self.spacegroup or SpaceGroup("P1", device=self.hkl.device) - sg = self.spacegroup if self.spacegroup else "P1" - - self._centric_flags = is_centric_from_hkl(self.hkl, sg) + self._centric_flags = sg.is_centric(self.hkl) return self._centric_flags @@ -3186,8 +3050,6 @@ def fill(self, d_min: Optional[float] = None) -> "ReflectionData": Missing reflections have F/I/phase/fom = 0.0, F_sigma/I_sigma = 1.0 and ``masks['missing'] = True``. """ - from torchref.symmetry.reciprocal_symmetry import complete_hkl - if self.hkl is None: raise ValueError("ReflectionData has no Miller indices loaded") if self.cell is None: @@ -3200,8 +3062,9 @@ def fill(self, d_min: Optional[float] = None) -> "ReflectionData": d_min = self.resolution.min().item() # Get complete HKL set with index mapping - filled_hkl, indices, missing = complete_hkl( - self.hkl, self.cell.data, self.spacegroup or "P1", d_min, device=self.device + sg = self.spacegroup or SpaceGroup("P1", device=self.device) + filled_hkl, indices, missing = sg.complete_hkl( + self.hkl, self.cell.data, d_min, device=self.device ) # Use remap to create the new dataset @@ -3241,15 +3104,13 @@ def expand_to_p1( shift, ``resolution`` is recomputed and ``bin_indices`` is cleared. ``source``/``last_op`` record the provenance. """ - from torchref.symmetry.reciprocal_symmetry import expand_hkl - if self.hkl is None: raise ValueError("ReflectionData has no Miller indices loaded") # Get expanded HKL set with index mapping and phase shifts - hkl_p1, indices, phase_shifts = expand_hkl( + sg = self.spacegroup or SpaceGroup("P1", device=self.device) + hkl_p1, indices, phase_shifts = sg.expand_hkl( self.hkl, - self.spacegroup or "P1", include_friedel=include_friedel, remove_absences=remove_absences, device=self.device, @@ -3302,16 +3163,15 @@ def reduce_to_spacegroup( ``validation_flags`` set if any equivalent is set. Any other per-reflection field takes its first valid equivalent. """ - from torchref.symmetry.reciprocal_symmetry import reduce_hkl from torchref.symmetry.spacegroup import SpaceGroup if self.hkl is None: raise ValueError("ReflectionData has no Miller indices loaded") # Get reduction mapping - hkl_asu, reduction_indices, phase_shifts = reduce_hkl( - self.hkl, spacegroup, include_friedel=include_friedel, device=self.device - ) + hkl_asu, reduction_indices, phase_shifts = SpaceGroup( + spacegroup, device=self.device + ).reduce_hkl(self.hkl, include_friedel=include_friedel, device=self.device) n_asu = len(hkl_asu) n_equiv = reduction_indices.shape[1] @@ -3466,8 +3326,6 @@ def _aggregate_any(tensor): name = f.name if name in _already_set or name in _recomputed: continue - if name in self._NON_PER_REFLECTION_TENSORS: - continue val = getattr(self, name) if not isinstance(val, torch.Tensor): continue @@ -3523,13 +3381,12 @@ def canonicalize(self, include_friedel: bool = True) -> "ReflectionData": ReflectionData New object with canonicalized, sorted Miller indices. """ - from torchref.symmetry.reciprocal_symmetry import canonicalize_hkl - if self.hkl is None: raise ValueError("ReflectionData has no Miller indices loaded") - canonical_hkl, phase_shifts, friedel_flags, sort_indices = canonicalize_hkl( - self.hkl, self.spacegroup or "P1", include_friedel, device=self.device + sg = self.spacegroup or SpaceGroup("P1", device=self.device) + canonical_hkl, phase_shifts, friedel_flags, sort_indices = sg.canonicalize_hkl( + self.hkl, include_friedel, device=self.device ) # Reorder all fields using __select__ @@ -3589,363 +3446,25 @@ def get_scattering_vectors(self) -> torch.Tensor: return math_torch.get_scattering_vectors(self.hkl, self.cell.data) - def get_radial_shells( - self, - n_shells: int = 20, - d_min: Optional[float] = None, - d_max: Optional[float] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Create uniform radial shells in 1/d space for normalization. - - Uniform in 1/d, unlike :meth:`get_bins` which makes equal-count bins. - Caches the result on ``self.radial_shell_indices``. - - Parameters - ---------- - n_shells : int - Number of radial shells. Default is 20. - d_min, d_max : float, optional - Resolution limits in Angstroms; default to the dataset's own - smallest / largest d. - - Returns - ------- - shell_edges : torch.Tensor - Shell boundaries in Angstroms^-1, shape (n_shells+1,). - shell_centers : torch.Tensor - Shell centers in Angstroms^-1, shape (n_shells,). - shell_indices : torch.Tensor - Shell index for each reflection, shape (N,). Values -1 for out-of-range. - """ - from torchref.base.normalization import ( - assign_to_shells, - compute_radial_shells, - ) - - if self.resolution is None: - self._calculate_resolution() - - # Get resolution limits - if d_min is None: - d_min = self.get_max_res() - if d_max is None: - d_max = self.get_min_res() - - # Compute shells - shell_edges, shell_centers = compute_radial_shells( - d_min, d_max, n_shells, device=self.device - ) - - # Get s-vectors and magnitudes - s_vectors = self.get_scattering_vectors() - s_mag = torch.linalg.norm(s_vectors, dim=1) - - # Assign to shells - shell_indices = assign_to_shells(s_mag, shell_edges) - - # Cache shell indices - self.radial_shell_indices = shell_indices - - return shell_edges, shell_centers, shell_indices - - def fit_anisotropy( - self, - n_shells: int = 20, - d_min: Optional[float] = None, - d_max: Optional[float] = None, - n_iterations: int = 100, - verbose: Optional[bool] = None, - ) -> torch.Tensor: - """ - Fit anisotropy correction parameters to minimize CV within shells. - - Optimizes U so that corrected F² values have minimal coefficient of - variation within each resolution shell. - - Parameters - ---------- - n_shells : int - Number of resolution shells for variance calculation. - d_min, d_max : float, optional - Resolution limits in Angstroms; default to the dataset's own - smallest / largest d. - n_iterations : int - Optimizer steps; see :func:`fit_anisotropy_correction` -- below 20 - nothing is optimized. - verbose : bool, optional - Print progress. If None, uses self.verbose. - - Returns - ------- - U : torch.Tensor - Fitted anisotropy parameters [u11, u22, u33, u12, u13, u23], shape (6,). - Also stored in self.U_aniso. - - Raises - ------ - ValueError - If no amplitude data is available. - """ - from torchref.base import fit_anisotropy_correction - - if self.F is None: - raise ValueError("No amplitude data loaded") - - if verbose is None: - verbose = self.verbose > 0 - - # Get F² values - F_squared = self.F**2 - - # Get s-vectors - s_vectors = self.get_scattering_vectors() - - # Get resolution limits - if d_min is None: - d_min = self.get_max_res() - if d_max is None: - d_max = self.get_min_res() - - # Fit anisotropy - U, final_cv = fit_anisotropy_correction( - F_squared, - s_vectors, - n_shells=n_shells, - d_min=d_min, - d_max=d_max, - n_iterations=n_iterations, - verbose=verbose, - ) - - # Store result - self.U_aniso = U - - return U - - def setup_anisotropy( - self, - U_aniso: Optional[torch.Tensor] = None, - ) -> None: - """ - Setup anisotropy correction parameters. - - Parameters - ---------- - U_aniso : torch.Tensor, optional - Anisotropic parameters [u11, u22, u33, u12, u13, u23], shape (6,). - If None, U_aniso is initialized to a zero (6,) tensor. - - Returns - ------- - ReflectionData - Self, for method chaining. - """ - - if U_aniso is None: - U_aniso = torch.zeros( - 6, device=self.device, dtype=dtypes.float, requires_grad=False - ) - else: - U_aniso = U_aniso.to( - device=self.device, dtype=dtypes.float, requires_grad=False - ) - self.U_aniso = U_aniso - - return self - - def apply_anisotropy_correction( - self, - U_aniso: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """ - Apply anisotropy correction to F² values. - - Parameters - ---------- - U_aniso : torch.Tensor, optional - Anisotropic parameters [u11, u22, u33, u12, u13, u23], shape (6,). - If None, uses self.U_aniso (must have called fit_anisotropy first). - - Returns - ------- - F_corrected: torch.Tensor - Anisotropy-corrected F values, shape (N,). - sigma_F_corrected: torch.Tensor - Uncertainties of corrected F values, shape (N,). - Raises - ------ - ValueError - If no U parameters available and none provided. - """ - from torchref.base import apply_anisotropy_correction - - if U_aniso is None: - U_aniso = self.U_aniso - if U_aniso is None: - raise ValueError( - "No anisotropy parameters available. " - "Call fit_anisotropy() first or provide U_aniso." - ) - - if self.F is None: - raise ValueError("No amplitude data loaded") - - # Get s-vectors - s_vectors = self.get_scattering_vectors() - # Use raw tensors directly to preserve gradient flow - # (MaskedTensor doesn't support autograd operations) - F = self.F - sigma = self.F_sigma - # Apply correction - F_corrected = apply_anisotropy_correction(F, s_vectors, U_aniso) - sigma_F_corrected = ( - apply_anisotropy_correction(sigma, s_vectors, U_aniso) - if sigma is not None - else None - ) - - return F_corrected, sigma_F_corrected - - def compute_e_values( - self, - n_shells: int = 20, - d_min: Optional[float] = None, - d_max: Optional[float] = None, - apply_anisotropy: bool = True, - fit_anisotropy: bool = True, - verbose: Optional[bool] = None, - ) -> torch.Tensor: - """ - Compute E-values with optional anisotropy correction. - - E-values are normalized structure factors where = 1 within each - resolution shell. Anisotropy correction can be applied first to account - for directional variation in diffraction. - - Parameters - ---------- - n_shells : int - Number of resolution shells for normalization. - d_min, d_max : float, optional - Resolution limits in Angstroms; default to the dataset's own - smallest / largest d. - apply_anisotropy : bool - If True, correct for anisotropy before normalizing. - fit_anisotropy : bool - If True, refit U first; if False, reuse the existing - ``self.U_aniso``. Ignored unless ``apply_anisotropy``. - verbose : bool, optional - Print progress. If None, uses self.verbose. - - Returns - ------- - E : torch.Tensor - E-values, shape (N,). Also stores ``self.E``, ``self.E_squared`` and - ``self.radial_shell_indices`` as a side effect. - - Raises - ------ - ValueError - If no amplitude data is available. - """ - from torchref.base import F_squared_to_E_values - - if self.F is None: - raise ValueError("No amplitude data loaded") - - if verbose is None: - verbose = self.verbose > 0 - - # Get resolution limits - if d_min is None: - d_min = self.get_max_res() - if d_max is None: - d_max = self.get_min_res() - - # Get F² values (possibly with anisotropy correction) - if apply_anisotropy: - if fit_anisotropy: - self.fit_anisotropy( - n_shells=n_shells, d_min=d_min, d_max=d_max, verbose=verbose - ) - F_squared = self.apply_anisotropy_correction()[0] ** 2 - else: - F_squared = self.F**2 - - # Get s-vectors - s_vectors = self.get_scattering_vectors() - - # Compute E-values - E, E_squared, shell_idx = F_squared_to_E_values( - F_squared, s_vectors, n_shells=n_shells, d_min=d_min, d_max=d_max - ) - - # Store results - self.E = E - self.E_squared = E_squared - self.radial_shell_indices = shell_idx - - if verbose: - print(f"E-value statistics:") - print(f" E range: [{E.min():.3f}, {E.max():.3f}]") - print(f" E mean: {E.mean():.3f}, std: {E.std():.3f}") - print(f" E² mean: {E_squared.mean():.3f} (should be ~1.0)") - return E - - def setup_scale(self, scale: Optional[float] = None) -> float: - """ - Set overall scale factor, parametrized in log space. - - Parameters - ---------- - scale : float, optional - If provided, sets the scale factor directly (stored as its log). - If None (default), the scale defaults to 1.0 (``log_scale = 0.0``). - - Returns - ------- - ReflectionData - Self, for method chaining. - """ - if scale is None: - self.log_scale = torch.tensor( - 0.0, device=self.device, requires_grad=False, dtype=dtypes.float - ) - else: - self.log_scale = torch.log( - torch.tensor( - scale, device=self.device, requires_grad=False, dtype=dtypes.float - ) - ) - return self - def get_corrected_data(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Return amplitudes and sigmas, shape (N,), in this dataset's units. + + Raw datasets return their measurements; ScaledDataset exposes the live + scale correction through the same observation attributes. """ - Get the anisotropy-corrected, scaled (F, F_sigma). + return self.F, self.F_sigma - Returns - ------- - Tuple[torch.Tensor, torch.Tensor] - Full-size F and F_sigma with ``exp(log_scale)`` and ``U_aniso`` - applied. + def get_corrected_intensities(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Return intensities and sigmas, shape (N,), in this dataset's units. Raises ------ ValueError - If ``setup_scale`` / ``setup_anisotropy`` have not run. + If no intensity observations are available. """ - - if not hasattr(self, "log_scale") or self.log_scale is None: - raise ValueError("Scale not set up. Call setup_scale() first.") - if not hasattr(self, "U_aniso") or self.U_aniso is None: - raise ValueError("Anisotropy not set up. Call setup_anisotropy() first.") - F_corrected, F_sigma_corrected = self.apply_anisotropy_correction() - scale_factor = torch.exp(self.log_scale) - F_scaled = F_corrected * scale_factor - F_sigma_scaled = F_sigma_corrected * scale_factor - - return F_scaled, F_sigma_scaled + if self.I is None: + raise ValueError("No intensities on this dataset (I/SIGI required)") + return self.I, self.I_sigma def generate_validation_set( self, @@ -4027,23 +3546,3 @@ def generate_validation_set( f"free={n_free} ({100*n_free/total:.1f}%), " f"val={n_val} ({100*n_val/total:.1f}%)" ) - - def parameters(self) -> List[Parameter]: - """ - The scaling tensors (``log_scale``, ``U_aniso``) to optimize. - - Despite the ``List[Parameter]`` annotation these are plain tensors with - ``requires_grad=False``; the caller must call ``requires_grad_(True)`` - before handing them to an optimizer (see ``DatasetCollection.scale``). - - Returns - ------- - list of torch.Tensor - Whichever of the two are set. - """ - params = [] - if self.log_scale is not None: - params.append(self.log_scale) - if self.U_aniso is not None: - params.append(self.U_aniso) - return params diff --git a/torchref/io/datasets/scaled_dataset.py b/torchref/io/datasets/scaled_dataset.py new file mode 100644 index 00000000..e8d44412 --- /dev/null +++ b/torchref/io/datasets/scaled_dataset.py @@ -0,0 +1,160 @@ +"""ReflectionData subclass exposing live scaler-owned observation corrections.""" + +from dataclasses import fields +from typing import TYPE_CHECKING + +import torch + +from .reflection_data import ReflectionData + +if TYPE_CHECKING: + from torchref.scaling.dataset_scaler import DatasetScaler + + +class _ScaledObservation: + """Keep dataclass initialization in raw storage and scale only public reads.""" + + def __init__(self, name, power): + self.name, self.power = name, power + + def __get__(self, obj, owner=None): + if obj is None: + return None + raw = obj.__dict__.get("_raw_" + self.name) + scaler = obj.__dict__.get("scaler") + if raw is None or scaler is None: + return raw + correction = scaler(obj.scale_key, obj.hkl).to(raw) + return raw * correction.pow(self.power) + + def __set__(self, obj, value): + if obj.__dict__.get("scaler") is not None: + raise AttributeError( + "Scaled observations are read-only; edit the raw dataset before scaling" + ) + obj.__dict__["_raw_" + self.name] = value + + +class ScaledDataset(ReflectionData): + """Expose a raw dataset through one live row of a shared DatasetScaler. + + Parameters + ---------- + data : ReflectionData + Measurements and metadata, copied without changing the source. Passing + a scaled dataset uses its raw observations, never a second correction. + scaler : DatasetScaler + Shared parameter owner, retained strongly even outside a collection. + key : str + Stable dataset key in the scaler. + + Notes + ----- + F/F_sigma (amplitude units) and I/I_sigma (intensity units), shape (N,), + are corrected read-only expressions. Explicit *_raw properties expose the + stored measurements. Selection/copy preserves the shared scaler; moving a + view also moves the shared scaler. Raw source datasets remain unchanged. + """ + + F = _ScaledObservation("F", 1) + F_sigma = _ScaledObservation("F_sigma", 1) + I = _ScaledObservation("I", 2) + I_sigma = _ScaledObservation("I_sigma", 2) + + def __init__(self, data: ReflectionData, scaler: "DatasetScaler", key: str) -> None: + if key not in scaler.keys: + raise KeyError(key) + raw = data.raw_data() if isinstance(data, ScaledDataset) else data + raw = raw.__select__(torch.arange(len(raw), device=raw.device)) + raw.spacegroup = raw.spacegroup.copy() + self._install_raw(raw) + self.scaler = scaler + self.scale_key = key + self.to(scaler.device) + + def _install_raw(self, raw): + self.scaler = None + super().__init__( + **{f.name: getattr(raw, f.name) for f in fields(ReflectionData)} + ) + self.masks = raw.masks + self.source = None + + @property + def F_raw(self) -> torch.Tensor | None: + """Uncorrected amplitudes, shape (N,), in input amplitude units.""" + return self._raw_F + + @property + def F_sigma_raw(self) -> torch.Tensor | None: + """Uncorrected amplitude sigmas, shape (N,), in input amplitude units.""" + return self._raw_F_sigma + + @property + def I_raw(self) -> torch.Tensor | None: + """Uncorrected intensities, shape (N,), in input intensity units.""" + return self._raw_I + + @property + def I_sigma_raw(self) -> torch.Tensor | None: + """Uncorrected intensity sigmas, shape (N,), in input intensity units.""" + return self._raw_I_sigma + + def raw_data(self) -> ReflectionData: + """Return an independent parameter-free copy of the raw observations.""" + values = { + f.name: getattr(self, f.name) + for f in fields(ReflectionData) + if f.name not in ("F", "F_sigma", "I", "I_sigma", "source") + } + values.update( + { + name: getattr(self, name + "_raw") + for name in ("F", "F_sigma", "I", "I_sigma") + } + ) + raw = ReflectionData(**values) + raw.masks = self.masks + result = raw.__select__(torch.arange(len(raw), device=raw.device)) + result.source = None + return result + + def __select__(self, indices: torch.Tensor, op=None) -> "ScaledDataset": + """Select reflection indices or a boolean mask, preserving live scaling.""" + raw = self.raw_data().__select__(indices, op=op) + return ScaledDataset(raw, self.scaler, self.scale_key) + + def copy(self) -> "ScaledDataset": + """Copy observations and metadata while sharing the scaler parameters.""" + return ScaledDataset(self.raw_data(), self.scaler, self.scale_key) + + def __deepcopy__(self, memo: dict) -> "ScaledDataset": + """Copy this view without duplicating the shared parameter owner.""" + result = self.copy() + memo[id(self)] = result + return result + + def validate_hkl( + self, hkl_ref: torch.Tensor, *, identity_hkl: torch.Tensor | None = None + ) -> "ScaledDataset": + """Align this view to HKL (H, 3), without modifying its source or scaler.""" + raw = self.raw_data().validate_hkl(hkl_ref, identity_hkl=identity_hkl) + scaler, key = self.scaler, self.scale_key + self._install_raw(raw) + self.scaler, self.scale_key = scaler, key + return self + + def _get_state(self) -> dict: + return { + "raw": self.raw_data()._get_state(), + "scaler": self.scaler.get_state(), + "key": self.scale_key, + } + + @classmethod + def _from_state(cls, state: dict, device=None) -> "ScaledDataset": + from torchref.scaling.dataset_scaler import DatasetScaler + + scaler = DatasetScaler.from_state(state["scaler"], device) + raw = ReflectionData._from_state(dict(state["raw"]), device) + return cls(raw, scaler, state["key"]) diff --git a/torchref/io/hkl.py b/torchref/io/hkl.py new file mode 100644 index 00000000..50864b15 --- /dev/null +++ b/torchref/io/hkl.py @@ -0,0 +1,142 @@ +"""Reader for CrystFEL partialator reflection lists (`.hkl`). + +CrystFEL partialator output format:: + + CrystFEL reflection list version 2.0 + Symmetry: 2/m_uab + h k l I phase sigma(I) nmeas + 0 0 8 15631.77 - 2997.88 499 + ... + +The file carries intensities, sigmas, nmeas counts (and optional phase +strings), but **not** unit-cell or space-group metadata — those +typically live alongside in a ``.cell`` file. Cell and spacegroup must +therefore be supplied by the caller. + +Usage mirrors :class:`torchref.io.mtz.MTZReader`:: + + reader = HKLReader(verbose=1).read("td1.hkl", + cell=[a,b,c,al,be,ga], + spacegroup="P 1 21 1") + data_dict, cell, spacegroup = reader() + +The returned ``data_dict`` conforms to the intensity-path contract of +:meth:`ReflectionData.load` — ``{"HKL": (N,3) int, "I": (N,) float, +"SIGI": (N,) float, "I_col": "I (CrystFEL)"}`` — so French-Wilson kicks +in automatically and amplitudes are derived downstream. +""" + +from typing import Any, Optional, Tuple, Union + +import numpy as np + + +class HKLReader: + """Reader for CrystFEL partialator `.hkl` reflection lists.""" + + def __init__(self, verbose: int = 0): + self.verbose = verbose + self.data: Optional[dict] = None + self.cell: Optional[np.ndarray] = None + self.spacegroup: Optional[str] = None + self.nmeas: Optional[np.ndarray] = None + + def read( + self, + filepath: str, + cell: Union[list, tuple, np.ndarray, Any], + spacegroup: Union[str, Any], + ) -> "HKLReader": + """Parse a CrystFEL `.hkl` file. + + Parameters + ---------- + filepath : str + Path to the CrystFEL reflection list. + cell : list | tuple | np.ndarray | torchref.symmetry.Cell | torch.Tensor + Unit cell (a, b, c, alpha, beta, gamma). If a ``Cell`` object + is passed, ``.data`` is extracted. + spacegroup : str | gemmi.SpaceGroup | torchref.symmetry.SpaceGroup + Space group. If a wrapper object is passed, its ``.hm`` or + ``.short_name()`` is used. + """ + if self.verbose > 1: + print(f"Reading CrystFEL hkl file: {filepath}") + + if hasattr(cell, "data"): # torchref.symmetry.Cell + cell = cell.data + if hasattr(cell, "detach"): # torch.Tensor + cell = cell.detach().cpu().numpy() + cell_arr = np.asarray(cell, dtype=float).reshape(-1) + if cell_arr.size != 6: + raise ValueError( + f"cell must have 6 entries (a, b, c, al, be, ga); got {cell_arr.size}" + ) + self.cell = cell_arr + + if isinstance(spacegroup, str): + self.spacegroup = spacegroup + elif hasattr(spacegroup, "hm"): # torchref.symmetry.SpaceGroup + self.spacegroup = spacegroup.hm + elif hasattr(spacegroup, "short_name"): # gemmi.SpaceGroup + self.spacegroup = spacegroup.short_name() + else: + raise ValueError( + f"Cannot normalize spacegroup of type {type(spacegroup)}" + ) + + h_list, k_list, l_list, I_list, sig_list, n_list = [], [], [], [], [], [] + in_header = True + with open(filepath) as f: + for line in f: + if in_header: + if line.strip().startswith("h "): + in_header = False + continue + s = line.split() + if len(s) < 7 or not s[0].lstrip("-").isdigit(): + continue + h_list.append(int(s[0])) + k_list.append(int(s[1])) + l_list.append(int(s[2])) + I_list.append(float(s[3])) + sig_list.append(float(s[5])) + n_list.append(int(s[6])) + + if not h_list: + raise ValueError(f"No reflections parsed from {filepath}") + + hkl = np.column_stack([h_list, k_list, l_list]).astype(np.int32) + I = np.asarray(I_list, dtype=np.float64) + sig = np.asarray(sig_list, dtype=np.float64) + self.nmeas = np.asarray(n_list, dtype=np.int32) + + self.data = { + "HKL": hkl, + "I": I, + "SIGI": sig, + "I_col": "I (CrystFEL)", + } + + if self.verbose > 0: + print( + f"Parsed {len(hkl)} reflections from {filepath} " + f"(cell={cell_arr.tolist()}, spacegroup='{self.spacegroup}')" + ) + return self + + def __call__(self) -> Tuple[dict, np.ndarray, str]: + """Return ``(data_dict, cell, spacegroup)`` for ``ReflectionData.load()``.""" + if self.data is None: + raise RuntimeError("Call .read(path, cell, spacegroup) first.") + return self.data, self.cell, self.spacegroup + + +def read( + filepath: str, + cell: Union[list, tuple, np.ndarray, Any], + spacegroup: Union[str, Any], + verbose: int = 0, +) -> HKLReader: + """Shortcut: ``HKLReader(verbose=verbose).read(filepath, cell, spacegroup)``.""" + return HKLReader(verbose=verbose).read(filepath, cell, spacegroup) diff --git a/torchref/io/ihm.py b/torchref/io/ihm.py index 30ca1c83..06d4de65 100644 --- a/torchref/io/ihm.py +++ b/torchref/io/ihm.py @@ -40,6 +40,42 @@ def _check_ihm_available(): ) +def _add_group_with_independent_populations(collection, name, fractions): + """Add an IHM model group, keeping its populations exactly as deposited. + + ``ModelCollection`` stores populations as one activation fraction shared across + timepoints plus a per-timepoint branching, which is the right model for a kinetic + series driven by a single pump. An IHM ensemble is not that: its model groups carry + arbitrary, independently deposited populations, and two groups may well disagree + about how much of the sample is in the reference state. + + So the group is registered at whatever activation the collection already holds and + its deposited fractions are installed as an override, which ``fractions`` returns + verbatim and ``write_ihm`` therefore round-trips unchanged. + """ + import torch + + try: + collection.add_timepoint(name, fractions=fractions) + return + except ValueError: + pass + + n = len(fractions) + placeholder = [0.0] * n + placeholder[0] = 1.0 - float(collection.alpha_mean) + if n > 1: + placeholder[1] = 1.0 - placeholder[0] + collection.add_timepoint(name, fractions=placeholder) + collection[name].set_fraction_override( + torch.tensor( + fractions, + dtype=collection._activation_logit.dtype, + device=collection._activation_logit.device, + ) + ) + + class IHMReader: """ Read IHM mmCIF files into torchref ModelCollection + IHMEnsembleMapping. @@ -507,7 +543,9 @@ def build_model_collection( if is_dark: collection.add_dark(fractions=fractions) else: - collection.add_timepoint(group.name, fractions=fractions) + _add_group_with_independent_populations( + collection, group.name, fractions + ) return collection diff --git a/torchref/io/ihm_mapping.py b/torchref/io/ihm_mapping.py index ba8853f7..d713769a 100644 --- a/torchref/io/ihm_mapping.py +++ b/torchref/io/ihm_mapping.py @@ -12,7 +12,7 @@ --------------- IHM state -> base model (ModelFT) in ModelCollection IHM model group -> timepoint entry (_SharedMixedModel) in ModelCollection -IHM population fraction -> fraction_params in _SharedMixedModel +IHM population fraction -> activation / branching on ModelCollection """ from dataclasses import dataclass, field diff --git a/torchref/io/metadata.py b/torchref/io/metadata.py index 224ef4a1..4a1f378c 100644 --- a/torchref/io/metadata.py +++ b/torchref/io/metadata.py @@ -12,10 +12,43 @@ from __future__ import annotations import json -from dataclasses import dataclass, field, fields, asdict +from dataclasses import asdict, dataclass, field, fields from datetime import date from typing import Any, Dict, List, Optional +#: Input records that describe the crystal, the sample and its chemistry. +#: Refinement moves atoms; it does not invalidate any of these, so they are +#: carried through. Listed in the order the PDB format mandates -- +#: ``render_pdb_header`` emits them in this sequence and relies on it. +#: Deliberately absent: AUTHOR and JRNL (they credit the deposited entry, not +#: this refinement), HELIX/SHEET (nothing here computes secondary structure, so +#: they could only be stale or wrong), and HEADER/REVDAT/OBSLTE/CAVEAT/SPLIT +#: (assertions about a PDB entry that this file is not). +_KEEP_RECORDS = ( + "TITLE", "COMPND", "SOURCE", "KEYWDS", "EXPDTA", "MDLTYP", + "DBREF", "DBREF1", "DBREF2", "SEQADV", "SEQRES", "MODRES", + "HET", "HETNAM", "HETSYN", "FORMUL", + "SSBOND", "LINK", "CISPEP", "SITE", +) + +#: REMARK numbers dropped from the input. 2 is the resolution, which we +#: regenerate; 3 is the refinement, which is ours to write and whose statistics +#: describe a model we just replaced; 500 lists geometry outliers of those same +#: superseded coordinates. +_DROP_REMARKS = {2, 3, 500} + +#: mmCIF categories carried over from an input CIF -- entity, sequence and +#: connectivity, i.e. the CIF counterpart of ``_KEEP_RECORDS``. The input's +#: ``_refine`` is NOT here: it is the previous program's statistics. +_KEEP_CIF_CATEGORIES = ( + "_entity", "_entity_poly", "_entity_poly_seq", "_struct_conn", + "_chem_comp", "_struct_ref", "_struct_ref_seq", "_exptl", +) + +#: Width of the label field in a REMARK 3 line, so the colons line up. Matches +#: the longest label we emit ("RESOLUTION RANGE HIGH (ANGSTROMS)"). +_REMARK3_LABEL_WIDTH = 33 + @dataclass class RefinementMetadata: @@ -28,6 +61,9 @@ class RefinementMetadata: ---------- program, program_version, refinement_method : str Refinement program identification. + target_function, optimizer : str + The function minimised and how, e.g. ``"MAXIMUM LIKELIHOOD"`` and + ``"LBFGS, 5 MACROCYCLES"``. Rendered only when set. resolution_high, resolution_low : float, optional Resolution limits ``d_min`` / ``d_max`` in Angstroms. n_reflections_work, n_reflections_test, n_reflections_all : int, optional @@ -46,16 +82,31 @@ class RefinementMetadata: Unit cell ``[a, b, c, alpha, beta, gamma]`` and space-group name. title, authors Structure title and author names. - passthrough_pdb_remarks, passthrough_cif_categories - Raw REMARK lines / mmCIF category items carried over from an input file. - custom_remarks : list of str - Extra REMARK 3 lines to append. + starting_model : str, optional + Input model this refinement started from (path or PDB ID). + rfree_selection : str, optional + Where the free-set flags came from. Filled from + ``ReflectionData.rfree_source``, so the values are that field's: + ``"MTZReader FreeR"`` for flags read from the input, or + ``"Generated (resolution-binned, ASU-grouped, seed N)"`` for a draw this + run made. Two refinements with different free sets have incomparable + R-free values, which is why it is recorded rather than inferred. + output_remarks : str + Author-supplied free text. Rendered only when non-empty. + software_chain : list of dict + Programs applied before this refinement, read from an input mmCIF's + ``_software`` loop, so ours appends to the chain instead of erasing it. + passthrough_pdb_remarks, passthrough_pdb_records, passthrough_cif_categories + Surviving REMARK lines, structural records keyed by record name, and + mmCIF category items carried over from an input file. """ # Program identification program: str = "TORCHREF" program_version: str = "" refinement_method: str = "" # e.g. "difference-refine", "LBFGS" + target_function: str = "" # e.g. "MAXIMUM LIKELIHOOD" + optimizer: str = "" # e.g. "LBFGS, 5 MACROCYCLES" # Resolution resolution_high: Optional[float] = None # d_min in Angstroms @@ -97,13 +148,25 @@ class RefinementMetadata: title: str = "" authors: List[str] = field(default_factory=list) - # Pass-through: raw header lines from input file + # Provenance + starting_model: Optional[str] = None + rfree_selection: Optional[str] = None + + # Author-supplied free text, rendered as REMARK 3 OTHER REFINEMENT REMARKS + # and _refine.details. Never generated -- if the author has nothing to say, + # the field stays empty and neither is emitted. + output_remarks: str = "" + + # Programs that touched the model before us, read from the input's + # _software loop so ours can be appended rather than replacing the chain. + software_chain: List[Dict[str, str]] = field(default_factory=list) + + # Pass-through from the input file: surviving REMARKs, structural records + # keyed by record name (see _KEEP_RECORDS), and mmCIF categories. passthrough_pdb_remarks: List[str] = field(default_factory=list) + passthrough_pdb_records: Dict[str, List[str]] = field(default_factory=dict) passthrough_cif_categories: Dict[str, Any] = field(default_factory=dict) - # Custom remarks - custom_remarks: List[str] = field(default_factory=list) - # ------------------------------------------------------------------ # # Serialization # ------------------------------------------------------------------ # @@ -142,6 +205,7 @@ def from_refinement(cls, refinement) -> RefinementMetadata: best-effort: anything unavailable is left unset, silently. """ import torch + from torchref import __version__ meta = cls(program_version=__version__) @@ -167,7 +231,7 @@ def from_refinement(cls, refinement) -> RefinementMetadata: try: rd = refinement.reflection_data with torch.no_grad(): - hkl, fobs, sigma, rfree_flags = rd() + hkl, fobs, sigma, rfree_flags = rd.hkl, rd.F, rd.F_sigma, rd.rfree_flags n_all = len(fobs) n_test = int(rfree_flags.sum().item()) if rfree_flags.dtype == torch.bool else int((~rfree_flags.bool()).sum().item()) n_work = n_all - n_test @@ -196,7 +260,7 @@ def from_refinement(cls, refinement) -> RefinementMetadata: # --- Geometry deviations (silently skip if no restraints) --- try: model = refinement.model - if model.initialized and model._restraints is not None: + if model.ctx.initialized and model._restraints is not None: restraints = model.restraints if hasattr(restraints, "bond_deviations"): with torch.no_grad(): @@ -238,6 +302,25 @@ def from_refinement(cls, refinement) -> RefinementMetadata: if getattr(refinement, "verbose", 0) > 0: print(f"Could not record solvent parameters: {exc}") + # --- Provenance --- + # Both are recorded on the objects already; the header just had no way + # to say them. rfree_source in particular is what distinguishes a test + # set read from the input file from one this run drew itself, and hence + # whether R-free is comparable to the number the input reported. + try: + input_file = refinement.model.ctx.input_file + if input_file: + meta.starting_model = str(input_file) + except Exception: + pass + + try: + source = refinement.reflection_data.rfree_source + if source: + meta.rfree_selection = source + except Exception: + pass + # --- Cell and spacegroup --- try: model = refinement.model @@ -255,45 +338,82 @@ def from_refinement(cls, refinement) -> RefinementMetadata: # ------------------------------------------------------------------ # @classmethod - def from_pdb_file(cls, filepath: str) -> RefinementMetadata: - """Extract header metadata from an existing PDB file. - - Captures TITLE, AUTHOR, and REMARK records for pass-through. + def from_pdb_file( + cls, filepath: str, *, supersede_refinement: bool = True + ) -> RefinementMetadata: + """Extract the carry-through header of an existing PDB file. + + Captures TITLE, the structural records in ``_KEEP_RECORDS`` and every + REMARK except those in ``_DROP_REMARKS``. AUTHOR and JRNL are absent + from ``_KEEP_RECORDS`` and so never collected: they credit whoever + deposited the entry, not this refinement. + + The input's REMARK 3 is dropped rather than carried, which is the whole + point -- a refined file that repeats the previous program's R-factors + alongside its own asserts two different refinements at once. + + Parameters + ---------- + supersede_refinement : bool, optional + Whether this file's refinement is about to be replaced -- the + default, and the case for refinement output. Pass ``False`` when + annotating a file without re-refining it: nothing supersedes the + existing REMARK 3 or AUTHOR records then, and dropping them would + lose statistics and credit that are still accurate. """ meta = cls() - remarks = [] + remarks: List[str] = [] + records: Dict[str, List[str]] = {} try: with open(filepath, "r") as f: for line in f: record = line[:6].strip() - if record in ("ATOM", "HETATM"): + # MODEL as well as the atoms: it opens the coordinate + # section in a multi-model file. + if record in ("ATOM", "HETATM", "MODEL"): break - if record == "TITLE": + if record == "AUTHOR" and not supersede_refinement: + for author in line[10:].strip().split(","): + if author.strip(): + meta.authors.append(author.strip()) + elif record == "TITLE": + # Kept on `title` rather than as a raw record so the + # --title override has something to override. title_text = line[10:].strip() - if meta.title: - meta.title += " " + title_text - else: - meta.title = title_text - elif record == "AUTHOR": - author_text = line[10:].strip() - # Authors are comma-separated in PDB - for author in author_text.split(","): - author = author.strip() - if author: - meta.authors.append(author) - elif record.startswith("REMARK"): + meta.title = ( + meta.title + " " + title_text if meta.title else title_text + ) + elif record == "REMARK": + try: + number = int(line[7:10]) + except ValueError: + # A REMARK with no parsable number is not one we can + # judge, so leave it out. + continue + if supersede_refinement and number in _DROP_REMARKS: + continue remarks.append(line.rstrip("\n")) + elif record in _KEEP_RECORDS: + records.setdefault(record, []).append(line.rstrip("\n")) meta.passthrough_pdb_remarks = remarks + meta.passthrough_pdb_records = records except Exception: pass return meta @classmethod def from_cif_file(cls, filepath: str) -> RefinementMetadata: - """Extract refinement metadata from an existing mmCIF file. + """Extract the carry-through metadata of an existing mmCIF file. - Captures ``_struct.title``, ``_audit_author.name``, and - ``_refine`` category items for pass-through. + Captures ``_struct.title``, the entity/sequence/connectivity categories + in ``_KEEP_CIF_CATEGORIES``, and the ``_software`` loop -- the last so + this refinement can append itself to the chain of programs rather than + presenting itself as the only one. + + The input's ``_refine`` category is deliberately NOT captured. It holds + the previous program's R-factors and resolution, which this refinement + supersedes; carrying them forward is the mmCIF form of the duplicated + REMARK 3 that ``from_pdb_file`` used to produce. """ meta = cls() try: @@ -302,40 +422,51 @@ def from_cif_file(cls, filepath: str) -> RefinementMetadata: doc = gemmi.cif.read(filepath) block = doc[0] - # Title title = block.find_value("_struct.title") if title and title != "?": meta.title = gemmi.cif.as_string(title) - # Authors - author_loop = block.find(["_audit_author.name"]) - if author_loop: - for row in author_loop: - name = gemmi.cif.as_string(row[0]) - if name and name != "?": - meta.authors.append(name) - - # _refine category pass-through - refine_cats = {} - for tag in block.find(["_refine."]): - # Collect all _refine.* pairs - pass - # Use find_values for individual items - for item_name in [ - "_refine.ls_R_factor_R_work", - "_refine.ls_R_factor_R_free", - "_refine.ls_d_res_high", - "_refine.ls_d_res_low", - "_refine.ls_number_reflns_R_work", - "_refine.ls_number_reflns_R_free", - "_refine.B_iso_mean", - ]: - val = block.find_value(item_name) - if val and val not in ("?", "."): - refine_cats[item_name] = gemmi.cif.as_string(val) - - if refine_cats: - meta.passthrough_cif_categories["_refine"] = refine_cats + # Prior programs, so ours lands at max(ordinal) + 1. + chain: List[Dict[str, str]] = [] + table = block.find( + "_software.", + ["name", "?version", "?classification", "?pdbx_ordinal", + "?description"], + ) + for row in table: + name = gemmi.cif.as_string(row[0]) + if not name or name in ("?", "."): + continue + entry = {"name": name} + for idx, key in ((1, "version"), (2, "classification"), + (3, "pdbx_ordinal"), (4, "description")): + if row.has(idx): + val = gemmi.cif.as_string(row[idx]) + if val and val not in ("?", "."): + entry[key] = val + chain.append(entry) + meta.software_chain = chain + + # Entity, sequence and connectivity categories, in whichever form + # the input used them (loop or key-value). + cats: Dict[str, Any] = {} + for item in block: + if item.loop is not None: + tags = list(item.loop.tags) + if tags[0].split(".")[0] not in _KEEP_CIF_CATEGORIES: + continue + cats[tags[0].split(".")[0]] = { + tag: [ + item.loop[r, c] for r in range(item.loop.length()) + ] + for c, tag in enumerate(tags) + } + elif item.pair is not None: + tag, val = item.pair + category = tag.split(".")[0] + if category in _KEEP_CIF_CATEGORIES: + cats.setdefault(category, {})[tag] = val + meta.passthrough_cif_categories = cats except Exception: pass @@ -372,9 +503,23 @@ def merge(self, other: RefinementMetadata) -> RefinementMetadata: a for a in other_val if a not in self_val ] setattr(merged, f.name, merged_authors) - elif f.name == "custom_remarks": - merged_remarks = list(self_val) + list(other_val) - setattr(merged, f.name, merged_remarks) + elif f.name == "passthrough_pdb_records": + merged_records = {k: list(v) for k, v in self_val.items()} + for key, rows in other_val.items(): + existing = merged_records.setdefault(key, []) + existing.extend([r for r in rows if r not in existing]) + setattr(merged, f.name, merged_records) + elif f.name == "software_chain": + # Keyed on name+version: re-refining with the same build should + # not add a second identical link to the chain. + merged_chain = list(self_val) + seen = {(e.get("name"), e.get("version")) for e in merged_chain} + for entry in other_val: + key = (entry.get("name"), entry.get("version")) + if key not in seen: + merged_chain.append(entry) + seen.add(key) + setattr(merged, f.name, merged_chain) else: # other takes precedence if non-None and non-default if other_val is not None and other_val != "" and other_val != []: @@ -388,56 +533,82 @@ def merge(self, other: RefinementMetadata) -> RefinementMetadata: # ------------------------------------------------------------------ # def render_pdb_header(self) -> str: - """Render metadata as PDB header records (REMARK 3, TITLE, AUTHOR). + """Render the header as PDB records, ready to precede CRYST1. - Returns - ------- - str - Multi-line string ready to insert into a PDB file. + Records come out in the order the PDB format mandates -- TITLE and the + entry-level records, then REMARKs in ascending numeric order with ours + slotted in at 3, then sequence, chemistry and connectivity. Only the + REMARK 3 block is generated; everything else is either carried through + from the input or supplied by the caller. """ lines: List[str] = [] + records = self.passthrough_pdb_records - # Pass-through remarks first (from input file) - for remark in self.passthrough_pdb_remarks: - lines.append(remark) + def _emit(*names: str) -> None: + for name in names: + lines.extend(records.get(name, [])) - # TITLE + # -- entry level -------------------------------------------------- # if self.title: _wrap_pdb_record(lines, "TITLE", self.title) + _emit("COMPND", "SOURCE", "KEYWDS", "EXPDTA", "MDLTYP") - # AUTHOR + # Only ever what the caller set explicitly: authors are not inherited + # from the input, since they credit that deposition and not this run. if self.authors: - author_str = ", ".join(self.authors) - _wrap_pdb_record(lines, "AUTHOR", author_str) + _wrap_pdb_record(lines, "AUTHOR", ", ".join(self.authors)) + + # -- REMARKs, ascending, ours at 3 -------------------------------- # + def _remark_number(line: str) -> int: + try: + return int(line[7:10]) + except ValueError: + return 0 + + passthrough = sorted(self.passthrough_pdb_remarks, key=_remark_number) + lines.extend(r for r in passthrough if _remark_number(r) < 3) + lines.extend(self._render_remark3()) + lines.extend(r for r in passthrough if _remark_number(r) > 3) - # REMARK 3 - Refinement statistics + # -- sequence, chemistry, connectivity ---------------------------- # + _emit("DBREF", "DBREF1", "DBREF2", "SEQADV", "SEQRES", "MODRES", + "HET", "HETNAM", "HETSYN", "FORMUL", + "SSBOND", "LINK", "CISPEP", "SITE") + + return "\n".join(lines) + "\n" + + def _render_remark3(self) -> List[str]: + """Build the REMARK 3 block: this refinement, and only this one.""" + lines: List[str] = [] lines.append("REMARK 3") lines.append("REMARK 3 REFINEMENT.") - lines.append( - f"REMARK 3 PROGRAM : {self.program} {self.program_version}".rstrip() - ) + _ident(lines, "PROGRAM", f"{self.program} {self.program_version}".strip()) if self.refinement_method: - lines.append( - f"REMARK 3 METHOD : {self.refinement_method}" - ) + _ident(lines, "METHOD", self.refinement_method) + if self.target_function: + _ident(lines, "TARGET", self.target_function) + if self.optimizer: + _ident(lines, "OPTIMIZER", self.optimizer) lines.append("REMARK 3") - # Data used in refinement lines.append("REMARK 3 DATA USED IN REFINEMENT.") _remark3(lines, "RESOLUTION RANGE HIGH (ANGSTROMS)", self.resolution_high, ".2f") _remark3(lines, "RESOLUTION RANGE LOW (ANGSTROMS)", self.resolution_low, ".2f") _remark3(lines, "NUMBER OF REFLECTIONS", self.n_reflections_all, "d") lines.append("REMARK 3") - # Fit to data lines.append("REMARK 3 FIT TO DATA USED IN REFINEMENT.") + # Where the free set came from, before the R-factors it conditions: + # R-free values from different test sets are not comparable, and the + # reader has no way to tell without this. + if self.rfree_selection: + _remark3(lines, "FREE R VALUE TEST SET SELECTION", self.rfree_selection) _remark3(lines, "R VALUE (WORKING SET)", self.r_work, ".4f") _remark3(lines, "FREE R VALUE", self.r_free, ".4f") _remark3(lines, "FREE R VALUE TEST SET SIZE (%)", self.percent_free, ".1f") _remark3(lines, "FREE R VALUE TEST SET COUNT", self.n_reflections_test, "d") lines.append("REMARK 3") - # B-values lines.append("REMARK 3 B VALUES.") # Wilson-plot B is not computed; passing None makes _remark3 render the # literal "NULL" here intentionally (not a bug). @@ -447,33 +618,35 @@ def render_pdb_header(self) -> str: _remark3(lines, "B MAX (A**2)", self.b_max, ".2f") lines.append("REMARK 3") - # RMS deviations lines.append("REMARK 3 RMS DEVIATIONS FROM IDEAL VALUES.") _remark3(lines, "BOND LENGTHS (A)", self.rmsd_bond_lengths, ".3f") _remark3(lines, "BOND ANGLES (DEGREES)", self.rmsd_bond_angles, ".2f") lines.append("REMARK 3") - # Model contents lines.append("REMARK 3 NUMBER OF NON-HYDROGEN ATOMS USED IN REFINEMENT.") _remark3(lines, "PROTEIN ATOMS", self.n_atoms_protein, "d") _remark3(lines, "SOLVENT ATOMS", self.n_atoms_solvent, "d") _remark3(lines, "TOTAL", self.n_atoms_total, "d") lines.append("REMARK 3") - # Solvent model if self.solvent_model_ksol is not None or self.solvent_model_bsol is not None: lines.append("REMARK 3 BULK SOLVENT MODELLING.") _remark3(lines, "K_SOL", self.solvent_model_ksol, ".4f") _remark3(lines, "B_SOL", self.solvent_model_bsol, ".2f") lines.append("REMARK 3") - # Custom remarks - for remark in self.custom_remarks: - lines.append(f"REMARK 3 {remark}") + if self.starting_model: + lines.append(f"REMARK 3 STARTING MODEL: {self.starting_model}") + lines.append("REMARK 3") - lines.append("REMARK 3") + # The only free text in the block, and the caller wrote all of it. + if self.output_remarks: + lines.append("REMARK 3 OTHER REFINEMENT REMARKS:") + for paragraph in self.output_remarks.splitlines(): + _wrap_remark3_text(lines, paragraph.strip()) + lines.append("REMARK 3") - return "\n".join(lines) + "\n" + return lines # ------------------------------------------------------------------ # # mmCIF rendering @@ -492,16 +665,40 @@ def render_cif_categories(self) -> Dict[str, Dict[str, str]]: """ cats: Dict[str, Dict[str, str]] = {} - # _software - sw = {} - sw["_software.name"] = self.program + # _software: every program applied to this model, in order, ours last. + # Always a loop, even with one entry -- that is what lets the next + # refinement append a link rather than overwrite the chain, which is the + # only record of prior work that mmCIF actually has room for. + chain = list(self.software_chain) + ordinals = [] + for entry in chain: + try: + ordinals.append(int(entry.get("pdbx_ordinal", 0))) + except (TypeError, ValueError): + pass + description = self.refinement_method or ", ".join( + part for part in (self.target_function, self.optimizer) if part + ) + ours = { + "name": self.program, + "classification": "refinement", + "pdbx_ordinal": str(max(ordinals, default=len(chain)) + 1), + } if self.program_version: - sw["_software.version"] = self.program_version - sw["_software.classification"] = "refinement" - if self.refinement_method: - sw["_software.description"] = self.refinement_method - sw["_software.pdbx_ordinal"] = "1" - cats["_software"] = sw + ours["version"] = self.program_version + if description: + ours["description"] = description + chain.append(ours) + columns = [ + key + for key in ("pdbx_ordinal", "name", "version", "classification", + "description") + if any(key in entry for entry in chain) + ] + cats["_software"] = { + f"_software.{key}": [entry.get(key, "?") for entry in chain] + for key in columns + } # _struct if self.title: @@ -541,9 +738,26 @@ def render_cif_categories(self) -> Dict[str, Dict[str, str]]: ref["_refine.solvent_model_param_ksol"] = f"{self.solvent_model_ksol:.4f}" if self.solvent_model_bsol is not None: ref["_refine.solvent_model_param_bsol"] = f"{self.solvent_model_bsol:.2f}" + # Free-set provenance sits beside the R-factors it conditions: the two + # R-free values either side of a changed test set are not comparable. + if self.rfree_selection: + ref["_refine.pdbx_R_Free_selection_details"] = self.rfree_selection + if self.starting_model: + ref["_refine.pdbx_starting_model"] = self.starting_model + if self.refinement_method: + ref["_refine.pdbx_method_to_determine_struct"] = self.refinement_method + if self.output_remarks: + ref["_refine.details"] = self.output_remarks if ref: cats["_refine"] = ref + # What this refinement started from. Standard category, and the only + # structured place to say it. + if self.starting_model: + cats["_pdbx_initial_refinement_model"] = _initial_model_category( + self.starting_model + ) + # _refine_ls_restr (geometry deviations, as loop) if self.rmsd_bond_lengths is not None or self.rmsd_bond_angles is not None: restr_types = [] @@ -582,6 +796,29 @@ def render_cif_categories(self) -> Dict[str, Dict[str, str]]: # ====================================================================== # +def _initial_model_category(starting_model: str) -> Dict[str, str]: + """Describe the starting model the way deposited entries do. + + A four-character stem that looks like a PDB ID (digit then three + alphanumerics, e.g. ``3GR5.pdb``) is reported as an accession code; anything + else is named in ``details`` and left unaccessioned rather than guessed at. + """ + import os + import re + + basename = os.path.basename(starting_model) + stem = os.path.splitext(basename)[0] + cat = { + "_pdbx_initial_refinement_model.id": "1", + "_pdbx_initial_refinement_model.type": "experimental model", + "_pdbx_initial_refinement_model.details": basename, + } + if re.fullmatch(r"[0-9][A-Za-z0-9]{3}", stem): + cat["_pdbx_initial_refinement_model.source_name"] = "PDB" + cat["_pdbx_initial_refinement_model.accession_code"] = stem.upper() + return cat + + def _remark3( lines: List[str], label: str, value: Any, fmt: str = "" ) -> None: @@ -594,8 +831,59 @@ def _remark3( formatted = f"{value:{fmt}}" else: formatted = "NULL" - line = f"REMARK 3 {label} : {formatted}" - lines.append(line) + # Pad the label so the colons align down the block. Callers used to have to + # pre-pad their own labels, and the ones that forgot rendered ragged. + lines.append(f"REMARK 3 {label:<{_REMARK3_LABEL_WIDTH}} : {formatted}") + + +def _ident(lines: List[str], label: str, value: str) -> None: + """Append a ``REMARK 3 LABEL : value`` line, wrapped if long. + + Overflow continues on a further line whose label field is blank and whose + colon stays in the same column, which is what REFMAC does with its own + long values:: + + REMARK 3 AUTHORS : MURSHUDOV,SKUBAK,LEBEDEV,PANNU,STEINER, + REMARK 3 : NICHOLLS,WINN,LONG,VAGIN + + Without this an optimizer description naming the cycles, mode, ADP model and + scale target runs past column 80. + """ + head = f"REMARK 3 {label:<12}: " + cont = f"REMARK 3 {'':<12}: " + width = 80 - len(head) + prefix, current = head, "" + for word in value.split(): + if current and len(current) + 1 + len(word) > width: + lines.append(prefix + current) + prefix, current = cont, word + else: + current = current + " " + word if current else word + if current or prefix is head: + lines.append(prefix + current) + + +def _wrap_remark3_text(lines: List[str], text: str) -> None: + """Append free text as continuation-free ``REMARK 3`` lines. + + REMARK records have no continuation-number field -- unlike TITLE or AUTHOR, + they simply repeat the same number -- so this wraps on width alone. An empty + paragraph becomes a bare ``REMARK 3`` spacer. + """ + prefix = "REMARK 3 " + if not text: + lines.append("REMARK 3") + return + width = 80 - len(prefix) + current = "" + for word in text.split(): + if current and len(current) + 1 + len(word) > width: + lines.append(prefix + current) + current = word + else: + current = current + " " + word if current else word + if current: + lines.append(prefix + current) def _wrap_pdb_record(lines: List[str], record: str, text: str) -> None: diff --git a/torchref/io/pdb.py b/torchref/io/pdb.py index 8e431e79..c7d749f8 100644 --- a/torchref/io/pdb.py +++ b/torchref/io/pdb.py @@ -410,6 +410,15 @@ def extract_pdb_headers(filepath: str) -> list: return headers +#: Columns of the LINK-record table that ``Model.load`` reads off a reader's ``.links``. +#: Shared by the PDB and mmCIF readers so the topology builder sees one schema. +LINK_COLUMNS = ( + "name1", "altloc1", "resname1", "chainid1", "resseq1", "icode1", + "name2", "altloc2", "resname2", "chainid2", "resseq2", "icode2", + "length", +) + + def extract_link_records(filepath: str, verbose: int = 0) -> pd.DataFrame: """Parse LINK records from a PDB file (PDB v3.3 format). @@ -476,14 +485,7 @@ def extract_link_records(filepath: str, verbose: int = 0) -> pd.DataFrame: if verbose > 1: print(f"Warning: skipping malformed LINK: {line.rstrip()}") - df = pd.DataFrame( - rows, - columns=[ - "name1", "altloc1", "resname1", "chainid1", "resseq1", "icode1", - "name2", "altloc2", "resname2", "chainid2", "resseq2", "icode2", - "length", - ], - ) + df = pd.DataFrame(rows, columns=list(LINK_COLUMNS)) if verbose > 0 and (len(df) or skipped_sym or skipped_bad): print( f"LINK records: parsed {len(df)}, " @@ -492,7 +494,7 @@ def extract_link_records(filepath: str, verbose: int = 0) -> pd.DataFrame: return df -def write(df: pd.DataFrame, filepath: str, template: str = None, metadata=None) -> None: +def write(df: pd.DataFrame, filepath: str, metadata=None) -> None: """ Write a DataFrame to a PDB file. @@ -504,9 +506,6 @@ def write(df: pd.DataFrame, filepath: str, template: str = None, metadata=None) tempfactor, element, charge. filepath : str Output PDB filename. - template : str, optional - PDB template file to copy header from. Deprecated in favour of - ``metadata``; no ``DeprecationWarning`` is emitted when it is used. metadata : RefinementMetadata, optional Metadata to render as PDB header (REMARK 3, TITLE, etc.). @@ -525,14 +524,6 @@ def write(df: pd.DataFrame, filepath: str, template: str = None, metadata=None) if metadata is not None: n.write(metadata.render_pdb_header()) - # Copy template header if provided (deprecated path) - if template is not None: - with open(template) as t: - for line in t: - if "REMARK" not in line and "ATOM" in line: - break - n.write(line) - # Write CRYST1 record if cell info available (directly before atoms) try: cell = df.attrs["cell"] @@ -613,8 +604,8 @@ def write(df: pd.DataFrame, filepath: str, template: str = None, metadata=None) s = ( f"{str(ATOM):<6}{int(serial):>5} {name_field}{str(altloc):>1}" f"{str(resname):>3}{str(chainid):>2}{int(resseq):>4}{str(icode):>4}" - f"{round(x, 3):>8}{round(y, 3):>8}{round(z_coord, 3):>8}" - f"{round(occupancy, 3):>6.2f}{round(tempfactor, 2):>6}" + f"{x:>8.3f}{y:>8.3f}{z_coord:>8.3f}" + f"{occupancy:>6.2f}{tempfactor:>6.2f}" f"{str(element):>12}{charge:>2}\n" ) n.write(s) @@ -730,8 +721,8 @@ def write_multi_model( s = ( f"{str(ATOM):<6}{int(serial):>5} {name_field}{altloc:>1}" f"{resname:>3}{chainid:>2}{resseq:>4}{icode:>4}" - f"{round(x, 3):>8}{round(y, 3):>8}{round(z_coord, 3):>8}" - f"{round(occupancy, 3):>6.2f}{round(tempfactor, 2):>6}" + f"{x:>8.3f}{y:>8.3f}{z_coord:>8.3f}" + f"{occupancy:>6.2f}{tempfactor:>6.2f}" f"{element:>12}{charge_str:>2}\n" ) f.write(s) diff --git a/torchref/maps/ded_weights.py b/torchref/maps/ded_weights.py new file mode 100644 index 00000000..e09a93b8 --- /dev/null +++ b/torchref/maps/ded_weights.py @@ -0,0 +1,253 @@ +"""Registered per-reflection weights for light-minus-dark difference coefficients. + +A weight scheme turns the observed differences and their uncertainties into one weight +per reflection, normalised to mean one so that maps built from different schemes sit on +comparable scales. Three schemes are registered: + +``none`` + Every reflection weighted equally. +``inverse_variance`` + ``1 / sigma_diff**2``. Weights by precision alone; the right rule for averaging + estimates of one quantity, and the default for difference maps. +``sigma_d`` + The Wiener weight ``S / (S + sigma_diff**2)`` with ``S`` the expected true difference + power from :mod:`torchref.refinement.model_error_estimation.sigma_d`. Weights by the + signal fraction of each coefficient, so strong reflections whose expected difference + is large keep their weight. ``S`` is ``mean(dF**2) - mean(sigma**2)`` per shell, so + it inherits any miscalibration of ``sigma_diff``: where the reported sigmas are too + large the estimate finds no power and the weight vanishes, which turns the scheme + into a resolution cut. The count of such shells is reported as ``n_s2_clamped``; + a large fraction means the sigmas, not the data, are deciding the map. + +Plain tensors in and out. The weights live on the device of ``delta_obs``. The sigma_D +estimator is imported inside the scheme that needs it so that :mod:`torchref.maps` does +not import :mod:`torchref.refinement` at module load. +""" + +import warnings +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import torch + +from torchref.base.reciprocal.basis import get_scattering_vectors +from torchref.base.targets.xray_likelihoods import SIGMA_FLOOR_ABS, SIGMA_FLOOR_FRAC + +if TYPE_CHECKING: + from torchref.refinement.model_error_estimation.sigma_d import SigmaDConfig + +#: The selectable schemes, in the order they are reported. +SCHEMES = ("none", "inverse_variance", "sigma_d") +#: Scheme applied when none is named. Inverse variance, because ``sigma_d`` depends on +#: calibrated sigmas: on the 15 Sep campaign TorchSX's TD1 sigmas were ~1.5x too large at +#: high resolution, ``sigma_d`` zeroed 60-90 % of the shells there and the map agreement +#: rose in the bulk solvent as much as in the region of interest. +DEFAULT_SCHEME = "inverse_variance" +#: MTZ column carrying each scheme's weight (type ``W``); ``none`` writes no column. +WEIGHT_COLUMNS = {"inverse_variance": "W_IVW", "sigma_d": "W_SD"} + + +class DedWeightFallbackWarning(UserWarning): + """A requested weight scheme could not be evaluated and another was applied.""" + + +@dataclass(frozen=True) +class DedWeights: + """One scheme's weights and how they came about. + + Attributes + ---------- + scheme + The scheme requested. + applied + The scheme whose weights are in ``weights``; differs from ``scheme`` only after a + fallback, which ``diagnostics["fallback_reason"]`` then names. + weights + Per-reflection weights, shape ``(N,)``, mean one over finite positive entries. + diagnostics + Scheme-specific record: for ``sigma_d`` the fitted exponent, shrinkage sd, + clamp counters and the per-shell table. + """ + + scheme: str + applied: str + weights: torch.Tensor + diagnostics: dict = field(default_factory=dict) + + +def normalise_mean_one(w: torch.Tensor) -> torch.Tensor: + """Divide by the mean over all entries so the column averages one; unchanged when + that mean is not positive. + + Non-finite entries become zero, so a coefficient without a usable uncertainty drops + out of the map rather than poisoning it. Zero weights (shells without difference + power) stay zero and count in the mean, so the column mean is one whatever fraction + of the reflections carries weight. + """ + w = torch.where(torch.isfinite(w), w, torch.zeros_like(w)) + if w.numel() == 0 or not bool((w > 0).any()): + return w + return w / w.mean() + + +def reflection_geometry(hkl, cell, spacegroup, device, dtype): + """``(epsilon, d_star_sq)`` for ``hkl``: the reflection multiplicity from + ``spacegroup`` (ones when ``None``) and ``1/d**2`` in A^-2 from ``cell``, both on + ``device`` in ``dtype``.""" + from torchref.refinement.model_error_estimation.sigma_a import epsilon_from_hkl + + hkl_t = torch.as_tensor(hkl, device=device) + cell_t = cell if torch.is_tensor(cell) else cell.data + cell_t = torch.as_tensor(cell_t, device=device, dtype=dtype) + s = get_scattering_vectors(hkl_t, cell_t) + dss = (s * s).sum(dim=1).to(dtype) + eps = epsilon_from_hkl(hkl_t, spacegroup).to(device=device, dtype=dtype) + return eps, dss + + +def _inverse_variance(sigma_diff: torch.Tensor) -> torch.Tensor: + """``1 / sigma**2`` with sigma floored at a tenth of its median, so a reported zero + uncertainty gives a large finite weight rather than an infinite one.""" + finite = torch.isfinite(sigma_diff) & (sigma_diff >= 0) + positive = finite & (sigma_diff > 0) + if not bool(positive.any()): + return torch.zeros_like(sigma_diff) + floor = (sigma_diff[positive].median() * SIGMA_FLOOR_FRAC).clamp_min( + SIGMA_FLOOR_ABS + ) + sig = torch.where(finite, sigma_diff, torch.full_like(sigma_diff, float("inf"))) + return 1.0 / sig.clamp(min=floor) ** 2 + + +def compute_ded_weights( + scheme: str, + *, + delta_obs: torch.Tensor, + sigma_diff: torch.Tensor, + hkl: torch.Tensor, + cell, + spacegroup, + f_dark: torch.Tensor | None = None, + fit_mask: torch.Tensor | None = None, + sigma_d_config: "SigmaDConfig | None" = None, +) -> DedWeights: + """Per-reflection weights for one scheme. + + Parameters + ---------- + scheme : str + One of :data:`SCHEMES`. + delta_obs, sigma_diff : torch.Tensor + Signed observed differences and their propagated uncertainty, shape ``(N,)``, on + one common amplitude scale. + hkl : torch.Tensor + Miller indices, shape ``(N, 3)``. + cell : Cell or torch.Tensor + Unit cell, as a :class:`~torchref.symmetry.Cell` or its six parameters in A and + degrees. + spacegroup : SpaceGroup or None + For the reflection multiplicity; ``None`` means ones. + f_dark : torch.Tensor, optional + Dark amplitudes, shape ``(N,)``, for the sigma_D amplitude power law. + fit_mask : torch.Tensor, optional + Reflections entering the sigma_D fit; default every finite one. + sigma_d_config : SigmaDConfig, optional + Exponent and shrinkage settings for ``sigma_d``. + + Returns + ------- + DedWeights + Mean-one weights on ``delta_obs.device``. When ``sigma_d`` finds no difference + power in any shell, the inverse-variance weights are returned with + ``applied="inverse_variance"`` and a :class:`DedWeightFallbackWarning`. + """ + if scheme not in SCHEMES: + raise ValueError(f"Unknown DED weight scheme {scheme!r}; choose from {SCHEMES}") + sigma_diff = sigma_diff.reshape(-1).to(delta_obs.device, delta_obs.dtype) + if scheme == "none": + return DedWeights(scheme, scheme, torch.ones_like(sigma_diff)) + if scheme == "inverse_variance": + return DedWeights( + scheme, scheme, normalise_mean_one(_inverse_variance(sigma_diff)) + ) + + from torchref.refinement.model_error_estimation.sigma_d import ( + SigmaDConfig, + estimate_sigma_d, + sigma_d_per_reflection, + ) + + config = sigma_d_config if sigma_d_config is not None else SigmaDConfig() + d = delta_obs.reshape(-1) + eps, dss = reflection_geometry(hkl, cell, spacegroup, d.device, d.dtype) + f = f_dark.reshape(-1).to(d.device, d.dtype) if f_dark is not None else None + mask = ( + fit_mask.reshape(-1).to(d.device, torch.bool) + if fit_mask is not None + else torch.isfinite(d) & torch.isfinite(sigma_diff) + ) + shells = estimate_sigma_d( + d, sigma_diff, eps, dss, f, mask, gamma=config.gamma, shrink=config.shrink + ) + est = sigma_d_per_reflection(shells, dss, eps, f, sigma_diff) + diagnostics = { + "gamma": shells.gamma, + "gamma_fitted": shells.gamma_fitted, + "tau": shells.tau, + "curve_a": shells.curve_a, + "curve_b": shells.curve_b, + "degenerate": shells.degenerate, + "all_zero": shells.all_zero, + **shells.diagnostics, + "shells": { + "d_star_sq": shells.bin_dss.detach().cpu().tolist(), + "counts": shells.counts.detach().cpu().tolist(), + "B": shells.B.detach().cpu().tolist(), + "S2": shells.S2.detach().cpu().tolist(), + "Sigma_N_raw": shells.Sigma_N_raw.detach().cpu().tolist(), + "Sigma_N": shells.Sigma_N.detach().cpu().tolist(), + }, + "weight_sigma_d_raw": est.w.detach(), + } + if shells.all_zero or shells.degenerate: + reason = ( + "no difference power above the measurement variance in any shell " + f"(n_s2_clamped={shells.diagnostics['n_s2_clamped']})" + if shells.all_zero + else "fewer than two usable reflections" + ) + warnings.warn( + f"sigma_d weights: {reason}; applying inverse-variance weights instead", + DedWeightFallbackWarning, + stacklevel=2, + ) + diagnostics["fallback_reason"] = reason + return DedWeights( + scheme, + "inverse_variance", + normalise_mean_one(_inverse_variance(sigma_diff)), + diagnostics, + ) + return DedWeights(scheme, scheme, normalise_mean_one(est.w), diagnostics) + + +def all_ded_weights(**kwargs) -> dict[str, DedWeights]: + """Every registered scheme on the same inputs, keyed by scheme name. + + Takes the keyword arguments of :func:`compute_ded_weights` except ``scheme``. Used for + side-by-side reporting and for writing every weight column at once. + """ + return {scheme: compute_ded_weights(scheme, **kwargs) for scheme in SCHEMES} + + +__all__ = [ + "DEFAULT_SCHEME", + "SCHEMES", + "WEIGHT_COLUMNS", + "DedWeightFallbackWarning", + "DedWeights", + "all_ded_weights", + "compute_ded_weights", + "normalise_mean_one", + "reflection_geometry", +] diff --git a/torchref/maps/difference_map.py b/torchref/maps/difference_map.py index 6051c5ed..4dc010ea 100644 --- a/torchref/maps/difference_map.py +++ b/torchref/maps/difference_map.py @@ -14,7 +14,7 @@ from torchref.base.reciprocal.grid_operations import place_on_grid from torchref.io.datasets.collection import DatasetCollection from torchref.maps.map import Map -from torchref.symmetry.reciprocal_symmetry import expand_hkl +from torchref.symmetry import SpaceGroup from torchref.utils.device_resolution import resolve_device @@ -65,8 +65,16 @@ class DifferenceMap(Map): (see :mod:`torchref.maps.map`). """ - def __init__(self, data, data_reference, model, gridsize=None, - device: Optional[torch.device] = None): + def __init__( + self, + data, + data_reference, + model, + gridsize=None, + device: Optional[torch.device] = None, + units: str = "normalized", + scale: Optional[torch.Tensor] = None, + ): # Pin all three inputs onto one device before constructing the # DatasetCollection / super().__init__ — both consume tensors # from data.hkl / model and would otherwise inherit whichever @@ -82,15 +90,21 @@ def __init__(self, data, data_reference, model, gridsize=None, ) self._collection.add_dataset("perturbed", data) self._collection.scale() + self.data_reference = self._collection["reference"] + self.data_perturbed = self._collection["perturbed"] # Use reference dataset for cell, spacegroup, hkl via super().__init__ super().__init__( - data=data_reference, + data=self.data_reference, model=model, gridsize=gridsize, map_type="Fcalc", # placeholder, calculate() is overridden device=resolved, + units=units, ) + # Per-reflection observed-to-model scale over the reference dataset's full + # reflection list; dividing by it puts the differences in electrons. + self.scale = scale def calculate(self) -> torch.Tensor: """Compute the isomorphous difference map. @@ -112,17 +126,18 @@ def calculate(self) -> torch.Tensor: # Expand to P1 without Friedel mates (expand_to_p1() would reset # scaling, so expand manually via expand_hkl) - sg = self.data_reference.spacegroup or "P1" - hkl_p1, orig_idx, _ = expand_hkl( - hkl_asu, sg, + sg = self.data_reference.spacegroup or SpaceGroup("P1", device=hkl_asu.device) + hkl_p1, orig_idx, _ = sg.expand_hkl( + hkl_asu, include_friedel=False, remove_absences=True, device=hkl_asu.device, ) # Map scaled amplitudes to P1 (amplitudes are invariant under symmetry) - fobs_ref_p1 = fobs_ref[orig_idx] - fobs_pert_p1 = fobs_pert[orig_idx] - delta_f_p1 = fobs_pert_p1 - fobs_ref_p1 + delta_f = fobs_pert - fobs_ref + if self.scale is not None: + delta_f = delta_f / self.scale.to(delta_f)[mask_combined] + delta_f_p1 = delta_f[orig_idx] # Compute Fcalc for P1 hkl (for phases) fcalc_p1 = self.model.get_structure_factor(hkl_p1) @@ -141,6 +156,8 @@ def calculate(self) -> torch.Tensor: grid = place_on_grid( hkl_p1, coefficients_p1, gridsize, enforce_hermitian=True ) - self._map = torch.fft.fftn(grid, dim=(0, 1, 2), norm="forward").real + self._map = self._to_units( + torch.fft.fftn(grid, dim=(0, 1, 2), norm="forward").real + ) return self._map diff --git a/torchref/maps/map.py b/torchref/maps/map.py index 62570125..204df60c 100644 --- a/torchref/maps/map.py +++ b/torchref/maps/map.py @@ -23,7 +23,6 @@ from torchref.base.reciprocal.grid_operations import place_on_grid from torchref.io.cif import write_map -from torchref.symmetry.grid_utils import calculate_optimal_grid_size from torchref.utils.device_mixin import DeviceMixin from torchref.utils.device_resolution import resolve_device @@ -45,6 +44,11 @@ class Map(DeviceMixin): Default is ``"2Fo-Fc"``. Note ``"2Fo-Fc"`` is a *plain* 2Fo-Fc map (no figure-of-merit ``m`` and no sigma-A coefficient ``D``; i.e. ``m=1``, ``D=1``), not a likelihood-weighted 2mFo-DFc map. + units : str, optional + ``"normalized"`` (default) keeps the FFT's ``1/N`` normalisation; + ``"electrons"`` gives ``(1/V) sum_h F(h) exp(-2 pi i h.x)``, electrons per + cubic Angstrom, which is meaningful only when the coefficients are on the + absolute scale. Attributes ---------- @@ -73,6 +77,7 @@ class Map(DeviceMixin): """ VALID_MAP_TYPES = ("2Fo-Fc", "Fcalc") + VALID_UNITS = ("normalized", "electrons") def __init__( self, @@ -81,11 +86,15 @@ def __init__( gridsize: Optional[Tuple[int, int, int]] = None, map_type: str = "2Fo-Fc", device: Optional[torch.device] = None, + units: str = "normalized", ): if map_type not in self.VALID_MAP_TYPES: raise ValueError( f"map_type must be one of {self.VALID_MAP_TYPES}, got '{map_type}'" ) + if units not in self.VALID_UNITS: + raise ValueError(f"units must be one of {self.VALID_UNITS}, got '{units}'") + self.units = units self.device = resolve_device(data, model, device=device) self.data = data self.model = model @@ -104,10 +113,8 @@ def map_data(self) -> Optional[torch.Tensor]: def _determine_gridsize(self) -> Tuple[int, int, int]: """Determine optimal grid size from cell, resolution, and spacegroup.""" - cell_params = self.data.cell.data max_res = float(self.data.resolution.min()) - spacegroup = self.data.spacegroup.name - return calculate_optimal_grid_size(cell_params, max_res, spacegroup) + return self.data.spacegroup.optimal_grid_size(self.data.cell, max_res) def _compute_map_coefficients( self, fobs: torch.Tensor, fcalc: torch.Tensor @@ -169,9 +176,17 @@ def calculate(self) -> torch.Tensor: # FFT to real space: ρ(r) = (1/N) * sum_h F(h) * exp(-2πi h·r) # (norm="forward" applies the 1/N normalization, N = grid points) self._map = torch.fft.fftn(grid, dim=(0, 1, 2), norm="forward").real + self._map = self._to_units(self._map) return self._map + def _to_units(self, real_map: torch.Tensor) -> torch.Tensor: + """Rescale a ``1/N``-normalised FFT map to the configured units.""" + if self.units == "electrons": + volume = self.data.cell.volume.to(real_map.dtype) + return real_map * (real_map.numel() / volume) + return real_map + def write(self, filepath: str) -> int: """Write the map to a CCP4 file. diff --git a/torchref/model/__init__.py b/torchref/model/__init__.py index 685c2cee..58f91bc5 100644 --- a/torchref/model/__init__.py +++ b/torchref/model/__init__.py @@ -1,19 +1,21 @@ """Atomic models: coordinates, ADPs, occupancies and their structure factors. -:class:`Model` holds the refinable atomic parameters; :class:`ModelFT` adds -structure-factor calculation on top, through :class:`SfFFT` (FFT) or +:class:`Model` holds the refinable atomic parameters, with the crystallographic +context, atom table and provenance split out into :class:`ModelContext`; +:class:`ModelFT` adds +structure-factor calculation on top, through :class:`SfFFT` or :class:`SfDS` (direct summation). :class:`MixedModel` combines ModelFT states by population fraction (e.g. dark/light), and :class:`ModelCollection` keys mixtures by timepoint (``_SharedMixedModel`` is its non-re-registering variant). The wrappers from :mod:`torchref.model.parameter_wrappers` -- :class:`MixedTensor` and its ``Positive`` / ``Cholesky`` / ``Occupancy`` subclasses plus :class:`RigidXYZTensor` -- are the parametrizations that decide -which parameters are refinable. ``FFT`` is a deprecated alias for -:class:`SfFFT`. +which parameters are refinable. """ -from torchref.model.sf_fft import SfFFT, FFT +from torchref.model.sf_fft import SfFFT from torchref.model.sf_ds import SfDS +from torchref.model.context import ModelContext from torchref.model.mixed_model import MixedModel from torchref.model.model import Model from torchref.model.model_ft import ModelFT @@ -28,11 +30,11 @@ from torchref.model.rigid_xyz import RigidXYZTensor __all__ = [ - "FFT", "SfFFT", "SfDS", "MixedModel", "Model", + "ModelContext", "ModelFT", "MixedTensor", "PositiveMixedTensor", diff --git a/torchref/model/context.py b/torchref/model/context.py new file mode 100644 index 00000000..1b9fc66e --- /dev/null +++ b/torchref/model/context.py @@ -0,0 +1,148 @@ +"""The information half of a :class:`~torchref.model.model.Model`. + +:class:`ModelContext` holds what a model *is loaded from* and *sits in* -- the unit +cell, the space group, the atom table, the link records and the provenance -- as +opposed to what is being refined, which stays on the model as parameter wrappers and +per-atom buffers. + +Splitting it out means the crystallographic context can be passed to code that needs +only that (structure-factor engines, scalers, most targets) without handing over the +refinable state, and it keeps the model's own surface to parameters and behaviour. + +Mutable by design; prefer :meth:`ModelContext.copy` over editing in place. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, List, Optional + +from torchref.utils.device_mixin import DeviceMixin + +if TYPE_CHECKING: + import pandas + + from torchref.symmetry import Cell, SpaceGroup + + +@dataclass(eq=False, repr=False) +class ModelContext(DeviceMixin): + """Crystallographic context, atom bookkeeping and provenance for one model. + + Parameters + ---------- + cell : Cell or None + Unit cell, or None before a structure is loaded. + spacegroup : SpaceGroup or None + Space group, or None before a structure is loaded. + pdb : pandas.DataFrame or None + The atom table. Refreshed from the model's tensors only by + ``Model.update_pdb``, so it is stale between refinement steps by design. + links : list or None + Link records from the reader, used to build inter-residue restraints. + altloc_pairs : list + Index groups of alternative conformations, rebuilt by + ``Model.register_alternative_conformations``. + input_file : str or None + Path the structure was loaded from. + cif_path : str or None + Restraint dictionary path, if one was set. + verbose : int, default 1 + Verbosity level. + strip_H : bool, default True + Whether hydrogens were stripped on load. + hydrogens_in_xray : bool, default True + Whether hydrogens enter the structure-factor calculation. Restraints and the + non-bonded term see them either way; the bulk-solvent mask never does. + add_hydrogens : bool, default False + Generate hydrogens on load for residues that arrive without them. Ignored when + ``strip_H`` is set, which removes them again. + hydrogen_mode : str, default "free" + How hydrogen rows are parametrised: ``"riding"`` (positions derived from the + parent heavy atoms each forward, not refined), ``"free"`` (ordinary refinable + atoms) or ``"none"`` (the table holds no hydrogens). + initialized : bool, default False + Whether a structure has been loaded. ``if model:`` tests this. + + Notes + ----- + Deliberately does **not** carry the device or float dtype. Those are live + :class:`~torchref.utils.device_mixin.DeviceMixin` trackers that the traversal + rewrites in place on the object that owns the tensors, so they stay on the model + rather than becoming a second source of truth here. + + Holds no refinable parameters, so this is a dataclass rather than an + ``nn.Module``. + """ + + cell: Optional["Cell"] = None + spacegroup: Optional["SpaceGroup"] = None + pdb: Optional["pandas.DataFrame"] = None + links: Optional[List[Any]] = None + altloc_pairs: List[Any] = field(default_factory=list) + input_file: Optional[str] = None + cif_path: Optional[str] = None + verbose: int = 1 + strip_H: bool = True + hydrogens_in_xray: bool = True + add_hydrogens: bool = False + hydrogen_mode: str = "free" + initialized: bool = False + + def copy(self) -> "ModelContext": + """An independent copy. + + The atom table is deep-copied and the cell and space group are cloned, so + nothing is shared with the original. Cloning the space group matters now that + it is a mutable dataclass: sharing the reference would let an edit through one + model's context reach every model that was copied from it. + + Returns + ------- + ModelContext + New context sharing no mutable state with this one. + """ + return ModelContext( + cell=self.cell.clone() if self.cell is not None else None, + spacegroup=( + self.spacegroup.copy() if self.spacegroup is not None else None + ), + pdb=self.pdb.copy(deep=True) if self.pdb is not None else None, + links=list(self.links) if self.links is not None else None, + altloc_pairs=[ + tuple(t.clone() for t in group) for group in self.altloc_pairs + ], + input_file=self.input_file, + cif_path=self.cif_path, + verbose=self.verbose, + strip_H=self.strip_H, + hydrogens_in_xray=self.hydrogens_in_xray, + add_hydrogens=self.add_hydrogens, + hydrogen_mode=self.hydrogen_mode, + initialized=self.initialized, + ) + + @property + def crystal_key(self): + """Value identity of the crystal, or None while cell or space group is unset. + + Returns + ------- + tuple or None + ``(cell.key, spacegroup.key)``; hashable, so anything derived from the + crystal alone can be cached against it. + """ + if self.cell is None or self.spacegroup is None: + return None + return (self.cell.key, self.spacegroup.key) + + def __repr__(self) -> str: + n_atoms = 0 if self.pdb is None else len(self.pdb) + sg = None if self.spacegroup is None else self.spacegroup.name + return ( + f"ModelContext(spacegroup={sg!r}, n_atoms={n_atoms}, " + f"initialized={self.initialized})" + ) + + +__all__ = ["ModelContext"] diff --git a/torchref/model/disorder_field.py b/torchref/model/disorder_field.py new file mode 100644 index 00000000..a2911d3f --- /dev/null +++ b/torchref/model/disorder_field.py @@ -0,0 +1,1177 @@ +"""Node-field parametrization of the atomic displacement parameters. + +A disorder field stores disorder parameters on a small set of **nodes** and gives each +atom a distance-weighted mean of the nodes near it, so the parameter count scales with +node count rather than atom count. + +This is :class:`~torchref.model.parameter_wrappers.OccupancyTensor`'s collapse-and-expand +with a soft, distance-derived expansion in place of a fixed integer assignment: storage +is ``(K, 2)`` per node, ``forward()`` returns one B per atom. Two index spaces therefore +meet in this class, and callers must not mix them --- masks handed to +:meth:`~DisorderFieldTensor.update_refinable_mask` are in ATOM space, while +``refinable_mask`` and :meth:`get_refinable_count` are in NODE space. + +A node's position is anchored, not free: it is the centroid of the atoms in its anchor +cluster, plus an optional refinable offset. Anchoring keeps a node inside the molecule +and moving with the model, so the offset says only where it sits *relative* to the atoms +it serves and cannot wander off into solvent. +""" + +import math +from typing import Callable, Optional, Tuple + +import torch +from torch import nn + +from torchref.config import get_float_dtype, get_int_dtype, normalize_device +from torchref.model.parameter_wrappers import ( + MixedTensor, + chol_param_count, + psd_to_raw, + raw6_to_u6, + raw_to_cholesky, + u6_to_matrix, + u6_to_raw6, +) +from torchref.utils.utils import ModuleReference + +__all__ = [ + "DisorderFieldTensor", + "NodePayload", + "IsotropicPayload", + "AnisotropicPayload", + "ModeCovariancePayload", + "MODE_SETS", + "PAYLOAD_CODES", + "payload_code", + "payload_from_code", + "farthest_point_anchors", + "density_anchor_rows", + "build_neighbor_list", +] + + +def farthest_point_anchors(xyz: torch.Tensor, n_nodes: int) -> torch.Tensor: + """Pick ``n_nodes`` well-spread anchor atoms, then relax them onto local density. + + Greedy farthest-point selection seeded from the atom nearest the centroid, followed + by Lloyd iterations that move each anchor to the atom closest to its cluster mean. + Farthest-point alone favours extremities; the Lloyd pass pulls the anchors back onto + where atoms actually are, which is what a disorder field wants. + + Deterministic end to end --- no RNG --- so two processes produce the same anchors. + Restraint row order was hash-seed dependent in this package once; node placement is + not going to be. + + Parameters + ---------- + xyz : torch.Tensor + ``(N, 3)`` atom coordinates. + n_nodes : int + Number of anchors to pick. Clamped to ``N``. + + Returns + ------- + torch.Tensor + ``(K,)`` int64 atom indices, sorted ascending. + """ + n_atoms = int(xyz.shape[0]) + n_nodes = max(1, min(int(n_nodes), n_atoms)) + + centroid = xyz.mean(dim=0, keepdim=True) + first = int(torch.cdist(centroid, xyz).argmin()) + + chosen = [first] + d2_nearest = ((xyz - xyz[first]) ** 2).sum(-1) + for _ in range(n_nodes - 1): + nxt = int(d2_nearest.argmax()) + chosen.append(nxt) + d2_nearest = torch.minimum(d2_nearest, ((xyz - xyz[nxt]) ** 2).sum(-1)) + + anchors = torch.tensor(chosen, dtype=torch.int64, device=xyz.device) # dtype-ok: anchor atom indices; torch indexing requires int64 + + # Lloyd relaxation, snapping to real atoms so an anchor is always an atom index. + for _ in range(10): + assign = torch.cdist(xyz, xyz[anchors]).argmin(dim=1) + moved = anchors.clone() + for j in range(anchors.shape[0]): + members = (assign == j).nonzero(as_tuple=True)[0] + if members.numel() == 0: + continue + mean = xyz[members].mean(dim=0, keepdim=True) + moved[j] = members[int(torch.cdist(mean, xyz[members]).argmin())] + moved = torch.unique(moved) + if moved.shape[0] == anchors.shape[0] and bool((moved == anchors).all()): + break + anchors = moved + + return torch.sort(anchors).values + + +def density_anchor_rows(xyz: torch.Tensor, n_nodes: int): + """Anchor each node on its whole density cluster rather than on one atom. + + :func:`farthest_point_anchors` snaps every anchor onto an atom, which puts each node + exactly *on* an atom -- so narrowing its kernel isolates the atom it is already + standing on, and the node ends up owning a single atom outright. Anchoring on the + cluster instead places the node at the cluster centroid, generally between atoms, so + there is no atom for it to fall onto. + + Returns + ------- + tuple of torch.Tensor + ``(atom index per entry, node index per entry)``, flat and ragged, suitable for + ``DisorderFieldTensor(anchor_rows=...)``. Every node keeps at least its seed + atom, so no node is left with an empty neighbourhood. + """ + seeds = farthest_point_anchors(xyz, n_nodes) + assign = torch.cdist(xyz, xyz[seeds]).argmin(dim=1) + atom_idx = torch.arange(xyz.shape[0], dtype=torch.int64, device=xyz.device) # dtype-ok: arange atom indices; index requires int64 + + # A seed whose cluster somehow came out empty still needs a position. + present = torch.bincount(assign, minlength=seeds.shape[0]) > 0 + if not bool(present.all()): + missing = (~present).nonzero(as_tuple=True)[0] + atom_idx = torch.cat([atom_idx, seeds[missing]]) + assign = torch.cat([assign, missing]) + + order = torch.argsort(assign) + return atom_idx[order], assign[order] + + +def build_neighbor_list( + xyz: torch.Tensor, node_pos: torch.Tensor, k: int +) -> torch.Tensor: + """The ``k`` nearest nodes to each atom. + + A dense ``cdist`` plus ``topk``. Node counts are small by construction --- the whole + point of the representation --- so ``(N, K)`` stays cheap and a spatial cell list + (``topology.nonbonded.build_cell_list``) would only add machinery. Revisit if node + counts ever approach atom counts. + + Parameters + ---------- + xyz : torch.Tensor + ``(N, 3)`` atom coordinates. + node_pos : torch.Tensor + ``(K, 3)`` node positions. + k : int + Candidates per atom. Clamped to ``K``. + + Returns + ------- + torch.Tensor + ``(N, k)`` int64 node indices, nearest first. + """ + k = max(1, min(int(k), int(node_pos.shape[0]))) + d = torch.cdist(xyz, node_pos) + return d.topk(k, dim=1, largest=False).indices.contiguous() + + +def _wrap_accessor(xyz_fn): + """Hold a coordinate accessor without registering it as a submodule. + + ``model.xyz`` is itself an ``nn.Module``, so a plain assignment would enrol it in + this wrapper's module tree and drag it into ``state_dict``, ``.to()`` and + ``deepcopy``. :class:`~torchref.utils.utils.ModuleReference` exists for exactly that + and is what the device-conformance walker already knows how to follow. A bare + callable needs no wrapping. + """ + if isinstance(xyz_fn, nn.Module): + return ModuleReference(xyz_fn) + return xyz_fn + + +# ---------------------------------------------------------------------------------- +# Payload strategies. A payload says what one node carries and how that becomes a +# per-atom quantity; it knows nothing about where nodes are or how weights arise. +# +# Deliberately stateless. Node parameters live in one flat leaf on the field itself, +# because ``Model.parameters_of_types`` reads a single ``refinable_params`` per type, +# and a strategy owning parameters would break that. Each strategy declares only how +# many columns of that leaf it interprets. +# ---------------------------------------------------------------------------------- + + +class NodePayload: + """What a node carries, and how it becomes a per-atom ADP. + + Attributes + ---------- + width : int + Columns of node storage this payload interprets. + out_width : int + Components of the per-atom output: 1 for an isotropic B, 6 for a U tensor. + """ + + width: int = 1 + out_width: int = 1 + + def contributions(self, payload, xyz, node_pos, neighbor_list): + """``(n_atoms, k, out_width)``: what each candidate node offers each atom. + + ``xyz`` and ``node_pos`` are passed even though the payloads here ignore them, + because a payload with an r-dependence (TLS: constant, linear and quadratic in + the displacement from the node) needs them, and giving it the arguments now + means adding one later touches no shared code. + """ + raise NotImplementedError + + def fit(self, target, w_dense, epsilon, xyz, node_pos, neighbor_list): + """``(K, width)`` payload whose field reproduces ``target`` as closely as it can. + + Takes the same geometric context as :meth:`contributions` and for the same + reason: an r-dependent payload cannot build its modes without it. ``w_dense`` is + the ``(n_atoms, K)`` weight matrix at the seeded kernel widths, which is what + makes the payload-only problem linear. + """ + raise NotImplementedError + + def log_magnitude(self, payload): + """``(K,)`` log of each node's ADP magnitude, for a magnitude restraint. + + Lets a restraint price node values without branching on payload type. + """ + raise NotImplementedError + + +class IsotropicPayload(NodePayload): + """One isotropic B per node, stored as ``log B`` so it stays positive. + + The per-atom B is a convex combination of positive node values, so it is positive + without a clamp. + """ + + width = 1 + out_width = 1 + + def contributions(self, payload, xyz, node_pos, neighbor_list): + return torch.exp(payload[:, 0])[neighbor_list].unsqueeze(-1) + + def fit(self, target, w_dense, epsilon, xyz, node_pos, neighbor_list): + b = _ridged_solve(w_dense, target.unsqueeze(-1)).squeeze(-1) + return torch.log(b.clamp(min=epsilon)).unsqueeze(-1) + + def log_magnitude(self, payload): + return payload[:, 0] + + +class AnisotropicPayload(NodePayload): + """A full U tensor per node, positive-definite by construction. + + Stored as the six free parameters of a Cholesky factor, so ``U = L L^T`` is PD for + any parameter value -- the same device + :class:`~torchref.model.parameter_wrappers.CholeskyMixedTensor` uses for per-atom + ADPs, and for the same reason: an indefinite U makes the anisotropic B-matrix + singular and the structure-factor FFT returns NaN. + + Positive-definiteness survives the combination for free: the per-atom U is a convex + combination of PD matrices. Averaging in U space is what buys that -- averaging the + Cholesky parameters instead would be a different object with no such guarantee. + """ + + width = 6 + out_width = 6 + + def __init__(self, epsilon: float = 1e-3): + # Floor on the Cholesky diagonal, which bounds the smallest eigenvalue of U + # from below. Same default and same meaning as the per-atom wrapper. + self.epsilon = float(epsilon) + + def contributions(self, payload, xyz, node_pos, neighbor_list): + return raw6_to_u6(payload, self.epsilon)[neighbor_list] + + def fit(self, target, w_dense, epsilon, xyz, node_pos, neighbor_list): + """Fit six U components at once, then re-encode as Cholesky parameters. + + The per-atom U is linear in each component independently, so this is the same + ridged solve as the isotropic case with a six-column right-hand side. The + least-squares result is not constrained to be PD, which is why it goes back + through ``u6_to_raw6`` -- that projects onto PD by clamping eigenvalues. + """ + if target.ndim == 1: # a B target: lift to the equivalent isotropic U + u_iso = target / (8.0 * math.pi**2) + zero = torch.zeros_like(u_iso) + target = torch.stack([u_iso, u_iso, u_iso, zero, zero, zero], dim=1) + return u6_to_raw6(_ridged_solve(w_dense, target), self.epsilon) + + def log_magnitude(self, payload): + u6 = raw6_to_u6(payload, self.epsilon) + b_eq = (8.0 * math.pi**2 / 3.0) * (u6[:, 0] + u6[:, 1] + u6[:, 2]) + return torch.log(b_eq.clamp(min=1e-6)) + + +# ---------------------------------------------------------------------------------- +# Displacement-mode generators. A gradient mode is a constant 3x3 matrix G acting on the +# displacement r from the node, giving the displacement field psi(r) = G r. Rotation, +# dilation and deviatoric strain together span every linear displacement field, and +# splitting them that way is what lets a mode set stop partway. +# ---------------------------------------------------------------------------------- + +_SQ2 = math.sqrt(2.0) +_SQ3 = math.sqrt(3.0) +_SQ6 = math.sqrt(6.0) + +# Rotations are NOT normalised: psi_i(r) = e_i x r exactly, so that the rigid mode set +# reproduces the textbook TLS formula with no stray factor. The others are Frobenius +# normalised, which is a conditioning choice and nothing more. +_GENERATORS = { + "rotation": [ + [[0.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]], # e1 x r + [[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [-1.0, 0.0, 0.0]], # e2 x r + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], # e3 x r + ], + "dilation": [ + [[1 / _SQ3, 0.0, 0.0], [0.0, 1 / _SQ3, 0.0], [0.0, 0.0, 1 / _SQ3]], + ], + "deviatoric": [ + [[1 / _SQ2, 0.0, 0.0], [0.0, -1 / _SQ2, 0.0], [0.0, 0.0, 0.0]], + [[1 / _SQ6, 0.0, 0.0], [0.0, 1 / _SQ6, 0.0], [0.0, 0.0, -2 / _SQ6]], + [[0.0, 1 / _SQ2, 0.0], [1 / _SQ2, 0.0, 0.0], [0.0, 0.0, 0.0]], + [[0.0, 0.0, 1 / _SQ2], [0.0, 0.0, 0.0], [1 / _SQ2, 0.0, 0.0]], + [[0.0, 0.0, 0.0], [0.0, 0.0, 1 / _SQ2], [0.0, 1 / _SQ2, 0.0]], + ], +} + +#: Named mode sets, in order of expressiveness. Three translations are always present; +#: each entry lists the gradient modes added on top. +MODE_SETS = { + "constant": (), + "rigid": ("rotation",), + "rigid_dilation": ("rotation", "dilation"), + "affine": ("rotation", "dilation", "deviatoric"), +} + + +class ModeCovariancePayload(NodePayload): + """A node carries the covariance of its displacement modes; TLS is one mode set. + + Instead of storing an ADP and averaging it, store the *displacement field* the node + represents and take its covariance. With ``q`` modes ``Psi(r) = [psi_1(r) ... psi_q(r)]`` + the node's disorder is ``u(r) = Psi(r) c`` for a random coefficient vector ``c``, and + the ADP an atom at displacement ``r`` receives is:: + + U(r) = Psi(r) Sigma Psi(r)^T, Sigma = , Sigma = L L^T + + Three properties follow, and they are the whole reason for the form: + + * **Positive-semidefinite at every r, unconditionally**, because + ``U = (Psi L)(Psi L)^T``. An arbitrary polynomial in ``r`` carries no such + guarantee and goes indefinite somewhere --- and "somewhere" is the edge of the + node's region, exactly where the softmax weights have not yet decayed. + * **Spatial variation becomes intra-node and smooth by construction.** A constant-U + node can only express variation by having neighbours, so detail costs nodes, and + every added node is another kernel that can collapse onto a single atom. Here one + node's U already varies across its whole region, and it cannot spike. + * ``U(r)`` is **linear in Sigma**, so fitting stays a linear problem. + + Mode sets, from :data:`MODE_SETS`, with ``q(q+1)/2`` parameters per node: + + ================== === ======= ========================================= + set q params model + ================== === ======= ========================================= + ``constant`` 3 6 one U per node; same model as + :class:`AnisotropicPayload` + ``rigid`` 6 21 **exactly TLS** (20 determinable; ``tr S`` + is the one flat direction) + ``rigid_dilation`` 7 28 TLS plus uniform breathing + ``affine`` 12 78 full linear displacement field: TLS plus + shear and extension + ================== === ======= ========================================= + + With the rigid set this reproduces ``U(r) = T + A S + S^T A^T + A L A^T`` identically, + ``A`` being the matrix whose columns are ``e_i x r``: the classical TLS expression is + what ``Psi Sigma Psi^T`` expands to when the modes are three translations and three + rotations. Releasing the antisymmetry of the gradient -- the ``dilation`` and + ``deviatoric`` rungs -- gives domains that breathe and shear as well as rotate. + + Displacements are divided by the node layout's own length scale (median + nearest-neighbour node distance, detached) before the modes are built. That is pure + conditioning: without it the gradient modes carry a factor of the domain size against + the translations, and the curvature ratio between them runs to several hundred. It + is detached and derived from the layout for the reason + :class:`~torchref.refinement.targets.adp.NodeSmoothnessTarget` uses the same + quantity: the length scale is a property of where the nodes are, not something the + optimiser should tune. + + Memory scales as ``n_atoms * k * q^2`` for the gathered node factors, so this payload + is meant for the small-``K`` regime it was designed for (a handful to a few dozen + expressive nodes, with ``k_neighbors`` set to ``K``). At ``q = 12`` and + ``k = 8`` that is ~90 MB for 20k atoms; a large ``K`` *and* a large ``k`` together + is what to avoid. + + Parameters + ---------- + mode_set : str, optional + Key of :data:`MODE_SETS`. Default ``"rigid"``, i.e. TLS. + epsilon : float, optional + Floor on the Cholesky diagonal of ``Sigma``, which bounds its smallest + eigenvalue. Also the value the non-translation modes start at, so a freshly + fitted field begins as the equivalent constant-U field. + """ + + out_width = 6 + + def __init__(self, mode_set: str = "rigid", epsilon: float = 1e-3): + if mode_set not in MODE_SETS: + raise ValueError( + f"Unknown mode set {mode_set!r}. Available: {sorted(MODE_SETS)}." + ) + self.mode_set = mode_set + self.epsilon = float(epsilon) + self._gradient_names = MODE_SETS[mode_set] + self.q = 3 + sum(len(_GENERATORS[n]) for n in self._gradient_names) + self.width = chol_param_count(self.q) + self._generator_cache = {} + + def __repr__(self): + return ( + f"ModeCovariancePayload({self.mode_set!r}, q={self.q}, " + f"params={self.width})" + ) + + # ------------------------------------------------------------------ + # Modes. + # ------------------------------------------------------------------ + + def _generators(self, dtype, device): + """``(q - 3, 3, 3)`` gradient generators, cached per dtype and device.""" + key = (dtype, str(device)) + G = self._generator_cache.get(key) + if G is None: + rows = [m for n in self._gradient_names for m in _GENERATORS[n]] + G = torch.tensor(rows, dtype=dtype, device=device).reshape(-1, 3, 3) + self._generator_cache[key] = G + return G + + @staticmethod + def _length_scale(node_pos): + """Median nearest-neighbour node distance, as a plain float. + + Deliberately outside the graph. A median's derivative is supported on whichever + single node pair sits at the median, which is an artifact of the layout rather + than a direction worth following, and leaving it connected would also let the + optimiser rescale its own modes by spreading the nodes apart. Node position + keeps its real gradient through the displacement ``r``. + """ + with torch.no_grad(): + if node_pos.shape[0] < 2: + return 1.0 + d = torch.cdist(node_pos, node_pos) + d.fill_diagonal_(float("inf")) + return max(float(d.min(dim=1).values.median()), 1e-3) + + def modes(self, r): + """``(..., 3, q)`` displacement modes at (already scaled) displacement ``r``.""" + eye = torch.eye(3, dtype=r.dtype, device=r.device).expand( + *r.shape[:-1], 3, 3 + ) + if not self._gradient_names: + return eye + G = self._generators(r.dtype, r.device) + grad = torch.einsum("sij,...j->...is", G, r) + return torch.cat([eye, grad], dim=-1) + + def sigma(self, payload): + """``(K, q, q)`` mode covariance of each node, positive-definite.""" + L = raw_to_cholesky(payload, self.q, self.epsilon) + return L @ L.transpose(-1, -2) + + # ------------------------------------------------------------------ + # NodePayload interface. + # ------------------------------------------------------------------ + + def contributions(self, payload, xyz, node_pos, neighbor_list): + r = (xyz.unsqueeze(1) - node_pos[neighbor_list]) / self._length_scale(node_pos) + Psi = self.modes(r) # (N, k, 3, q) + L = raw_to_cholesky(payload, self.q, self.epsilon) # (K, q, q) + A = Psi @ L[neighbor_list] # (N, k, 3, q) + U = A @ A.transpose(-1, -2) # (N, k, 3, 3) + return torch.stack( + [U[..., 0, 0], U[..., 1, 1], U[..., 2, 2], + U[..., 0, 1], U[..., 0, 2], U[..., 1, 2]], + dim=-1, + ) + + def fit(self, target, w_dense, epsilon, xyz, node_pos, neighbor_list): + """Seed the translation block from the constant-U solve, floor the rest. + + The full joint solve is available in principle -- ``U(r)`` is linear in + ``Sigma``, so it is one least-squares problem in ``K * q(q+1)/2`` unknowns -- but + it is not what is wanted here. Seeding only the translations makes the field + start as the equivalent constant-U field, which is a state whose R-factor is + already known, so entering this parametrisation cannot make the model worse and + refinement can only move away from a sane point. It also sidesteps the joint + solve's normal equations, which stop being cheap well before ``K`` does. + """ + if target.ndim == 1: # a B target: lift to the equivalent isotropic U + u_iso = target / (8.0 * math.pi**2) + zero = torch.zeros_like(u_iso) + target = torch.stack([u_iso, u_iso, u_iso, zero, zero, zero], dim=1) + u6 = _ridged_solve(w_dense, target) # (K, 6) + sigma = u6.new_zeros(u6.shape[0], self.q, self.q) + sigma[:, :3, :3] = u6_to_matrix(u6) + # psd_to_raw clamps every eigenvalue to epsilon^2, so the gradient modes come + # out at the floor rather than at zero -- non-degenerate, and negligible against + # a real U. + return psd_to_raw(sigma, self.epsilon) + + def log_magnitude(self, payload): + """Log ``B_eq`` of ``U(0)``: the translation block, which is the node's own ADP. + + Evaluated at the node rather than averaged over its region, so the number means + the same thing for every mode set and a magnitude restraint can price nodes + without knowing which one is in use. + """ + T = self.sigma(payload)[:, :3, :3] + b_eq = (8.0 * math.pi**2 / 3.0) * (T[:, 0, 0] + T[:, 1, 1] + T[:, 2, 2]) + return torch.log(b_eq.clamp(min=1e-6)) + + +def _ridged_solve(w_dense, target): + """Least squares ``min ||W x - target||`` through the ridged normal equations. + + Non-finite target rows are dropped rather than carried in: deposited models have + NaN ADPs, and one NaN row propagates through the normal equations and takes every + node with it. Those atoms still receive a fitted value on output. + """ + finite = torch.isfinite(target).all(dim=-1) + if not bool(finite.all()): + w_dense, target = w_dense[finite], target[finite] + if target.shape[0] == 0: + return torch.zeros( + w_dense.shape[1], target.shape[-1], + dtype=w_dense.dtype, device=w_dense.device, + ) + gram = w_dense.T @ w_dense + ridge = 1e-6 * torch.diagonal(gram).mean().clamp(min=1e-30) + eye = torch.eye(gram.shape[0], dtype=gram.dtype, device=gram.device) + return torch.linalg.solve(gram + ridge * eye, w_dense.T @ target) + + +#: Stable integer code per payload, so a saved field can rebuild the one it had. +#: ``state_dict`` holds tensors only, and the payload is a plain object that never +#: reaches it, so without this a restore has to guess -- and guessing wrong is not a +#: clean failure: a mode payload restored as a constant-U one has the wrong storage +#: width and only shows up as a shape mismatch, or worse, silently different ADPs. +#: +#: **Append, never renumber.** A saved state dict holds the number. +PAYLOAD_CODES = { + "isotropic": 0, + "anisotropic": 1, + "modes:constant": 2, + "modes:rigid": 3, + "modes:rigid_dilation": 4, + "modes:affine": 5, +} + + +def payload_code(payload: "NodePayload") -> int: + """Code identifying ``payload`` well enough to rebuild it.""" + if isinstance(payload, ModeCovariancePayload): + key = f"modes:{payload.mode_set}" + elif isinstance(payload, AnisotropicPayload): + key = "anisotropic" + else: + key = "isotropic" + if key not in PAYLOAD_CODES: + raise ValueError( + f"Payload {key!r} has no code in PAYLOAD_CODES, so a field carrying it " + "cannot be saved and restored. Add one (appending, never renumbering)." + ) + return PAYLOAD_CODES[key] + + +def payload_from_code(code: int, epsilon: float = 1e-3) -> "NodePayload": + """Rebuild the payload a saved ``code`` names.""" + names = {v: k for k, v in PAYLOAD_CODES.items()} + key = names.get(int(code)) + if key is None: + raise ValueError( + f"Unknown payload code {code!r}. It was written by a newer TorchRef than " + f"this one, which knows {sorted(PAYLOAD_CODES)}." + ) + if key == "isotropic": + return IsotropicPayload() + if key == "anisotropic": + return AnisotropicPayload(epsilon=epsilon) + return ModeCovariancePayload(key.split(":", 1)[1], epsilon=epsilon) + + +class DisorderFieldTensor(MixedTensor): + """Per-atom ADPs from a small set of nodes, each atom a weighted mean of its k nearest. + + Storage is ``(K, 2)``: ``[log B, log sigma]`` per node. ``forward()`` returns + ``(n_atoms,)`` isotropic B, so this drops into the ``model.adp`` slot and every + consumer of ``adp()`` keeps working unchanged. + + The weight of node ``j`` at atom ``i`` is ``softmax_j(-d_ij^2 / 2 sigma_j^2)`` over + that atom's candidate list, so weights are non-negative and sum to one and B is a + convex combination of positive node values --- positive for free, with no clamping. + + Coordinates come from an accessor injected at construction rather than being passed + per call, which keeps ``forward()`` argument-free. That makes the inherited forward + cache incorrect on its own, since :class:`~torchref.utils.caching.CachedForwardMixin` + fingerprints parameters, buffers and call *arguments* --- and a borrowed accessor's + output is none of those. :meth:`_fingerprint_state` closes that by folding the + accessor's output into the key. + + Parameters + ---------- + initial_values : torch.Tensor, optional + ``(n_atoms,)`` isotropic B to fit the field to. Omit for an empty shell ready + for ``load_state_dict``. + xyz_fn : callable, optional + Returns the current ``(n_atoms, 3)`` coordinates. Typically ``model.xyz``. Held + by reference and deliberately invisible to ``state_dict``, device traversal and + ``copy``; re-attach with :meth:`set_xyz_fn` after a state-dict load. + n_nodes : int, optional + Number of nodes. Default 32. + k_neighbors : int, optional + Candidate nodes per atom. Default 12. Doubles as the skin margin that makes a + slightly stale candidate list harmless, so prefer generous over tight. + anchor_rows : tuple of torch.Tensor, optional + ``(flat atom indices, node index per entry)`` defining each node's anchor + neighbourhood. Omit to anchor every node at a single atom, which is what a model + without a topology gets. + node_values : torch.Tensor, optional + ``(K, 2)`` storage to adopt directly instead of fitting to ``initial_values``. + Used by :meth:`copy`. + refinable_mask : torch.Tensor, optional + Boolean mask. Interpreted in ATOM space unless ``mask_in_node_space``. + mask_in_node_space : bool, optional + Treat ``refinable_mask`` as already collapsed to ``(K,)``. Default False. + requires_grad : bool, optional + Whether node parameters carry gradients. Default True. + dtype, device : optional + Floating dtype and device. + name : str, optional + Wrapper name. Defaults to ``"adp"`` so ``Model`` consumers find it. + epsilon : float, optional + Floor on ``sigma`` and on fitted node B, in the same units as each. Default 1e-3. + """ + + def __init__( + self, + initial_values: Optional[torch.Tensor] = None, + xyz_fn: Optional[Callable[[], torch.Tensor]] = None, + n_nodes: int = 32, + k_neighbors: int = 12, + refine_positions: bool = False, + payload: Optional["NodePayload"] = None, + anchor_rows: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + node_values: Optional[torch.Tensor] = None, + refinable_mask: Optional[torch.Tensor] = None, + mask_in_node_space: bool = False, + requires_grad: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + name: Optional[str] = "adp", + epsilon: float = 1e-3, + ): + self.epsilon = epsilon + self._k_neighbors = int(k_neighbors) + self._refine_positions = bool(refine_positions) + # Storage columns are [payload | log sigma | offset]. Payload first keeps the + # isotropic layout unchanged, so an existing state dict still loads. + self._payload = payload if payload is not None else IsotropicPayload() + object.__setattr__(self, "_xyz_fn", _wrap_accessor(xyz_fn)) + + if initial_values is None and node_values is None: + device = normalize_device(device) + dtype = dtype if dtype is not None else get_float_dtype() + super().__init__(None, None, requires_grad, dtype, device, name) + self._full_shape = 0 + self.register_buffer("neighbor_list", None) + self.register_buffer("anchor_atom", None) + self.register_buffer("anchor_node", None) + # device=device so the empty path lands on the requested device, + # matching the populated path below (it used to omit it and land on CPU). + self.register_buffer( + "payload_code", + torch.tensor( + payload_code(self._payload), + dtype=get_int_dtype(), + device=device, + ), + ) + return + + if xyz_fn is None: + raise ValueError( + "DisorderFieldTensor needs xyz_fn to place its nodes; pass the model's " + "coordinate wrapper (e.g. model.xyz)." + ) + + xyz = xyz_fn().detach() + if dtype is None: + dtype = initial_values.dtype if initial_values is not None else xyz.dtype + if device is None: + device = xyz.device + xyz = xyz.to(dtype=dtype, device=device) + + n_atoms = int(xyz.shape[0]) + + if anchor_rows is None: + anchor_atom = farthest_point_anchors(xyz, n_nodes) + anchor_node = torch.arange( + anchor_atom.shape[0], dtype=torch.int64, device=device # dtype-ok: arange anchor indices; index requires int64 + ) + else: + anchor_atom, anchor_node = anchor_rows + anchor_atom = anchor_atom.to(device=device, dtype=torch.int64) # dtype-ok: anchor_atom indices cast; index requires int64 + anchor_node = anchor_node.to(device=device, dtype=torch.int64) # dtype-ok: anchor_node indices cast; index requires int64 + + n_k = int(anchor_node.max()) + 1 + node_pos = self._segment_mean(xyz, anchor_atom, anchor_node, n_k) + neighbor_list = build_neighbor_list(xyz, node_pos, self._k_neighbors) + + if node_values is None: + node_values = self._fit_nodes( + initial_values.to(dtype=dtype, device=device), + xyz, + node_pos, + neighbor_list, + ) + node_values = node_values.to(dtype=dtype, device=device) + + if refinable_mask is None: + node_mask = torch.ones(n_k, dtype=torch.bool, device=device) + elif mask_in_node_space: + node_mask = refinable_mask.to(device=device, dtype=torch.bool) + else: + node_mask = self._collapse_mask( + refinable_mask.to(device=device, dtype=torch.bool), neighbor_list, n_k + ) + + # ``register_buffer`` needs ``nn.Module.__init__`` to have run, which happens + # inside this call, so every buffer below is registered after it. + super().__init__( + initial_values=node_values, + refinable_mask=node_mask, + requires_grad=requires_grad, + dtype=dtype, + device=device, + name=name, + ) + self._full_shape = n_atoms + self.register_buffer("anchor_atom", anchor_atom) + self.register_buffer("anchor_node", anchor_node) + self.register_buffer("neighbor_list", neighbor_list) + # Which payload this field carries, so a restore rebuilds it rather than + # inferring it from the storage width. + self.register_buffer( + "payload_code", + torch.tensor( + payload_code(self._payload), dtype=get_int_dtype(), device=device + ), + ) + + # ------------------------------------------------------------------ + # Construction helpers. + # ------------------------------------------------------------------ + + @staticmethod + def _segment_mean( + xyz: torch.Tensor, + anchor_atom: torch.Tensor, + anchor_node: torch.Tensor, + n_nodes: int, + ) -> torch.Tensor: + """Mean coordinate of each node's anchor atoms, ``(K, 3)``. + + Differentiable in ``xyz``, which is what makes a node move with the model. + """ + acc = xyz.new_zeros(n_nodes, 3) + acc = acc.index_add(0, anchor_node, xyz[anchor_atom]) + counts = torch.zeros(n_nodes, dtype=xyz.dtype, device=xyz.device) + counts = counts.index_add( + 0, anchor_node, torch.ones_like(anchor_node, dtype=xyz.dtype) + ) + return acc / counts.clamp(min=1.0).unsqueeze(-1) + + @staticmethod + def _weights( + xyz: torch.Tensor, + node_pos: torch.Tensor, + neighbor_list: torch.Tensor, + log_sigma: torch.Tensor, + ) -> torch.Tensor: + """Softmax weights over each atom's candidate nodes, ``(n_atoms, k)``. + + Normalised across the candidates, so rows sum to one and no atom can end up + without support even when every node is far away. + """ + cand = node_pos[neighbor_list] + d2 = ((xyz.unsqueeze(1) - cand) ** 2).sum(-1) + sigma2 = torch.exp(2.0 * log_sigma)[neighbor_list] + return torch.softmax(-d2 / (2.0 * sigma2), dim=1) + + def _fit_nodes( + self, + target: torch.Tensor, + xyz: torch.Tensor, + node_pos: torch.Tensor, + neighbor_list: torch.Tensor, + ) -> torch.Tensor: + """Node storage whose field reproduces ``target`` as closely as it can. + + Kernel width is seeded at half the median nearest-neighbour node spacing, which + makes it a property of the node layout rather than a tuned constant. With the + weights fixed at that seed the payload is linear in the target, so the payload + half is a closed-form solve rather than an optimisation loop -- delegated, + because what "linear in the target" means differs between a scalar B and a + six-component U. + + Returns + ------- + torch.Tensor + ``(K, payload.width + 1 + 3*refine_positions)`` node storage. + """ + n_k = int(node_pos.shape[0]) + if n_k > 1: + dnode = torch.cdist(node_pos, node_pos) + dnode.fill_diagonal_(float("inf")) + spacing = float(dnode.min(dim=1).values.median()) + else: + spacing = float(xyz.std()) * 2.0 + sigma0 = max(spacing / 2.0, 10.0 * self.epsilon) + log_sigma = torch.full( + (n_k,), math.log(sigma0), dtype=xyz.dtype, device=xyz.device + ) + + w_sparse = self._weights(xyz, node_pos, neighbor_list, log_sigma) + w_dense = torch.zeros(xyz.shape[0], n_k, dtype=xyz.dtype, device=xyz.device) + w_dense.scatter_(1, neighbor_list, w_sparse) + + payload = self._payload.fit( + target, w_dense, self.epsilon, xyz, node_pos, neighbor_list + ) + + columns = [payload, log_sigma.unsqueeze(-1)] + if self._refine_positions: + columns.append(torch.zeros_like(node_pos)) + return torch.cat(columns, dim=1) + + @staticmethod + def _collapse_mask( + atom_mask: torch.Tensor, neighbor_list: torch.Tensor, n_nodes: int + ) -> torch.Tensor: + """Atom-space mask to node space: a node is refinable if any atom it serves is.""" + acc = torch.zeros(n_nodes, dtype=torch.bool, device=atom_mask.device) + served = neighbor_list[atom_mask] + if served.numel(): + acc[served.reshape(-1)] = True + return acc + + # ------------------------------------------------------------------ + # Public surface. + # ------------------------------------------------------------------ + + @property + def payload(self) -> "NodePayload": + """What each node carries and how it becomes a per-atom ADP.""" + return self._payload + + @property + def out_width(self) -> int: + """Components of the per-atom output: 1 for isotropic B, 6 for a U tensor.""" + return self._payload.out_width + + def _split(self, raw): + """Storage columns as ``(payload, log sigma, offset or None)``.""" + w = self._payload.width + offset = raw[:, w + 1 : w + 4] if self._refine_positions else None + return raw[:, :w], raw[:, w], offset + + def log_magnitude(self, raw=None) -> torch.Tensor: + """``(K,)`` log ADP magnitude per node, whatever the payload. + + Lets a magnitude restraint price node values without knowing the layout. + """ + if raw is None: + raw = super().forward() + return self._payload.log_magnitude(self._split(raw)[0]) + + @property + def shape(self): + """Shape of the FULL per-atom tensor, not the node storage.""" + return (self._full_shape,) + + @property + def node_shape(self): + """Shape of the node storage.""" + return tuple(self.fixed_values.shape) + + @property + def n_nodes(self) -> int: + """Number of nodes.""" + return int(self.fixed_values.shape[0]) + + def set_xyz_fn(self, xyz_fn: Callable[[], torch.Tensor]) -> None: + """Attach the coordinate accessor. + + Needed after ``load_state_dict``, which cannot carry a callable. + """ + object.__setattr__(self, "_xyz_fn", _wrap_accessor(xyz_fn)) + + @property + def refines_positions(self) -> bool: + """Whether node positions carry a refinable offset.""" + return self._refine_positions + + def node_positions(self, xyz: Optional[torch.Tensor] = None) -> torch.Tensor: + """Node positions, ``(K, 3)``: anchored centroid plus any refinable offset.""" + if xyz is None: + xyz = self._xyz_fn() + return self._node_positions_from(xyz, super().forward()) + + def _node_positions_from(self, xyz, raw): + """Anchor centroid, displaced by the refinable offset when there is one. + + Keeping the centroid as the base rather than storing absolute coordinates means + a node still travels with the model under xyz refinement; the offset only says + where it sits *relative* to the atoms it belongs to. + """ + base = self._segment_mean( + xyz, self.anchor_atom, self.anchor_node, raw.shape[0] + ) + offset = self._split(raw)[2] + return base if offset is None else base + offset + + def weights(self, xyz: Optional[torch.Tensor] = None) -> torch.Tensor: + """Per-atom weights over candidate nodes, ``(n_atoms, k)``. Rows sum to one.""" + if xyz is None: + xyz = self._xyz_fn() + raw = super().forward() + return self._weights( + xyz, + self._node_positions_from(xyz, raw), + self.neighbor_list, + self._split(raw)[1], + ) + + def node_load(self, xyz: Optional[torch.Tensor] = None) -> torch.Tensor: + """Total weight each node carries across all atoms, ``(K,)``. + + Not obtainable by summing :meth:`weights` over atoms: that returns + ``(n_atoms, k)`` over each atom's CANDIDATES, so a column is a candidate *slot* + shared by different nodes for different atoms, and summing it gives a length-k + vector with no meaning. The weights have to be scattered back into node space + through ``neighbor_list`` first, which is what this does. + + Load is what identifies a node that has stopped doing useful work: a node that + narrows onto a handful of atoms carries almost none. + """ + W = self.weights(xyz) + load = W.new_zeros(self.n_nodes) + return load.index_add(0, self.neighbor_list.reshape(-1), W.reshape(-1)) + + def smallest_candidate_weight(self, xyz: Optional[torch.Tensor] = None) -> float: + """Largest per-atom minimum candidate weight --- the list-adequacy invariant. + + Each atom sees only its ``k`` candidate nodes. While the weakest of those + candidates carries negligible weight, the list brackets the real neighbourhood + and a node drifting in or out of it cannot move any ADP appreciably. When this + rises, the list no longer brackets it and should be rebuilt with + :meth:`rebuild_neighbor_list` or a larger ``k``. + """ + with torch.no_grad(): + return float(self.weights(xyz).min(dim=1).values.max()) + + def rebuild_neighbor_list( + self, xyz: Optional[torch.Tensor] = None, k_neighbors: Optional[int] = None + ) -> None: + """Recompute which nodes each atom sees, at the current coordinates. + + The candidate list is the slowly-varying combinatorial half of the field and is + never refreshed implicitly: membership is piecewise constant in position, so a + rebuild is a discrete jump and belongs at a point the caller chooses. + """ + if xyz is None: + xyz = self._xyz_fn() + xyz = xyz.detach() + if k_neighbors is not None: + self._k_neighbors = int(k_neighbors) + self.neighbor_list = build_neighbor_list( + xyz, self.node_positions(xyz).detach(), self._k_neighbors + ) + self.reset_forward_cache() + + def evaluate(self, xyz: torch.Tensor, raw: torch.Tensor) -> torch.Tensor: + """Per-atom B from explicit coordinates and node storage, ``(n_atoms,)``. + + The field's arithmetic, with no accessor and no cache in the way, so it can be + differentiated and checked on its own. ``forward()`` is this plus the plumbing + that fetches both arguments. + + Parameters + ---------- + xyz : torch.Tensor + ``(n_atoms, 3)`` coordinates. + raw : torch.Tensor + ``(K, 2)`` node storage, ``[log b, log sigma]``. + """ + payload, log_sigma, _ = self._split(raw) + node_pos = self._node_positions_from(xyz, raw) + W = self._weights(xyz, node_pos, self.neighbor_list, log_sigma) + contrib = self._payload.contributions( + payload, xyz, node_pos, self.neighbor_list + ) + out = (W.unsqueeze(-1) * contrib).sum(dim=1) + # A scalar payload reports per-atom B as (N,), not (N, 1): that is the shape + # every consumer of ``adp()`` expects. + return out.squeeze(-1) if self._payload.out_width == 1 else out + + def forward(self) -> torch.Tensor: + """Per-atom isotropic B, ``(n_atoms,)``. + + A convex combination of positive node values, so strictly positive without a + clamp. Translation-invariant by construction: node positions are centroids of + atom coordinates, so a rigid shift of the model moves the nodes with it and + leaves every distance, and therefore every weight, unchanged. + """ + return self.evaluate(self._xyz_fn(), super().forward()) + + def node_values(self) -> torch.Tensor: + """The assembled node storage ``(K, 2)`` in raw ``[log b, log sigma]`` space.""" + return super().forward() + + def _fingerprint_state(self): + """Fold the accessor's coordinates into the forward-cache key. + + Without this the cache would be keyed on parameters and buffers alone and would + serve a per-atom B computed at coordinates that have since moved --- the + coordinates reach ``forward()`` through the accessor, not through an argument, + so the mixin cannot see them by itself. + """ + base = super()._fingerprint_state() + if self._xyz_fn is None: + return base + xyz = self._xyz_fn() + return base + ((xyz.data_ptr(), xyz._version),) + + def _set_values(self, key, value: torch.Tensor) -> None: + """Rejected: a node field cannot represent arbitrary per-atom values. + + Assigning per-atom ADPs would silently be a projection onto the field rather + than the write the caller asked for. Use :meth:`refit` to move the field toward + a per-atom target, or switch the model back to a per-atom representation. + """ + raise NotImplementedError( + "DisorderFieldTensor stores K nodes, not per-atom values, so per-atom " + "assignment is not representable. Use refit() to fit the field to a " + "per-atom target, or Model.set_adp_mode('isotropic') to leave field mode." + ) + + def refit(self, target_b: torch.Tensor) -> None: + """Re-fit the node values to a per-atom target, in place. + + The target is per-atom B for a scalar payload, or per-atom U6 for a tensor one; + an anisotropic payload also accepts a B target and lifts it to ``U_iso * I``. + + Replaces ``refinable_params``, so any optimizer state held for it is stale. + """ + xyz = self._xyz_fn().detach() + target_b = target_b.to(dtype=self.dtype, device=self.device) + node_values = self._fit_nodes( + target_b, xyz, self.node_positions(xyz).detach(), self.neighbor_list + ) + self.fixed_values = node_values.clone().detach() + refinable = node_values[self.refinable_mask].clone().detach() + self.refinable_params = nn.Parameter( + refinable, requires_grad=self.refinable_params.requires_grad + ) + self._build_index_cache() + self.reset_forward_cache() + + def update_refinable_mask( + self, new_mask: torch.Tensor, in_node_space: bool = False + ): + """Repartition refinable/fixed nodes, keeping the raw node values. + + Parameters + ---------- + new_mask : torch.Tensor + Boolean mask, ``(n_atoms,)`` in atom space or ``(K,)`` when + ``in_node_space``. An atom-space mask collapses with OR: a node is refinable + if any atom it serves is. + in_node_space : bool, optional + Whether ``new_mask`` is already in node space. Default False. + """ + new_mask = new_mask.to(device=self.device, dtype=torch.bool) + if in_node_space: + if new_mask.shape[0] != self.n_nodes: + raise ValueError( + f"Node-space mask must have shape ({self.n_nodes},), " + f"got {tuple(new_mask.shape)}" + ) + node_mask = new_mask + else: + if new_mask.shape[0] != self._full_shape: + raise ValueError( + f"Atom-space mask must have shape ({self._full_shape},), " + f"got {tuple(new_mask.shape)}" + ) + node_mask = self._collapse_mask( + new_mask, self.neighbor_list, self.n_nodes + ) + + current = self.fixed_values.clone() + if self.refinable_mask is not None and bool(self.refinable_mask.any()): + current[self.refinable_mask] = self.refinable_params.data + + self.fixed_values = current.clone().detach() + if bool(node_mask.any()): + self.refinable_params = nn.Parameter( + current[node_mask].clone().detach(), + requires_grad=self.refinable_params.requires_grad, + ) + else: + self.refinable_params = nn.Parameter( + torch.empty(0, 2, dtype=self.dtype, device=self.device), + requires_grad=False, + ) + self.refinable_mask = node_mask + self.fixed_mask = ~node_mask + self._build_index_cache() + self.reset_forward_cache() + + def copy(self) -> "DisorderFieldTensor": + """Independent copy sharing no parameter storage. + + The coordinate accessor is carried by REFERENCE, never deep-copied: it holds a + cache that can contain a graph-attached tensor, which ``deepcopy`` refuses to + walk. + """ + accessor = self._xyz_fn + if isinstance(accessor, ModuleReference): + accessor = accessor.module + new = DisorderFieldTensor( + initial_values=None, + xyz_fn=accessor, + k_neighbors=self._k_neighbors, + refine_positions=self._refine_positions, + payload=self._payload, + anchor_rows=(self.anchor_atom.clone(), self.anchor_node.clone()), + node_values=self.node_values().detach().clone(), + refinable_mask=self.refinable_mask.clone(), + mask_in_node_space=True, + requires_grad=self.refinable_params.requires_grad, + dtype=self.dtype, + device=self.device, + name=self._name, + epsilon=self.epsilon, + ) + new.neighbor_list = self.neighbor_list.clone() + new._full_shape = self._full_shape + return new + + def __repr__(self) -> str: + name_str = f"'{self.name}', " if self.name is not None else "" + return ( + f"DisorderFieldTensor({name_str}atoms={self._full_shape}, " + f"nodes={self.n_nodes}, k={self._k_neighbors}, " + f"payload={type(self._payload).__name__}, dtype={self.dtype}, " + f"device={self.device}, refinable={self.get_refinable_count()})" + ) diff --git a/torchref/model/mixed_model.py b/torchref/model/mixed_model.py index c23d3ca5..91ced770 100644 --- a/torchref/model/mixed_model.py +++ b/torchref/model/mixed_model.py @@ -178,10 +178,14 @@ def dtype_float(self): # Grid infrastructure (delegates to constituent models) # ========================================================================= + def real_space_grid(self) -> torch.Tensor: + """Build the Cartesian grid from the first model (shared cell → same grid).""" + return self.models[0].real_space_grid() + @property - def real_space_grid(self) -> Optional[torch.Tensor]: - """Real-space coordinate grid from first model (shared cell → same grid).""" - return self.models[0].real_space_grid + def grid_shape(self) -> Optional[tuple]: + """Map dimensions (nx, ny, nz) from the first model.""" + return self.models[0].grid_shape @property def fft(self): @@ -193,11 +197,6 @@ def gridsize(self) -> Optional[torch.Tensor]: """Grid dimensions (nx, ny, nz) from first model.""" return self.models[0].gridsize - @property - def map_symmetry(self): - """Map symmetry operator from first model.""" - return self.models[0].map_symmetry - @property def inv_fractional_matrix(self) -> torch.Tensor: """Inverse fractionalization (orthogonalization) matrix.""" @@ -226,24 +225,6 @@ def setup_grid(self, max_res=None, gridsize=None): for model in self.models: model.setup_grid(max_res=max_res, gridsize=gridsize) - def get_radius(self, min_radius_Angstrom: float = 4.0) -> int: - """ - Get the radius in voxels for density calculation. - - Delegates to first model (same grid → same voxel size). - - Parameters - ---------- - min_radius_Angstrom : float, optional - Minimum radius in Angstroms. Default is 4.0. - - Returns - ------- - int - Radius in voxels. - """ - return self.models[0].get_radius(min_radius_Angstrom) - def build_complete_map(self) -> torch.Tensor: """ Mixed electron density ``density_mixed = Σ w_i density_i``. diff --git a/torchref/model/model.py b/torchref/model/model.py index 12c6fd58..56d9a8db 100644 --- a/torchref/model/model.py +++ b/torchref/model/model.py @@ -14,13 +14,21 @@ from typing import Dict, Iterable, List, Optional, Tuple, Union +import warnings + import gemmi import torch import torch.nn as nn from torchref.base import math_torch -from torchref.config import get_float_dtype, normalize_device +from torchref.config import ( + canonical_device, + get_default_device, + get_float_dtype, + normalize_device, +) from torchref.io import cif, pdb +from torchref.model.context import ModelContext from torchref.model.parameter_wrappers import ( CholeskyMixedTensor, MixedTensor, @@ -69,10 +77,12 @@ class Model(DeviceMovementMixin, DebugMixin, nn.Module): """ Base model class for atomic structure models using PyTorch. - Owns the atomic data -- coordinates, atomic displacement parameters and + Owns the refinable atomic data -- coordinates, atomic displacement parameters and occupancies -- each held in a parameter wrapper that decides which atoms are - refinable. Build it empty (``Model()`` then ``load_pdb`` / ``load_cif`` / - ``load_state_dict``); ``if model:`` tests *initialization*, not existence. + refinable. Everything the structure was *loaded from* rather than refined lives on + :attr:`ctx`, a :class:`~torchref.model.context.ModelContext`. Build the model empty + (``Model()`` then ``load_pdb`` / ``load_cif`` / ``load_state_dict``); ``if model:`` + tests *initialization*, not existence. Parameters ---------- @@ -83,7 +93,15 @@ class Model(DeviceMovementMixin, DebugMixin, nn.Module): device : torch.device, optional Computation device. Defaults to the configured device.current. strip_H : bool, optional - Whether to strip hydrogen atoms when loading. Default is True. + Whether to strip hydrogen atoms when loading. Default False: hydrogens are kept + where the file has them. + add_hydrogens : bool, optional + Generate missing hydrogens on load when True. Default False; + ignored when ``strip_H`` is set. + cif_path : str or list of str, optional + Restraint dictionary file(s) for residues the monomer library does not know, or + whose library entry should be overridden. Given here rather than after loading so + that hydrogen generation on load reads the same dictionary the restraints will. Attributes ---------- @@ -96,15 +114,20 @@ class Model(DeviceMovementMixin, DebugMixin, nn.Module): positive-definite by construction. Isotropic atoms carry ``U = NaN``. occupancy : OccupancyTensor Atomic occupancies with values in [0, 1]. + ctx : ModelContext + The unit cell, space group, atom table, link records, provenance and + configuration. The fields not forwarded below are reached through it, e.g. + ``model.ctx.strip_H`` and ``model.ctx.initialized``. pdb : pandas.DataFrame - DataFrame containing atomic model data. Only refreshed from the tensors - by :meth:`update_pdb`. + Atom table, forwarded to :attr:`ctx`. Only refreshed from the tensors by + :meth:`update_pdb`. cell : Cell - Unit cell object with parameters [a, b, c, alpha, beta, gamma]. - spacegroup, symmetry : SpaceGroup - Space group object; ``symmetry`` is the same object under its old name. - initialized : bool - Whether the model has been initialized with data. + Unit cell, forwarded to :attr:`ctx`. + spacegroup : SpaceGroup + Space group, forwarded to :attr:`ctx`. + device : torch.device + Where the tensors live. Kept on the model rather than the context because the + device-movement machinery rewrites it in place. """ def __init__( @@ -112,7 +135,10 @@ def __init__( dtype_float=None, verbose=1, device=None, - strip_H: bool = True, + strip_H: bool = False, + add_hydrogens: bool = False, + cif_path: Optional[Union[str, List[str]]] = None, + hydrogens_in_xray: bool = True, ): """ Initialize an empty Model shell. @@ -129,7 +155,18 @@ def __init__( device : torch.device, optional Computation device. Defaults to the configured device.current. strip_H : bool, optional - Whether to strip hydrogen atoms when loading. Default is True. + Whether to strip hydrogen atoms when loading. Default False: hydrogens are + kept where the file has them. + add_hydrogens : bool, optional + Generate missing hydrogens on load when True. Default False; + ignored when ``strip_H`` is set. + cif_path : str or list of str, optional + Restraint dictionary file(s); see the class docstring. :meth:`set_restraints_cif` + can still change it after loading, but generation on load only sees the value + given here. + hydrogens_in_xray : bool, optional + Whether hydrogens contribute to the structure factors. Default True. They + stay in the restraints either way; see :attr:`hydrogens_in_xray`. """ super().__init__() # Resolve dtype/device at call time (not import time) so a runtime @@ -137,21 +174,21 @@ def __init__( if dtype_float is None: dtype_float = get_float_dtype() device = normalize_device(device) + # ``device`` and ``dtype_float`` stay here rather than moving into the context: + # they are live ``DeviceMixin`` trackers, rewritten in place by the traversal on + # whichever object owns the tensors. self.dtype_float = dtype_float - self.verbose = verbose self.device = device - self.strip_H = strip_H - self._exclude_H_from_sf = False - - # State tracking - self.initialized = False - self.altloc_pairs = [] - # These will be set during load() or load_state_dict() - self.pdb = None - self.links = None - self._cell: Optional[Cell] = None - self._spacegroup: Optional[SpaceGroup] = None + # Everything the model is loaded from and sits in, as opposed to what is + # refined. Populated by load() / create_from_state_dict(). + self.ctx = ModelContext( + verbose=verbose, + strip_H=strip_H, + add_hydrogens=add_hydrogens, + cif_path=cif_path, + hydrogens_in_xray=hydrogens_in_xray, + ) # Submodules (created during load or load_state_dict) self.xyz = None @@ -164,7 +201,6 @@ def __init__( # Restraints (built lazily on first access) self._restraints = None - self._cif_path = None def __bool__(self): """Return the initialization status when used in boolean context. @@ -173,86 +209,181 @@ def __bool__(self): an uninitialized (but non-``None``) model is falsy. Use ``if model is not None`` when you mean an existence check. """ - return self.initialized + return self.ctx.initialized + + @property + def hydrogens_in_xray(self) -> bool: + """Whether hydrogens enter ``get_iso()`` / ``get_aniso()`` and so Fcalc. + + Restraints and the non-bonded term see the hydrogens either way, and the + bulk-solvent mask never does. Default True. Changing it re-keys the + iso/aniso partition on the next access; no cache needs clearing. + """ + return self.ctx.hydrogens_in_xray + + @hydrogens_in_xray.setter + def hydrogens_in_xray(self, value: bool): + self.ctx.hydrogens_in_xray = bool(value) @property def exclude_H_from_sf(self) -> bool: - """Drop H from ``get_iso()`` / ``get_aniso()`` (so from Fcalc) while - keeping them in the geometry and VDW restraints. Default False. + """Inverse of :attr:`hydrogens_in_xray`. + + .. deprecated:: + Use ``hydrogens_in_xray`` instead. """ - return self._exclude_H_from_sf + warnings.warn( + "exclude_H_from_sf is deprecated; use hydrogens_in_xray", + DeprecationWarning, + stacklevel=2, + ) + return not self.ctx.hydrogens_in_xray @exclude_H_from_sf.setter def exclude_H_from_sf(self, value: bool): - self._exclude_H_from_sf = bool(value) - # The cached iso/aniso indices encode the H choice, so rebuild them. - if self.initialized and self.pdb is not None: - self._rebuild_sf_indices() - - def _rebuild_sf_indices(self): - """Rebuild cached iso/aniso index arrays from aniso_flag and H mask.""" - iso_mask = ~self.aniso_flag - aniso_mask = self.aniso_flag - - if self._exclude_H_from_sf and self.pdb is not None: - if not hasattr(self, "_heavy_atom_mask"): - h_mask = torch.tensor( - (self.pdb["element"].str.strip() != "H").values, + warnings.warn( + "exclude_H_from_sf is deprecated; use hydrogens_in_xray", + DeprecationWarning, + stacklevel=2, + ) + self.ctx.hydrogens_in_xray = not bool(value) + + def _sf_atom_mask(self) -> Optional[torch.Tensor]: + """Atoms that enter Fcalc, or None when every atom does. + + Boolean ``(N,)`` over the atom table. Built lazily as the ``_heavy_atom_mask`` + buffer, which is dropped with the other per-atom caches when the atom set + changes, so it never outlives the table it was built for. + """ + if self.ctx.hydrogens_in_xray or self.pdb is None: + return None + if getattr(self, "_heavy_atom_mask", None) is None: + self.register_buffer( + "_heavy_atom_mask", + torch.tensor( + (self.pdb["element"].str.strip().str.upper() != "H").values, dtype=torch.bool, device=self.device, - ) - self.register_buffer("_heavy_atom_mask", h_mask) - iso_mask = iso_mask & self._heavy_atom_mask - aniso_mask = aniso_mask & self._heavy_atom_mask - - self._iso_indices = iso_mask.nonzero(as_tuple=True)[0] - self._aniso_indices = aniso_mask.nonzero(as_tuple=True)[0] + ), + ) + return self._heavy_atom_mask + + # -- iso/aniso partition, derived on access --------------------------- + # + # These four are a cache over ``aniso_flag`` and the H choice, and caches in + # this codebase are recomputed on access rather than copied. Keying them on a + # fingerprint of their inputs means there is no invalidation to remember: + # every mutation that could change them changes the fingerprint, including an + # in-place edit of ``aniso_flag`` (``_version`` moves) and a whole-tensor + # replacement (``data_ptr`` moves). + # + # Eager rebuilding is what made ``copy()`` fragile. A fresh copy is + # constructed, then has its context replaced and its buffers cloned, so + # indices built during construction describe the wrong ``aniso_flag`` -- + # and they do not raise, they silently gather the wrong atoms, with + # ``_aniso_is_empty`` able to skip anisotropic atoms outright. + + def _sf_partition(self): + """``(iso_idx, aniso_idx, iso_covers_all, aniso_is_empty)``, cached.""" + flag = self.aniso_flag + heavy = getattr(self, "_heavy_atom_mask", None) + fp = ( + (flag.data_ptr(), flag._version) if flag is not None else None, + bool(self.ctx.hydrogens_in_xray), + None if heavy is None else (heavy.data_ptr(), heavy._version), + 0 if self.pdb is None else len(self.pdb), + ) + cached = getattr(self, "_sf_partition_cache", None) + if cached is not None and self._sf_partition_fp == fp: + return cached + + iso_mask = ~flag + aniso_mask = flag + sf_atoms = self._sf_atom_mask() + if sf_atoms is not None: + # The mask is part of the key, so re-key after building it. + fp = (fp[0], fp[1], (sf_atoms.data_ptr(), sf_atoms._version), fp[3]) + iso_mask = iso_mask & sf_atoms + aniso_mask = aniso_mask & sf_atoms + + iso_idx = iso_mask.nonzero(as_tuple=True)[0] + aniso_idx = aniso_mask.nonzero(as_tuple=True)[0] # Fast-path flags: an everywhere-True iso_mask lets ``get_iso()`` skip the # gather (and its ``index_put_`` backward) entirely, and - # ``_aniso_is_empty`` lets ``get_aniso()`` short-circuit — the typical + # ``_aniso_is_empty`` lets ``get_aniso()`` short-circuit -- the typical # macromolecular case. - self._iso_covers_all = bool(iso_mask.all().item()) - self._aniso_is_empty = int(self._aniso_indices.numel()) == 0 + out = (iso_idx, aniso_idx, + bool(iso_mask.all().item()), int(aniso_idx.numel()) == 0) + self._sf_partition_cache = out + self._sf_partition_fp = fp + return out + + @property + def _iso_indices(self) -> torch.Tensor: + return self._sf_partition()[0] + + @property + def _aniso_indices(self) -> torch.Tensor: + return self._sf_partition()[1] + + @property + def _iso_covers_all(self) -> bool: + return self._sf_partition()[2] + + @property + def _aniso_is_empty(self) -> bool: + return self._sf_partition()[3] + # ========================================================================= # Cell, SpaceGroup, and Symmetry properties # ========================================================================= + @property + def pdb(self) -> Optional["pandas.DataFrame"]: + """Atom table. Only refreshed from the tensors by :meth:`update_pdb`.""" + return self.ctx.pdb + + @pdb.setter + def pdb(self, value): + self.ctx.pdb = value + @property def cell(self) -> Optional[Cell]: """Unit cell object with parameters [a, b, c, alpha, beta, gamma].""" - return self._cell + return self.ctx.cell @cell.setter def cell(self, value: Cell): """Set the unit cell.""" - self._cell = value + self.ctx.cell = value @property - def spacegroup(self) -> Optional[gemmi.SpaceGroup]: + def spacegroup(self) -> Optional[SpaceGroup]: """Space group object, or None if not set.""" - return self._spacegroup + return self.ctx.spacegroup @spacegroup.setter def spacegroup(self, value): - """Set the space group from a SpaceGroup, gemmi object, name or number.""" - if value is not None: + """Set the space group from a SpaceGroup, gemmi object, name or number. + + The model owns its space group: an incoming ``SpaceGroup`` is copied rather + than shared, because ``.to()`` moves in place and would otherwise relocate + the caller's object. The copy lands on the model's device and float dtype. + """ + if value is None: + self.ctx.spacegroup = None + elif isinstance(value, SpaceGroup): + self.ctx.spacegroup = value.copy().to( + device=self.device, dtype=self.dtype_float + ) + else: # ``device=self.device``: SpaceGroup falls back to the global # default otherwise, so setting a spacegroup on a CPU-pinned Model # would silently plant accelerator-resident matrices on it. - self._spacegroup = SpaceGroup(value, device=self.device) - else: - self._spacegroup = None - - @property - def symmetry(self) -> Optional[SpaceGroup]: - """The same object as :attr:`spacegroup`, under its older name.""" - return self._spacegroup - - @symmetry.setter - def symmetry(self, value: Optional[SpaceGroup]): - """Set the space group object directly (no coercion, unlike ``spacegroup``).""" - self._spacegroup = value + self.ctx.spacegroup = SpaceGroup( + value, dtype=self.dtype_float, device=self.device + ) # ========================================================================= # Crystallographic matrix properties (delegated to Cell) @@ -290,7 +421,7 @@ def _build_z_tensor(self) -> torch.Tensor: if hasattr(self, "_Z") and self._Z is not None: return self._Z - if not self.initialized or self.pdb is None: + if not self.ctx.initialized or self.pdb is None: raise RuntimeError( "Cannot build Z tensor: model not initialized. " "Load data first with load_pdb() or load_cif()." @@ -304,7 +435,7 @@ def _build_z_tensor(self) -> torch.Tensor: for elem in self.pdb["element"] ] self.register_buffer( - "_Z", torch.tensor(z_values, dtype=torch.int32, device=self.device) + "_Z", torch.tensor(z_values, dtype=torch.int32, device=self.device) # dtype-ok: atomic-number Z categorical codes buffer; fixed int32 lookup keys ) return self._Z @@ -320,13 +451,13 @@ def _build_parametrization(self): if self._parametrization is not None: return self._parametrization - if not self.initialized or self.pdb is None: + if not self.ctx.initialized or self.pdb is None: raise RuntimeError( "Cannot build parametrization: model not initialized. " "Load data first with load_pdb() or load_cif()." ) - if self.verbose > 1: + if self.ctx.verbose > 1: print("Building ITC92 parametrization via table lookup...") from torchref.base.scattering.scattering_table import get_scattering_params_by_z @@ -351,11 +482,11 @@ def _build_parametrization(self): B[idx : idx + 1], ) - if self.verbose > 0: + if self.ctx.verbose > 0: print( f"Parametrization built for {len(self._parametrization)} unique atom types" ) - if self.verbose > 1: + if self.ctx.verbose > 1: print("Elements with parametrization:", list(self._parametrization.keys())) return self._parametrization @@ -378,9 +509,8 @@ def get_scattering_params_iso(self): Notes ----- - ``n_iso_atoms`` honors ``exclude_H_from_sf``: when H exclusion is - active the isotropic count is the H-excluded count (mirroring - :meth:`get_iso`). + ``n_iso_atoms`` honors ``hydrogens_in_xray``: when hydrogens are excluded + the isotropic count is the heavy-atom count (mirroring :meth:`get_iso`). """ self._build_parametrization() idx = self._iso_indices @@ -399,9 +529,8 @@ def get_scattering_params_aniso(self): Notes ----- - ``n_aniso_atoms`` honors ``exclude_H_from_sf``: when H exclusion is - active the anisotropic count is the H-excluded count (mirroring - :meth:`get_aniso`). + ``n_aniso_atoms`` honors ``hydrogens_in_xray``: when hydrogens are excluded + the anisotropic count is the heavy-atom count (mirroring :meth:`get_aniso`). """ self._build_parametrization() idx = self._aniso_indices @@ -425,39 +554,39 @@ def set_restraints_cif(self, cif_path): Model Self, for method chaining. """ - self._cif_path = cif_path + self.ctx.cif_path = cif_path # Reset restraints so they will be rebuilt on next access self._restraints = None return self def _build_restraints(self): - """Build and cache ``RestraintsNew`` over this model's DataFrame, wiring in + """Build and cache ``Restraints`` over this model's DataFrame, wiring in the live ``xyz`` / ``adp`` / ``vdw_radii`` callables. """ if self._restraints is not None: return self._restraints - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError( "Cannot build restraints: model not initialized. " "Load data first with load_pdb() or load_cif()." ) - from torchref.restraints.restraints import RestraintsNew + from torchref.topology.restraints import Restraints - if self.verbose > 0: + if self.ctx.verbose > 0: print("Building restraints...") - self._restraints = RestraintsNew( + self._restraints = Restraints( pdb=self.pdb, - cif_path=self._cif_path, + cif_path=self.ctx.cif_path, xyz_fn=self.xyz, adp_fn=self.adp, vdw_radii_fn=self.get_vdw_radii, - cell=self._cell, - spacegroup=self._spacegroup, - links=self.links, - verbose=self.verbose, + cell=self.ctx.cell, + spacegroup=self.ctx.spacegroup, + links=self.ctx.links, + verbose=self.ctx.verbose, ) return self._restraints @@ -512,7 +641,32 @@ def torsion_deviations_with_sigmas(self): """ return self.restraints.torsion_deviations_with_sigmas(self.xyz()) - def load(self, reader): + #: Per-atom buffers built lazily on first use and cached. Each is sized to the atom + #: table, so all of them go stale the moment the atom set changes. + _ATOM_DERIVED_BUFFERS = ( + "vdw_radii", + "_Z", + "_A", + "_B", + "_heavy_atom_mask", + ) + + def _invalidate_atom_derived_caches(self) -> None: + """Drop the lazily-cached per-atom buffers. + + Each is guarded by ``hasattr`` and returned as-is once built, so a load that + changes the atom count would otherwise hand back a buffer sized for the previous + one. That surfaced when hydrogen generation began extending the table in place: + the van der Waals radii stayed at the heavy-atom count while the pair list + indexed the full set, and the non-bonded build raised ``IndexError``. Rebuilding + a new model each time had hidden it. + """ + for name in self._ATOM_DERIVED_BUFFERS: + if hasattr(self, name): + delattr(self, name) + self._parametrization = None + + def load(self, reader, add_hydrogens: bool = None): """ Populate the model from a reader callable. @@ -528,7 +682,11 @@ def load(self, reader): ---------- reader : callable Zero-argument callable returning ``(pdb_df, cell, spacegroup)``. An - optional ``.links`` attribute on it is stored on ``self.links``. + optional ``.links`` attribute on it is stored on ``self.ctx.links``. + add_hydrogens : bool, optional + Whether to top up missing hydrogens once the model is built. Defaults to the + context's setting, and is forced off for the re-entry that + :meth:`_add_missing_hydrogens` makes, so generation happens once per load. Returns ------- @@ -541,15 +699,25 @@ def load(self, reader): ``aniso_flag`` buffer, the four wrappers, the default masks, the altloc registration and ``initialized = True``. """ + if add_hydrogens is None: + add_hydrogens = self.ctx.add_hydrogens and not self.ctx.strip_H + self._invalidate_atom_derived_caches() self.pdb, cell, spacegroup = reader() - self.links = getattr(reader, "links", None) + self.ctx.links = getattr(reader, "links", None) self.pdb = ( self.pdb.loc[self.pdb["element"] != "H"].reset_index(drop=True) - if self.strip_H + if self.ctx.strip_H else self.pdb ) self.pdb.dropna(subset=["x", "y", "z", "tempfactor", "occupancy"], inplace=True) + # Reindex before deriving the ``index`` column: every consumer uses it to + # address length-N per-atom tensors positionally (see + # ``_create_occupancy_groups``), so a gapped index from the drop above sends + # them past the end. Only the strip_H branch reset, so a model losing rows to + # the dropna instead -- an atom with no coordinates or no B -- raised + # IndexError at load. Hit on roughly one PDB-REDO entry in six. + self.pdb.reset_index(drop=True, inplace=True) self.pdb["index"] = self.pdb.index.to_numpy(dtype=int) self.cell = Cell(cell, dtype=self.dtype_float, device=self.device) @@ -563,8 +731,6 @@ def load(self, reader): self.pdb["anisou_flag"].values, dtype=torch.bool, device=self.device ), ) - # Pre-compute integer indices for SF calculation (respects exclude_H_from_sf) - self._rebuild_sf_indices() self.xyz = MixedTensor( torch.tensor(self.pdb[["x", "y", "z"]].values, dtype=self.dtype_float), @@ -604,9 +770,67 @@ def load(self, reader): self.set_default_masks() self.register_alternative_conformations() - self.initialized = True + self.ctx.initialized = True + + if add_hydrogens: + self._add_missing_hydrogens() return self + def _add_missing_hydrogens(self) -> None: + """Top up the hydrogens the atom table is missing, in place. + + Per parent, not per file: a structure deposited with some hydrogens gets the + rest, because the plan only ever proposes a hydrogen the template names and the + model does not have. 1AK5 arrives with 675 of roughly 2500, and a + does-it-have-any test would have left it there. + + Re-enters :meth:`load` on the augmented atom table, which rebuilds the parameter + wrappers and per-atom buffers at the new size. The re-entry is told not to + consider hydrogens again, so this runs once per load rather than recursing to a + fixed point. + + Costs a restraint build that is then discarded, because the plan needs the + topology and the topology is built over the atoms as loaded. Loading invokes + this only when ``add_hydrogens=True`` is requested. + """ + from torchref.topology.hydrogens import ( + augment_atom_table, + optimise_free_torsions, + plan_hydrogens, + ) + + restraints = self.restraints + xyz = self.xyz().detach() + plan = plan_hydrogens( + restraints.topology, restraints.cif_dict, xyz, verbose=self.ctx.verbose + ) + if self.ctx.verbose > 0 and restraints.missing_residues: + print( + "No restraint dictionary for " + f"{sorted(restraints.missing_residues)}: not hydrogenated. Pass one " + "with cif_path / --cif." + ) + if plan.n_hydrogens == 0: + return + optimise_free_torsions(plan, restraints.topology, xyz) + augmented = augment_atom_table(self.pdb, plan, restraints.topology) + + if self.ctx.verbose > 0: + print(f"Generated {plan.n_hydrogens} hydrogens") + + # The topology and every per-atom tensor are sized for the old atom set. + self._restraints = None + cell, spacegroup = self.cell, self.spacegroup + links = self.ctx.links + + def reader(): + return augmented, cell.data.cpu().numpy(), spacegroup + + # Carried explicitly: ``load`` reads links off the reader, so a bare callable + # would drop the LINK records the first read resolved. + reader.links = links + self.load(reader, add_hydrogens=False) + def load_pdb(self, file): """ Load atomic model from PDB file. @@ -621,8 +845,8 @@ def load_pdb(self, file): Model Self, for method chaining. """ - self._input_file = str(file) - reader = pdb.PDBReader(verbose=self.verbose).read(file) + self.ctx.input_file = str(file) + reader = pdb.PDBReader(verbose=self.ctx.verbose).read(file) return self.load(reader) def load_cif(self, file): @@ -639,8 +863,8 @@ def load_cif(self, file): Model Self, for method chaining. """ - self._input_file = str(file) - if self.verbose > 0: + self.ctx.input_file = str(file) + if self.ctx.verbose > 0: print(f"Loading CIF file: {file}") # Read CIF file @@ -722,7 +946,7 @@ def _create_occupancy_groups(self, pdb_df, initial_occ): altloc_groups = [] refinable_mask = torch.zeros(n_atoms, dtype=torch.bool) - sharing_groups_tensor = torch.arange(n_atoms, dtype=torch.long) + sharing_groups_tensor = torch.arange(n_atoms, dtype=torch.long) # dtype-ok: arange atom indices (sharing groups); index requires long collapsed_idx = 0 # First pass: altlocs. ALL atoms of one conformation must share a collapsed @@ -791,14 +1015,14 @@ def _create_occupancy_groups(self, pdb_df, initial_occ): # Compact to contiguous indices 0..n_collapsed-1. unique_indices = torch.unique(sharing_groups_tensor, sorted=True) - index_map = torch.zeros(n_atoms, dtype=torch.long) + index_map = torch.zeros(n_atoms, dtype=torch.long) # dtype-ok: index_map atom-index remap; indexing requires long for new_idx, old_idx in enumerate(unique_indices): mask = sharing_groups_tensor == old_idx sharing_groups_tensor[mask] = new_idx n_collapsed = len(unique_indices) - if self.verbose > 1: + if self.ctx.verbose > 1: n_groups = n_collapsed n_independent = n_atoms - n_collapsed n_refinable = refinable_mask.sum().item() @@ -817,10 +1041,14 @@ def update_pdb(self): """ Write the current refinable parameters back into ``self.pdb``. - Copies the live values of ``xyz`` (x/y/z), ``u`` (u11..u23), ``adp`` - (tempfactor), and ``occupancy`` from the parameter wrappers into the - corresponding columns of the ``self.pdb`` DataFrame. Called by every - writer and by ``hydrogenate`` / ``generate_hydrogens`` before output. + Copies the live values of ``xyz`` (x/y/z), ``u`` (u11..u23) and + ``occupancy`` from the parameter wrappers into the corresponding columns of + the ``self.pdb`` DataFrame. Called by every writer and by ``hydrogenate`` + before output. + + ``tempfactor`` is the equivalent isotropic B whenever any atom is + anisotropic, so the column agrees with the ANISOU records written beside it; + with no anisotropic atoms it is the isotropic wrapper directly. Returns ------- @@ -837,7 +1065,19 @@ def update_pdb(self): self.pdb.loc[:, ["u11", "u22", "u33", "u12", "u13", "u23"]] = ( self.u().cpu().detach().numpy() ) - self.pdb.loc[:, "tempfactor"] = self.adp().cpu().detach().numpy() + # The B column must agree with the ANISOU records beside it: for an + # anisotropic atom the PDB convention is B_eq = (8 pi^2 / 3) tr(U), not + # whatever the isotropic wrapper still happens to hold. That wrapper stops + # being refined the moment an atom goes anisotropic, so writing it directly + # emits a stale B alongside a live U. + if getattr(self, "_aniso_is_empty", True): + self.pdb.loc[:, "tempfactor"] = self.adp().cpu().detach().numpy() + else: + from torchref.base.targets.adp import u6_b_eq + + self.pdb.loc[:, "tempfactor"] = ( + u6_b_eq(self.adp_u6()).cpu().detach().numpy() + ) self.pdb.loc[:, "occupancy"] = self.occupancy().cpu().detach().numpy() return self.pdb @@ -894,46 +1134,44 @@ def _after_device_apply( self, old_device, new_device, old_dtype, new_dtype, *, device_changed, dtype_changed, ): - """Regenerate the iso/aniso index tensors on the new device. + """Report the move. - The movement hook, not a ``to()`` override (``_apply`` bypasses ``to()``) - and not ``reset_cache()`` (which fires after every optimizer step). + This used to regenerate the iso/aniso index tensors, which a device move + would otherwise leave on the old device. It no longer has to: the + partition is derived on access and keyed on ``aniso_flag``'s identity, + and ``nn.Module._apply`` replaces the buffer rather than mutating it, so + the move invalidates the cache by itself. """ - if getattr(self, "aniso_flag", None) is not None: - self._rebuild_sf_indices() - if self.verbose > 0: + if self.ctx.verbose > 0: print(f"Model moved to device: {self.device}") def copy(self): """ Create a deep copy of the Model. - Creates a complete independent copy including all registered buffers, - module parameters, PDB DataFrame, and spacegroup information. + Independent in every part: the context is copied via + :meth:`~torchref.model.context.ModelContext.copy`, buffers are cloned and each + parameter wrapper is copied through its own ``copy`` so its parametrization + survives. Returns ------- Model A new, fully independent Model instance with copied data. """ - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError("Cannot copy an uninitialized Model. Load data first.") model_copy = Model( dtype_float=self.dtype_float, - verbose=self.verbose, + verbose=self.ctx.verbose, device=self.device, - strip_H=self.strip_H, + strip_H=self.ctx.strip_H, ) - model_copy.pdb = self.pdb.copy(deep=True) - - # Setter also sets symmetry; gemmi.SpaceGroup is immutable, so shared. - model_copy.spacegroup = self.spacegroup - model_copy.initialized = True - - if self.cell is not None: - model_copy.cell = self.cell.clone() + # One call carries the atom table, cell, space group, altloc groups and + # provenance, each deep-copied or cloned -- see ``ModelContext.copy``. + model_copy.ctx = self.ctx.copy() for buffer_name, buffer_value in self._buffers.items(): if buffer_value is not None: @@ -945,14 +1183,13 @@ def copy(self): if module is not None and hasattr(module, "copy"): setattr(model_copy, module_name, module.copy()) - if hasattr(self, "altloc_pairs") and self.altloc_pairs: - model_copy.altloc_pairs = [ - tuple(tensor.clone() for tensor in group) for group in self.altloc_pairs - ] - else: - model_copy.altloc_pairs = [] + # Anything that borrows the coordinates -- the ADP node field, the restraints' + # pair-list maintenance -- carries the reference through its own ``copy`` and + # still points at THIS model's ``xyz``. Re-point it, or the two models silently + # share coordinates and the copy is not independent. + model_copy._repoint_coordinate_accessors() - if self.verbose > 0: + if self.ctx.verbose > 0: print(f"✓ Model copied successfully ({len(model_copy.pdb)} atoms)") return model_copy @@ -992,7 +1229,7 @@ def get_iso(self): Return per-atom parameters for the isotropic atom subset. Selects atoms whose ADP is a single scalar ``b``: ``~self.aniso_flag``, - intersected with the heavy-atom mask when ``exclude_H_from_sf`` is on. + intersected with the heavy-atom mask when ``hydrogens_in_xray`` is off. Returns ------- @@ -1059,17 +1296,20 @@ def parameters_of_types(self, types: Iterable[str]) -> List[nn.Parameter]: Returns ------- list of nn.Parameter - The ``refinable_params`` leaf for each requested type, in the - order the types were given. + Leaves for each requested type, in the order the types were given. + Coordinate wrappers may expose additional torsion and rotation leaves. """ out: List[nn.Parameter] = [] for t in types: wrapper = getattr(self, t, None) if wrapper is None: continue - rp = getattr(wrapper, "refinable_params", None) - if rp is not None: - out.append(rp) + if hasattr(wrapper, "optimization_parameters"): + out.extend(wrapper.optimization_parameters()) + else: + rp = getattr(wrapper, "refinable_params", None) + if rp is not None: + out.append(rp) return out def freeze(self, target: str): @@ -1131,7 +1371,16 @@ def unfreeze(self, target: str): self.occupancy_mask, in_compressed_space=False ) - def set_adp_mode(self, mode: str = "isotropic", aniso_selection: str = None): + def set_adp_mode( + self, + mode: str = "isotropic", + aniso_selection: str = None, + n_nodes: int = None, + k_neighbors: int = 12, + refine_node_positions: bool = True, + mode_set: str = None, + init: str = "fit", + ): """Set the atomic displacement parameter (ADP) parametrization. Repartitions atoms between isotropic (a single B in ``adp``) and @@ -1146,21 +1395,94 @@ def set_adp_mode(self, mode: str = "isotropic", aniso_selection: str = None): Parameters ---------- - mode : {"isotropic", "anisotropic"}, optional + mode : {"isotropic", "anisotropic", "field", "field_aniso", "preserve"}, optional ``"isotropic"`` (default) converts every atom, previously anisotropic ones to ``B_eq = (8 pi^2 / 3)(U11 + U22 + U33)``. ``"anisotropic"`` converts those matching ``aniso_selection``, expanding isotropic atoms - to ``U = (B / 8 pi^2) I``. + to ``U = (B / 8 pi^2) I``. ``"field"`` replaces the per-atom isotropic B + with a :class:`~torchref.model.disorder_field.DisorderFieldTensor`, whose + node values are least-squares fitted to the B it replaces, so the atom + count stops setting the ADP parameter count. ``"field_aniso"`` is the same + representation carrying a full U per node, which takes over ``u`` rather + than ``adp``. ``"preserve"`` is a no-op: the ADPs stay exactly as the file + supplied them, anisotropic where the file was anisotropic. aniso_selection : str, optional Phenix-style selection for ``mode="anisotropic"``, default ``"not resname HOH and not element H"``; ignored otherwise. + n_nodes : int, optional + Nodes for ``mode="field"``. Defaults to one per 25 atoms, floored at 4. + k_neighbors : int, optional + Candidate nodes per atom for ``mode="field"``. Default 12. + refine_node_positions : bool, optional + Give each node a refinable offset from its anchor centroid, at three extra + parameters per node. On by default: it is what lets the load-balancing + restraint move a node toward atoms instead of only widening its kernel. + init : {"fit", "flat"}, optional + What a field mode fits its nodes to: ``"fit"`` (default) the model's current + per-atom ADPs, ``"flat"`` a single level with their spatial structure + discarded. See :meth:`_install_disorder_field`. + mode_set : str, optional + For ``mode="field_aniso"``, a key of + :data:`~torchref.model.disorder_field.MODE_SETS` --- ``"rigid"`` is TLS, + ``"affine"`` adds shear and extension. The node then stores the covariance + of its displacement modes, so the U it gives an atom depends on where that + atom sits inside the node's region rather than being constant across it. + Default ``None`` keeps the constant-U payload. Notes ----- Run once at model setup, before scaling / restraints / targets. The isotropic result matches a freshly-loaded isotropic-only model. + + Leaving ``"field"`` needs no special case: the conversion reads ``adp()``, + which a field evaluates per atom, so the field materialises into a per-atom + wrapper on the way out. """ - if not getattr(self, "initialized", False) or self.pdb is None: + if not self.ctx.initialized or self.pdb is None: + return + if mode == "preserve": + # Leave the ADPs exactly as loaded. Constructing a Refinement otherwise + # reparametrises them before anything else runs, which silently discards a + # deposited model's anisotropy -- use this when the starting model's own + # ADPs are the thing being measured. + return + if mode in ("field", "field_aniso"): + aniso = mode == "field_aniso" + # Run the partition first either way: it owns every buffer keyed off the + # iso/aniso split, and it converts the stored values in the right direction + # (B -> U_iso*I entering anisotropic, U -> B_eq entering isotropic), so the + # field is fitted to a target that is already in its own representation. + if aniso: + # Every atom, unless the caller narrows it. A node field is not the + # per-atom parametrisation that "not water, not hydrogen" exists to + # ration -- its cost is set by node count, not atom count -- and a + # partial selection would leave half the ADPs coming from the field and + # half from the per-atom wrapper, which is not a representation anyone + # asked for. + if aniso_selection is None: + target_mask = torch.ones( + len(self.pdb), dtype=torch.bool, device=self.device + ) + else: + from torchref.utils.utils import create_selection_mask + + target_mask = torch.as_tensor( + create_selection_mask(aniso_selection, self.pdb), + dtype=torch.bool, + ).to(self.device) + else: + target_mask = torch.zeros( + len(self.pdb), dtype=torch.bool, device=self.device + ) + self._apply_adp_partition(target_mask) + self._install_disorder_field( + n_nodes=n_nodes, + k_neighbors=k_neighbors, + refine_node_positions=refine_node_positions, + anisotropic=aniso, + mode_set=mode_set, + init=init, + ) return if mode == "isotropic": aniso_mask = torch.zeros( @@ -1175,10 +1497,157 @@ def set_adp_mode(self, mode: str = "isotropic", aniso_selection: str = None): ).to(self.device) else: raise ValueError( - f"Unknown ADP mode: {mode!r}. Use 'isotropic' or 'anisotropic'." + f"Unknown ADP mode: {mode!r}. Use 'isotropic', 'anisotropic', " + "'field', 'field_aniso' or 'preserve'." ) self._apply_adp_partition(aniso_mask) + @property + def adp_is_field(self) -> bool: + """Whether either ADP slot holds a node field rather than a per-atom wrapper.""" + from torchref.model.disorder_field import DisorderFieldTensor + + return isinstance(self.adp, DisorderFieldTensor) or isinstance( + self.u, DisorderFieldTensor + ) + + @property + def adp_field(self): + """The node field driving the ADPs, or ``None`` if neither slot holds one.""" + from torchref.model.disorder_field import DisorderFieldTensor + + for wrapper in (self.u, self.adp): + if isinstance(wrapper, DisorderFieldTensor): + return wrapper + return None + + def _install_disorder_field( + self, + n_nodes: int = None, + k_neighbors: int = 12, + refine_node_positions: bool = False, + anisotropic: bool = False, + mode_set: str = None, + init: str = "fit", + ): + """Replace a per-atom ADP wrapper with a node field fitted to it. + + The field lands in the slot its payload feeds: an isotropic payload takes over + ``adp`` and leaves the model isotropic, an anisotropic one takes over ``u`` and + the model refines every selected atom anisotropically. Both expect the partition + to have run first, which :meth:`set_adp_mode` arranges. + + ``mode_set`` selects a displacement-mode payload in place of the constant-U one, + which is the difference between a node holding a single ADP and a node holding a + motion whose ADP varies across its region. + + ``init`` chooses what the field is fitted to: + + ``"fit"`` + The per-atom ADPs the model currently holds. Right when those mean something + --- a deposited or already-refined model --- because the field then starts + from a state whose R-factor is known. + ``"flat"`` + A single value, the median of those ADPs. Right when they do not mean + anything. 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 and that is not worth holding: measured on 2A25, the + fitted field starts 0.025 R-free WORSE than a flat one, before any + refinement. The level is kept because it is close to right and the scaler + owns it anyway; only the spatial structure is discarded. + """ + from torchref.model.disorder_field import ( + AnisotropicPayload, + DisorderFieldTensor, + IsotropicPayload, + ModeCovariancePayload, + density_anchor_rows, + ) + + if mode_set is not None and not anisotropic: + raise ValueError( + "mode_set describes an anisotropic displacement field and has no " + "isotropic form; use mode='field_aniso'." + ) + + with torch.no_grad(): + xyz = self.xyz().detach() + # The fit target is whatever the partition just produced: per-atom U6 for + # the anisotropic payload, per-atom B for the isotropic one. + target = ( + self.adp_u6().detach().clone() + if anisotropic + else self.adp().detach().clone() + ) + if init == "flat": + # Flatten through the equivalent isotropic B, and hand the payload a 1-D + # target: its ``fit`` lifts that to U_iso * I. Taking a median over all + # six U components instead would set the off-diagonals equal to the + # diagonals, giving eigenvalues (3L, 0, 0) -- singular, and NaN once the + # Cholesky encode takes log(diag - epsilon). + b = ( + (8.0 * math.pi**2 / 3.0) * target[:, :3].sum(dim=1) + if target.ndim == 2 + else target + ) + finite = torch.isfinite(b) + if not bool(finite.any()): + raise ValueError("cannot flatten an all-NaN ADP target") + level = b[finite].median() + target = torch.where(finite, level.expand_as(b), b) + elif init != "fit": + raise ValueError( + f"init={init!r}; expected 'fit' (use the model's own ADPs) or " + "'flat' (discard their spatial structure, keep the level)." + ) + B = target + if n_nodes is None: + n_nodes = max(4, int(round(len(self.pdb) / 25.0))) + + # Anchor on density clusters, not 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. + anchor_rows = density_anchor_rows(xyz, min(n_nodes, len(self.pdb))) + + if mode_set is not None: + payload = ModeCovariancePayload(mode_set) + elif anisotropic: + payload = AnisotropicPayload() + else: + payload = IsotropicPayload() + + field = DisorderFieldTensor( + initial_values=target.to(self.dtype_float), + xyz_fn=self.xyz, + n_nodes=n_nodes, + refine_positions=refine_node_positions, + payload=payload, + anchor_rows=anchor_rows, + k_neighbors=k_neighbors, + name="aniso_U" if anisotropic else "adp", + dtype=self.dtype_float, + device=self.device, + ) + if anisotropic: + self.u = field + # The mask is in atom space either way; the field collapses it onto nodes. + self.u.update_refinable_mask(self.u_mask) + else: + self.adp = field + self.adp.update_refinable_mask(self.adp_mask) + + if self.ctx.verbose > 0: + kind = mode_set if mode_set else ("aniso U" if anisotropic else "iso B") + was = len(self.pdb) * (6 if anisotropic else 1) + print( + f"ADP field ({kind}): {field.n_nodes} nodes, k={k_neighbors}, " + f"{int(field.get_refinable_count())} refinable nodes, " + f"{int(field.refinable_params.numel())} parameters " + f"(was {was} per-atom)" + ) + if hasattr(self, "reset_cache"): + self.reset_cache() + def _apply_adp_partition(self, aniso_mask: torch.Tensor): """Convert ADP storage to match a target anisotropic-atom mask. @@ -1220,7 +1689,6 @@ def _apply_adp_partition(self, aniso_mask: torch.Tensor): # Update the per-atom iso/aniso split and everything keyed off it. self.aniso_flag = aniso_mask.clone() - self._rebuild_sf_indices() # Clean partition: isotropic atoms refine B (adp), anisotropic atoms refine U. self.adp_mask = ~aniso_mask self.u_mask = aniso_mask.clone() @@ -1303,7 +1771,7 @@ def update_mask_from_selection( setattr(self, mask_name, updated_mask) - if self.verbose > 0: + if self.ctx.verbose > 0: n_selected = selection_mask.sum().item() n_refinable = updated_mask.sum().item() action = "frozen" if freeze else "unfrozen" @@ -1347,7 +1815,7 @@ def apply_mask_to_parameter(self, target: str): f"Invalid target: '{target}'. Must be 'xyz', 'adp', 'u', or 'occupancy'" ) - if self.verbose > 0: + if self.ctx.verbose > 0: n_refinable = getattr(self, f"{target}_mask").sum().item() print(f" Applied mask to {target}: {n_refinable} atoms refinable") @@ -1427,7 +1895,7 @@ def get_aniso(self): Selects atoms whose ADP is the 6-element tensor ``u = (u11, u22, u33, u12, u13, u23)``: ``self.aniso_flag``, intersected - with the heavy-atom mask when ``exclude_H_from_sf`` is on. + with the heavy-atom mask when ``hydrogens_in_xray`` is off. Returns ------- @@ -1536,14 +2004,14 @@ def print_parameters_info(self): def register_alternative_conformations(self): """ - Rebuild ``self.altloc_pairs`` from the ``altloc`` column. + Rebuild ``self.ctx.altloc_pairs`` from the ``altloc`` column. One tuple per residue that has multiple conformations, holding one index tensor per conformation (in sorted altloc order), e.g. ``[(tensor([100, 101]), tensor([110, 111])), ...]``. Overwrites any previous content, so call it after the atom numbering changes. """ - self.altloc_pairs = [] + self.ctx.altloc_pairs = [] pdb_with_altlocs = self.pdb[self.pdb["altloc"] != ""] @@ -1561,11 +2029,11 @@ def register_alternative_conformations(self): for altloc in unique_altlocs: altloc_atoms = group[group["altloc"] == altloc] indices = torch.tensor( - altloc_atoms["index"].tolist(), dtype=torch.long + altloc_atoms["index"].tolist(), dtype=torch.long # dtype-ok: altloc atom indices; indexing requires long ) conformation_tensors.append(indices) - self.altloc_pairs.append(tuple(conformation_tensors)) + self.ctx.altloc_pairs.append(tuple(conformation_tensors)) def shake_coords(self, stddev: float): """ @@ -1578,9 +2046,14 @@ def shake_coords(self, stddev: float): new_xyz = xyz + torch.normal( mean=0.0, std=stddev, size=xyz.shape, device=self.device ) - self.xyz = MixedTensor( - new_xyz, refinable_mask=self.xyz.refinable_mask, name="xyz" - ) + if hasattr(self.xyz, "with_values"): + # A riding wrapper keeps its frames; only the stored rows take the noise. + self.xyz = self.xyz.with_values(new_xyz) + else: + self.xyz = MixedTensor( + new_xyz, refinable_mask=self.xyz.refinable_mask, name="xyz" + ) + self._repoint_coordinate_accessors() def shake_adp(self, stddev: float): """ @@ -1597,139 +2070,24 @@ def shake_adp(self, stddev: float): new_adp, refinable_mask=self.adp.refinable_mask, name="adp" ) - def generate_hydrogens(self, mon_lib_path: str = None) -> "Model": - """ - Generate hydrogen atoms for the current model using gemmi. - Places hydrogens at ideal geometry using the CCP4 monomer library and - gemmi's topology engine. Returns a new Model instance with hydrogens - added; the original model is not modified. + def _new_model_from_df(self, df, *, strip_H=None, add_hydrogens=False): + """Build a fresh model of the same class from a DataFrame. - Parameters - ---------- - mon_lib_path : str, optional - Path to CCP4 monomer library directory. If None, uses the monomer - library bundled with torchref (covers standard amino acids and - common small molecules). - - Returns - ------- - Model - A new Model instance with hydrogen atoms added (strip_H=False). - Unknown residues are skipped silently. - - Notes - ----- - Reads the *current* coordinates (via :meth:`update_pdb`), so run it after - any coordinate change that should be reflected in the H positions. + ``add_hydrogens`` defaults to False: the caller has + already settled which atoms the table holds, and generating more would fight + that. :meth:`hydrogenate` passes an already-augmented table for the same reason. """ - import os - import tempfile - - import gemmi - - from torchref import PATH_TORCHREF_DATA - - # ``mgr`` is set when we fall back to TorchRef's auto-fetching monomer - # library manager; per-residue CIFs are then resolved through it (which - # downloads/caches on demand) rather than from ``mon_lib_path`` directly. - mgr = None - if mon_lib_path is None: - import os as _os - - # In priority order: CCP4's own env var, a library bundled next to the - # repo, then the partial one shipped inside torchref. - candidates = [ - _os.environ.get("CLIBD_MON", ""), - str(PATH_TORCHREF_DATA.parent.parent / "external_monomer_library"), - str(PATH_TORCHREF_DATA / "monomer_library"), - ] - mon_lib_path = None - for c in candidates: - if c and _os.path.isfile(_os.path.join(c, "ener_lib.cif")): - mon_lib_path = c - break - if mon_lib_path is None: - # No complete CCP4 library: fall back to TorchRef's manager, which - # ships standard residues and auto-downloads the rest. This is the - # normal path — a CCP4 install is not required. - from torchref.restraints.library import get_library_manager - - mgr = get_library_manager(verbose=self.verbose) - mon_lib_path = str(mgr.ensure_gemmi_base()) - - # gemmi reads from a file, so the live tensors must reach the DataFrame. - self.update_pdb() - - with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as f: - tmp_heavy = f.name - with tempfile.NamedTemporaryFile(suffix=".pdb", delete=False) as f: - tmp_with_h = f.name - - try: - from torchref.io import pdb as io_pdb - from torchref.utils.utils import sanitize_pdb_dataframe - - pdb_out = sanitize_pdb_dataframe(self.pdb.copy()) - pdb_out.attrs["spacegroup"] = ( - self.spacegroup.hm if self.spacegroup else "P 1" - ) - io_pdb.write(pdb_out, tmp_heavy) - - st = gemmi.read_structure(tmp_heavy) - st.setup_entities() - - # Per-residue CIFs come from the manager (bundled → cache → download) - # when we fell back to it, else from the explicit library directory. - monlib = gemmi.read_monomer_lib(mon_lib_path, []) - resnames = set(r.name for m in st for c in m for r in c) - for rn in resnames: - if mgr is not None: - cif = mgr.get_cif_file(rn) - cif_path = str(cif) if cif is not None else None - else: - cif_path = os.path.join(mon_lib_path, rn[0].lower(), rn + ".cif") - if not os.path.exists(cif_path): - cif_path = None - if cif_path is None: - continue - doc = gemmi.cif.read(cif_path) - for block in doc: - if block.name == rn or block.name.startswith("comp_" + rn): - monlib.add_monomer_if_present(block) - break - - gemmi.prepare_topology(st, monlib, h_change=gemmi.HydrogenChange.ReAdd) - st.write_pdb(tmp_with_h) - - # strip_H=False, or the hydrogens we just placed would be dropped again. - new_model = self.__class__( - dtype_float=self.dtype_float, - verbose=self.verbose, - device=self.device, - strip_H=False, - ) - new_model.load_pdb(tmp_with_h) - - finally: - for p in (tmp_heavy, tmp_with_h): - try: - os.unlink(p) - except OSError: - pass - - return new_model - - def _new_model_from_df(self, df, *, strip_H=None): - """Build a fresh model of the same class from a DataFrame.""" import inspect - sh = self.strip_H if strip_H is None else strip_H + sh = self.ctx.strip_H if strip_H is None else strip_H ctor_kw = dict( dtype_float=self.dtype_float, verbose=0, device=self.device, strip_H=sh, + add_hydrogens=add_hydrogens, + cif_path=self.ctx.cif_path, ) sig = inspect.signature(self.__class__.__init__) for pname, param in sig.parameters.items(): @@ -1737,19 +2095,18 @@ def _new_model_from_df(self, df, *, strip_H=None): continue if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD): continue + if pname == "gridsize": + # The constructor argument is the explicit override, not the + # derived grid a ``gridsize`` attribute would return. + if hasattr(self, "explicit_gridsize"): + ctor_kw[pname] = self.explicit_gridsize + continue if hasattr(self, pname): ctor_kw[pname] = getattr(self, pname) - if "gridsize" in sig.parameters and hasattr(self, "_explicit_gridsize"): - ctor_kw["gridsize"] = self._explicit_gridsize new_model = self.__class__(**ctor_kw) sg_str = self.spacegroup.xhm if self.spacegroup else "P 1" new_model.load(lambda: (df, self.pdb.attrs.get("cell"), sg_str)) - if hasattr(new_model, "setup_grid"): - new_model.setup_grid() - # Propagate CIF restraint paths so restraints are rebuilt correctly - if self._cif_path is not None: - new_model._cif_path = self._cif_path return new_model def strip_altlocs(self) -> "Model": @@ -1818,616 +2175,52 @@ def strip_hydrogens(self) -> "Model": filtered.attrs = pdb.attrs.copy() return self._new_model_from_df(filtered, strip_H=True) - # Module-level cache for CIF monomer data (shared across calls) - _hydrogenate_cif_cache = {} - - def hydrogenate( - self, - verbose: int = 0, - optimize: bool = False, - lbfgs_steps: int = 3, - max_iter: int = 20, - ) -> "Model": - """ - Return a new model with hydrogen atoms placed via Kabsch alignment. + def hydrogenate(self, verbose: int = 0, optimize: bool = True) -> "Model": + """Return a new model with hydrogens added from the monomer templates. - Uses torchref's monomer library to identify missing H atoms, places - them by SVD-aligning ideal monomer coordinates onto the current model - coordinates, then corrects each H to sit at ideal bond length from its - parent atom. The original model is not modified. + Hydrogen generation is template instantiation over the topology: each residue's + library template is aligned onto the heavy atoms present and its hydrogens read + off, and the bond graph decides how many hydrogens a parent can carry and which + of them have a free torsion. Missing HOH hydrogens use the water dictionary + geometry with a random initial orientation, controlled by ``torch.manual_seed``. + Existing hydrogen coordinates are retained. The original model is not modified. Parameters ---------- - verbose : int, optional - Verbosity level (0=silent, 1=summary, 2=detailed). Default 0. - optimize : bool, optional - If True, run a short LBFGS geometry optimization on H positions - after placement. Default False (Kabsch placement only). - lbfgs_steps : int, optional - Number of LBFGS outer steps (only when optimize=True). Default 3. - max_iter : int, optional - Max line-search iterations per LBFGS step. Default 20. + verbose : int, default 0 + Verbosity level. + optimize : bool, default True + Scan each free torsion -- hydroxyl, thiol, amine, methyl -- for the + least-clashing angle. The template's dihedral for those is arbitrary, so this + is on by default; it is a rotation about one bond and costs little. Returns ------- Model - New model with hydrogen atoms added. - All parameters are unfrozen in the returned model. + New model with hydrogens, built with ``strip_H=False`` so they survive the + load. """ - import numpy as np - import pandas as pd - - from torchref.restraints.library import MonomerLibraryManager + from torchref.topology.hydrogens import ( + augment_atom_table, + optimise_free_torsions, + plan_hydrogens, + ) - # Sync current coordinates into DataFrame self.update_pdb() + restraints = self.restraints # builds the topology this reads + xyz = self.xyz().detach() - lib = MonomerLibraryManager(verbose=0) - cache = Model._hydrogenate_cif_cache - - # --- Phase A: build per-residue-type lookup tables (cached) --- - for rn in self.pdb["resname"].unique(): - rn_str = str(rn).strip() - if not rn_str: - continue - if rn_str in cache: - if cache[rn_str] is None or "heavy_neighbor_map" in cache[rn_str]: - continue - del cache[rn_str] # Stale entry, re-read - cif_path = lib.get_cif_file(rn_str) - if cif_path is None: - cache[rn_str] = None - continue - try: - from torchref.io.cif_readers import RestraintCIFReader - - reader = RestraintCIFReader(str(cif_path)) - all_data = reader.get_all_restraints() - comp_data = all_data.get(rn_str) or all_data.get(rn_str.upper()) - if comp_data is None: - cache[rn_str] = None - continue - atom_df = comp_data.get("atoms", comp_data.get("atom")) - bond_df = comp_data.get("bonds", comp_data.get("bond")) - if atom_df is None or atom_df.empty or "x" not in atom_df.columns: - cache[rn_str] = None - continue - except Exception: - cache[rn_str] = None - continue - - ids = atom_df["atom_id"].astype(str).str.strip().values - elems = atom_df["type_symbol"].astype(str).str.strip().values - coords = atom_df[["x", "y", "z"]].values.astype(np.float64) - is_h = np.array([e.upper() == "H" for e in elems]) - id_to_idx = {n: i for i, n in enumerate(ids)} - - # H→parent map + ideal bond lengths + heavy adjacency - parent_map = {} # h_name -> parent_name - ideal_bl = {} # h_name -> ideal bond length (Angstrom) - heavy_neighbor_map = {} # heavy_name -> [bonded heavy names] - if bond_df is not None and not bond_df.empty: - a1s = bond_df["atom1"].astype(str).str.strip().values - a2s = bond_df["atom2"].astype(str).str.strip().values - vals = pd.to_numeric(bond_df["value"], errors="coerce").values - h_set = set(ids[is_h]) - for i in range(len(a1s)): - b1, b2 = a1s[i], a2s[i] - if b1 in h_set and b2 in id_to_idx and not is_h[id_to_idx[b2]]: - parent_map[b1] = b2 - if np.isfinite(vals[i]): - ideal_bl[b1] = float(vals[i]) - elif b2 in h_set and b1 in id_to_idx and not is_h[id_to_idx[b1]]: - parent_map[b2] = b1 - if np.isfinite(vals[i]): - ideal_bl[b2] = float(vals[i]) - # Heavy-atom adjacency for local Kabsch - i1, i2 = id_to_idx.get(b1), id_to_idx.get(b2) - if ( - i1 is not None - and i2 is not None - and not is_h[i1] - and not is_h[i2] - ): - heavy_neighbor_map.setdefault(b1, []).append(b2) - heavy_neighbor_map.setdefault(b2, []).append(b1) - - cache[rn_str] = { - "ids": ids, - "elems": elems, - "coords": coords, - "is_h": is_h, - "id_to_idx": id_to_idx, - "heavy_names": ids[~is_h], - "heavy_coords": coords[~is_h], - "h_names": ids[is_h], - "h_coords": coords[is_h], - "parent_map": parent_map, - "ideal_bl": ideal_bl, - "heavy_neighbor_map": heavy_neighbor_map, - } - - # Filter to available residue types - available = { - rn: cache[rn] - for rn in self.pdb["resname"].unique() - if str(rn).strip() in cache and cache.get(str(rn).strip()) is not None - } - if not available: - if verbose > 0: - print("No monomer library data found; returning copy.") - return self.copy() - - # --- Phase B: place H atoms via Kabsch alignment --- - model_names_arr = self.pdb["name"].astype(str).str.strip().values - model_xyz_arr = self.pdb[["x", "y", "z"]].values.astype(np.float64) - model_occ_arr = self.pdb["occupancy"].values.astype(np.float64) - model_bfac_arr = self.pdb["tempfactor"].values.astype(np.float64) - model_atom_type_arr = self.pdb["ATOM"].values - model_altloc_arr = self.pdb["altloc"].values.astype(str) - - group_cols = ["chainid", "resseq", "icode", "resname"] - group_keys = self.pdb[group_cols].values - changes = np.zeros(len(group_keys), dtype=bool) - changes[0] = True - for c in range(4): - changes[1:] |= group_keys[1:, c] != group_keys[:-1, c] - group_starts = np.nonzero(changes)[0] - group_ends = np.append(group_starts[1:], len(group_keys)) - - # Pre-allocate lists for H atom data columns - h_x, h_y, h_z = [], [], [] - h_names_out, h_altlocs, h_resnames = [], [], [] - h_chainids, h_resseqs, h_icodes = [], [], [] - h_occ, h_bfac, h_atom_types = [], [], [] - h_insert_after = [] - - max_bond_dist = 1.5 # Reject H atoms placed > this from parent - _std_val = {"C": 4, "N": 3, "O": 2, "S": 2} - - # Heavy-atom mask for distance-based neighbor detection - model_elem_arr = self.pdb["element"].astype(str).str.strip().values - model_heavy_mask_full = np.array([e.upper() != "H" for e in model_elem_arr]) - - for gi in range(len(group_starts)): - s, e = group_starts[gi], group_ends[gi] - rn = str(group_keys[s, 3]).strip() - info = cache.get(rn) - if info is None: - continue - chainid = group_keys[s, 0] - resseq = group_keys[s, 1] - icode = group_keys[s, 2] - - names_in_model = set(model_names_arr[s:e]) - h_to_add_mask = np.array( - [n not in names_in_model for n in info["h_names"]], dtype=bool - ) - if not h_to_add_mask.any(): - continue - h_names_add = info["h_names"][h_to_add_mask] - h_coords_ideal = info["h_coords"][h_to_add_mask] - - # Altloc handling - altlocs_in_res = set(model_altloc_arr[s:e]) - altloc_list = ( - [""] - if altlocs_in_res <= {""} - else sorted(a for a in altlocs_in_res if a != "") - ) - - for altloc in altloc_list: - if altloc == "": - mask = np.ones(e - s, dtype=bool) - else: - al = model_altloc_arr[s:e] - mask = (al == altloc) | (al == "") - - conf_names = model_names_arr[s:e][mask] - conf_xyz = model_xyz_arr[s:e][mask] - conf_occ = model_occ_arr[s:e][mask] - conf_bfac = model_bfac_arr[s:e][mask] - conf_atom_type = model_atom_type_arr[s:e][mask] - - # Name→index lookup for this conformer - name_to_idx = {} - for j, cn in enumerate(conf_names): - if cn not in name_to_idx: - name_to_idx[cn] = j - - conf_name_set = set(conf_names) - common_mask = np.array( - [n in conf_name_set for n in info["heavy_names"]], - dtype=bool, - ) - n_common = common_mask.sum() - - # Global Kabsch when ≥ 3 matching heavy atoms - R_global = t_global = None - if n_common >= 3: - P = info["heavy_coords"][common_mask] - Q = np.array( - [ - conf_xyz[name_to_idx[n]] - for n in info["heavy_names"][common_mask] - ], - dtype=np.float64, - ) - cp, cq = P.mean(0), Q.mean(0) - Hm = (P - cp).T @ (Q - cq) - U, S, Vt = np.linalg.svd(Hm) - d = np.linalg.det(Vt.T @ U.T) - sign_d = np.diag([1.0, 1.0, 1.0 if d > 0 else -1.0]) - R_global = Vt.T @ sign_d @ U.T - t_global = cq - R_global @ cp - - # Group H atoms by parent for placement - parent_to_hi = {} - for hi, h_name in enumerate(h_names_add): - pn = info["parent_map"].get(h_name) - if pn is not None and pn in name_to_idx: - parent_to_hi.setdefault(pn, []).append(hi) - - hnm = info.get("heavy_neighbor_map", {}) - id2i = info["id_to_idx"] - all_coords = info["coords"] - mask_idx = np.where(mask)[0] # conformer indices in [s:e] - - for par_name, hi_list in parent_to_hi.items(): - pidx = name_to_idx[par_name] - parent_pos = conf_xyz[pidx] - parent_full = s + mask_idx[pidx] - - # Heavy neighbors in the model (distance-based, - # includes cross-residue bonds like C-N peptide) - dvec = model_xyz_arr - model_xyz_arr[parent_full] - dists_sq = (dvec**2).sum(1) - bonded = np.where( - (dists_sq > 0.09) & (dists_sq < 3.61) & model_heavy_mask_full - )[0] - bonded = bonded[bonded != parent_full] - n_model_heavy = len(bonded) - - # Expected H count from standard valence - par_elem = info["elems"][id2i[par_name]].upper() - expected_h = max( - 0, - _std_val.get(par_elem, 4) - n_model_heavy, - ) - - # --- Step 1: local Kabsch for initial placement --- - local_set = {par_name} - for nb in hnm.get(par_name, []): - local_set.add(nb) - for nb2 in hnm.get(nb, []): - local_set.add(nb2) - local_names = [ - n for n in local_set if n in name_to_idx and n in id2i - ] - - if len(local_names) >= 3: - Pl = np.array([all_coords[id2i[n]] for n in local_names]) - Ql = np.array([conf_xyz[name_to_idx[n]] for n in local_names]) - cpl, cql = Pl.mean(0), Ql.mean(0) - Hl = (Pl - cpl).T @ (Ql - cql) - Ul, _, Vtl = np.linalg.svd(Hl) - dl = np.linalg.det(Vtl.T @ Ul.T) - sl = np.diag([1.0, 1.0, 1.0 if dl > 0 else -1.0]) - R_use = Vtl.T @ sl @ Ul.T - t_use = cql - R_use @ cpl - elif R_global is not None: - R_use, t_use = R_global, t_global - else: - R_use = None # Will use random placement - - # Kabsch-place and filter by distance - valid_h = [] - if R_use is not None: - for hi in hi_list: - h_name = h_names_add[hi] - h_cif = all_coords[id2i[h_name]] - h_pos = R_use @ h_cif + t_use - direction = h_pos - parent_pos - dist = np.linalg.norm(direction) - if dist < 1e-6 or dist > max_bond_dist: - continue - bl = info["ideal_bl"].get(h_name, dist) - h_pos = parent_pos + direction * (bl / dist) - valid_h.append((h_name, h_pos, bl)) - else: - # Random-rotation placement (< 3 matching atoms) - # Apply a random SO(3) rotation to ideal CIF - # geometry so internal angles are preserved. - # Random rotation via QR decomposition. - M = np.random.randn(3, 3) - Q_r, _ = np.linalg.qr(M) - if np.linalg.det(Q_r) < 0: - Q_r[:, 0] = -Q_r[:, 0] - par_cif = all_coords[id2i[par_name]] - for hi in hi_list: - h_name = h_names_add[hi] - h_cif = all_coords[id2i[h_name]] - bl = info["ideal_bl"].get(h_name, 0.97) - d_ideal = h_cif - par_cif - d_rot = Q_r @ d_ideal - dn = np.linalg.norm(d_rot) - if dn > 1e-6: - d_rot = d_rot * (bl / dn) - else: - d_rot = np.array([bl, 0.0, 0.0]) - valid_h.append((h_name, parent_pos + d_rot, bl)) - - # Limit to expected count (removes terminal H) - if len(valid_h) > expected_h: - valid_h.sort(key=lambda x: x[0]) # alphabetical - valid_h = valid_h[:expected_h] - - # --- Step 2: geometric re-placement --- - if n_model_heavy >= 2: - nvecs = model_xyz_arr[bonded] - model_xyz_arr[parent_full] - svec = nvecs.sum(0) - snorm = np.linalg.norm(svec) - - if len(valid_h) == 1 and snorm > 1e-6: - # Single H: place opposite to neighbors - h_nm, _, bl = valid_h[0] - h_pos = parent_pos - bl * svec / snorm - valid_h[0] = (h_nm, h_pos, bl) - - elif len(valid_h) == 2 and n_model_heavy == 2 and snorm > 1e-6: - # CH2-like: sp3 tetrahedral placement - v1, v2 = nvecs[0], nvecs[1] - base = -svec / snorm - perp = np.cross(v1, v2) - pn = np.linalg.norm(perp) - if pn > 1e-6: - perp = perp / pn - n1 = np.linalg.norm(v1) - n2 = np.linalg.norm(v2) - c12 = np.dot(v1, v2) / (n1 * n2) - denom = 3.0 * np.sqrt(max(1e-12, (1 + c12) / 2)) - a = min(1.0, 1.0 / denom) - b = np.sqrt(max(0, 1 - a * a)) - d_up = a * base + b * perp - d_dn = a * base - b * perp - # Assign Kabsch-nearest to each - _, pos0, bl0 = valid_h[0] - _, pos1, bl1 = valid_h[1] - g_up = parent_pos + bl0 * d_up - g_dn = parent_pos + bl1 * d_dn - if pos0 is not None and pos1 is not None: - d_same = np.linalg.norm( - pos0 - g_up - ) + np.linalg.norm(pos1 - g_dn) - d_swap = np.linalg.norm( - pos0 - g_dn - ) + np.linalg.norm(pos1 - g_up) - if d_swap < d_same: - g_up, g_dn = g_dn, g_up - valid_h[0] = (valid_h[0][0], g_up, bl0) - valid_h[1] = (valid_h[1][0], g_dn, bl1) - - elif n_model_heavy == 1: - # One heavy neighbor: place H opposite to it - nvec = model_xyz_arr[bonded[0]] - model_xyz_arr[parent_full] - nn = np.linalg.norm(nvec) - if nn > 1e-6: - d_opp = -nvec / nn - for vi in range(len(valid_h)): - if valid_h[vi][1] is None: - nm, _, bl = valid_h[vi] - valid_h[vi] = (nm, parent_pos + bl * d_opp, bl) - - # Fill remaining None positions with random dirs - for vi in range(len(valid_h)): - if valid_h[vi][1] is not None: - continue - nm, _, bl = valid_h[vi] - # Random unit vector via Marsaglia method - while True: - u = np.random.uniform(-1, 1, 3) - n2 = (u * u).sum() - if 0.01 < n2 < 1.0: - break - d = u / np.sqrt(n2) - # Push away from already-placed H siblings - for vj in range(len(valid_h)): - if vj == vi or valid_h[vj][1] is None: - continue - sep = parent_pos + bl * d - valid_h[vj][1] - if np.linalg.norm(sep) < 0.5 * bl: - d = -d # flip to other hemisphere - break - valid_h[vi] = (nm, parent_pos + bl * d, bl) - - # --- Step 3: emit placed H atoms --- - for h_nm, h_pos, _ in valid_h: - h_x.append(h_pos[0]) - h_y.append(h_pos[1]) - h_z.append(h_pos[2]) - h_names_out.append(h_nm) - h_altlocs.append(altloc) - h_resnames.append(rn) - h_chainids.append(chainid) - h_resseqs.append(resseq) - h_icodes.append(icode) - h_occ.append(conf_occ[pidx]) - h_bfac.append(conf_bfac[pidx]) - h_atom_types.append(conf_atom_type[pidx]) - h_insert_after.append(e - 1) - - n_h_placed = len(h_x) - if n_h_placed == 0: - if verbose > 0: - print("No hydrogen atoms to add; returning copy.") - return self.copy() - - if verbose > 0: - print(f"Placing {n_h_placed} hydrogen atoms...") - - # Build H DataFrame in one shot - h_df = pd.DataFrame( - { - "ATOM": h_atom_types, - "serial": 0, - "name": h_names_out, - "altloc": h_altlocs, - "resname": h_resnames, - "chainid": h_chainids, - "resseq": h_resseqs, - "icode": h_icodes, - "x": h_x, - "y": h_y, - "z": h_z, - "occupancy": h_occ, - "tempfactor": h_bfac, - "element": "H", - "charge": 0, - "anisou_flag": False, - "u11": 0.0, - "u22": 0.0, - "u33": 0.0, - "u12": 0.0, - "u13": 0.0, - "u23": 0.0, - } - ) - insert_after = np.array(h_insert_after) - - # Interleave: assign sort keys - n_orig = len(self.pdb) - sort_key = np.empty(n_orig + n_h_placed, dtype=np.float64) - sort_key[:n_orig] = np.arange(n_orig, dtype=np.float64) - _, inv, counts = np.unique( - insert_after, return_inverse=True, return_counts=True - ) - cumcount = np.zeros(n_h_placed, dtype=np.float64) - group_running = np.zeros(len(counts), dtype=np.float64) - for i in range(n_h_placed): - g = inv[i] - cumcount[i] = group_running[g] - group_running[g] += 1 - sort_key[n_orig:] = ( - insert_after + 0.5 + cumcount * (0.4 / np.maximum(counts[inv], 1)) - ) - - augmented_df = pd.concat([self.pdb, h_df], ignore_index=True) - augmented_df = augmented_df.iloc[ - np.argsort(sort_key, kind="stable") - ].reset_index(drop=True) - augmented_df["serial"] = np.arange(1, len(augmented_df) + 1) - augmented_df["index"] = np.arange(len(augmented_df)) - - for col in ( - "x", - "y", - "z", - "occupancy", - "tempfactor", - "u11", - "u22", - "u33", - "u12", - "u13", - "u23", - ): - augmented_df[col] = pd.to_numeric( - augmented_df[col], errors="coerce" - ).astype(float) - augmented_df["serial"] = augmented_df["serial"].astype(int) - augmented_df["resseq"] = augmented_df["resseq"].astype(int) - augmented_df["charge"] = augmented_df["charge"].fillna(0).astype(int) - augmented_df["anisou_flag"] = augmented_df["anisou_flag"].astype(bool) - augmented_df[["altloc", "icode"]] = augmented_df[["altloc", "icode"]].fillna("") - augmented_df["element"] = ( - augmented_df["element"].astype(str).str.strip().str.capitalize() + plan = plan_hydrogens( + restraints.topology, restraints.cif_dict, xyz, verbose=verbose ) - augmented_df.attrs["cell"] = self.pdb.attrs.get("cell") - augmented_df.attrs["spacegroup"] = self.pdb.attrs.get("spacegroup", "P 1") - - new_model = self._new_model_from_df(augmented_df, strip_H=False) - - if verbose > 0: - n_h = (new_model.pdb["element"] == "H").sum() - print(f" New model: {len(new_model.pdb)} atoms ({n_h} H)") - - # --- Phase C (optional): LBFGS geometry optimization --- if optimize: - new_model.freeze_all() - new_model.unfreeze_selection("element H", targets="xyz") - refinable_params = [p for p in new_model.parameters() if p.numel() > 0] - if refinable_params: - try: - from torchref.refinement.targets.combined import ( - TotalGeometryTarget, - ) - - geom_target = TotalGeometryTarget(new_model, verbose=0) - targets = { - n: geom_target[n] - for n in ("bond", "angle", "torsion", "chiral") - } - - def _geom_loss(): - total = torch.tensor(0.0, device=self.device) - for t in targets.values(): - val = t() - if torch.isfinite(val): - total = total + val - return total - - if verbose > 0: - with torch.no_grad(): - init_l = _geom_loss() - print(f" Geometry loss before: {init_l.item():.4f}") - for m in new_model.modules(): - if hasattr(m, "reset_forward_cache"): - m.reset_forward_cache() - - opt = torch.optim.LBFGS( - refinable_params, - lr=0.1, - max_iter=max_iter, - history_size=100, - line_search_fn="strong_wolfe", - ) - best_loss = float("inf") - best_params = [p.data.clone() for p in refinable_params] - - def closure(): - opt.zero_grad() - loss = _geom_loss() - if loss.requires_grad and torch.isfinite(loss): - loss.backward() - for p in refinable_params: - if p.grad is not None: - p.grad.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0) - return loss - - for _ in range(lbfgs_steps): - opt.step(closure) - with torch.no_grad(): - cur = _geom_loss() - if torch.isfinite(cur) and cur.item() < best_loss: - best_loss = cur.item() - best_params = [p.data.clone() for p in refinable_params] - with torch.no_grad(): - for p, bp in zip(refinable_params, best_params): - p.data.copy_(bp) - if verbose > 0: - with torch.no_grad(): - fin_l = _geom_loss() - print(f" Geometry loss after: {fin_l.item():.4f}") - except Exception as e: - if verbose > 0: - print(f" Warning: optimization failed: {e}") - new_model.set_default_masks() - new_model.unfreeze_all() + optimise_free_torsions(plan, restraints.topology, xyz) if verbose > 0: - print(" Hydrogenation complete.") + print(f"Adding {plan.n_hydrogens} hydrogens") + augmented = augment_atom_table(self.pdb, plan, restraints.topology) + return self._new_model_from_df(augmented, strip_H=False) - return new_model def state_dict(self, destination=None, prefix="", keep_vars=False): """ @@ -2456,19 +2249,18 @@ def state_dict(self, destination=None, prefix="", keep_vars=False): destination=destination, prefix=prefix, keep_vars=keep_vars ) - state[prefix + "pdb"] = ( - self.pdb.copy() if hasattr(self, "pdb") and self.pdb is not None else None - ) + state[prefix + "pdb"] = self.pdb.copy() if self.pdb is not None else None state[prefix + "cell"] = self.cell.data.cpu() if self.cell is not None else None # As a string: gemmi.SpaceGroup is not picklable. state[prefix + "spacegroup"] = self.spacegroup.xhm if self.spacegroup else None - state[prefix + "initialized"] = self.initialized + state[prefix + "initialized"] = self.ctx.initialized state[prefix + "dtype_float"] = self.dtype_float state[prefix + "device"] = self.device - state[prefix + "strip_H"] = self.strip_H - state[prefix + "altloc_pairs"] = ( - self.altloc_pairs if hasattr(self, "altloc_pairs") else [] - ) + state[prefix + "strip_H"] = self.ctx.strip_H + state[prefix + "cif_path"] = self.ctx.cif_path + state[prefix + "altloc_pairs"] = self.ctx.altloc_pairs + state[prefix + "hydrogens_in_xray"] = self.ctx.hydrogens_in_xray + state[prefix + "hydrogen_mode"] = self.ctx.hydrogen_mode return state @@ -2482,10 +2274,10 @@ def save_state(self, path: str): Path to save the state dictionary to. """ torch.save(self.state_dict(), path) - if self.verbose > 0: + if self.ctx.verbose > 0: print(f"Saved model state to {path}") - def load_state(self, path: str, strict: bool = True): + def load_state(self, path: str, strict: bool = True, device=None): """ Load the complete state of the model from a file. @@ -2496,16 +2288,203 @@ def load_state(self, path: str, strict: bool = True): strict : bool, optional Accepted for signature compatibility; the restore goes through :meth:`create_from_state_dict`, which is never strict. + device : torch.device, optional + Device to restore onto. Defaults to this model's current device, so an + in-place reload keeps its placement; pass one to restore elsewhere. """ - state_dict = torch.load(path, map_location=self.device, weights_only=False) + target_device = self.device if device is None else device + state_dict = torch.load(path, map_location=target_device, weights_only=False) loaded = type(self).create_from_state_dict( - state_dict, device=self.device, verbose=self.verbose + state_dict, device=target_device, verbose=self.ctx.verbose ) # Adopt the fully-built model's state wholesale. self.__dict__.update(loaded.__dict__) - if self.verbose > 0: + if self.ctx.verbose > 0: print(f"Loaded model state from {path}") + @staticmethod + def _restore_adp_slot(prefix, state_dict, pdb, saved_dtype, xyz_wrapper): + """Rebuild the ``adp`` or ``u`` wrapper, as a node field when the state was one. + + Built from the PDB for its shapes and masks only; ``load_state_dict`` overwrites + every value afterwards. + + A saved :class:`~torchref.model.disorder_field.DisorderFieldTensor` is recognised + by its ``neighbor_list``, not by the shape of its storage: the ``u`` slot holds a + 2-D tensor either way, so shape alone cannot tell a ``(K, 10)`` node field from a + ``(n_atoms, 6)`` per-atom U. + + Parameters + ---------- + prefix : {"adp", "u"} + Which slot to rebuild. ``"u"`` carries the anisotropic representation. + state_dict : dict + The state being restored, read but not consumed. + pdb : pandas.DataFrame + Atom table supplying the initial values. + saved_dtype : torch.dtype + Float dtype the state was saved in. + xyz_wrapper : MixedTensor + The already-rebuilt coordinate wrapper; a node field derives its node + positions from it. + """ + from torchref.model.parameter_wrappers import ( + CholeskyMixedTensor, + PositiveMixedTensor, + ) + + aniso = prefix == "u" + name = "aniso_U" if aniso else "adp" + mask = state_dict.get(f"{prefix}.refinable_mask") + if aniso: + initial = torch.tensor( + pdb[["u11", "u22", "u33", "u12", "u13", "u23"]].values, + dtype=saved_dtype, + ) + else: + initial = torch.tensor(pdb["tempfactor"].values, dtype=saved_dtype) + + saved_nl = state_dict.get(f"{prefix}.neighbor_list") + if saved_nl is None: + # Match load(): the anisotropic U is a CholeskyMixedTensor so a restored + # model refines it in the same positive-definite-by-construction + # parametrization as a freshly-loaded one. + wrapper = CholeskyMixedTensor if aniso else PositiveMixedTensor + return wrapper(initial, refinable_mask=mask, name=name) + + from torchref.model.disorder_field import ( + AnisotropicPayload, + DisorderFieldTensor, + IsotropicPayload, + payload_from_code, + ) + + # The saved code names the payload exactly. Fall back to inferring it from the + # slot for state dicts written before the code existed, where the only payloads + # were the two the slot already implies. + saved_code = state_dict.get(f"{prefix}.payload_code") + if saved_code is not None: + payload = payload_from_code(int(saved_code)) + else: + payload = AnisotropicPayload() if aniso else IsotropicPayload() + saved_values = state_dict[f"{prefix}.fixed_values"] + # Rebuild with the SAVED anchor rows: cluster anchoring makes these length + # n_atoms where single-atom anchoring makes them length K, so reconstructing + # them from scratch would shape-mismatch on load. + saved_anchor_atom = state_dict.get(f"{prefix}.anchor_atom") + saved_anchor_node = state_dict.get(f"{prefix}.anchor_node") + return DisorderFieldTensor( + initial_values=initial, + xyz_fn=xyz_wrapper, + n_nodes=int(saved_values.shape[0]), + k_neighbors=int(saved_nl.shape[1]), + payload=payload, + # Storage is [payload | log sigma | offset], so the extra three columns + # say whether node positions carry a refinable offset. + refine_positions=bool(saved_values.shape[1] == payload.width + 4), + anchor_rows=( + (saved_anchor_atom, saved_anchor_node) + if saved_anchor_atom is not None + else None + ), + refinable_mask=mask, + mask_in_node_space=True, + name=name, + dtype=saved_dtype, + ) + + @classmethod + def _rebuild_wrappers_from_pdb(cls, instance, pdb, state_dict, saved_dtype, device): + """Give ``instance`` parameter wrappers and per-atom buffers of the right shape. + + The half of :meth:`create_from_state_dict` that every subclass needs + identically, so subclasses call this rather than restating it: a per-class copy + drifts, and a restore that rebuilds the wrong wrapper type fails on a shape + mismatch rather than on anything that names the real cause. + + Values are placeholders throughout --- the caller's ``load_state_dict`` is what + puts the saved numbers in. Only shapes, masks and dtypes matter here. + """ + from torchref.model.parameter_wrappers import MixedTensor, OccupancyTensor + + n_atoms = len(pdb) + + xyz_values = torch.tensor(pdb[["x", "y", "z"]].values, dtype=saved_dtype) + if state_dict.get("xyz.h_row") is not None: + # A saved riding wrapper is recognised by its frame buffers, never by + # shape: its storage is (n_base, 3), a plain wrapper's (n_atoms, 3), and + # both are 2-D. The frames restore from the buffers, so no topology is + # needed here. + from torchref.model.riding_xyz import RidingXYZTensor + from torchref.topology.hydrogens import HydrogenFrames + + frames = HydrogenFrames.from_tensors( + state_dict["xyz.h_row"], + state_dict["xyz.parent_row"], + state_dict["xyz.n1_row"], + state_dict["xyz.n2_row"], + state_dict["xyz.frame_valid"], + state_dict.get("xyz.torsion_group"), + state_dict.get("xyz.rotation_group"), + ) + instance.xyz = RidingXYZTensor( + xyz_values, + frames, + refinable_mask=state_dict.get("xyz.refinable_mask"), + mask_in_base_space=True, + name="xyz", + ) + else: + instance.xyz = MixedTensor( + xyz_values, + refinable_mask=state_dict.get("xyz.refinable_mask"), + name="xyz", + ) + instance.adp = cls._restore_adp_slot( + "adp", state_dict, pdb, saved_dtype, instance.xyz + ) + instance.u = cls._restore_adp_slot( + "u", state_dict, pdb, saved_dtype, instance.xyz + ) + + initial_occ = torch.tensor(pdb["occupancy"].values, dtype=saved_dtype) + sharing_groups, altloc_groups, refinable_mask = ( + instance._create_occupancy_groups(pdb, initial_occ) + ) + # A saved mask is in group space; expand it back over atoms. + saved_occ_mask = state_dict.get("occupancy.refinable_mask") + if saved_occ_mask is not None: + if saved_occ_mask.device != sharing_groups.device: + saved_occ_mask = saved_occ_mask.to(sharing_groups.device) + refinable_mask = saved_occ_mask[sharing_groups] + + instance.occupancy = OccupancyTensor( + initial_values=initial_occ, + sharing_groups=sharing_groups, + altloc_groups=altloc_groups, + refinable_mask=refinable_mask, + dtype=saved_dtype, + device=device, + name="occupancy", + ) + + if "aniso_flag" not in instance._buffers or instance.aniso_flag is None: + instance.register_buffer( + "aniso_flag", + torch.tensor(pdb["anisou_flag"].values, dtype=torch.bool), + ) + for mask_name in ("xyz_mask", "adp_mask", "u_mask", "occupancy_mask"): + instance.register_buffer( + mask_name, torch.ones(n_atoms, dtype=torch.bool, device=device) + ) + + # Note: inv_fractional_matrix, fractional_matrix and recB are properties + # delegating to Cell, so they are not registered as buffers. + if state_dict.get("vdw_radii") is not None: + instance.register_buffer( + "vdw_radii", torch.zeros_like(state_dict["vdw_radii"], device=device) + ) + @classmethod def create_from_state_dict( cls, @@ -2525,7 +2504,10 @@ def create_from_state_dict( state_dict : dict State dictionary from torch.save(model.state_dict(), ...). device : torch.device, optional - Device to place tensors on. Defaults to the configured device.current. + Move the restored model here once it is built. The restore itself always + runs on CPU; ``None`` then moves it to the configured default device + (``get_default_device()``), so a round-trip lands beside a same-config + model rather than stranding itself on CPU. Pass a device to override. verbose : int, optional Verbosity level. Default is 1. dtype_float : torch.dtype, optional @@ -2542,9 +2524,16 @@ def create_from_state_dict( anisotropic ``u`` is rebuilt as a :class:`CholeskyMixedTensor`, matching :meth:`load`, so the positive-definite parametrization round-trips. """ - # Resolve dtype/device at call time so the fallbacks below use the - # current config, not the import-time default. - device = normalize_device(device) + # Build on CPU throughout, then move once at the end -- to the caller's device + # if they named one, otherwise to the configured default device, so a restore + # lands beside a same-config model instead of stranding itself on CPU. One + # device for the whole model is the invariant that matters: the wrappers are + # built from the atom table and land on CPU whatever is asked for, so resolving + # an accelerator up front splits the model rather than placing it. + target_device = ( + canonical_device(device) if device is not None else get_default_device() + ) + device = torch.device("cpu") if dtype_float is None: dtype_float = get_float_dtype() pdb = state_dict.pop("pdb", None) @@ -2552,17 +2541,30 @@ def create_from_state_dict( spacegroup = state_dict.pop("spacegroup", None) initialized = state_dict.pop("initialized", False) saved_dtype = state_dict.pop("dtype_float", dtype_float) - saved_device = state_dict.pop("device", device) + state_dict.pop("device", None) # popped so it never reaches load_state_dict strip_H = state_dict.pop("strip_H", True) + cif_path = state_dict.pop("cif_path", None) altloc_pairs = state_dict.pop("altloc_pairs", []) + hydrogens_in_xray = state_dict.pop("hydrogens_in_xray", True) + hydrogen_mode = state_dict.pop("hydrogen_mode", None) instance = cls( - dtype_float=saved_dtype, verbose=verbose, device=device, strip_H=strip_H + dtype_float=saved_dtype, + verbose=verbose, + device=device, + strip_H=strip_H, + cif_path=cif_path, + hydrogens_in_xray=hydrogens_in_xray, ) + if hydrogen_mode is None: + # Older checkpoints: riding wrappers did not exist, so any hydrogens + # present were free parameters. + hydrogen_mode = "riding" if state_dict.get("xyz.h_row") is not None else "free" + instance.ctx.hydrogen_mode = hydrogen_mode instance.pdb = pdb - instance.initialized = initialized - instance.altloc_pairs = altloc_pairs + instance.ctx.initialized = initialized + instance.ctx.altloc_pairs = altloc_pairs # Setter also sets symmetry. instance.spacegroup = spacegroup @@ -2573,89 +2575,7 @@ def create_from_state_dict( # The wrappers are built from the PDB purely to get the right shapes and # masks; load_state_dict below overwrites their values. if pdb is not None: - n_atoms = len(pdb) - - xyz_mask = state_dict.get("xyz.refinable_mask") - adp_mask = state_dict.get("adp.refinable_mask") - u_mask = state_dict.get("u.refinable_mask") - - instance.xyz = MixedTensor( - torch.tensor(pdb[["x", "y", "z"]].values, dtype=saved_dtype), - refinable_mask=xyz_mask, - name="xyz", - ) - instance.adp = PositiveMixedTensor( - torch.tensor(pdb["tempfactor"].values, dtype=saved_dtype), - refinable_mask=adp_mask, - name="adp", - ) - # Match load(): the anisotropic U is a CholeskyMixedTensor so the - # restored model refines it in the same positive-definite-by- - # construction parametrization as a freshly-loaded one. - instance.u = CholeskyMixedTensor( - torch.tensor( - pdb[["u11", "u22", "u33", "u12", "u13", "u23"]].values, - dtype=saved_dtype, - ), - refinable_mask=u_mask, - name="aniso_U", - ) - - # Create OccupancyTensor - initial_occ = torch.tensor(pdb["occupancy"].values, dtype=saved_dtype) - sharing_groups, altloc_groups, refinable_mask = ( - instance._create_occupancy_groups(pdb, initial_occ) - ) - - # Override mask if present in state_dict - saved_occ_mask = state_dict.get("occupancy.refinable_mask") - if saved_occ_mask is not None: - if saved_occ_mask.device != sharing_groups.device: - saved_occ_mask = saved_occ_mask.to(sharing_groups.device) - refinable_mask = saved_occ_mask[sharing_groups] - - instance.occupancy = OccupancyTensor( - initial_values=initial_occ, - sharing_groups=sharing_groups, - altloc_groups=altloc_groups, - refinable_mask=refinable_mask, - dtype=saved_dtype, - device=device, - name="occupancy", - ) - - # Register buffers that are needed - if "aniso_flag" not in instance._buffers or instance.aniso_flag is None: - instance.register_buffer( - "aniso_flag", - torch.tensor(pdb["anisou_flag"].values, dtype=torch.bool), - ) - # Pre-compute SF indices (respects exclude_H_from_sf) - instance._rebuild_sf_indices() - - # Register mask buffers - instance.register_buffer( - "xyz_mask", torch.ones(n_atoms, dtype=torch.bool, device=device) - ) - instance.register_buffer( - "adp_mask", torch.ones(n_atoms, dtype=torch.bool, device=device) - ) - instance.register_buffer( - "u_mask", torch.ones(n_atoms, dtype=torch.bool, device=device) - ) - instance.register_buffer( - "occupancy_mask", torch.ones(n_atoms, dtype=torch.bool, device=device) - ) - - # Register other buffers based on state_dict - # Note: inv_fractional_matrix, fractional_matrix, recB are now properties - # delegating to Cell, so they're not registered as buffers - buffer_names = ["vdw_radii"] - for name in buffer_names: - if name in state_dict and state_dict[name] is not None: - instance.register_buffer( - name, torch.zeros_like(state_dict[name], device=device) - ) + cls._rebuild_wrappers_from_pdb(instance, pdb, state_dict, saved_dtype, device) # Drop only empty-in-dim-0 tensors (placeholders from an atom-less state); # scalars and non-tensor entries must survive for load_state_dict. @@ -2666,6 +2586,11 @@ def create_from_state_dict( } instance.load_state_dict(state_dict, strict=False) + # Always placed: target_device is the caller's device or the configured default, + # never None. Without this the restore used to stay on CPU and split a + # round-trip's restored model from its (default-device) source. + instance.to(target_device) + if verbose > 0: n_atoms = len(instance.pdb) if instance.pdb is not None else 0 print(f"Created Model from state_dict: {n_atoms} atoms") @@ -2707,7 +2632,7 @@ def get_selection_mask(self, selection: str) -> torch.Tensor: """ from torchref.utils.utils import parse_phenix_selection - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError( "Cannot get selection mask from an uninitialized Model. Load data first." ) @@ -2756,7 +2681,7 @@ def select(self, selection: str) -> "Model": """ from torchref.utils.utils import parse_phenix_selection - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError( "Cannot select from an uninitialized Model. Load data first." ) @@ -2772,9 +2697,11 @@ def select(self, selection: str) -> "Model": # type(self), so a subclass returns its own type. selected_model = type(self)( dtype_float=self.dtype_float, - verbose=self.verbose, + verbose=self.ctx.verbose, device=self.device, - strip_H=self.strip_H, + strip_H=self.ctx.strip_H, + cif_path=self.ctx.cif_path, + hydrogens_in_xray=self.ctx.hydrogens_in_xray, ) # ``index`` must be renumbered: the occupancy grouping below reads it. @@ -2783,7 +2710,7 @@ def select(self, selection: str) -> "Model": selected_model.pdb = selected_model.pdb.reset_index(drop=True) selected_model.pdb["index"] = selected_model.pdb.index.to_numpy(dtype=int) - # Setter also sets symmetry; gemmi.SpaceGroup is immutable, so shared. + # The setter rebuilds a SpaceGroup, so the selection gets its own. selected_model.spacegroup = self.spacegroup # The fractional / reciprocal matrices are properties over the Cell, so @@ -2795,18 +2722,21 @@ def select(self, selection: str) -> "Model": selected_model.register_buffer( "aniso_flag", self.aniso_flag[selection_mask].clone() ) - # Pre-compute SF indices (respects exclude_H_from_sf) - selected_model._rebuild_sf_indices() - selected_model.xyz = MixedTensor( - self.xyz()[selection_mask].clone().detach(), - refinable_mask=( - self.xyz.refinable_mask[selection_mask] - if self.xyz.refinable_mask is not None - else None - ), - name="xyz", - ) + if hasattr(self.xyz, "select_rows"): + # Riding wrapper: frames are remapped, a hydrogen whose parent is cut + # becomes an ordinary row. + selected_model.xyz = self.xyz.select_rows(selection_mask) + else: + selected_model.xyz = MixedTensor( + self.xyz()[selection_mask].clone().detach(), + refinable_mask=( + self.xyz.refinable_mask[selection_mask] + if self.xyz.refinable_mask is not None + else None + ), + name="xyz", + ) selected_model.adp = PositiveMixedTensor( self.adp()[selection_mask].clone().detach(), @@ -2845,9 +2775,10 @@ def select(self, selection: str) -> "Model": selected_model.set_default_masks() selected_model.register_alternative_conformations() - selected_model.initialized = True + selected_model.ctx.initialized = True + selected_model.ctx.hydrogen_mode = self.ctx.hydrogen_mode - if self.verbose > 0: + if self.ctx.verbose > 0: print(f"Selected {n_selected}/{len(self.pdb)} atoms with '{selection}'") return selected_model @@ -2864,7 +2795,7 @@ def xyz_fractional(self) -> torch.Tensor: torch.Tensor Tensor of shape (n_atoms, 3) with fractional coordinates. """ - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError( "Model must be initialized to compute fractional coordinates." ) @@ -2901,7 +2832,7 @@ def rotate( Model Self, for method chaining. """ - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError("Model must be initialized to apply rotation.") xyz = self.xyz() @@ -2945,7 +2876,7 @@ def translate(self, translation: torch.Tensor, fractional: bool = False) -> "Mod model.translate(torch.tensor([5.0, 0.0, 0.0])) # 5 Å in x model.translate(torch.tensor([0.5, 0.5, 0.5]), fractional=True) # half cell """ - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError("Model must be initialized to apply translation.") xyz = self.xyz() @@ -2974,11 +2905,233 @@ def get_centroid(self) -> torch.Tensor: torch.Tensor Centroid coordinates with shape (3,). """ - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError("Model must be initialized to compute centroid.") return self.xyz().mean(dim=0) + # ------------------------------------------------------------------ + # Hydrogen parametrisation + # ------------------------------------------------------------------ + + @property + def hydrogen_mode(self) -> str: + """``"riding"``, ``"free"`` or ``"none"``; see :class:`ModelContext`.""" + return self.ctx.hydrogen_mode + + def hydrogen_frames(self): + """Which rows ride on which heavy atoms, for the current atom table. + + Read off the riding coordinate wrapper when one is installed, else derived + from the bond graph, which costs a restraint build the first time. + + Returns + ------- + HydrogenFrames + """ + if hasattr(self.xyz, "hydrogen_frames"): + return self.xyz.hydrogen_frames() + frames = getattr(self, "_hydrogen_frames", None) + if frames is not None and frames.n_hydrogens >= 0: + return frames + from torchref.topology.hydrogens import hydrogen_frames + + return hydrogen_frames(self.restraints.topology) + + def _repoint_coordinate_accessors(self) -> None: + """Make every borrowed coordinate accessor read the current ``xyz`` wrapper. + + The restraints keep ``xyz_fn`` for pair-list maintenance and the ADP node + field borrows the coordinates through ``set_xyz_fn``; after the wrapper slot + is replaced both would otherwise keep reading a dead module. + """ + restraints = self._restraints + if restraints is not None: + restraints._xyz_fn = self.xyz + restraints._adp_fn = self.adp + restraints._vdw_radii_fn = self.get_vdw_radii + for module in self._modules.values(): + if module is not None and hasattr(module, "set_xyz_fn"): + module.set_xyz_fn(self.xyz) + + def _complete_riding_waters(self, frames): + """Complete HOH residues once and remap frames and refinement selections.""" + from dataclasses import fields + + import numpy as np + + from torchref.topology.hydrogens import ( + HydrogenFrames, + augment_atom_table_with_maps, + hydrogen_frames, + plan_hydrogens, + ) + + if not self.ctx.add_hydrogens or self.ctx.strip_H: + return frames + if not self.pdb["resname"].str.strip().eq("HOH").any(): + return frames + restraints = self.restraints + dictionaries = { + key: value for key, value in restraints.cif_dict.items() if key == "HOH" + } + plan = plan_hydrogens(restraints.topology, dictionaries, self.xyz().detach()) + if plan.n_hydrogens == 0: + return frames + + generated = hydrogen_frames(restraints.topology, plan) + if frames is not None: + # Water groups include both existing and planned H atoms; keep custom + # frames for the rest of the table and give the water groups fresh IDs. + water = self.pdb["resname"].str.strip().eq("HOH").to_numpy() + keep = ~water[frames.parent_row] + take = water[generated.parent_row] + arrays = {} + for field in fields(HydrogenFrames): + existing = getattr(frames, field.name)[keep] + added = getattr(generated, field.name)[take].copy() + if field.name in ("torsion_group", "rotation_group"): + added[added >= 0] += int(existing.max(initial=-1)) + 1 + arrays[field.name] = np.concatenate((existing, added)) + generated = HydrogenFrames(**arrays) + + self.update_pdb() + augmented, old_rows, new_rows = augment_atom_table_with_maps( + self.pdb, plan, restraints.topology + ) + frames = generated.remap(old_rows).fill_planned_rows(new_rows) + source = torch.empty(len(augmented), dtype=torch.long, device=self.device) + old_index = torch.as_tensor(old_rows, device=self.device) + new_index = torch.as_tensor(new_rows, device=self.device) + source[old_index] = torch.arange(len(self.pdb), device=self.device) + source[new_index] = torch.as_tensor(plan.parent, device=self.device) + xyz = ( + self.xyz.to_mixed_tensor() + if hasattr(self.xyz, "to_mixed_tensor") + else self.xyz + ) + masks = { + "xyz": xyz.refinable_mask[source], + "occupancy": self.occupancy.get_refinable_atoms()[source], + } + from torchref.model.disorder_field import DisorderFieldTensor + + adp_fields = {} + for name in ("adp", "u"): + wrapper = getattr(self, name) + if isinstance(wrapper, DisorderFieldTensor): + adp_fields[name] = wrapper + else: + masks[name] = wrapper.refinable_mask[source] + if "u" in masks: + masks["u"][new_index] = False + gradients = { + name: getattr(self, name).refinable_params.requires_grad for name in masks + } + adp = self.adp().detach() + cell, spacegroup, links = self.cell, self.spacegroup, self.ctx.links + + def reader(): + return augmented, cell.data.cpu().numpy(), spacegroup + + reader.links = links + strip_h = self.ctx.strip_H + self.ctx.strip_H = False + self._restraints = None + try: + self.load(reader, add_hydrogens=False) + finally: + self.ctx.strip_H = strip_h + if "adp" in masks: + self.adp[old_index] = adp + for name, mask in masks.items(): + wrapper = getattr(self, name) + wrapper.update_refinable_mask(mask) + wrapper.refinable_params.requires_grad_(gradients[name]) + for name, field in adp_fields.items(): + field.anchor_atom = old_index[field.anchor_atom] + field.neighbor_list = field.neighbor_list[source] + field._full_shape = len(augmented) + field.set_xyz_fn(self.xyz) + setattr(self, name, field) + self._hydrogen_frames = frames + return frames + + def set_hydrogen_mode(self, mode: str, frames=None) -> "Model": + """Switch the hydrogen parametrisation of the current atom table. + + Parameters + ---------- + mode : str + ``"riding"``: hydrogen coordinates derive from their parents each forward; + rotatable groups retain shared torsion or orientation parameters. + Missing HOH hydrogens are completed only when ``ctx.add_hydrogens`` + is True and ``ctx.strip_H`` is False. With hydrogen generation disabled, + the atom table is unchanged. + ``"free"``: hydrogens are ordinary refinable atoms again. + ``"none"`` is a different atom table; use + :meth:`strip_hydrogens`. + frames : HydrogenFrames, optional + Riding frames for the current table; default :meth:`hydrogen_frames`. + Water frames are completed and row indices remapped if atoms are added. + + Returns + ------- + Model + Self, for chaining. + + Notes + ----- + Replaces the ``xyz`` wrapper, so any optimizer or ``LossState`` built over the + old parameters is stale; :meth:`Refinement.set_hydrogen_mode` does the + engine-side reset. The refinable set carries over row for row (a hydrogen + released to ``"free"`` follows its parent's mask). Existing atom coordinates + are preserved. Completing waters rebuilds the per-atom wrappers and + restraints; new hydrogens inherit their oxygen's refinement selections. + Water initialization follows ``torch.manual_seed`` and never runs in forward. + """ + from torchref.model.riding_xyz import RidingXYZTensor + + if not self.ctx.initialized: + raise RuntimeError("Load a structure before setting the hydrogen mode.") + if mode == "none": + raise ValueError( + "hydrogen_mode 'none' changes the atom table; use strip_hydrogens()" + ) + if mode not in ("riding", "free"): + raise ValueError(f"unknown hydrogen_mode {mode!r}") + + if mode == "riding": + frames = self._complete_riding_waters(frames) + if isinstance(self.xyz, RidingXYZTensor) and frames is None: + return self + if frames is None: + frames = self.hydrogen_frames() + if isinstance(self.xyz, RidingXYZTensor): + current = self.xyz.to_mixed_tensor() + else: + current = self.xyz + new_xyz = RidingXYZTensor.from_mixed_tensor(current, frames) + else: + if isinstance(self.xyz, RidingXYZTensor): + frames = self.xyz.hydrogen_frames() + new_xyz = self.xyz.to_mixed_tensor() + else: + new_xyz = self.xyz + + if new_xyz is not self.xyz: + # Pop first so the new wrapper registers as a fresh submodule. + self._modules.pop("xyz") + self.xyz = new_xyz + self._repoint_coordinate_accessors() + self._hydrogen_frames = frames + self.ctx.hydrogen_mode = mode + if hasattr(self, "reset_cache"): + self.reset_cache() + if self.ctx.verbose > 0: + print(f"Hydrogen mode: {mode} ({self.xyz})") + return self + def use_rigid_xyz(self) -> "Model": """ Swap ``self.xyz`` for a per-chain :class:`RigidXYZTensor`. @@ -2999,7 +3152,7 @@ def use_rigid_xyz(self) -> "Model": """ from torchref.model.rigid_xyz import RigidXYZTensor - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError( "Model must be initialized before use_rigid_xyz(). " "Load data first with load_pdb() or load_cif()." @@ -3036,7 +3189,7 @@ def use_rigid_xyz(self) -> "Model": "Polymer filter removed every atom — cannot build rigid bodies." ) mobile_mask = torch.from_numpy(mobile_arr).to(device=self.device) - if self.verbose > 0 and int(drop.sum()) > 0: + if self.ctx.verbose > 0 and int(drop.sum()) > 0: n_water = int(is_water.sum()) n_ion = int((is_single_atom & ~is_std & ~is_water).sum()) print( @@ -3062,6 +3215,7 @@ def use_rigid_xyz(self) -> "Model": # submodule cleanly rather than colliding with the old one. self._rigid_original_xyz_container = self._modules.pop("xyz") self.xyz = rigid_xyz + self._repoint_coordinate_accessors() # Snapshot which groups were refinable BEFORE freezing them, so the # restore re-enables exactly those and leaves already-frozen ones alone. @@ -3080,7 +3234,7 @@ def use_rigid_xyz(self) -> "Model": if hasattr(self, "reset_cache"): self.reset_cache() - if self.verbose > 0: + if self.ctx.verbose > 0: print( f"Switched to rigid-body parametrization: {rigid_xyz} " f"({rigid_xyz.n_chains} chain(s))" @@ -3113,7 +3267,13 @@ def restore_xyz_from_rigid(self, commit: bool = True) -> "Model": if commit: with torch.no_grad(): current = self.xyz().detach().clone() - new_xyz = MixedTensor(current, name="xyz", device=self.device) + stashed = getattr(self, "_rigid_original_xyz_container", None) + if stashed is not None and hasattr(stashed, "with_values"): + # A riding wrapper keeps its frames; a rigid motion leaves every + # local offset unchanged. + new_xyz = stashed.with_values(current) + else: + new_xyz = MixedTensor(current, name="xyz", device=self.device) self._modules.pop("xyz", None) self.xyz = new_xyz xyz_mask = getattr(self, "xyz_mask", None) @@ -3132,6 +3292,7 @@ def restore_xyz_from_rigid(self, commit: bool = True) -> "Model": if hasattr(self, "_rigid_original_xyz_container"): del self._rigid_original_xyz_container + self._repoint_coordinate_accessors() # Re-enable exactly the groups use_rigid_xyz() froze, so subsequent # per-atom / ADP refinement has parameters to optimize. diff --git a/torchref/model/model_collection.py b/torchref/model/model_collection.py index d56d6c9a..18dc8c96 100644 --- a/torchref/model/model_collection.py +++ b/torchref/model/model_collection.py @@ -1,9 +1,26 @@ """ Model collection for time-resolved kinetic refinement. -Provides ModelCollection — a named dictionary of MixedModel instances at -different timepoints that share the same base structural models (ModelFT). -Keys match DatasetCollection keys so targets can automatically pair them. +Provides ModelCollection — a named dictionary of mixed models at different timepoints +that share the same base structural models (ModelFT). Keys match DatasetCollection keys +so targets can automatically pair them. + +Populations are stored **factorised**, not as a free vector per timepoint:: + + w(t) = (1 - alpha) * e_ref + alpha * q(t) + +with one mean activation ``alpha`` shared across every timepoint and a per-timepoint +branching ``q(t)`` over the non-reference components. This is the statement that only +the overall degree of activation varies from crystal to crystal, while the branching +among excited states is conserved -- and it makes the mixture exactly linear in +``alpha``, so :meth:`ModelCollection.activation_jacobian` is constant and a second +moment of the activation distribution costs no extra structure-factor evaluation. + +``ModelCollection`` owns the population parameters; each timepoint is a view onto one +row (:class:`_SharedMixedModel`). One consequence worth knowing: freezing or unfreezing +fractions is collection-wide, because a single activation cannot be frozen for one +timepoint alone. Timepoints that genuinely need independent populations are driven +through ``set_fraction_override`` instead. """ from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple @@ -13,42 +30,52 @@ from torchref.utils.device_mixin import DeviceMovementMixin from torchref.utils.device_resolution import resolve_device +from torchref.utils.utils import ModuleReference if TYPE_CHECKING: from torchref.model.model_ft import ModelFT - from torchref.model.mixed_model import MixedModel + +#: Keep activation logits finite by clamping fractions away from 0 and 1. +_FRACTION_EPS = 1e-6 + + +def _logit(p: float) -> float: + """Inverse sigmoid, clamped away from the infinities at 0 and 1.""" + p = min(max(float(p), _FRACTION_EPS), 1.0 - _FRACTION_EPS) + return float(torch.log(torch.tensor(p / (1.0 - p)))) class _SharedMixedModel(DeviceMovementMixin, nn.Module): """ - MixedModel variant that references shared base models without re-registering them. + One timepoint's view of a :class:`ModelCollection`. - Standard MixedModel wraps models in nn.ModuleList, which causes - double-registration when the same ModelFT objects appear in multiple - timepoints. This class stores the shared models as a plain list - (no ownership) and only owns its own fraction parameters. + Owns nothing. The shared base models are held as a plain list and the population + parameters live on the parent collection, so neither is re-registered here -- + the same ownership pattern in both cases, and what keeps a base model's + parameters from appearing once per timepoint in ``parameters()``. - An external fraction override (via ``set_fraction_override``) can replace - the softmax-derived fractions; while active, ``fractions`` and ``forward`` - use the override tensor instead of ``softmax(fraction_params)``. + An external fraction override (via ``set_fraction_override``) replaces the + derived fractions; while active, ``fractions`` and ``forward`` use the override + tensor, and gradients flow to whatever produced it. Parameters ---------- base_models : List[ModelFT] Shared structural models (not re-registered as submodules here). - initial_fractions : List[float] - Initial population fractions (must sum to 1). - frozen_fractions : bool - If True, fractions are excluded from optimization. + collection : ModelCollection + Owner of the activation, branching and dispersion parameters. Referenced + without registration. + index : int + This timepoint's row in the collection's insertion order. device : torch.device, optional - Device for fraction parameters. + Device to reconcile the base models onto. """ def __init__( self, base_models: List["ModelFT"], - initial_fractions: List[float], - frozen_fractions: bool = False, + collection: "ModelCollection", + index: int, device: Optional[torch.device] = None, ): super().__init__() @@ -56,48 +83,35 @@ def __init__( # Store as plain list — the parent ModelCollection owns the ModuleList self._base_models = base_models - n = len(base_models) - if len(initial_fractions) != n: - raise ValueError( - f"Number of fractions ({len(initial_fractions)}) must match " - f"number of models ({n})." - ) - total = sum(initial_fractions) - if abs(total - 1.0) > 1e-3: - raise ValueError(f"Initial fractions must sum to 1.0, got {total:.6f}.") + # Parent reference, deliberately not a submodule: the population parameters + # are the collection's, shared across every timepoint. + self._collection_ref = ModuleReference(collection) + self._index = index - # Normalize to handle floating point drift - initial_fractions = [f / total for f in initial_fractions] - - # Reconcile across *all* base models, not just the first: otherwise a - # mixed-device list stays unreconciled and ``fractions_tensor`` below can - # land on a device the later models are not on. - device = resolve_device(*base_models, device=device) - - # Match base models' float dtype (consistent under a float64 config). - fractions_tensor = torch.tensor( - initial_fractions, dtype=base_models[0].dtype_float, device=device - ) - theta = torch.log(fractions_tensor.clamp(min=1e-6)) - self.fraction_params = nn.Parameter(theta, requires_grad=not frozen_fractions) + # Reconcile across *all* base models, not just the first, so a mixed-device + # list does not stay unreconciled. + resolve_device(*base_models, device=device) # Optional override: when set, fractions property returns this tensor - # instead of softmax(fraction_params). Used by refine_kinetics() to + # instead of the collection's derived row. Used by refine_kinetics() to # route kinetic model predictions directly into the F_calc computation. self._fraction_override: Optional[torch.Tensor] = None - # ------------------------------------------------------------------ - # Properties - # ------------------------------------------------------------------ + + @property + def collection(self) -> "ModelCollection": + """The owning collection.""" + return self._collection_ref.module @property def fractions(self) -> torch.Tensor: """Normalized population fractions -- the override tensor while one is - set (see ``set_fraction_override``), else ``softmax(fraction_params)``. + set (see ``set_fraction_override``), else this timepoint's row of the + parent's :meth:`ModelCollection.fractions_matrix`. """ if self._fraction_override is not None: return self._fraction_override - return torch.softmax(self.fraction_params, dim=0) + return self.collection.fractions_matrix()[self._index] @property def models(self) -> List["ModelFT"]: @@ -120,9 +134,12 @@ def device(self): def dtype_float(self): return self._base_models[0].dtype_float - @property def real_space_grid(self): - return self._base_models[0].real_space_grid + return self._base_models[0].real_space_grid() + + @property + def grid_shape(self): + return self._base_models[0].grid_shape @property def fft(self): @@ -132,10 +149,6 @@ def fft(self): def gridsize(self): return self._base_models[0].gridsize - @property - def map_symmetry(self): - return self._base_models[0].map_symmetry - @property def inv_fractional_matrix(self): return self.cell.inv_fractional_matrix.to(dtype=self.dtype_float) @@ -144,17 +157,11 @@ def inv_fractional_matrix(self): def fractional_matrix(self): return self.cell.fractional_matrix.to(dtype=self.dtype_float) - # ------------------------------------------------------------------ - # Grid / density helpers (delegate to base models) - # ------------------------------------------------------------------ def setup_grid(self, max_res=None, gridsize=None): for model in self._base_models: model.setup_grid(max_res=max_res, gridsize=gridsize) - def get_radius(self, min_radius_Angstrom: float = 4.0) -> int: - return self._base_models[0].get_radius(min_radius_Angstrom) - def build_complete_map(self) -> torch.Tensor: """Mixed electron density: sum_i w_i * density_i.""" fractions = self.fractions @@ -164,9 +171,6 @@ def build_complete_map(self) -> torch.Tensor: density = weighted if density is None else density + weighted return density - # ------------------------------------------------------------------ - # Forward: weighted structure factors - # ------------------------------------------------------------------ def forward(self, hkl: torch.Tensor, recalc: bool = False) -> torch.Tensor: """ @@ -192,32 +196,34 @@ def forward(self, hkl: torch.Tensor, recalc: bool = False) -> torch.Tensor: f_mixed = weighted_f if f_mixed is None else f_mixed + weighted_f return f_mixed - # ------------------------------------------------------------------ - # Freeze / unfreeze - # ------------------------------------------------------------------ def freeze_fractions(self): - self.fraction_params.requires_grad = False + """Freeze the population parameters. + + **Collection-wide.** The mean activation is a single parameter shared by every + timepoint, so it cannot be frozen for one timepoint alone; this delegates to + :meth:`ModelCollection.freeze_all_fractions`. + """ + self.collection.freeze_all_fractions() def unfreeze_fractions(self): - self.fraction_params.requires_grad = True + """Unfreeze the population parameters. Collection-wide; see + :meth:`freeze_fractions`.""" + self.collection.unfreeze_all_fractions() def set_fraction_override(self, fractions: torch.Tensor): """Override fractions with an external tensor (e.g. from kinetic model). - While active, ``self.fractions`` returns this tensor instead of - ``softmax(fraction_params)``, allowing gradients to flow through - the external source. + While active, ``self.fractions`` returns this tensor instead of the + collection's derived row, allowing gradients to flow through the external + source. """ self._fraction_override = fractions def clear_fraction_override(self): - """Remove fraction override, reverting to softmax(fraction_params).""" + """Remove the fraction override, reverting to the collection's derived row.""" self._fraction_override = None - # ------------------------------------------------------------------ - # Convenience - # ------------------------------------------------------------------ def get_vdw_radii(self): return self._base_models[0].get_vdw_radii() @@ -233,8 +239,12 @@ def get_individual_fcalc(self, hkl, recalc=True): def __repr__(self): fracs = self.fractions.detach().tolist() frac_str = ", ".join(f"{f:.3f}" for f in fracs) - frozen_str = "frozen" if not self.fraction_params.requires_grad else "learnable" - return f"_SharedMixedModel({len(self._base_models)} models, fractions=[{frac_str}], {frozen_str})" + learnable = self.collection._activation_logit.requires_grad + frozen_str = "learnable" if learnable else "frozen" + return ( + f"_SharedMixedModel({len(self._base_models)} models, " + f"fractions=[{frac_str}], {frozen_str})" + ) class ModelCollection(DeviceMovementMixin, nn.Module): @@ -284,19 +294,38 @@ def __init__( # Register base models as owned submodules (single source of truth) self._base_models = nn.ModuleList(base_models) - # Per-timepoint mixed models (own only fraction params) + # Per-timepoint views (own nothing; see _SharedMixedModel) self._timepoints = nn.ModuleDict() self._order: List[str] = [] + device = resolve_device(*base_models) + dtype = base_models[0].dtype_float + + # Factor populations as (1 - alpha) * e_ref + alpha * q(t), with one + # shared activation and a branching distribution per timepoint. Refining + # these parameters is opt-in. + self._activation_logit = nn.Parameter( + torch.tensor(_logit(1e-6), dtype=dtype, device=device), + requires_grad=False, + ) + self._branching_logits = nn.ParameterList() + self._branching_rows: Dict[str, int] = {} + + # Dispersion of the activation across crystals, as the fraction of its + # maximum: sigma_alpha^2 = alpha (1 - alpha) * lambda, so 0 <= lambda <= 1 + # holds by construction. Stored as a plain float while fixed, because + # sigmoid can never return exactly 0 and lambda = 0 is what reproduces the + # single-moment (coherent) model. + self._lambda_logit = nn.Parameter( + torch.zeros((), dtype=dtype, device=device), requires_grad=False + ) + self._lambda_fixed: Optional[float] = 0.0 + if self.verbose > 0: print( f"ModelCollection initialized with {len(base_models)} base models" ) - # ------------------------------------------------------------------ - # Add timepoints - # ------------------------------------------------------------------ - def add_timepoint( self, name: str, @@ -326,21 +355,84 @@ def add_timepoint( n = len(self._base_models) if fractions is None: fractions = [1.0 / n] * n + if len(fractions) != n: + raise ValueError( + f"Number of fractions ({len(fractions)}) must match " + f"number of models ({n})." + ) + total = sum(fractions) + if abs(total - 1.0) > 1e-3: + raise ValueError(f"Initial fractions must sum to 1.0, got {total:.6f}.") + fractions = [f / total for f in fractions] + + index = len(self._order) + self._install_populations(name, fractions) mixed = _SharedMixedModel( base_models=list(self._base_models), - initial_fractions=fractions, - frozen_fractions=frozen_fractions, + collection=self, + index=index, ) self._timepoints[name] = mixed self._order.append(name) + if frozen_fractions: + self.freeze_all_fractions() + if self.verbose > 0: frac_str = ", ".join(f"{f:.3f}" for f in fractions) print(f" Added timepoint '{name}': fractions=[{frac_str}]") return self + def _install_populations(self, name: str, fractions: List[float]) -> None: + """Invert requested fractions into the (activation, branching) factorisation. + + The reference component's weight is ``1 - alpha`` by construction, so a + timepoint that is pure reference carries no branching row and leaves the + activation alone. Every other timepoint pins the shared activation; a second + one asking for a different value cannot be represented and is rejected rather + than silently projected. + + Raises + ------ + ValueError + If ``fractions`` implies an activation incompatible with one already set + by an earlier timepoint. + """ + alpha = 1.0 - fractions[0] + + if alpha <= _FRACTION_EPS: + # Pure reference: this is the dark / ground state, i.e. the alpha = 0 + # evaluation of the same parametrisation. No branching row. + return + + current = float(self.alpha_mean) + if self._branching_rows: + if abs(alpha - current) > 1e-3: + established = ", ".join(sorted(self._branching_rows)) + raise ValueError( + f"Timepoint {name!r} asks for activation {alpha:.4f}, but " + f"{established} already set it to {current:.4f}. One activation " + f"fraction is shared across all timepoints -- only the branching " + f"among excited components varies with time. To drive timepoints " + f"with independent populations, use set_fraction_override() on " + f"each one instead of passing fractions here." + ) + else: + with torch.no_grad(): + self._activation_logit.fill_(_logit(alpha)) + + # Branching over the K-1 non-reference components, renormalised within alpha. + excited = torch.tensor( + [f / alpha for f in fractions[1:]], + dtype=self._activation_logit.dtype, + device=self._activation_logit.device, + ) + logits = torch.log(excited.clamp(min=_FRACTION_EPS)) + self._branching_rows[name] = len(self._branching_logits) + self._branching_logits.append(nn.Parameter(logits, requires_grad=False)) + def add_dark( self, fractions: Optional[List[float]] = None ) -> "ModelCollection": @@ -363,11 +455,9 @@ def add_dark( n = len(self._base_models) fractions = [0.0] * n fractions[0] = 1.0 - return self.add_timepoint(self._dark_key, fractions, frozen_fractions=True) - - # ------------------------------------------------------------------ - # Class methods - # ------------------------------------------------------------------ + # No frozen_fractions here: the reference is the alpha = 0 evaluation of the + # shared parametrisation, so it owns nothing that could be frozen. + return self.add_timepoint(self._dark_key, fractions) @classmethod def from_kinetics( @@ -412,10 +502,6 @@ def from_kinetics( return collection - # ------------------------------------------------------------------ - # IHM I/O - # ------------------------------------------------------------------ - @classmethod def from_ihm( cls, @@ -476,10 +562,6 @@ def write_ihm(self, filepath: str, mapping=None, datasets=None) -> None: ) writer.write(filepath) - # ------------------------------------------------------------------ - # Dict-like access - # ------------------------------------------------------------------ - def __getitem__(self, name: str) -> "_SharedMixedModel": return self._timepoints[name] @@ -505,10 +587,6 @@ def items(self) -> List[Tuple[str, "_SharedMixedModel"]]: def get(self, name: str, default=None): return self._timepoints.get(name, default) - # ------------------------------------------------------------------ - # Convenience properties - # ------------------------------------------------------------------ - @property def dark_key(self) -> str: return self._dark_key @@ -544,38 +622,290 @@ def spacegroup(self): def device(self): return self._base_models[0].device - # ------------------------------------------------------------------ - # Fractions inspection - # ------------------------------------------------------------------ - def get_all_fractions(self) -> Dict[str, torch.Tensor]: """Current fractions for each timepoint (including dark).""" return {name: self._timepoints[name].fractions for name in self._order} def get_fractions_matrix(self) -> torch.Tensor: + """All fractions as a matrix ``[n_timepoints, n_models]``, in insertion order. + + Alias of :meth:`fractions_matrix`, kept because it is the established name. + """ + return self.fractions_matrix() + + @property + def alpha_mean(self) -> torch.Tensor: + """Mean activation fraction, shared across all timepoints.""" + return torch.sigmoid(self._activation_logit) + + @property + def lambda_twin(self) -> torch.Tensor: + """Activation dispersion as a fraction of its maximum, in ``[0, 1]``. + + Zero is the coherent single-moment model. One means every crystal is either + fully activated or fully dark. Exactly representable while fixed; once + refinement is enabled it is ``sigmoid`` of a parameter and therefore strictly + interior. + """ + if self._lambda_fixed is not None: + return torch.tensor( + self._lambda_fixed, + dtype=self._lambda_logit.dtype, + device=self._lambda_logit.device, + ) + return torch.sigmoid(self._lambda_logit) + + @property + def sigma_alpha_sq(self) -> torch.Tensor: + """Variance of the activation across crystals. + + ``alpha (1 - alpha) * lambda``, so ``0 <= sigma_alpha_sq <= alpha (1 - alpha)`` + holds by construction -- the upper bound being the Bernoulli case. + """ + alpha = self.alpha_mean + return alpha * (1.0 - alpha) * self.lambda_twin + + def branching(self) -> torch.Tensor: + """Per-timepoint distribution over the non-reference components. + + Returns + ------- + torch.Tensor + Shape ``(n_branching_rows, n_base_models - 1)``, rows summing to 1. + Empty when no non-reference timepoint has been added. + """ + if not len(self._branching_logits): + return torch.zeros( + (0, max(len(self._base_models) - 1, 0)), + dtype=self._activation_logit.dtype, + device=self._activation_logit.device, + ) + return torch.stack( + [torch.softmax(row, dim=0) for row in self._branching_logits], dim=0 + ) + + def activation_jacobian(self) -> torch.Tensor: + """``d(fractions) / d(alpha)`` for every timepoint. + + Shape ``(n_timepoints, n_base_models)``. Reference-only rows are exactly zero, + so the reference dataset carries no activation gradient. Every other row is + ``q(t) - e_ref``, whose entries sum to zero because the fractions stay on the + simplex. + + This is what makes a second moment computable through the same machinery as the + first: the mixture is exactly linear in ``alpha``, so this Jacobian is constant + in ``alpha`` and can be scaled by the same affine scaler as the mixture itself. + """ + n_models = len(self._base_models) + dtype = self._activation_logit.dtype + device = self._activation_logit.device + + q_all = self.branching() + rows = [] + for name in self._order: + row = torch.zeros(n_models, dtype=dtype, device=device) + if name in self._branching_rows: + q = q_all[self._branching_rows[name]] + row = torch.cat( + [torch.full((1,), -1.0, dtype=dtype, device=device), q] + ) + rows.append(row) + if not rows: + return torch.zeros((0, n_models), dtype=dtype, device=device) + return torch.stack(rows, dim=0) + + def fractions_matrix(self) -> torch.Tensor: + """Population fractions for every timepoint, ``[n_timepoints, n_models]``. + + ``e_ref + alpha * activation_jacobian``. Reference-only rows come out as exactly + ``e_ref`` with no gradient path to the activation. + """ + n_models = len(self._base_models) + dtype = self._activation_logit.dtype + device = self._activation_logit.device + + e_ref = torch.zeros(n_models, dtype=dtype, device=device) + e_ref[0] = 1.0 + return e_ref.unsqueeze(0) + self.alpha_mean * self.activation_jacobian() + + def fraction_parameters(self) -> List[nn.Parameter]: + """The population parameters, for handing to an optimizer. + + The shared activation and every branching row, plus the dispersion when it is + refinable. Replaces reaching into a per-timepoint parameter. + """ + params: List[nn.Parameter] = [self._activation_logit] + params.extend(self._branching_logits) + if self._lambda_fixed is None: + params.append(self._lambda_logit) + return params + + def set_activation(self, alpha: float) -> "ModelCollection": + """Set the shared mean activation fraction, in place and without gradient.""" + if not 0.0 <= float(alpha) <= 1.0: + raise ValueError(f"alpha must lie in [0, 1]; got {alpha}") + with torch.no_grad(): + self._activation_logit.fill_(_logit(alpha)) + return self + + def set_branching(self, name: str, q: torch.Tensor) -> "ModelCollection": + """Set one timepoint's branching distribution, in place and without gradient. + + Parameters + ---------- + name : str + Timepoint key. Must be a non-reference timepoint. + q : torch.Tensor + Weights over the ``n_base_models - 1`` non-reference components. Normalised + internally; need not sum to 1. """ - All fractions as a matrix [n_timepoints, n_models]. + if name not in self._branching_rows: + raise KeyError( + f"{name!r} has no branching row -- it is the reference timepoint, " + f"whose fractions are fixed at the alpha = 0 evaluation." + ) + q = torch.as_tensor( + q, dtype=self._activation_logit.dtype, device=self._activation_logit.device + ) + q = q / q.sum() + with torch.no_grad(): + self._branching_logits[self._branching_rows[name]].copy_( + torch.log(q.clamp(min=_FRACTION_EPS)) + ) + return self + + def set_lambda_twin( + self, value: Optional[float], refinable: bool = False + ) -> "ModelCollection": + """Set the activation dispersion, fixed or refinable. + + Parameters + ---------- + value : float or None + Dispersion in ``[0, 1]``. ``None`` keeps the current value and only changes + refinability. + refinable : bool, optional + If True, ``lambda_twin`` becomes ``sigmoid`` of a live parameter and joins + :meth:`fraction_parameters`. Default False, which stores an exact float -- + the only way ``lambda_twin`` can be exactly 0. + """ + if value is not None: + if not 0.0 <= float(value) <= 1.0: + raise ValueError(f"lambda_twin must lie in [0, 1]; got {value}") + with torch.no_grad(): + self._lambda_logit.fill_(_logit(value)) + if refinable: + self._lambda_fixed = None + self._lambda_logit.requires_grad_(True) + else: + self._lambda_fixed = ( + float(value) if value is not None else float(self.lambda_twin) + ) + self._lambda_logit.requires_grad_(False) + return self + + def compute_component_fcalcs( + self, hkl: torch.Tensor, recalc: bool = False + ) -> torch.Tensor: + """Per-base-model structure factors, stacked. + + Each base model is evaluated once, so a caller that needs several fraction + mixtures of the same models pays for the structure factors once rather than + once per mixture. + + Parameters + ---------- + hkl : torch.Tensor + Miller indices of shape (n_reflections, 3). These reach the models + unchanged, so pass the *signed* indices when Bijvoet mates must be + distinguished -- or go through + :meth:`~torchref.io.datasets.collection.DatasetCollection.component_structure_factors`, + which handles the convention. + recalc : bool, optional + Force recomputation rather than reusing each model's cached SF. - Rows are ordered by ``self._order`` (i.e. insertion order). + Returns + ------- + torch.Tensor + Complex structure factors of shape ``(n_base_models, n_reflections)``. """ return torch.stack( - [self._timepoints[n].fractions for n in self._order], dim=0 + [m(hkl, recalc=recalc) for m in self._base_models], dim=0 + ) + + def mix_component_fcalcs( + self, component_fcalcs: torch.Tensor, weights: torch.Tensor + ) -> torch.Tensor: + """Contract stacked per-component SFs with a weight matrix. + + ``weights [T, K] @ component_fcalcs [K, R] -> [T, R]``. Separated from + :meth:`compute_component_fcalcs` because the same component stack is contracted + with more than one weight matrix -- the fractions themselves, and any derivative + of them with respect to a shared parameter. + + Parameters + ---------- + component_fcalcs : torch.Tensor + Complex SFs of shape ``(K, n_reflections)``. + weights : torch.Tensor + Real weights of shape ``(T, K)``. + + Returns + ------- + torch.Tensor + Complex SFs of shape ``(T, n_reflections)``. + """ + return torch.einsum( + "tk,kr->tr", weights.to(component_fcalcs.dtype), component_fcalcs ) - # ------------------------------------------------------------------ - # Freeze / unfreeze helpers - # ------------------------------------------------------------------ + def compute_all_fcalc( + self, hkl: torch.Tensor, recalc: bool = False + ) -> torch.Tensor: + """Mixed ``F_calc`` for every timepoint at once. + + Equivalent to calling each timepoint's ``forward`` in turn, but evaluates each + shared base model once instead of once per timepoint. Rows follow + :meth:`get_fractions_matrix`, i.e. insertion order. + + Parameters + ---------- + hkl : torch.Tensor + Miller indices of shape (n_reflections, 3); see + :meth:`compute_component_fcalcs` on the index convention. + recalc : bool, optional + Force recomputation rather than reusing each model's cached SF. + + Returns + ------- + torch.Tensor + Complex SFs of shape ``(n_timepoints, n_reflections)``. + """ + component_fcalcs = self.compute_component_fcalcs(hkl, recalc=recalc) + return self.mix_component_fcalcs( + component_fcalcs, self.get_fractions_matrix() + ) def freeze_all_fractions(self): - """Freeze fractions at all timepoints.""" - for _, mixed in self: - mixed.freeze_fractions() + """Exclude the population parameters from optimization. + + Acts on the shared activation and every branching row. There is nothing + per-timepoint to freeze: one activation serves all of them, and the reference + timepoint has no parameters at all. + """ + self._activation_logit.requires_grad_(False) + for row in self._branching_logits: + row.requires_grad_(False) def unfreeze_all_fractions(self): - """Unfreeze fractions at all timepoints (except dark).""" - for name, mixed in self: - if name != self._dark_key: - mixed.unfreeze_fractions() + """Include the population parameters in optimization. + + The dispersion ``lambda_twin`` is *not* affected; enable it explicitly with + :meth:`set_lambda_twin` so it can never be refined by accident. + """ + self._activation_logit.requires_grad_(True) + for row in self._branching_logits: + row.requires_grad_(True) def freeze_structures(self): """Freeze xyz and adp on all base models.""" @@ -589,10 +919,6 @@ def unfreeze_structures(self): model.unfreeze("xyz") model.unfreeze("b") - # ------------------------------------------------------------------ - # I/O - # ------------------------------------------------------------------ - def write_pdbs(self, outdir: str): """ Write each base model to a PDB file in *outdir*. @@ -612,10 +938,6 @@ def write_pdbs(self, outdir: str): if self.verbose > 0: print(f" Wrote {path}") - # ------------------------------------------------------------------ - # Repr - # ------------------------------------------------------------------ - def __repr__(self): tp_names = ", ".join(self._order[:4]) if len(self._order) > 4: diff --git a/torchref/model/model_ft.py b/torchref/model/model_ft.py index 85ec6e13..0c3299a1 100644 --- a/torchref/model/model_ft.py +++ b/torchref/model/model_ft.py @@ -1,8 +1,8 @@ """ModelFT -- a :class:`~torchref.model.Model` that can compute structure factors. -Adds the electron-density / FFT path (via an :class:`~torchref.model.SfFFT` -submodule created as soon as both cell and space group are set), the ITC92 -scattering parametrization, and the anomalous f' / f'' correction. +Adds the electron-density / FFT path (an :class:`~torchref.model.SfFFT` submodule +that reads the crystal off the model's context and sizes its grid lazily), the +ITC92 scattering parametrization, and the anomalous f' / f'' correction. """ import math @@ -13,11 +13,10 @@ import torch from torchref.base.fourier import fft, ifft -from torchref.config import dtypes, get_float_dtype, normalize_device +from torchref.config import canonical_device, dtypes, get_default_device, get_float_dtype from torchref.model.model import Model from torchref.model.sf_fft import SfFFT from torchref.symmetry import SpaceGroup -from torchref.symmetry.map_symmetry import MapSymmetry from torchref.utils.caching import CachedForwardMixin @@ -53,15 +52,16 @@ class ModelFT(CachedForwardMixin, Model): ---------- max_res, wavelength, anomalous_threshold : float The constructor arguments above, readable back as attributes. - gridsize, real_space_grid : torch.Tensor - Grid dimensions ``(nx, ny, nz)`` and coordinate grid - ``(nx, ny, nz, 3)``; both live on the ``SfFFT`` submodule. + gridsize : torch.Tensor or None + Grid dimensions ``(nx, ny, nz)``, derived by the ``SfFFT`` submodule from + the cell, space group, ``max_res`` and ``explicit_gridsize`` on first use + and re-derived when any of them changes. A coordinate grid is not stored; + :meth:`real_space_grid` builds one on demand for the few callers that want + the Cartesian positions themselves. map : torch.Tensor or None Most recently computed electron density map. parametrization : dict ITC92 parametrization dictionary {element: (A, B)}. - map_symmetry : MapSymmetry - Symmetry operator for map calculations. """ def __init__( @@ -109,8 +109,17 @@ def __init__( """ super().__init__(*args, **kwargs) - self.max_res = max_res - self._explicit_gridsize = gridsize + # The engine reads cell and space group off ``self.ctx`` as they are set; + # its grid is derived on first use and re-derived when the crystal, + # ``max_res`` or ``explicit_gridsize`` change. + self._fft = SfFFT( + ctx=self.ctx, + max_res=max_res, + explicit_gridsize=gridsize, + dtype_float=self.dtype_float, + device=self.device, + verbose=self.ctx.verbose, + ) self.wavelength = wavelength self.anomalous_threshold = anomalous_threshold @@ -118,52 +127,61 @@ def __init__( # so it round-trips through state_dict and follows .to(device). f' is always # applied when wavelength is set; f'' only when this is True (unmerged data). self.register_buffer( - "anomalous_bijvoet", torch.tensor(bool(apply_bijvoet)), persistent=True + "anomalous_bijvoet", + torch.tensor(bool(apply_bijvoet), device=self.device), + persistent=True, ) self._anomalous_cache = None # Will hold (mask, f_prime, f_double_prime) self._anomalous_elements_hash = ( None # Hash of element list for cache invalidation ) - self._fft = None + + # ========================================================================= + # Engine binding and grid inputs + # ========================================================================= + + @property + def fft(self) -> SfFFT: + """The SfFFT submodule, bound to this model's context. + + ``copy()`` and ``load_state`` replace the context object itself; re-pointing + the engine here keeps ``fft.ctx is self.ctx`` on every path. + """ + fft = self._fft + if fft.ctx is not self.ctx: + fft.ctx = self.ctx + return fft + + @property + def max_res(self) -> Optional[float]: + """Maximum resolution in Angstroms that sizes the grid; owned by the engine.""" + return self._fft.max_res + + @max_res.setter + def max_res(self, value) -> None: + self._fft.max_res = None if value is None else float(value) @property - def cell(self): - """Unit cell object with parameters [a, b, c, alpha, beta, gamma].""" - return self._cell + def explicit_gridsize(self) -> Optional[Tuple[int, int, int]]: + """Fixed grid dimensions overriding ``max_res``, or None.""" + return self._fft.explicit_gridsize - @cell.setter - def cell(self, value): - """Set the unit cell; also builds the FFT once the spacegroup is set.""" - self._cell = value - self._maybe_initialize_fft() + @explicit_gridsize.setter + def explicit_gridsize(self, value) -> None: + self._fft.explicit_gridsize = value @property - def spacegroup(self): - """Space group object.""" - return self._spacegroup - - @spacegroup.setter - def spacegroup(self, value): - """Set the space group (SpaceGroup, gemmi.SpaceGroup, name or number); - also builds the FFT once the cell is set. + def grid_key(self): + """What the grid is derived from; see :attr:`SfFFT.grid_key`.""" + return self.fft.grid_key + + def _fingerprint_state(self): + """Fold the grid key into the forward-cache key. + + Parameters and buffers alone would miss a cell, space-group or resolution + change that leaves the grid buffers untouched until the next forward. """ - if value is not None: - self._spacegroup = SpaceGroup( - value, dtype=self.dtype_float, device=self.device - ) - else: - self._spacegroup = None - self._maybe_initialize_fft() - - def _maybe_initialize_fft(self): - """(Re)build the SfFFT submodule once both cell and spacegroup are set.""" - if self._cell is not None and self._spacegroup is not None: - self._fft = SfFFT( - cell=self._cell, - spacegroup=self._spacegroup, - device=self.device, - max_res=self.max_res, - ) + return super()._fingerprint_state() + (self.fft.grid_key,) def load_pdb(self, filename): """ @@ -180,8 +198,6 @@ def load_pdb(self, filename): Self, for method chaining. """ super().load_pdb(filename) - # FFT is now initialized via cell/spacegroup setters in parent load() - self.setup_grid() return self def select(self, selection): @@ -189,9 +205,8 @@ def select(self, selection): Return a new ModelFT containing only the selected atoms. Extends :meth:`Model.select` with the FT-specific setup: rebuilding - the ITC92 parametrization and the real-space grid for the reduced - atom set. The FFT itself is initialized via the cell/spacegroup - setters during the base ``select``. + the ITC92 parametrization and carrying ``max_res`` and + ``explicit_gridsize`` across, so the selection sizes its grid the same way. Parameters ---------- @@ -205,15 +220,14 @@ def select(self, selection): Notes ----- - The ModelFT-specific constructor arguments -- ``max_res``, - ``wavelength``, ``anomalous_threshold``, ``gridsize`` -- are **not** - propagated: :meth:`Model.select` passes only the base kwargs, so the - returned model silently carries the ModelFT defaults for those. + ``wavelength`` and ``anomalous_threshold`` are **not** propagated: + :meth:`Model.select` passes only the base kwargs, so the returned model + carries the ModelFT defaults for those. """ selection = super().select(selection) selection._build_parametrization() - # FFT is initialized via cell/spacegroup setters in parent select() - selection.setup_grid() + selection.max_res = self.max_res + selection.explicit_gridsize = self.explicit_gridsize return selection def load_cif(self, filename): @@ -232,36 +246,8 @@ def load_cif(self, filename): """ super().load_cif(filename) self._build_parametrization() - # FFT is now initialized via cell/spacegroup setters in parent load() - self.setup_grid() return self - def setup_gridsize(self, max_res=None): - """ - Compute optimal grid dimensions. - - Delegates to FFT.compute_grid_size(). - - Parameters - ---------- - max_res : float, optional - Maximum resolution in Angstroms. If None, uses self.max_res. - - Returns - ------- - torch.Tensor - Grid dimensions (nx, ny, nz) as int32 tensor. - """ - if max_res is not None: - self.max_res = max_res - self._fft.max_res = max_res - - if self.verbose > 1: - print(f"Defining grid size for max_res={self.max_res} Å") - - gridsize = self.cell.compute_grid_size(self.max_res) - return torch.tensor(gridsize, dtype=dtypes.int, device=self.device) - def _build_parametrization(self): """Build the ITC92 parametrization (delegates to :class:`Model`).""" return super()._build_parametrization() @@ -283,48 +269,39 @@ def B(self) -> torch.Tensor: return self._B # ========================================================================= - # Backward-compatible properties for FFT grid attributes + # Grid, resolved by the engine # ========================================================================= @property def gridsize(self) -> Optional[torch.Tensor]: - """Grid dimensions (nx, ny, nz).""" - return self._fft.gridsize + """Grid dimensions (nx, ny, nz), or None until cell and space group are set.""" + return self.fft.gridsize - @gridsize.setter - def gridsize(self, value): - """Set grid size (for backward compatibility).""" - self._fft.gridsize = value + def real_space_grid(self) -> torch.Tensor: + """Build the Cartesian coordinate of every grid point, ``(nx, ny, nz, 3)``. - @property - def real_space_grid(self) -> Optional[torch.Tensor]: - """Real-space coordinate grid with shape (nx, ny, nz, 3).""" - return self._fft.real_space_grid + Not stored: at ``12 * nx * ny * nz`` bytes it is the largest tensor a model + would hold, and no structure-factor path reads it -- every splat derives a + voxel's position from its index. Built here for the callers that genuinely + want the coordinates, and discarded when they are done with it. + """ + from torchref.base.fourier import get_real_grid - @real_space_grid.setter - def real_space_grid(self, value): - """Set real space grid (for backward compatibility).""" - self._fft.real_space_grid = value + return get_real_grid( + fractional_matrix=self.cell.fractional_matrix, + gridsize=self.gridsize, + device=self.device, + ) @property - def voxel_size(self) -> Optional[torch.Tensor]: - """Voxel dimensions.""" - return self._fft.voxel_size - - @voxel_size.setter - def voxel_size(self, value): - """Set voxel size (for backward compatibility).""" - self._fft.voxel_size = value + def grid_shape(self) -> Optional[tuple]: + """Map dimensions ``(nx, ny, nz)``, or None until cell and space group are set.""" + return self.fft.grid_shape @property - def map_symmetry(self) -> Optional[MapSymmetry]: - """Symmetry operator for map calculations.""" - return self._fft.map_symmetry - - @map_symmetry.setter - def map_symmetry(self, value): - """Set map symmetry (for backward compatibility).""" - self._fft.map_symmetry = value + def voxel_size(self) -> Optional[torch.Tensor]: + """Voxel edge vector sum, or None until cell and space group are set.""" + return self.fft.voxel_size def get_iso(self): """ @@ -378,69 +355,22 @@ def get_aniso(self): return xyz, u, occupancy, A, B - def setup_grid(self, max_res=None, gridsize=None): + def setup_grid(self, *, max_res=None, gridsize=None): """ - Setup real-space grid for electron density calculation. + Override the grid's inputs explicitly and resolve the grid now. - Delegates to FFT.setup_grid() using the stored cell and spacegroup. + Not needed on the normal path: the engine sizes its grid from the cell, + space group and ``max_res`` on first use and follows any later change. Parameters ---------- max_res : float, optional - Maximum resolution for grid spacing in Angstroms. - If None, uses self.max_res. + New maximum resolution in Angstroms. None leaves the current value. gridsize : tuple of int, optional - Explicit grid size (nx, ny, nz). If None, computed automatically - using Cell.compute_grid_size() and SpaceGroup.suggest_grid_size(). - """ - if max_res is not None: - self.max_res = max_res - self._fft.max_res = max_res - - if self.verbose > 1: - print(f"Setting up grids with max_res={self.max_res} Å") - - gridsize_to_use = gridsize or self._explicit_gridsize - - self._fft.setup_grid( - gridsize=gridsize_to_use, - max_res=self.max_res, - ) - - if self.verbose > 2: - print(f"Grid shape: {self._fft.real_space_grid.shape[:-1]}") - print(f"Voxel size: {self._fft.voxel_size}") - - def get_radius(self, min_radius_Angstrom: float = 4.0): + Fixed grid size (nx, ny, nz). None leaves :attr:`explicit_gridsize` + unchanged. """ - Get a single fixed splat radius in voxels for the given minimum. - - Vestigial: the density path truncates each atom at its own - ``torchref.sigma_cutoff_ed * sigma_eff`` radius and never consults this. - - Parameters - ---------- - min_radius_Angstrom : float, optional - Minimum radius in Angstroms. Default is 4.0. - - Returns - ------- - int - Radius in voxels. - """ - if not hasattr(self, "real_space_grid") or self.real_space_grid is None: - self.setup_grid() - voxel_size = self.real_space_grid[1, 1, 1] - self.real_space_grid[0, 0, 0] - min_radius = ( - torch.ceil(min_radius_Angstrom / torch.min(voxel_size)) - .to(dtypes.int) - .item() - ) - if self.verbose > 1: - print( - f"Calculated radius for density calculation: {min_radius} voxels (voxel size: {voxel_size}), this corresponds to at least {min_radius_Angstrom} Å" - ) - return min_radius + self.fft.setup_grid(max_res=max_res, gridsize=gridsize) def build_complete_map(self, radius=None, apply_symmetry=True): """ @@ -466,7 +396,7 @@ def build_complete_map(self, radius=None, apply_symmetry=True): """ self.map = self.build_initial_map(apply_symmetry=apply_symmetry) - if self.verbose > 2: + if self.ctx.verbose > 2: print( f"Density map built. Sum: {self.map.sum():.2f}, Max: {self.map.max():.4f}" ) @@ -488,15 +418,12 @@ def build_initial_map(self, apply_symmetry=True): torch.Tensor Electron density map with shape (nx, ny, nz). """ - if self._fft.real_space_grid is None: - self.setup_grid() - - if self.verbose > 2: + if self.ctx.verbose > 2: print("Building density map (per-atom variable radius)...") xyz_iso, adp_iso, occ_iso, A_iso, B_iso = self.get_iso() - if self.verbose > 3: + if self.ctx.verbose > 3: assert torch.all( torch.isfinite(A_iso) ), "Non-finite values found in A_iso during map building." @@ -529,7 +456,7 @@ def build_initial_map(self, apply_symmetry=True): apply_symmetry=apply_symmetry, ) - if self.verbose > 3: + if self.ctx.verbose > 3: assert torch.all( torch.isfinite(self.map) ), "Non-finite values found in map." @@ -558,7 +485,7 @@ def save_map(self, filename): np_map = self.map.detach().cpu().numpy().astype(np.float32) cell = self.cell.tolist() - if self.verbose > 0: + if self.ctx.verbose > 0: print(f"Saving map to {filename}") print(f" Map shape: {self.map.shape}") print(f" Map sum: {self.map.sum():.2f}") @@ -571,7 +498,7 @@ def save_map(self, filename): map_ccp.setup(0.0) map_ccp.update_ccp4_header() map_ccp.write_ccp4_map(filename) - if self.verbose > 0: + if self.ctx.verbose > 0: print("Map saved successfully") def get_map_statistics(self): @@ -641,7 +568,7 @@ def _get_anomalous_cache( unique_elements, self.wavelength, self.anomalous_threshold ) - if self.verbose > 1 and significant: + if self.ctx.verbose > 1 and significant: print( f"Anomalous scatterers at {self.wavelength:.4f} Å: " f"{list(significant.keys())}" @@ -750,14 +677,6 @@ def get_structure_factor( """ return self(hkl, recalc=recalc, apply_anomalous=apply_anomalous) - @property - def fft(self): - """The SfFFT submodule, built on first access (needs cell + spacegroup).""" - if self._fft is None: - self._maybe_initialize_fft() - - return self._fft - def _check_forward_dtype(self, hkl: torch.Tensor) -> None: """Fail fast on a model/input float-dtype mismatch, which would otherwise surface as a cryptic matmul or Triton-compile error deep in the kernels. @@ -820,7 +739,7 @@ def forward(self, hkl, apply_anomalous: bool = True) -> torch.Tensor: sf, hkl, include_fdp=bool(self.anomalous_bijvoet) ) - if self.verbose > 2: + if self.ctx.verbose > 2: assert torch.all( torch.isfinite(sf) ), "Non-finite values found while calculating fcalc." @@ -832,8 +751,9 @@ def copy(self, detach: bool = True) -> "ModelFT": Create a deep copy of the ModelFT. Creates a complete independent copy including all Model base class data, - FFT submodule state (gridsize, real_space_grid, voxel_size, map_symmetry), - ITC92 parametrization, and scalar attributes. + the grid inputs (``max_res``, ``explicit_gridsize``; the grid itself is + re-derived from the copied context), the ITC92 parametrization, and + scalar attributes. Cache is reset to empty. Parameters @@ -841,39 +761,29 @@ def copy(self, detach: bool = True) -> "ModelFT": detach : bool, optional If True, the copy's parameters will be detached from the computation graph (default: True). - Returns ------- ModelFT A new, fully independent ModelFT instance with copied data. """ - if not self.initialized: + if not self.ctx.initialized: raise RuntimeError("Cannot copy an uninitialized ModelFT. Load data first.") model_copy = ModelFT( dtype_float=self.dtype_float, - verbose=self.verbose, + verbose=self.ctx.verbose, device=self.device, - strip_H=self.strip_H, + strip_H=self.ctx.strip_H, max_res=self.max_res, - gridsize=self._explicit_gridsize, + gridsize=self.explicit_gridsize, wavelength=self.wavelength, anomalous_threshold=self.anomalous_threshold, ) - model_copy.pdb = self.pdb.copy(deep=True) - - if self._spacegroup is not None: - model_copy._spacegroup = self._spacegroup.copy() - else: - model_copy._spacegroup = None - - model_copy.initialized = True + # Carries the atom table, cell, space group, altloc groups and provenance. + model_copy.ctx = self.ctx.copy() - if self.cell is not None: - model_copy.cell = self.cell.clone() - - # Own buffers only; the FFT submodule's are handled by its copy() below. + # Own buffers only; the engine's grid buffers are derived, not copied. for buffer_name, buffer_value in self._buffers.items(): if buffer_value is not None: if detach: @@ -883,36 +793,29 @@ def copy(self, detach: bool = True) -> "ModelFT": else: model_copy.register_buffer(buffer_name, buffer_value.clone()) - # Parameter wrappers via their own .copy(); _fft / _spacegroup are separate. - skip_modules = {"_fft", "_spacegroup", "spacegroup", "_symmetry", "symmetry"} + # Parameter wrappers via their own .copy(); the engine came from the ctor. + skip_modules = {"_fft"} for module_name, module in self._modules.items(): if module_name in skip_modules: continue if module is not None and hasattr(module, "copy"): setattr(model_copy, module_name, module.copy()) - # Copy alternative conformation pairs - if hasattr(self, "altloc_pairs") and self.altloc_pairs: - model_copy.altloc_pairs = [ - tuple(tensor.clone() for tensor in group) for group in self.altloc_pairs - ] - else: - model_copy.altloc_pairs = [] - if hasattr(self, "_parametrization") and self._parametrization is not None: import copy as copy_module model_copy._parametrization = copy_module.deepcopy(self._parametrization) - if self._fft is not None: - model_copy._fft = self._fft.copy() - if self._fft.real_space_grid is not None: - model_copy.setup_grid(max_res=self.max_res) + # Borrowed coordinate accessors (restraints, ADP node field) still point at + # THIS model's wrappers after their own ``copy``; re-point them. + model_copy._repoint_coordinate_accessors() # Don't share cached structure factors with the original. model_copy.reset_cache() + # The iso/aniso partition is derived state, not a buffer, so it is not + # carried by the buffer loop above; get_iso()/get_aniso() read it. - if self.verbose > 0: + if self.ctx.verbose > 0: print(f"✓ ModelFT copied successfully ({len(model_copy.pdb)} atoms)") return model_copy @@ -922,8 +825,9 @@ def state_dict(self, destination=None, prefix="", keep_vars=False): Return a dictionary containing the complete state of the ModelFT. Extends parent Model.state_dict() with FT-specific parameters: - ``max_res``, ``wavelength``, and ``anomalous_threshold``. Grid state - is handled by the FFT submodule. + ``max_res``, ``explicit_gridsize``, ``wavelength`` and + ``anomalous_threshold``. The grid is derived from these and the crystal, + so it is not stored. Parameters ---------- @@ -939,12 +843,13 @@ def state_dict(self, destination=None, prefix="", keep_vars=False): dict Complete state dictionary. """ - # Parent covers _A/_B and the FFT submodule's buffers. + # Parent covers _A/_B; the engine's grid buffers are non-persistent. state = super().state_dict( destination=destination, prefix=prefix, keep_vars=keep_vars ) state[prefix + "max_res"] = self.max_res + state[prefix + "explicit_gridsize"] = self.explicit_gridsize state[prefix + "wavelength"] = self.wavelength state[prefix + "anomalous_threshold"] = self.anomalous_threshold @@ -971,7 +876,9 @@ def create_from_state_dict( state_dict : dict State dictionary from torch.save(model.state_dict(), ...). device : torch.device, optional - Device to place tensors on. Defaults to the configured device.current. + Move the restored model here once it is built. The restore itself always + runs on CPU; ``None`` then moves it to the configured default device; see + :meth:`Model.create_from_state_dict`. verbose : int, optional Verbosity level. Default is 1. dtype_float : torch.dtype, optional @@ -989,15 +896,19 @@ def create_from_state_dict( The anisotropic ``u`` is rebuilt as a :class:`CholeskyMixedTensor`, as in :meth:`load`, so the positive-definite parametrization round-trips. """ - from torchref.symmetry import SpaceGroup - - # Resolve dtype/device at call time so the fallback below uses the - # current config rather than an import-time default. - device = normalize_device(device) + # Build on CPU throughout and move once at the end, as Model does; the grid + # setup below otherwise sizes an accelerator allocation before the model is + # placed. The final target is the caller's device, or the configured default + # when they name none, so a restore lands beside a same-config model. + target_device = ( + canonical_device(device) if device is not None else get_default_device() + ) + device = torch.device("cpu") if dtype_float is None: dtype_float = get_float_dtype() max_res = state_dict.pop("max_res", 1.0) + explicit_gridsize = state_dict.pop("explicit_gridsize", None) state_dict.pop("radius_angstrom", None) # legacy key, no longer used wavelength = state_dict.pop("wavelength", 1.0) anomalous_threshold = state_dict.pop("anomalous_threshold", 0.5) @@ -1010,11 +921,17 @@ def create_from_state_dict( state_dict.pop("device", None) # Remove but don't use (use provided device) strip_H = state_dict.pop("strip_H", True) altloc_pairs = state_dict.pop("altloc_pairs", []) - - # FFT submodule buffers are prefixed "_fft."; older checkpoints are flat. - gridsize = state_dict.pop("_fft.gridsize", None) - if gridsize is None: - gridsize = state_dict.pop("gridsize", None) + hydrogens_in_xray = state_dict.pop("hydrogens_in_xray", True) + hydrogen_mode = state_dict.pop("hydrogen_mode", None) + + # Checkpoints written while the grid was stored state carry its buffers + # ("_fft." prefixed, or flat in older ones). The size is adopted below only + # when it differs from what the crystal and max_res give. + legacy_gridsize = state_dict.pop("_fft.gridsize", None) + if legacy_gridsize is None: + legacy_gridsize = state_dict.pop("gridsize", None) + state_dict.pop("_fft.voxel_size", None) + state_dict.pop("voxel_size", None) instance = cls( dtype_float=saved_dtype, @@ -1022,15 +939,20 @@ def create_from_state_dict( device=device, strip_H=strip_H, max_res=max_res, + gridsize=explicit_gridsize, wavelength=wavelength, anomalous_threshold=anomalous_threshold, + hydrogens_in_xray=hydrogens_in_xray, ) instance.pdb = pdb - instance.initialized = initialized - instance.altloc_pairs = altloc_pairs + instance.ctx.initialized = initialized + instance.ctx.altloc_pairs = altloc_pairs + if hydrogen_mode is None: + hydrogen_mode = "riding" if state_dict.get("xyz.h_row") is not None else "free" + instance.ctx.hydrogen_mode = hydrogen_mode - # Setter also sets symmetry; the cell setter below then builds the FFT. + # The engine reads both off the context; nothing further to build. instance.spacegroup = spacegroup_str from torchref.symmetry import Cell @@ -1038,87 +960,12 @@ def create_from_state_dict( if cell_tensor is not None: instance.cell = Cell(cell_tensor, dtype=saved_dtype, device=device) - # If PDB exists, create the parameter wrappers with correct shapes + # Wrappers and per-atom buffers: shared with Model so the two restores cannot + # drift apart again. ModelFT adds only its own scattering buffers below. if pdb is not None: - from torchref.model.parameter_wrappers import ( - CholeskyMixedTensor, - MixedTensor, - OccupancyTensor, - PositiveMixedTensor, - ) - - n_atoms = len(pdb) - - xyz_mask = state_dict.get("xyz.refinable_mask") - adp_mask = state_dict.get("adp.refinable_mask") - u_mask = state_dict.get("u.refinable_mask") - - instance.xyz = MixedTensor( - torch.tensor(pdb[["x", "y", "z"]].values, dtype=saved_dtype), - refinable_mask=xyz_mask, - name="xyz", - ) - instance.adp = PositiveMixedTensor( - torch.tensor(pdb["tempfactor"].values, dtype=saved_dtype), - refinable_mask=adp_mask, - name="adp", - ) - instance.u = CholeskyMixedTensor( - torch.tensor( - pdb[["u11", "u22", "u33", "u12", "u13", "u23"]].values, - dtype=saved_dtype, - ), - refinable_mask=u_mask, - name="aniso_U", - ) - - initial_occ = torch.tensor(pdb["occupancy"].values, dtype=saved_dtype) - sharing_groups, altloc_groups, refinable_mask = ( - instance._create_occupancy_groups(pdb, initial_occ) - ) - - saved_occ_mask = state_dict.get("occupancy.refinable_mask") - if saved_occ_mask is not None: - if saved_occ_mask.device != sharing_groups.device: - saved_occ_mask = saved_occ_mask.to(sharing_groups.device) - refinable_mask = saved_occ_mask[sharing_groups] - - instance.occupancy = OccupancyTensor( - initial_values=initial_occ, - sharing_groups=sharing_groups, - altloc_groups=altloc_groups, - refinable_mask=refinable_mask, - dtype=saved_dtype, - device=device, - name="occupancy", - ) - - if "aniso_flag" not in instance._buffers or instance.aniso_flag is None: - instance.register_buffer( - "aniso_flag", - torch.tensor(pdb["anisou_flag"].values, dtype=torch.bool), - ) - - # Register mask buffers - instance.register_buffer( - "xyz_mask", torch.ones(n_atoms, dtype=torch.bool, device=device) - ) - instance.register_buffer( - "adp_mask", torch.ones(n_atoms, dtype=torch.bool, device=device) + cls._rebuild_wrappers_from_pdb( + instance, pdb, state_dict, saved_dtype, device ) - instance.register_buffer( - "u_mask", torch.ones(n_atoms, dtype=torch.bool, device=device) - ) - instance.register_buffer( - "occupancy_mask", torch.ones(n_atoms, dtype=torch.bool, device=device) - ) - - # Register vdw_radii if present - if "vdw_radii" in state_dict and state_dict["vdw_radii"] is not None: - instance.register_buffer( - "vdw_radii", - torch.zeros_like(state_dict["vdw_radii"], device=device), - ) # Scattering buffers: accept both old-style (A, B) and new (_A, _B). a_key = "_A" if "_A" in state_dict else "A" if "A" in state_dict else None @@ -1133,13 +980,17 @@ def create_from_state_dict( "_B", torch.zeros_like(state_dict[b_key], device=device) ) - if gridsize is not None and cell_tensor is not None: - if isinstance(gridsize, torch.Tensor): - gs_tuple = tuple(int(x) for x in gridsize.tolist()) - else: - gs_tuple = tuple(int(x) for x in gridsize) - - instance.setup_grid(gridsize=gs_tuple) + if ( + legacy_gridsize is not None + and explicit_gridsize is None + and instance.ctx.crystal_key is not None + and instance.max_res is not None + ): + if isinstance(legacy_gridsize, torch.Tensor): + legacy_gridsize = legacy_gridsize.tolist() + legacy = tuple(int(x) for x in legacy_gridsize) + if legacy != instance.fft.compute_optimal_gridsize(instance.max_res): + instance.explicit_gridsize = legacy # Drop empty placeholders, remapping old-style A/B keys to _A/_B. filtered_state_dict = {} @@ -1154,6 +1005,9 @@ def create_from_state_dict( instance.load_state_dict(filtered_state_dict, strict=False) + # Always placed: target_device is the caller's device or the configured default. + instance.to(target_device) + instance.reset_cache() if verbose > 0: diff --git a/torchref/model/parameter_wrappers.py b/torchref/model/parameter_wrappers.py index f2beddc4..1991cabd 100644 --- a/torchref/model/parameter_wrappers.py +++ b/torchref/model/parameter_wrappers.py @@ -299,12 +299,24 @@ def __setitem__(self, key, value) -> None: self._set_values(key, value) + @property + def _storage_rows(self) -> int: + """Rows of the stored tensor. Equal to ``shape[0]`` unless a subclass derives + rows it does not store, in which case masks handed to the mutation methods + below are in storage space.""" + return 0 if self.fixed_values is None else int(self.fixed_values.shape[0]) + + def _storage_values(self) -> torch.Tensor: + """The stored rows assembled, in public units. Equal to ``forward()`` unless a + subclass derives extra rows.""" + return self.forward() + def _set_values(self, key, value: torch.Tensor) -> None: """Write already-cast values into the storage; override to re-encode. Rebuilds ``fixed_values`` and re-extracts ``refinable_params``. """ - current_full = self.forward().detach() + current_full = self._storage_values().detach() current_full[key] = value self.fixed_values = current_full.clone() @@ -445,13 +457,13 @@ def update_refinable_mask( If True, also re-baseline ``fixed_values`` to the current values. Default is False. """ - if new_mask.shape[0] != self.shape[0]: + if new_mask.shape[0] != self._storage_rows: raise ValueError( f"new_mask shape {new_mask.shape} must match " f"tensor shape {self.shape}" ) - current_full = self.forward().detach() + current_full = self._storage_values().detach() new_mask = self._normalize_refinable_mask(new_mask) self.refinable_mask = new_mask @@ -539,7 +551,7 @@ def refine( If True, re-baseline ``fixed_values`` to the current values first. Default is False. """ - current_full = self.forward().detach() + current_full = self._storage_values().detach() # Union of the current refinable mask with the new selection. new_mask = self.refinable_mask.clone() @@ -547,7 +559,7 @@ def refine( if isinstance(selection, torch.Tensor): if selection.dtype == torch.bool: if len(self.shape) > 1: - if selection.shape[0] != self.shape[0] or len(selection.shape) != 1: + if selection.shape[0] != self._storage_rows or len(selection.shape) != 1: raise ValueError( f"Boolean selection shape {selection.shape} must be 1D " f"matching first dimension {self.shape[0]} for multi-dimensional " @@ -600,7 +612,7 @@ def fix( If True (default), freeze at the current values; if False, the selected elements revert to the stored ``fixed_values``. """ - current_full = self.forward().detach() + current_full = self._storage_values().detach() # Current refinable mask minus the selection. new_mask = self.refinable_mask.clone() @@ -608,7 +620,7 @@ def fix( if isinstance(selection, torch.Tensor): if selection.dtype == torch.bool: if len(self.shape) > 1: - if selection.shape[0] != self.shape[0] or len(selection.shape) != 1: + if selection.shape[0] != self._storage_rows or len(selection.shape) != 1: raise ValueError( f"Boolean selection shape {selection.shape} must be 1D " f"matching first dimension {self.shape[0]} for multi-dimensional " @@ -859,7 +871,7 @@ def set(self, values: torch.Tensor, mask: torch.Tensor) -> None: ValueError If the shapes disagree or any value is non-positive. """ - if mask.shape[0] != self.shape[0]: + if mask.shape[0] != self._storage_rows: raise ValueError( f"Mask shape {mask.shape} must match tensor's first dimension {self.shape[0]}" ) @@ -972,6 +984,137 @@ def __str__(self) -> str: ) +# ---------------------------------------------------------------------------------- +# U <-> Cholesky transforms. Free functions because two unrelated holders need them: +# CholeskyMixedTensor for per-atom ADPs, and the node-field anisotropic payload for +# per-node ones. Both operate on (..., 6) tensors and pass NaN rows through untouched. +# ---------------------------------------------------------------------------------- + + +def u6_to_matrix(U: torch.Tensor) -> torch.Tensor: + """``(..., 6)`` U components to a symmetric ``(..., 3, 3)`` matrix.""" + M = U.new_zeros(*U.shape[:-1], 3, 3) + M[..., 0, 0] = U[..., 0] + M[..., 1, 1] = U[..., 1] + M[..., 2, 2] = U[..., 2] + M[..., 0, 1] = M[..., 1, 0] = U[..., 3] + M[..., 0, 2] = M[..., 2, 0] = U[..., 4] + M[..., 1, 2] = M[..., 2, 1] = U[..., 5] + return M + + +def raw6_to_u6(raw: torch.Tensor, epsilon: float) -> torch.Tensor: + """Cholesky free parameters to U components, ``U = L L^T``. + + Positive-definite for any input: the diagonal of ``L`` is ``exp(x) + epsilon``, so + ``epsilon`` bounds the smallest eigenvalue of ``U`` from below. No factorisation + happens here, which is what makes this safe to call in a forward pass. + """ + diag, off = raw[..., :3], raw[..., 3:] + L11 = torch.exp(diag[..., 0]) + epsilon + L22 = torch.exp(diag[..., 1]) + epsilon + L33 = torch.exp(diag[..., 2]) + epsilon + L21, L31, L32 = off[..., 0], off[..., 1], off[..., 2] + return torch.stack( + [ + L11 * L11, + L21 * L21 + L22 * L22, + L31 * L31 + L32 * L32 + L33 * L33, + L21 * L11, + L31 * L11, + L31 * L21 + L32 * L22, + ], + dim=-1, + ) + + +def u6_to_raw6(U: torch.Tensor, epsilon: float) -> torch.Tensor: + """U components to Cholesky free parameters, projecting onto positive-definite. + + A least-squares or deposited U need not be PD, so the matrix is symmetrised and its + eigenvalues clamped before factorising. Runs at construction and on mask changes, + never in a forward pass. Forced onto the CPU: cuSolver's batched kernels fail on + the large degenerate batches an isotropic model produces, while LAPACK handles them. + """ + finite = torch.isfinite(U).all(dim=-1) + M = u6_to_matrix(torch.nan_to_num(U, nan=0.0)) + eye = torch.eye(3, dtype=M.dtype, device=M.device).expand_as(M) + M = torch.where(finite[..., None, None], M, eye) + M = 0.5 * (M + M.transpose(-1, -2)) + + src_device = M.device + M = M.cpu() + w, V = torch.linalg.eigh(M) + w = w.clamp(min=epsilon * epsilon) + M = (V * w.unsqueeze(-2)) @ V.transpose(-1, -2) + L = torch.linalg.cholesky(M) + diag = torch.stack([L[..., 0, 0], L[..., 1, 1], L[..., 2, 2]], dim=-1) + off = torch.stack([L[..., 1, 0], L[..., 2, 0], L[..., 2, 1]], dim=-1) + raw_diag = torch.log((diag - epsilon).clamp(min=1e-12)) + raw = torch.cat([raw_diag, off], dim=-1).to(src_device) + return torch.where(finite.unsqueeze(-1), raw, torch.full_like(raw, float("nan"))) + + +# ---------------------------------------------------------------------------------- +# The same transform at arbitrary size, for a covariance that is not a 3x3 U tensor. +# A node of the disorder field carries the covariance of its displacement modes, which +# is q x q for q modes; the pair above is the q = 3 case with the indexing unrolled. +# Kept as the general form rather than replacing the unrolled pair, which is a forward +# hot path. +# ---------------------------------------------------------------------------------- + + +def chol_param_count(q: int) -> int: + """Free parameters in a ``q x q`` lower-triangular factor.""" + return q * (q + 1) // 2 + + +def raw_to_cholesky(raw: torch.Tensor, q: int, epsilon: float) -> torch.Tensor: + """Free parameters to a lower-triangular ``(..., q, q)`` factor. + + Layout is ``[log diagonal (q) | strict lower triangle (q(q-1)/2), row major]``, and + the diagonal is ``exp(x) + epsilon``, so ``L L^T`` is positive-definite for any + input and ``epsilon`` bounds its smallest eigenvalue from below. At ``q = 3`` this + is the same layout and the same convention as :func:`raw6_to_u6`. + + No factorisation happens here, which is what makes it safe in a forward pass. + """ + rows, cols = torch.tril_indices(q, q, offset=-1, device=raw.device) + L = raw.new_zeros(*raw.shape[:-1], q, q) + diag = torch.exp(raw[..., :q]) + epsilon + idx = torch.arange(q, device=raw.device) + L[..., idx, idx] = diag + if rows.numel(): + L[..., rows, cols] = raw[..., q:] + return L + + +def psd_to_raw(M: torch.Tensor, epsilon: float) -> torch.Tensor: + """Symmetric ``(..., q, q)`` matrix to Cholesky free parameters, projecting onto PSD. + + The inverse of :func:`raw_to_cholesky`, with the same eigenvalue clamp and the same + CPU-forced ``eigh`` as :func:`u6_to_raw6`: a least-squares or seeded covariance need + not be positive-definite, and cuSolver's batched kernels fail on the degenerate + batches a near-isotropic model produces. Runs at construction, never in a forward + pass. + """ + q = M.shape[-1] + src_device = M.device + M = 0.5 * (M + M.transpose(-1, -2)) + M = M.cpu() + w, V = torch.linalg.eigh(M) + w = w.clamp(min=epsilon * epsilon) + M = (V * w.unsqueeze(-2)) @ V.transpose(-1, -2) + L = torch.linalg.cholesky(M) + idx = torch.arange(q) + rows, cols = torch.tril_indices(q, q, offset=-1) + raw_diag = torch.log((L[..., idx, idx] - epsilon).clamp(min=1e-12)) + parts = [raw_diag] + if rows.numel(): + parts.append(L[..., rows, cols]) + return torch.cat(parts, dim=-1).to(src_device) + + class CholeskyMixedTensor(MixedTensor): """A MixedTensor for anisotropic ADPs (U tensors) kept positive-definite. @@ -1029,57 +1172,16 @@ def __init__( # ------------------------------------------------------------------ @staticmethod def _u6_to_matrix(U: torch.Tensor) -> torch.Tensor: - M = U.new_zeros(*U.shape[:-1], 3, 3) - M[..., 0, 0] = U[..., 0] - M[..., 1, 1] = U[..., 1] - M[..., 2, 2] = U[..., 2] - M[..., 0, 1] = M[..., 1, 0] = U[..., 3] - M[..., 0, 2] = M[..., 2, 0] = U[..., 4] - M[..., 1, 2] = M[..., 2, 1] = U[..., 5] - return M + """Delegate to :func:`u6_to_matrix`.""" + return u6_to_matrix(U) def _u6_to_raw6(self, U: torch.Tensor) -> torch.Tensor: - """U components -> Cholesky free parameters [log(L_ii - eps); L_offdiag].""" - eps = self.epsilon - finite = torch.isfinite(U).all(dim=-1) - M = self._u6_to_matrix(torch.nan_to_num(U, nan=0.0)) - eye = torch.eye(3, dtype=M.dtype, device=M.device).expand_as(M) - M = torch.where(finite[..., None, None], M, eye) - # Project to positive-definite: symmetrise, clamp eigenvalues off zero. - # No-op for well-conditioned deposited U; rescues marginally non-PD input. - M = 0.5 * (M + M.transpose(-1, -2)) - # eigh + Cholesky forced onto the CPU: cuSolver's *batched* kernels fail - # (CUSOLVER_STATUS_INVALID_VALUE) on the large degenerate batches an - # isotropic ensemble produces (U ≡ 0), while LAPACK handles them. This - # runs only at load / mask change, never per optimizer step. - src_device = M.device - M = M.cpu() - w, V = torch.linalg.eigh(M) - w = w.clamp(min=eps * eps) - M = (V * w.unsqueeze(-2)) @ V.transpose(-1, -2) - L = torch.linalg.cholesky(M) - diag = torch.stack([L[..., 0, 0], L[..., 1, 1], L[..., 2, 2]], dim=-1) - off = torch.stack([L[..., 1, 0], L[..., 2, 0], L[..., 2, 1]], dim=-1) - raw_diag = torch.log((diag - eps).clamp(min=1e-12)) # invert exp(x)+eps - raw = torch.cat([raw_diag, off], dim=-1).to(src_device) - nan = torch.full_like(raw, float("nan")) - return torch.where(finite.unsqueeze(-1), raw, nan) + """U components -> Cholesky free parameters. See :func:`u6_to_raw6`.""" + return u6_to_raw6(U, self.epsilon) def _raw6_to_u6(self, raw: torch.Tensor) -> torch.Tensor: - """Cholesky free parameters -> U components (U = L Lᵀ). PD by construction.""" - eps = self.epsilon - diag, off = raw[..., :3], raw[..., 3:] - L11 = torch.exp(diag[..., 0]) + eps - L22 = torch.exp(diag[..., 1]) + eps - L33 = torch.exp(diag[..., 2]) + eps - L21, L31, L32 = off[..., 0], off[..., 1], off[..., 2] - U11 = L11 * L11 - U22 = L21 * L21 + L22 * L22 - U33 = L31 * L31 + L32 * L32 + L33 * L33 - U12 = L21 * L11 - U13 = L31 * L11 - U23 = L31 * L21 + L32 * L22 - return torch.stack([U11, U22, U33, U12, U13, U23], dim=-1) + """Cholesky free parameters -> U components. See :func:`raw6_to_u6`.""" + return raw6_to_u6(raw, self.epsilon) def forward(self) -> torch.Tensor: """Return the full U tensor (positive-definite per finite row).""" @@ -1087,7 +1189,7 @@ def forward(self) -> torch.Tensor: def _set_values(self, key, value: torch.Tensor) -> None: """Set U-space values at ``key``; stored internally as Cholesky params.""" - current = self.forward().detach() + current = self._storage_values().detach() current[key] = value raw = self._u6_to_raw6(current) self.fixed_values = raw.clone() @@ -1101,14 +1203,14 @@ def fix(self, mask: torch.Tensor, freeze_at_current: bool = True): """Freeze rows, storing their current value in Cholesky space.""" if freeze_at_current: with torch.no_grad(): - raw = self._u6_to_raw6(self.forward()) + raw = self._u6_to_raw6(self._storage_values()) self.fixed_values[mask] = raw[mask] super().fix(mask, freeze_at_current=False) def refine(self, mask: torch.Tensor): """Make rows refinable, preserving their current value in Cholesky space.""" with torch.no_grad(): - raw = self._u6_to_raw6(self.forward()) + raw = self._u6_to_raw6(self._storage_values()) self.fixed_values[mask] = raw[mask] super().refine(mask) @@ -1126,12 +1228,12 @@ def update_refinable_mask( storage); convert to Cholesky parameters first, mirroring :meth:`PositiveMixedTensor.update_refinable_mask`. """ - if new_mask.shape[0] != self.shape[0]: + if new_mask.shape[0] != self._storage_rows: raise ValueError( f"new_mask shape {new_mask.shape} must match tensor shape {self.shape}" ) with torch.no_grad(): - current_raw = self._u6_to_raw6(self.forward()) + current_raw = self._u6_to_raw6(self._storage_values()) new_mask = self._normalize_refinable_mask(new_mask) self.refinable_mask = new_mask self.fixed_mask = ~new_mask @@ -1372,11 +1474,11 @@ def _setup_sharing_groups_and_expansion( # Use sharing_groups directly as the expansion mask if sharing_groups is None: # No sharing - each atom maps to its own index - expansion_mask = torch.arange(n_atoms, dtype=torch.long, device=device) + expansion_mask = torch.arange(n_atoms, dtype=torch.long, device=device) # dtype-ok: arange expansion_mask atom indices; index requires long self._collapsed_shape = n_atoms else: # Use the provided index tensor - expansion_mask = sharing_groups.to(device=device, dtype=torch.long) + expansion_mask = sharing_groups.to(device=device, dtype=torch.long) # dtype-ok: expansion_mask atom/group indices for scatter; requires long self._collapsed_shape = expansion_mask.max().item() + 1 self.register_buffer("expansion_mask", expansion_mask) @@ -1398,10 +1500,10 @@ def _setup_sharing_groups_and_expansion( for conf_atoms in conf_groups: if isinstance(conf_atoms, (list, tuple)): conf_atoms = torch.tensor( - conf_atoms, dtype=torch.long, device=device + conf_atoms, dtype=torch.long, device=device # dtype-ok: conf_atoms atom indices; indexing requires long ) else: - conf_atoms = conf_atoms.to(device=device, dtype=torch.long) + conf_atoms = conf_atoms.to(device=device, dtype=torch.long) # dtype-ok: conf_atoms atom indices cast; indexing requires long # Get collapsed index for first atom collapsed_idx = expansion_mask[conf_atoms[0]].item() @@ -1429,7 +1531,7 @@ def _setup_sharing_groups_and_expansion( # Store as dictionary with keys like 'linked_occ_2', 'linked_occ_3', etc. for n_conf, groups in linked_occupancies.items(): # Shape: (N_groups, n_conf) - tensor = torch.tensor(groups, dtype=torch.long, device=device) + tensor = torch.tensor(groups, dtype=torch.long, device=device) # dtype-ok: linked-occupancy group index buffer; indexing requires long self.register_buffer(f"linked_occ_{n_conf}", tensor) # Store which sizes we have @@ -1437,7 +1539,7 @@ def _setup_sharing_groups_and_expansion( # Create count buffer for vectorized collapse operations # counts[i] = number of atoms that map to collapsed index i - counts = torch.zeros(self._collapsed_shape, dtype=torch.long, device=device) + counts = torch.zeros(self._collapsed_shape, dtype=torch.long, device=device) # dtype-ok: count accumulator; scatter_add source is long ones, dtype must match counts.scatter_add_(0, expansion_mask, torch.ones_like(expansion_mask)) self.register_buffer("collapse_counts", counts) @@ -1915,7 +2017,7 @@ def from_residue_groups( grouped = pdb_dataframe.groupby(["resname", "resseq", "chainid", "altloc"]) n_atoms = len(initial_values) - sharing_groups_tensor = torch.arange(n_atoms, dtype=torch.long) + sharing_groups_tensor = torch.arange(n_atoms, dtype=torch.long) # dtype-ok: arange atom indices (sharing groups); index requires long # Singletons keep their arange ids (0..n_atoms-1); start multi-atom # group ids past that range so a group id can never collide with a # singleton's leftover arange id (the torch.unique compaction below diff --git a/torchref/model/riding_xyz.py b/torchref/model/riding_xyz.py new file mode 100644 index 00000000..12771a67 --- /dev/null +++ b/torchref/model/riding_xyz.py @@ -0,0 +1,816 @@ +"""Coordinate wrapper whose hydrogen rows ride on their parents. + +A :class:`RidingXYZTensor` looks like a :class:`~torchref.model.parameter_wrappers.MixedTensor` +over the whole atom table -- ``forward()`` returns ``(N, 3)``, masks are given in atom +space -- but only the non-riding rows are stored and refined. Each riding hydrogen is a +reference offset in a frame built from its parent and two reference heavy atoms +(:mod:`torchref.base.coordinates.local_frame`), rebuilt from the current heavy +coordinates on every forward. A force on a hydrogen therefore lands on the atoms that +carry it, which is the riding-hydrogen convention of Phenix and Refmac. + +Two index spaces meet here and callers must not mix them: everything public -- masks +passed to ``update_refinable_mask`` / ``refine`` / ``fix`` / ``set``, the result of +``forward()``, ``__getitem__`` / ``__setitem__`` -- is in FULL atom space, while the +inherited ``refinable_mask`` / ``fixed_values`` / ``refinable_params`` and the counts +from ``get_refinable_count()`` are in STORAGE space (the non-riding rows). Independent +angles are held in ``torsions`` and ``rotations`` and exposed together with the stored +coordinates by ``optimization_parameters()``. This is the +same contract :class:`~torchref.model.parameter_wrappers.OccupancyTensor` uses for its +collapsed groups. +""" + +from typing import Iterator, Optional, Union + +import numpy as np +import torch +import torch.nn as nn +from torchref.base.coordinates.local_frame import ( + frame_is_degenerate, + local_frame_coordinates, + place_local_frame, + rotate_vectors, +) +from torchref.model.parameter_wrappers import MixedTensor +from torchref.topology.hydrogens import HydrogenFrames + + +class _DerivedRowsMixin: + """Bookkeeping shared by wrappers that store some rows and derive the rest. + + Registers the row maps as buffers and maintains the storage-space gathers the + derivation needs. Subclasses call :meth:`_register_rows` after the parent + constructor has run and :meth:`_rebuild_row_cache` from ``_build_index_cache``. + """ + + #: Full-space row of every stored row, ``(N_base,)``. + base_row: torch.Tensor + #: Full-space row of every derived (riding) row, ``(H,)``. + h_row: torch.Tensor + #: Frame atoms per riding row, full-space, ``(H,)`` each. Absent -> clamped to 0 + #: and masked through ``frame_valid``. + parent_row: torch.Tensor + n1_row: torch.Tensor + n2_row: torch.Tensor + frame_valid: torch.Tensor + + def _register_rows(self, n_full: int, frames: HydrogenFrames, device) -> None: + h = np.asarray(frames.h_row, dtype=np.int64) + if len(h) and (h.min() < 0 or h.max() >= n_full): + raise ValueError("frames carry a hydrogen row outside the atom table") + is_riding = np.zeros(n_full, dtype=bool) + is_riding[h] = True + base = np.nonzero(~is_riding)[0] + for name in ("parent_row", "n1_row", "n2_row"): + rows = np.asarray(getattr(frames, name), dtype=np.int64) + if len(rows) and is_riding[rows[rows >= 0]].any(): + raise ValueError(f"{name} must reference stored rows, not riding ones") + + long = dict( + dtype=torch.int64, device=device + ) # dtype-ok: row index buffers; int64 index required + self.register_buffer("base_row", torch.as_tensor(base, **long)) + self.register_buffer("h_row", torch.as_tensor(h, **long)) + self.register_buffer( + "parent_row", + torch.as_tensor(np.asarray(frames.parent_row, dtype=np.int64), **long), + ) + self.register_buffer( + "n1_row", torch.as_tensor(np.asarray(frames.n1_row, dtype=np.int64), **long) + ) + self.register_buffer( + "n2_row", torch.as_tensor(np.asarray(frames.n2_row, dtype=np.int64), **long) + ) + self.register_buffer( + "frame_valid", + torch.as_tensor(np.asarray(frames.frame_valid, dtype=bool), device=device), + ) + self._rebuild_row_cache() + + def _rebuild_row_cache(self) -> None: + """Derive the storage-space gathers from the row buffers.""" + base = getattr(self, "base_row", None) + if base is None or getattr(self, "h_row", None) is None: + self._n_full = 0 + return + device = base.device + n_full = int(base.numel() + self.h_row.numel()) + self._n_full = n_full + full_to_base = torch.full( + (max(n_full, 1),), -1, dtype=torch.int64, device=device + ) # dtype-ok: index map; int64 + full_to_base[base] = torch.arange( + base.numel(), dtype=torch.int64, device=device + ) # dtype-ok: index map; int64 + self._parent_bidx = full_to_base[self.parent_row.clamp(min=0)].clamp(min=0) + self._n1_bidx = full_to_base[self.n1_row.clamp(min=0)].clamp(min=0) + self._n2_bidx = full_to_base[self.n2_row.clamp(min=0)].clamp(min=0) + # ``cat([base, derived])[gather]`` lays the full table out in one gather. + order = torch.empty( + n_full, dtype=torch.int64, device=device + ) # dtype-ok: gather index; int64 + order[base] = torch.arange( + base.numel(), dtype=torch.int64, device=device + ) # dtype-ok: gather index; int64 + order[self.h_row] = base.numel() + torch.arange( + self.h_row.numel(), + dtype=torch.int64, + device=device, # dtype-ok: gather index; int64 + ) + self._gather_order = order + + @property + def n_hydrogens(self) -> int: + """How many rows ride.""" + return 0 if getattr(self, "h_row", None) is None else int(self.h_row.numel()) + + @property + def n_base(self) -> int: + """How many rows are stored.""" + return ( + 0 if getattr(self, "base_row", None) is None else int(self.base_row.numel()) + ) + + def hydrogen_frames(self) -> HydrogenFrames: + """The frames as a CPU record, full-space rows.""" + return HydrogenFrames.from_tensors( + self.h_row, + self.parent_row, + self.n1_row, + self.n2_row, + self.frame_valid, + self.torsion_group, + self.rotation_group, + ) + + def _to_full_bool(self, selection) -> torch.Tensor: + """Any selection (bool mask, slice, indices) as a full-space bool mask.""" + if isinstance(selection, torch.Tensor) and selection.dtype == torch.bool: + if selection.ndim != 1 or selection.shape[0] != self._n_full: + raise ValueError( + f"Boolean selection shape {tuple(selection.shape)} must be " + f"({self._n_full},)" + ) + return selection.to(device=self.base_row.device) + mask = torch.zeros(self._n_full, dtype=torch.bool, device=self.base_row.device) + mask[selection] = True + return mask + + def _project(self, full_mask: torch.Tensor) -> torch.Tensor: + """Storage-space view of a full-space bool mask (riding rows dropped).""" + return full_mask.to(device=self.base_row.device, dtype=torch.bool)[ + self.base_row + ] + + def _expand_mask(self, base_mask: torch.Tensor) -> torch.Tensor: + """Full-space bool mask from a storage-space one; riding rows False.""" + out = torch.zeros(self._n_full, dtype=torch.bool, device=base_mask.device) + out[self.base_row] = base_mask + return out + + +class RidingXYZTensor(_DerivedRowsMixin, MixedTensor): + """Coordinates with riding hydrogen rows derived from their parents. + + Parameters + ---------- + initial_values : torch.Tensor, optional + Full atom table, Cartesian Angstroms, shape ``(N, 3)``. The riding rows fix the + local offsets; every other row is stored. None gives the empty shell that + ``load_state_dict`` fills. + frames : HydrogenFrames, optional + Which rows ride and on which atoms; required with ``initial_values``. + refinable_mask : torch.Tensor, optional + Boolean ``(N,)`` in full atom space, or + ``(N_base,)`` in storage space with ``mask_in_base_space``. None: every stored + row and every orientation refinable. + mask_in_base_space : bool, default False + Whether ``refinable_mask`` is already in storage space, as a saved state dict + hands it back. + requires_grad, dtype, device, name + As for :class:`~torchref.model.parameter_wrappers.MixedTensor`. + eps : float + Norm floor of the frame kernel, Angstroms. + + Notes + ----- + ``shape`` is the full ``(N, 3)``; ``refinable_mask``, ``fixed_values`` and + ``refinable_params`` are storage-space (``N_base`` rows). Frames whose reference + atoms are collinear or missing fall back to a Cartesian ``parent + offset``. + ``torsions`` holds one angle in radians per freely rotating bonded group; + ``rotations`` holds one Cartesian rotation vector in radians per unanchored + group (for example water). Both are :class:`MixedTensor` submodules. Selecting + a parent or any member hydrogen selects that group's orientation. Coordinate + assignment adopts the supplied orientation as the reference and zeros angles; + copies and checkpoints preserve the parameter values and references. + """ + + def __init__( + self, + initial_values: Optional[torch.Tensor] = None, + frames: Optional[HydrogenFrames] = None, + refinable_mask: Optional[torch.Tensor] = None, + *, + mask_in_base_space: bool = False, + requires_grad: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + name: Optional[str] = "xyz", + eps: float = 1e-8, + ): + self._eps = float(eps) + if initial_values is None: + super().__init__( + None, requires_grad=requires_grad, dtype=dtype, device=device, name=name + ) + for buffer in ("base_row", "h_row", "parent_row", "n1_row", "n2_row"): + self.register_buffer( + buffer, + torch.zeros( + 0, dtype=torch.int64, device=self.device + ), # dtype-ok: empty row-index buffer; int64 + ) + self.register_buffer( + "frame_valid", torch.zeros(0, dtype=torch.bool, device=self.device) + ) + self.register_buffer( + "local_offset", torch.zeros(0, 3, dtype=self.dtype, device=self.device) + ) + self.register_buffer( + "rigid_offset", torch.zeros(0, 3, dtype=self.dtype, device=self.device) + ) + self.register_buffer( + "virtual_reference", + torch.zeros(0, 3, dtype=self.dtype, device=self.device), + ) + self._initialize_orientations(HydrogenFrames.empty(), None) + self.register_load_state_dict_post_hook(self._after_load) + self._build_index_cache() + return + + if frames is None: + raise ValueError("frames are required with initial_values") + if initial_values.ndim != 2 or initial_values.shape[1] != 3: + raise ValueError( + f"initial_values must be (N, 3), got {tuple(initial_values.shape)}" + ) + dtype = dtype if dtype is not None else initial_values.dtype + device = device if device is not None else initial_values.device + values = initial_values.detach().to(dtype=dtype, device=device) + n_full = values.shape[0] + + is_riding = np.zeros(n_full, dtype=bool) + is_riding[np.asarray(frames.h_row, dtype=np.int64)] = True + base_rows = torch.as_tensor( + np.nonzero(~is_riding)[0], dtype=torch.int64, device=device + ) # dtype-ok: row index; int64 + + if refinable_mask is None: + base_mask = None + elif mask_in_base_space: + base_mask = refinable_mask.to(device=device, dtype=torch.bool) + else: + if refinable_mask.shape[0] != n_full: + raise ValueError( + f"refinable_mask has {refinable_mask.shape[0]} rows, table has {n_full}" + ) + base_mask = refinable_mask.to(device=device, dtype=torch.bool)[base_rows] + + super().__init__( + values.index_select(0, base_rows), + base_mask, + requires_grad=requires_grad, + dtype=dtype, + device=device, + name=name, + ) + self._register_rows(n_full, frames, device) + self.register_buffer( + "local_offset", torch.zeros(self.n_hydrogens, 3, dtype=dtype, device=device) + ) + self.register_buffer( + "rigid_offset", torch.zeros(self.n_hydrogens, 3, dtype=dtype, device=device) + ) + self.register_buffer("virtual_reference", torch.zeros_like(self.local_offset)) + full_mask = None if refinable_mask is None else refinable_mask.to(device=device) + if full_mask is not None and mask_in_base_space: + full_mask = self._expand_mask(full_mask) + self._initialize_orientations(frames, full_mask) + self.refresh_offsets(values) + self.register_load_state_dict_post_hook(self._after_load) + self._build_index_cache() + + # ------------------------------------------------------------------ + # Assembly + # ------------------------------------------------------------------ + + def _build_index_cache(self): + super()._build_index_cache() + self._rebuild_row_cache() + if hasattr(self, "torsion_group"): + self._rebuild_orientation_cache() + + def _initialize_orientations(self, frames, full_mask): + for name in ("torsion_group", "rotation_group"): + labels = np.asarray(getattr(frames, name)) + selected = labels >= 0 + unique, inverse = np.unique(labels[selected], return_inverse=True) + compact = np.full(len(labels), -1, dtype=np.int64) + compact[selected] = inverse + for group in unique: + members = np.flatnonzero(labels == group) + if len(np.unique(frames.parent_row[members])) != 1: + raise ValueError("An orientation group must share one parent") + if name == "torsion_group" and ( + (frames.n1_row[members] < 0).any() + or len(np.unique(frames.n1_row[members])) != 1 + ): + raise ValueError("A torsion group must share one bonded axis") + self.register_buffer(name, torch.as_tensor(compact, device=self.device)) + self._rebuild_orientation_cache() + requires_grad = self.refinable_params.requires_grad + self.torsions = MixedTensor( + torch.zeros( + self._torsion_parents.numel(), dtype=self.dtype, device=self.device + ), + requires_grad=requires_grad, + name="hydrogen_torsions", + ) + self.rotations = MixedTensor( + torch.zeros( + self._rotation_parents.numel(), 3, dtype=self.dtype, device=self.device + ), + requires_grad=requires_grad, + name="hydrogen_rotations", + ) + if full_mask is not None: + for wrapper, mask in zip( + (self.torsions, self.rotations), self._orientation_selection(full_mask) + ): + wrapper.update_refinable_mask(mask) + + def _rebuild_orientation_cache(self): + self._virtual_frame = (self.torsion_group >= 0) & (self.n2_row < 0) + self._has_virtual_frames = bool(self._virtual_frame.any()) + for kind in ("torsion", "rotation"): + labels = getattr(self, kind + "_group") + rows = (labels >= 0).nonzero(as_tuple=True)[0] + groups = labels[rows] + if rows.numel(): + order = torch.argsort(groups, stable=True) + sorted_groups = groups[order] + first = torch.cat( + [ + torch.ones(1, dtype=torch.bool, device=groups.device), + sorted_groups[1:] != sorted_groups[:-1], + ] + ) + first_rows = rows[order[first]] + parents = self.parent_row[first_rows] + else: + first_rows = rows + parents = self.parent_row[:0] + setattr(self, "_" + kind + "_h", rows) + setattr(self, "_" + kind + "_inverse", groups) + setattr(self, "_" + kind + "_parents", parents) + setattr(self, "_" + kind + "_first", first_rows) + + def _orientation_selection(self, full_mask): + selections = [] + for kind in ("torsion", "rotation"): + parents = getattr(self, "_" + kind + "_parents") + rows = getattr(self, "_" + kind + "_h") + groups = getattr(self, "_" + kind + "_inverse") + selected = full_mask[parents].to(torch.int32) + selected.index_add_(0, groups, full_mask[self.h_row[rows]].to(torch.int32)) + selections.append(selected > 0) + return selections + + def parameters(self, recurse: bool = True) -> Iterator[nn.Parameter]: + """Yield stored-coordinate and orientation leaves, including frozen shells.""" + return nn.Module.parameters(self, recurse=recurse) + + def optimization_parameters(self) -> list[nn.Parameter]: + """Return coordinate, torsion and rotation leaves for the xyz optimizer.""" + return [ + self.refinable_params, + self.torsions.refinable_params, + self.rotations.refinable_params, + ] + + @property + def _storage_rows(self) -> int: + return 0 if self.fixed_values is None else int(self.fixed_values.shape[0]) + + def _storage_values(self) -> torch.Tensor: + """The stored rows assembled, ``(N_base, 3)``.""" + return MixedTensor.forward(self) + + def evaluate( + self, + base_xyz: torch.Tensor, + torsions: Optional[torch.Tensor] = None, + rotations: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full coordinates from stored ones: the pure, differentiable part of forward. + + Parameters + ---------- + base_xyz : torch.Tensor + Stored Cartesian coordinates in Å, shape ``(N_base, 3)``. + torsions : torch.Tensor, optional + Full group angles in radians, shape ``(n_torsions,)``. Defaults to + the stored torsion parameters. + rotations : torch.Tensor, optional + Full group rotation vectors in radians, shape ``(n_rotations, 3)``. + Defaults to the stored orientation parameters. + + Returns + ------- + torch.Tensor + Shape ``(N, 3)``; riding rows placed from their frames. Gradients on any + row reach ``base_xyz`` through the frame Jacobian. + """ + if self.n_hydrogens == 0: + return base_xyz + local = self.local_offset + if self._torsion_h.numel(): + angles = self.torsions.forward() if torsions is None else torsions + cs = torch.stack((angles.cos(), angles.sin()), dim=-1)[ + self._torsion_inverse + ] + offsets = local[self._torsion_h] + x, y, z = offsets.unbind(-1) + c, sn = cs.unbind(-1) + turned = torch.stack((x, c * y - sn * z, sn * y + c * z), dim=-1) + local = local.index_copy(0, self._torsion_h, turned) + p = base_xyz.index_select(0, self._parent_bidx) + n2 = base_xyz.index_select(0, self._n2_bidx) + if self._has_virtual_frames: + n2 = torch.where( + self._virtual_frame[:, None], p + self.virtual_reference, n2 + ) + h = place_local_frame( + p, + base_xyz.index_select(0, self._n1_bidx), + n2, + local, + self.frame_valid, + self.rigid_offset, + eps=self._eps, + ) + if self._rotation_h.numel(): + vectors = self.rotations.forward() if rotations is None else rotations + offsets = rotate_vectors( + self.rigid_offset[self._rotation_h], vectors[self._rotation_inverse] + ) + h = h.index_copy(0, self._rotation_h, p[self._rotation_h] + offsets) + return torch.cat([base_xyz, h], dim=0).index_select(0, self._gather_order) + + def forward(self) -> torch.Tensor: + """The full ``(N, 3)`` table, riding rows derived from the stored ones.""" + return self.evaluate(self._storage_values()) + + @property + def shape(self): + """Full-space shape ``(N, 3)``.""" + if self.fixed_values is None: + return () + return (self._n_full, int(self.fixed_values.shape[1])) + + @property + def base_shape(self): + """Storage-space shape ``(N_base, 3)``.""" + return () if self.fixed_values is None else tuple(self.fixed_values.shape) + + @property + def full_refinable_mask(self) -> torch.Tensor: + """Refinable rows in full atom space; riding rows are never refinable.""" + return self._expand_mask(self.refinable_mask) + + # ------------------------------------------------------------------ + # Offsets + # ------------------------------------------------------------------ + + @torch.no_grad() + def refresh_offsets(self, full_xyz: Optional[torch.Tensor] = None) -> None: + """Re-derive every local offset from full-space coordinates. + + Parameters + ---------- + full_xyz : torch.Tensor, optional + Shape ``(N, 3)``; defaults to the current ``forward()``, which leaves the + coordinates unchanged while rebasing the angular parameters. Pass the + table after an external re-placement (a + torsion re-scan, a hydrogen written by ``__setitem__``) to adopt it. + + Notes + ----- + Frames that are geometrically degenerate at these coordinates are demoted to + the rigid fallback. Rewrites buffers in place, so the forward cache is + invalidated automatically. + """ + if self.n_hydrogens == 0: + return + if full_xyz is None: + full_xyz = self.forward() + full_xyz = full_xyz.detach().to(dtype=self.dtype, device=self.device) + p = full_xyz.index_select(0, self.parent_row) + n1 = full_xyz.index_select(0, self.n1_row.clamp(min=0)) + n2 = full_xyz.index_select(0, self.n2_row.clamp(min=0)) + h = full_xyz.index_select(0, self.h_row) + if self._has_virtual_frames: + first = self._torsion_first[self._torsion_inverse] + self.virtual_reference[self._torsion_h] = (h - p)[first] + n2 = torch.where( + self._virtual_frame[:, None], p + self.virtual_reference, n2 + ) + topological = (self.n1_row >= 0) & ((self.n2_row >= 0) | self._virtual_frame) + valid = topological & ~frame_is_degenerate(p, n1, n2) + local = local_frame_coordinates(p, n1, n2, h, eps=self._eps) + self.frame_valid.copy_(valid) + self.local_offset.copy_( + torch.where(valid.unsqueeze(-1), local, torch.zeros_like(local)) + ) + self.rigid_offset.copy_(h - p) + for orientation in (self.torsions, self.rotations): + orientation.refinable_params.zero_() + orientation.fixed_values.zero_() + orientation.reset_forward_cache() + + def set_hydrogen_positions(self, h_xyz: torch.Tensor) -> None: + """Adopt new positions for the riding rows, in ``h_row`` order, ``(H, 3)``.""" + full = self.forward().detach() + full[self.h_row] = h_xyz.to(dtype=self.dtype, device=self.device) + self.refresh_offsets(full) + + # ------------------------------------------------------------------ + # Mutation in full space + # ------------------------------------------------------------------ + + def _set_values(self, key, value: torch.Tensor) -> None: + """Write full-space values; stored rows update, riding rows become new offsets.""" + full = self.forward().detach() + full[key] = value + super()._set_values(slice(None), full.index_select(0, self.base_row)) + self.refresh_offsets(full) + + def set(self, values: torch.Tensor, mask: torch.Tensor) -> None: + """Write ``values`` at the True rows of a full-space ``mask``.""" + if mask.ndim != 1 or mask.shape[0] != self._n_full: + raise ValueError( + f"Mask shape {tuple(mask.shape)} must be ({self._n_full},)" + ) + mask = mask.to(device=self.device, dtype=torch.bool) + n_selected = int(mask.sum().item()) + if tuple(values.shape) != (n_selected, 3): + raise ValueError( + f"Values shape {tuple(values.shape)} doesn't match ({n_selected}, 3)" + ) + self._set_values(mask, values.to(dtype=self.dtype, device=self.device)) + + def update_refinable_mask( + self, new_mask: torch.Tensor, reset_refinable: bool = False + ): + """Repartition coordinates and orientations with a full- or storage-space mask.""" + full_mask = new_mask.to(device=self.device) + if new_mask.shape[0] == self._n_full: + new_mask = self._project(new_mask) + elif new_mask.shape[0] != self._storage_rows: + raise ValueError( + f"new_mask has {new_mask.shape[0]} rows; expected {self._n_full} " + f"(atom space) or {self._storage_rows} (storage space)" + ) + if full_mask.shape[0] != self._n_full: + full_mask = self._expand_mask(full_mask) + super().update_refinable_mask(new_mask, reset_refinable=reset_refinable) + for wrapper, mask in zip( + (self.torsions, self.rotations), self._orientation_selection(full_mask) + ): + wrapper.update_refinable_mask(mask, reset_refinable=reset_refinable) + + def refine( + self, selection: Union[slice, torch.Tensor, tuple], reset_values: bool = False + ): + """Add a full-space selection to the refinable set.""" + full_mask = self._to_full_bool(selection) + super().refine(self._project(full_mask), reset_values) + for wrapper, mask in zip( + (self.torsions, self.rotations), self._orientation_selection(full_mask) + ): + wrapper.refine(mask, reset_values) + + def fix( + self, + selection: Union[slice, torch.Tensor, tuple], + freeze_at_current: bool = True, + ): + """Remove a full-space selection from the refinable set.""" + full_mask = self._to_full_bool(selection) + super().fix(self._project(full_mask), freeze_at_current) + for wrapper, mask in zip( + (self.torsions, self.rotations), self._orientation_selection(full_mask) + ): + wrapper.fix(mask, freeze_at_current) + + def refine_all(self): + """Make every stored row and orientation refinable.""" + self.refine(torch.ones(self._n_full, dtype=torch.bool, device=self.device)) + + def fix_all(self, freeze_at_current: bool = True): + """Fix every stored row and orientation.""" + self.fix( + torch.ones(self._n_full, dtype=torch.bool, device=self.device), + freeze_at_current=freeze_at_current, + ) + + def update_fixed_values(self, new_values: torch.Tensor): + """Replace the stored rows' fixed buffer from a full-space ``(N, 3)`` table.""" + if tuple(new_values.shape) == self.shape: + new_values = new_values.index_select(0, self.base_row.to(new_values.device)) + super().update_fixed_values(new_values) + + # ------------------------------------------------------------------ + # Conversions and copies + # ------------------------------------------------------------------ + + def to_mixed_tensor(self) -> MixedTensor: + """Materialise as a plain per-atom wrapper; hydrogens follow their parent's mask.""" + mask = self.full_refinable_mask.clone() + mask[self.h_row] = mask[self.parent_row] + return MixedTensor( + self.forward().detach(), + mask, + requires_grad=self.refinable_params.requires_grad, + dtype=self.dtype, + device=self.device, + name=self.name, + ) + + @classmethod + def from_mixed_tensor( + cls, xyz: MixedTensor, frames: HydrogenFrames, **kwargs + ) -> "RidingXYZTensor": + """Wrap an existing per-atom coordinate tensor with riding frames.""" + return cls( + xyz.forward().detach(), + frames, + refinable_mask=xyz.refinable_mask, + requires_grad=xyz.refinable_params.requires_grad, + dtype=xyz.dtype, + device=xyz.device, + name=xyz.name, + **kwargs, + ) + + def with_values(self, full_xyz: torch.Tensor) -> "RidingXYZTensor": + """Same frames and mask, new coordinates ``(N, 3)``.""" + result = RidingXYZTensor( + full_xyz, + self.hydrogen_frames(), + refinable_mask=self.refinable_mask.clone(), + mask_in_base_space=True, + requires_grad=self.refinable_params.requires_grad, + dtype=self.dtype, + device=self.device, + name=self.name, + eps=self._eps, + ) + for name in ("torsions", "rotations"): + getattr(result, name).update_refinable_mask( + getattr(self, name).refinable_mask.clone() + ) + return result + + def select_rows(self, keep: torch.Tensor) -> "RidingXYZTensor": + """The wrapper over the rows where ``keep`` is True, frames remapped. + + A hydrogen whose parent is not kept becomes an ordinary stored row. + """ + keep_np = np.asarray(keep.detach().cpu().numpy(), dtype=bool) + old_to_new = np.full(len(keep_np), -1, dtype=np.int64) + old_to_new[keep_np] = np.arange(int(keep_np.sum())) + frames = self.hydrogen_frames().remap(old_to_new) + keep_t = torch.as_tensor(keep_np, device=self.device) + result = RidingXYZTensor( + self.forward().detach()[keep_t], + frames, + refinable_mask=self.full_refinable_mask[keep_t], + requires_grad=self.refinable_params.requires_grad, + dtype=self.dtype, + device=self.device, + name=self.name, + eps=self._eps, + ) + for name, labels in ( + ("torsions", frames.torsion_group), + ("rotations", frames.rotation_group), + ): + retained = torch.as_tensor( + np.unique(labels[labels >= 0]), device=self.device + ) + getattr(result, name).update_refinable_mask( + getattr(self, name).refinable_mask[retained] + ) + return result + + def clone(self) -> "RidingXYZTensor": + """Independent copy, offsets carried over bit for bit.""" + out = self.with_values(self.forward().detach()) + with torch.no_grad(): + out.local_offset.copy_(self.local_offset) + out.rigid_offset.copy_(self.rigid_offset) + out.frame_valid.copy_(self.frame_valid) + out.virtual_reference.copy_(self.virtual_reference) + out.torsions = self.torsions.copy() + out.rotations = self.rotations.copy() + return out + + def copy(self) -> "RidingXYZTensor": + """Alias for :meth:`clone`.""" + return self.clone() + + def clip(self, min_value=None, max_value=None) -> "RidingXYZTensor": + """Clip the full table; riding rows re-derive from the clipped heavy atoms.""" + full = self.forward().detach() + if min_value is not None: + full = torch.clamp(full, min=min_value) + if max_value is not None: + full = torch.clamp(full, max=max_value) + return self.with_values(full) + + def _after_load(self, module, incompatible_keys): + self._build_index_cache() + self.reset_forward_cache() + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + if prefix + "torsion_group" not in state_dict: + # A checkpoint without group metadata describes fixed orientations. + hydrogen_rows = state_dict[prefix + "h_row"] + defaults = self.state_dict() + for name in ("torsion_group", "rotation_group"): + defaults[name] = torch.full_like(hydrogen_rows, -1) + defaults["virtual_reference"] = torch.zeros_like( + state_dict[prefix + "rigid_offset"] + ) + for name in ("torsions", "rotations"): + shape = (0,) if name == "torsions" else (0, 3) + empty = MixedTensor( + torch.empty(shape, dtype=self.dtype, device=self.device) + ) + defaults.update( + { + name + "." + key: value + for key, value in empty.state_dict().items() + } + ) + for name, value in defaults.items(): + if name in ( + "torsion_group", + "rotation_group", + "virtual_reference", + ) or name.startswith(("torsions.", "rotations.")): + state_dict.setdefault(prefix + name, value) + for name, buffer in list(self._buffers.items()): + saved = state_dict.get(prefix + name) + if saved is not None and (buffer is None or saved.shape != buffer.shape): + value = ( + torch.empty_like(saved, device=self.device) + if buffer is None + else buffer.new_empty(saved.shape) + ) + setattr(self, name, value) + saved_params = state_dict.get(prefix + "refinable_params") + if ( + saved_params is not None + and saved_params.shape != self.refinable_params.shape + ): + self.refinable_params = nn.Parameter( + self.refinable_params.new_empty(saved_params.shape), + requires_grad=self.refinable_params.requires_grad, + ) + for name in ("torsions", "rotations"): + saved = state_dict.get(prefix + name + ".fixed_values") + if saved is not None: + mask = state_dict[prefix + name + ".refinable_mask"].to(self.device) + setattr( + self, + name, + MixedTensor( + saved.to(device=self.device, dtype=self.dtype), + mask, + requires_grad=self.refinable_params.requires_grad, + name="hydrogen_" + name, + ), + ) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def __repr__(self) -> str: + name_str = f"'{self.name}', " if self.name is not None else "" + return ( + f"RidingXYZTensor({name_str}shape={self.shape}, dtype={self.dtype}, " + f"device={self.device}, refinable={self.get_refinable_count()}, " + f"fixed={self.get_fixed_count()}, riding_h={self.n_hydrogens})" + ) + + +__all__ = ["RidingXYZTensor"] diff --git a/torchref/model/rigid_xyz.py b/torchref/model/rigid_xyz.py index d74c8fe9..0d0a304a 100644 --- a/torchref/model/rigid_xyz.py +++ b/torchref/model/rigid_xyz.py @@ -9,6 +9,15 @@ centroid and translating it. XYZ Euler matches Phenix's default ``euler_angle_convention`` and keeps the rotation Jacobian full-rank at the origin (no gimbal lock when angles reset to zero after ``bake()``). + +**``euler_angles`` is NOT in radians.** It is stored pre-multiplied by +``angle_scale``, the per-chain radius of gyration in Angstroms, so a unit step in +an angle and a unit step in a translation displace atoms comparably. Without it +the rotation block of the Hessian carries ~``Rg**2`` the curvature of the +translation block and L-BFGS needs an order of magnitude more iterations to place +six parameters. ``forward()`` divides the scale out; +:attr:`RigidXYZTensor.rotation_radians` returns the physical angle. Setting +``angle_scale`` to ones gives the unscaled parametrization. """ from typing import Optional, Sequence @@ -64,13 +73,14 @@ def __init__( dtype = dtype if dtype is not None else get_float_dtype() self.register_buffer("original_xyz", torch.empty(0, 3, device=device, dtype=dtype)) self.register_buffer( - "chain_indices", torch.empty(0, dtype=torch.long, device=device) + "chain_indices", torch.empty(0, dtype=torch.long, device=device) # dtype-ok: empty chain_indices buffer; indexing requires long ) self.register_buffer("chain_centers", torch.empty(0, 3, device=device, dtype=dtype)) self.register_buffer( "mobile_mask", torch.empty(0, dtype=torch.bool, device=device) ) self.register_buffer("atom_weights", torch.empty(0, device=device, dtype=dtype)) + self.register_buffer("angle_scale", torch.empty(0, device=device, dtype=dtype)) self.euler_angles = nn.Parameter(torch.empty(0, 3, device=device, dtype=dtype)) self.translations = nn.Parameter(torch.empty(0, 3, device=device, dtype=dtype)) self._n_chains = 0 @@ -162,6 +172,12 @@ def __init__( self.register_buffer("chain_centers", chain_centers) self.register_buffer("mobile_mask", mobile_t) self.register_buffer("atom_weights", atom_weights_t) + self.register_buffer( + "angle_scale", + self._compute_angle_scale( + original_xyz_t, chain_centers, mobile_idx, mobile_t, n_chains + ), + ) self.euler_angles = nn.Parameter( torch.zeros((n_chains, 3), dtype=dtype, device=device) @@ -173,6 +189,28 @@ def __init__( self._n_chains = n_chains self._chain_id_order = list(chain_id_order) + # ----------------------------------------------------------------------- + # Angle scaling (preconditioning) + # ----------------------------------------------------------------------- + @staticmethod + def _compute_angle_scale(xyz, centers, mobile_idx, mobile_mask, n_chains): + """Per-chain radius of gyration about the rotation centre, in Angstroms. + + The lever arm converting radians into Angstroms of atom displacement. + Floored at 1 A so a two-atom body cannot divide by ~0. + """ + d = xyz[mobile_mask] - centers[mobile_idx] + sq = torch.zeros(n_chains, dtype=xyz.dtype, device=xyz.device) + sq.index_add_(0, mobile_idx, d.pow(2).sum(dim=1)) + counts = torch.zeros(n_chains, dtype=xyz.dtype, device=xyz.device) + counts.index_add_(0, mobile_idx, torch.ones_like(d[:, 0])) + return (sq / counts.clamp(min=1.0)).sqrt().clamp(min=1.0) + + @property + def rotation_radians(self) -> torch.Tensor: + """The physical per-chain XYZ-Euler angles, in radians.""" + return self.euler_angles / self.angle_scale.unsqueeze(1) + # ----------------------------------------------------------------------- # Forward — reconstruct full xyz # ----------------------------------------------------------------------- @@ -185,7 +223,7 @@ def forward(self) -> torch.Tensor: # the angles are exactly zero): XYZ keeps the Jacobian full-rank # at the origin, while ZYZ has a gimbal-lock singularity there # (dR/dα_1 and dR/dα_3 collapse onto z-axis rotations when β=0). - R = rotation_matrix_euler_xyz(self.euler_angles) # (n_chains, 3, 3) + R = rotation_matrix_euler_xyz(self.rotation_radians) # (n_chains, 3, 3) per_atom_R = R[self.chain_indices] # (N, 3, 3) per_atom_center = self.chain_centers[self.chain_indices] # (N, 3) @@ -333,6 +371,14 @@ def update_fixed_values(self, new_values: torch.Tensor): ) centers.index_add_(0, mobile_idx, mobile_xyz * mobile_w.unsqueeze(1)) self.chain_centers.copy_(centers / w_sum.unsqueeze(1).clamp(min=1e-12)) + # The scale follows the chain geometry. A rigid re-pose leaves it + # unchanged, but any coordinates are accepted here, so recompute. + self.angle_scale.copy_( + self._compute_angle_scale( + self.original_xyz, self.chain_centers, mobile_idx, mobile, + self._n_chains, + ) + ) self.euler_angles.zero_() self.translations.zero_() self.reset_forward_cache() @@ -355,6 +401,11 @@ def copy(self) -> "RigidXYZTensor": atom_weights=self.atom_weights.clone(), ) with torch.no_grad(): + # Carry the scale rather than letting the constructor re-derive it: + # the angles are copied raw, so a scale that has been overridden + # (``angle_scale.fill_(1.0)``) would otherwise give the copy a + # different pose from the original. + new.angle_scale.copy_(self.angle_scale) new.euler_angles.copy_(self.euler_angles) new.translations.copy_(self.translations) return new diff --git a/torchref/model/sf_ds.py b/torchref/model/sf_ds.py index 003bb878..d4962cad 100644 --- a/torchref/model/sf_ds.py +++ b/torchref/model/sf_ds.py @@ -5,7 +5,7 @@ crystallographic symmetry in reciprocal space. """ -from typing import Optional, Tuple +from typing import TYPE_CHECKING, Optional, Tuple import torch import torch.nn as nn @@ -18,12 +18,14 @@ get_scattering_vectors, reciprocal_basis_matrix, ) -from torchref.config import dtypes, get_complex_dtype, get_default_device +from torchref.config import dtypes, get_complex_dtype from torchref.symmetry import Cell, SpaceGroup -from torchref.symmetry.spacegroup import SpaceGroupLike from torchref.utils.device_mixin import DeviceMovementMixin from torchref.utils.device_resolution import require_cell_dtype, resolve_device +if TYPE_CHECKING: + from torchref.model.context import ModelContext + class SfDS(DeviceMovementMixin, nn.Module): """ @@ -37,36 +39,38 @@ class SfDS(DeviceMovementMixin, nn.Module): Parameters ---------- - cell : Cell, optional - Unit cell object containing cell parameters. - spacegroup : SpaceGroupLike, optional - Space group specification (string, int, or gemmi.SpaceGroup). - If None, defaults to P1. + ctx : ModelContext + The crystallographic context to read the cell and space group from. dtype_float : torch.dtype, optional Data type for floating point tensors. Default is dtypes.float. device : torch.device, optional - Computation device. Defaults to the configured default device - (``get_default_device()``). + Computation device. Defaults to the cell's device when the context has one, + else the configured default device. verbose : int, optional Verbosity level for logging. Default is 0. max_memory_gb : float, optional Maximum memory to use for intermediate tensors in GB. Default is 2.0. Set to None to disable batching. + force_portable : bool, optional + Pin the portable reference path instead of the fastest usable backend, + per instance. ``None`` (default) defers to the process-wide setting, + so ``with use_portable():`` steers an unconfigured instance. Attributes ---------- cell, spacegroup : Cell, SpaceGroup - The unit cell and the space group as an nn.Module carrying its symmetry - matrices and translations; setting ``cell`` drops the cached - reciprocal basis. + Read through from the context. The reciprocal basis is memoised against + the cell's value key, so replacing the context's cell needs no further call. Examples -------- Standalone usage:: - from torchref.symmetry import Cell - cell = Cell([50, 60, 70, 90, 90, 90]) - sf_ds = SfDS(cell, spacegroup='P212121') + from torchref.model.context import ModelContext + from torchref.symmetry import Cell, SpaceGroup + ctx = ModelContext(cell=Cell([50, 60, 70, 90, 90, 90]), + spacegroup=SpaceGroup('P212121')) + sf_ds = SfDS(ctx) sf, _ = sf_ds.compute_structure_factors( hkl, xyz_iso, adp_iso, occ_iso, A_iso, B_iso ) @@ -74,142 +78,81 @@ class SfDS(DeviceMovementMixin, nn.Module): def __init__( self, - cell: Optional[Cell] = None, - spacegroup: SpaceGroupLike = None, + ctx: "ModelContext", + *, dtype_float: torch.dtype = None, device: torch.device = None, verbose: int = 0, max_memory_gb: float = 2.0, force_portable: Optional[bool] = None, ): - """ - Initialize the SfDS module with cell and spacegroup. - - Parameters - ---------- - cell : Cell, optional - Unit cell object. If None, must be set later. - spacegroup : SpaceGroupLike, optional - Space group specification. If None, defaults to P1. - dtype_float : torch.dtype, optional - Data type for floating point tensors. Default is dtypes.float. - device : torch.device, optional - Computation device. Defaults to the configured device.current. - verbose : int, optional - Verbosity level for logging. Default is 0. - max_memory_gb : float, optional - Maximum memory for intermediate tensors in GB. Default is 2.0. - force_portable : bool, optional - Pin the portable reference path instead of the fastest usable backend, - per instance. ``None`` (default) defers to the process-wide setting, - so ``with use_portable():`` steers an unconfigured instance. - """ super().__init__() - if dtype_float is None: - dtype_float = dtypes.float - self.dtype_float = dtype_float - # Derive from ``cell`` when no device is given, instead of jumping to + self.ctx = ctx + self.dtype_float = dtypes.float if dtype_float is None else dtype_float + # Derive from the cell when no device is given, instead of jumping to # the global default and leaving a caller-supplied cell behind on - # another device. An explicit ``device`` still wins and moves the cell. - self.device = resolve_device(cell, device=device) + # another device. An explicit ``device`` wins and moves the crystal as + # a whole. + if device is not None: + ctx.to(device) + self.device = resolve_device(ctx.cell, device=device) self.verbose = verbose self.max_memory_gb = max_memory_gb self.force_portable = force_portable - # Store cell and spacegroup - self._cell = cell - self._spacegroup = None - - if spacegroup is not None or cell is not None: - self._spacegroup = SpaceGroup( - spacegroup, dtype=dtype_float, device=self.device - ) - - # Cache reciprocal basis matrix + # Reciprocal basis, memoised against the cell it was computed from. self._recB: Optional[torch.Tensor] = None + self._recB_key = None # ========================================================================= - # Cell and SpaceGroup properties + # Crystal, read through from the context # ========================================================================= @property def cell(self) -> Optional[Cell]: - """Unit cell object.""" - return self._cell - - @cell.setter - def cell(self, value: Cell): - """Set unit cell and invalidate cached reciprocal basis matrix.""" - self._cell = value - self._recB = None # Invalidate cache + """The context's unit cell.""" + return self.ctx.cell @property def spacegroup(self) -> Optional[SpaceGroup]: - """Space group object (SpaceGroup nn.Module).""" - return self._spacegroup - - @spacegroup.setter - def spacegroup(self, value: SpaceGroupLike): - """Set space group.""" - if value is not None: - self._spacegroup = SpaceGroup( - value, dtype=self.dtype_float, device=self.device - ) - else: - self._spacegroup = None + """The context's space group.""" + return self.ctx.spacegroup @property def fractional_matrix(self) -> Optional[torch.Tensor]: - """Get fractionalization matrix from cell.""" - if self._cell is not None: - return self._cell.fractional_matrix - return None + """Fractionalization matrix from the cell.""" + cell = self.ctx.cell + return None if cell is None else cell.fractional_matrix @property def inv_fractional_matrix(self) -> Optional[torch.Tensor]: - """Get orthogonalization matrix from cell.""" - if self._cell is not None: - return self._cell.inv_fractional_matrix - return None - - def set_cell_and_spacegroup(self, cell: Cell, spacegroup: SpaceGroupLike = None): - """ - Set cell and spacegroup for this SfDS instance. - - Parameters - ---------- - cell : Cell - Unit cell object. - spacegroup : SpaceGroupLike, optional - Space group specification. - - Notes - ----- - Receiver wins: an incoming cell on another device is moved to match - this module rather than the other way round. - """ - self.device = resolve_device(self, cell) - self._cell = cell - self._recB = None # Invalidate cache - self.spacegroup = spacegroup + """Orthogonalization matrix from the cell.""" + cell = self.ctx.cell + return None if cell is None else cell.inv_fractional_matrix # ========================================================================= # Internal helper methods # ========================================================================= - def _get_reciprocal_basis_matrix(self) -> torch.Tensor: - """Cached ``(3, 3)`` reciprocal basis (a*, b*, c* as rows). - - Raises ``RuntimeError`` if no cell is set, and refuses a cell whose dtype - differs from ``self.dtype_float``. - """ - if self._cell is None: - raise RuntimeError("Cell not set. Call set_cell_and_spacegroup() first.") - require_cell_dtype(self._cell, self.dtype_float, type(self).__name__) - - if self._recB is None: - self._recB = reciprocal_basis_matrix(self._cell.data) + def _require_cell(self) -> Cell: + """The context's cell, refusing a missing one or a dtype mismatch.""" + cell = self.ctx.cell + if cell is None: + raise RuntimeError( + f"{type(self).__name__}: the context {self.ctx!r} has no cell. Load a " + "structure, or pass a context whose cell is set." + ) + # Refused, not reconciled: a dtype cast is lossy, so the cell is the + # caller's to fix. See ``require_cell_dtype``. + require_cell_dtype(cell, self.dtype_float, type(self).__name__) + return cell + def _get_reciprocal_basis_matrix(self) -> torch.Tensor: + """``(3, 3)`` reciprocal basis (a*, b*, c* as rows), memoised per cell value.""" + cell = self._require_cell() + if self._recB is None or self._recB_key != cell.key: + self._recB = reciprocal_basis_matrix(cell.data) + self._recB_key = cell.key return self._recB def _compute_scattering_factors( @@ -232,45 +175,14 @@ def _compute_scattering_factors( return f - def _get_spacegroup_callable(self): - """Symmetry-application callable for the direct-summation kernels. - - They expect ``(3, N)`` fractional coordinates in and ``(3, N, n_ops)`` - out; identity-only when no space group is set. - """ - if self._spacegroup is None: - # P1 symmetry - identity operation only - def p1_symmetry(coords_3N): - # coords_3N: (3, N) -> (3, N, 1) - return coords_3N.unsqueeze(2) - - return p1_symmetry - - def apply_symmetry(coords_3N): - # coords_3N: (3, N) -> coords_N3: (N, 3) - coords_N3 = coords_3N.T - - # Apply symmetry: (N, 3) -> (N, 3, ops) - transformed = self._spacegroup.apply(coords_N3) - - # Reorder to (3, N, ops) - # (N, 3, ops) -> (3, N, ops) - result = transformed.permute(1, 0, 2) - - return result - - return apply_symmetry - def _cartesian_to_fractional(self, xyz_cartesian: torch.Tensor) -> torch.Tensor: """``(N, 3)`` Cartesian coordinates to fractional; needs a cell whose dtype matches ``self.dtype_float``. """ - if self._cell is None: - raise RuntimeError("Cell not set. Call set_cell_and_spacegroup() first.") - require_cell_dtype(self._cell, self.dtype_float, type(self).__name__) + cell = self._require_cell() # fractional = cartesian @ inv_frac_matrix.T - return torch.matmul(xyz_cartesian, self.inv_fractional_matrix.T) + return torch.matmul(xyz_cartesian, cell.inv_fractional_matrix.T) # ========================================================================= # Structure Factor Computation @@ -322,11 +234,7 @@ def compute_structure_factors( None Second return value is None (for API compatibility with SfFFT). """ - if self._cell is None: - raise RuntimeError("Cell not set. Call set_cell_and_spacegroup() first.") - # Refused, not reconciled: unlike the device normalization just below, a dtype cast - # is lossy, so the cell is the caller's to fix. See ``require_cell_dtype``. - require_cell_dtype(self._cell, self.dtype_float, type(self).__name__) + self._require_cell() # Normalize the input hkl onto this module's device. The symmetry # helpers derive equiv_hkls/phases from hkl.device while sf_total is @@ -348,7 +256,7 @@ def compute_structure_factors( ) # No symmetry: compute F_P1 directly - if not apply_symmetry or self._spacegroup is None: + if not apply_symmetry or self.ctx.spacegroup is None: sf_p1 = self._compute_p1_sf( hkl, xyz_frac_iso, @@ -365,20 +273,9 @@ def compute_structure_factors( return sf_p1, None # Apply late symmetry: F_sym(h) = Σ_ops exp(2πi h.t) * F_P1(R^T @ h) - from torchref.base.reciprocal import ( - compute_symmetry_equivalent_hkls, - compute_translation_phases, - ) - - n_ops = self._spacegroup.n_ops - rotation_matrices = self._spacegroup.matrices - translations = self._spacegroup.translations - - # Compute equivalent HKLs: (n_ops, N, 3) - equiv_hkls = compute_symmetry_equivalent_hkls(hkl, rotation_matrices) - - # Compute translation phase shifts: (n_ops, N) - phases = compute_translation_phases(hkl, translations) + n_ops = self.ctx.spacegroup.n_ops + equiv_hkls = self.ctx.spacegroup.expand_reciprocal(hkl) # (n_ops, N, 3) + phases = self.ctx.spacegroup.phase_factors(hkl) # (n_ops, N) # Compute F_P1 at each equivalent HKL and combine sf_total = torch.zeros( @@ -430,7 +327,7 @@ def _compute_p1_sf( """ # Get reciprocal basis matrix and compute scattering vectors recB = self._get_reciprocal_basis_matrix() - s_vectors = get_scattering_vectors(hkl, self._cell.data, recB) + s_vectors = get_scattering_vectors(hkl, self.ctx.cell.data, recB) s = torch.norm(s_vectors, dim=1) sf_total = torch.zeros( @@ -474,34 +371,6 @@ def _compute_p1_sf( # ========================================================================= def reset_cache(self) -> None: - """Drop the cached reciprocal-basis matrix; recomputed on next use.""" + """Drop the memoised reciprocal basis; recomputed on next use.""" self._recB = None - - def copy(self) -> "SfDS": - """Create a deep copy of this SfDS module. - - Returns - ------- - SfDS - A new SfDS instance with cloned cell and spacegroup. - """ - # Clone the cell - new_cell = self._cell.clone() if self._cell is not None else None - - # Copy the spacegroup - new_spacegroup = ( - self._spacegroup.copy() if self._spacegroup is not None else None - ) - - # Create new SfDS with copied components - new_ds = SfDS( - cell=new_cell, - spacegroup=new_spacegroup, - dtype_float=self.dtype_float, - device=self.device, - verbose=self.verbose, - max_memory_gb=self.max_memory_gb, - force_portable=self.force_portable, - ) - - return new_ds + self._recB_key = None diff --git a/torchref/model/sf_fft.py b/torchref/model/sf_fft.py index 12c727da..81a76863 100644 --- a/torchref/model/sf_fft.py +++ b/torchref/model/sf_fft.py @@ -1,223 +1,307 @@ """SfFFT -- structure factors via FFT. -An nn.Module owning the real-space grid setup, the electron-density build from -atomic parameters, and the FFT to structure factors. Usable standalone or as -``ModelFT``'s submodule. ``FFT`` is a deprecated alias for :class:`SfFFT`. +An nn.Module that reads the crystal off a :class:`~torchref.model.context.ModelContext`, +sizes its real-space grid lazily from that crystal and the resolution, builds the +electron density from atomic parameters, and transforms it to structure factors. +Usable standalone or as ``ModelFT``'s submodule. """ -from typing import Optional, Tuple, Union +from typing import TYPE_CHECKING, Optional, Tuple import torch import torch.nn as nn -from torchref.base.fourier import get_real_grid, ifft +from torchref.base.fourier import ifft from torchref.base.reciprocal import extract_structure_factor_from_grid -from torchref.config import dtypes, get_default_device +from torchref.config import dtypes from torchref.symmetry import Cell, SpaceGroup -from torchref.symmetry.map_symmetry import MapSymmetry -from torchref.symmetry.spacegroup import SpaceGroupLike -from torchref.utils.caching import ParameterFingerprint from torchref.utils.device_mixin import DeviceMovementMixin -from torchref.utils.device_resolution import resolve_device +from torchref.utils.device_resolution import require_cell_dtype, resolve_device + +if TYPE_CHECKING: + from torchref.model.context import ModelContext class SfFFT(DeviceMovementMixin, nn.Module): """ Structure Factor calculator using FFT (Fast Fourier Transform). - Built from a Cell and optionally a SpaceGroup, which drive the grid. Call - :meth:`setup_grid` before :meth:`build_density_map`; the higher-level - :meth:`compute_structure_factors` does both. + The engine does not own a cell or a space group: it holds the + :class:`~torchref.model.context.ModelContext` it was given and reads + ``ctx.cell`` and ``ctx.spacegroup`` whenever it needs them. The grid is derived + from that crystal, ``max_res`` and ``explicit_gridsize`` on first use and kept + until any of them changes (:attr:`grid_key`), so assigning a new cell or + resolution to the context needs no further call. Parameters ---------- - cell : Cell - Unit cell object containing cell parameters. - spacegroup : SpaceGroupLike, optional - Space group specification (string, int, or gemmi.SpaceGroup). - If None, defaults to P1. + ctx : ModelContext + The crystallographic context to read the cell and space group from. May be + incomplete at construction; the grid is sized once both are set. max_res : float, optional - Maximum resolution for grid spacing in Angstroms. Default is 1.5. + Maximum resolution for grid spacing in Angstroms. Default is 1.0. + explicit_gridsize : tuple of int, optional + Fixed grid dimensions ``(nx, ny, nz)``; overrides the resolution-derived size. dtype_float : torch.dtype, optional Data type for floating point tensors. Default is dtypes.float. device : torch.device, optional - Computation device. Defaults to the configured default device - (``get_default_device()``). + Computation device. Defaults to the cell's device when the context has one, + else the configured default device. verbose : int, optional Verbosity level for logging. Default is 0. + use_late_symmetry : bool, optional + Apply symmetry in reciprocal space after the FFT when the grid permits + exact indexing (default); otherwise symmetrise the density map first. Attributes ---------- cell, spacegroup : Cell, SpaceGroup - The unit cell and the space group as an nn.Module carrying its symmetry - matrices and translations; ``symmetry`` is an alias for ``spacegroup``. - gridsize, real_space_grid, voxel_size : torch.Tensor or None - Grid dimensions ``(nx, ny, nz)``, coordinate grid ``(nx, ny, nz, 3)`` and - voxel dimensions -- all ``None`` until :meth:`setup_grid` runs. - map_symmetry : MapSymmetry or None - Symmetry operator for map calculations. + Read through from the context. + gridsize, voxel_size : torch.Tensor or None + Grid dimensions ``(nx, ny, nz)`` and the voxel edge vector sum, resolved on + access; ``None`` while the context has no cell or space group. No coordinate + grid is stored: the splats derive a voxel's Cartesian position from its + index, so materialising one would cost ``12 * nx * ny * nz`` bytes that + nothing reads. Call :func:`torchref.base.fourier.get_real_grid` if you + genuinely need one. """ def __init__( self, - cell: Optional[Cell] = None, - spacegroup: SpaceGroupLike = None, - max_res: float = 1.5, + ctx: "ModelContext", + *, + max_res: Optional[float] = 1.0, + explicit_gridsize: Optional[Tuple[int, int, int]] = None, dtype_float: torch.dtype = None, device: Optional[torch.device] = None, verbose: int = 0, use_late_symmetry: bool = True, ): - """ - Initialize the SfFFT module with cell and spacegroup. - - Parameters - ---------- - cell : Cell, optional - Unit cell object. If None, must be set later via set_cell(). - spacegroup : SpaceGroupLike, optional - Space group specification. If None, defaults to P1. - max_res : float, optional - Maximum resolution for grid spacing in Angstroms. Default is 1.5. - dtype_float : torch.dtype, optional - Data type for floating point tensors. Default is dtypes.float. - device : torch.device, optional - Computation device. Default is None (uses cell's device). If Cell is also None, defaults to CPU. - verbose : int, optional - Verbosity level for logging. Default is 0. - use_late_symmetry : bool, optional - If True (default), apply symmetry in reciprocal space after FFT - ("late symmetry") for faster structure factor calculation. - If False, apply symmetry to density map before FFT ("early symmetry"). - """ super().__init__() + self.ctx = ctx self.max_res = max_res - if dtype_float is None: - dtype_float = dtypes.float - self.dtype_float = dtype_float + self.explicit_gridsize = explicit_gridsize + self.dtype_float = dtypes.float if dtype_float is None else dtype_float - # One device for the module and everything it builds. ``resolve_device`` - # also moves ``cell`` when an explicit ``device`` disagrees with it, so - # the cell and the SpaceGroup below cannot end up split. - self.device = resolve_device(cell, device=device) + # One device for the module and everything it builds. An explicit + # ``device`` moves the crystal as a whole, so cell, space group and the + # grid buffers cannot end up split; without one the cell's device wins, + # and an empty context falls back to the configured default. + if device is not None: + ctx.to(device) + self.device = resolve_device(ctx.cell, device=device) self.verbose = verbose self.use_late_symmetry = use_late_symmetry - # Store cell and spacegroup - self._cell = cell - self._spacegroup = None + # Derived from the crystal on first use. Non-persistent: rebuilt from the + # context on restore rather than read back from a checkpoint. + self.register_buffer("_gridsize", None, persistent=False) + self.register_buffer("_voxel_size", None, persistent=False) + self._late_symmetry_compatible: Optional[bool] = None + self._grid_key = None - if spacegroup is not None or cell is not None: - # ``self.device``, not the raw ``device`` argument: the latter is - # ``None`` on the derive-from-cell path, which would silently put - # the symmetry matrices on the global default instead. - self._spacegroup = SpaceGroup( - spacegroup, dtype=dtype_float, device=self.device - ) + # ========================================================================= + # Crystal, read through from the context + # ========================================================================= - # Buffers (registered during setup_grid) - self.register_buffer("gridsize", None) - self.register_buffer("real_space_grid", None) - self.register_buffer("voxel_size", None) + @property + def cell(self) -> Optional[Cell]: + """The context's unit cell.""" + return self.ctx.cell - # Map symmetry operator (set during setup_grid) - self.map_symmetry: Optional[MapSymmetry] = None + @property + def spacegroup(self) -> Optional[SpaceGroup]: + """The context's space group.""" + return self.ctx.spacegroup - # Late symmetry compatibility flag (set during setup_grid) - self._late_symmetry_compatible: Optional[bool] = None + @property + def fractional_matrix(self) -> Optional[torch.Tensor]: + """Fractionalization matrix from the cell, on this module's device/dtype.""" + cell = self.ctx.cell + if cell is None: + return None + # Move device first, then cast: a combined ``.to(device=cpu, + # dtype=float64)`` from an MPS-resident cell raises because MPS + # rejects the transient float64 view (MPS has no float64). + return cell.fractional_matrix.to(device=self.device).to(dtype=self.dtype_float) - # Cached reciprocal symmetry extractor (precomputed flat indices). - # Keyed on a (data_ptr, _version, numel) fingerprint of the HKL tensor - # rather than id(hkl): a garbage-collected tensor reallocated at the - # same address cannot silently alias a stale extractor. - self._sym_extractor = None - self._sym_extractor_hkl_fp: Optional[ParameterFingerprint] = None + @property + def inv_fractional_matrix(self) -> Optional[torch.Tensor]: + """Orthogonalization matrix from the cell, on this module's device/dtype.""" + cell = self.ctx.cell + if cell is None: + return None + # Move device first, then cast (see ``fractional_matrix``). + return cell.inv_fractional_matrix.to(device=self.device).to( + dtype=self.dtype_float + ) # ========================================================================= - # Cell and SpaceGroup properties + # Grid # ========================================================================= @property - def cell(self) -> Optional[Cell]: - """Unit cell object.""" - return self._cell - - @cell.setter - def cell(self, value: Cell): - """Set unit cell.""" - self._cell = value + def explicit_gridsize(self) -> Optional[Tuple[int, int, int]]: + """Fixed grid dimensions, or None to size the grid from ``max_res``.""" + return self._explicit_gridsize + + @explicit_gridsize.setter + def explicit_gridsize(self, value) -> None: + self._explicit_gridsize = ( + None if value is None else tuple(int(x) for x in value) + ) @property - def spacegroup(self) -> Optional[SpaceGroup]: - """Space group object (SpaceGroup nn.Module).""" - return self._spacegroup - - @spacegroup.setter - def spacegroup(self, value: SpaceGroupLike): - """Set space group.""" - if value is not None: - self._spacegroup = SpaceGroup( - value, dtype=self.dtype_float, device=self.device + def grid_key(self): + """Everything the grid is derived from, as a hashable tuple. + + ``(cell.key, spacegroup.key, max_res, explicit_gridsize)``, or ``None`` + while the context has no cell or space group. Callers that cache anything + grid-shaped can compare against it. + """ + crystal = self.ctx.crystal_key + if crystal is None: + return None + return ( + *crystal, + None if self.max_res is None else float(self.max_res), + self.explicit_gridsize, + ) + + def ensure_grid(self) -> bool: + """Bring the grid in line with :attr:`grid_key`. + + Returns + ------- + bool + True when the grid was (re)built, False when it was already current or + the context has no crystal yet. + + Raises + ------ + RuntimeError + If neither ``max_res`` nor ``explicit_gridsize`` is set, or the cell or + space group disagree with this module's dtype or device. + """ + key = self.grid_key + if key == self._grid_key: + return False + if key is None: + # The crystal went away; drop the grid derived from the old one. + self._gridsize = None + self._voxel_size = None + self._late_symmetry_compatible = None + self._grid_key = None + return False + + cell, spacegroup = self.ctx.cell, self.ctx.spacegroup + require_cell_dtype(cell, self.dtype_float, type(self).__name__) + if spacegroup.matrices.dtype != self.dtype_float: + raise RuntimeError( + f"{type(self).__name__} was built for {self.dtype_float} but its " + f"space group holds {spacegroup.matrices.dtype}. Rebuild the space " + "group at the module's dtype." + ) + if spacegroup.matrices.device != cell.data.device: + raise RuntimeError( + f"{type(self).__name__}: cell on {cell.data.device} but space group " + f"on {spacegroup.matrices.device}. Move the context as a whole." ) + + if self.explicit_gridsize is not None: + gridsize = self.explicit_gridsize + elif self.max_res is not None: + gridsize = self.compute_optimal_gridsize(self.max_res) else: - self._spacegroup = None + raise RuntimeError( + f"{type(self).__name__} cannot size its grid: set max_res or " + "explicit_gridsize." + ) + shape = tuple(int(n) for n in gridsize) - @property - def symmetry(self) -> Optional[SpaceGroup]: - """Symmetry operations handler (alias for spacegroup).""" - return self._spacegroup + if self.verbose > 1: + print(f"Setting up grids with max_res={self.max_res} Å") - @property - def fractional_matrix(self) -> Optional[torch.Tensor]: - """Get fractionalization matrix from cell, on this module's device/dtype.""" - if self._cell is not None: - # Move device first, then cast: a combined ``.to(device=cpu, - # dtype=float64)`` from an MPS-resident cell raises because MPS - # rejects the transient float64 view (MPS has no float64). - return self._cell.fractional_matrix.to(device=self.device).to( - dtype=self.dtype_float + previous = self._gridsize + self._gridsize = torch.tensor(shape, dtype=dtypes.int, device=self.device) + + # The step between diagonally adjacent grid points, i.e. the sum of the three + # cell edge vectors each divided by its own sampling count. Equal to the true + # per-axis voxel edge lengths only for an orthogonal cell; kept because that is + # what the previous grid-differencing definition produced. + self._voxel_size = self.fractional_matrix @ ( + 1.0 / self._gridsize.to(self.dtype_float) + ) + + # Every symmetry-equivalent HKL lands on an integer grid point exactly when + # the grid admits direct indexing, which the space group answers without + # building an operator. + self._late_symmetry_compatible = spacegroup.can_index_directly(shape) + if self.verbose > 0 and self.use_late_symmetry: + if self._late_symmetry_compatible: + print("SfFFT: Using late symmetry (reciprocal space)") + else: + print( + "SfFFT: Late symmetry disabled - grid not compatible " + "(falling back to early symmetry)" + ) + + # The space group memoises its map operator and reciprocal extractor per + # grid shape. They are keyed on the shape, so a same-shape rebuild keeps + # them; a different shape drops them so the old operator's sampling grids + # do not stay resident. + if previous is not None and tuple(int(n) for n in previous.tolist()) != shape: + spacegroup.reset_cache() + + if self.verbose > 2: + print(f"Grid shape: {shape}") + print(f"Voxel size: {self._voxel_size}") + + # Last, so a failure above leaves the old key in place and the next call + # tries again. + self._grid_key = key + return True + + def _require_grid(self) -> None: + """Resolve the grid, refusing to proceed without a crystal.""" + self.ensure_grid() + if self._gridsize is None: + raise RuntimeError( + f"{type(self).__name__} has no crystal to size its grid: the context " + f"{self.ctx!r} has no cell or space group. Load a structure, or pass a " + "context whose cell and spacegroup are set." ) - return None @property - def inv_fractional_matrix(self) -> Optional[torch.Tensor]: - """Get orthogonalization matrix from cell, on this module's device/dtype.""" - if self._cell is not None: - # Move device first, then cast (see ``fractional_matrix``). - return self._cell.inv_fractional_matrix.to(device=self.device).to( - dtype=self.dtype_float - ) - return None + def gridsize(self) -> Optional[torch.Tensor]: + """Grid dimensions ``(nx, ny, nz)``, or None without a crystal.""" + self.ensure_grid() + return self._gridsize - def set_cell_and_spacegroup(self, cell: Cell, spacegroup: SpaceGroupLike = None): - """ - Set cell and spacegroup for this SfFFT instance. + @property + def voxel_size(self) -> Optional[torch.Tensor]: + """Voxel edge vector sum, or None without a crystal.""" + self.ensure_grid() + return self._voxel_size - Parameters - ---------- - cell : Cell - Unit cell object. - spacegroup : SpaceGroupLike, optional - Space group specification. - - Notes - ----- - Receiver wins: this module may already own grid buffers, so an incoming - cell on another device is moved to match rather than dragging the - module after it. - """ - self.device = resolve_device(self, cell) - self._cell = cell - self.spacegroup = spacegroup + @property + def grid_shape(self) -> Optional[Tuple[int, int, int]]: + """Map dimensions ``(nx, ny, nz)`` as Python ints, or None without a crystal.""" + gridsize = self.gridsize + if gridsize is None: + return None + return tuple(int(n) for n in gridsize) - # ========================================================================= - # Grid Setup Methods - # ========================================================================= + @property + def late_symmetry_compatible(self) -> Optional[bool]: + """Whether the current grid admits reciprocal-space symmetrisation.""" + self.ensure_grid() + return self._late_symmetry_compatible def compute_optimal_gridsize(self, max_res: Optional[float] = None) -> tuple: """ - Compute optimal grid dimensions using the stored cell and spacegroup. + Compute optimal grid dimensions from the context's cell and space group. Uses Cell.compute_grid_size() for base calculation and Symmetry.suggest_grid_size() for symmetry optimization. @@ -235,160 +319,56 @@ def compute_optimal_gridsize(self, max_res: Optional[float] = None) -> tuple: Raises ------ RuntimeError - If cell has not been set. + If the context has no cell or space group. """ - if self._cell is None: - raise RuntimeError("Cell not set. Call set_cell_and_spacegroup() first.") + cell, spacegroup = self.ctx.cell, self.ctx.spacegroup + if cell is None or spacegroup is None: + raise RuntimeError( + f"{type(self).__name__}: the context {self.ctx!r} has no cell or " + "space group to size a grid from." + ) resolution = max_res if max_res is not None else self.max_res - from torchref.symmetry.spacegroup import suggest_grid_size - # Use Cell's method for base grid size calculation - gridsize_initial = self._cell.compute_grid_size(resolution) + gridsize_initial = cell.compute_grid_size(resolution) if self.verbose > 1: print(f"Initial grid size from cell: {gridsize_initial}") # Optimize for symmetry and FFT-friendliness - gridsize_optimized = suggest_grid_size( - gridsize_initial, self._spacegroup, make_fft_friendly=True + gridsize_optimized = spacegroup.suggest_grid_size( + gridsize_initial, make_fft_friendly=True ) if self.verbose > 1 and gridsize_optimized != gridsize_initial: print( f"Optimized grid size from {gridsize_initial} to {gridsize_optimized} " f"(symmetry + FFT friendly)" ) - return gridsize_optimized - - @staticmethod - def compute_real_space_grid( - fractional_matrix: torch.Tensor, - gridsize: torch.Tensor, - device: torch.device = None, - ) -> torch.Tensor: - """ - Generate the real-space coordinate grid. - - Parameters - ---------- - fractional_matrix : torch.Tensor - Fractionalization matrix mapping Cartesian to fractional - coordinates, with shape (3, 3). - gridsize : torch.Tensor - Grid dimensions (nx, ny, nz). - device : torch.device, optional - Target device. Defaults to the configured default device. - - Returns - ------- - torch.Tensor - Real-space grid with shape (nx, ny, nz, 3). - """ - # Forward ``device`` as-is, including ``None``: ``get_real_grid`` infers - # from ``fractional_matrix`` when no device is given, and resolving the - # global default here would preempt that. - return get_real_grid( - fractional_matrix=fractional_matrix, gridsize=gridsize, device=device - ) + return tuple(int(n) for n in gridsize_optimized) def setup_grid( self, - gridsize: Optional[Tuple[int, int, int]] = None, + *, max_res: Optional[float] = None, - ): + gridsize: Optional[Tuple[int, int, int]] = None, + ) -> None: """ - Setup the real-space grid for electron density calculation. - - This method initializes and stores the grid state for subsequent - density map calculations. Uses the stored cell and spacegroup. + Override the grid's inputs explicitly and resolve it now. Parameters ---------- - gridsize : tuple of int, optional - Explicit grid size (nx, ny, nz). If None, computed automatically - using Cell.compute_grid_size() and Symmetry.suggest_grid_size(). max_res : float, optional - Maximum resolution in Angstroms. If None, uses self.max_res. - - Raises - ------ - RuntimeError - If cell has not been set. + New maximum resolution in Angstroms. None leaves the current value. + gridsize : tuple of int, optional + Fixed grid size (nx, ny, nz), kept until cleared through + :attr:`explicit_gridsize`. None leaves the current value. """ - if self._cell is None: - raise RuntimeError("Cell not set. Call set_cell_and_spacegroup() first.") - if max_res is not None: - self.max_res = max_res - - if self.verbose > 1: - print(f"Setting up grids with max_res={self.max_res} Å") - - # Compute or use provided grid size + self.max_res = float(max_res) if gridsize is not None: - self.gridsize = torch.tensor(gridsize, dtype=dtypes.int, device=self.device) - else: - optimal_gridsize = self.compute_optimal_gridsize(self.max_res) - self.gridsize = torch.tensor( - optimal_gridsize, dtype=dtypes.int, device=self.device - ) - - # Compute real space grid - self.real_space_grid = self.compute_real_space_grid( - self._cell.fractional_matrix, self.gridsize, self.device - ) - - # Compute voxel size - self.voxel_size = self.real_space_grid[2, 2, 2] - self.real_space_grid[1, 1, 1] - - # Initialize map symmetry operator if space group is set - if self._spacegroup is not None: - self.map_symmetry = MapSymmetry( - space_group=self._spacegroup, - map_shape=self.real_space_grid.shape[:-1], - cell_params=self._cell.data, - verbose=self.verbose, - device=self.device, - ) - - # Check late symmetry compatibility - self._late_symmetry_compatible = self._check_late_symmetry_compatible() - - if self.use_late_symmetry and self._late_symmetry_compatible: - if self.verbose > 0: - print( - "SfFFT: Using late symmetry (reciprocal space)" - ) - elif self.use_late_symmetry and not self._late_symmetry_compatible: - if self.verbose > 0: - print( - "SfFFT: Late symmetry disabled - grid not compatible " - "(falling back to early symmetry)" - ) - else: - self.map_symmetry = None - self._late_symmetry_compatible = False - - # Invalidate cached symmetry extractor (grid shape changed) - self._sym_extractor = None - self._sym_extractor_hkl_fp = None - - if self.verbose > 2: - print(f"Grid shape: {self.real_space_grid.shape[:-1]}") - print(f"Voxel size: {self.voxel_size}") - - def _check_late_symmetry_compatible(self) -> bool: - """True when every symmetry-equivalent HKL lands on an integer grid - point, which the MapSymmetry factory signals by returning a - ``MapSymmetryDirect`` (direct indexing, no interpolation). - """ - if self.map_symmetry is None: - return False - - from torchref.symmetry.map_symmetry import MapSymmetryDirect - - return isinstance(self.map_symmetry, MapSymmetryDirect) + self.explicit_gridsize = gridsize + self.ensure_grid() # ========================================================================= # Density Map Building Methods @@ -411,8 +391,6 @@ def build_density_map( """ Build electron density map from atomic parameters. - Calls :meth:`setup_grid` itself if no grid has been set up yet. - Parameters ---------- xyz_iso, adp_iso, occ_iso : torch.Tensor @@ -433,13 +411,13 @@ def build_density_map( torch.Tensor Electron density map with shape (nx, ny, nz). """ - if self.real_space_grid is None: - self.setup_grid() + self._require_grid() from torchref.base.electron_density.main import build_electron_density density_map = build_electron_density( - real_space_grid=self.real_space_grid, + grid_shape=self.grid_shape, + device=self.device, xyz_iso=xyz_iso, adp_iso=adp_iso, occ_iso=occ_iso, @@ -447,7 +425,6 @@ def build_density_map( B_iso=B_iso, inv_frac_matrix=self.inv_fractional_matrix, frac_matrix=self.fractional_matrix, - voxel_size=self.voxel_size, xyz_aniso=xyz_aniso, u_aniso=u_aniso, occ_aniso=occ_aniso, @@ -456,9 +433,8 @@ def build_density_map( dtype=self.dtype_float, ) - # Apply symmetry if requested - if apply_symmetry and self.map_symmetry is not None: - density_map = self.map_symmetry(density_map) + if apply_symmetry: + density_map = self.ctx.spacegroup.symmetrize_map(density_map) return density_map @@ -483,37 +459,23 @@ def map_to_structure_factors( hkl : torch.Tensor Miller indices with shape (n_reflections, 3). apply_symmetry : bool, optional - If True (default) and late symmetry is enabled/compatible, apply - symmetry in reciprocal space. If False, the density map is assumed - to already have symmetry applied (early symmetry path). + If True (default), apply symmetry in reciprocal space. If False, the + density map is assumed to already have symmetry applied (early + symmetry path). Returns ------- torch.Tensor Complex structure factors with shape (n_reflections,). """ - reciprocal_space_grid = ifft(density_map, self.cell.volume) + self._require_grid() + reciprocal_space_grid = ifft(density_map, self.ctx.cell.volume) - # Use late symmetry if enabled, compatible, and requested if apply_symmetry: - # Lazily build / reuse cached extractor (precomputed flat indices) - if self._sym_extractor is None or ( - self._sym_extractor_hkl_fp is None - or not self._sym_extractor_hkl_fp.matches([hkl]) - ): - from torchref.base.reciprocal import ReciprocalSymmetryExtractor - - grid_shape = tuple(int(x) for x in self.gridsize) - self._sym_extractor = ReciprocalSymmetryExtractor( - hkl, - self.spacegroup, - grid_shape, - device=reciprocal_space_grid.device, - ) - self._sym_extractor_hkl_fp = ParameterFingerprint([hkl]) - return self._sym_extractor.extract_from_grid(reciprocal_space_grid) - else: - return extract_structure_factor_from_grid(reciprocal_space_grid, hkl) + # Memoised on the space group per (hkl, grid shape). + extractor = self.ctx.spacegroup.reciprocal_extractor(hkl, self.grid_shape) + return extractor.extract_from_grid(reciprocal_space_grid) + return extract_structure_factor_from_grid(reciprocal_space_grid, hkl) def compute_structure_factors( self, @@ -563,8 +525,9 @@ def compute_structure_factors( Electron density map with shape (nx, ny, nz). Note: When using late symmetry, this is the P1 map (without symmetry). """ - # Late symmetry: build a P1 map, symmetrize in reciprocal space. - # Early symmetry: symmetrize the density map before the FFT. + # Resolve the grid first: the late-symmetry flag belongs to the grid the + # density is about to be built on. + self._require_grid() use_late = ( apply_symmetry and self.use_late_symmetry and self._late_symmetry_compatible ) @@ -588,55 +551,3 @@ def compute_structure_factors( apply_symmetry=use_late, # Late symmetry ) return sf, density_map - - # ========================================================================= - # Device Movement - # ========================================================================= - - def reset_cache(self) -> None: - """Drop the cached symmetry extractor; recomputed on next use.""" - self._sym_extractor = None - self._sym_extractor_hkl_fp = None - - def copy(self) -> "SfFFT": - """Create a deep copy of this SfFFT module. - - Returns - ------- - SfFFT - A new SfFFT instance with cloned cell, spacegroup, and buffers. - """ - # Clone the cell - new_cell = self._cell.clone() if self._cell is not None else None - - # Copy the spacegroup - new_spacegroup = ( - self._spacegroup.copy() if self._spacegroup is not None else None - ) - - # Create new SfFFT with copied components - new_fft = SfFFT( - cell=new_cell, - spacegroup=new_spacegroup, - max_res=self.max_res, - dtype_float=self.dtype_float, - device=self.device, - verbose=self.verbose, - use_late_symmetry=self.use_late_symmetry, - ) - - return new_fft - - -# Backward compatibility alias — deprecated, use SfFFT directly -def FFT(*args, **kwargs): - """Deprecated: use SfFFT instead.""" - import warnings - - warnings.warn( - "FFT is deprecated, use SfFFT instead. " - "FFT will be removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - return SfFFT(*args, **kwargs) diff --git a/torchref/refinement/base_refinement.py b/torchref/refinement/base_refinement.py index 4eaf78e2..96ee5799 100644 --- a/torchref/refinement/base_refinement.py +++ b/torchref/refinement/base_refinement.py @@ -4,6 +4,7 @@ from typing import Any, Dict, Optional +import math import torch from torch.nn import Module as nnModule @@ -53,8 +54,27 @@ # neighbours, and the log-normal KL term it replaced was a single intensive # scalar. Pending the R_free weight scan, 1.0 leaves it at the group weight. "adp/sigd": 1.0, + # Load balancing for the node-field ADP representation. Sub-weight on the adp + # group, and inert on the per-atom path, so it only acts in field mode. Set + # above the group weight because it is a barrier against a degenerate direction + # rather than a prior competing with the data. + "adp/node_load": 10.0, + # Magnitude prior on the node values. Off pending its own measurement: the load + # barrier acts only on the weights, so this is what actually bounds an extreme + # node B, but it has not been screened yet. Same convention as + # 'geometry/ramachandran'. + "adp/node_smoothness": 0.0, } +#: Weight overrides a node-field ADP representation needs, applied by +#: :meth:`BaseRefinement.set_adp_representation`. +#: +#: Work reflections per ADP parameter that :meth:`set_adp_representation` targets when +#: sizing a field. PDB-REDO holds ~7 across its whole resolution range and switches model +#: form to stay there; measured on 179 of their entries, 7 is also where a node field +#: peaks, and both directions from it are worse. +DEFAULT_REFLECTIONS_PER_ADP_PARAMETER = 7.0 + class Refinement(DeviceMixin, DebugMixin, nnModule): """ @@ -109,11 +129,16 @@ def __init__( french_wilson: bool = True, anomalous: Optional[bool] = None, adp_mode: str = "isotropic", + adp_mode_set: str = None, + n_nodes: int = None, + reflections_per_adp_parameter: float = DEFAULT_REFLECTIONS_PER_ADP_PARAMETER, xray_mode: str = "ml", sigma_a_max: float = SIGMA_A_MAX, shrink: bool = SHRINK_ENABLED, scale_target: str = DEFAULT_SCALE_TARGET, aniso_selection: Optional[str] = None, + add_hydrogens: bool = False, + hydrogens_in_xray: bool = True, ): """Initialize Refinement, fully if ``data_file`` and ``pdb`` are given. @@ -126,8 +151,9 @@ def __init__( Path to the MTZ or CIF file holding reflection data. pdb : str, optional Path to the PDB or CIF file holding the initial model. - cif : str, optional - Path to a CIF file of restraints (monomer library). + cif : str or list of str, optional + Restraint dictionary file(s) for residues the monomer library lacks. Given to + the model at construction so hydrogen generation on load reads it too. verbose : int, optional Verbosity level. Default 1. max_res : float, optional @@ -157,6 +183,19 @@ def __init__( ADP parametrization: ``"isotropic"`` (default) refines a per-atom B-factor, ``"anisotropic"`` a 6-component U tensor for the atoms selected by ``aniso_selection`` (see :meth:`Model.set_adp_mode`). + ``"field"`` / ``"field_aniso"`` replace it with a node field, sized and + reweighted by :meth:`set_adp_representation`; ``"preserve"`` leaves the + file's own ADPs untouched. + adp_mode_set : str, optional + Displacement-mode set for ``adp_mode="field_aniso"`` --- ``"rigid"`` is TLS, + ``"rigid_dilation"`` adds uniform breathing. See + :data:`~torchref.model.disorder_field.MODE_SETS`. + n_nodes : int, optional + Explicit node count for a field mode. Default None sizes it from the data + through ``reflections_per_adp_parameter``. + reflections_per_adp_parameter : float, optional + Work reflections per ADP parameter a field is sized to hold. Default 7, + which is where a node field peaks and what PDB-REDO holds. xray_mode : str, optional X-ray target taxonomy row; see :meth:`set_xray_target_mode`. sigma_a_max, shrink : optional @@ -168,6 +207,12 @@ def __init__( aniso_selection : str, optional Phenix-style selection of atoms refined anisotropically when ``adp_mode="anisotropic"``. Defaults to all non-water heavy atoms. + add_hydrogens : bool, optional + Generate missing hydrogens when loading the model. Default False. + Hydrogens already present in the input are retained either way. + hydrogens_in_xray : bool, optional + Whether hydrogens contribute to the structure factors. Default True. They + take part in the restraints either way. """ super().__init__() # Refinement constructs its own submodules from file paths, so @@ -197,6 +242,11 @@ def __init__( # model right after load, before scaling/restraints/targets. self.adp_mode = adp_mode self.aniso_selection = aniso_selection + # Node-field settings. adp_mode_set names the displacement-mode set; n_nodes + # None means "size it from the data", which is what set_adp_representation does. + self.adp_mode_set = adp_mode_set + self.n_nodes = n_nodes + self.reflections_per_adp_parameter = reflections_per_adp_parameter # Everything the x-ray targets are built from must be set BEFORE # _init_targets() further down this __init__ (it also calls get_scales()). # They are read back through _xray_target_kwargs(), which is the single @@ -233,6 +283,9 @@ def __init__( device=self.device, wavelength=self.wavelength, anomalous_threshold=self.anomalous_threshold, + add_hydrogens=add_hydrogens, + cif_path=cif, + hydrogens_in_xray=hydrogens_in_xray, ) self.scaler = Scaler( verbose=self.verbose, device=self.device, nbins=self.nbins, @@ -279,6 +332,10 @@ def __init__( device=self.device, wavelength=self.wavelength, anomalous_threshold=self.anomalous_threshold, + add_hydrogens=add_hydrogens, + hydrogens_in_xray=hydrogens_in_xray, + # Before load, not after: generation on load reads this dictionary. + cif_path=cif, # Apply the f'' (Bijvoet) term only when the data were loaded as # explicit Friedel pairs; merged data gate it off. apply_bijvoet=not self.reflection_data.friedel_merged, @@ -293,12 +350,19 @@ def __init__( ) self._sync_model_cell_to_data() - # Set ADP parametrization (iso/aniso) before scaling/restraints/targets - # so all structure-factor evaluation sees the chosen representation. - self.model.set_adp_mode(self.adp_mode, self.aniso_selection) + # Set the ADP parametrization before scaling/restraints/targets so all + # structure-factor evaluation sees the chosen representation. Routed through + # set_adp_representation rather than straight to the model: a field mode has + # to be sized from the reflection count and reweighted, and the model can do + # neither. Targets do not exist yet, so it will not try to rebuild them. + self.set_adp_representation( + self.adp_mode, + mode_set=self.adp_mode_set, + n_nodes=self.n_nodes, + reflections_per_parameter=self.reflections_per_adp_parameter, + ) self.setup_scaler() - # Configure CIF path for lazy restraint building (restraints built on first access) - self.model.set_restraints_cif(cif) + # The CIF path went in at construction; build the restraints over it now. self.model._build_restraints() self._freeze_unrestrained_residues() @@ -377,7 +441,7 @@ def mark(idx): return # 4. freeze xyz of those atoms (same path as freeze_selection) - model.xyz_mask[torch.tensor(freeze_idx, dtype=torch.long)] = False + model.xyz_mask[torch.tensor(freeze_idx, dtype=torch.long)] = False # dtype-ok: freeze index used to index xyz_mask; PyTorch requires int64 model.apply_mask_to_parameter("xyz") if self.verbose > 0: shown = frozen_res[:20] + (["..."] if len(frozen_res) > 20 else []) @@ -415,6 +479,207 @@ def _build_xray_targets(self, mode: str) -> None: ) self.xray_mode = mode + # ------------------------------------------------------------------ + # ADP representation. + # ------------------------------------------------------------------ + + FIELD_MODES = ("field", "field_aniso") + + def _field_parameters_per_node(self, mode, mode_set, refine_node_positions): + """Storage columns one node costs: payload + log sigma + optional offset. + + Read off the payload rather than tabulated, so a new payload cannot silently + desynchronise the budget arithmetic from what the field actually allocates. + """ + from torchref.model.disorder_field import ( + AnisotropicPayload, + IsotropicPayload, + ModeCovariancePayload, + ) + + if mode_set is not None: + payload = ModeCovariancePayload(mode_set) + elif mode == "field_aniso": + payload = AnisotropicPayload() + else: + payload = IsotropicPayload() + return payload.width + 1 + (3 if refine_node_positions else 0) + + def nodes_for_reflection_budget( + self, + mode: str = "field_aniso", + mode_set: str = None, + reflections_per_parameter: float = DEFAULT_REFLECTIONS_PER_ADP_PARAMETER, + refine_node_positions: bool = True, + ) -> int: + """Node count giving ``reflections_per_parameter`` work reflections per ADP parameter. + + The reason this lives on the refinement and not on :class:`Model`: the model has + no idea how much data there is, and node count is set by the data rather than by + the structure. Measured on 179 PDB-REDO entries, node count correlates with + reflection count far more strongly than with atom count, and the model's own + default (one node per 25 atoms) is unrelated to either. + + The work set is the denominator because it is what the refinement fits, and it is + what PDB-REDO's ``NREFCNT`` counts, so the ratio is comparable to theirs. + + Returns + ------- + int + At least 2 --- a single node has no spatial structure to express. + """ + per_node = self._field_parameters_per_node( + mode, mode_set, refine_node_positions + ) + n_work = int(self.data.work.n) + budget = n_work / float(reflections_per_parameter) + return max(2, int(round(budget / per_node))) + + def flatten_adp_field(self) -> bool: + """Discard the field's spatial structure, keeping its level. Returns whether it ran. + + A node field fits its structure once, at the moment it is installed, and then only + refines from there. Early in a refinement that structure is derived against + coordinates that are still wrong, and nothing later re-derives it -- the same shape + of mistake as fitting bulk solvent to the starting model and never revisiting it, + which cost 11.5% error by cycle 4. Calling this between macro cycles throws away + the accumulated structure so the data rebuilds it against the coordinates as they + now are. + + The level is preserved: only the spatial variation is reset. Deliberately a hard + reset rather than a pull toward flat, because a soft version is another weight to + tune and the point is to test whether re-deriving helps at all. + + No-op when the model is not in field mode, so a driver can call it unconditionally. + """ + field = self.model.adp_field + if field is None: + return False + with torch.no_grad(): + per_atom = field().detach() + if per_atom.ndim == 2: + # A U6 field. Flatten through the equivalent isotropic B, NOT by taking a + # median over all six components: setting the off-diagonals to the same + # value as the diagonals gives a matrix with eigenvalues (3L, 0, 0), which + # is singular, and the Cholesky encode of it is NaN. refit lifts a 1-D B + # target to U_iso * I, which is the flat U that is actually meant. + b = (8.0 * math.pi**2 / 3.0) * per_atom[:, :3].sum(dim=1) + else: + b = per_atom + finite = torch.isfinite(b) + if not bool(finite.any()): + return False + level = b[finite].median() + target = torch.where(finite, level.expand_as(b), b) + # refit replaces refinable_params, so any cached leaf set or optimizer state + # referring to the old tensor is stale. + field.refit(target) + self.reset_loss_state() + if self.verbose > 0: + print(f"Flattened the ADP field to a level of {float(level):.2f}") + return True + + def set_adp_representation( + self, + mode: str, + mode_set: str = None, + n_nodes: int = None, + reflections_per_parameter: float = DEFAULT_REFLECTIONS_PER_ADP_PARAMETER, + k_neighbors: int = 12, + refine_node_positions: bool = True, + aniso_selection: str = None, + ): + """Switch the ADP parametrization, sizing and reweighting it for this data set. + + :meth:`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. It also cannot swap the ADP restraint set, which is a property of the + representation rather than a weight to tune. + + The loss is **not** rebalanced for a field. The point of the representation is that + smoothness comes from the parametrisation, so a field should need *less* + regularisation than a per-atom model, not a reweighted version of the same + priors. :data:`DEFAULT_GROUP_WEIGHTS` already carries everything a field needs, + and an earlier attempt to raise the ``adp`` group for field mode had two side + effects worth remembering: ``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 the field's configuration differ from every per-atom baseline in a way that + had nothing to do with ADPs. + + Safe to call after construction: the targets and scales are rebuilt afterwards, + which is what the model's own "run once at setup" caveat is about. + + Parameters + ---------- + mode : str + Any mode :meth:`Model.set_adp_mode` accepts. ``"field"`` and + ``"field_aniso"`` are sized and reweighted; the per-atom modes just pass + through, with any field weight overrides removed again. + mode_set : str, optional + Displacement-mode set for ``mode="field_aniso"``; see + :data:`~torchref.model.disorder_field.MODE_SETS`. + n_nodes : int, optional + Explicit node count, bypassing the reflection budget entirely. + reflections_per_parameter : float, optional + Target work reflections per ADP parameter when ``n_nodes`` is not given. + + Returns + ------- + dict + What was applied: mode, mode set, node count, parameter count and the + reflections-per-parameter actually achieved. Worth logging --- the achieved + ratio differs from the requested one by the integer rounding of node count. + """ + is_field = mode in self.FIELD_MODES + if mode_set is not None and mode != "field_aniso": + raise ValueError( + f"mode_set={mode_set!r} describes an anisotropic displacement field; " + 'use mode="field_aniso".' + ) + + if is_field and n_nodes is None: + n_nodes = self.nodes_for_reflection_budget( + mode, mode_set, reflections_per_parameter, refine_node_positions + ) + + self.model.set_adp_mode( + mode, + aniso_selection if aniso_selection is not None else self.aniso_selection, + n_nodes=n_nodes, + k_neighbors=min(k_neighbors, n_nodes) if is_field else k_neighbors, + refine_node_positions=refine_node_positions, + mode_set=mode_set, + ) + self.adp_mode = mode + self.adp_mode_set = mode_set + + # Targets hold per-atom index tensors keyed off the old parametrization, the + # scales were fitted against the old F_calc, and which ADP restraints even apply + # is a property of the representation -- so the component set changes, not just + # the weights. reset_loss_state is what makes the next access register the new + # set; it also drops the Logger, which holds a reference to the old state and + # would otherwise keep recording into it. + if getattr(self, "adp_target", None) is not None: + self._init_targets() + self.reset_loss_state() + + n_par = sum(p.numel() for p in self.model.parameters_of_types(("adp", "u"))) + applied = dict( + mode=mode, mode_set=mode_set, n_nodes=n_nodes, n_adp_parameters=int(n_par), + reflections_per_parameter=( + float(self.data.work.n) / n_par if n_par else float("inf") + ), + ) + if self.verbose > 0: + label = mode if mode_set is None else f"{mode}/{mode_set}" + print( + f"ADP representation: {label}" + + (f", {n_nodes} nodes" if is_field else "") + + f", {n_par} parameters, " + f"{applied['reflections_per_parameter']:.1f} work reflections each" + ) + return applied + def _init_targets(self, xray_mode: str = None): """Build the x-ray, geometry and ADP targets and initialise the scales. @@ -487,6 +752,38 @@ def reset_loss_state(self) -> None: self._loss_state = None self._logger = None + def set_hydrogen_mode(self, mode: str) -> "Refinement": + """Switch the model's hydrogen parametrisation and reset the engine state. + + Parameters + ---------- + mode : str + ``"riding"`` or ``"free"``; see :meth:`Model.set_hydrogen_mode`. + + Returns + ------- + Refinement + Self, for chaining. + + Notes + ----- + The coordinate wrapper is replaced, so cached optimizers and the persistent + ``LossState`` are dropped and rebuilt on the next step. Call between macro + cycles, never inside one. + """ + n_atoms = len(self.model.pdb) + self.model.set_hydrogen_mode(mode) + if ( + len(self.model.pdb) != n_atoms + and getattr(self, "adp_target", None) is not None + ): + self._init_targets() + persistent = getattr(self, "_persistent_optimizers", None) + if persistent is not None: + persistent.clear() + self.reset_loss_state() + return self + def refine_scaler(self): """Refit the scaler against the current model. @@ -596,7 +893,7 @@ def _sync_model_cell_to_data( f" data cell: {d}", stacklevel=2, ) - self.model.cell = self.reflection_data.cell + self.model.cell = self.reflection_data.cell.clone() self.model.reset_cache() def parameters(self, recurse: bool = True): @@ -851,8 +1148,8 @@ def collect_deposition_metadata(self, metadata=None): return metadata.merge(refinement_meta) # Merge with pass-through headers from input file - if hasattr(self.model, "_input_file") and self.model._input_file: - input_file = self.model._input_file + if self.model.ctx.input_file: + input_file = self.model.ctx.input_file if input_file.endswith(".pdb"): input_meta = RefinementMetadata.from_pdb_file(input_file) elif input_file.endswith((".cif", ".mmcif")): @@ -980,7 +1277,7 @@ def extract_submodule_state(state_dict: dict, prefix: str) -> dict: scaler = Scaler(model, reflection_data, verbose=verbose, device=device) # Create Restraints with model (required for proper setup) - from torchref.restraints import Restraints + from torchref.topology.restraints import Restraints restraints = Restraints(model, verbose=verbose) @@ -1015,7 +1312,7 @@ def extract_submodule_state(state_dict: dict, prefix: str) -> dict: instance.scaler.set_model_and_data(instance.model, instance.reflection_data) # Initialize targets if model is available - if instance.model is not None and instance.model.initialized: + if instance.model is not None and instance.model.ctx.initialized: try: instance._init_targets() except Exception as e: diff --git a/torchref/refinement/loss_state.py b/torchref/refinement/loss_state.py index 9df09fd0..dee6b00f 100644 --- a/torchref/refinement/loss_state.py +++ b/torchref/refinement/loss_state.py @@ -21,7 +21,7 @@ import torch from torch import nn -from torchref.config import canonical_device, get_default_device +from torchref.config import canonical_device, get_default_device, get_float_dtype from torchref.utils.autograd_introspection import collect_loss_leaves, _iter_roots from torchref.utils.device_mixin import DeviceMovementMixin from torchref.utils.loss_validation import validate_loss @@ -349,7 +349,7 @@ def compile_aggregate(self, **compile_kwargs) -> "LossState": device = self.device def _compiled_fn(): - total = torch.tensor(0.0, device=device) + total = torch.tensor(0.0, dtype=get_float_dtype(), device=device) for fn, w in zip(fns, weights): total = total + w * fn() return total @@ -407,7 +407,7 @@ def aggregate(self, log_values: bool = False) -> torch.Tensor: self.new_entry() self._losses.clear() - total = torch.tensor(0.0, device=self.device) + total = torch.tensor(0.0, dtype=get_float_dtype(), device=self.device) # --- compiled group --- # Skipped when log_values=True: the fused closure does not expose diff --git a/torchref/refinement/model_error_estimation/_shells.py b/torchref/refinement/model_error_estimation/_shells.py new file mode 100644 index 00000000..3d6e9a7d --- /dev/null +++ b/torchref/refinement/model_error_estimation/_shells.py @@ -0,0 +1,187 @@ +"""Resolution-shell machinery shared by the model-error estimators. + +Equal-count shells over ``d*^2``, atomic-free segment sums, linear interpolation of +per-shell values back to reflections, and DerSimonian-Laird shrinkage of noisy per-shell +estimates toward a weighted straight line. :mod:`.sigma_a` and :mod:`.sigma_d` both +build on these; ``estimate_beta`` keeps its own module-level aliases so that its body +resolves the same globals it always did. + +Plain tensors in and out. Every result lives on the device of its inputs, and float +work happens in the dtype of the inputs, so callers control both by what they pass. +""" + +from functools import lru_cache + +import torch + + +@lru_cache(maxsize=8) +def segment_layout(lengths: tuple[int, ...], device_str: str): + """``(index, mask)`` placing contiguous segments on a padded ``(n_seg, max_len)`` grid. + + Cached: the sigma_A solve reduces ``n_grid * n_stages`` times over one layout. + ``lengths`` is a tuple so it can be a cache key. + """ + device = torch.device(device_str) + # dtype-ok: segment lengths for cumsum offsets/gather index; PyTorch requires int64 + L = torch.tensor(lengths, dtype=torch.long, device=device) + total = int(L.sum()) + max_len = int(L.max()) if L.numel() else 0 + # dtype-ok: zero offset concatenated into gather index; PyTorch requires int64 + zero = torch.zeros(1, dtype=torch.long, device=device) + starts = torch.cat([zero, L.cumsum(0)[:-1]]) + ar = torch.arange(max_len, device=device).reshape(1, max_len) + # Clamp keeps the gather in bounds for the padding slots; `mask` zeroes them anyway. + index = (starts.reshape(-1, 1) + ar).clamp(max=max(total - 1, 0)) + mask = ar < L.reshape(-1, 1) + return index, mask + + +def segsum(x: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: + """Sum ``x`` over contiguous segments, reducing along a padded trailing axis. + + Replaces ``torch.segment_reduce``, which is unimplemented on MPS. Keeps the properties + that op was chosen for: atomic-free, one fixed reduction order per segment, so the + result is bit-stable run to run and does not depend on ``scatter_add``'s CUDA atomicAdd + accumulation order (see ``tests/unit/refinement/test_estimate_beta_determinism.py``). + + Deliberately NOT ``cumsum[end] - cumsum[start]``, the usual contiguous-segment trick: + that recovers each shell sum by subtracting two running totals of the whole array, + reintroducing the large-minus-large these estimators are written to avoid. + + ``x`` reduces over its last axis, so a leading batch dimension is handled in one call. + Segments differ in length by at most one element, so the padding overhead is at most + ``n_seg`` slots. + """ + index, mask = segment_layout(tuple(int(v) for v in lengths), str(x.device)) + return (x[..., index] * mask.to(x.dtype)).sum(dim=-1) + + +def interp_in_dss( + dss_all: torch.Tensor, bin_dss: torch.Tensor, vals: torch.Tensor +) -> torch.Tensor: + """Linear interpolation of per-bin ``vals`` (at ``bin_dss``) to all reflections by + their ``d_star_sq``; clamp-to-edge outside the range.""" + n_bins = bin_dss.numel() + if n_bins == 1: + return torch.full_like(dss_all, float(vals[0])) + idx = torch.searchsorted(bin_dss, dss_all).clamp(1, n_bins - 1) + x0 = bin_dss[idx - 1] + x1 = bin_dss[idx] + wlin = ((dss_all - x0) / (x1 - x0).clamp(min=1e-30)).clamp(0.0, 1.0) + return (1 - wlin) * vals[idx - 1] + wlin * vals[idx] + + +def equal_count_shells( + dss: torch.Tensor, *, per_bin: int, min_bins: int, min_per_bin: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Equal-count resolution shells over ``dss``, sorted ascending. + + Parameters + ---------- + dss : torch.Tensor + ``d*^2`` of the reflections entering the fit, shape ``(n,)``, in A^-2. + per_bin : int + Target reflections per shell. + min_bins, min_per_bin : int + Floor on the shell count for sparse sets: at least ``min_bins`` shells as long + as each still holds ``min_per_bin`` reflections. + + Returns + ------- + tuple + ``(order, seg, seg_lengths, n_bins)``. ``order`` sorts ``dss`` ascending (a stable + sort, so tied values bin identically on every backend); ``seg`` is the shell index + of each sorted reflection, a non-decreasing ramp; ``seg_lengths`` the count per + shell. + """ + n = int(dss.numel()) + order = torch.argsort(dss, stable=True) + n_by_count = max(1, n // per_bin) + n_cap = max(1, n // min_per_bin) + n_bins = max(n_by_count, min(min_bins, n_cap)) + seg = ( + torch.arange(n, device=dss.device) * n_bins + ) // n # dtype-ok: bincount input; PyTorch requires int64 + seg_lengths = torch.bincount(seg, minlength=n_bins) + return order, seg, seg_lengths, n_bins + + +def dl_shrink_to_line( + y: torch.Tensor, + var: torch.Tensor, + x: torch.Tensor, + *, + slope_min: float | None = None, + slope_max: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float, float]: + """Shrink noisy per-shell values toward a weighted straight line in ``x``. + + DerSimonian-Laird shrinkage toward a two-parameter line fitted across all shells, + which the two parameters determine far better than any one shell is determined:: + + fit y = a + b*x, weights 1/var + tau^2 = DL between-shell variance about the line, weights 1/var + w_i = var_i / (var_i + tau^2) + y_i <- (1 - w_i)*y_i + w_i*line_i + + ``tau^2`` is the size of the dataset-specific residual the line does not capture, so + ``w_i -> 0`` where that residual is real and large and ``w_i -> 1`` where the shell is + badly determined. One shot, no iteration: the target is a fixed line. Weights are + ``1/var``, never counts, because count weighting lets high-``var`` shells dominate + ``Q`` and veto shrinkage entirely. + + Parameters + ---------- + y, var, x : torch.Tensor + Per-shell value, its sampling variance and the abscissa (``d*^2``), shape + ``(k,)``. A shell with non-finite ``y`` or ``var`` (or ``var <= 0``) takes no part + in the fit and is replaced by the line outright (``w = 1``). + slope_min, slope_max : float, optional + Clamp on the fitted slope ``b``; ``a`` is refitted after the clamp so the line + still passes through the weighted centroid. + + Returns + ------- + tuple + ``(y_shrunk, w, tau_sq, a, b)``. With fewer than four usable shells, or when the + slope is unidentifiable (all shells at one ``x``), the input is returned + unchanged with ``w = 0``, ``tau_sq = 0`` and NaN line coefficients. + """ + nan = float("nan") + usable = torch.isfinite(y) & torch.isfinite(var) & (var > 0) + k = int(usable.sum()) + # Two fitted parameters need at least two residual degrees of freedom. + if k < 4: + return y, torch.zeros_like(y), y.new_zeros(()), nan, nan + + wt = torch.where(usable, 1.0 / var.clamp(min=1e-30), torch.zeros_like(var)) + yz = torch.where(usable, y, torch.zeros_like(y)) + S = wt.sum() + Sx = (wt * x).sum() + Sxx = (wt * x * x).sum() + Sy = (wt * yz).sum() + Sxy = (wt * x * yz).sum() + det = S * Sxx - Sx * Sx + # Relative, not absolute: `det` is a difference of two ~`S**2 * x**2` terms, so on a + # degenerate input it lands at the cancellation floor, not near zero. + if float(det.abs()) <= 1e-12 * float((S * Sxx).abs()): + return y, torch.zeros_like(y), y.new_zeros(()), nan, nan + b = (S * Sxy - Sx * Sy) / det + if slope_min is not None: + b = b.clamp(min=slope_min) + if slope_max is not None: + b = b.clamp(max=slope_max) + a = (wt * (yz - b * x)).sum() / S.clamp(min=1e-30) + line = a + b * x + + resid = torch.where(usable, yz - line, torch.zeros_like(y)) + Q = (wt * resid * resid).sum() + dof = float(k - 2) # two parameters were fitted + c = (S - (wt * wt).sum() / S.clamp(min=1e-30)).clamp(min=1e-30) + # Q < k-2 means the scatter about the line is SMALLER than the noise alone predicts, + # i.e. no evidence of structure the line is missing -> tau^2 = 0 -> take the line. + tau_sq = ((Q - dof) / c).clamp(min=0.0) + w = torch.where(usable, var / (var + tau_sq).clamp(min=1e-30), torch.ones_like(var)) + out = (1.0 - w) * torch.where(usable, y, line) + w * line + return out, w, tau_sq, float(a), float(b) diff --git a/torchref/refinement/model_error_estimation/sigma_a.py b/torchref/refinement/model_error_estimation/sigma_a.py index 79f33405..dab6d53a 100644 --- a/torchref/refinement/model_error_estimation/sigma_a.py +++ b/torchref/refinement/model_error_estimation/sigma_a.py @@ -17,14 +17,41 @@ import math from dataclasses import dataclass -from functools import lru_cache -from typing import Optional, Tuple +from typing import Optional import torch -from torchref.base.french_wilson import epsilon_from_hkl # noqa: F401 (re-export) from torchref.config import get_float_dtype +from ._shells import interp_in_dss as _interp_in_dss +from ._shells import segment_layout as _segment_layout # noqa: F401 +from ._shells import segsum as _segsum + +def epsilon_from_hkl(hkl: torch.Tensor, spacegroup) -> torch.Tensor: + """Per-reflection epsilon, tolerating a missing space group. + + Thin adapter over :meth:`~torchref.symmetry.symmetry.Symmetry.epsilon`, which owns + the multiplicity count. It exists because reflection data may carry no space group + at all, and every consumer here would otherwise repeat the same guard. + + Parameters + ---------- + hkl : torch.Tensor + Miller indices, shape ``(N, 3)``. + spacegroup : Symmetry or None + The group. ``None`` means no symmetry information, which yields ones -- the + same answer P1 gives. + + Returns + ------- + torch.Tensor + Multiplicities, shape ``(N,)``, at the configured float dtype, on ``hkl``'s + device. + """ + if spacegroup is None or not hasattr(spacegroup, "epsilon"): + return torch.ones(hkl.shape[0], device=hkl.device, dtype=get_float_dtype()) + return spacegroup.epsilon(hkl) + # --- sigma_A estimator constants ------------------------------------------------- #: Upper bound on the per-shell ``sigma_A``, i.e. the floor on the model-error variance at #: ``(1 - SIGMA_A_MAX**2) * Sigma_N``. @@ -230,47 +257,6 @@ def _rice_nll_reduced( return torch.where(centric, cen, acen) -@lru_cache(maxsize=8) -def _segment_layout(lengths: Tuple[int, ...], device_str: str): - """``(index, mask)`` placing contiguous segments on a padded ``(n_seg, max_len)`` grid. - - Cached: ``_solve_sigma_a`` reduces ``n_grid * n_stages`` times over one layout. - ``lengths`` is a tuple so it can be a cache key. - """ - device = torch.device(device_str) - L = torch.tensor(lengths, dtype=torch.long, device=device) - total = int(L.sum()) - max_len = int(L.max()) if L.numel() else 0 - zero = torch.zeros(1, dtype=torch.long, device=device) - starts = torch.cat([zero, L.cumsum(0)[:-1]]) - ar = torch.arange(max_len, device=device).reshape(1, max_len) - # Clamp keeps the gather in bounds for the padding slots; `mask` zeroes them anyway. - index = (starts.reshape(-1, 1) + ar).clamp(max=max(total - 1, 0)) - mask = ar < L.reshape(-1, 1) - return index, mask - - -def _segsum(x: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: - """Sum ``x`` over contiguous segments, reducing along a padded trailing axis. - - Replaces ``torch.segment_reduce``, which is unimplemented on MPS. Keeps the properties - that op was chosen for: atomic-free, one fixed reduction order per segment, so the - result is bit-stable run to run and does not depend on ``scatter_add``'s CUDA atomicAdd - accumulation order (the original GPU non-determinism bug -- see - ``tests/unit/refinement/test_estimate_beta_determinism.py``). - - Deliberately NOT ``cumsum[end] - cumsum[start]``, the usual contiguous-segment trick: - that recovers each shell sum by subtracting two running totals of the whole array, - reintroducing the large-minus-large this module is written to avoid. - - ``x`` reduces over its last axis, so a leading batch dimension (the grid candidates) - is handled in one call. Segments here differ in length by at most one element, so the - padding overhead is at most ``n_seg`` slots. - """ - index, mask = _segment_layout(tuple(int(v) for v in lengths), str(x.device)) - return (x[..., index] * mask.to(x.dtype)).sum(dim=-1) - - def _grid_ladder(n: int, ratio: float, device, dtype) -> torch.Tensor: """``ratio ** (k - (n-1)/2)`` for ``k`` in ``[0, n)``, built from Python floats. @@ -523,7 +509,7 @@ def estimate_beta( out_dtype = F_obs.dtype dtype = torch.promote_types(get_float_dtype(), out_dtype) - if dtype == torch.float64 and device.type == "mps": + if dtype == torch.float64 and device.type == "mps": # dtype-ok: MPS capability guard, not an allocation raise RuntimeError( "MPS has no float64; set the defaults float dtype to float32 or use CPU" ) @@ -710,19 +696,6 @@ def segsum(x): ) -def _interp_in_dss(dss_all, bin_dss, vals): - """Linear interpolation of per-bin ``vals`` (at ``bin_dss``) to all - reflections by their ``d_star_sq``; clamp-to-edge outside the range.""" - n_bins = bin_dss.numel() - if n_bins == 1: - return torch.full_like(dss_all, float(vals[0])) - idx = torch.searchsorted(bin_dss, dss_all).clamp(1, n_bins - 1) - x0 = bin_dss[idx - 1] - x1 = bin_dss[idx] - wlin = ((dss_all - x0) / (x1 - x0).clamp(min=1e-30)).clamp(0.0, 1.0) - return (1 - wlin) * vals[idx - 1] + wlin * vals[idx] - - # ===================================================================== # Stateful estimator (owned by the consuming target, not the scaler) # ===================================================================== diff --git a/torchref/refinement/model_error_estimation/sigma_d.py b/torchref/refinement/model_error_estimation/sigma_d.py new file mode 100644 index 00000000..8ead3617 --- /dev/null +++ b/torchref/refinement/model_error_estimation/sigma_d.py @@ -0,0 +1,795 @@ +"""Difference-driven error estimation: the expected difference power ``sigma_D``. + +A light-minus-dark difference coefficient ``dF_obs = dF_true + noise`` carries a true +signal whose power ``S = E[dF_true**2]`` varies with resolution and with the dark +amplitude, and a measurement noise ``sigma_diff**2`` the merge reports. The best linear +estimate of ``dF_true`` from ``dF_obs`` is ``w * dF_obs`` with the Wiener weight +``w = S / (S + sigma_diff**2)``, so a difference map needs ``S`` per reflection. Inverse +variance alone weights by precision and treats every reflection as carrying the same +expected difference, which suppresses the strong reflections whose difference power is +ten to seventy times that of weak ones. + +``S`` needs no half datasets. Per resolution shell the second moment of the observed +differences is ``B = S + S2`` with ``S2`` the mean measurement variance, so +``Sigma_N = B - S2`` is the expected true difference power, the same identity +:mod:`.sigma_a` uses for amplitudes. Within a shell the power follows the dark amplitude +as ``(F_dark / )**gamma`` with one fitted exponent, carried by the per-reflection +multiplier ``epsilon`` exactly as the reflection multiplicity is. With a difference model +``dF_calc`` the shell moments also give the Gaussian coupling ``alpha = / +`` and the unexplained power ``beta_model = Sigma_N - alpha**2 Sigma_P``, the +extra variance a difference likelihood adds to ``sigma_diff**2``. + +Differences are signed and small, so the statistics are Gaussian throughout: there is no +Rice branch and no centric distinction. Plain tensors in and out, no ``ReflectionData`` +or ``Scaler`` coupling, so :mod:`torchref.maps` and :mod:`torchref.cli` can import this +module without closing an import cycle. Every result lives on the device of its inputs. +""" + +import math +from dataclasses import dataclass + +import torch + +from torchref.config import get_float_dtype + +from ._shells import equal_count_shells, interp_in_dss, segsum +from .sigma_a import SHRINK_ENABLED + +# --- sigma_D estimator constants ------------------------------------------------- +#: Shell construction, matching the ``estimate_beta`` defaults so a sigma_A and a sigma_D +#: fit on the same reflections use the same shells. +PER_BIN = 140 +MIN_BINS = 5 +MIN_PER_BIN = 40 +#: Exponent of the dark-amplitude power law when it cannot be fitted. The fitted value on +#: the small-molecule, TD1 and bacteriorhodopsin A/B campaigns was 0.8-1.2, so 1.0 (power +#: proportional to the amplitude) is the informed default; 0.0 would be the pure shell model. +GAMMA_DEFAULT = 1.0 +#: Bounds on the fitted exponent. Outside [0, 2] the class means are dominated by one +#: amplitude decile and the regression is on noise; 2 is proportionality to the intensity. +GAMMA_BOUNDS = (0.0, 2.0) +#: Dark-amplitude quantile classes per shell for the exponent regression. Four keeps at +#: least 35 reflections per class at the default shell size; more classes did not move +#: the fitted exponent on the campaign data. +N_F_CLASSES = 4 +#: Minimum number of usable (shell, class) cells before the exponent is fitted at all. +GAMMA_MIN_CLASSES = 8 +#: Minimum reflections in a class for its moment to enter the exponent regression. +MIN_PER_CLASS = 5 +#: Largest standard error at which a fitted exponent is used. Above it the class +#: moments are noise (a null dataset gives se ~ 1 or more), and the default is safer than +#: a random exponent that would redistribute weight between amplitude classes. +GAMMA_SE_MAX = 0.5 +#: Gauss-Newton iterations for the decaying-curve fit; the problem is two-parameter and +#: well conditioned, so this is far more than it needs. +CURVE_ITERS = 60 +#: Floor on ``F_dark`` relative to its shell mean before the power law is evaluated, so a +#: zero or near-zero dark amplitude cannot delete the expected difference power. +F_FLOOR_FRAC = 0.05 +#: Positive floor used where a logarithm of a clamped-to-zero power is needed. +_TINY = 1e-30 + + +@dataclass(frozen=True) +class SigmaDConfig: + """The estimator's knobs, as one value. + + ``gamma=None`` fits the dark-amplitude exponent; a float fixes it. ``shrink=None`` + means the module default shared with sigma_A, normalised here so consumers never + handle ``None``. Frozen, so two consumers sharing a config cannot drift apart. + """ + + gamma: float | None = None + shrink: bool | None = None + + def __post_init__(self): + if self.gamma is not None: + g = float(self.gamma) + lo, hi = GAMMA_BOUNDS + if not (lo <= g <= hi): + raise ValueError(f"gamma must lie in {GAMMA_BOUNDS}, got {g}") + object.__setattr__(self, "gamma", g) + object.__setattr__( + self, "shrink", bool(SHRINK_ENABLED if self.shrink is None else self.shrink) + ) + + +@dataclass(frozen=True) +class SigmaDShells: + """Per-shell output of :func:`estimate_sigma_d`. + + All power quantities are ``epsilon``-reduced and in ``F**2`` units of the input. + + Attributes + ---------- + B, S2 + Raw second moment of the observed differences and mean measurement variance. + Sigma_N_raw, Sigma_N + Expected true difference power ``(B - S2)`` clamped at zero, before and after the + shrinkage of the signed value toward a decaying curve ``exp(a + b d*^2)`` fitted + to every shell. ``Sigma_N`` is what the weights use. + Sigma_P, C, alpha, beta_model + Model power, cross moment, Gaussian coupling ``C / Sigma_P`` and unexplained power + ``(Sigma_N - alpha**2 Sigma_P)`` clamped at zero. Without a model ``Sigma_P`` and + ``C`` are zero, ``alpha`` one and ``beta_model == Sigma_N``. + counts, bin_dss + Reflections per shell and its mean ``d*^2`` in A^-2, the interpolation abscissa. + bin_log_fbar, bin_log_z + Log of the shell-mean dark amplitude and of the shell mean of + ``(F / Fbar)**gamma``, so the per-reflection multiplier + ``(F / Fbar)**gamma / Z`` has shell mean one and ``Sigma_N`` stays the shell mean + of the per-reflection power. Zero when no dark amplitude was supplied. + shrink_w, tau, curve_a, curve_b + Shrinkage weight per shell, the between-shell sd about the fitted curve and its + coefficients ``exp(a + b d*^2)`` (NaN when no curve was fitted; the curve is then + zero everywhere). + gamma, gamma_fitted + The exponent used and whether it was fitted rather than fixed or defaulted. + has_model, degenerate, all_zero + Whether ``dF_calc`` was supplied, whether fewer than two usable reflections + existed, and whether every shell's ``Sigma_N`` is zero (weights would vanish). + diagnostics + Counters: ``n_dropped, n_fit, n_shell, n_s2_clamped, n_beta_clamped, + n_f_floored, n_class_dropped, n_class_used, gamma_se, gamma_at_bound, + gamma_reason``. + """ + + B: torch.Tensor + S2: torch.Tensor + Sigma_N_raw: torch.Tensor + Sigma_N: torch.Tensor + Sigma_P: torch.Tensor + C: torch.Tensor + alpha: torch.Tensor + beta_model: torch.Tensor + counts: torch.Tensor + bin_dss: torch.Tensor + bin_log_fbar: torch.Tensor + bin_log_z: torch.Tensor + shrink_w: torch.Tensor + tau: float + curve_a: float + curve_b: float + gamma: float + gamma_fitted: bool + has_model: bool + degenerate: bool + all_zero: bool + diagnostics: dict + + +@dataclass(frozen=True) +class SigmaDEstimate: + """Everything a consumer needs from one estimate, per reflection and detached. + + Attributes + ---------- + S + Expected true difference power ``epsilon * Sigma_N(d*^2) * g(F_dark)``. + sigma_sq + The measurement variance the weight was formed with (``sigma_diff**2``). + w + Wiener weight ``S / (S + sigma_sq)`` in ``[0, 1)``, not normalised. + alpha, beta_model + Coupling and unexplained power, interpolated per shell; ``beta_model`` carries + the same ``epsilon * g`` multiplier as ``S``. + epsilon + The multiplicity actually applied. + shells + The :class:`SigmaDShells` this was interpolated from. + """ + + S: torch.Tensor + sigma_sq: torch.Tensor + w: torch.Tensor + alpha: torch.Tensor + beta_model: torch.Tensor + epsilon: torch.Tensor + shells: SigmaDShells + + +def _working_dtype(t: torch.Tensor) -> torch.dtype: + dtype = torch.promote_types(get_float_dtype(), t.dtype) + # dtype-ok: MPS capability guard, not an allocation + if dtype == torch.float64 and t.device.type == "mps": + raise RuntimeError( + "MPS has no float64; set the defaults float dtype to float32 or use CPU" + ) + return dtype + + +def _degenerate( + delta_obs: torch.Tensor, gamma: float, has_model: bool, diagnostics: dict, out_dtype +) -> SigmaDShells: + """One conservative shell: the mean squared difference as the power, alpha one.""" + ok = torch.isfinite(delta_obs) + b = (delta_obs[ok] ** 2).mean() if bool(ok.any()) else delta_obs.new_ones(()) + one = torch.ones(1, device=delta_obs.device, dtype=out_dtype) + zero = torch.zeros(1, device=delta_obs.device, dtype=out_dtype) + b1 = (one * b).to(out_dtype) + return SigmaDShells( + B=b1, + S2=zero, + Sigma_N_raw=b1, + Sigma_N=b1, + Sigma_P=zero, + C=zero, + alpha=one, + beta_model=b1, + counts=zero, + bin_dss=zero, + bin_log_fbar=zero, + bin_log_z=zero, + shrink_w=zero, + tau=0.0, + curve_a=float("nan"), + curve_b=float("nan"), + gamma=gamma, + gamma_fitted=False, + has_model=has_model, + degenerate=True, + all_zero=False, + diagnostics=diagnostics, + ) + + +def _fit_gamma( + d2e: torch.Tensor, + s2e: torch.Tensor, + log_ratio: torch.Tensor, + seg: torch.Tensor, + n_bins: int, +) -> tuple[float, float, bool, int, int, str]: + """Fit the dark-amplitude exponent from within-shell amplitude classes. + + Each shell is split into ``N_F_CLASSES`` quantile classes of the dark amplitude. A + class contributes ``log( - )`` against its mean log amplitude + ratio when that difference power is positive. One slope is fitted across all shells + with the shell means removed (fixed effects), weighted by ``n_c / 2``: the log of a + mean of ``n`` squared Gaussians has variance ``2 / n``. + + Returns ``(gamma, gamma_se, at_bound, n_used, n_dropped, reason)``; ``reason`` is + ``"fitted"`` or names why the default was taken. + """ + xs, ys, ws, shell_id = [], [], [], [] + n_dropped = 0 + for k in range(n_bins): + in_shell = torch.nonzero(seg == k, as_tuple=True)[0] + n_k = int(in_shell.numel()) + if n_k < N_F_CLASSES * MIN_PER_CLASS: + n_dropped += N_F_CLASSES + continue + order = torch.argsort(log_ratio[in_shell], stable=True) + idx = in_shell[order] + cls = ( + torch.arange(n_k, device=seg.device) * N_F_CLASSES + ) // n_k # dtype-ok: bincount input; PyTorch requires int64 + lengths = torch.bincount(cls, minlength=N_F_CLASSES).to(d2e.dtype) + m = (segsum(d2e[idx], lengths) - segsum(s2e[idx], lengths)) / lengths + xc = segsum(log_ratio[idx], lengths) / lengths + keep = (m > 0) & (lengths >= MIN_PER_CLASS) + n_dropped += int((~keep).sum()) + if int(keep.sum()) < 2: + continue + xs.append(xc[keep]) + ys.append(torch.log(m[keep])) + ws.append(lengths[keep] / 2.0) + shell_id.append(torch.full_like(xc[keep], float(k))) + if not xs: + return GAMMA_DEFAULT, float("nan"), False, 0, n_dropped, "too_few_classes" + x = torch.cat(xs) + y = torch.cat(ys) + w = torch.cat(ws) + sid = torch.cat(shell_id) + n_used = int(x.numel()) + if n_used < GAMMA_MIN_CLASSES: + return GAMMA_DEFAULT, float("nan"), False, n_used, n_dropped, "too_few_classes" + # Remove each shell's weighted mean from x and y: the slope is then estimated from + # within-shell contrasts only, so shell-to-shell differences in power cannot leak in. + xc = x.clone() + yc = y.clone() + n_shells_used = 0 + for k in torch.unique(sid): + s = sid == k + n_shells_used += 1 + wk = w[s] + xc[s] = x[s] - (wk * x[s]).sum() / wk.sum() + yc[s] = y[s] - (wk * y[s]).sum() / wk.sum() + sxx = (w * xc * xc).sum() + if float(sxx) <= 0.0: + return ( + GAMMA_DEFAULT, + float("nan"), + False, + n_used, + n_dropped, + "no_amplitude_spread", + ) + gamma = float((w * xc * yc).sum() / sxx) + dof = n_used - n_shells_used - 1 + if dof > 0: + resid = yc - gamma * xc + s2 = float((w * resid * resid).sum() / dof) + gamma_se = math.sqrt(max(s2, 0.0) / float(sxx)) + else: + gamma_se = float("nan") + if not math.isfinite(gamma_se) or gamma_se > GAMMA_SE_MAX: + return GAMMA_DEFAULT, gamma_se, False, n_used, n_dropped, "too_uncertain" + lo, hi = GAMMA_BOUNDS + clamped = min(max(gamma, lo), hi) + return clamped, gamma_se, clamped != gamma, n_used, n_dropped, "fitted" + + +def _fit_decay(y: torch.Tensor, var: torch.Tensor, x: torch.Tensor): + """Weighted fit of ``exp(a + b x)``, ``b <= 0``, to signed per-shell power. + + Works on the signed ``B - S2`` of every shell, so shells whose power is zero or + negative by sampling noise pull the curve down instead of being ignored: on a null + dataset the curve goes to zero rather than to the winner's curse of the positive + shells. Gauss-Newton with step halving on the weighted least squares; the fit is + two-parameter and well conditioned. + + Returns ``(curve, a, b)``; ``curve`` is zeros with NaN coefficients when fewer than + four shells are usable or the weighted mean power is not positive. + """ + nan = float("nan") + usable = torch.isfinite(y) & torch.isfinite(var) & (var > 0) + if int(usable.sum()) < 4: + return torch.zeros_like(y), nan, nan + w = torch.where(usable, 1.0 / var.clamp(min=_TINY), torch.zeros_like(var)) + yz = torch.where(usable, y, torch.zeros_like(y)) + mean = float((w * yz).sum() / w.sum()) + if mean <= 0.0: + return torch.zeros_like(y), nan, nan + a = torch.tensor(math.log(mean), dtype=y.dtype, device=y.device) + b = torch.zeros((), dtype=y.dtype, device=y.device) + + def loss(a_, b_): + r = yz - torch.exp(a_ + b_ * x) + return float((w * r * r).sum()) + + current = loss(a, b) + for _ in range(CURVE_ITERS): + f = torch.exp(a + b * x) + r = yz - f + # Jacobian of f with respect to (a, b): f and f*x. + j_a, j_b = f, f * x + g = torch.stack([(w * r * j_a).sum(), (w * r * j_b).sum()]) + h = torch.stack( + [ + torch.stack([(w * j_a * j_a).sum(), (w * j_a * j_b).sum()]), + torch.stack([(w * j_a * j_b).sum(), (w * j_b * j_b).sum()]), + ] + ) + h = ( + h + + 1e-12 * torch.eye(2, dtype=h.dtype, device=h.device) * h.diagonal().max() + ) + step = torch.linalg.solve(h, g) + scale = 1.0 + improved = False + for _ in range(12): + a_new = a + scale * step[0] + b_new = (b + scale * step[1]).clamp(max=0.0) + new = loss(a_new, b_new) + if new < current: + a, b, current, improved = a_new, b_new, new, True + break + scale *= 0.5 + if not improved or float(step.abs().max()) < 1e-9: + break + return torch.exp(a + b * x), float(a), float(b) + + +def estimate_sigma_d( + delta_obs: torch.Tensor, + sigma_diff: torch.Tensor, + epsilon: torch.Tensor | None, + d_star_sq: torch.Tensor, + f_dark: torch.Tensor | None, + fit_mask: torch.Tensor, + *, + delta_calc: torch.Tensor | None = None, + gamma: float | None = None, + shrink: bool | None = None, + per_bin: int = PER_BIN, + min_bins: int = MIN_BINS, + min_per_bin: int = MIN_PER_BIN, +) -> SigmaDShells: + """Per-shell expected difference power, with the dark-amplitude exponent and, + given a difference model, its coupling and unexplained power. + + Runs under ``torch.no_grad()``. The working dtype is the wider of the configured + float dtype and ``delta_obs.dtype``; results are cast back to ``delta_obs.dtype``. + + Parameters + ---------- + delta_obs : torch.Tensor + Signed observed differences ``F_light - F_dark``, shape ``(N,)``, on one common + amplitude scale. + sigma_diff : torch.Tensor + Propagated uncertainty of ``delta_obs``, shape ``(N,)``, same units. + epsilon : torch.Tensor or None + Reflection multiplicity, shape ``(N,)``; ``None`` means ones. + d_star_sq : torch.Tensor + ``1/d**2`` per reflection, shape ``(N,)``, in A^-2. + f_dark : torch.Tensor or None + Dark amplitude for the power law, shape ``(N,)``. ``None`` disables the amplitude + dependence (``gamma`` reported as the default with reason ``"no_f_dark"``). + fit_mask : torch.Tensor + Boolean ``(N,)``: which reflections enter the fit. + delta_calc : torch.Tensor, optional + Model differences ``|F_calc_light| - |F_calc_dark|``, shape ``(N,)``, on the + observed scale. Enables ``alpha`` and ``beta_model``. + gamma : float, optional + Fix the exponent instead of fitting it. + shrink : bool, optional + Shrink the signed shell power toward a decaying curve in ``d*^2``; default the + module setting. + per_bin, min_bins, min_per_bin : int, optional + Shell construction, see :func:`~._shells.equal_count_shells`. + + Returns + ------- + SigmaDShells + One frozen record of per-shell quantities plus counters. + """ + device = delta_obs.device + out_dtype = delta_obs.dtype + dtype = _working_dtype(delta_obs) + shrink = bool(SHRINK_ENABLED if shrink is None else shrink) + if gamma is not None: + lo, hi = GAMMA_BOUNDS + if not (lo <= float(gamma) <= hi): + raise ValueError(f"gamma must lie in {GAMMA_BOUNDS}, got {gamma}") + + with torch.no_grad(): + d_all = delta_obs.reshape(-1).to(dtype) + s_all = sigma_diff.reshape(-1).to(dtype) + x_all = d_star_sq.reshape(-1).to(dtype) + e_all = ( + epsilon.reshape(-1).to(dtype) + if epsilon is not None + else torch.ones_like(d_all) + ) + has_f = f_dark is not None + f_all = f_dark.reshape(-1).to(dtype) if has_f else None + has_model = delta_calc is not None + c_all = delta_calc.reshape(-1).to(dtype) if has_model else None + + finite = ( + torch.isfinite(d_all) + & torch.isfinite(s_all) + & torch.isfinite(x_all) + & torch.isfinite(e_all) + & (s_all >= 0.0) + & (e_all > 0.0) + ) + if has_f: + finite &= torch.isfinite(f_all) + if has_model: + finite &= torch.isfinite(c_all) + fit = fit_mask.reshape(-1).to(torch.bool) + usable = fit & finite + n_dropped = int((fit & ~finite).sum()) + idx = torch.nonzero(usable, as_tuple=True)[0] + n_fit = int(idx.numel()) + + diagnostics = { + "n_dropped": n_dropped, + "n_fit": n_fit, + "n_shell": 0, + "n_s2_clamped": 0, + "n_beta_clamped": 0, + "n_f_floored": 0, + "n_class_dropped": 0, + "n_class_used": 0, + "gamma_se": float("nan"), + "gamma_at_bound": False, + "gamma_reason": "degenerate", + } + if n_fit < 2: + g = float(gamma) if gamma is not None else GAMMA_DEFAULT + return _degenerate(d_all, g, has_model, diagnostics, out_dtype) + + order, seg, seg_lengths, n_bins = equal_count_shells( + x_all[idx], per_bin=per_bin, min_bins=min_bins, min_per_bin=min_per_bin + ) + sel = idx[order] + d, s, e, x = d_all[sel], s_all[sel], e_all[sel], x_all[sel] + counts = seg_lengths.to(dtype) + d2e = d * d / e + s2e = s * s / e + + B = segsum(d2e, seg_lengths) / counts + S2 = segsum(s2e, seg_lengths) / counts + Sigma_N_raw = (B - S2).clamp(min=0.0) + n_s2_clamped = int((S2 >= B).sum()) + bin_dss = segsum(x, seg_lengths) / counts + + # --- dark-amplitude power law ------------------------------------------- + if has_f: + f = f_all[sel] + fbar = (segsum(f, seg_lengths) / counts).clamp(min=_TINY) + fbar_h = fbar[seg] + floor = F_FLOOR_FRAC * fbar_h + n_f_floored = int((f < floor).sum()) + f_fl = torch.maximum(f, floor) + log_ratio = torch.log(f_fl) - torch.log(fbar_h) + bin_log_fbar = torch.log(fbar) + else: + n_f_floored = 0 + log_ratio = torch.zeros_like(d) + bin_log_fbar = torch.zeros_like(bin_dss) + + if gamma is not None: + g_used, g_se, at_bound, n_used, n_cls_dropped, reason = ( + float(gamma), + float("nan"), + False, + 0, + 0, + "fixed", + ) + fitted = False + elif not has_f: + g_used, g_se, at_bound, n_used, n_cls_dropped, reason = ( + GAMMA_DEFAULT, + float("nan"), + False, + 0, + 0, + "no_f_dark", + ) + fitted = False + else: + g_used, g_se, at_bound, n_used, n_cls_dropped, reason = _fit_gamma( + d2e, s2e, log_ratio, seg, n_bins + ) + fitted = reason == "fitted" + + if has_f: + g_raw = torch.exp(g_used * log_ratio) + Z = (segsum(g_raw, seg_lengths) / counts).clamp(min=_TINY) + bin_log_z = torch.log(Z) + else: + bin_log_z = torch.zeros_like(bin_dss) + + # --- difference model ------------------------------------------------------ + if has_model: + c = c_all[sel] + Sigma_P = segsum(c * c / e, seg_lengths) / counts + C = segsum(d * c / e, seg_lengths) / counts + alpha = C / Sigma_P.clamp(min=_TINY) + else: + Sigma_P = torch.zeros_like(B) + C = torch.zeros_like(B) + alpha = torch.ones_like(B) + + # --- stability shrinkage of the signed power toward a decaying curve --------- + # The signed B - S2 keeps every shell as evidence: a shell below zero by noise + # says the power there is small, and it must count. var(B) is 2 B**2 / n for + # Gaussian differences, so that is the sampling variance of each shell's value. + signed = B - S2 + var_s = 2.0 * B * B / counts + if shrink: + curve, curve_a, curve_b = _fit_decay(signed, var_s, bin_dss) + resid = signed - curve + prec = 1.0 / var_s.clamp(min=_TINY) + Q = (prec * resid * resid).sum() + dof = float(max(int(signed.numel()) - 2, 1)) + c = (prec.sum() - (prec * prec).sum() / prec.sum().clamp(min=_TINY)).clamp( + min=_TINY + ) + # Q < dof means the shells scatter no more than their noise: take the curve. + tau_sq = ((Q - dof) / c).clamp(min=0.0) + shrink_w = var_s / (var_s + tau_sq).clamp(min=_TINY) + Sigma_N = ((1.0 - shrink_w) * signed + shrink_w * curve).clamp(min=0.0) + else: + Sigma_N = Sigma_N_raw + shrink_w, tau_sq = torch.zeros_like(B), B.new_zeros(()) + curve_a = curve_b = float("nan") + Sigma_N = torch.where(torch.isfinite(Sigma_N), Sigma_N, torch.zeros_like(B)) + + beta_model = (Sigma_N - alpha * alpha * Sigma_P).clamp(min=0.0) + n_beta_clamped = int(((Sigma_N - alpha * alpha * Sigma_P) < 0.0).sum()) + all_zero = bool((Sigma_N <= 0.0).all()) + + diagnostics.update( + n_shell=int(n_bins), + n_s2_clamped=n_s2_clamped, + n_beta_clamped=n_beta_clamped, + n_f_floored=n_f_floored, + n_class_dropped=n_cls_dropped, + n_class_used=n_used, + gamma_se=g_se, + gamma_at_bound=at_bound, + gamma_reason=reason, + ) + + to = lambda t: t.to(out_dtype) + return SigmaDShells( + B=to(B), + S2=to(S2), + Sigma_N_raw=to(Sigma_N_raw), + Sigma_N=to(Sigma_N), + Sigma_P=to(Sigma_P), + C=to(C), + alpha=to(alpha), + beta_model=to(beta_model), + counts=to(counts), + bin_dss=to(bin_dss), + bin_log_fbar=to(bin_log_fbar), + bin_log_z=to(bin_log_z), + shrink_w=to(shrink_w), + tau=float(tau_sq.clamp(min=0.0).sqrt()), + curve_a=curve_a, + curve_b=curve_b, + gamma=float(g_used), + gamma_fitted=fitted, + has_model=has_model, + degenerate=False, + all_zero=all_zero, + diagnostics=diagnostics, + ) + + +def sigma_d_per_reflection( + shells: SigmaDShells, + d_star_sq: torch.Tensor, + epsilon: torch.Tensor | None, + f_dark: torch.Tensor | None, + sigma_diff: torch.Tensor, +) -> SigmaDEstimate: + """Interpolate a shell estimate onto reflections and form the Wiener weights. + + Parameters + ---------- + shells : SigmaDShells + The shell estimate. + d_star_sq : torch.Tensor + ``1/d**2`` of the output reflections, shape ``(M,)``, in A^-2. + epsilon : torch.Tensor or None + Multiplicity of the output reflections, shape ``(M,)``; ``None`` means ones. + f_dark : torch.Tensor or None + Dark amplitude of the output reflections for the power law; reflections with a + missing or non-finite value get a multiplier of one. + sigma_diff : torch.Tensor + Propagated uncertainty of the output differences, shape ``(M,)``. Non-finite + entries give a weight of zero. + + Returns + ------- + SigmaDEstimate + Per-reflection, detached fields all of length ``M``. + """ + with torch.no_grad(): + dtype = shells.Sigma_N.dtype + grid = d_star_sq.reshape(-1).to(dtype) + eps = ( + epsilon.reshape(-1).to(dtype) + if epsilon is not None + else torch.ones_like(grid) + ) + sig = sigma_diff.reshape(-1).to(dtype) + if shells.degenerate or shells.bin_dss.numel() == 0: + sigma_n = torch.full_like(grid, float(shells.Sigma_N[0])) + alpha = torch.full_like(grid, float(shells.alpha[0])) + beta_model = torch.full_like(grid, float(shells.beta_model[0])) + g = torch.ones_like(grid) + else: + log_sn = interp_in_dss( + grid, shells.bin_dss, torch.log(shells.Sigma_N.clamp(min=_TINY)) + ) + sigma_n = torch.exp(log_sn) + sigma_n = torch.where( + sigma_n > 10.0 * _TINY, sigma_n, torch.zeros_like(sigma_n) + ) + alpha = interp_in_dss(grid, shells.bin_dss, shells.alpha) + log_bm = interp_in_dss( + grid, shells.bin_dss, torch.log(shells.beta_model.clamp(min=_TINY)) + ) + beta_model = torch.exp(log_bm) + beta_model = torch.where( + beta_model > 10.0 * _TINY, beta_model, torch.zeros_like(beta_model) + ) + if f_dark is not None and shells.gamma != 0.0: + f = f_dark.reshape(-1).to(dtype) + log_fbar = interp_in_dss(grid, shells.bin_dss, shells.bin_log_fbar) + log_z = interp_in_dss(grid, shells.bin_dss, shells.bin_log_z) + fbar = torch.exp(log_fbar) + f_fl = torch.maximum(f, F_FLOOR_FRAC * fbar) + g = torch.exp(shells.gamma * (torch.log(f_fl) - log_fbar) - log_z) + g = torch.where(torch.isfinite(f) & (fbar > 0), g, torch.ones_like(g)) + else: + g = torch.ones_like(grid) + S = eps * sigma_n * g + sigma_sq = sig * sig + w = torch.where( + torch.isfinite(sigma_sq), + S / (S + sigma_sq).clamp(min=_TINY), + torch.zeros_like(S), + ) + return SigmaDEstimate( + S=S.detach(), + sigma_sq=sigma_sq.detach(), + w=w.detach(), + alpha=alpha.detach(), + beta_model=(eps * beta_model * g).detach(), + epsilon=eps.detach(), + shells=shells, + ) + + +class SigmaDEstimator: + """Lazy, cached difference-power estimate. + + Thin stateful wrapper around :func:`estimate_sigma_d` and + :func:`sigma_d_per_reflection`: caches the detached estimate and re-estimates only + after :meth:`reset`. **The owning target must call :meth:`reset` from its + ``maintenance()`` hook**, otherwise the estimate is frozen for the whole run. Holds + no tensors of its own beyond the cache, so it has no device to move. + + Parameters + ---------- + config : SigmaDConfig, optional + Exponent and shrinkage settings; the module defaults when omitted. + """ + + def __init__(self, config: SigmaDConfig | None = None): + self.config = config if config is not None else SigmaDConfig() + self._cache: SigmaDEstimate | None = None + self._shells: SigmaDShells | None = None + + def reset(self) -> None: + """Invalidate the cache so the next :meth:`get` re-estimates.""" + self._cache = None + + @property + def shells(self) -> SigmaDShells | None: + """Last shell estimate, for diagnostics; ``None`` until the first call.""" + return self._shells + + def get( + self, + delta_obs: torch.Tensor, + sigma_diff: torch.Tensor, + epsilon: torch.Tensor | None, + d_star_sq: torch.Tensor, + f_dark: torch.Tensor | None, + fit_mask: torch.Tensor, + *, + delta_calc: torch.Tensor | None = None, + target_dss: torch.Tensor | None = None, + out_epsilon: torch.Tensor | None = None, + out_f_dark: torch.Tensor | None = None, + out_sigma_diff: torch.Tensor | None = None, + ) -> SigmaDEstimate: + """Return the cached-or-recomputed :class:`SigmaDEstimate`. + + The fit inputs may be a pooled, flattened set (several datasets end to end); the + ``target_*`` / ``out_*`` arguments map the result onto another reflection list, + defaulting to the fit inputs themselves. + """ + if self._cache is not None: + return self._cache + shells = estimate_sigma_d( + delta_obs, + sigma_diff, + epsilon, + d_star_sq, + f_dark, + fit_mask, + delta_calc=delta_calc, + gamma=self.config.gamma, + shrink=self.config.shrink, + ) + self._shells = shells + self._cache = sigma_d_per_reflection( + shells, + d_star_sq if target_dss is None else target_dss, + epsilon if out_epsilon is None else out_epsilon, + f_dark if out_f_dark is None else out_f_dark, + sigma_diff if out_sigma_diff is None else out_sigma_diff, + ) + return self._cache diff --git a/torchref/refinement/model_error_estimation/sigma_m.py b/torchref/refinement/model_error_estimation/sigma_m.py index a5817dbc..20bd6300 100644 --- a/torchref/refinement/model_error_estimation/sigma_m.py +++ b/torchref/refinement/model_error_estimation/sigma_m.py @@ -133,7 +133,7 @@ def prepare( s_half_sq = s_half_sq.to(device=device, dtype=dtype) s_sq = 4.0 * s_half_sq - valid_f = validity.to(torch.bool).to(device).to(dtype) + valid_f = validity.to(torch.bool).to(device=device, dtype=dtype) n_valid = valid_f.sum().clamp(min=1.0) sigma_obs = sigma_obs.to(device=device, dtype=dtype) self.sigma_d_mean = (sigma_obs * valid_f).sum() / n_valid diff --git a/torchref/refinement/optimizers/curvature.py b/torchref/refinement/optimizers/curvature.py index b3192e19..74b5bbae 100644 --- a/torchref/refinement/optimizers/curvature.py +++ b/torchref/refinement/optimizers/curvature.py @@ -36,7 +36,7 @@ def _sample_probe( """Draw one Hutchinson probe vector of length ``numel``.""" if probe == "rademacher": r = torch.randint( - 0, 2, (numel,), generator=generator, device=device, dtype=torch.int64 + 0, 2, (numel,), generator=generator, device=device, dtype=torch.int64 # dtype-ok: randint {0,1} bernoulli draw, immediately cast to float dtype; width irrelevant ) return r.to(dtype).mul_(2.0).sub_(1.0) # {0,1} -> {-1,+1} if probe == "gaussian": diff --git a/torchref/refinement/rigid_body_refinement.py b/torchref/refinement/rigid_body_refinement.py index 70bb8d30..5718c19b 100644 --- a/torchref/refinement/rigid_body_refinement.py +++ b/torchref/refinement/rigid_body_refinement.py @@ -8,10 +8,12 @@ below it switches to ``ml`` with the normal Scaler. """ +import copy from typing import List, Optional import torch +from torchref.refinement.loss_state import LossState from torchref.scaling.scaler import Scaler @@ -88,14 +90,57 @@ def _xray_mode_for_cutoff(cls, d_min: float) -> str: # ----------------------------------------------------------------------- # Run # ----------------------------------------------------------------------- + @staticmethod + def _sandbox(ref): + """A shallow clone of ``ref`` that shares its model but owns its namespace. + + Every cutoff calls :meth:`_rebind_for_data`, which assigns + ``reflection_data``, ``scaler`` and the x-ray targets for that cutoff's + resolution and target mode. Run against the real Refinement those + assignments are destructive -- the caller gets its data and scaler + silently replaced by whatever the last cutoff used. + + Directing the step at a clone confines all of it. The real Refinement is + never written to, so there is nothing to restore and no window in which + it is inconsistent. The step builds its own single-target + :class:`~torchref.refinement.loss_state.LossState` rather than borrowing + the refinement's, so the caller's targets and any weights registered on + them are untouched as well. + + The model is deliberately shared, not copied: ``use_rigid_xyz`` swaps its + xyz container in place, so refined coordinates reach the caller by object + identity and need no copy-back. + + ``nn.Module`` keeps submodules in ``_modules``; copying ``__dict__`` + alone would leave that dict shared, and a submodule assignment on the + clone would write straight through to the original. + """ + sandbox = copy.copy(ref) + sandbox.__dict__ = dict(ref.__dict__) + for slot in ("_modules", "_parameters", "_buffers"): + if slot in sandbox.__dict__: + sandbox.__dict__[slot] = dict(ref.__dict__[slot]) + return sandbox + def run(self): """Step through every cutoff coarse to fine and return ``[(d_min, LossState), ...]``. - Restores the original ``reflection_data`` on exit, and bakes the final transform - back - into a plain ``ModelFT`` unless ``commit=False``. + Runs against a sandbox clone of the refinement (see :meth:`_sandbox`), so + the caller's targets, weights and ``reflection_data`` are left untouched. + Refined coordinates still reach the caller: the model is shared. + + Bakes the final transform back into a plain ``ModelFT`` unless + ``commit=False``. """ + real = self.refinement + self.refinement = self._sandbox(real) + try: + return self._run() + finally: + self.refinement = real + + def _run(self): ref = self.refinement original_data = ref.reflection_data @@ -106,6 +151,20 @@ def run(self): else self.default_cutoffs(native_dmin) ) + # ``cut_res`` masks in place and returns ``self``, so each cutoff below + # stamps ``masks["resolution"]`` on the caller's own object and rebinding + # restores nothing. Snapshot it (or its absence) to put back. + had_resolution_mask = "resolution" in original_data.masks + saved_resolution_mask = ( + original_data.masks["resolution"].clone() if had_resolution_mask else None + ) + + def restore_resolution_mask(): + if had_resolution_mask: + original_data.masks["resolution"] = saved_resolution_mask + else: + original_data.masks.pop("resolution", None) + # Swap the model's xyz container in place for a RigidXYZTensor. ref.model.use_rigid_xyz() @@ -120,7 +179,9 @@ def run(self): step_state = self._run_one_cutoff(d_min) history.append((float(d_min), step_state)) finally: - # Restore full-resolution data view. + # Before rebinding, so the scaler and targets are built against the + # data the caller has. + restore_resolution_mask() self._rebind_for_data(original_data) if self.commit: @@ -162,9 +223,14 @@ def _rebind_for_data(self, data, model=None, xray_mode=None): device=ref.device, ) ref.scaler.initialize() - ref._init_targets(xray_mode=xray_mode) + # Only the x-ray half of _init_targets: the step optimizes against x-ray data + # alone, and TotalGeometryTarget / TotalADPTarget would be constructed here + # purely to be left unused -- NonBondedTarget's pair list among them. + # get_scales() still runs, because the x-ray target reads the scaler's + # parameters. + ref._build_xray_targets(xray_mode) + ref.get_scales() - ref.reset_loss_state() # Clear cached LBFGS optimizers (they were built over the old model's # parameters and would now point at stale leaves). if hasattr(ref, "_persistent_optimizers"): @@ -175,81 +241,71 @@ def _run_one_cutoff(self, d_min: float): ref = self.refinement rigid_model = ref.model - state = ref.complete_loss_state() + # Active targets during rigid-body refinement: x-ray only. Internal bonded + # geometry is rigid by construction; ADP / occupancy are frozen. Inter-chain vdW + # is intentionally off -- Phenix runs rigid-body without atomistic restraints, + # and we have measured that vdW adds no signal here and destabilizes the + # coarsest cutoff. + # + # A state of its own rather than the refinement's: nothing here has to be undone + # afterwards, no maintenance() hook of a target we are not using can fire (in + # particular NonBondedTarget rebuilds its VDW pair list whenever atoms drift + # >1 A, seconds of work for a term that is not in the sum), and the caller's + # state keeps its targets and any weights registered on it. + # + # Weight 1.0: with one term the weight is a scalar on the whole objective, and + # 1.0 is what DEFAULT_GROUP_WEIGHTS gives x-ray anyway. + state = LossState(device=ref.device) + state.register_target("xray", ref.xray_target_work) + state.set_weight("xray", 1.0) + state.cache_losses() + + rigid_params = [ + rigid_model.xyz.euler_angles, + rigid_model.xyz.translations, + ] + + # Decide whether to use the inner-cycle (mask-refresh) loop. + # Unsatisfiable as it stands: nothing sets ``c_iso.requires_grad = + # False``, so ``_run_inner_cycles`` does not run. + use_inner_cycles = ( + ref.scaler is not None + and getattr(ref.scaler, "solvent", None) is not None + and getattr(ref.scaler, "c_iso", None) is not None + and ref.scaler.c_iso.requires_grad is False + ) - # Snapshot weights so we can restore after the step. - original_weights = dict(state.weights) - try: - # Active targets during rigid-body refinement: xray only. Internal - # bonded geometry is rigid by construction; ADP / occupancy are - # frozen. Inter-chain vdW is intentionally off — Phenix runs - # rigid-body without atomistic restraints and we've observed vdW - # adds no signal here and destabilizes the coarsest cutoff. - # - # We DROP non-xray targets from state entirely (not just zero - # their weight) so their maintenance() hooks don't fire during - # the rigid-body LBFGS. In particular ``NonBondedTarget`` - # rebuilds its VDW pair list whenever atoms drift >1 Å, a - # multi-second recomputation that's pure waste when the - # target weight is 0. State is rebuilt fresh on the next - # cutoff via _rebind_for_data → reset_loss_state → - # _init_targets, so this drop is local to this cutoff. - keep_names = {"xray"} - for name in list(state.targets.keys()): - if name not in keep_names: - state.targets.pop(name, None) - original_weights.pop(name, None) - - rigid_params = [ - rigid_model.xyz.euler_angles, - rigid_model.xyz.translations, - ] - - # Decide whether to use the inner-cycle (mask-refresh) loop. - # Triggered when the scaler has a bulk-solvent component whose - # mask depends on atom positions (ls_wunit_k1 path here). For - # the ml path the scaler is fully refit between cutoffs - # and co-optimized with rigid params in a single LBFGS. - use_inner_cycles = ( - ref.scaler is not None - and getattr(ref.scaler, "solvent", None) is not None - and getattr(ref.scaler, "c_iso", None) is not None - and ref.scaler.c_iso.requires_grad is False - ) + if use_inner_cycles: + self._run_inner_cycles(d_min, state, rigid_params, n_inner=5) + else: + # Rigid parameters only. The body target centres on + # ``alpha*|F_calc|`` and ``alpha`` absorbs a rescaling of ``F_calc`` + # exactly, so the scale has a flat direction here -- the rule + # ``SCALE_TARGETS`` states for the scale fit. ``refine_scaler`` + # (objective ``ls``) owns the scale, between cutoffs. Omitting them + # is enough: ``LossState.run`` freezes leaves the optimizer lacks. + opt_params = rigid_params - if use_inner_cycles: - self._run_inner_cycles(d_min, state, rigid_params, n_inner=5) - else: - # Single-shot path: rigid params + scaler params co-optimized. - if ref.scaler is not None: - opt_params = rigid_params + list(ref.scaler.parameters()) - else: - opt_params = rigid_params - - rigid_model.reset_cache() - opt = torch.optim.LBFGS( - opt_params, - max_iter=self.iterations_per_step, - **self.DEFAULT_LBFGS_KWARGS, - ) - state.step( - opt, - context=f"rigid_body[d_min={d_min:.2f}]", + rigid_model.reset_cache() + opt = torch.optim.LBFGS( + opt_params, + max_iter=self.iterations_per_step, + **self.DEFAULT_LBFGS_KWARGS, + ) + state.step( + opt, + context=f"rigid_body[d_min={d_min:.2f}]", + ) + if ref.verbose > 0: + try: + rwork, rfree = ref.get_rfactor() + print( + f" rigid-body d_min={d_min:.2f} " + f"(lbfgs, iters={self.iterations_per_step}): " + f"Rwork={rwork:.4f} Rfree={rfree:.4f}" ) - if ref.verbose > 0: - try: - rwork, rfree = ref.get_rfactor() - print( - f" rigid-body d_min={d_min:.2f} " - f"(lbfgs, iters={self.iterations_per_step}): " - f"Rwork={rwork:.4f} Rfree={rfree:.4f}" - ) - except Exception: - pass - finally: - # Restore weights. - for name, w in original_weights.items(): - state.set_weight(name, w) + except Exception: + pass return state def _run_inner_cycles(self, d_min, state, rigid_params, n_inner: int = 5): diff --git a/torchref/refinement/targets/__init__.py b/torchref/refinement/targets/__init__.py index 52736fb0..4bdaba50 100644 --- a/torchref/refinement/targets/__init__.py +++ b/torchref/refinement/targets/__init__.py @@ -5,8 +5,8 @@ """ from .adp import ( - ADPSigdTarget, ADPLocalityTarget, + ADPSigdTarget, ADPSimilarityTarget, ADPTarget, RigidBondTarget, @@ -20,9 +20,12 @@ von_mises_nll, ) from .collection import ( + COLLECTION_XRAY_TARGETS, + CollectionDifferenceIntensityTarget, + CollectionDifferenceSigmaDTarget, CollectionDifferenceTarget, CollectionMLTarget, - CollectionRiceTarget, + CollectionTwoMomentIntensityTarget, MultiModelADPTarget, MultiModelGeometryTarget, ) @@ -31,6 +34,7 @@ TotalADPTarget, TotalGeometryTarget, ) +from .dataset_scaling import DatasetScalingTarget from .difference import ( DifferenceXrayTarget, PhaseInformedDifferenceTarget, @@ -64,6 +68,7 @@ ) __all__ = [ + "DatasetScalingTarget", # Base classes "Target", "ModelTarget", @@ -86,8 +91,11 @@ "create_xray_target", # Collection (multi-dataset) targets "CollectionDifferenceTarget", - "CollectionRiceTarget", + "CollectionTwoMomentIntensityTarget", "CollectionMLTarget", + "CollectionDifferenceIntensityTarget", + "CollectionDifferenceSigmaDTarget", + "COLLECTION_XRAY_TARGETS", "MultiModelGeometryTarget", "MultiModelADPTarget", # Difference targets diff --git a/torchref/refinement/targets/adp/__init__.py b/torchref/refinement/targets/adp/__init__.py index b4638eeb..c7f3ec79 100644 --- a/torchref/refinement/targets/adp/__init__.py +++ b/torchref/refinement/targets/adp/__init__.py @@ -3,6 +3,8 @@ from .rigid_bond import RigidBondTarget from .sigd import ADPSigdTarget from .locality import ADPLocalityTarget +from .node_load import NodeLoadTarget +from .node_smoothness import NodeSmoothnessTarget from .scaler_log_scale import ScalerLogScaleTrendTarget from .scaler_u import ScalerURegularizationTarget @@ -12,6 +14,8 @@ "RigidBondTarget", "ADPSigdTarget", "ADPLocalityTarget", + "NodeLoadTarget", + "NodeSmoothnessTarget", "ScalerURegularizationTarget", "ScalerLogScaleTrendTarget", ] diff --git a/torchref/refinement/targets/adp/node_load.py b/torchref/refinement/targets/adp/node_load.py new file mode 100644 index 00000000..559e725a --- /dev/null +++ b/torchref/refinement/targets/adp/node_load.py @@ -0,0 +1,127 @@ +"""Load balancing for a node-field ADP representation.""" + +import torch +from typing import TYPE_CHECKING, Dict + +from torchref.utils.stats import ( + VERBOSITY_DEBUG, + VERBOSITY_DETAILED, + VERBOSITY_STANDARD, + StatEntry, + stat, +) + +from .base import ADPTarget + +if TYPE_CHECKING: + from torchref.model.model import Model + + +class NodeLoadTarget(ADPTarget): + """Keep every disorder-field node carrying a fair share of atoms. + + A node's load is the total weight it holds across all atoms, + :meth:`~torchref.model.disorder_field.DisorderFieldTensor.node_load`, and the + weights are a partition of unity, so the loads sum to the atom count and their mean + is ``n_atoms / K`` whatever the model does. + + Without this the field has a degenerate direction: a node can narrow its kernel + until it holds a single atom, then take whatever value fits that atom. Measured, a + collapsed node ends up with a load near or below one atom against a healthy median + of seven, and sets its atom's B into the hundreds or thousands. One node fitting one + atom is per-atom refinement wearing a node's clothes, which is the thing the + representation exists to avoid. + + The penalty is **one-sided**, ``softplus(-log(load / mean_load))``: it grows as a + node is abandoned, and flattens to zero once a node carries its share. That + asymmetry is deliberate. The symmetric choice -- maximising the entropy of the load + distribution -- is optimal at *uniform* load, so it would also penalise a broad node + that legitimately covers more atoms than its neighbours. Fitted fields span nearly + two orders of magnitude in kernel width within a single structure, and that spread + is the representation working, not failing. + + Acts through the weights, so its gradient reaches node positions and kernel widths + but never the node values: it removes the *opportunity* to place an extreme B rather + than penalising the B itself. It therefore composes with, rather than duplicates, + the restraints that act on the values. + + Inert unless the model is in field mode, so it can be registered unconditionally. + + Parameters + ---------- + model : Model, optional + Reference to the Model object. + sharpness : float, optional + Softplus temperature in log-load units. Smaller is a harder barrier. Default + 0.5, which leaves a node at the mean load contributing about 0.1 and a node at a + tenth of the mean about 2.3. + verbose : int, optional + Verbosity level. Default is 0. + """ + + #: Hierarchical key this target registers under. Required, not cosmetic: + #: LossState.register_targets takes the key from ``.name``, so without it the + #: target inherits ``Target.name`` ("model_target"), registers under that, + #: collides with every other unnamed target, and no ``adp/...`` weight can + #: reach it -- the term is then built, callable, and never in the loss. + name: str = "adp/node_load" + + def __init__( + self, + model: "Model" = None, + sharpness: float = 0.5, + verbose: int = 0, + device=None, + **kwargs, + ): + super().__init__(model, verbose, device=device, **kwargs) + self.sharpness = float(sharpness) + + @property + def _field(self): + """The disorder field, or ``None`` when the model is not in field mode. + + Reads ``Model.adp_field`` rather than the ``adp`` slot directly: an anisotropic + payload lives in ``u`` instead, and looking only at ``adp`` would leave this + target silently inert in exactly the mode with the most node parameters to + collapse. + """ + return getattr(self.model, "adp_field", None) + + def _relative_load(self) -> torch.Tensor: + """Each node's load as a multiple of the mean load, ``(K,)``.""" + field = self._field + load = field.node_load() + # Mean load is n_atoms / K exactly, because the weights sum to one per atom. + return load / (load.sum().detach() / load.shape[0]).clamp(min=1e-12) + + def forward(self) -> torch.Tensor: + """Summed one-sided load deficit over nodes, or zero outside field mode.""" + field = self._field + if field is None: + return torch.zeros((), device=self.device) + rel = self._relative_load() + deficit = -torch.log(rel.clamp(min=1e-12)) / self.sharpness + return torch.nn.functional.softplus(deficit).sum() * self.sharpness + + def stats(self) -> Dict[str, any]: + """Load distribution across nodes, and how much of it the barrier sees.""" + field = self._field + if field is None: + return {"node_load_active": stat(0.0, VERBOSITY_DEBUG)} + with torch.no_grad(): + rel = self._relative_load() + loss = self.forward() + return { + "node_load_loss": stat(float(loss), VERBOSITY_STANDARD), + "n_nodes": stat(int(rel.numel()), VERBOSITY_STANDARD), + "load_min_rel": stat(float(rel.min()), VERBOSITY_STANDARD), + "load_median_rel": stat(float(rel.median()), VERBOSITY_DETAILED), + "load_max_rel": stat(float(rel.max()), VERBOSITY_DETAILED), + # The population the barrier exists for. + "n_below_quarter_share": stat( + int((rel < 0.25).sum()), VERBOSITY_STANDARD + ), + "load_cv": stat(float(rel.std() / rel.mean().clamp(min=1e-12)), + VERBOSITY_DEBUG), + } diff --git a/torchref/refinement/targets/adp/node_smoothness.py b/torchref/refinement/targets/adp/node_smoothness.py new file mode 100644 index 00000000..6ae5634b --- /dev/null +++ b/torchref/refinement/targets/adp/node_smoothness.py @@ -0,0 +1,154 @@ +"""Magnitude prior on the node values of a disorder field.""" + +import torch +from typing import TYPE_CHECKING, Dict + +from torchref.utils.stats import ( + VERBOSITY_DEBUG, + VERBOSITY_DETAILED, + VERBOSITY_STANDARD, + StatEntry, + stat, +) + +from .base import ADPTarget + +if TYPE_CHECKING: + from torchref.model.model import Model + + +class NodeSmoothnessTarget(ADPTarget): + """Penalise a node whose B departs from the nodes around it. + + The companion to :class:`~torchref.refinement.targets.adp.NodeLoadTarget`, which + acts on the weights and therefore cannot reach the node *values* at all. Blocking a + node from narrowing does not stop it taking an extreme B -- measured, it makes the + consequence broader rather than smaller, because the extreme value can no longer be + confined to the single atom the node had isolated. So the two terms close different + halves: one denies the opportunity, this one prices the magnitude. + + The penalty is a distance-weighted sum over node pairs:: + + L = sum_{k 1 else 1.0 + lam = max(lam, 1e-3) + + w = torch.exp(-(d**2) / (2.0 * lam * lam)) + w = torch.triu(w, diagonal=1) + diff2 = (log_b[:, None] - log_b[None, :]) ** 2 + return w, diff2, lam + + def forward(self) -> torch.Tensor: + """Weighted mean squared log-B difference between nearby nodes.""" + field = self._field + if field is None or field.n_nodes < 2: + return torch.zeros((), device=self.device) + w, diff2, _ = self._pair_terms() + total = w.sum() + if float(total.detach()) <= 0.0: + return torch.zeros((), device=self.device) + return (w * diff2).sum() / total + + def stats(self) -> Dict[str, any]: + """Spread of the node values, and how localised the departures are.""" + field = self._field + if field is None or field.n_nodes < 2: + return {"node_smoothness_active": stat(0.0, VERBOSITY_DEBUG)} + with torch.no_grad(): + w, diff2, lam = self._pair_terms() + loss = self.forward() + log_b = field.log_magnitude() + b = torch.exp(log_b) + return { + "node_smoothness_loss": stat(float(loss), VERBOSITY_STANDARD), + "node_b_median": stat(float(b.median()), VERBOSITY_STANDARD), + "node_b_max": stat(float(b.max()), VERBOSITY_STANDARD), + "node_log_b_sd": stat(float(log_b.std()), VERBOSITY_DETAILED), + "node_pair_length_scale": stat(float(lam), VERBOSITY_DETAILED), + # How far the worst node sits above its own neighbourhood. + "node_b_max_over_median": stat( + float(b.max() / b.median().clamp(min=1e-12)), VERBOSITY_STANDARD + ), + } diff --git a/torchref/refinement/targets/adp/rigid_bond.py b/torchref/refinement/targets/adp/rigid_bond.py index c8b69dd8..07a639a7 100644 --- a/torchref/refinement/targets/adp/rigid_bond.py +++ b/torchref/refinement/targets/adp/rigid_bond.py @@ -122,7 +122,7 @@ def _bond_pairs(self) -> torch.Tensor: chunks.append(idx_) if chunks: return torch.cat(chunks, dim=0).contiguous() - return torch.empty(0, 2, dtype=torch.long, device=self.model.xyz().device) + return torch.empty(0, 2, dtype=torch.long, device=self.model.xyz().device) # dtype-ok: empty (0,2) atom-pair index tensor; PyTorch requires int64 def _compute_aniso_rigid_bond(self) -> torch.Tensor: """Rigid-bond NLL from ``Δz = l^T U_1 l - l^T U_2 l`` along each bond. diff --git a/torchref/refinement/targets/adp/sigd.py b/torchref/refinement/targets/adp/sigd.py index aaf4cb2f..95aace4d 100644 --- a/torchref/refinement/targets/adp/sigd.py +++ b/torchref/refinement/targets/adp/sigd.py @@ -118,6 +118,7 @@ def stats(self) -> Dict[str, any]: beta = float((adp - self._b_shift).clamp(min=1e-3).mean()) * (alpha - 1.0) # std(log B) = sqrt(trigamma(alpha)); torch.polygamma(1, .) is trigamma. implied_std = math.sqrt( + # dtype-ok: deliberate float64 for a scalar polygamma; extracted via float() float(torch.polygamma(1, torch.tensor(alpha, dtype=torch.float64))) ) diff --git a/torchref/refinement/targets/adp/similarity.py b/torchref/refinement/targets/adp/similarity.py index d8619131..191d2379 100644 --- a/torchref/refinement/targets/adp/similarity.py +++ b/torchref/refinement/targets/adp/similarity.py @@ -94,7 +94,7 @@ def _get_pair_indices(self) -> torch.Tensor: if chunks: cached = torch.cat(chunks, dim=0).contiguous() else: - cached = torch.empty(0, 2, dtype=torch.long, + cached = torch.empty(0, 2, dtype=torch.long, # dtype-ok: empty (0,2) atom-pair index tensor; PyTorch requires int64 device=self.model.xyz().device) self._simu_pair_indices_cache = cached return cached diff --git a/torchref/refinement/targets/base.py b/torchref/refinement/targets/base.py index 0a4b62e2..884e349b 100644 --- a/torchref/refinement/targets/base.py +++ b/torchref/refinement/targets/base.py @@ -385,6 +385,38 @@ def get_F_calc_scaled(self, hkl=None, recalc=False, fcalc=None): """ return torch.abs(self.get_fcalc_scaled(hkl, recalc=recalc, fcalc=fcalc)) + def get_I_calc_scaled(self, hkl=None, recalc=False, fcalc=None): + """ + Compute scaled structure factor intensities ``|F_calc|**2``. + + The intensity sibling of :meth:`get_F_calc_scaled`, and the reason the observable + is a choice rather than an assumption: both are one line over the same complex + ``get_fcalc_scaled``, so nothing upstream of here knows which observable a target + fits. + + Squaring the *scaled* amplitude is what makes this correct -- the scale and the + anisotropy factor both enter squared, matching + :meth:`ReflectionData.get_corrected_intensities` on the observation side. Squaring + an unscaled ``F_calc`` and scaling afterwards with the amplitude factors would be + wrong by that factor, which is resolution-dependent and so reads as a scale or B + error rather than as a bug. + + Parameters + ---------- + hkl : torch.Tensor, optional + Miller indices. If None, uses data's hkl. + recalc : bool, optional + Force recalculation. Default is False. + fcalc : torch.Tensor, optional + Pre-computed structure factors. If provided, skips model computation. + + Returns + ------- + torch.Tensor + Scaled structure factor intensities ``|F_calc|**2``. + """ + return self.get_fcalc_scaled(hkl, recalc=recalc, fcalc=fcalc).abs() ** 2 + # ============================================================================= # Utility Functions for NLL Computation diff --git a/torchref/refinement/targets/collection/__init__.py b/torchref/refinement/targets/collection/__init__.py index f4302a77..adea453f 100644 --- a/torchref/refinement/targets/collection/__init__.py +++ b/torchref/refinement/targets/collection/__init__.py @@ -8,18 +8,38 @@ """ from ._util import _scale_fcalc -from .base import CollectionXrayTarget +from .base import ( + CollectionLossInputs, + CollectionSigmaALossInputs, + CollectionSigmaDLossInputs, + CollectionXrayTarget, +) +from .intensity import CollectionTwoMomentIntensityTarget from .multimodel import MultiModelADPTarget, MultiModelGeometryTarget from .xray import ( + CollectionDifferenceIntensityTarget, + CollectionDifferenceSigmaDTarget, CollectionDifferenceTarget, CollectionMLTarget, - CollectionRiceTarget, +) +from ._specs import ( # noqa: E402 (imports the rows above) + COLLECTION_XRAY_TARGETS, + CollectionXrayTargetSpec, + CollectionXrayTargetTable, ) __all__ = [ + "COLLECTION_XRAY_TARGETS", + "CollectionXrayTargetSpec", + "CollectionXrayTargetTable", "CollectionXrayTarget", + "CollectionLossInputs", + "CollectionSigmaALossInputs", + "CollectionSigmaDLossInputs", + "CollectionTwoMomentIntensityTarget", "CollectionDifferenceTarget", - "CollectionRiceTarget", + "CollectionDifferenceIntensityTarget", + "CollectionDifferenceSigmaDTarget", "CollectionMLTarget", "MultiModelGeometryTarget", "MultiModelADPTarget", diff --git a/torchref/refinement/targets/collection/_specs.py b/torchref/refinement/targets/collection/_specs.py new file mode 100644 index 00000000..2d1b9436 --- /dev/null +++ b/torchref/refinement/targets/collection/_specs.py @@ -0,0 +1,161 @@ +"""The collection X-ray target taxonomy, as data. + +The multi-dataset mirror of :mod:`torchref.refinement.targets.xray._specs`, with the same +invariants checked the same way at import: unique names, and **one class per row**, so +dispatch is ``spec.target_cls(**kwargs)`` with nothing to branch on. + +Five rows over two axes -- what the loss compares (a difference from the collection mean, +or each dataset absolutely) and in which observable: + +==================== =========== ============================================== +row observable compares +==================== =========== ============================================== +``difference`` amplitude ``F_i - F_mean`` against the model's own spread +``difference_i`` intensity the same, in intensities +``difference_sd`` amplitude ``F_i - F_mean`` against ``alpha dF_calc``, variance + ``beta_model + sigma^2`` from a sigma_D fit +``two_moment`` intensity ``|F(alpha)|^2 + sigma_alpha^2 |dF|^2`` +``ml`` amplitude each dataset absolutely, at a shared Luzzati beta +==================== =========== ============================================== + +The absolute ``ml`` channel constrains the overall level when all component models +are free. Rice likelihoods describe amplitudes; intensity rows use Gaussian losses. +""" + +from dataclasses import dataclass, field +from typing import Dict, Tuple + +from .base import CollectionXrayTarget +from .intensity import CollectionTwoMomentIntensityTarget +from .xray import ( + CollectionDifferenceIntensityTarget, + CollectionDifferenceSigmaDTarget, + CollectionDifferenceTarget, + CollectionMLTarget, +) + +@dataclass(frozen=True) +class CollectionXrayTargetSpec: + """One selectable collection x-ray target: a name, and the class implementing it. + + Attributes + ---------- + name + The row name, as used in ``LossState`` keys (``xray/``). + target_cls + The class. **One class per row**, checked by :class:`CollectionXrayTargetTable`. + doc + One line, for ``--help`` and the loss breakdown. + observable + ``"amplitude"`` or ``"intensity"``. Checked against the class, so a spec and its + implementation cannot disagree -- a row advertising intensities while reading + amplitudes would be wrong by ``2|F|``, which is resolution-dependent and so reads + as a scale or B error rather than as a bug. + """ + + name: str + target_cls: type + doc: str + observable: str = "amplitude" + + def __post_init__(self): + if not ( + isinstance(self.target_cls, type) + and issubclass(self.target_cls, CollectionXrayTarget) + ): + raise TypeError( + f"{self.name}: target_cls {self.target_cls!r} is not a " + f"CollectionXrayTarget subclass" + ) + if self.observable not in ("amplitude", "intensity"): + raise ValueError( + f"{self.name}: observable must be 'amplitude' or 'intensity', " + f"got {self.observable!r}" + ) + declared = getattr(self.target_cls, "observable", "amplitude") + if declared != self.observable: + raise ValueError( + f"{self.name}: spec says observable={self.observable!r} but " + f"{self.target_cls.__name__} says {declared!r}" + ) + + +@dataclass(frozen=True) +class CollectionXrayTargetTable: + """The taxonomy, with uniqueness checked at import.""" + + specs: Tuple[CollectionXrayTargetSpec, ...] + _by_name: Dict[str, CollectionXrayTargetSpec] = field( + init=False, repr=False, default=None + ) + + def __post_init__(self): + lookup: Dict[str, CollectionXrayTargetSpec] = {} + for spec in self.specs: + if spec.name in lookup: + raise ValueError(f"duplicate collection x-ray target name {spec.name!r}") + lookup[spec.name] = spec + by_cls: Dict[type, CollectionXrayTargetSpec] = {} + for spec in self.specs: + if spec.target_cls in by_cls: + raise ValueError( + f"{spec.name} and {by_cls[spec.target_cls].name} both map to " + f"{spec.target_cls.__name__}. One class per row is the invariant this " + f"table exists to enforce: a class serving two rows has to branch on " + f"something at runtime." + ) + by_cls[spec.target_cls] = spec + object.__setattr__(self, "_by_name", lookup) + + @property + def names(self) -> Tuple[str, ...]: + """Canonical names, in table order.""" + return tuple(s.name for s in self.specs) + + def by_name(self, name: str) -> CollectionXrayTargetSpec: + spec = self._by_name.get(name) + if spec is None: + raise ValueError( + f"Unknown collection X-ray target: {name!r}. " + f"Available: {', '.join(self.names)}" + ) + return spec + + +COLLECTION_XRAY_TARGETS = CollectionXrayTargetTable( + specs=( + CollectionXrayTargetSpec( + name="difference", + target_cls=CollectionDifferenceTarget, + doc="Gaussian on each dataset's amplitude difference from the collection " + "mean, with the dataset/mean covariance propagated.", + ), + CollectionXrayTargetSpec( + name="difference_i", + target_cls=CollectionDifferenceIntensityTarget, + observable="intensity", + doc="As 'difference' but on intensities, skipping the French-Wilson " + "conversion that reshapes the weak tail.", + ), + CollectionXrayTargetSpec( + name="difference_sd", + target_cls=CollectionDifferenceSigmaDTarget, + doc="As 'difference', centred on alpha * dF_calc with the unexplained " + "difference power beta_model (sigma_D, fitted on the free set) added to " + "the measurement variance.", + ), + CollectionXrayTargetSpec( + name="two_moment", + target_cls=CollectionTwoMomentIntensityTarget, + observable="intensity", + doc="Merged intensities as |F(alpha)|^2 + sigma_alpha^2 |dF|^2, accounting " + "for crystal-to-crystal spread in activation.", + ), + CollectionXrayTargetSpec( + name="ml", + target_cls=CollectionMLTarget, + doc="Read MLF per dataset at one shared Luzzati beta fitted on the pooled " + "free reflections. The absolute channel.", + ), + ) +) diff --git a/torchref/refinement/targets/collection/_util.py b/torchref/refinement/targets/collection/_util.py index 7611ccf1..1740335e 100644 --- a/torchref/refinement/targets/collection/_util.py +++ b/torchref/refinement/targets/collection/_util.py @@ -12,3 +12,16 @@ def _scale_fcalc(scaler, fcalc, model): if hasattr(scaler, "forward_mixed") and hasattr(model, "fractions"): return scaler.forward_mixed(fcalc, model.fractions) return scaler(fcalc) + + +def common_geom(data): + """``(epsilon, d_star_sq)`` on a dataset's HKL: multiplicity and ``1/d**2`` in A^-2.""" + import torch + + from torchref.base.reciprocal import get_scattering_vectors + from torchref.refinement.model_error_estimation.sigma_a import epsilon_from_hkl + + eps = epsilon_from_hkl(data.hkl, getattr(data, "spacegroup", None)) + s = get_scattering_vectors(data.hkl, data.cell) + dss = (torch.norm(s, dim=1) ** 2).to(eps.dtype) + return eps, dss diff --git a/torchref/refinement/targets/collection/base.py b/torchref/refinement/targets/collection/base.py index 064cd971..a274ca68 100644 --- a/torchref/refinement/targets/collection/base.py +++ b/torchref/refinement/targets/collection/base.py @@ -1,23 +1,25 @@ -"""Shared base for collection (multi-dataset) X-ray targets. - -:class:`CollectionXrayTarget` gives them the same subset and R-factor contract as -the single-dataset -:class:`~torchref.refinement.targets.xray.base.XrayTarget`: the 3-way ``use_set`` -selector over each member's ``data.work``/``free``/``validation`` accessors, the one -shared :func:`~torchref.base.metrics.rfactor.rfactor_work_free` computed through the -same scaling the loss sees, and the standard ``loss``/``n``/``rwork``/``rfree`` -``stats()`` dict. - -Since every member is expanded onto one common HKL grid, per-dataset R-factors form a -distribution: headline ``rwork``/``rfree`` are its median, with the 10/25/75/90 -percentiles at higher verbosity. +"""Shared observation access, reduction and reporting for collection X-ray targets. + +Targets declare an amplitude or intensity observable. ``_loss_inputs`` gathers +observations, predictions, uncertainties and masks on the common HKL grid; +``_per_refl`` returns unreduced losses used by both ``forward`` and ``residuals``. +Difference targets intersect member masks so each fitted reflection is present +in every dataset's selected work, free or validation subset. + +R-factors use the same scaled predictions as the loss. Reporting gives the median +across datasets, with the 10/25/75/90 percentiles at higher verbosity. """ -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, NamedTuple import torch from torchref.base.metrics.rfactor import rfactor_work_free +from torchref.base.targets.xray_likelihoods import ( + SIGMA_FLOOR_ABS, + SIGMA_FLOOR_FRAC, +) +from torchref.config import get_float_dtype from torchref.refinement.targets.base import Target from torchref.utils.stats import ( VERBOSITY_DEBUG, @@ -40,6 +42,71 @@ _R_PCT_LABELS = ("p10", "p25", "p50", "p75", "p90") +class CollectionLossInputs(NamedTuple): + """What a collection row's :meth:`CollectionXrayTarget._per_refl` reads. + + Every tensor is ``(N, n_hkl)`` on the collection's common HKL grid, with ``N`` the + number of matched datasets in ``keys`` order -- full size rather than compact, because + the members are already expanded onto one grid and a compact form would need a + different index map per dataset. + + ``mask`` has already been intersected with finiteness of ``obs`` and ``sigma``, and + those two have been substituted where non-finite. That order matters: masking the + *loss* is not enough, because ``torch.where`` selects the finite branch for the value + while still backpropagating NaN through the branch it discarded. Real reflection files + carry non-finite intensities (excluded rows, and rows French-Wilson rejected), so this + is load-bearing rather than defensive. + """ + + obs: torch.Tensor + model: torch.Tensor + sigma: torch.Tensor + mask: torch.Tensor + keys: List[str] + + +class CollectionSigmaALossInputs(NamedTuple): + """:class:`CollectionLossInputs` plus one shared model-error estimate. + + The collection twin of + :class:`~torchref.refinement.targets.xray.sigma_a.SigmaALossInputs`. ``beta`` and + ``epsilon`` live on the **common HKL**, shape ``(n_hkl,)``, and broadcast over the + dataset axis: they are fitted once on the pooled free reflections of every data-model + pair, so one per-reflection variance serves every member. + + The two shapes never mix, because each class pairs its own ``_loss_inputs`` with its + own ``_per_refl``. + """ + + obs: torch.Tensor + model: torch.Tensor + sigma: torch.Tensor + mask: torch.Tensor + keys: List[str] + centric: torch.Tensor = None + beta: torch.Tensor = None + epsilon: torch.Tensor = None + + +class CollectionSigmaDLossInputs(NamedTuple): + """:class:`CollectionLossInputs` plus the sigma_D difference-error estimate. + + ``alpha`` and ``beta_model`` live on the **common HKL**, shape ``(n_hkl,)``, and + broadcast over the dataset axis: the coupling of the model difference to the true + one and the difference power the model leaves unexplained, fitted once on the pooled + free reflections of the timepoint rows. Detached, so gradients reach the models only + through ``model``. + """ + + obs: torch.Tensor + model: torch.Tensor + sigma: torch.Tensor + mask: torch.Tensor + keys: List[str] + alpha: torch.Tensor = None + beta_model: torch.Tensor = None + + class CollectionXrayTarget(Target): """Base class for multi-dataset X-ray targets. @@ -63,6 +130,21 @@ class CollectionXrayTarget(Target): name: str = "collection_xray" + #: Which measured column this row fits: ``"amplitude"`` or ``"intensity"``. Declared + #: rather than passed, for the same reason as the single-dataset table -- see + #: :mod:`torchref.refinement.targets.xray.observable`. + observable: str = "amplitude" + + #: Fewest matched datasets for the loss to mean anything. The difference targets need + #: two (there is no difference from a single dataset); the per-dataset rows need one. + min_datasets: int = 1 + + #: Multiplies the work-set loss. Rows carrying a likelihood whose magnitude differs + #: from its siblings' set this so the term neither swamps nor is swamped by the + #: restraints; :meth:`CollectionTwoMomentIntensityTarget.calibrate_base_weight` fits + #: it against a reference target's gradient norm. + base_weight: float = 1.0 + def __init__( self, dataset_collection: "DatasetCollection", @@ -83,14 +165,10 @@ def __init__( self.use_set = use_set self.use_work_set = use_set == "work" - # ------------------------------------------------------------------ - # Dataset / model / subset plumbing - # ------------------------------------------------------------------ - def _keys(self) -> List[str]: """Matched dataset keys this target fits: dark + present timepoints. Targets - fitting only part of the collection override it (``CollectionRiceTarget`` - drops the dark reference). + fitting only part of the collection override it (a target fitting only the + excited timepoints drops the dark reference). """ dc = self._dataset_collection mc = self._model_collection @@ -125,9 +203,114 @@ def _scaled_amp_full(self, data, model, recalc: bool = True) -> torch.Tensor: fcalc = data.structure_factors(model, recalc=recalc) return torch.abs(_scale_fcalc(self._scaler, fcalc, model)) - # ------------------------------------------------------------------ - # R-factor reporting (shared source of truth) - # ------------------------------------------------------------------ + def _stack_observations(self, keys: List[str]): + """``(obs, sigma)``, each ``(N, n_hkl)``, in this row's observable. + + Collection accessors read live dataset views and name any member missing + an intensity column. + """ + dc = self._dataset_collection + if self.observable == "intensity": + return dc.stack_I_obs(keys), dc.stack_I_sigma(keys) + return dc.stack_F_obs(keys), dc.stack_F_sigma(keys) + + def _stack_model(self, keys: List[str], recalc: bool = False) -> torch.Tensor: + """The model prediction, ``(N, n_hkl)``, in this row's observable. + + Default is the per-dataset scaled amplitude (squared for an intensity row). Rows + whose prediction is not a function of one dataset at a time -- the two-moment + model, which mixes shared components across timepoints -- override this. + """ + dc = self._dataset_collection + mc = self._model_collection + amp = torch.stack( + [self._scaled_amp_full(dc[k], mc[k], recalc=recalc) for k in keys] + ) + return amp**2 if self.observable == "intensity" else amp + + def _sigma_floor(self, sigma: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Floor for ``sigma``, at :data:`SIGMA_FLOOR_FRAC` of its median over ``mask``. + + A merged sigma can be reported as exactly zero; unfloored, one such reflection + dominates the whole sum. Detached, because it is a numerical safeguard rather than + a fitted quantity -- a gradient through a median would make the loss depend on the + ordering of near-equal sigmas. + + Taken over the fitted rows only: the members are reindexed onto a common grid, so + the rows a dataset does not own carry filler that would move the median. + """ + selected = sigma[mask] + if selected.numel() == 0: + return torch.as_tensor(1e-6, device=sigma.device, dtype=sigma.dtype) + floor = torch.median(selected).detach() * SIGMA_FLOOR_FRAC + return floor.clamp(min=SIGMA_FLOOR_ABS) + + def _loss_inputs(self, recalc: bool = False) -> CollectionLossInputs: + """Gather this row's observations, model, sigma and mask -- see + :class:`CollectionLossInputs` for the shapes and the NaN discipline. + + Rows narrow the mask here (the difference targets require a reflection to be in + the subset of every dataset) rather than inside :meth:`_per_refl`, so that + :meth:`forward`'s sum and :meth:`residuals`' array agree on which reflections + count. + """ + keys = self._keys() + obs, sigma = self._stack_observations(keys) + model = self._stack_model(keys, recalc=recalc) + mask = self._dataset_collection.stack_masks(keys, use_set=self.use_set) + + obs = obs.to(model.dtype) + sigma = sigma.to(model.dtype) + + # Sanitise into the mask BEFORE the graph, not after: see CollectionLossInputs. + valid = torch.isfinite(obs) & torch.isfinite(sigma) + mask = mask & valid + obs = torch.where(valid, obs, torch.zeros_like(obs)) + sigma = torch.where(valid, sigma, torch.ones_like(sigma)) + + return CollectionLossInputs(obs, model, sigma, mask, keys) + + def _per_refl(self, ctx: CollectionLossInputs) -> torch.Tensor: + """The likelihood, per reflection and **unreduced**, shape ``(N, n_hkl)``. + + One per selectable row; no row branches. :meth:`forward` is the masked sum of + this. + """ + raise NotImplementedError + + def forward(self) -> torch.Tensor: + """Masked sum of :meth:`_per_refl`, with ``base_weight`` on the work set only. + + Cache reset first: a preceding no-grad ``stats()`` or ``get_rfactor()`` call can + leave a detached tensor in a base model's cache, which would silently kill the + loss backward. + """ + keys = self._keys() + if len(keys) < self.min_datasets: + return torch.zeros((), device=self._dataset_collection.hkl.device) + + self._reset_model_caches() + ctx = self._loss_inputs(recalc=False) + total = (self._per_refl(ctx) * ctx.mask).sum() + # Work set only: the free-set value is a diagnostic and has to stay comparable + # across weightings. + if self.use_work_set and self.base_weight != 1.0: + total = self.base_weight * total + return total + + def residuals(self) -> torch.Tensor: + """:meth:`_per_refl` over every reflection, ``(N, n_hkl)``, unsummed and unmasked. + + The unreduced :meth:`forward`: same observable, same model, same variance. Masked + reflections still get a value, so the array can be used to ask *why* one was + excluded rather than only reflecting the answer back, and non-finite values + survive because here a NaN is a finding rather than a nuisance. + """ + keys = self._keys() + if len(keys) < self.min_datasets: + dc = self._dataset_collection + return torch.zeros((0, len(dc.hkl)), device=dc.hkl.device) + return self._per_refl(self._loss_inputs(recalc=True)) def get_rfactor(self) -> Dict[str, object]: """Per-dataset R-work / R-free plus percentile summaries. @@ -172,14 +355,11 @@ def _percentiles(values: List[float]) -> Dict[str, float]: """10/25/50/75/90 percentiles of a list of per-dataset R-factors.""" if not values: return {} - t = torch.tensor(values, dtype=torch.float64) - q = torch.quantile(t, torch.tensor(_R_PERCENTILES, dtype=torch.float64)) + dtype = get_float_dtype() + t = torch.tensor(values, dtype=dtype) + q = torch.quantile(t, torch.tensor(_R_PERCENTILES, dtype=dtype)) return {lbl: q[i].item() for i, lbl in enumerate(_R_PCT_LABELS)} - # ------------------------------------------------------------------ - # Stats - # ------------------------------------------------------------------ - def _n_reflections(self) -> int: """Total reflections in this target's subset across all datasets.""" dc = self._dataset_collection diff --git a/torchref/refinement/targets/collection/intensity.py b/torchref/refinement/targets/collection/intensity.py new file mode 100644 index 00000000..0469448a --- /dev/null +++ b/torchref/refinement/targets/collection/intensity.py @@ -0,0 +1,370 @@ +"""Two-moment intensity target for time-resolved collections. + +Merged Bragg intensities see the crystal-to-crystal activation distribution only through +its first two moments. With the branching among excited components conserved, the mixture +is exactly linear in the activation fraction, so the intensity is exactly quadratic and:: + + = |F(alpha_mean)|^2 + sigma_alpha^2 |dF/dalpha|^2 + +holds for *any* activation distribution and any number of components -- an identity, not a +truncation. The first term is what every existing target models; the second is the variance +the coherent model discards, and it is strictly positive, phase-blind, and largest exactly +where the difference signal is. + +The target works in **intensities** rather than amplitudes on purpose: the French-Wilson +conversion reshapes precisely the quadratic information the second moment lives in, so an +amplitude formulation would fit a distorted version of the quantity it is trying to measure. +Members must therefore carry ``I``/``SIGI``; there is no ``F**2`` fallback, because that +would silently reintroduce the distortion. +""" + +from typing import TYPE_CHECKING, Dict, List + +import torch + +from torchref.base.metrics.rfactor import rfactor_work_free +from torchref.base.targets.xray_likelihoods import ( + gaussian_per_refl, + intensity_var_from_sigma_obs, +) +from torchref.utils.stats import ( + VERBOSITY_DEBUG, + VERBOSITY_ESSENTIAL, + VERBOSITY_STANDARD, + StatEntry, + stat, +) + +from .base import CollectionXrayTarget + +if TYPE_CHECKING: + from torchref.io.datasets.collection import DatasetCollection + from torchref.model.model_collection import ModelCollection + from torchref.scaling.scaler_base import ScalerBase + + +class CollectionTwoMomentIntensityTarget(CollectionXrayTarget): + """ + Gaussian intensity likelihood at the two-moment forward model. + + Inherits the subset selector, the cache-reset discipline and the stats shape from + :class:`~torchref.refinement.targets.collection.base.CollectionXrayTarget`, so it + cannot disagree with the amplitude targets about which reflections it fits. + + The forward model is built from **one** set of per-component structure factors, + contracted twice: once with the fractions to get the mean, once with the activation + Jacobian to get its derivative. Both contractions go through the shared scaler, which + is affine in ``F_calc`` and mixes the bulk solvent linearly in the weights -- so the + second contraction returns the correctly scaled derivative rather than needing a + separate differentiation path. + + With ``lambda_twin`` fixed at zero the variance branch is not built at all. That makes + the coherent limit identical rather than merely equal: multiplying a live ``dF`` branch + by exactly zero would still propagate a non-finite ``F_calc`` into the loss. + + Parameters + ---------- + dataset_collection : DatasetCollection + Members must all carry intensities. + model_collection : ModelCollection + Supplies the components, the fractions and the activation moments. + scaler : ScalerBase, optional + Shared scaler. Needs ``forward_batched`` to scale the batch in one pass; without + one the unscaled mixture is used. + use_work_set : bool, optional + Legacy bool; superseded by ``use_set``. + use_set : str, optional + Canonical 3-way subset selector ``"work"``/``"free"``/``"val"``. + verbose : int, optional + Verbosity level. + base_weight : float, optional + Multiplies the summed loss on the work set. Intensities are squared amplitudes, + so this target's loss and gradient are on a completely different scale from the + amplitude targets it sits beside -- left at 1.0 it swamps them and the geometry + restraints with it. Default 1.0; use :meth:`calibrate_base_weight` to set it + from the data rather than by hand. + + Raises + ------ + ValueError + On construction, if any fitted dataset carries no intensities. + """ + + name: str = "collection_two_moment_intensity" + + #: Fits the merged INTENSITIES, so the base reads ``I``/``sigI``. The point of the + #: row: French-Wilson reshapes precisely the quadratic information the second moment + #: lives in, and there is deliberately no ``F**2`` fallback. + observable: str = "intensity" + + def __init__( + self, + dataset_collection: "DatasetCollection", + model_collection: "ModelCollection", + scaler: "ScalerBase" = None, + use_work_set: bool = True, + use_set: str = None, + verbose: int = 0, + base_weight: float = 1.0, + ): + super().__init__( + dataset_collection, + model_collection, + scaler=scaler, + use_work_set=use_work_set, + use_set=use_set, + verbose=verbose, + ) + self.base_weight = float(base_weight) + # Fail here rather than inside the first loss evaluation: LossState probes a + # target's forward at registration, and a traceback from there is much harder to + # trace back to "this MTZ had no intensity columns". + missing = [ + key for key in self._keys() if dataset_collection[key].I is None + ] + if missing: + raise ValueError( + f"Datasets {missing} carry no intensities. The two-moment target fits " + f"merged intensities directly -- converting amplitudes back with F**2 " + f"would reintroduce the French-Wilson distortion it exists to avoid. " + f"Supply reflection files with I/SIGI columns." + ) + + + def _row_indices(self, keys: List[str]) -> List[int]: + """Rows of the collection's fraction matrix corresponding to ``keys``.""" + order = self._model_collection.keys() + return [order.index(k) for k in keys] + + def _scale_batch(self, fcalc_batch, weights): + """Apply the shared scaler to a ``[T, R]`` batch with per-row weights.""" + scaler = self._scaler + if scaler is None: + return fcalc_batch + return scaler.forward_batched(fcalc_batch, weights) + + def intensity_model(self, recalc: bool = False) -> torch.Tensor: + """The two-moment predicted intensities, shape ``(n_datasets, n_reflections)``. + + Parameters + ---------- + recalc : bool, optional + Force recomputation of the component structure factors. + + Returns + ------- + torch.Tensor + Predicted intensities, rows aligned with :meth:`_keys`. + """ + dc, mc = self._dataset_collection, self._model_collection + keys = self._keys() + rows = self._row_indices(keys) + + components = dc.component_structure_factors(mc, recalc=recalc) + weights = mc.fractions_matrix()[rows] + + mean = self._scale_batch( + mc.mix_component_fcalcs(components, weights), weights + ) + intensity = mean.abs() ** 2 + + sigma_alpha_sq = mc.sigma_alpha_sq + if self._variance_is_live(sigma_alpha_sq): + jacobian = mc.activation_jacobian()[rows] + derivative = self._scale_batch( + mc.mix_component_fcalcs(components, jacobian), jacobian + ) + intensity = intensity + sigma_alpha_sq * derivative.abs() ** 2 + return intensity + + def _variance_is_live(self, sigma_alpha_sq) -> bool: + """Whether the second moment contributes. + + False only when the dispersion is *exactly* zero and not refinable, in which case + the derivative branch is skipped entirely rather than multiplied by zero. + """ + mc = self._model_collection + if mc._lambda_fixed is None: + return True + return bool(sigma_alpha_sq.detach().ne(0).any()) + + def _stack_model(self, keys, recalc: bool = False) -> torch.Tensor: + """The two-moment intensity, not a per-dataset squared amplitude. + + Overridden because this row's prediction is **not** a function of one dataset at a + time: the mean mixes shared components across timepoints and the variance term is + built from the activation Jacobian over the same components. ``keys`` is accepted + for the base's signature; :meth:`intensity_model` derives the rows itself. + """ + return self.intensity_model(recalc=recalc) + + def _per_refl(self, ctx) -> torch.Tensor: + """Per-reflection Gaussian NLL of the observed intensities under the model.""" + # The residual is formed and masked BEFORE the Gaussian, so a masked-out row + # contributes an exact zero rather than a value that merely gets multiplied by + # zero. That matters if the model is ever non-finite on an unfitted row: here the + # `where` discards it, whereas `nll * mask` would propagate NaN into the sum. + residual = torch.where( + ctx.mask, ctx.obs - ctx.model, torch.zeros_like(ctx.obs) + ) + var = intensity_var_from_sigma_obs( + ctx.sigma, floor=self._sigma_floor(ctx.sigma, ctx.mask) + ) + return gaussian_per_refl( + residual, torch.zeros_like(residual), var, var_floor=0.0 + ) + + + def calibrate_base_weight( + self, reference, parameters, ratio: float = 1.0, floor: float = 1e-12 + ) -> float: + """Set ``base_weight`` so this target pushes as hard as ``reference``. + + Matched on the **gradient norm** with respect to the refined parameters, not on + the loss value. A large loss with a flat gradient moves nothing, so loss + magnitude is the wrong thing to equalise; what competes with the geometry and + similarity restraints is the size of the step this term asks for. + + This is the per-cycle gradient-ratio weighting the collection targets' base + weights were a stopgap for, applied once at setup rather than every cycle -- + enough to put an intensity target and an amplitude target on the same footing, + which is otherwise a several-orders-of-magnitude mismatch. + + Parameters + ---------- + reference : Target + The target to match, normally the difference target already driving the + refinement. + parameters : iterable of torch.nn.Parameter + The parameters actually being refined; only those with ``requires_grad`` + are used. + ratio : float, optional + Desired ratio of this target's gradient norm to the reference's. Default + 1.0 (equal footing); below 1 makes this target the junior partner. + floor : float, optional + Guard for a vanishing reference gradient. + + Returns + ------- + float + The ``base_weight`` that was set. + """ + params = [p for p in parameters if p.requires_grad] + if not params: + raise ValueError("No refinable parameters given; cannot calibrate.") + + def _grad_norm(target, scale_out=1.0): + grads = torch.autograd.grad( + target.forward(), params, retain_graph=False, allow_unused=True + ) + total = sum( + float((g.detach() ** 2).sum()) for g in grads if g is not None + ) + return (total**0.5) / scale_out + + saved = self.base_weight + self.base_weight = 1.0 + try: + own = _grad_norm(self) + finally: + self.base_weight = saved + + ref = _grad_norm(reference) + if own <= floor: + if self.verbose: + print( + " two-moment calibration: own gradient is ~0, leaving " + f"base_weight at {self.base_weight:.4g}" + ) + return self.base_weight + + self.base_weight = float(ratio * max(ref, floor) / own) + if self.verbose: + print( + f" two-moment weight calibration: |grad_ref|={ref:.4g}, " + f"|grad_self|={own:.4g} -> base_weight={self.base_weight:.4g}" + ) + return self.base_weight + + + def get_rfactor(self) -> Dict[str, object]: + """Per-dataset R-work / R-free against the two-moment amplitudes. + + The reported amplitude is ``sqrt(I_model)``, the RMS amplitude the two-moment + model actually predicts -- not ``|F(alpha_mean)|``, which is only its first term + and would not correspond to the loss being minimised. + + Overrides the base implementation, which is per-pair and would recompute the + component stack once per dataset. + """ + dc = self._dataset_collection + keys = self._keys() + per_dataset: Dict[str, tuple] = {} + rworks: List[float] = [] + rfrees: List[float] = [] + + with torch.no_grad(): + model = self.intensity_model(recalc=True) + amplitudes = model.clamp(min=0.0).sqrt() + for row, key in enumerate(keys): + rwork, rfree = rfactor_work_free(dc[key], amplitudes[row]) + per_dataset[key] = (rwork, rfree) + rworks.append(rwork) + rfrees.append(rfree) + + return { + "per_dataset": per_dataset, + "rwork_pct": self._percentiles(rworks), + "rfree_pct": self._percentiles(rfrees), + } + + def stats(self) -> Dict[str, StatEntry]: + """Base collection X-ray stats plus the activation moments. + + ``dI_frac`` is the mean fraction of the predicted intensity carried by the second + moment. It is what separates "the dispersion refined to zero" from "the dispersion + was never refined", which are otherwise indistinguishable in the summary. + """ + out = super().stats() + mc = self._model_collection + + with torch.no_grad(): + alpha = float(mc.alpha_mean) + lam = float(mc.lambda_twin) + sigma_sq = float(mc.sigma_alpha_sq) + + out["base_weight"] = stat(self.base_weight, VERBOSITY_STANDARD) + out["alpha_mean"] = stat(alpha, VERBOSITY_ESSENTIAL) + out["lambda_twin"] = stat(lam, VERBOSITY_ESSENTIAL) + out["sigma_alpha_sq"] = stat(sigma_sq, VERBOSITY_STANDARD) + out["alpha_sd"] = stat(sigma_sq**0.5, VERBOSITY_STANDARD) + + keys = self._keys() + if keys and self._variance_is_live(mc.sigma_alpha_sq): + rows = self._row_indices(keys) + components = self._dataset_collection.component_structure_factors( + mc, recalc=True + ) + jacobian = mc.activation_jacobian()[rows] + derivative = self._scale_batch( + mc.mix_component_fcalcs(components, jacobian), jacobian + ) + variance_term = mc.sigma_alpha_sq * derivative.abs() ** 2 + total = self.intensity_model(recalc=False) + mask = self._dataset_collection.stack_masks( + keys, use_set=self.use_set + ) + denom = total[mask].abs().clamp(min=1e-12) + out["dI_frac"] = stat( + float((variance_term[mask] / denom).mean()), VERBOSITY_STANDARD + ) + else: + out["dI_frac"] = stat(0.0, VERBOSITY_STANDARD) + + branching = mc.branching() + for name, row in mc._branching_rows.items(): + for k in range(branching.shape[1]): + out[f"q_{name}_{k + 1}"] = stat( + float(branching[row, k]), VERBOSITY_DEBUG + ) + return out diff --git a/torchref/refinement/targets/collection/xray.py b/torchref/refinement/targets/collection/xray.py index f673c962..2b4587ec 100644 --- a/torchref/refinement/targets/collection/xray.py +++ b/torchref/refinement/targets/collection/xray.py @@ -2,27 +2,52 @@ One X-ray likelihood across a paired ``DatasetCollection`` + ``ModelCollection``, keys matched so each timepoint dataset meets its own mixed model: -:class:`CollectionDifferenceTarget` (mean-based differences, the primary optimization -driver), :class:`CollectionRiceTarget` (per-timepoint Rice at ``beta = sigma**2``) and -:class:`CollectionMLTarget` (Rice with one shared Luzzati ``beta`` pooled over all -datasets' free reflections, owned by the target rather than the scaler). -All three inherit the single-dataset subset/masking/R-factor contract from -:class:`~torchref.refinement.targets.collection.base.CollectionXrayTarget`. +:class:`CollectionDifferenceTarget` + Mean-based differences on amplitudes; the primary optimization driver. +:class:`CollectionDifferenceIntensityTarget` + The same on intensities -- the whole class is one ``observable`` declaration. +:class:`CollectionDifferenceSigmaDTarget` + The amplitude difference centred on ``alpha * dF_calc`` with the unexplained + difference power ``beta_model`` added to the measurement variance, both from a + sigma_D fit on the free set. +:class:`CollectionMLTarget` + Read MLF per dataset at one shared Luzzati ``beta``, pooled over every dataset's free + reflections and owned by the target rather than the scaler. The absolute channel. + +All of them get their observations, model, mask and likelihood seam from +:class:`~torchref.refinement.targets.collection.base.CollectionXrayTarget`, so each row is +a ``_per_refl`` and nothing else. The selectable set is +:data:`~torchref.refinement.targets.collection._specs.COLLECTION_XRAY_TARGETS`. + """ from typing import TYPE_CHECKING, Dict -import numpy as np import torch from torchref.base.reciprocal import get_scattering_vectors -from torchref.base.targets.xray_likelihoods import complex_var_from_beta, rice_math -from torchref.refinement.model_error_estimation.sigma_a import SigmaAEstimator, epsilon_from_hkl +from torchref.base.targets.xray_likelihoods import ( + complex_var_from_beta, + gaussian_per_refl, + rice_per_refl, +) +from torchref.refinement.model_error_estimation.sigma_a import ( + SigmaAEstimator, + epsilon_from_hkl, +) +from torchref.refinement.model_error_estimation.sigma_d import ( + SigmaDConfig, + SigmaDEstimator, +) from torchref.utils.stats import VERBOSITY_STANDARD, StatEntry, stat -from ._util import _LOG_2PI, _scale_fcalc -from .base import CollectionXrayTarget +from ._util import common_geom +from .base import ( + CollectionSigmaALossInputs, + CollectionSigmaDLossInputs, + CollectionXrayTarget, +) if TYPE_CHECKING: from torchref.io.datasets.collection import DatasetCollection @@ -77,6 +102,9 @@ class CollectionDifferenceTarget(CollectionXrayTarget): name: str = "difference_xray" + #: There is no difference from a single dataset. + min_datasets: int = 2 + def __init__( self, dataset_collection: "DatasetCollection", @@ -97,113 +125,101 @@ def __init__( ) self.normalize = normalize - def forward(self) -> torch.Tensor: - """Summed Gaussian NLL of the difference-from-mean; 0.0 if fewer than 2 sets.""" - dc = self._dataset_collection - mc = self._model_collection - - all_keys = self._keys() - N = len(all_keys) - if N < 2: - return torch.tensor(0.0, device=dc.hkl.device) - - # Clear caches so a preceding no-grad stats()/get_rfactor() call cannot - # leave a detached tensor that breaks the loss backward. - self._reset_model_caches() - - F_obs_list, sigma_list, mask_list, F_calc_list = [], [], [], [] + def _loss_inputs(self, recalc: bool = False): + """The base's stack, with the mask narrowed across datasets. - for key in all_keys: - data = dc[key] - model = mc[key] + A reflection counts only if it is in this target's subset in **every** dataset: + the per-reflection mean ties them together, so a reflection missing from one + member would silently shift the reference for all the others. Narrowed here + rather than inside :meth:`_per_refl` so ``forward``'s sum and ``residuals``' + array agree on which reflections count. + """ + ctx = super()._loss_inputs(recalc=recalc) + mask_all = ctx.mask.all(dim=0, keepdim=True).expand_as(ctx.mask) + return ctx._replace(mask=mask_all) - F_obs, sigma = data.get_corrected_data() - F_calc = self._scaled_amp_full(data, model, recalc=False) - # Validity + work/free/val selection, validation carved out of both. - mask = self._subset(data).mask + def _per_refl(self, ctx) -> torch.Tensor: + """Gaussian NLL of the difference-from-mean, per reflection and unreduced.""" + N = len(ctx.keys) + mask_all = ctx.mask[0] # (n_hkl,) -- every row is the same after _loss_inputs - F_obs_list.append(F_obs) - sigma_list.append(sigma) - mask_list.append(mask) - F_calc_list.append(F_calc) + delta_obs = ctx.obs - ctx.obs.mean(dim=0) + delta_calc = ctx.model - ctx.model.mean(dim=0) - F_obs_stack = torch.stack(F_obs_list) # (N, n_hkl) - sigma_stack = torch.stack(sigma_list) # (N, n_hkl) - mask_stack = torch.stack(mask_list) # (N, n_hkl) - F_calc_stack = torch.stack(F_calc_list) # (N, n_hkl) + # Var(F_i - F_mean) = sigma_i^2 (1 - 2/N) + (sum_j sigma_j^2) / N^2 + sum_sigma_sq = (ctx.sigma**2).sum(dim=0) + sigma_diff_sq = ctx.sigma**2 * (1 - 2.0 / N) + sum_sigma_sq / (N**2) + sigma_diff = torch.sqrt(sigma_diff_sq.clamp(min=1e-12)) - # A reflection must be in this subset in ALL datasets. - mask_all = mask_stack.all(dim=0) # (n_hkl,) - - F_mean_obs = F_obs_stack.mean(dim=0) # (n_hkl,) - F_calc_mean = F_calc_stack.mean(dim=0) # (n_hkl,) + # Mask via torch.where, not boolean indexing: no nonzero() device sync. + delta_obs = torch.where(mask_all, delta_obs, torch.zeros_like(delta_obs)) + delta_calc = torch.where(mask_all, delta_calc, torch.zeros_like(delta_calc)) + sigma_diff = torch.where(mask_all, sigma_diff, torch.ones_like(sigma_diff)) - delta_F_obs = F_obs_stack - F_mean_obs - delta_F_calc = F_calc_stack - F_calc_mean + # Floored on the PROPAGATED difference sigma, not the raw measurement sigma -- + # that is the quantity dividing the residual here. + sigma_safe = sigma_diff.clamp(min=self._sigma_floor(sigma_diff, ctx.mask)) - # Var(F_i - F_mean) = σ_i²·(1 - 2/N) + (Σ_j σ_j²) / N² - sum_sigma_sq = (sigma_stack**2).sum(dim=0) # (n_hkl,) - sigma_diff_sq = sigma_stack**2 * (1 - 2.0 / N) + sum_sigma_sq / (N**2) - sigma_diff = torch.sqrt(sigma_diff_sq.clamp(min=1e-12)) # (N, n_hkl) + nll = gaussian_per_refl(delta_obs, delta_calc, sigma_safe**2, var_floor=0.0) + # A single NaN would poison the whole gradient; 1e6 lets the step be rejected. + return torch.where(torch.isfinite(nll), nll, torch.full_like(nll, 1e6)) - # Mask via torch.where, not boolean indexing: no nonzero() device sync. - delta_F_obs = torch.where(mask_all, delta_F_obs, torch.zeros_like(delta_F_obs)) - delta_F_calc = torch.where( - mask_all, delta_F_calc, torch.zeros_like(delta_F_calc) - ) - sigma_diff = torch.where(mask_all, sigma_diff, torch.ones_like(sigma_diff)) - # Floor sigma at 10% of its median so a zero sigma cannot blow up. - eps = ( - torch.median(sigma_diff[:, mask_all].reshape(-1)) * 1e-1 - if mask_all.any() - else 1e-3 - ) - sigma_safe = sigma_diff.clamp(min=eps) +class CollectionDifferenceIntensityTarget(CollectionDifferenceTarget): + """The difference-from-mean target on **intensities** instead of amplitudes. - diff = delta_F_obs - delta_F_calc - nll = 0.5 * (diff / sigma_safe) ** 2 + torch.log(sigma_safe) + 0.5 * _LOG_2PI + A one-line row: the base reads ``I``/``sigI`` and predicts ``|F_calc|**2``, and the + difference-from-mean algebra is observable-agnostic -- + ``Var(x_i - x_mean) = sigma_i^2 (1 - 2/N) + (sum_j sigma_j^2)/N^2`` holds for any + quantity whose members share a mean. So the whole row is the ``observable`` + declaration, which is the point of having the axis at all. - # A single NaN would poison the whole gradient; 1e6 lets the step be rejected. - nll = torch.where(torch.isfinite(nll), nll, torch.full_like(nll, 1e6)) + Prefer it over :class:`CollectionDifferenceTarget` when the difference signal is weak + relative to the measurement error, because ``F_obs`` on a merged dataset is a + French-Wilson posterior: strictly positive, so it reshapes exactly the weak reflections + a small difference lives in, and it cannot represent a negative intensity at all. + Prefer the amplitude row when the output difference *map* is the product, since the DED + coefficients are amplitudes and keeping the loss and the map in one space is one fewer + conversion to get wrong. - total_nll = (nll * mask_all).sum() + Both rows are offered rather than one chosen: which wins is a property of a dataset's + signal-to-noise, not something to settle once in the library. + """ - return total_nll + name: str = "difference_intensity_xray" + observable: str = "intensity" # ========================================================================= -# CollectionRiceTarget +# CollectionDifferenceSigmaDTarget # ========================================================================= -class CollectionRiceTarget(CollectionXrayTarget): - """ - Multi-timepoint Rice maximum-likelihood amplitude target. +class CollectionDifferenceSigmaDTarget(CollectionDifferenceTarget): + """The difference-from-mean Gaussian with a sigma_D error model. - Rice NLL for acentrics and the folded-normal form for centrics, per timepoint. - The per-timepoint losses are independent, so reflections and masks are - concatenated across timepoints and masked once rather than reduced in a - Python loop. + The parent compares ``dF_obs`` with ``dF_calc`` under the measurement variance + alone. Here the likelihood is centred on ``alpha * dF_calc`` and its variance is + ``beta_model + sigma_diff**2``: ``alpha`` is the Gaussian coupling of the model + difference to the true one and ``beta_model`` the difference power the model does + not explain, both per resolution shell from + :class:`~torchref.refinement.model_error_estimation.sigma_d.SigmaDEstimator` fitted + on the pooled **free** reflections of the timepoint rows, with the dark-amplitude + power law carried per reflection. A poor light model therefore inflates the + variance where it fails instead of pulling the coordinates toward noise. + + At ``N = 2`` the timepoint row's difference from the mean is half the dark + subtraction; ``S`` and ``sigma_diff**2`` scale together, so the estimate is + invariant to that factor. The estimate is cached until :meth:`maintenance`, which + ``LossState`` calls after each optimizer-step block. Parameters ---------- - dataset_collection : DatasetCollection - model_collection : ModelCollection - scaler : ScalerBase, optional - Single scaler applied to each timepoint's F_calc. - normalize : bool - Unused placeholder. ``forward`` always returns the unnormalised summed - NLL regardless of this flag. - use_work_set : bool - Legacy bool; superseded by ``use_set``. If True, loss on the work set. - use_set : str, optional - Canonical 3-way subset selector ``"work"``/``"free"``/``"val"``. - verbose : int - Verbosity level. + sigma_d_config : SigmaDConfig, optional + Exponent and shrinkage settings; the module defaults when omitted. """ - name: str = "collection_rice_xray" + name: str = "difference_sigma_d_xray" def __init__( self, @@ -214,89 +230,113 @@ def __init__( use_work_set: bool = True, use_set: str = None, verbose: int = 0, + sigma_d_config: SigmaDConfig = None, ): super().__init__( dataset_collection, model_collection, scaler=scaler, + normalize=normalize, use_work_set=use_work_set, use_set=use_set, verbose=verbose, ) - self.normalize = normalize + # Constructed once; the cache lives until maintenance() resets it. + self._sigma_d = SigmaDEstimator(sigma_d_config) + self._eps_common: torch.Tensor = None + self._dss_common: torch.Tensor = None + self._geom_key: int = None - def _keys(self): - """Rice fits the timepoints only — the dark reference is excluded.""" - mc = self._model_collection - dc = self._dataset_collection - return [n for n in mc.timepoint_names if n in dc] + def _common_geom(self): + """``(epsilon, d_star_sq)`` on the common HKL, cached per dark dataset.""" + data = self._dataset_collection[self._model_collection.dark_key] + key = id(data) + if self._eps_common is None or self._geom_key != key: + self._eps_common, self._dss_common = common_geom(data) + self._geom_key = key + return self._eps_common, self._dss_common - def forward(self) -> torch.Tensor: - """Summed Rice NLL over every timepoint's reflections in this subset.""" + @staticmethod + def _difference_terms(ctx): + """Difference-from-mean observations, model and propagated sigma, ``(N, n_hkl)``.""" + N = len(ctx.keys) + delta_obs = ctx.obs - ctx.obs.mean(dim=0) + delta_calc = ctx.model - ctx.model.mean(dim=0) + sum_sigma_sq = (ctx.sigma**2).sum(dim=0) + sigma_diff_sq = ctx.sigma**2 * (1 - 2.0 / N) + sum_sigma_sq / (N**2) + return delta_obs, delta_calc, torch.sqrt(sigma_diff_sq.clamp(min=1e-12)) + + def _loss_inputs(self, recalc: bool = False): + """The parent's stack plus ``alpha`` and ``beta_model`` on the common HKL. + + The estimator sees the timepoint rows only (the dark row is the reference the + differences are taken against), their free reflections, and a detached model + difference, so gradients reach the models only through ``ctx.model``. + """ + ctx = super()._loss_inputs(recalc=recalc) + delta_obs, delta_calc, sigma_diff = self._difference_terms(ctx) + delta_calc = delta_calc.detach() + dark = ctx.keys.index(self._model_collection.dark_key) + rows = [i for i in range(len(ctx.keys)) if i != dark] or [dark] dc = self._dataset_collection - mc = self._model_collection + eps, dss = self._common_geom() + dtype = ctx.obs.dtype + eps, dss = eps.to(dtype), dss.to(dtype) + f_dark = ctx.obs[dark] + # The free set, independent of this target's own subset; the estimator drops + # non-finite observations itself. + fit_mask = torch.cat( + [dc[ctx.keys[i]].free.mask.to(ctx.mask.device) for i in rows] + ) + n_rows = len(rows) + est = self._sigma_d.get( + torch.cat([delta_obs[i] for i in rows]), + torch.cat([sigma_diff[i] for i in rows]), + eps.repeat(n_rows), + dss.repeat(n_rows), + f_dark.repeat(n_rows), + fit_mask, + delta_calc=torch.cat([delta_calc[i] for i in rows]), + target_dss=dss, + out_epsilon=eps, + out_f_dark=f_dark, + out_sigma_diff=sigma_diff[rows[0]], + ) + return CollectionSigmaDLossInputs( + *ctx, alpha=est.alpha.to(dtype), beta_model=est.beta_model.to(dtype) + ) - tp_names = self._keys() - if not tp_names: - return torch.tensor(0.0, device=mc.device) - - self._reset_model_caches() - - fo_parts, fc_parts, sig_parts, cen_parts, mask_parts = [], [], [], [], [] - for tp_name in tp_names: - data = dc[tp_name] - model = mc[tp_name] - - F_obs, sigma = data.get_corrected_data() - F_calc = self._scaled_amp_full(data, model, recalc=False) - centric = data.centric - if centric is None: - centric = torch.zeros( - len(F_obs), dtype=torch.bool, device=F_obs.device - ) - - fo_parts.append(F_obs) - fc_parts.append(F_calc) - sig_parts.append(sigma) - cen_parts.append(centric) - mask_parts.append(self._subset(data).mask) - - # Mask the flat arrays once: compact, no per-dataset Python reduction. - mask = torch.cat(mask_parts) - F_obs = torch.cat(fo_parts)[mask] - F_calc = torch.cat(fc_parts)[mask] - sigma = torch.cat(sig_parts)[mask] - centric = torch.cat(cen_parts)[mask] - - if F_obs.numel() == 0: - return torch.tensor(0.0, device=mc.device) - - # Plain Rice: the model-error variance IS the measurement variance here. - beta = sigma**2 - eb = beta.clamp(min=1e-6) - - # --- Acentric Rice NLL --- - term1 = -torch.log(2 * F_obs / eb + 1e-12) - term2 = F_obs**2 / eb - term3 = F_calc**2 / eb - arg_bessel = (2 * F_obs * F_calc / eb).clamp(max=1e6) - term4 = -(torch.log(torch.special.i0e(arg_bessel) + 1e-12) + arg_bessel) - loss_acentric = term1 + term2 + term3 + term4 - - # --- Centric NLL --- - term1_c = -0.5 * torch.log(2 / (np.pi * eb) + 1e-12) - term2_c = F_obs**2 / (2 * eb) - term3_c = F_calc**2 / (2 * eb) - term4_c = -(F_obs * F_calc) / eb - arg_exp = (-2 * F_obs * F_calc / eb).clamp(min=-80.0, max=80.0) - term5_c = -torch.log((1 + torch.exp(arg_exp)) / 2 + 1e-12) - loss_centric = term1_c + term2_c + term3_c + term4_c + term5_c - - loss = torch.where(centric, loss_centric, loss_acentric) + def _per_refl(self, ctx) -> torch.Tensor: + """Gaussian NLL of the difference about ``alpha * dF_calc`` with variance + ``beta_model + sigma_diff**2``, per reflection and unreduced.""" + mask_all = ctx.mask[0] + delta_obs, delta_calc, sigma_diff = self._difference_terms(ctx) + delta_obs = torch.where(mask_all, delta_obs, torch.zeros_like(delta_obs)) + delta_calc = torch.where(mask_all, delta_calc, torch.zeros_like(delta_calc)) + sigma_diff = torch.where(mask_all, sigma_diff, torch.ones_like(sigma_diff)) + sigma_safe = sigma_diff.clamp(min=self._sigma_floor(sigma_diff, ctx.mask)) + var = ctx.beta_model.unsqueeze(0) + sigma_safe**2 + mean = ctx.alpha.unsqueeze(0) * delta_calc + nll = gaussian_per_refl(delta_obs, mean, var, var_floor=0.0) # A single NaN would poison the whole gradient; 1e6 lets the step be rejected. - loss = torch.where(torch.isfinite(loss), loss, torch.full_like(loss, 1e6)) + return torch.where(torch.isfinite(nll), nll, torch.full_like(nll, 1e6)) - return loss.sum() + def maintenance(self) -> None: + """Invalidate the sigma_D estimate so it is refitted from the updated models on + the next forward (``LossState`` calls this after each optimizer-step block).""" + self._sigma_d.reset() + + def stats(self) -> Dict[str, StatEntry]: + """Base collection X-ray stats plus the sigma_D fit summary.""" + out = super().stats() + sh = self._sigma_d.shells + if sh is not None: + out["sigma_d_gamma"] = stat(float(sh.gamma), VERBOSITY_STANDARD) + out["sigma_d_tau"] = stat(float(sh.tau), VERBOSITY_STANDARD) + out["sigma_d_shells_without_power"] = stat( + float(sh.diagnostics["n_s2_clamped"]), VERBOSITY_STANDARD + ) + return out # ========================================================================= @@ -394,89 +434,63 @@ def _common_geom(self): self._eps_common, self._dss_common, self._geom_key = eps, dss, key return self._eps_common, self._dss_common - def forward(self) -> torch.Tensor: - """Summed Read-MLF loss over all datasets at the shared beta, work-set weighted.""" - dc = self._dataset_collection - mc = self._model_collection + def _loss_inputs(self, recalc: bool = False): + """The base's stack, plus the shared model-error variance this row needs. - all_keys = self._keys() - if not all_keys: - return torch.tensor(0.0, device=mc.device) + ``beta`` and ``epsilon`` are fitted **once** on the pooled free reflections of + every data-model pair and mapped onto the common HKL, so one per-reflection + variance serves every dataset -- and since they live on that common grid they + broadcast over the dataset axis rather than being tiled. + Detached, so gradients reach the models only through ``F_calc``. Fitted on the + free set (``data.free`` excludes validation), and cached until + :meth:`maintenance` resets it. + """ + ctx = super()._loss_inputs(recalc=recalc) + dc = self._dataset_collection eps_common, dss_common = self._common_geom() - dtype = dss_common.dtype - - # Clear any structure-factor cache populated under no_grad (e.g. by the - # scaler's joint initialization or a preceding stats() call) so the - # F_calc computed below carries a live graph. - self._reset_model_caches() - - # Each pair's F_calc is computed once WITH grad for the loss; detached copies - # feed the pooled single-beta estimate. - fo_parts, fc_parts, cen_parts, msk_parts = [], [], [], [] - eps_parts, dss_parts, free_parts, sig_parts = [], [], [], [] - for key in all_keys: - data = dc[key] - model = mc[key] - - F_obs, sig_obs = data.get_corrected_data() - F_obs = F_obs.to(dtype) - sig_parts.append(sig_obs.to(dtype).reshape(-1)) - F_calc = self._scaled_amp_full(data, model, recalc=False).to(dtype) - centric = data.centric - if centric is None: - centric = torch.zeros( - len(F_obs), dtype=torch.bool, device=F_obs.device - ) - - fo_parts.append(F_obs) - fc_parts.append(F_calc) - cen_parts.append(centric) - msk_parts.append(self._subset(data).mask) - - # Beta is estimated on the free set; data.free excludes validation. - free_parts.append(data.free.mask) - eps_parts.append(eps_common.to(dtype)) - dss_parts.append(dss_common.to(dtype)) - - # One shared (beta, epsilon) for all datasets, mapped onto the common HKL via - # target_dss. Detached; cached until maintenance() resets it. - _est = self._sigma_a.get( - torch.cat(fo_parts), - torch.cat([fc.detach() for fc in fc_parts]), # beta needs no gradient - torch.cat(cen_parts), - torch.cat(eps_parts), - torch.cat(dss_parts), - torch.cat(free_parts), + dtype = ctx.obs.dtype + + centric = dc.get_centric_flags() + if centric is None: + centric = torch.zeros( + ctx.obs.shape[-1], dtype=torch.bool, device=ctx.obs.device + ) + + n_ds = len(ctx.keys) + free = torch.cat([dc[k].free.mask for k in ctx.keys]) + est = self._sigma_a.get( + ctx.obs.reshape(-1).to(dtype), + # beta needs no gradient. + ctx.model.detach().reshape(-1).to(dtype), + centric.repeat(n_ds), + eps_common.to(dtype).repeat(n_ds), + dss_common.to(dtype).repeat(n_ds), + free, out_epsilon=eps_common.to(dtype), target_dss=dss_common, # Always passed, as at every other call site: it is what makes sigma_A the # correlation with the noise-free amplitudes rather than with the noisy data. - sigma_obs=torch.cat(sig_parts), + sigma_obs=ctx.sigma.reshape(-1).to(dtype), ) - beta, eps = _est.beta, _est.epsilon - - # beta/eps live on the common HKL, so they are tiled once per dataset to line - # up with the concatenation order. - n_ds = len(all_keys) - F_obs_cat = torch.cat(fo_parts) - F_calc_cat = torch.cat(fc_parts) - centric_cat = torch.cat(cen_parts) - mask_cat = torch.cat(msk_parts) - beta_cat = beta.to(F_obs_cat.dtype).repeat(n_ds) - eps_cat = eps.to(F_obs_cat.dtype).repeat(n_ds) if eps is not None else None - - # TOTAL variance (`est.beta`, not `beta_model`): this likelihood does not - # account for sigma_obs itself, so the measurement variance must stay inside beta. - total = rice_math( - F_obs_cat, F_calc_cat, complex_var_from_beta(beta_cat, eps_cat), - centric_cat, mask=mask_cat, + return CollectionSigmaALossInputs( + *ctx, + centric=centric, + beta=est.beta.to(dtype), + epsilon=None if est.epsilon is None else est.epsilon.to(dtype), ) - # Base weight drives refinement; applied on the work set only. - if self.use_work_set: - total = self.base_weight * total - return total + def _per_refl(self, ctx) -> torch.Tensor: + """Read MLF per reflection: Rice for acentrics, folded normal for centrics. + + TOTAL variance (``est.beta``, not ``beta_model``): this likelihood does not + account for ``sigma_obs`` itself, so the measurement variance stays inside + ``beta``. + """ + Sigma = complex_var_from_beta(ctx.beta, ctx.epsilon) + nll = rice_per_refl(ctx.obs, ctx.model, Sigma, ctx.centric) + # A single NaN would poison the whole gradient; 1e6 lets the step be rejected. + return torch.where(torch.isfinite(nll), nll, torch.full_like(nll, 1e6)) def maintenance(self) -> None: """Invalidate the shared beta so it is re-estimated from the updated diff --git a/torchref/refinement/targets/combined.py b/torchref/refinement/targets/combined.py index 545201a2..032d8e9f 100644 --- a/torchref/refinement/targets/combined.py +++ b/torchref/refinement/targets/combined.py @@ -20,6 +20,8 @@ ) from torchref.refinement.targets.adp import ( ADPSimilarityTarget, ADPLocalityTarget, ADPSigdTarget, + NodeLoadTarget, + NodeSmoothnessTarget, ) from torchref.utils.stats import ( VERBOSITY_DETAILED, @@ -363,15 +365,43 @@ class TotalADPTarget(CombinedModelTargets): """ def _create_targets(self) -> Dict[str, Target]: - """Build the three ADP component targets.""" - print("Initializing TotalADPTarget with component targets...") - return { - "simu": ADPSimilarityTarget(self.model, verbose=self.verbose), - "locality": ADPLocalityTarget( - self.model, verbose=self.verbose - ), - "sigd": ADPSigdTarget(self.model, verbose=self.verbose), - } + """Build the ADP component targets that apply to the model's representation. + + Only the applicable ones are registered, rather than registering all of them and + zero-weighting the inapplicable half. ``simu`` and ``locality`` restrain by + penalty exactly the spatial smoothness a node field enforces by construction, so + in field mode they are not a weak prior but a duplicate of the parametrisation; + and ``node_load`` / ``node_smoothness`` have nothing to act on off it. + + Registering-then-zeroing would cost nothing at run time --- ``LossState.aggregate`` + skips a zero-weight target --- but it leaves the correctness of the setup resting + on a weight, so anyone who touches the ``adp`` group weight for their own reasons + silently re-enables a double-counted restraint. Whether a term applies is a + property of the representation, not a number to be tuned. + + ``sigd`` applies either way: it is a prior on the marginal B distribution, which + a field constrains no more than a per-atom parametrisation does. + """ + if self.model.adp_is_field: + targets = { + "sigd": ADPSigdTarget(self.model, verbose=self.verbose), + "node_load": NodeLoadTarget(self.model, verbose=self.verbose), + "node_smoothness": NodeSmoothnessTarget( + self.model, verbose=self.verbose + ), + } + else: + targets = { + "simu": ADPSimilarityTarget(self.model, verbose=self.verbose), + "locality": ADPLocalityTarget(self.model, verbose=self.verbose), + "sigd": ADPSigdTarget(self.model, verbose=self.verbose), + } + if self.verbose > 0: + print( + "Initializing TotalADPTarget with component targets: " + + ", ".join(targets) + ) + return targets def print_statistics(self) -> None: """ diff --git a/torchref/refinement/targets/dataset_scaling.py b/torchref/refinement/targets/dataset_scaling.py new file mode 100644 index 00000000..542b55c4 --- /dev/null +++ b/torchref/refinement/targets/dataset_scaling.py @@ -0,0 +1,42 @@ +"""Joint observed-dataset scaling target, independent of structural models.""" + +from typing import TYPE_CHECKING + +import torch + +from torchref.base.targets.dataset_scaling import dataset_scaling_loss +from torchref.refinement.targets.base import Target + +if TYPE_CHECKING: + from torchref.scaling.dataset_scaler import DatasetScaler + + +class DatasetScalingTarget(Target): + """Fit a shared consensus to the scaler's training observations. + + Parameters + ---------- + scaler : DatasetScaler + Owner of the centered log-scale and anisotropy parameters. Its prepared + observations exclude held-out reflections from every participating dataset. + """ + + name = "dataset_scaling" + + def __init__(self, scaler: "DatasetScaler") -> None: + super().__init__(device=scaler.device) + self.scaler = scaler + self._adopt_device(scaler) + + def forward(self) -> torch.Tensor: + """Return the dimensionless loss per independent training contrast.""" + scaler = self.scaler + return ( + dataset_scaling_loss( + scaler.amplitudes, + scaler.sigmas, + scaler.log_corrections(scaler.hkl), + scaler.fit_mask, + ) + / scaler.n_contrasts + ) diff --git a/torchref/refinement/targets/difference.py b/torchref/refinement/targets/difference.py index 33870325..d0a75ab7 100644 --- a/torchref/refinement/targets/difference.py +++ b/torchref/refinement/targets/difference.py @@ -204,10 +204,10 @@ def _match_reflections(self): device = hkl_light.device self._matched_indices_light = torch.tensor( - matched_light, dtype=torch.long, device=device + matched_light, dtype=torch.long, device=device # dtype-ok: matched atom indices used for indexing; PyTorch requires int64 ) self._matched_indices_dark = torch.tensor( - matched_dark, dtype=torch.long, device=device + matched_dark, dtype=torch.long, device=device # dtype-ok: matched atom indices used for indexing; PyTorch requires int64 ) # Store common HKL (using light indices, they should be identical) diff --git a/torchref/refinement/targets/geometry/chiral.py b/torchref/refinement/targets/geometry/chiral.py index e10f567a..1cc6dcee 100644 --- a/torchref/refinement/targets/geometry/chiral.py +++ b/torchref/refinement/targets/geometry/chiral.py @@ -82,7 +82,7 @@ def get_violations(self, threshold: float = 0.5) -> Dict[str, torch.Tensor]: if "chiral" not in self.restraints.restraints: return { - "indices": torch.tensor([], dtype=torch.long, device=device).reshape( + "indices": torch.tensor([], dtype=torch.long, device=device).reshape( # dtype-ok: empty restraint index tensor; PyTorch requires int64 for indexing 0, 4 ), "volumes": torch.tensor([], device=device), diff --git a/torchref/refinement/targets/geometry/non_bonded.py b/torchref/refinement/targets/geometry/non_bonded.py index 5d56fbe4..637be377 100644 --- a/torchref/refinement/targets/geometry/non_bonded.py +++ b/torchref/refinement/targets/geometry/non_bonded.py @@ -241,7 +241,7 @@ def _compute_positions( return pos1, pos2, min_distances cell = self.model.cell - sg = self.model.symmetry + sg = self.model.spacegroup mate_source = xyz[indices[:, 1]] # (N_pairs, 3) -- gradients flow frac = cell.cartesian_to_fractional(mate_source) @@ -285,8 +285,8 @@ def forward(self) -> torch.Tensor: vdw_data["min_distances"], vdw_data.get("symop_indices"), vdw_data.get("cell_offsets"), - self.model.symmetry.matrices, - self.model.symmetry.translations, + self.model.spacegroup.matrices, + self.model.spacegroup.translations, self.model.cell.fractional_matrix, self.model.cell.inv_fractional_matrix, self._c_rep, self._r_exp, @@ -346,7 +346,7 @@ def get_violations(self, threshold: float = 0.0) -> Dict[str, torch.Tensor]: if "vdw" not in self.restraints.restraints: return { - "indices": torch.tensor([], dtype=torch.long, device=device).reshape( + "indices": torch.tensor([], dtype=torch.long, device=device).reshape( # dtype-ok: empty restraint index tensor; PyTorch requires int64 for indexing 0, 2 ), "violations": torch.tensor([], device=device), @@ -359,7 +359,7 @@ def get_violations(self, threshold: float = 0.0) -> Dict[str, torch.Tensor]: if indices is None or len(indices) == 0: return { - "indices": torch.tensor([], dtype=torch.long, device=device).reshape( + "indices": torch.tensor([], dtype=torch.long, device=device).reshape( # dtype-ok: empty restraint index tensor; PyTorch requires int64 for indexing 0, 2 ), "violations": torch.tensor([], device=device), diff --git a/torchref/refinement/targets/geometry/non_bonded_h.py b/torchref/refinement/targets/geometry/non_bonded_h.py index 39f6a1d8..8a249b24 100644 --- a/torchref/refinement/targets/geometry/non_bonded_h.py +++ b/torchref/refinement/targets/geometry/non_bonded_h.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from torchref.model.model import Model - from torchref.restraints.hydrogen_topology import HydrogenTopology + from torchref.topology.riding import HydrogenTopology class NonBondedHTarget(NonBondedTarget): @@ -93,7 +93,7 @@ def _compute_h_vdw_loss( ordering to do identity and real symmetry transforms in one pass. Other modes take the inline eager path below. """ - from torchref.restraints.hydrogen_topology import place_riding_hydrogens + from torchref.topology.riding import place_riding_hydrogens device = xyz.device @@ -113,8 +113,8 @@ def _compute_h_vdw_loss( return nonbonded_heavy_math( xyz_all, indices, h_topo.cand_min_dist, h_topo.cand_symop_idx, h_topo.cand_cell_offset, - self.model.symmetry.matrices, - self.model.symmetry.translations, + self.model.spacegroup.matrices, + self.model.spacegroup.translations, self.model.cell.fractional_matrix, self.model.cell.inv_fractional_matrix, self._c_rep, self._r_exp, @@ -123,7 +123,7 @@ def _compute_h_vdw_loss( # Slow path: gaussian / soft modes, inline eager. pos_i = xyz_all[h_topo.cand_idx_i] - n_asu = getattr(h_topo, 'n_asu_candidates', n_cand) + n_asu = h_topo.n_asu_candidates n_sym = n_cand - n_asu min_dist = h_topo.cand_min_dist @@ -134,7 +134,7 @@ def _compute_h_vdw_loss( if n_sym > 0: cell = self.model.cell - sg = self.model.symmetry + sg = self.model.spacegroup sym_source = xyz_all[h_topo.cand_idx_j[n_asu:]] frac = cell.cartesian_to_fractional(sym_source) R = sg.matrices[h_topo.cand_symop_idx[n_asu:]].to(frac.dtype) @@ -200,7 +200,7 @@ def get_violations(self, threshold: float = 0.0) -> Dict[str, torch.Tensor]: ``xyz_all`` -- so symmetry-mate H contacts are reported at their intra-ASU separation, unlike in the loss. """ - from torchref.restraints.hydrogen_topology import place_riding_hydrogens + from torchref.topology.riding import place_riding_hydrogens result = super().get_violations(threshold) @@ -234,7 +234,7 @@ def get_violations(self, threshold: float = 0.0) -> Dict[str, torch.Tensor]: def stats(self) -> Dict[str, any]: """Get statistics including H-VDW contacts.""" - from torchref.restraints.hydrogen_topology import place_riding_hydrogens + from torchref.topology.riding import place_riding_hydrogens result = super().stats() diff --git a/torchref/refinement/targets/similarity.py b/torchref/refinement/targets/similarity.py index 526ebc21..d01c3ef8 100644 --- a/torchref/refinement/targets/similarity.py +++ b/torchref/refinement/targets/similarity.py @@ -68,10 +68,10 @@ def __init__( # path (the one ``load_state_dict`` uses) would have no such buffers at all. # ``_build_atom_map`` overwrites them rather than creating them. self.register_buffer( - "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) + "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) # dtype-ok: index buffer for gather/index_select; PyTorch requires int64 ) self.register_buffer( - "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) + "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) # dtype-ok: index buffer for gather/index_select; PyTorch requires int64 ) if model_dark is not None and model_light is not None: self._build_atom_map() @@ -140,10 +140,10 @@ def _build_atom_map(self): "dark and light models" ) self.register_buffer( - "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) + "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) # dtype-ok: index buffer for gather/index_select; PyTorch requires int64 ) self.register_buffer( - "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) + "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) # dtype-ok: index buffer for gather/index_select; PyTorch requires int64 ) return @@ -166,13 +166,13 @@ def _build_atom_map(self): self.register_buffer( "_idx_dark", torch.tensor( - merged["_idx_dark"].values, dtype=torch.long, device=self.device + merged["_idx_dark"].values, dtype=torch.long, device=self.device # dtype-ok: atom index tensor used for indexing; PyTorch requires int64 ), ) self.register_buffer( "_idx_light", torch.tensor( - merged["_idx_light"].values, dtype=torch.long, device=self.device + merged["_idx_light"].values, dtype=torch.long, device=self.device # dtype-ok: atom index tensor used for indexing; PyTorch requires int64 ), ) diff --git a/torchref/refinement/targets/xray/__init__.py b/torchref/refinement/targets/xray/__init__.py index 6189e6ac..cfe98148 100644 --- a/torchref/refinement/targets/xray/__init__.py +++ b/torchref/refinement/targets/xray/__init__.py @@ -15,6 +15,7 @@ from .nll import NLLXrayTarget from .nll_beta import NLLBetaXrayTarget from .rice import RiceXrayTarget +from .observable import IntensityObservableMixin, NLLIntensityXrayTarget from .sigma_a import AlphaCentredMixin, SigmaALossInputs, SigmaAXrayTarget __all__ = [ @@ -23,6 +24,8 @@ "SigmaAXrayTarget", "SigmaALossInputs", "AlphaCentredMixin", + "IntensityObservableMixin", + "NLLIntensityXrayTarget", # the five selectable likelihood rows "NLLXrayTarget", "NLLBetaXrayTarget", diff --git a/torchref/refinement/targets/xray/_specs.py b/torchref/refinement/targets/xray/_specs.py index 98940dcb..1058ba0d 100644 --- a/torchref/refinement/targets/xray/_specs.py +++ b/torchref/refinement/targets/xray/_specs.py @@ -7,6 +7,31 @@ following :mod:`torchref.utils.backends`; string literals validated against a table, not enums, is the house convention. +## The observable + +Every row shares one forward model -- the scaled complex ``F_calc`` -- and each declares +which measured column it compares against, via ``spec.observable``: + +* **amplitude** ``F_obs`` vs ``|F_calc|`` (all the sigma_A rows, and ``ls``) +* **intensity** ``I_obs`` vs ``|F_calc|**2`` (``nll_i``) + +The intensity rows exist because ``F_obs`` on a merged dataset is a French-Wilson posterior +rather than a measurement: strictly positive, so it reshapes the weak tail and erases +negative intensities. A row whose signal lives in the *quadratic* part of the data reads +``I_obs`` directly. See :mod:`torchref.refinement.targets.xray.observable`, which is the +whole of that axis -- one ``get_data`` override, no runtime branch anywhere. + +Note the axis is **not square**: there is no intensity Rice, because Rice and the folded +normal are distributions *of an amplitude* and the intensity analogue is the exponential / +chi-square_1 Wilson distribution -- a different primitive, not a different variance. +R-factors stay on amplitudes for every row regardless, so they remain comparable across the +whole table. + +Intensity rows are **not** admissible as scale targets: +:data:`~torchref.scaling.scaler_base.SCALE_TARGETS` normalises its objective by +``1/sum(F_obs**2)``, which is dimensionally wrong for an ``O(F**4)`` loss, and that tuple +fails closed on anything it does not list. + ## The sigma_A family Each is a choice of (distribution) x (variance) x (mean): @@ -54,6 +79,7 @@ from .ml_noalpha import MLNoAlphaXrayTarget from .nll import NLLXrayTarget from .nll_beta import NLLBetaXrayTarget +from .observable import IntensityObservableMixin, NLLIntensityXrayTarget # noqa: F401 #: The mode built when none is given. DEFAULT_XRAY_MODE = "ml" @@ -75,18 +101,38 @@ class XrayTargetSpec: aliases Retired spellings kept working; resolving one emits a ``DeprecationWarning``. No row carries one at present, so the tests exercise this with their own table. + observable + Which measured column the row fits: ``"amplitude"`` or ``"intensity"``. Declarative + rather than a constructor flag, because it is a property of the row -- see + :mod:`torchref.refinement.targets.xray.observable`. Checked here against the class, + so a spec and its implementation cannot disagree. """ name: str target_cls: type doc: str aliases: Tuple[str, ...] = () + observable: str = "amplitude" def __post_init__(self): if not (isinstance(self.target_cls, type) and issubclass(self.target_cls, XrayTarget)): raise TypeError( f"{self.name}: target_cls {self.target_cls!r} is not an XrayTarget subclass" ) + if self.observable not in ("amplitude", "intensity"): + raise ValueError( + f"{self.name}: observable must be 'amplitude' or 'intensity', " + f"got {self.observable!r}" + ) + # The class declares its own observable (the mixin sets it); the spec must agree. + # Otherwise a row could advertise intensities while reading `sub.F`, which no test + # downstream of here would notice -- the loss would simply be wrong by 2|F|. + declared = getattr(self.target_cls, "observable", "amplitude") + if declared != self.observable: + raise ValueError( + f"{self.name}: spec says observable={self.observable!r} but " + f"{self.target_cls.__name__} says {declared!r}" + ) @dataclass(frozen=True) @@ -171,6 +217,13 @@ def by_name(self, name: str) -> XrayTargetSpec: doc="Gaussian amplitude NLL weighted by the experimental sigma only. No " "model-error term, so it does not control overfitting.", ), + XrayTargetSpec( + name="nll_i", + target_cls=NLLIntensityXrayTarget, + observable="intensity", + doc="Gaussian NLL on the observed INTENSITIES weighted by sigma(I). As 'nll' " + "but skips the French-Wilson conversion, which reshapes the weak tail.", + ), XrayTargetSpec( name="ls", target_cls=LeastSquaresXrayTarget, diff --git a/torchref/refinement/targets/xray/nll.py b/torchref/refinement/targets/xray/nll.py index edeb2d98..4524486d 100644 --- a/torchref/refinement/targets/xray/nll.py +++ b/torchref/refinement/targets/xray/nll.py @@ -1,6 +1,7 @@ -import torch from typing import TYPE_CHECKING +import torch + from torchref.base.targets.xray_likelihoods import ( amplitude_var_from_sigma_obs, nll_per_refl, @@ -24,13 +25,9 @@ class NLLXrayTarget(XrayTarget): here that does not. Was ``GaussianXrayTarget``; the taxonomy names the row, and "Gaussian" named the distribution, which ``nll_beta`` shares. - **Not a** :class:`SigmaAXrayTarget`, and deliberately so. Beyond needing no estimate, it - reads its amplitudes through :meth:`XrayTarget.get_data`, which goes via - ``ReflectionData._corrected_or_raw`` and falls back to **raw** amplitudes when the scaler - has not run; the sigma_A path calls ``get_corrected_data()``, which raises instead. - Moving this target onto that path would turn a silent fallback into a hard failure on - unscaled data -- a behaviour change, not a refactor. It would also lose the fused Triton - kernel and the ``median(sigma)*0.1`` clamp, neither of which the beta-variance path has. + Read observations through the dataset subset accessors, which expose live + corrections for ScaledDataset. The target uses a fused Triton kernel where + available and floors uncertainties at one tenth of their median. Attributes ---------- diff --git a/torchref/refinement/targets/xray/observable.py b/torchref/refinement/targets/xray/observable.py new file mode 100644 index 00000000..b9d54d74 --- /dev/null +++ b/torchref/refinement/targets/xray/observable.py @@ -0,0 +1,128 @@ +"""Select measured intensity observations for Gaussian X-ray targets. + +``IntensityObservableMixin`` reads ``I_obs`` and ``sigma(I)`` and predicts +``|F_calc|**2``, retaining negative intensities that French-Wilson amplitude +conversion reshapes. Likelihoods, subsets and masks come from the target class; +reported R-factors remain in amplitude space. Rice targets describe amplitudes +and therefore have no intensity variant. +""" + +from typing import Tuple + +import torch + +from torchref.base.targets.xray_likelihoods import ( + SIGMA_FLOOR_ABS, + SIGMA_FLOOR_FRAC, + gaussian_per_refl, + intensity_var_from_sigma_obs, + _masked_sum, +) + +from .base import XrayTarget + + +class IntensityObservableMixin: + """Read ``I_obs``/``sigma(I)`` and predict ``|F_calc|**2``. + + Mix in **before** an :class:`~.base.XrayTarget` subclass. The only override is + :meth:`get_data`, which is the single place the observable is chosen -- so a row + composed with this mixin cannot end up fitting intensities while reporting statistics + on amplitudes, and no method needs a runtime branch. + + Note there is deliberately no ``_scaled_F_calc_full`` override. That method feeds + :meth:`XrayTarget.get_rfactor`, and for a ``|F_calc|**2`` model its correct value is + ``sqrt(I_calc) == |F_calc|`` -- exactly what the inherited implementation returns. So + **R-factors stay on amplitudes for every row**, comparable across the whole table + regardless of which observable drove the loss. A row whose intensity model is *not* + the square of an amplitude (the two-moment model, where it is + ``|F|**2 + var*|dF|**2``) must override it to report ``sqrt`` of its own model. + """ + + #: Declared for the taxonomy table, and readable off any constructed target. + observable: str = "intensity" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Checked at construction rather than at the first forward: LossState probes + # `forward()` once when a target is registered, so a missing column would + # otherwise surface as a failure deep inside setup with no mention of the cause. + data = getattr(self, "_data", None) + if data is not None and getattr(data, "I", None) is None: + raise ValueError( + f"{type(self).__name__} fits intensities, but this dataset carries none. " + "Load an MTZ/mmCIF with an I column (`I-obs`/`intensity_meas`), or select " + "an amplitude row such as `nll`." + ) + + def get_data( + self, fcalc: torch.Tensor = None, sub=None + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, object]: + """``(I_obs, I_calc, sigma_I, centric, sub)`` -- the intensity twin of + :meth:`XrayTarget.get_data`, same tuple shape and same subset semantics. + + Both observation columns are the *corrected* views (``sub.I``/``sub.sigI``), in + which the scale and the anisotropy factor enter squared, so they are on the same + footing as the squared model amplitude from :meth:`get_I_calc_scaled`. + """ + if sub is None: + sub = self._subset() + + I_obs = sub.I + sigma = sub.sigI + centric = sub.centric + + if fcalc is not None: + I_calc_full = self.get_I_calc_scaled(fcalc=fcalc) + else: + I_calc_full = self.get_I_calc_scaled(recalc=False) + I_calc = sub.select(I_calc_full) + + return I_obs, I_calc, sigma, centric, sub + + def _sigma_floor(self) -> torch.Tensor: + """The intensity-sigma floor, taken from THIS target's own fitted subset. + + Computed here rather than inside the variance builder so it does not depend on + which reflections a particular call happens to pass. ``forward`` evaluates on the + target's subset while ``residuals`` evaluates on every reflection; a floor derived + from the argument therefore differs between them, and the same reflection scores + differently in the two -- 0.09% on a 1DAW work set, 1.8% on its free set, because + sigma(I) spans orders of magnitude where sigma(F) does not. + + Detached: it is a numerical safeguard, not a fitted quantity, and letting a + gradient run back through a median would make the loss depend on the ordering of + near-equal sigmas. + """ + sigma = self._subset().sigI + if sigma is None or sigma.numel() == 0: + return torch.as_tensor(1e-6) + return (torch.median(sigma).detach() * SIGMA_FLOOR_FRAC).clamp(min=SIGMA_FLOOR_ABS) + + +class NLLIntensityXrayTarget(IntensityObservableMixin, XrayTarget): + """``--xray-mode nll_i``: Gaussian intensity NLL weighted by the experimental sigma. + + NLL = 0.5*(I_obs - |F_calc|**2)**2/sigma_I**2 + log(sigma_I) + 0.5*log(2*pi) + + The intensity counterpart of :class:`~.nll.NLLXrayTarget`, and like it carries no + model-error term, so it does **not** control overfitting. + + Subclasses :class:`~.base.XrayTarget` directly rather than ``NLLXrayTarget``, because + that row's ``forward`` calls the fused Triton amplitude kernel + (``nll_sigma_obs_math``) which has no intensity counterpart. Here ``forward`` is the + structural ``_masked_sum(_per_refl(...))``, so it and :meth:`residuals` are the same + expression by construction rather than by test. + """ + + target_value: float = 1.0 + + def forward(self, fcalc: torch.Tensor = None) -> torch.Tensor: + """Summed Gaussian NLL of the observed intensities on this target's set.""" + return _masked_sum(self._per_refl(self._loss_inputs(fcalc=fcalc))) + + def _per_refl(self, ctx) -> torch.Tensor: + """Per-reflection Gaussian on the intensity. See :func:`gaussian_per_refl`.""" + I_obs, I_calc, sigma, _, _ = ctx + var = intensity_var_from_sigma_obs(sigma, floor=self._sigma_floor()) + return gaussian_per_refl(I_obs, I_calc, var, var_floor=0.0) diff --git a/torchref/refinement/targets/xray/rice.py b/torchref/refinement/targets/xray/rice.py index 728e881c..a5599e3a 100644 --- a/torchref/refinement/targets/xray/rice.py +++ b/torchref/refinement/targets/xray/rice.py @@ -22,22 +22,20 @@ class RiceXrayTarget(XrayTarget): empirically it was the worst-behaved target measured, destroying geometry (bond RMSZ 28.0 where every other target sat near 1.3). See ``_specs.py``'s module docstring. - This class survives for exactly one caller: - :class:`torchref.experimental.alignment.rigid_body.RigidBodyRefinement`, the FFT-direct - rigid-body aligner in the molecular-replacement pipeline, which constructs it directly - rather than through the factory. - - **Why it was not simply repointed at ``nll``** during the 2026-08 target refactor, as - originally planned: that aligner has **no test coverage whatsoever** - (``tests/integration/test_rigid_body_refinement.py`` exercises the *other* rigid-body - module, ``refinement/rigid_body_refinement.py``). Swapping a Rice likelihood for a - Gaussian there would be an untested numerical change in a live MR path, so the - likelihood was kept and only its *implementation* was de-duplicated -- the body now - calls the shared :func:`~torchref.base.targets.xray_likelihoods.rice_math` primitive - instead of a second copy of the Rice in the deleted ``xray_ml`` module. - - Whoever gives that aligner a test should revisit this: ``nll`` or ``ml_noalpha`` is - almost certainly the better objective, and then this class can go. + **It now has no caller at all, and is a deletion candidate.** It survived the + 2026-08 target refactor for exactly one: the FFT-direct rigid-body aligner in the + molecular-replacement pipeline, which constructed it directly rather than through + the factory. That aligner had no test coverage + (``tests/integration/test_rigid_body_refinement.py`` exercises the *other* + rigid-body module, ``refinement/rigid_body_refinement.py``), so repointing it at + ``nll`` would have been an untested numerical change in a live path -- the + likelihood was kept and only its *implementation* de-duplicated onto the shared + :func:`~torchref.base.targets.xray_likelihoods.rice_math` primitive. + + The MR pipeline no longer polishes placements, so that aligner is gone and the + constraint with it. What remains is this class, three unit tests of it, and an + export. Removing all of that is a ``refinement/`` change and belongs in a + ``refinement/`` commit, not an alignment one. """ #: ``epsilon * beta`` was clamped here in the implementation this replaced. Preserved diff --git a/torchref/refinement/targets/xray/sigma_a.py b/torchref/refinement/targets/xray/sigma_a.py index 1934b393..b9e4b0a1 100644 --- a/torchref/refinement/targets/xray/sigma_a.py +++ b/torchref/refinement/targets/xray/sigma_a.py @@ -175,10 +175,6 @@ def _loss_inputs( eps_full, dss_full = self._geom() eps_full = eps_full.to(F_calc_full.dtype) dss_full = dss_full.to(F_calc_full.dtype) - # ONE data path. `sub.F` / `sub.sigF` go through `_corrected_or_raw()`, which - # silently falls back to RAW amplitudes when the scaler has not run, while the - # estimator below is fed `get_corrected_data()`, which raises instead. Mixing the two - # can put raw amplitudes and a scaled-data variance in the same loss. F_obs_full, sigma_full = self._data.get_corrected_data() F_obs_full = F_obs_full.to(F_calc_full.dtype).reshape(-1) centric_full = self._data.centric diff --git a/torchref/restraints/__init__.py b/torchref/restraints/__init__.py deleted file mode 100644 index c9a1c1c7..00000000 --- a/torchref/restraints/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Geometry restraints (bonds, angles, torsions, planes, chirals, VDW contacts). - -``Restraints`` (an alias of ``RestraintsNew``) builds and holds them all from CIF -dictionaries resolved through :func:`get_library_manager`; ``MONOMER_LIB_PATH`` -resolves lazily to that manager's ``monomer_dir``. Ideal values come from the CCP4 -Monomer Library (Long et al. 2017, Acta Cryst. D73, 112-122). - -The builder classes and topology helpers (``build_all_restraints``, -``HydrogenTopology``, ``build_hydrogen_topology``, the intra-/inter-residue -builders) are used across the package but deliberately *not* re-exported here -- -import them from their defining submodules. -""" - -from torchref.restraints.library import get_library_manager -from torchref.restraints.restraints import RestraintsNew as Restraints - - -def __getattr__(name): - """Lazy access to MONOMER_LIB_PATH for backward compatibility.""" - if name == "MONOMER_LIB_PATH": - return get_library_manager().monomer_dir - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = [ - "Restraints", - "MONOMER_LIB_PATH", - "get_library_manager", -] diff --git a/torchref/scaling/__init__.py b/torchref/scaling/__init__.py index 27d4f8a2..1d42dffb 100644 --- a/torchref/scaling/__init__.py +++ b/torchref/scaling/__init__.py @@ -1,21 +1,26 @@ -"""Scaling calculated structure factors onto observed data. +"""Scale observed datasets and calculated structure factors. Per-bin overall scale, anisotropic correction and bulk-solvent contribution. :class:`ScalerBase` is model-independent -- every method that needs ``F_calc`` takes it as an argument; :class:`Scaler` holds a :class:`~torchref.model.Model` and computes ``F_calc`` itself; :class:`CollectionScaler` fits one shared set of scales jointly across a dataset/model collection. :class:`SolventModel` supplies -the flat bulk-solvent term (k_sol, B_sol). +the flat bulk-solvent term (k_sol, B_sol). ``DatasetScaler`` independently fits +relative observed-data corrections; ``WilsonNormaliser`` supplies E values. """ +from torchref.scaling.collection_scaler import CollectionScaler +from torchref.scaling.dataset_scaler import DatasetScaler from torchref.scaling.scaler import Scaler from torchref.scaling.scaler_base import ScalerBase from torchref.scaling.solvent import SolventModel -from torchref.scaling.collection_scaler import CollectionScaler +from torchref.scaling.wilson import WilsonNormaliser __all__ = [ "Scaler", + "DatasetScaler", "ScalerBase", "SolventModel", "CollectionScaler", + "WilsonNormaliser", ] diff --git a/torchref/scaling/basis.py b/torchref/scaling/basis.py new file mode 100644 index 00000000..50f39f77 --- /dev/null +++ b/torchref/scaling/basis.py @@ -0,0 +1,74 @@ +"""Chebyshev basis in resolution, shared by everything that fits a smooth curve in |s|. + +Two things in here carry argument rather than convention, and both were settled +by the scaler rework: + +* **The abscissa is ``sin(theta)/lambda``, not ``s**2``.** The modulation a + resolution-dependent scale has to represent is gentle through the bulk of the + range and has real structure in the first few percent of ``s**2``; a basis + uniform in ``s**2`` spends nearly all its resolution where nothing happens. +* **The basis is prefix-nested.** ``chebyshev_design(x, k)`` equals + ``chebyshev_design(x, n)[:, :k]`` for ``k <= n``, so raising the order adds + detail without redefining the terms already fitted, and a caller can slice + instead of rebuilding. + +``lo``/``hi`` exist for **extrapolation**, and the reason is the clamp rather +than the mapping. An affine remap does not change the space a polynomial basis +spans, so two fits over different ranges recover the same *function* where their +data overlap -- only the coefficients and the conditioning differ. What does +differ is outside the fitted range: ``u`` saturates at the ends, so every column +goes constant and the curve is frozen at its endpoint value. + +So a fit over one resolution range, evaluated somewhere else, silently returns a +flat extrapolation. That is the case a shared ``lo``/``hi`` is for -- fitting on +one reflection set and using the curve on another, which is what comparing two +fits, or fitting on a crystal lattice and evaluating on a dense sampling, +actually requires. ``ScalerBase`` does not need it: it builds one design over +all reflections and slices rows. +""" + +from __future__ import annotations + +from typing import Optional, Union + +import torch + +__all__ = ["chebyshev_design"] + + +def chebyshev_design( + x: torch.Tensor, + n_coeff: int, + lo: Optional[Union[float, torch.Tensor]] = None, + hi: Optional[Union[float, torch.Tensor]] = None, +) -> torch.Tensor: + """``(N, n_coeff)`` Chebyshev design matrix in ``x``. + + Parameters + ---------- + x : torch.Tensor + ``(N,)`` abscissa, normally ``sin(theta)/lambda``. + n_coeff : int + Number of Chebyshev terms. ``1`` gives a single constant column, i.e. a + global scale with no resolution dependence. + lo, hi : float or torch.Tensor, optional + Range to map onto ``[-1, 1]``. Both default to ``x``'s own extremes, + which is right for a single dataset and wrong the moment two fits have + to be compared -- see the module docstring. + + Returns + ------- + torch.Tensor + ``(N, n_coeff)``, column 0 all ones, every entry in ``[-1, 1]``. + """ + if n_coeff < 1: + raise ValueError(f"n_coeff must be at least 1, got {n_coeff}") + lo = x.min() if lo is None else torch.as_tensor(lo, dtype=x.dtype, device=x.device) + hi = x.max() if hi is None else torch.as_tensor(hi, dtype=x.dtype, device=x.device) + u = (2 * (x - lo) / (hi - lo).clamp(min=1e-12) - 1).clamp(-1.0, 1.0) + cols = [torch.ones_like(u), u] + for _ in range(2, n_coeff): + cols.append(2 * u * cols[-1] - cols[-2]) # Chebyshev recurrence + # The slice is what makes ``n_coeff == 1`` work: the loop does not run and + # the pre-seeded linear column is dropped. + return torch.stack(cols[:n_coeff], dim=1) diff --git a/torchref/scaling/collection_scaler.py b/torchref/scaling/collection_scaler.py index 585b4cb4..2b7d918d 100644 --- a/torchref/scaling/collection_scaler.py +++ b/torchref/scaling/collection_scaler.py @@ -8,16 +8,18 @@ combination at the same population fractions as the structural models. """ -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Dict import torch import torch.nn as nn from torchref.base.metrics.rfactor import rfactor_work_free -from torchref.base.reciprocal import get_scattering_vectors -from torchref.base.targets.xray_likelihoods import complex_var_from_beta, rice_math from torchref.config import get_float_dtype -from torchref.scaling.scaler_base import ScalerBase +from torchref.scaling.scaler_base import ( + DEFAULT_SCALE_TARGET, + SCALE_TARGETS, + ScalerBase, +) from torchref.scaling.solvent import SS_HALF_BOUNDS, SolventModel from torchref.utils.utils import ModuleReference @@ -26,6 +28,44 @@ from torchref.model.model_collection import ModelCollection +class _DatasetScalerView(nn.Module): + """One dataset's view of a shared :class:`CollectionScaler`. + + Exists so the scale fit can hand a plain scaler to a taxonomy row. A row calls + ``self._scaler(fcalc)`` and knows nothing else about scaling, but the collection's + bulk solvent is a fraction-weighted mixture that depends on *which* dataset is being + scaled -- information :meth:`ScalerBase.forward` has no way to carry. This forwards to + :meth:`CollectionScaler.forward_mixed` with the fractions bound. + + The parent is held through :class:`~torchref.utils.utils.ModuleReference`, so its + parameters are **not** re-registered here: the optimiser is still built from the one + scaler, and a row holding this view contributes no leaves of its own. + """ + + def __init__(self, parent: "CollectionScaler", fractions: torch.Tensor): + super().__init__() + self._parent = ModuleReference(parent) + self.register_buffer("_fractions", fractions.detach().clone()) + + def __getattr__(self, name): + """Anything this view does not own belongs to the parent. + + ``nn.Module.__getattr__`` resolves parameters, buffers and submodules first; a + row reading some other scaler attribute (bin edges, resolution limits) should see + the parent's, not an AttributeError. + """ + try: + return super().__getattr__(name) + except AttributeError: + parent = self.__dict__.get("_parent") + if parent is None: + raise + return getattr(parent, name) + + def forward(self, fcalc: torch.Tensor) -> torch.Tensor: + return self._parent.forward_mixed(fcalc, self._fractions) + + class CollectionScaler(ScalerBase): """ Joint scaler for DatasetCollection + ModelCollection. @@ -152,7 +192,7 @@ def _calc_initial_scale_joint(self): pos_mask = torch.ones_like(fobs, dtype=torch.bool) mask = (work_mask & pos_mask).to(torch.bool) - bins = self.bins[mask].to(torch.int64) + bins = self.bins[mask].to(torch.int64) # dtype-ok: bin indices for scatter/index_select; PyTorch requires int64 log_ratios = ( torch.log(fobs_clamped[mask]) - torch.log(fcalc_amp[mask]) ).to(self.device) @@ -165,7 +205,7 @@ def _calc_initial_scale_joint(self): per_bin = scales / (counts + 1e-6) with torch.no_grad(): - target = per_bin.detach()[self.bins.to(torch.int64)] + target = per_bin.detach()[self.bins.to(torch.int64)] # dtype-ok: bin indices for advanced indexing; PyTorch requires int64 design = self._iso_design.to(target.dtype) coeff = torch.linalg.lstsq(design, target.unsqueeze(1)).solution.squeeze(1) self.c_iso = nn.Parameter(coeff.detach()) @@ -283,6 +323,61 @@ def forward_mixed( f_sol_raw_mixed = self.get_mixed_solvent_raw(fractions) return super().forward(fcalc, f_sol_override=f_sol_raw_mixed) + def compute_component_solvent_raw(self) -> torch.Tensor: + """Raw (un-damped) complex solvent SFs for every component, stacked. + + The solvent counterpart of + :meth:`~torchref.model.model_collection.ModelCollection.compute_component_fcalcs`. + Each component's mask FFT is cached, so repeated calls are cheap. + + Returns + ------- + torch.Tensor + Complex tensor of shape ``(n_components, n_reflections)``. + """ + return torch.stack( + [ + self._get_component_f_sol_raw(i) + for i in range(len(self._component_solvent_models)) + ], + dim=0, + ) + + def forward_batched( + self, + fcalc_batch: torch.Tensor, + fractions_matrix: torch.Tensor, + ) -> torch.Tensor: + """Scale a batch of mixtures, each with its own fraction-weighted solvent. + + The batched form of :meth:`forward_mixed`: one shared set of scale parameters + applied to ``T`` mixtures at once, with the bulk solvent mixed per row by the + same weights. Since ``ScalerBase.forward`` is affine in ``fcalc`` and the mixed + solvent is linear in the weights, passing a *derivative* of the fractions in + place of the fractions returns the corresponding derivative of the scaled + structure factors. + + Parameters + ---------- + fcalc_batch : torch.Tensor + Complex structure factors of shape ``(T, n_reflections)``. + fractions_matrix : torch.Tensor + Weights of shape ``(T, n_components)``, one row per member of the batch. + + Returns + ------- + torch.Tensor + Scaled complex structure factors of shape ``(T, n_reflections)``. + """ + component_sol_raw = self.compute_component_solvent_raw() + # Keep real fraction multiplication and component accumulation identical + # to forward_mixed; complex GEMM changes rounding near solvent cancellation. + f_sol_batch = sum( + fractions_matrix[:, i, None] * component_sol_raw[i] + for i in range(component_sol_raw.shape[0]) + ) + return super().forward(fcalc_batch, f_sol_override=f_sol_batch) + # ------------------------------------------------------------------ # Joint LBFGS refinement # ------------------------------------------------------------------ @@ -294,12 +389,29 @@ def refine_lbfgs_joint( max_iter: int = 200, history_size: int = 10, verbose: bool = True, + scale_target: str = DEFAULT_SCALE_TARGET, ) -> dict: """ - Refine scale parameters using LBFGS against **all** datasets. - - The closure sums the NLL across every matched dataset–model pair, - so a single set of scale parameters is fitted jointly. + Refine the shared scale parameters against **all** datasets jointly. + + One set of scale parameters serves every matched dataset-model pair, so the + closure sums a per-dataset objective. Each dataset's term is built from a row of + :data:`~torchref.refinement.targets.xray._specs.XRAY_TARGETS`, exactly as + :meth:`ScalerBase.refine_lbfgs` builds its single-dataset one -- so both scale + fits evaluate the same likelihood code, and neither carries a private copy of it. + + The row sees this dataset's own **mixed** bulk solvent, via a + :class:`_DatasetScalerView` that shares the parent's parameters and applies + :meth:`forward_mixed`. That is why the scaler cannot simply be handed to the + target: the solvent depends on which dataset's fractions are in play, and the + plain :meth:`ScalerBase.forward` has no way to know. + + Amplitudes throughout, whatever observable the *refinement* target fits. + Unit-weight least squares on intensities would put leverage where the data is + strongest: the residual goes as ``2 F dF``, so the squared residual carries an + extra factor of ``F**2`` and a global scale plus B plus anisotropy would be + determined almost entirely by the strongest low-resolution reflections, leaving + high resolution unconstrained. Parameters ---------- @@ -313,117 +425,119 @@ def refine_lbfgs_joint( LBFGS history size. verbose : bool Print progress. + scale_target : str, optional + Objective, one of :data:`~torchref.scaling.scaler_base.SCALE_TARGETS`. + Defaults to :data:`~torchref.scaling.scaler_base.DEFAULT_SCALE_TARGET` + (unit-weight ``ls``) -- the same default and the same reason as the + single-dataset fit. Returns ------- dict Refinement metrics (steps, rwork, rfree of dark dataset). + + Raises + ------ + ValueError + If ``scale_target`` is not in + :data:`~torchref.scaling.scaler_base.SCALE_TARGETS`. """ + if scale_target not in SCALE_TARGETS: + raise ValueError( + f"scale_target must be one of {SCALE_TARGETS}, got {scale_target!r}" + ) # Local import, deliberately: `torchref.refinement` imports # `torchref.scaling` at module scope, so hoisting these to module level # closes an import cycle. Do not "tidy" them up. - from torchref.refinement.model_error_estimation.sigma_a import SigmaAEstimator, epsilon_from_hkl from torchref.refinement.loss_state import LossState + from torchref.refinement.targets.xray import create_xray_target dc = self._dataset_collection mc = self._model_collection all_keys = [mc.dark_key] + mc.timepoint_names - # Pre-compute all fcalc (detached) plus the per-dataset σ_A model-error - # variance (beta/epsilon). beta is estimated ONCE on each dataset's free - # set from the currently-scaled |F_calc| and held detached during the fit, - # as in the single-dataset ``ScalerBase.refine_lbfgs``; that variance is - # what stops the scale collapsing toward zero in weak shells. + # Per dataset: a detached F_calc, and a table row that scales it through this + # dataset's solvent view. `model=None`, so the row never recomputes structure + # factors and the only leaves in the graph are this scaler's own parameters. + # Rows that need a model-error estimate build and own one themselves, as under + # the body refinement -- there is nothing to precompute here. fcalc_cache = {} fractions_cache = {} - beta_cache = {} - eps_cache = {} - work_cache = {} - centric_cache = {} + terms = [] for name in all_keys: if name not in dc: continue data = dc[name] model = mc[name] - hkl = data.hkl - fobs, sigma = data.get_corrected_data() with torch.no_grad(): - fc = model(hkl).detach() - fracs = model.fractions.detach() - f_sol_raw = self.get_mixed_solvent_raw(fracs) - scaled0 = super(CollectionScaler, self).forward( - fc, f_sol_override=f_sol_raw + fcalc_cache[name] = model(data.hkl).detach() + fractions_cache[name] = model.fractions.detach() + view = _DatasetScalerView(self, fractions_cache[name]) + terms.append( + ( + create_xray_target( + data=data, + model=None, + scaler=view, + mode=scale_target, + use_set="work", + verbose=0, + device=self.device, + ), + fcalc_cache[name], ) - fc_amp0 = torch.abs(scaled0).reshape(-1) - fobs0 = fobs.to(fc_amp0.dtype).reshape(-1) - eps0 = epsilon_from_hkl( - hkl, getattr(data, "spacegroup", None) - ).to(fc_amp0.dtype) - s = get_scattering_vectors(hkl, data.cell) - dss0 = (torch.norm(s, dim=1) ** 2).to(fc_amp0.dtype) - # sigma_obs must be passed, as at every other call site: it is what - # makes sigma_A the correlation with the noise-free amplitudes. - _est = SigmaAEstimator().get( - fobs0, fc_amp0, data.centric, eps0, dss0, data.free.mask, - sigma_obs=sigma.to(fc_amp0.dtype).reshape(-1), - ) - # TOTAL variance (`beta`, not `beta_model`): this scale fit uses the same - # likelihood as `ml`, which does not account for sigma_obs separately. - beta0, eps0 = _est.beta, _est.epsilon - fcalc_cache[name] = fc - fractions_cache[name] = fracs - beta_cache[name] = beta0 - eps_cache[name] = eps0 - work_cache[name] = data.work - centric_cache[name] = data.centric - - # Wrap the joint σ_A ML loss + U-penalty as a LossState target, reusing its - # NaN/Inf rejection. fcalc is detached, so the only leaves in the graph are - # the scaler's own parameters. + ) + + # Use one fixed normalizer: LBFGS uses absolute tolerances, and a large + # float32 loss can round away the decrease sought by the line search. + with torch.no_grad(): + ssq = sum( + float(dc[n].work.F.detach().pow(2).sum()) + for n in fcalc_cache + ) + _norm = 1.0 / max(ssq, 1e-30) scaler_self = self class _CollectionScalerJointTarget(nn.Module): + """The table rows, closed over their detached ``fcalc``.""" + name = "scaler/joint" def forward(self): - total = torch.tensor(0.0, device=scaler_self.device) - n = 0 - for nm in all_keys: - if nm not in fcalc_cache: - continue - fc = fcalc_cache[nm] - fracs = fractions_cache[nm] - f_sol_raw = scaler_self.get_mixed_solvent_raw(fracs) - scaled = super(CollectionScaler, scaler_self).forward( - fc, f_sol_override=f_sol_raw - ) - # σ_A (Read MLF) scale-fit on the WORK set, with detached - # free-set beta/epsilon — same likelihood the body - # refinement uses. - amp = torch.abs(scaled).reshape(-1) - work = work_cache[nm] - F_obs = work.F.to(amp.dtype) - Fc = work.select(amp) - beta_w = work.select(beta_cache[nm]).to(F_obs.dtype) - eps_w = ( - work.select(eps_cache[nm]).to(F_obs.dtype) - if eps_cache[nm] is not None - else None - ) - centric_w = work.select(centric_cache[nm]) - loss = rice_math( - F_obs, Fc, complex_var_from_beta(beta_w, eps_w), centric_w - ) + total = torch.zeros((), device=scaler_self.device) + for target, fc in terms: + loss = target(fcalc=fc) + # Skip a dataset whose term went non-finite rather than poisoning + # the whole joint gradient with it. if torch.isfinite(loss): total = total + loss - n += 1 - if n > 0: - total = total / n - u_penalty = torch.sum(scaler_self.U**2) - return total + u_penalty + return total * _norm + + def maintenance(self): + """Forward the hook so sigma_A rows drop their ``beta`` cache after a + step block, as they do under the body refinement.""" + for target, _ in terms: + maint = getattr(target, "maintenance", None) + if maint is not None: + maint() + + class _CollectionScalerUPenalty(nn.Module): + """``sum(U**2)`` on the anisotropic scale tensor. + + Its normaliser is pinned to **amplitudes** rather than following the + objective. Sharing ``_norm`` would make the penalty's weight relative to the + likelihood depend on which objective was selected, which is a silent change + of regularisation strength dressed up as a change of objective. + """ + + name = "scaler/u_penalty" + + def forward(self): + return torch.sum(scaler_self.U**2) * _norm state = LossState(device=self.device) state.register_target("scaler/joint", _CollectionScalerJointTarget()) + state.register_target("scaler/u_penalty", _CollectionScalerUPenalty()) optimizer = torch.optim.LBFGS( self.parameters(), diff --git a/torchref/scaling/dataset_scaler.py b/torchref/scaling/dataset_scaler.py new file mode 100644 index 00000000..3d220270 --- /dev/null +++ b/torchref/scaling/dataset_scaler.py @@ -0,0 +1,305 @@ +"""Joint relative scaling of observed datasets without a privileged reference. + +DatasetScaler owns all fitted corrections. ScaledDataset exposes one correction +through the reflection-data interface; no optimization state lives on raw datasets. +""" + +from collections.abc import Mapping +from math import isfinite +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from torchref.io import ReflectionData + +import torch +from torch import nn + +from torchref.base.targets.xray_likelihoods import SIGMA_FLOOR_ABS, SIGMA_FLOOR_FRAC +from torchref.config import get_float_dtype, get_int_dtype, normalize_device +from torchref.utils.device_mixin import DeviceMixin + + +def _identity_hkl(data): + """Keep Bijvoet observations separate while matching dataset identities.""" + return data.hkl if data.friedel_merged else data._hkl_for_sf() + + +class DatasetScaler(DeviceMixin, nn.Module): + """Fit N observed datasets to a shared sigma-weighted amplitude consensus. + + Parameters + ---------- + datasets : Mapping[str, ReflectionData] + At least two raw datasets in the same space-group setting and Friedel + convention. Sources are copied without mutation; membership is fixed. + device : torch.device or str, optional + Computation device, defaulting to the first dataset's device. Prepared + fitting arrays and owned copies move here without moving source data. + + Notes + ----- + Corrections have zero mean in log space over datasets, including anisotropy. + The six quadratic coefficients use dimensionless, normalized Miller indices; + they are not Cartesian atomic displacement parameters. Every corrected read + depends on all parameter rows through centering. ``fit`` freezes parameters + when it finishes; call ``requires_grad_(True)`` for custom differentiable use. + """ + + def __init__( + self, + datasets: Mapping[str, "ReflectionData"], + device: torch.device | str | None = None, + ) -> None: + super().__init__() + if len(datasets) < 2: + raise ValueError("Dataset scaling requires at least two datasets") + self.keys = tuple(datasets) + self.datasets = {} + for key, data in datasets.items(): + raw = data.raw_data() if hasattr(data, "raw_data") else data + owned = raw.__select__(torch.arange(len(raw), device=raw.device)) + owned.source = None + owned.spacegroup = raw.spacegroup.copy() + self.datasets[key] = owned + + self.device = normalize_device( + device if device is not None else next(iter(datasets.values())).device + ) + self.dtype_float = get_float_dtype() + self.raw_parameters = nn.Parameter( + torch.zeros((len(self.keys), 7), device=self.device, dtype=self.dtype_float) + ) + for name in ("hkl", "amplitudes", "sigmas", "fit_mask", "hkl_scale"): + self.register_buffer(name, None) + self.n_contrasts = 0 + self._initialized = False + self.to(self.device) + self.prepare() + + @property + def corrections(self) -> torch.Tensor: + """Centered coefficients, shape (N, 7), in dimensionless log units.""" + return self.raw_parameters - self.raw_parameters.mean(dim=0, keepdim=True) + + def design(self, hkl: torch.Tensor) -> torch.Tensor: + """Return the log-correction basis, shape (H, 7), for integer HKL (H, 3).""" + q = hkl.to(device=self.device, dtype=self.raw_parameters.dtype) / self.hkl_scale + h, k, l = q.unbind(dim=-1) + return torch.stack( + (torch.ones_like(h), h * h, k * k, l * l, 2 * h * k, 2 * h * l, 2 * k * l), + dim=-1, + ) + + def log_corrections(self, hkl: torch.Tensor) -> torch.Tensor: + """Return dimensionless log amplitude corrections, shape (N, H).""" + return self.corrections @ self.design(hkl).T + + def forward(self, key: str, hkl: torch.Tensor) -> torch.Tensor: + """Return positive amplitude factors (H,) for dataset key and HKL (H, 3).""" + row = self.keys.index(key) + return (self.design(hkl) @ self.corrections[row]).exp() + + def prepare(self) -> None: + """Prepare training arrays and reject disconnected or unidentified fits. + + This reads current source masks and measurements. No free/validation + observation enters initialization, uncertainty floors or the objective. + Changes to reflection sets or fit masks require a new scaler once fitted. + """ + data = list(self.datasets.values()) + if tuple(self.datasets) != self.keys: + raise ValueError("Dataset membership changed; construct a new scaler") + symmetry = {d.spacegroup.xhm for d in data} + if len(symmetry) != 1 or len({d.friedel_merged for d in data}) != 1: + raise ValueError( + "Datasets require compatible symmetry settings and Friedel conventions" + ) + hkls = [_identity_hkl(d).to(self.device) for d in data] + hkl, inverse = torch.unique(torch.cat(hkls), dim=0, return_inverse=True) + shape = (len(data), len(hkl)) + amplitudes = torch.zeros(shape, device=self.device, dtype=self.dtype_float) + sigmas = torch.ones_like(amplitudes) + valid = torch.zeros(shape, device=self.device, dtype=torch.bool) + held_out = torch.zeros(len(hkl), device=self.device, dtype=torch.bool) + start = 0 + for row, ds in enumerate(data): + if ds.F_raw is None or ds.F_sigma_raw is None: + raise ValueError(f"Dataset {self.keys[row]!r} requires F and SIGF") + idx = inverse[start : start + len(ds)] + start += len(ds) + if len(torch.unique(idx)) != len(idx): + raise ValueError( + f"Dataset {self.keys[row]!r} contains duplicate reflection identities" + ) + f = ds.F_raw.detach().to(amplitudes) + sigma = ds.F_sigma_raw.detach().to(sigmas) + present = ds.masks().to(device=self.device, dtype=torch.bool) + usable = present & torch.isfinite(f) & torch.isfinite(sigma) & (sigma > 0) + work = ds.work.mask.to(self.device) + held_out[idx] |= present & ~work + valid[row, idx] = usable + amplitudes[row, idx] = torch.where(usable, f, torch.zeros_like(f)) + sigmas[row, idx] = torch.where(usable, sigma, torch.ones_like(sigma)) + mask = valid & ~held_out.unsqueeze(0) + mask &= mask.sum(dim=0, keepdim=True) >= 2 + reached = {0} + for _ in data: + reached |= { + j + for i in tuple(reached) + for j in range(len(data)) + if bool((mask[i] & mask[j]).any()) + } + if len(reached) != len(data): + raise ValueError( + "Dataset work-set overlap is disconnected; relative scales are unidentified" + ) + for row in range(len(data)): + floor = (sigmas[row, mask[row]].median() * SIGMA_FLOOR_FRAC).clamp_min( + SIGMA_FLOOR_ABS + ) + sigmas[row] = sigmas[row].clamp_min(floor) + if self.hkl_scale is None: + self.hkl_scale = ( + hkl[mask.any(dim=0)].to(amplitudes).abs().amax(dim=0).clamp_min(1) + ) + self.hkl, self.amplitudes, self.sigmas, self.fit_mask = ( + hkl, + amplitudes, + sigmas, + mask, + ) + self.n_contrasts = int((mask.sum(dim=0) - 1).clamp_min(0).sum()) + # Rank concerns only geometry and overlap, not observed amplitudes. The + # small-column check runs on CPU because MPS has no SVD implementation. + design = self.design(hkl).detach().cpu() + mask_cpu = mask.cpu() + first = mask_cpu.to(get_int_dtype()).argmax(dim=0) + blocks = [] + for row in range(len(data)): + selected = mask_cpu[row] & (first != row) + if not bool(selected.any()): + continue + block = torch.zeros((int(selected.sum()), len(data), 7), dtype=design.dtype) + block[:, row] = design[selected] + block[torch.arange(len(block)), first[selected]] = -design[selected] + blocks.append(block[:, :-1].reshape(len(block), -1)) + rank = int(torch.linalg.matrix_rank(torch.cat(blocks))) if blocks else 0 + if rank != 7 * (len(data) - 1): + raise ValueError( + "Overlapping reflections do not identify overall scale and six anisotropic coefficients" + ) + + def initialize(self) -> None: + """Seed overall log scales from robust pairwise log-amplitude ratios.""" + rows, values = [], [] + n = len(self.keys) + for i in range(n): + for j in range(i): + mask = ( + self.fit_mask[i] + & self.fit_mask[j] + & (self.amplitudes[i] > 0) + & (self.amplitudes[j] > 0) + ) + if not bool(mask.any()): + continue + row = torch.zeros(n, dtype=self.dtype_float) + row[i], row[j] = 1, -1 + rows.append(row) + values.append( + (self.amplitudes[j, mask].log() - self.amplitudes[i, mask].log()) + .median() + .cpu() + ) + rows.append(torch.ones(n, dtype=self.dtype_float)) + values.append(torch.zeros((), dtype=self.dtype_float)) + initial = torch.linalg.lstsq(torch.stack(rows), torch.stack(values)).solution + with torch.no_grad(): + self.raw_parameters.zero_() + self.raw_parameters[:, 0].copy_(initial.to(self.raw_parameters)) + self._initialized = True + + def fit(self, nsteps: int = 10, max_iter: int = 100) -> dict: + """Fit joint corrections with L-BFGS and freeze the resulting parameters. + + Parameters + ---------- + nsteps : int + Number of outer L-BFGS steps. + max_iter : int + Maximum iterations per outer step. + + Returns + ------- + dict + Initial/final normalized loss, contrast count and centered coefficients. + """ + from torchref.refinement.loss_state import LossState + from torchref.refinement.targets.dataset_scaling import DatasetScalingTarget + + if nsteps < 1 or max_iter < 1: + raise ValueError("nsteps and max_iter must be positive") + self.prepare() + if not self._initialized: + self.initialize() + saved = self.raw_parameters.detach().clone() + self.requires_grad_(True) + target = DatasetScalingTarget(self) + state = LossState(device=self.device) + state.register_target("scaling/datasets", target) + before = float(target().detach()) + optimizer = torch.optim.LBFGS( + self.parameters(), max_iter=max_iter, line_search_fn="strong_wolfe" + ) + try: + state.run(optimizer, nsteps=nsteps, log=False, context="dataset_scaler.fit") + after = float(target().detach()) + if not torch.isfinite(self.raw_parameters).all() or not isfinite(after): + raise RuntimeError( + "Dataset scaling produced non-finite parameters or loss" + ) + except Exception: + with torch.no_grad(): + self.raw_parameters.copy_(saved) + raise + finally: + self.requires_grad_(False) + return { + "loss_before": before, + "loss_after": after, + "n_contrasts": self.n_contrasts, + "corrections": dict( + zip(self.keys, self.corrections.detach().cpu().tolist()) + ), + } + + def get_state(self) -> dict: + """Return raw source states and the shared fitted parameter state.""" + return { + "datasets": {k: d._get_state() for k, d in self.datasets.items()}, + "parameters": self.raw_parameters.detach().cpu(), + "hkl_scale": self.hkl_scale.detach().cpu(), + "initialized": self._initialized, + } + + @classmethod + def from_state( + cls, state: dict, device: torch.device | str | None = None + ) -> "DatasetScaler": + """Restore a shared scaler and its raw sources on the requested device.""" + from torchref.io.datasets.reflection_data import ReflectionData + + obj = cls( + { + k: ReflectionData._from_state(dict(v), device) + for k, v in state["datasets"].items() + }, + device=device, + ) + with torch.no_grad(): + obj.raw_parameters.copy_(state["parameters"].to(obj.raw_parameters)) + obj.hkl_scale.copy_(state["hkl_scale"].to(obj.hkl_scale)) + obj._initialized = state["initialized"] + obj.requires_grad_(False) + return obj diff --git a/torchref/scaling/scaler_base.py b/torchref/scaling/scaler_base.py index 9fd9119f..b8350e48 100644 --- a/torchref/scaling/scaler_base.py +++ b/torchref/scaling/scaler_base.py @@ -12,6 +12,7 @@ import torch.nn as nn from torchref.base.math_torch import U_to_matrix +from torchref.scaling.basis import chebyshev_design from torchref.base.metrics import ( binwise_scale, nll_xray, @@ -146,19 +147,16 @@ def __init__( def _build_iso_design(self) -> torch.Tensor: """``(N, n_iso_coeff)`` Chebyshev design matrix for the isotropic scale. - The abscissa is ``sqrt(s_half_sq)``, i.e. ``sin(theta)/lambda``, mapped onto - ``[-1, 1]``. That coordinate rather than ``s**2`` because the modulation is gentle - through the bulk of the resolution range but has real structure in the first few - percent of ``s**2``; a basis uniform in ``s**2`` spends nearly all its resolution - where nothing happens. + The abscissa is ``sqrt(s_half_sq)``, i.e. ``sin(theta)/lambda``; see + :func:`torchref.scaling.basis.chebyshev_design` for why that coordinate. + + No explicit range: this design is built once over all reflections and + then *sliced* wherever a subset is needed (``forward`` does exactly + that), so the mapping is the same everywhere it is used. """ - x = torch.sqrt(self._s_half_sq.clamp(min=0)) - lo, hi = x.min(), x.max() - u = (2 * (x - lo) / (hi - lo).clamp(min=1e-12) - 1).clamp(-1.0, 1.0) - cols = [torch.ones_like(u), u] - for _ in range(2, self.n_iso_coeff): - cols.append(2 * u * cols[-1] - cols[-2]) # Chebyshev recurrence - return torch.stack(cols[: self.n_iso_coeff], dim=1) + return chebyshev_design( + torch.sqrt(self._s_half_sq.clamp(min=0)), self.n_iso_coeff, + ) def iso_log_scale(self, design: Optional[torch.Tensor] = None) -> torch.Tensor: """Per-reflection isotropic log scale ``design @ c_iso``, clamped to ``[-10, 10]``. @@ -266,7 +264,7 @@ def calc_initial_scale(self, fcalc: torch.Tensor): initial_log_scale.detach().cpu().numpy(), ) with torch.no_grad(): - target = initial_log_scale.detach().to(self.device)[self.bins.to(torch.int64)] + target = initial_log_scale.detach().to(self.device)[self.bins.to(torch.int64)] # dtype-ok: bin indices for advanced indexing; PyTorch requires int64 design = self._iso_design.to(target.dtype) coeff = torch.linalg.lstsq(design, target.unsqueeze(1)).solution.squeeze(1) self.c_iso = nn.Parameter(coeff.detach().to(self.device)) @@ -348,6 +346,34 @@ def get_scale(self) -> float: return torch.exp(self.iso_log_scale().mean()).item() return 1.0 + def multiplicative_scale(self) -> torch.Tensor: + """Per-reflection factor taking model amplitudes to the observed scale. + + ``K_overall * b_overall * anisotropy``: every multiplicative component + :meth:`forward` applies and none of the additive bulk-solvent term, so dividing + observed amplitudes by it returns them to the model's absolute scale, electrons. + Components not yet set up contribute ones. + + Returns + ------- + torch.Tensor + Shape ``(N,)`` over the scaler's full reflection list, detached, on + ``self.device`` in the scale parameters' dtype. + """ + c_iso = getattr(self, "c_iso", None) + dtype = c_iso.dtype if c_iso is not None else get_float_dtype() + factor = torch.ones(int(self.bins.numel()), device=self.device, dtype=dtype) + with torch.no_grad(): + if hasattr(self, "U"): + factor = factor * self.anisotropy_correction().to(factor) + if c_iso is not None: + factor = factor * torch.exp(self.iso_log_scale(self._iso_design)).to( + factor + ) + if getattr(self, "bin_wise_bfactor", None) is not None: + factor = factor * self.bin_wise_bfactor_correction().to(factor) + return factor.detach() + def setup_bin_wise_bfactor(self): """Initialize bin-wise B-factor correction parameters.""" self.bin_wise_bfactor = nn.Parameter( @@ -396,7 +422,7 @@ def get_binwise_mean_intensity(self, fcalc: torch.Tensor): mean_calc_intensity = torch.zeros(self.nbins, device=self.device, dtype=fobs.dtype) counts = torch.zeros(self.nbins, device=self.device, dtype=fobs.dtype) counts_vals = torch.ones_like(F_calc, device=self.device, dtype=fobs.dtype) - bins_sel = self.bins.to(torch.int64)[sel] + bins_sel = self.bins.to(torch.int64)[sel] # dtype-ok: bin indices for advanced indexing; PyTorch requires int64 mean_obs_intensity = torch.scatter_add( mean_obs_intensity, 0, bins_sel, intensities[sel] ) @@ -782,9 +808,10 @@ def forward( Deprecated and inert -- never read. Masking follows the input shape, so ``use_mask=False`` does *not* disable it. f_sol_override : torch.Tensor, optional - Raw solvent structure factors replacing the cached ``_f_sol_raw`` (k_sol / B_sol - / phase damping still applied). **Overwrites the cache**, so it persists into - later calls until invalidated. Used by ``CollectionScaler``. + Raw solvent structure factors used instead of the cached ``_f_sol_raw`` for this + call only (k_sol / B_sol / phase damping still applied); the cache is left + untouched. Shape ``(N,)`` or ``(B, N)`` -- a batched override keeps the batch + axis of the result. Used by ``CollectionScaler``. Returns ------- @@ -814,20 +841,25 @@ def forward( else: aniso_correction = torch.tensor(1.0, device=self.device, dtype=fcalc.dtype) - if f_sol_override is not None: - self._f_sol_raw = f_sol_override + # An override is consumed locally and never displaces the cache: it may carry a + # leading batch axis, and it belongs to one caller's fraction mixture rather than + # to this scaler's solvent model. + f_sol_raw_local = f_sol_override if hasattr(self, "solvent") and self.solvent is not None: # Lazily cache raw solvent SFs (FFT of mask) — only recomputed # when invalidated via _f_sol_raw = None (e.g. after update_solvent) - if self._f_sol_raw is None: - # The solvent mask is real density with no anomalous term, so - # F_sol(-h) is exactly conj(F_sol(h)) and evaluating on the - # canonical index already matches the canonical fcalc below. - self._f_sol_raw = self.solvent.get_rec_solvent(self.hkl) - + if f_sol_raw_local is None: + if self._f_sol_raw is None: + # The solvent mask is real density with no anomalous term, so + # F_sol(-h) is exactly conj(F_sol(h)) and evaluating on the + # canonical index already matches the canonical fcalc below. + self._f_sol_raw = self.solvent.get_rec_solvent(self.hkl) + f_sol_raw_local = self._f_sol_raw + + # Index the reflection axis, which is last for both (N,) and (B, N). f_sol_raw = ( - self._f_sol_raw[mask] if apply_internal_mask else self._f_sol_raw + f_sol_raw_local[..., mask] if apply_internal_mask else f_sol_raw_local ) if hasattr(self, "log_kmask"): @@ -869,10 +901,13 @@ def forward( else: b_overall = torch.tensor(1.0, device=self.device, dtype=fcalc.dtype) + # f_sol already carries the batch axis when it came from a batched override; + # only a per-reflection (N,) solvent needs one added to broadcast. + f_sol_expanded = f_sol if f_sol.ndim >= 2 else f_sol.unsqueeze(0) fcalc = ( K_overall.unsqueeze(0) * b_overall.unsqueeze(0) - * (aniso_correction.unsqueeze(0) * fcalc + f_sol.unsqueeze(0)) + * (aniso_correction.unsqueeze(0) * fcalc + f_sol_expanded) ) if not batched: diff --git a/torchref/scaling/solvent.py b/torchref/scaling/solvent.py index a226545e..87f906ee 100644 --- a/torchref/scaling/solvent.py +++ b/torchref/scaling/solvent.py @@ -143,6 +143,7 @@ def __init__( verbose=1, float_type=None, device=None, + ignore_hydrogens=True, ): """ Initialize SolventModel. @@ -171,6 +172,8 @@ def __init__( Initial phase offset in radians. verbose : int, default 1 Verbosity level. + ignore_hydrogens : bool, default True + Build the mask from heavy atoms only, whatever the model carries. float_type : torch.dtype, optional Float dtype. ``None`` (default) resolves at runtime to ``get_float_dtype()``, not a hard-wired ``torch.float32``. @@ -190,6 +193,9 @@ def __init__( self.solvent_radius = radius self.erosion_radius = erosion_radius self.optimize_phase = optimize_phase + # Heavy-atom radii already stand in for the hydrogens they carry, so a mask + # built over hydrogen rows too would exclude solvent twice. + self.ignore_hydrogens = bool(ignore_hydrogens) self._cache = TensorDict() # Empty initialization @@ -220,8 +226,6 @@ def __init__( self.model = ModuleReference(model) # Store reference to model self.model.get_vdw_radii() # Ensure VdW radii are available assert self.model, "Model is not initialized" - if model.real_space_grid == None: - model.setup_grid() # Phenix-style parameters self.solvent_radius = radius # For dilation (accessible surface) @@ -353,14 +357,23 @@ def get_solvent_mask(self): xyz = self.model.xyz() # (N_atoms, 3) vdw_radii = self.model.get_vdw_radii() # (N_atoms,) - self.real_space_grid = self.model.real_space_grid + if self.ignore_hydrogens: + # Heavy-atom radii are calibrated for masks built without hydrogens, so + # adding hydrogen spheres on top would exclude solvent twice. + heavy = torch.as_tensor( + (self.model.pdb["element"].str.strip().str.upper() != "H").values, + device=xyz.device, + ) + if not bool(heavy.all()): + xyz = xyz[heavy] + vdw_radii = vdw_radii[heavy] inv_frac = self.model.inv_fractional_matrix frac = self.model.fractional_matrix with torch.no_grad(): spacegroup = self.model.fft.spacegroup n_ops = spacegroup.n_ops - grid_shape = self.real_space_grid.shape[:-1] + grid_shape = self.model.grid_shape device = self.model.device n_atoms = xyz.shape[0] @@ -374,7 +387,7 @@ def get_solvent_mask(self): # grids, where the SF code's 1024 would OOM (denser intermediates). ATOM_CHUNK = 256 - grid_dims = torch.tensor(grid_shape, dtype=torch.long, device=device) + grid_dims = torch.tensor(grid_shape, dtype=torch.long, device=device) # dtype-ok: grid dims for voxel index arithmetic; PyTorch requires int64 grid_shape_float = grid_dims.float() inv_grid = 1.0 / grid_shape_float G = frac.T @ frac # metric tensor: r²_cart = diff_frac · G · diff_frac @@ -451,12 +464,12 @@ def get_solvent_mask(self): protein_voxels = ( torch.cat(protein_chunks, dim=0) if protein_chunks - else torch.empty((0, 3), dtype=torch.long, device=device) + else torch.empty((0, 3), dtype=torch.long, device=device) # dtype-ok: empty (0,3) voxel index tensor; PyTorch requires int64 for indexing ) boundary_voxels = ( torch.cat(boundary_chunks, dim=0) if boundary_chunks - else torch.empty((0, 3), dtype=torch.long, device=device) + else torch.empty((0, 3), dtype=torch.long, device=device) # dtype-ok: empty (0,3) voxel index tensor; PyTorch requires int64 for indexing ) del protein_chunks, boundary_chunks diff --git a/torchref/scaling/weighting.py b/torchref/scaling/weighting.py new file mode 100644 index 00000000..febe6a5f --- /dev/null +++ b/torchref/scaling/weighting.py @@ -0,0 +1,226 @@ +"""How much should a reflection count? The other half of the scaling/weighting split. + +:mod:`torchref.scaling.wilson` answers *what* we compare -- it removes the +resolution trend and leaves `` = 1``. That question turns out to be gauge +for a correlation: any per-shell scaling is absorbed downstream, which is why +twelve normalisation conventions moved a rotation function's truth rank by +nothing. This module answers the question that is not gauge. + +The weight has two sources: + +* **Measurement error**, per reflection, from ``I/sigma_I``. Two reflections at + the same resolution can differ enormously in how well they were measured, and + this is the only part that varies *within* a shell. That matters more than it + sounds -- a weight constant within a shell is a per-shell weight, and those + are exactly what a correlation absorbs. +* **Model error**, per resolution, through ``sigma_A``, which is smooth in + ``|s|`` and has no per-reflection content at all. + +Both come out of weighting by inverse variance, +``w = 1/(sigma_meas^2 + sigma_model^2)`` with ``sigma_model^2 = eps - +sigma_A^2``, the standard MLHL budget. The saturation usually added by hand is +already in there: once model error dominates, extra measurement precision buys +nothing, because what is wrong is the model and not the data. + +**They do not separate, and that was measured rather than assumed.** Writing the +weight as a product of a ``snr`` term and a ``sigma_A`` term looks natural and +fails: the ``sigma_A`` half is then dominated by its own singularity as +``sigma_A -> 1`` and stops depending on the model error it is named after. The +measurement term is what regularises it, so the two belong in one denominator. +:func:`inverse_variance_weight` is the form that works; +:func:`information_weight` is the pure measurement half, kept because it is the +right answer when no ``sigma_A`` is available and because it is the control that +says what the coupling is worth. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +__all__ = [ + "information_weight", + "inverse_variance_weight", + "snr_from_amplitude", + "normalise_weight", + "empirical_sigma_a", +] + +#: ``I/sigma_I`` at which measurement error stops being the limiting term. +#: Beyond it a reflection is no better determined for the purpose of comparing +#: against a model, because the model is what limits. Free parameter in +#: practice: the crossover really sits wherever ``sigma_A`` puts it, and +#: ``sigma_A`` before placement is assumed rather than fitted. +DEFAULT_SNR_CAP = 5.0 + +#: Backstop on the inverse-variance weight, for the case where measurement and +#: model variance both vanish. Not the working mechanism: the measurement term +#: is what bounds the weight at low resolution, where ``sigma_A -> 1`` would +#: otherwise send it to infinity on the strongest reflections. If this binds on +#: real data, ``sigma_A`` is wrong rather than the cap being too low. +DEFAULT_TRUST_CAP = 100.0 + + +def snr_from_amplitude( + F: torch.Tensor, sig_F: torch.Tensor, floor: float = 1e-12, +) -> torch.Tensor: + """``I/sigma_I`` from an amplitude and its error. + + With ``I = F^2`` the error propagates as ``sigma_I = 2 F sigma_F``, so the + intensity signal-to-noise is ``F / (2 sigma_F)`` -- half the amplitude's. + The factor is worth being explicit about: it only rescales the cap, but + quoting a cap against the wrong one silently doubles it. + """ + return (F.abs() / (2.0 * sig_F.abs().clamp(min=floor))).clamp(min=0.0) + + +def information_weight( + snr: torch.Tensor, *, cap: float = DEFAULT_SNR_CAP, +) -> torch.Tensor: + """Saturating measurement-information weight, ``snr^2 / (snr^2 + cap^2)``. + + Rises as ``(snr/cap)^2`` while measurement error dominates and flattens to 1 + once it does not. This is not a sigmoid chosen for its shape -- it is + ``1 / (1 + sigma_meas^2/sigma_model^2)`` rewritten, with ``cap`` the + signal-to-noise at which the two are equal. The saturation is a consequence + of the variance budget rather than a clip applied on top of one. + + Structurally this is what Phaser's ``DFAC`` already does: a monotone + function of signal-to-noise, in ``(0, 1)``, tending to 1 for well-measured + reflections. The difference is that its saturation point falls out of a Rice + moment calculation and cannot be moved, and this one is a number that can be + screened. + + Parameters + ---------- + snr : torch.Tensor + ``(N,)`` ``I/sigma_I``. Negative or zero values give weight 0, which is + the right answer for a measurement consistent with nothing. + cap : float, optional + Signal-to-noise at which the weight reaches 1/2. + """ + s2 = snr.clamp(min=0.0) ** 2 + return s2 / (s2 + float(cap) ** 2) + + +def inverse_variance_weight( + snr: torch.Tensor, + sigma_a: torch.Tensor, + *, + eps: Optional[torch.Tensor] = None, + cap: float = DEFAULT_TRUST_CAP, +) -> torch.Tensor: + """``1 / (1/snr^2 + eps - sigma_A^2)`` -- the two error sources, together. + + The obvious design is to factorise: a per-reflection term in ``snr`` times a + resolution term in ``sigma_A``. It does not work, and the reason is worth + keeping. + + Taken alone, ``sigma_A/(eps - sigma_A^2)`` is dominated by its own + singularity. On a realistic Luzzati falloff it runs 10.1 at the lowest + resolution shell against 1.0 at the next, and it comes out *identical* for + a 0.5 A and a 1.0 A coordinate error -- because as ``sigma_A -> 1`` the + shape is set entirely by ``1/(1 - sigma_A^2)`` and the model error it was + supposed to encode drops out. A weight carrying no information about the + thing it is named after is not a weight worth having. + + What regularises it is the term the factorisation threw away. Those + low-resolution reflections are strong but not infinitely well measured, so + ``1/snr^2`` is what stops the variance reaching zero. The two sources have + to sit in one denominator; they do not separate. + + ``1/snr^2`` is the measurement variance expressed in the same units as + ``eps - sigma_A^2``, i.e. relative to a normalised `` = 1``. The cap + is a backstop for the pathological case where both terms vanish, not the + working mechanism -- if it binds on real data, ``sigma_A`` is wrong. + + Parameters + ---------- + snr : torch.Tensor + ``(N,)`` ``I/sigma_I``. Zero gives zero weight. + sigma_a : torch.Tensor + ``(N,)`` model reliability in ``[0, 1)``, evaluated at each reflection. + eps : torch.Tensor, optional + ``(N,)`` multiplicity; ``None`` means 1. + cap : float, optional + Ceiling on the weight before normalisation. + """ + sa = sigma_a.clamp(min=0.0, max=1.0 - 1e-6) + e = torch.ones_like(sa) if eps is None else eps.to(sa.dtype).clamp(min=1.0) + v_meas = 1.0 / (snr.clamp(min=1e-8) ** 2) + v_model = (e - sa * sa).clamp(min=0.0) + w = 1.0 / (v_meas + v_model).clamp(min=1e-12) + return w.clamp(max=float(cap)) + + +def normalise_weight(w: torch.Tensor) -> torch.Tensor: + """Scale a weight to mean 1. + + Cosmetic for a correlation, where an overall factor cancels, and not + cosmetic for anything that compares scores across runs or reads a sigma + level off them. Doing it here means the cap is a number about *relative* + weighting rather than one entangled with whatever scale the inputs had. + """ + return w / w.mean().clamp(min=1e-30) + + +def empirical_sigma_a( + sigma_obs: torch.Tensor, + sigma_calc: torch.Tensor, + *, + floor: float = 1e-3, +) -> torch.Tensor: + """Model reliability measured, rather than assumed, from two Wilson curves. + + ``sigma_A`` in a rotation search is normally a *prior*: a Luzzati falloff + from a coordinate error guessed off the residue count, patched at low + resolution by Babinet's two universal constants. It never sees a residual. + + It does not have to. Total scattering per shell is **rotation-invariant**, + so the resolution-dependent disagreement between model and data is + measurable before the molecule is placed, even though the per-reflection + disagreement is not. Normalise both sides to `` = 1`` and their fitted + curves' ratio is exactly that disagreement: + + R(s) = Sigma_obs(s) / Sigma_calc(s) + + ``R < 1`` means the model predicts more scattering than is there, which at + low resolution is the bulk solvent it does not have; ``R > 1`` means it + predicts less. Either way the shared fraction is bounded by + ``min(R, 1/R)``, and ``sigma_A`` is its square root because ``sigma_A^2`` is + the fraction of intensity the model accounts for. + + **This is safe to estimate from the data being scored**, which normally it + would not be: the quantity is identical for every candidate orientation, so + it shifts all scores together and cannot bias the ranking toward any of + them. + + What it conflates -- solvent, an overall B mismatch, missing atoms, genuine + coordinate error -- it conflates deliberately. For deciding how far to trust + a resolution range the cause does not matter, only the size. What it cannot + see is *completeness*: forcing both sides to unit mean absorbs a uniform + factor, so a model that is half the asymmetric unit looks like a model that + is all of it, and only the tilt survives. + + That uniform factor is removed here, by dividing each curve by its geometric + mean over the points supplied. The two fits carry their own absolute + scales -- the data's arbitrary one and the model's electron scale -- and + without this step the ratio's *level* set the answer rather than its shape: + measured 0.02-0.06 on 1DAW and 2DQ6 and 8-12 on 3K7M, giving a flat + ``sigma_A`` of 0.15-0.35 that said nothing about resolution. + + Parameters + ---------- + sigma_obs, sigma_calc : torch.Tensor + ``(N,)`` fitted Wilson curves evaluated at the same ``|s|``. They must + come from fits sharing an abscissa, or each is frozen flat outside its + own range and the ratio is meaningless there. + floor : float, optional + Lower bound on the returned ``sigma_A``. + """ + log_r = (sigma_obs.clamp(min=1e-30).log() + - sigma_calc.clamp(min=1e-30).log()) + log_r = log_r - log_r.mean() # unit geometric mean: scale-free + shared = torch.exp(-log_r.abs()) # min(R, 1/R) + return shared.sqrt().clamp(min=float(floor), max=1.0 - 1e-6) diff --git a/torchref/scaling/wilson.py b/torchref/scaling/wilson.py new file mode 100644 index 00000000..e16b048d --- /dev/null +++ b/torchref/scaling/wilson.py @@ -0,0 +1,481 @@ +"""Absolute Wilson normalisation: fit ``Sigma(s)`` and divide it out. + +Distinct from :class:`~torchref.scaling.scaler_base.ScalerBase`, which is a +*relative* scaler -- it puts ``F_calc`` onto ``F_obs`` and every target it can +minimise compares the two. This one takes a single dataset and answers "what is +the expected intensity at this resolution", so that dividing by it leaves +`` = 1``. One dataset in, one curve out, no second dataset anywhere in the +objective. + +**Why this exists as one shared class.** The repo grew at least five private +answers to the same question -- ``base/wilson_outliers.robust_mean_intensity``, +``base/french_wilson.estimate_mean_intensity_by_resolution``, +``ReflectionData._calculate_wilson_b``, the ``Sigma_N`` estimator in +``refinement/model_error_estimation/sigma_a``, and a per-shell one inside the +alignment package -- differing in whether they use means or medians, whether +they divide out ``epsilon``, whether they separate centrics, and where they put +their shell edges. Consumers that disagree about what E means cannot be compared +with each other, which is exactly what went wrong between the rotation function +and its own rescore. + +**Scaling, not weighting.** This class answers *what* we compare. It says +nothing about how much any reflection should count -- no ``sigI``, no model +error, no solvent. Those belong to a weight, and mixing them in here is what +made the previous convention object impossible to reason about: it returned a +normalisation and a weight together, so sweeping it moved a gauge quantity and a +real one at the same time. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from torchref.config import get_float_dtype +from torchref.scaling.basis import chebyshev_design + +__all__ = ["WilsonNormaliser"] + +#: Chebyshev terms. Enough to follow a Wilson plot's curvature and the +#: low-resolution solvent deficit without chasing shell-to-shell noise. +#: Provisional: the order has never been chosen against a metric sensitive to +#: it, so screen on this class's own residual trend rather than on anything +#: downstream. +DEFAULT_N_COEFF = 6 + +#: Bound on ``log Sigma`` relative to its own constant term. A polynomial is +#: unbounded at the ends of its range, so without this a single extreme +#: reflection at the resolution limit can carry an arbitrary scale -- the same +#: reason ``ScalerBase.iso_log_scale`` clamps per reflection. +LOG_CLAMP = 10.0 + +#: Step halvings allowed per IRLS iteration before the step is abandoned. +MAX_HALVINGS = 30 + +#: IRLS iterations allowed before the fit is declared failed. Generous, because +#: it should never bind: at :data:`DEFAULT_RTOL` the fit converges in single +#: digits. It is a runaway guard, not a budget. +DEFAULT_MAX_ITER = 100 + +#: Floor on the fitted mean, to keep ``y/mu`` finite if a step overshoots. +#: Must be representable in the working dtype -- ``1e-300`` is a float64 +#: constant and flushes to zero in float32, which turns the guard into the +#: division by zero it exists to prevent. +_MU_FLOOR = 1e-30 + +#: Relative convergence tolerance -- see :meth:`WilsonNormaliser._irls` for why +#: it is relative to the improvement so far rather than to the objective. +#: +#: This is a normalisation curve, not a refined parameter. The quantity it +#: decides is ``E = F / sqrt(Sigma)``, which is then compared against a model +#: that is wrong by tens of percent, so four digits is already far past what +#: anything downstream can use. +DEFAULT_RTOL = 1e-4 + + +class WilsonNormaliser: + """``Sigma(s)`` by Gamma GLM, so that `` = 1`` by construction. + + The model is `` = eps_h * Sigma(s_h)`` with ``log Sigma`` a Chebyshev + polynomial in ``sin(theta)/lambda``, fitted by maximum likelihood under + + acentric I ~ Exp(Sigma) (Gamma, shape 1) + centric I ~ Sigma * chi^2_1 (Gamma, shape 1/2) + + i.e. a Gamma GLM with a log link and the shape as the prior weight. + + **Unit mean is an identity of the fit, not a normalisation step.** The + constant basis column's score equation is ``sum_h k_h (I_h/mu_h - 1) = 0``, + which is exactly `` = 1`` in the shape-weighted sense. Nothing is + rescaled afterwards and nothing can drift -- which is what makes a + downstream ``E^2 - 1`` a true centring rather than an approximate one. + + Least squares on ``log I`` would be the obvious alternative and is wrong: + ``E[log Gamma]`` carries a digamma offset, and with a constant term present + it is absorbed into the curve's shape rather than into the level. That is + the defect the overall-anisotropy fit was carrying. + + Parameters + ---------- + I : torch.Tensor + ``(N,)`` intensities. **Intensities, not amplitudes** -- Wilson + statistics are exact on I and awkward on F, and measurement error is + near-Gaussian on I but badly behaved on F for weak reflections, which is + the whole reason the French-Wilson posterior exists. Negative values are + allowed and kept: they are meaningful, unbiased measurements. They are + excluded from the *fit* (the Gamma likelihood has no support there) but + still receive a ``Sigma`` and a signed ``E_squared``. + s_mag : torch.Tensor + ``(N,)`` scattering-vector magnitude ``|s| = 1/d``, in inverse Angstrom. + eps : torch.Tensor, optional + ``(N,)`` reflection multiplicity. Divides the intensity before the fit, + because axial reflections are systematically stronger. ``None`` means 1 + everywhere, which is correct for a molecular transform sampled in a P1 + box -- multiplicity is a property of crystal symmetry and there is none + there. + centric : torch.Tensor, optional + ``(N,)`` bool, setting the Gamma shape. ``None`` means all acentric. + n_coeff : int, optional + Chebyshev terms. ``1`` gives a single global scale. + s_lo, s_hi : float, optional + ``|s|`` range mapped onto the basis. Defaults to this dataset's own + extremes. **Pass both explicitly whenever the curve will be evaluated + outside the fitted data's range** -- comparing two fits over different + ranges, or fitting on a crystal lattice and evaluating on a dense + sampling. The basis saturates at the ends, so beyond the fitted range + the curve is frozen flat rather than extrapolated. + fit_mask : torch.Tensor, optional + ``(N,)`` bool selecting which reflections *inform* the fit. Everything + still receives a ``Sigma``, because the curve is smooth and evaluable + anywhere. Use it to hold out systematic absences -- see + :meth:`from_hkl`, which does exactly that. + + Attributes + ---------- + coefficients : torch.Tensor + ``(n_coeff,)`` fitted Chebyshev coefficients of ``log Sigma``. + sigma_wilson : torch.Tensor + ``(N,)`` fitted ``Sigma(s)``. Deliberately not called ``sigma``: this + package also carries ``sig_F``, a measurement error, and ``sigma_a``, a + correlation coefficient, and the three are not interchangeable. + mean_intensity : torch.Tensor + ``(N,)`` ``eps * Sigma(s)``, the expected intensity of each reflection. + E_squared : torch.Tensor + ``(N,)`` ``I / mean_intensity``. **Signed** -- negative observations stay + negative. + E : torch.Tensor + ``(N,)`` ``sqrt(max(E_squared, 0))``. + """ + + MAX_HALVINGS = MAX_HALVINGS + + def __init__( + self, + I: torch.Tensor, + s_mag: torch.Tensor, + *, + eps: Optional[torch.Tensor] = None, + centric: Optional[torch.Tensor] = None, + n_coeff: int = DEFAULT_N_COEFF, + s_lo: Optional[float] = None, + s_hi: Optional[float] = None, + fit_mask: Optional[torch.Tensor] = None, + max_iter: int = DEFAULT_MAX_ITER, + rtol: float = DEFAULT_RTOL, + ) -> None: + if I.ndim != 1: + raise ValueError(f"I must be 1-D, got {tuple(I.shape)}") + if s_mag.shape != I.shape: + raise ValueError( + f"s_mag {tuple(s_mag.shape)} does not match I {tuple(I.shape)}" + ) + self.dtype = I.dtype + self.n_coeff = int(n_coeff) + self._I = I + self._s_mag = s_mag + self._eps = eps + self.s_lo = float(s_mag.min()) if s_lo is None else float(s_lo) + self.s_hi = float(s_mag.max()) if s_hi is None else float(s_hi) + + # The configured float dtype, not float64. This is a six-coefficient + # fit of a smooth curve whose answer is compared against a model wrong + # by tens of percent; it does not need double, and hardcoding it here + # would be the only double-precision path in the scaling package. + work = get_float_dtype() + eps_w = ( + torch.ones_like(I, dtype=work) if eps is None + else eps.to(work).clamp(min=1.0) + ) + # Shape 1 acentric (exponential), 1/2 centric. Enters as the IRLS weight + # because for a Gamma with shape k the variance is mu^2/k, so the + # log-link working weight is k itself. + k = ( + torch.ones_like(I, dtype=work) if centric is None + else torch.where(centric.to(torch.bool), 0.5, 1.0).to(work) + ) + + I_reduced = I.to(work) / eps_w + # The Gamma likelihood has no support at or below zero. Absences and + # negative measurements are held out of the fit and given a Sigma from + # the curve like everything else -- excluding them from the *estimate* + # is not the same as refusing to normalise them. + usable = torch.isfinite(I_reduced) & torch.isfinite(s_mag) & (I_reduced > 0) + if fit_mask is not None: + usable = usable & fit_mask.to(torch.bool) + if int(usable.sum()) < self.n_coeff + 1: + raise ValueError( + f"only {int(usable.sum())} usable reflections for a " + f"{self.n_coeff}-coefficient fit; need at least {self.n_coeff + 1}" + ) + self.n_fitted = int(usable.sum()) + + design = chebyshev_design( + (s_mag * 0.5).to(work), self.n_coeff, + lo=self.s_lo * 0.5, hi=self.s_hi * 0.5, + ) + self.coefficients, self.n_iter = self._irls( + design[usable], I_reduced[usable], k[usable], max_iter, rtol, + ) + + log_sigma = self._eval_log_sigma(design) + self.sigma_wilson = torch.exp(log_sigma).to(self.dtype) + self.mean_intensity = ( + torch.exp(log_sigma) * eps_w + ).clamp(min=1e-30).to(self.dtype) + self.E_squared = I / self.mean_intensity + self.E = self.E_squared.clamp(min=0.0).sqrt() + + # -- fitting ----------------------------------------------------------- + + def _eval_log_sigma(self, design: torch.Tensor) -> torch.Tensor: + c = self.coefficients + return (design @ c).clamp( + min=-LOG_CLAMP + float(c[0]), max=LOG_CLAMP + float(c[0]), + ) + + @staticmethod + def _solve_intercept( + beta: torch.Tensor, X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, + ) -> torch.Tensor: + """Put the intercept exactly on its score equation, closed form. + + The intercept's stationarity condition is ``sum_h k_h (I_h/mu_h - 1) = + 0``, which is `` = 1`` -- the identity this class exists to + provide. Shifting ``beta[0]`` by ``d`` scales every ``mu`` by ``e^d``, + so the ``d`` that satisfies it is available in one line: + + e^d = sum_h k_h (I_h/mu_h) / sum_h k_h + + Doing this explicitly decouples the identity from how tightly the SHAPE + converged. Without it `` = 1`` is only as good as the overall fit + tolerance -- at ``rtol = 1e-4`` it came out at 1 - 1e-5 -- and the + identity is not the kind of claim that should degrade with a stopping + rule. The remaining coefficients are untouched, so this changes the + curve's level and not its shape. + + In the working dtype like everything else here. The point is to make the + identity independent of the *stopping rule*, not to chase digits: it + lands within about 1e-6 of one, which is two orders inside anything that + reads it. + """ + eta = X @ beta + mu = torch.exp(eta.clamp(min=-LOG_CLAMP + float(beta[0]), + max=LOG_CLAMP + float(beta[0]))).clamp(min=_MU_FLOOR) + ratio = ((w * (y / mu)).sum() / w.sum()).clamp(min=_MU_FLOOR) + out = beta.clone() + out[0] = out[0] + torch.log(ratio) + return out + + + def _irls( + self, + X: torch.Tensor, + y: torch.Tensor, + w: torch.Tensor, + max_iter: int, + rtol: float, + ) -> Tuple[torch.Tensor, int]: + """Gamma GLM with a log link, by iteratively reweighted least squares. + + IRLS rather than a generic optimiser: for this link and family the + working weight does not depend on ``mu``, so each step is one weighted + least-squares solve and there is no step size, no line search and no + absolute tolerance to fail against an unnormalised objective. + + Convergence and step control both use the objective itself, + ``L = sum_h k_h (y_h/mu_h + log mu_h)`` -- the negative log-likelihood + with the terms not involving ``beta`` dropped. + + That choice is forced by what the alternatives do on real data. The + *coefficients* are underdetermined whenever the data occupy part of the + basis range, which is the normal case once an explicit ``s_lo``/``s_hi`` + is passed, so they wander in the flat directions long after the fit has + settled. The *deviance* carries a ``-log(y/mu)`` term that diverges as + ``y -> 0``, and calculated amplitudes have near-zeros at the nodes of + the molecular transform, so a few tiny intensities dominate it. And the + *fitted mean* cannot be compared as a ratio because it is floored, so a + collapsed fit reads as a converged one -- which is exactly how an early + version of this reported success while returning zeros. + + ``L`` has none of those problems: the ``log y`` term that breaks the + deviance is constant in ``beta`` and simply absent here. + + Step halving is the other half. IRLS on a log link can overshoot into + ``mu`` underflow, after which the working response ``y/mu`` explodes and + the next step is worse. Rejecting any step that does not improve ``L`` + and halving it is the standard remedy and makes the fit robust to the + ill-conditioning a partial basis range creates. + """ + # Seed at the constant curve, which is the exact MLE when Sigma has no + # resolution dependence. Every later iteration only adds shape. + beta = torch.zeros(self.n_coeff, dtype=X.dtype, device=X.device) + beta[0] = torch.log(((w * y).sum() / w.sum()).clamp(min=1e-30)) + + def objective(b): + eta = (X @ b).clamp( + min=-LOG_CLAMP + float(b[0]), max=LOG_CLAMP + float(b[0]), + ) + mu = torch.exp(eta).clamp(min=_MU_FLOOR) + return float((w * (y / mu + eta)).sum()), eta, mu + + L, eta, mu = objective(beta) + L0 = L # the constant-curve seed, for the ratio below + + # Built and factorised ONCE. For a Gamma with a log link the IRLS + # working weight is the shape k, which does not depend on mu -- so + # `X^T W X` is the same matrix at every iteration and only the working + # response changes. Rebuilding it per iteration costs an O(N n^2) pass + # over every reflection for an answer that cannot have changed. + XtW = X.transpose(0, 1) * w.unsqueeze(0) + A = XtW @ X + # Ridge proportional to the matrix's own scale: the high-order + # Chebyshev columns go near-singular when the data cover only part + # of the basis range. + A = A + torch.eye(self.n_coeff, dtype=A.dtype, device=A.device) * ( + 1e-10 * float(torch.diagonal(A).abs().max().clamp(min=1e-30)) + ) + # Factorised once, and by Cholesky rather than LU. `A` is `X^T W X` + # plus a ridge with positive IRLS weights, so it is symmetric positive + # definite by construction -- and MPS implements neither `lu_solve` nor + # `cholesky_solve` (torch 2.9.1), which left the per-iteration solve to + # a CPU round trip: 1015 us against 114 us for two triangular solves, on + # a 200k x 6 problem whose unavoidable `XtW @ z` is 966 us. Same + # arithmetic -- over the 16 datasets in ``tests/files/mtz`` the two + # agree to 1e-5 relative in float32 and 1e-14 in float64, with identical + # iteration counts on every one. + # + # `cholesky_ex` reports rather than raises, because a fully collinear + # basis is a thing this fit sees: the high-order Chebyshev columns go + # near-singular when the data cover only part of the basis range, and + # the ridge does not always rescue that. LU carried no definiteness + # requirement, so that case falls back to a general solve instead of + # failing. + chol = torch.linalg.cholesky_ex(A) + L_A = chol.L if int(chol.info) == 0 else None + + def _solve(rhs): + if L_A is None: + return torch.linalg.solve(A, rhs) + return torch.linalg.solve_triangular( + L_A.mT, + torch.linalg.solve_triangular(L_A, rhs, upper=False), + upper=True, + ) + + for it in range(1, max_iter + 1): + z = eta + (y - mu) / mu # working response + step = _solve((XtW @ z).unsqueeze(-1)).squeeze(-1) - beta + if not torch.isfinite(step).all(): + raise RuntimeError( + f"Wilson fit diverged at iteration {it}: the IRLS solve " + f"returned non-finite coefficients." + ) + + # Halve until the step actually improves the objective. + accepted = False + for _ in range(self.MAX_HALVINGS): + L_try, eta_try, mu_try = objective(beta + step) + if L_try <= L: + beta = beta + step + accepted = True + break + step = step * 0.5 + if not accepted: + # No downhill direction left: already at the optimum. + return self._solve_intercept(beta, X, y, w), it + + # Relative to the improvement achieved so far, not to |L|. + # + # |dL|/|L| is not usable here: under I -> cI the optimum is just + # beta[0] -> beta[0] + log c, so the fit is exactly scale invariant, + # but L picks up an additive `log c * sum(k)` and the ratio would + # mean something different at every scale. That additive term + # cancels in any DIFFERENCE, so a ratio of two differences is both + # relative and scale invariant -- which is what this is. + # + # The denominator is the total distance travelled from the constant + # seed, so the test reads "the last step moved us less than rtol of + # the way we have come". It is bounded below so a fit that starts at + # its own optimum (Sigma genuinely flat) terminates rather than + # dividing by zero. + step_gain = abs(L - L_try) + total_gain = max(abs(L0 - L_try), 1e-30) + L, eta, mu = L_try, eta_try, mu_try + if step_gain <= rtol * total_gain: + return self._solve_intercept(beta, X, y, w), it + raise RuntimeError( + f"Wilson fit did not converge in {max_iter} IRLS iterations " + f"(last step still worth {step_gain / total_gain:.2e} of the total " + f"improvement, against rtol={rtol:.0e}). Raising rather than " + f"falling back to a coarser estimate: a normaliser that silently " + f"becomes a different normaliser on hard cases is two normalisers " + f"wearing one name." + ) + + # -- evaluation elsewhere --------------------------------------------- + + def evaluate(self, s_mag: torch.Tensor) -> torch.Tensor: + """``Sigma(s)`` at arbitrary ``|s|``, on the basis this fit was built on. + + The curve is smooth, so it can be fitted on one reflection set and used + on another -- which is what makes a fit on the crystal lattice usable on + a dense sampling of the same transform. **Only inside ``[s_lo, s_hi]``**: + the basis saturates at the ends, so outside that range this returns the + endpoint value, flat, rather than an extrapolation. + """ + design = chebyshev_design( + (s_mag * 0.5).to(self.coefficients.dtype), self.n_coeff, + lo=self.s_lo * 0.5, hi=self.s_hi * 0.5, + ) + return torch.exp(self._eval_log_sigma(design)).to(self.dtype) + + # -- construction from crystallography -------------------------------- + + @classmethod + def from_hkl( + cls, + I: torch.Tensor, + hkl: torch.Tensor, + spacegroup, + cell, + **kwargs, + ) -> "WilsonNormaliser": + """Build from Miller indices, deriving ``|s|``, ``eps`` and centricity. + + The core takes plain tensors because not every caller has crystal + reflections -- a molecular transform sampled in a P1 box has no ``hkl`` + at all, and there ``eps`` is 1 with nothing centric. This constructor is + for the case that does. + + ``epsilon(friedel=False)``: Wilson's `` = eps * Sigma`` counts the + operations mapping ``h -> h``, which add coherently and set the mean. + The Friedel-folded count changes the *distribution* instead, and that is + centricity -- which enters here as the Gamma shape, separately. The two + branches feed two different parameters of the same likelihood. + """ + work = get_float_dtype() + hkl_l = hkl.to(torch.long) # dtype-ok: Miller indices are integers + # The cell may carry the configured default device while the reflections + # are somewhere else; the caller should not have to reconcile them. + rec = cell.reciprocal_basis_matrix.to(device=hkl_l.device, dtype=work) + s_mag = (hkl_l.to(work) @ rec).norm(dim=-1).to(I.dtype) + eps = spacegroup.epsilon(hkl_l, friedel=False).to(work) + centric = spacegroup.is_centric(hkl_l).to(torch.bool) + # Systematically absent reflections are zero by symmetry, not by + # measurement, so they carry no information about Sigma and would drag + # the Gamma fit toward zero. + fit_mask = ~spacegroup.is_absent(hkl_l).to(torch.bool) + user_mask = kwargs.pop("fit_mask", None) + if user_mask is not None: + fit_mask = fit_mask & user_mask.to(torch.bool) + return cls( + I, s_mag, eps=eps, centric=centric, fit_mask=fit_mask, **kwargs, + ) + + def __repr__(self) -> str: # pragma: no cover - display + return ( + f"{type(self).__name__}(N={self._I.numel()}, " + f"n_coeff={self.n_coeff}, n_fitted={self.n_fitted}, " + f"iters={self.n_iter})" + ) diff --git a/torchref/scripts/extract_ener_lib.py b/torchref/scripts/extract_ener_lib.py new file mode 100644 index 00000000..e0edc1b9 --- /dev/null +++ b/torchref/scripts/extract_ener_lib.py @@ -0,0 +1,97 @@ +"""Extract the per-atom-type table from the CCP4 energy library into a bundled CSV. + +The monomer library types every atom (``_chem_comp_atom.type_energy``: NH1, OC, CH3, +...) and ``ener_lib.cif`` says what each type is: its element, whether it donates or +accepts hydrogen bonds, and its van der Waals radius with and without the hydrogens it +normally carries. TorchRef reads that table from ``torchref/data/ener_lib_atoms.csv``; +this script regenerates the CSV from the library so the two cannot drift apart +unnoticed. + +Run as ``python -m torchref.scripts.extract_ener_lib [path/to/ener_lib.cif]``. Without a +path the library is fetched through the monomer-library manager. The hydrogen-bond +distance table is printed for inspection; it is the source of the contact-policy +defaults and is not bundled. +""" + +import csv +import sys +from pathlib import Path + +import gemmi + +from torchref import PATH_TORCHREF_DATA + +_ATOM_COLUMNS = ( + "type", + "weight", + "hb_type", + "vdw_radius", + "vdwh_radius", + "ion_radius", + "element", + "valency", + "sp", +) + +_OUT_COLUMNS = ("type", "element", "hb_type", "vdw_radius", "vdwh_radius", "ion_radius") + + +def _null(value: str) -> str: + return "" if value in (".", "?") else value + + +def extract(ener_lib: Path, out_csv: Path) -> int: + """Write the ``_lib_atom`` loop of ``ener_lib`` to ``out_csv``; return the row count.""" + block = gemmi.cif.read_file(str(ener_lib))[0] + table = block.find("_lib_atom.", list(_ATOM_COLUMNS)) + rows = [] + for row in table: + record = dict(zip(_ATOM_COLUMNS, (str(v) for v in row))) + vdw = _null(record["vdw_radius"]) + vdwh = _null(record["vdwh_radius"]) or vdw + rows.append( + { + "type": record["type"], + "element": record["element"], + "hb_type": record["hb_type"], + "vdw_radius": vdw, + "vdwh_radius": vdwh, + "ion_radius": _null(record["ion_radius"]), + } + ) + with open(out_csv, "w", newline="") as handle: + handle.write( + "# Per-energy-type atom properties from the CCP4 monomer library " + "ener_lib.cif (_lib_atom loop).\n" + "# hb_type: N neither, D donor, A acceptor, B both, " + "H hydrogen able to hydrogen-bond.\n" + "# vdw_radius: contact radius in Angstrom; vdwh_radius: radius to use when the " + "atom's own hydrogens are not modelled.\n" + "# Regenerate with python -m torchref.scripts.extract_ener_lib\n" + ) + writer = csv.DictWriter(handle, fieldnames=list(_OUT_COLUMNS)) + writer.writeheader() + writer.writerows(rows) + + hbond = block.find("_lib_hbond.", ["atom_type_1", "atom_type_2", "min", "dist"]) + print(f"{len(rows)} atom types written to {out_csv}") + print("hydrogen-bond distance table (type_1, type_2, well depth, distance):") + for row in hbond: + print(" ", " ".join(str(v) for v in row)) + return len(rows) + + +def main(argv=None) -> int: + argv = sys.argv[1:] if argv is None else argv + if argv: + ener_lib = Path(argv[0]) + else: + from torchref.topology.monomer.library import MonomerLibraryManager + + ener_lib = Path(MonomerLibraryManager(verbose=0).ensure_gemmi_base()) / "ener_lib.cif" + extract(ener_lib, Path(PATH_TORCHREF_DATA) / "ener_lib_atoms.csv") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/torchref/symmetry/__init__.py b/torchref/symmetry/__init__.py index 09014038..b8e3b862 100644 --- a/torchref/symmetry/__init__.py +++ b/torchref/symmetry/__init__.py @@ -1,83 +1,35 @@ -"""Crystallographic symmetry: space groups, unit cells, map and HKL symmetry. +"""Crystallographic symmetry: symmetry groups, space groups and unit cells. -:class:`SpaceGroup` (``nn.Module`` holding the operations as buffers) is the entry -point, with ``Symmetry`` a bare alias for it. :func:`MapSymmetry` handles real-space -density and :func:`ReciprocalSymmetry` structure-factor grids; all three accept a -space group as a string, an int 1-230, or a gemmi object. :class:`Cell` is separate --- it wraps the six cell parameters, not a space group. +:class:`Symmetry` holds a group as rotation matrices and fractional translations and +owns every verb derivable from the operations alone -- expansion of positions and +Miller indices, translation phases, the reflection predicates, symmetry-compatible grid +sizes, and map symmetrization. Nothing in it is crystallographic, so a group built from +a raw operation list serves non-crystallographic symmetry too. -The grid utilities re-exported here come from ``grid_utils``, which delegates to -``spacegroup``. ``spacegroup`` also defines its own same-named copies, which are -the source of truth and are *not* re-exported. +:class:`SpaceGroup` specialises it with the crystallographic identity (Hermann-Mauguin +naming, number, point group, crystal system) and the CCP4 asymmetric-unit verbs +(``expand_hkl``, ``reduce_hkl``, ``complete_hkl``, ``canonicalize_hkl``). It accepts a +name, a number 1-230, a ``gemmi.SpaceGroup``, another instance, or None for P1. + +:class:`Cell` is separate: it wraps the six cell parameters, not a symmetry group. + +Map and reciprocal-grid operators are reached through :class:`Symmetry` +(:meth:`~Symmetry.symmetrize_map`, :meth:`~Symmetry.reciprocal_extractor`), which owns +their caching -- the operator classes themselves are private. """ -from .cell import Cell, CellTensor -from .grid_utils import ( - calculate_optimal_grid_size, - check_grid_compatibility, - find_fft_friendly_size, - get_symmetry_grid_requirements, - is_fft_friendly, - recommend_grid_size, -) -from .map_symmetry import MapSymmetry, MapSymmetryDirect -from .reciprocal_symmetry import ( - ReciprocalSymmetry, - ReciprocalSymmetryGrid, - canonicalize_hkl, - complete_hkl, - expand_hkl, - expand_reciprocal_grid, - expand_reflections, - reduce_hkl, -) -from .spacegroup import ( - SpaceGroup, - SpaceGroupLike, - get_crystal_system, - get_operations_as_tensors, - get_point_group, - get_symmetry_operations, - is_centrosymmetric, - is_same_spacegroup, - n_operations, - spacegroup_to_str, -) -from .symmetry import Symmetry +from .cell import Cell +from .spacegroup import SpaceGroup, SpaceGroupLike +from .symmetry import Symmetry, find_fft_friendly_size, is_fft_friendly __all__ = [ # Unit cell "Cell", - # Space group utilities + # Symmetry groups + "Symmetry", "SpaceGroup", "SpaceGroupLike", - "spacegroup_to_str", - "get_symmetry_operations", - "get_operations_as_tensors", - "is_same_spacegroup", - "get_point_group", - "get_crystal_system", - "is_centrosymmetric", - "n_operations", - # Base symmetry - "Symmetry", - # Real space map symmetry - "MapSymmetry", - "MapSymmetryDirect", - # Reciprocal space symmetry - "ReciprocalSymmetry", - "ReciprocalSymmetryGrid", - "expand_hkl", - "complete_hkl", - "reduce_hkl", - "canonicalize_hkl", - "expand_reflections", - "expand_reciprocal_grid", - # Grid utilities - "get_symmetry_grid_requirements", - "check_grid_compatibility", - "recommend_grid_size", - "find_fft_friendly_size", + # Grid sizing helpers (group-independent) "is_fft_friendly", - "calculate_optimal_grid_size", + "find_fft_friendly_size", ] diff --git a/torchref/symmetry/cell.py b/torchref/symmetry/cell.py index cc0eebf0..63a47e94 100644 --- a/torchref/symmetry/cell.py +++ b/torchref/symmetry/cell.py @@ -20,7 +20,7 @@ from torchref.utils.device_mixin import _NonModuleDeviceMixin -@dataclass +@dataclass(eq=False) class Cell(_NonModuleDeviceMixin): """ Dataclass for crystallographic unit cells with cached derived quantities. @@ -33,6 +33,12 @@ class Cell(_NonModuleDeviceMixin): access and cached. The cache is cleared when the cell is moved to a different device or dtype. + Two cells compare and hash equal when their six parameters are equal + (:attr:`key`), independent of device and dtype, so a cell can key a dict or + a cache of quantities derived from it. A cell is therefore a value: editing + its parameter tensor in place is refused at the next derived read. Build a + new ``Cell`` and assign it instead. + Examples -------- >>> cell = Cell([50, 60, 70, 90, 90, 90]) @@ -45,6 +51,7 @@ class Cell(_NonModuleDeviceMixin): _data: torch.Tensor _cache: dict = field(default_factory=dict, repr=False) + _stamp: tuple = field(default=None, repr=False) def __init__( self, @@ -52,7 +59,6 @@ def __init__( *, dtype: torch.dtype = None, device: torch.device | str = None, - requires_grad: bool = False, ) -> None: """ Create a new Cell. @@ -66,8 +72,6 @@ def __init__( Desired data type. Defaults to the configured ``dtypes.float``. device : torch.device or str, optional Desired device. Defaults to the configured ``device.current``. - requires_grad : bool, optional - Whether to track gradients. Defaults to False. Raises ------ @@ -80,6 +84,10 @@ def __init__( # Convert to tensor first to get shape if isinstance(data, torch.Tensor): tensor = data.to(dtype=dtype, device=device) + if tensor is data: + # ``to`` returned the caller's tensor unchanged; own a copy so + # their later edits cannot reach into this cell. + tensor = tensor.clone() else: tensor = torch.tensor(data, dtype=dtype, device=device) @@ -93,11 +101,9 @@ def __init__( # Ensure 1D shape tensor = tensor.reshape(6) - if requires_grad: - tensor = tensor.requires_grad_(True) - object.__setattr__(self, "_data", tensor) object.__setattr__(self, "_cache", {}) + self._stamp_data() # ========================================================================= # Device/dtype movement methods @@ -108,23 +114,35 @@ def __init__( # and any cached tensor values) and then call ``reset_cache`` below. def reset_cache(self) -> None: - """Clear cached derived quantities (fractional matrix, volume, etc.).""" - object.__setattr__(self, "_cache", {}) - - def detach(self) -> "Cell": - """ - Return a new Cell with detached tensor (no gradient tracking). + """Clear cached derived quantities (fractional matrix, volume, etc.). - Returns - ------- - Cell - New Cell with detached data. + Also re-stamps the parameter tensor, so a device or dtype move (which + rebinds it and then calls this) is not mistaken for an in-place edit. """ - new_data = self._data.detach() - new_cell = Cell.__new__(Cell) - object.__setattr__(new_cell, "_data", new_data) - object.__setattr__(new_cell, "_cache", {}) - return new_cell + object.__setattr__(self, "_cache", {}) + self._stamp_data() + + def _stamp_data(self) -> None: + object.__setattr__(self, "_stamp", (id(self._data), self._data._version)) + + def _assert_unmodified(self) -> None: + """Refuse to serve derived quantities from a tensor edited in place.""" + stamp = getattr(self, "_stamp", None) + if stamp is None: + self._stamp_data() + return + if (id(self._data), self._data._version) == stamp: + return + cached = self._cache.get("key") + held = f" It held {cached}." if cached is not None else "" + raise RuntimeError( + "This Cell was edited in place after it was built." + held + " Cells are " + "values shared by reference (model context, structure-factor engine, " + "scaler), so an in-place edit changes the crystal under every holder. " + "Please don't edit Cell objects, create a new one -- " + "Cell([a, b, c, alpha, beta, gamma], dtype=cell.dtype, device=cell.device) " + "-- and assign it, e.g. model.cell = new_cell." + ) def clone(self) -> "Cell": """ @@ -139,8 +157,41 @@ def clone(self) -> "Cell": new_cell = Cell.__new__(Cell) object.__setattr__(new_cell, "_data", new_data) object.__setattr__(new_cell, "_cache", {}) + new_cell._stamp_data() return new_cell + # ========================================================================= + # Value identity + # ========================================================================= + + @property + def key(self) -> tuple: + """The six parameters as a tuple of Python floats. + + Read off the tensor once and cached alongside the derived quantities, so + the device synchronisation happens once per construction or + :meth:`reset_cache` rather than on every comparison. + + Returns + ------- + tuple of float + ``(a, b, c, alpha, beta, gamma)``. + """ + self._assert_unmodified() + key = self._cache.get("key") + if key is None: + key = tuple(float(v) for v in self._data.tolist()) + self._cache["key"] = key + return key + + def __hash__(self) -> int: + return hash(self.key) + + def __eq__(self, other) -> bool: + if not isinstance(other, Cell): + return NotImplemented + return self.key == other.key + # ========================================================================= # Basic properties # ========================================================================= @@ -160,11 +211,6 @@ def data(self) -> torch.Tensor: """Return the underlying tensor (for buffer registration).""" return self._data - @property - def requires_grad(self) -> bool: - """Return whether gradients are tracked.""" - return self._data.requires_grad - # ========================================================================= # Convenience properties for cell parameters # ========================================================================= @@ -218,6 +264,7 @@ def fractional_matrix(self) -> torch.Tensor: torch.Tensor Shape (3, 3) orthogonalization matrix. """ + self._assert_unmodified() if "fractional_matrix" not in self._cache: self._cache["fractional_matrix"] = self._compute_fractional_matrix() return self._cache["fractional_matrix"] @@ -234,6 +281,7 @@ def inv_fractional_matrix(self) -> torch.Tensor: torch.Tensor Shape (3, 3) fractionalization matrix. """ + self._assert_unmodified() if "inv_fractional_matrix" not in self._cache: self._cache["inv_fractional_matrix"] = torch.linalg.inv( self.fractional_matrix @@ -250,6 +298,7 @@ def volume(self) -> torch.Tensor: torch.Tensor Scalar tensor with the cell volume. """ + self._assert_unmodified() if "volume" not in self._cache: self._cache["volume"] = self._compute_volume() return self._cache["volume"] @@ -264,6 +313,7 @@ def reciprocal_basis_matrix(self) -> torch.Tensor: torch.Tensor Shape (3, 3) matrix where rows are the reciprocal basis vectors. """ + self._assert_unmodified() if "reciprocal_basis_matrix" not in self._cache: self._cache["reciprocal_basis_matrix"] = ( self._compute_reciprocal_basis_matrix() @@ -409,7 +459,3 @@ def __getitem__(self, idx: int) -> torch.Tensor: def __len__(self) -> int: """Return 6 (number of cell parameters).""" return 6 - - -# Keep CellTensor as an alias for backward compatibility -CellTensor = Cell diff --git a/torchref/symmetry/grid_utils.py b/torchref/symmetry/grid_utils.py deleted file mode 100644 index 902cf887..00000000 --- a/torchref/symmetry/grid_utils.py +++ /dev/null @@ -1,134 +0,0 @@ -"""FFT- and symmetry-compatible grid sizes. - -Interpolation-free symmetry expansion needs grid dimensions divisible by what the -screw axes demand, and radix-2,3,5 FFTs want factors of 2, 3, 5 only. - -These are thin wrappers over ``spacegroup``, which holds the canonical -implementations -- including its own ``is_fft_friendly`` / -``find_fft_friendly_size`` pair. Prefer ``spacegroup`` for new code. -""" - -import numpy as np -import torch - -from torchref.config import NYQUIST_OVERSAMPLING - - -def get_symmetry_grid_requirements(space_group: str) -> dict: - """Per-axis divisibility ``{'nx_mod', 'ny_mod', 'nz_mod'}`` for ``space_group``. - - Wrapper over :func:`~torchref.symmetry.spacegroup.get_grid_requirements`. - """ - # Import here to avoid circular imports - from torchref.symmetry.spacegroup import get_grid_requirements - - return get_grid_requirements(space_group) - - -def find_fft_friendly_size(n: int, divisibility: int = 1) -> int: - """Smallest size >= ``n`` factoring into 2, 3, 5 and divisible by ``divisibility``. - - Parameters - ---------- - n : int - Minimum grid size. - divisibility : int, default 1 - Required divisibility (e.g. 2 for a screw axis). - - Returns - ------- - int - Optimal grid size. - """ - candidate = n - - if candidate % divisibility != 0: - candidate = ((candidate // divisibility) + 1) * divisibility - - while not is_fft_friendly(candidate): - candidate += divisibility - - return candidate - - -def is_fft_friendly(n: int) -> bool: - """ - Check if a number has only factors of 2, 3, and 5. - - These are optimal for radix-2,3,5 FFT algorithms. - """ - if n <= 0: - return False - - # Remove all factors of 2, 3, 5 - while n % 2 == 0: - n //= 2 - while n % 3 == 0: - n //= 3 - while n % 5 == 0: - n //= 5 - - # If we're left with 1, the number is FFT-friendly - return n == 1 - - -def calculate_optimal_grid_size(cell_params, max_res: float, space_group: str) -> tuple: - """ - Optimal grid for a unit cell and space group. - - Satisfies Shannon-Nyquist sampling at - :data:`torchref.config.NYQUIST_OVERSAMPLING`, the screw-axis divisibility, and - FFT-friendliness (factors of 2, 3, 5 only). - - Parameters - ---------- - cell_params : array-like, shape (6,) - Unit cell [a, b, c, alpha, beta, gamma]. - max_res : float - Maximum resolution in Angstroms. - space_group : str - Space group symbol. - - Returns - ------- - tuple - Optimal grid dimensions (nx, ny, nz). - """ - # Import here to avoid circular imports - from torchref.symmetry.spacegroup import suggest_grid_size - - if isinstance(cell_params, torch.Tensor): - cell_params = cell_params.cpu().numpy() - - a, b, c = cell_params[:3] - - # Shannon-Nyquist: sample at NYQUIST_OVERSAMPLING × the maximum frequency - nx_min = int(np.floor(a / max_res * NYQUIST_OVERSAMPLING)) - ny_min = int(np.floor(b / max_res * NYQUIST_OVERSAMPLING)) - nz_min = int(np.floor(c / max_res * NYQUIST_OVERSAMPLING)) - - # Use spacegroup module to suggest optimal size - return suggest_grid_size((nx_min, ny_min, nz_min), space_group, make_fft_friendly=True) - - -def check_grid_compatibility(grid_shape: tuple, space_group: str) -> dict: - """Check ``(nx, ny, nz)`` against the space group symmetry and the FFT. - - Wrapper over - :func:`~torchref.symmetry.spacegroup.check_grid_compatibility`, which - documents the report dict. - """ - # Import here to avoid circular imports - from torchref.symmetry.spacegroup import ( - check_grid_compatibility as sg_check_grid_compatibility, - ) - - return sg_check_grid_compatibility(grid_shape, space_group) - - -def recommend_grid_size(current_shape: tuple, space_group: str) -> tuple: - """Smallest symmetry- and FFT-compatible grid at or above ``current_shape``.""" - # Import here to avoid circular imports - from torchref.symmetry.spacegroup import suggest_grid_size - - return suggest_grid_size(current_shape, space_group, make_fft_friendly=True) diff --git a/torchref/symmetry/map_symmetry.py b/torchref/symmetry/map_symmetry.py index a9d1852d..afbeb4a9 100644 --- a/torchref/symmetry/map_symmetry.py +++ b/torchref/symmetry/map_symmetry.py @@ -1,238 +1,250 @@ -"""Map-level symmetry operations for electron density maps. - -Applying symmetry to the map is far cheaper than generating symmetry mates per -atom. :func:`MapSymmetry` is a factory, not a class: it returns -:class:`MapSymmetryDirect` (exact integer indexing) when the grid allows, else the -interpolating implementation from ``map_symmetry_interpolation``. Space groups -accept strings, ints 1-230 or ``gemmi.SpaceGroup``. +"""Real-space map symmetrization, selected by grid compatibility. + +Two operators apply a :class:`~torchref.symmetry.symmetry.Symmetry` to a density map: +:class:`_MapSymmetryDirect` indexes symmetry mates at exact integers, and +:class:`~torchref.symmetry.map_symmetry_interpolation._MapSymmetryInterpolation` +falls back to ``grid_sample`` when the grid does not admit that. +:func:`build_map_operator` picks between them. + +Reach these through :meth:`~torchref.symmetry.symmetry.Symmetry.symmetrize_map` rather +than directly: it owns the caching, and it is the reason the choice of operator does +not leak into calling code. A grid that forces interpolation costs accuracy silently, +so ask :meth:`~torchref.symmetry.symmetry.Symmetry.can_index_directly` before +committing to a grid, and +:meth:`~torchref.symmetry.symmetry.Symmetry.suggest_grid_size` to fix one. + +Neither operator needs the unit cell: symmetry acts on fractional coordinates, so the +cell metric never enters. """ +from __future__ import annotations + import torch -import torch.nn as nn -from torchref.config import get_float_dtype, normalize_device -from torchref.symmetry.spacegroup import SpaceGroup, SpaceGroupLike from torchref.utils.device_mixin import DeviceMixin -def MapSymmetry( - space_group: SpaceGroupLike, - map_shape, - cell_params, - dtype_float=None, - verbose=1, - device=None, -): - """ - Build the appropriate MapSymmetry implementation for ``map_shape``. - - Returns :class:`MapSymmetryDirect` when the grid permits exact integer - indexing, otherwise the interpolating fallback -- so the return *type* - depends on the grid, and a mis-sized grid costs accuracy silently (at - ``verbose > 0`` a compatible grid is suggested). +def build_map_operator(symmetry, map_shape: tuple): + """Build the operator suited to ``map_shape``. Parameters ---------- - space_group : str, int, or gemmi.SpaceGroup - Space group specification (e.g., 'P21', 4, gemmi.SpaceGroup('P 21')). + symmetry : Symmetry + The group to apply. map_shape : tuple of int - Shape of the density map (nx, ny, nz). - cell_params : torch.Tensor, shape (6,) - Unit cell parameters [a, b, c, alpha, beta, gamma] in Å and degrees. - dtype_float : torch.dtype, optional - Floating point precision to use. Defaults to the configured - ``dtypes.float`` (``get_float_dtype()``, float32 in production). - verbose : int, default 1 - Verbosity level (0=silent, 1=info, 2=debug). - device : torch.device, default: configured device.current - Device to use for computation. + Density map dimensions ``(nx, ny, nz)``. Returns ------- - MapSymmetryDirect or MapSymmetryInterpolation - The appropriate implementation based on grid compatibility. + _MapSymmetryDirect or _MapSymmetryInterpolation + Direct integer indexing when the grid permits it, otherwise the interpolating + fallback. """ - if dtype_float is None: - dtype_float = get_float_dtype() - # ``cell_params`` is documented as a tensor, so follow it when no device is - # given rather than jumping to the global default and leaving the caller's - # cell behind. - if device is None and isinstance(cell_params, torch.Tensor): - device = cell_params.device - device = normalize_device(device) - symmetry = SpaceGroup(space_group, dtype=dtype_float, device=device) - compat = symmetry.check_grid_compatibility(map_shape) - - if compat["can_use_direct_indexing"]: - if verbose > 0: - print( - f"MapSymmetry: Using direct indexing (no interpolation) for {space_group}" - ) - return MapSymmetryDirect( - space_group, map_shape, cell_params, dtype_float, verbose, device - ) - else: - if verbose > 0: - print("MapSymmetry: Grid not compatible with direct indexing") - print(f" Using interpolation-based fallback for {space_group}") - if compat["issues"]: - for issue in compat["issues"]: - print(f" - {issue}") - suggested = symmetry.suggest_grid_size(map_shape, make_fft_friendly=True) - print(f" Suggested grid for direct indexing: {suggested}") - - from torchref.symmetry.map_symmetry_interpolation import ( - MapSymmetry as MapSymmetryInterpolation, - ) + if symmetry.can_index_directly(map_shape): + return _MapSymmetryDirect(symmetry, map_shape) - return MapSymmetryInterpolation( - space_group, map_shape, cell_params, dtype_float, verbose, device - ) + # Imported here, not at module scope: the interpolation module imports this one for + # the shared operator contract. + from torchref.symmetry.map_symmetry_interpolation import ( + _MapSymmetryInterpolation, + ) + return _MapSymmetryInterpolation(symmetry, map_shape) -class MapSymmetryDirect(DeviceMixin, nn.Module): - """ - Fast direct-indexing implementation of crystallographic symmetry operations. - Computes symmetry mates one operation at a time (streaming) so that - memory usage is O(grid) regardless of the number of symmetry operations, - rather than O(n_ops * grid) for storing precomputed index grids. +def _combine(mates: torch.Tensor, combine: str) -> torch.Tensor: + """Reduce stacked symmetry mates. - NOTE: Do not instantiate this class directly. Use the MapSymmetry() factory - function instead, which will automatically select the appropriate implementation. + Parameters + ---------- + mates : torch.Tensor + Stacked mates, shape ``(n_ops, nx, ny, nz)``. + combine : {'sum', 'max'} + ``'sum'`` for electron density, ``'max'`` for masks and boolean data. + + Returns + ------- + torch.Tensor + Shape ``(nx, ny, nz)``. + + Raises + ------ + ValueError + For an unknown mode. """ + if combine == "sum": + return mates.sum(dim=0) + if combine == "max": + return mates.max(dim=0)[0] + raise ValueError(f"Unknown combine mode: {combine}. Use 'sum' or 'max'.") - def __init__( - self, - space_group, - map_shape, - cell_params, - dtype_float=None, - verbose=1, - device=None, - ): - super().__init__() - if dtype_float is None: - dtype_float = get_float_dtype() - if device is None and isinstance(cell_params, torch.Tensor): - device = cell_params.device - self.dtype_float = dtype_float - self.space_group = space_group - self.map_shape = tuple(map_shape) - self.verbose = verbose - self.device = normalize_device(device) - # Coerce ``cell_params`` onto this module's device: it is a plain attribute, - # so ``DeviceMixin`` cannot see it and would never repair a mismatch. - if isinstance(cell_params, torch.Tensor): - cell_params = cell_params.to(device=self.device, dtype=self.dtype_float) - else: - cell_params = torch.as_tensor( - cell_params, device=self.device, dtype=self.dtype_float - ) - self.cell_params = cell_params - self.symmetry = SpaceGroup( - space_group, dtype=self.dtype_float, device=self.device - ) - self.n_ops = self.symmetry.matrices.shape[0] - self.can_use_direct_indexing = True +class _MapSymmetryDirect(DeviceMixin): + """Symmetrize maps by exact integer indexing, one operation at a time. - if self.verbose > 0: - print(f"MapSymmetryDirect initialized for {space_group}") - print(f" Number of symmetry operations: {self.n_ops}") - print(f" Map shape: {self.map_shape}") + Valid only on a grid whose dimensions satisfy the group's divisibility, so every + symmetry mate falls on a grid point. :func:`build_map_operator` enforces that. + + Parameters + ---------- + symmetry : Symmetry + The group to apply. + map_shape : tuple of int + Density map dimensions ``(nx, ny, nz)``. - # ------------------------------------------------------------------ - # Core: compute index grid for a single symmetry operation - # ------------------------------------------------------------------ + Notes + ----- + Holds no precomputed grids. Index grids are recomputed per operation so peak memory + stays at one index grid plus two density maps, rather than scaling with the number + of operations. + """ - def _compute_index_grid(self, op_index: int) -> torch.Tensor: - """Integer index grid (nx, ny, nz, 3) int64 for one op; not cached.""" + def __init__(self, symmetry, map_shape: tuple): + self.symmetry = symmetry + self.map_shape = tuple(int(n) for n in map_shape) + + @property + def n_ops(self) -> int: + """Number of symmetry operations.""" + return self.symmetry.n_ops + + @property + def device(self) -> torch.device: + """Device the operations live on.""" + return self.symmetry.device + + def _index_grid(self, op_index: int) -> torch.Tensor: + """Integer index grid for one operation. + + Parameters + ---------- + op_index : int + Operation index. + + Returns + ------- + torch.Tensor + Shape ``(nx, ny, nz, 3)``, dtype ``int64``. + + Notes + ----- + Deliberately does not use the batched + :meth:`~torchref.symmetry.symmetry.Symmetry.expand_positions`: that would + transform every operation at once, which is exactly the O(n_ops * grid) memory + this operator exists to avoid. + """ nx, ny, nz = self.map_shape - device = self.symmetry.matrices.device + symmetry = self.symmetry + dtype = symmetry.dtype + device = symmetry.device - fx = torch.arange(nx, dtype=self.dtype_float, device=device) / nx - fy = torch.arange(ny, dtype=self.dtype_float, device=device) / ny - fz = torch.arange(nz, dtype=self.dtype_float, device=device) / nz + fx = torch.arange(nx, dtype=dtype, device=device) / nx + fy = torch.arange(ny, dtype=dtype, device=device) / ny + fz = torch.arange(nz, dtype=dtype, device=device) / nz gx, gy, gz = torch.meshgrid(fx, fy, fz, indexing="ij") grid_flat = torch.stack([gx, gy, gz], dim=-1).reshape(-1, 3) - transformed = torch.matmul(self.symmetry.matrices[op_index], grid_flat.T).T - transformed = transformed + self.symmetry.translations[op_index] + transformed = torch.matmul(symmetry.matrices[op_index], grid_flat.T).T + transformed = transformed + symmetry.translations[op_index] transformed = transformed - torch.floor(transformed) - shape_t = torch.tensor([nx, ny, nz], dtype=self.dtype_float, device=device) - indices = torch.round(transformed * shape_t).to(torch.int64) + shape_t = torch.tensor([nx, ny, nz], dtype=dtype, device=device) + indices = torch.round(transformed * shape_t).to(torch.int64) # dtype-ok: rounded voxel grid indices; int64 index tensor required indices[:, 0] %= nx indices[:, 1] %= ny indices[:, 2] %= nz return indices.reshape(nx, ny, nz, 3) - # ------------------------------------------------------------------ - # Public API - # ------------------------------------------------------------------ - - def get_symmetry_mate(self, density_map, operation_index): - """Apply a single symmetry operation via direct indexing.""" - if operation_index < 0 or operation_index >= self.n_ops: + def _check_shape(self, density_map: torch.Tensor) -> None: + """Reject a map whose shape this operator was not built for.""" + if tuple(density_map.shape) != self.map_shape: raise ValueError( - f"Operation index {operation_index} out of range [0, {self.n_ops-1}]" + f"Map shape {tuple(density_map.shape)} does not match the operator's " + f"{self.map_shape}" ) - if density_map.shape != self.map_shape: - raise ValueError( - f"Map shape {density_map.shape} doesn't match expected {self.map_shape}" - ) - ig = self._compute_index_grid(operation_index) - return density_map[ig[..., 0], ig[..., 1], ig[..., 2]] - def forward(self, density_map, apply_symmetry=True, combine_mode="sum"): - """Apply symmetry operations to density map. + def mate(self, density_map: torch.Tensor, op_index: int) -> torch.Tensor: + """One symmetry mate of ``density_map``. - Computes one symmetry mate at a time and accumulates into the - result, so peak memory is only 1 index grid + 2 density maps - regardless of the number of symmetry operations. + Parameters + ---------- + density_map : torch.Tensor + Density, shape ``(nx, ny, nz)``. + op_index : int + Operation index in ``[0, n_ops)``. + + Returns + ------- + torch.Tensor + Shape ``(nx, ny, nz)``. """ - if not apply_symmetry or self.n_ops == 1: - return density_map - - ig = self._compute_index_grid(0) - if combine_mode == "sum": - result = density_map[ig[..., 0], ig[..., 1], ig[..., 2]] - for i in range(1, self.n_ops): - ig = self._compute_index_grid(i) - result = result + density_map[ig[..., 0], ig[..., 1], ig[..., 2]] - elif combine_mode == "max": - result = density_map[ig[..., 0], ig[..., 1], ig[..., 2]] - for i in range(1, self.n_ops): - ig = self._compute_index_grid(i) - result = torch.max( - result, density_map[ig[..., 0], ig[..., 1], ig[..., 2]] - ) - else: + if op_index < 0 or op_index >= self.n_ops: raise ValueError( - f"Unknown combine_mode: {combine_mode}. Use 'sum' or 'max'." + f"Operation index {op_index} out of range [0, {self.n_ops - 1}]" ) + self._check_shape(density_map) + ig = self._index_grid(op_index) + return density_map[ig[..., 0], ig[..., 1], ig[..., 2]] - return result + def all_mates(self, density_map: torch.Tensor) -> torch.Tensor: + """Every symmetry mate, stacked. + + Parameters + ---------- + density_map : torch.Tensor + Density, shape ``(nx, ny, nz)``. - def __call__(self, density_map, apply_symmetry=True, combine_mode="sum"): - """Make the class callable like a PyTorch module.""" - return self.forward( - density_map, apply_symmetry=apply_symmetry, combine_mode=combine_mode + Returns + ------- + torch.Tensor + Shape ``(n_ops, nx, ny, nz)``. + """ + self._check_shape(density_map) + return torch.stack( + [self.mate(density_map, i) for i in range(self.n_ops)], dim=0 ) - def get_symmetry_info(self): - """Get information about symmetry operations.""" - return { - "space_group": self.space_group, - "n_operations": self.n_ops, - "matrices": self.symmetry.matrices, - "translations": self.symmetry.translations, - } + def symmetrize( + self, density_map: torch.Tensor, combine: str = "sum" + ) -> torch.Tensor: + """Apply every operation and reduce the mates. + + Parameters + ---------- + density_map : torch.Tensor + Density, shape ``(nx, ny, nz)``. + combine : {'sum', 'max'}, default 'sum' + Reduction across mates. + + Returns + ------- + torch.Tensor + Shape ``(nx, ny, nz)``. + + Notes + ----- + Accumulates one mate at a time instead of stacking, keeping peak memory + independent of the operation count. + """ + self._check_shape(density_map) + result = self.mate(density_map, 0) + for i in range(1, self.n_ops): + mate = self.mate(density_map, i) + if combine == "sum": + result = result + mate + elif combine == "max": + result = torch.max(result, mate) + else: + raise ValueError( + f"Unknown combine mode: {combine}. Use 'sum' or 'max'." + ) + return result - def __repr__(self): + def __repr__(self) -> str: return ( - f"MapSymmetryDirect(space_group='{self.space_group}', " - f"n_ops={self.n_ops}, map_shape={self.map_shape})" + f"_MapSymmetryDirect(n_ops={self.n_ops}, map_shape={self.map_shape})" ) + + +__all__ = ["build_map_operator"] diff --git a/torchref/symmetry/map_symmetry_interpolation.py b/torchref/symmetry/map_symmetry_interpolation.py index 0e3ae55a..16fb7171 100644 --- a/torchref/symmetry/map_symmetry_interpolation.py +++ b/torchref/symmetry/map_symmetry_interpolation.py @@ -1,299 +1,173 @@ -""" -Map-level symmetry operations for electron density maps. +"""Map symmetrization by trilinear interpolation, for grids that need it. + +The fallback behind :func:`~torchref.symmetry.map_symmetry.build_map_operator` when a +grid does not satisfy the group's divisibility, so symmetry mates land between grid +points. Interpolating costs accuracy that exact indexing does not, which is why +:meth:`~torchref.symmetry.symmetry.Symmetry.suggest_grid_size` exists -- prefer fixing +the grid over landing here. -This module provides efficient symmetry operations applied directly to density maps, -which is much faster than applying symmetry to individual atoms. +Reach this through :meth:`~torchref.symmetry.symmetry.Symmetry.symmetrize_map`. """ -import numpy as np +from __future__ import annotations + import torch -import torch.nn as nn import torch.nn.functional as F -from torchref.config import get_float_dtype, normalize_device -from torchref.symmetry.spacegroup import SpaceGroup +from torchref.symmetry.map_symmetry import _combine from torchref.utils.device_mixin import DeviceMixin -class MapSymmetry(DeviceMixin, nn.Module): - """ - Applies crystallographic symmetry operations to electron density maps. +class _MapSymmetryInterpolation(DeviceMixin): + """Symmetrize maps by resampling with ``grid_sample``. - Takes an asymmetric-unit density map, applies each operation in fractional - coordinates, interpolates via ``grid_sample`` and sums the mates -- cheaper - than generating symmetry mates per atom and recalculating density. The - per-operation sampling grids are precomputed in ``__init__`` for the fixed - ``map_shape``. - - Attributes + Parameters ---------- - space_group : str - Space group name. - map_shape : tuple of int - Shape of the density map (nx, ny, nz). - cell_params : numpy.ndarray - Unit cell parameters. symmetry : Symmetry - Symmetry operations handler. - n_ops : int - Number of symmetry operations. - - Examples - -------- - :: - - map_sym = MapSymmetry(space_group='P21', map_shape=(64, 64, 64), cell_params=cell) - asymmetric_map = model.build_density_map() - symmetric_map = map_sym(asymmetric_map) + The group to apply. + map_shape : tuple of int + Density map dimensions ``(nx, ny, nz)``. + + Notes + ----- + Precomputes one sampling grid per operation, shape + ``(n_ops, nx, ny, nz, 3)`` -- hundreds of megabytes at production grid sizes. That + is why :class:`~torchref.symmetry.symmetry.Symmetry` memoizes only the most recent + shape and drops the cache on any device move. """ - def __init__( - self, - space_group, - map_shape, - cell_params, - dtype_float=None, - verbose=1, - device=None, - ): - """ - Initialize map symmetry operator. + def __init__(self, symmetry, map_shape: tuple): + self.symmetry = symmetry + self.map_shape = tuple(int(n) for n in map_shape) + self.sampling_grids = self._build_sampling_grids() - Parameters - ---------- - space_group : str - Space group name (e.g., 'P1', 'P21', 'P-1', etc.). - map_shape : tuple of int - Shape of the density map (nx, ny, nz). - cell_params : array-like, shape (6,) - Unit cell parameters [a, b, c, alpha, beta, gamma] in Å and degrees. - dtype_float : torch.dtype, optional - Floating point precision to use. Defaults to the configured - ``dtypes.float`` (``get_float_dtype()``, float32 in production). - verbose : int, default 1 - Verbosity level. - device : torch.device, default: configured device.current - Device to use for computation. - """ - super().__init__() - if dtype_float is None: - dtype_float = get_float_dtype() - device = normalize_device(device) - self.dtype_float = dtype_float - self.space_group = space_group - self.map_shape = tuple(map_shape) - self.cell_params = np.array(cell_params) - self.verbose = verbose - self.device = device - self.symmetry = SpaceGroup( - space_group, dtype=self.dtype_float, device=self.device - ) - self.n_ops = self.symmetry.matrices.shape[0] - if self.verbose > 0: - print(f"MapSymmetry initialized for {space_group}") - print(f" Number of symmetry operations: {self.n_ops}") - print(f" Map shape: {self.map_shape}") + @property + def n_ops(self) -> int: + """Number of symmetry operations.""" + return self.symmetry.n_ops - self._setup_fractional_grid() - self._setup_symmetry_grids() + @property + def device(self) -> torch.device: + """Device the sampling grids live on.""" + return self.sampling_grids.device - def _setup_fractional_grid(self): - """Fractional grid with voxels at edges i/N (CCTBX/gemmi convention).""" - nx, ny, nz = self.map_shape - - fx = torch.arange(nx, dtype=self.dtype_float, device=self.device) / nx - fy = torch.arange(ny, dtype=self.dtype_float, device=self.device) / ny - fz = torch.arange(nz, dtype=self.dtype_float, device=self.device) / nz - - # indexing='ij' so the result is (nx, ny, nz, 3) with the last dim [fx, fy, fz]. - grid_fx, grid_fy, grid_fz = torch.meshgrid(fx, fy, fz, indexing="ij") - grid_frac = torch.stack([grid_fx, grid_fy, grid_fz], dim=-1) - - self.register_buffer("grid_frac", grid_frac) + def _build_sampling_grids(self) -> torch.Tensor: + """Precompute per-operation ``grid_sample`` coordinates in ``[-1, 1]``. - def _setup_symmetry_grids(self): - """Precompute per-operation ``grid_sample`` coordinates in [-1, 1].""" + Returns + ------- + torch.Tensor + Shape ``(n_ops, nx, ny, nz, 3)``. + """ nx, ny, nz = self.map_shape - - grid_flat = self.grid_frac.reshape(-1, 3) - - sampling_grids_list = [] - - for i in range(self.n_ops): - # R @ coords + t, on (3, nx*ny*nz) - transformed = torch.matmul(self.symmetry.matrices[i], grid_flat.T) - transformed = transformed.T # (nx*ny*nz, 3) - transformed = transformed + self.symmetry.translations[i] - - # Wrap to [0, 1) for periodic boundary conditions - transformed = transformed - torch.floor(transformed) - grid_shape_tensor = torch.tensor( - [nx, ny, nz], dtype=self.dtype_float, device=transformed.device - ) - # grid_coord = -1 + 2*N/(N-1) * frac, per dimension - sampling_coords = ( - -1.0 + 2.0 * grid_shape_tensor / (grid_shape_tensor - 1.0) * transformed + symmetry = self.symmetry + dtype = symmetry.dtype + device = symmetry.device + + # Voxels at fractional edges i/N, the CCTBX/gemmi convention. + fx = torch.arange(nx, dtype=dtype, device=device) / nx + fy = torch.arange(ny, dtype=dtype, device=device) / ny + fz = torch.arange(nz, dtype=dtype, device=device) / nz + gx, gy, gz = torch.meshgrid(fx, fy, fz, indexing="ij") + grid_flat = torch.stack([gx, gy, gz], dim=-1).reshape(-1, 3) + + transformed = symmetry.expand_positions(grid_flat) # (n_ops, N, 3) + # Wrap into [0, 1) for periodic boundaries. + transformed = transformed - torch.floor(transformed) + + shape_t = torch.tensor([nx, ny, nz], dtype=dtype, device=device) + # grid_coord = -1 + 2*N/(N-1) * frac, per dimension. + sampling = -1.0 + 2.0 * shape_t / (shape_t - 1.0) * transformed + sampling = sampling.reshape(self.n_ops, nx, ny, nz, 3) + + # grid_sample reads the last axis as [x, y, z] -> [W, H, D], i.e. the REVERSE + # of our [fx, fy, fz] -> [D, H, W]. Dropping this reorder still interpolates, + # silently against the wrong axes. + return sampling[..., [2, 1, 0]].contiguous() + + def _check_shape(self, density_map: torch.Tensor) -> None: + """Reject a map whose shape this operator was not built for.""" + if tuple(density_map.shape) != self.map_shape: + raise ValueError( + f"Map shape {tuple(density_map.shape)} does not match the operator's " + f"{self.map_shape}" ) - sampling_grid = sampling_coords.reshape(nx, ny, nz, 3) - - # grid_sample reads the last axis as [x, y, z] -> [W, H, D], i.e. the - # REVERSE of our [fx, fy, fz] -> [D, H, W]. Dropping this reorder still - # interpolates, silently against the wrong axes. - sampling_grid = sampling_grid[ - ..., [2, 1, 0] - ] # [fx, fy, fz] -> [fz, fy, fx] - - sampling_grids_list.append(sampling_grid) - - sampling_grids_stacked = torch.stack(sampling_grids_list, dim=0) - - self.register_buffer("sampling_grids", sampling_grids_stacked) - - def get_symmetry_mate(self, density_map, operation_index): - """ - Apply a single symmetry operation to get one symmetry mate. + def mate(self, density_map: torch.Tensor, op_index: int) -> torch.Tensor: + """One symmetry mate of ``density_map``. Parameters ---------- - density_map : torch.Tensor, shape (nx, ny, nz) - Electron density map (typically from asymmetric unit). - operation_index : int - Index of the symmetry operation to apply (0 to n_ops-1). + density_map : torch.Tensor + Density, shape ``(nx, ny, nz)``. + op_index : int + Operation index in ``[0, n_ops)``. Returns ------- - torch.Tensor, shape (nx, ny, nz) - Density map after applying the symmetry operation. + torch.Tensor + Shape ``(nx, ny, nz)``. """ - if operation_index < 0 or operation_index >= self.n_ops: - raise ValueError( - f"Operation index {operation_index} out of range [0, {self.n_ops-1}]" - ) - - # Ensure map is correct shape - if density_map.shape != self.map_shape: + if op_index < 0 or op_index >= self.n_ops: raise ValueError( - f"Map shape {density_map.shape} doesn't match expected {self.map_shape}" + f"Operation index {op_index} out of range [0, {self.n_ops - 1}]" ) - - # Prepare for grid_sample - map_5d = density_map.unsqueeze(0).unsqueeze(0) # (1, 1, nx, ny, nz) - - # Get sampling grid for this operation - sampling_grid = self.sampling_grids[operation_index] - sampling_grid_batch = sampling_grid.unsqueeze(0) - - # Interpolate map at transformed coordinates - # align_corners=True ensures that: - # -1 maps to index 0 (fractional coord 0) - # +1 maps to index N-1 (fractional coord (N-1)/N) - # This matches the grid-edge convention (voxels at i/N) - # padding_mode='border' handles periodic boundary conditions via the wrapping - # we did in _setup_symmetry_grids - transformed_map = F.grid_sample( - map_5d, - sampling_grid_batch, - mode="bilinear", # Trilinear interpolation for 3D - padding_mode="border", # Use border mode since we pre-wrapped coordinates - align_corners=True, # Critical: matches grid-edge convention + self._check_shape(density_map) + + # align_corners=True maps -1 to index 0 and +1 to index N-1, matching the + # grid-edge convention above; padding_mode='border' is safe only because + # _build_sampling_grids already wrapped the coordinates. + transformed = F.grid_sample( + density_map.unsqueeze(0).unsqueeze(0), + self.sampling_grids[op_index].unsqueeze(0), + mode="bilinear", + padding_mode="border", + align_corners=True, ) + return transformed.squeeze(0).squeeze(0) - # Remove batch and channel dimensions - transformed_map = transformed_map.squeeze(0).squeeze(0) - - return transformed_map - - def get_all_symmetry_mates(self, density_map): - """ - Get all symmetry mates as a list. + def all_mates(self, density_map: torch.Tensor) -> torch.Tensor: + """Every symmetry mate, stacked. Parameters ---------- - density_map : torch.Tensor, shape (nx, ny, nz) - Electron density map (typically from asymmetric unit). + density_map : torch.Tensor + Density, shape ``(nx, ny, nz)``. Returns ------- - list of torch.Tensor - List of symmetry-related maps, one for each operation. + torch.Tensor + Shape ``(n_ops, nx, ny, nz)``. """ - mates = [] - for i in range(self.n_ops): - mates.append(self.get_symmetry_mate(density_map, i)) - return mates + self._check_shape(density_map) + return torch.stack( + [self.mate(density_map, i) for i in range(self.n_ops)], dim=0 + ) - def forward(self, density_map, apply_symmetry=True, combine_mode="sum"): - """ - Apply symmetry operations to density map. + def symmetrize( + self, density_map: torch.Tensor, combine: str = "sum" + ) -> torch.Tensor: + """Apply every operation and reduce the mates. Parameters ---------- - density_map : torch.Tensor, shape (nx, ny, nz) - Electron density map (typically from asymmetric unit). - apply_symmetry : bool, default True - If True, apply all symmetry operations and combine them. - If False, return input map unchanged (useful for P1 or debugging). - combine_mode : str, default 'sum' - How to combine symmetry mates: - - - 'sum': Sum all symmetry mates (for electron density) - - 'max': Take maximum across symmetry mates (for masks/boolean data) + density_map : torch.Tensor + Density, shape ``(nx, ny, nz)``. + combine : {'sum', 'max'}, default 'sum' + Reduction across mates. Returns ------- - torch.Tensor, shape (nx, ny, nz) - Symmetry-expanded density map (combined symmetry mates). + torch.Tensor + Shape ``(nx, ny, nz)``. """ - if not apply_symmetry or self.n_ops == 1: - # No symmetry or P1 - return density_map - - # Get all symmetry mates - mates = self.get_all_symmetry_mates(density_map) - mates_stacked = torch.stack(mates, dim=0) - - # Combine according to mode - if combine_mode == "sum": - symmetric_map = mates_stacked.sum(dim=0) - elif combine_mode == "max": - symmetric_map = mates_stacked.max(dim=0)[0] # max returns (values, indices) - else: - raise ValueError( - f"Unknown combine_mode: {combine_mode}. Use 'sum' or 'max'." - ) + return _combine(self.all_mates(density_map), combine) - return symmetric_map - - def __call__(self, density_map, apply_symmetry=True, combine_mode="sum"): - """Make the class callable like a PyTorch module.""" - return self.forward( - density_map, apply_symmetry=apply_symmetry, combine_mode=combine_mode + def __repr__(self) -> str: + return ( + f"_MapSymmetryInterpolation(n_ops={self.n_ops}, " + f"map_shape={self.map_shape})" ) - def get_symmetry_info(self): - """ - Get information about symmetry operations. - Returns - ------- - dict - Dictionary with the following keys: - - - 'space_group' : str - - 'n_operations' : int - - 'matrices' : torch.Tensor, shape (n_ops, 3, 3) - - 'translations' : torch.Tensor, shape (n_ops, 3) - """ - return { - "space_group": self.space_group, - "n_operations": self.n_ops, - "matrices": self.symmetry.matrices, - "translations": self.symmetry.translations, - } - - def __repr__(self): - return ( - f"MapSymmetry(space_group='{self.space_group}', " - f"n_ops={self.n_ops}, map_shape={self.map_shape})" - ) +__all__ = ["_MapSymmetryInterpolation"] diff --git a/torchref/symmetry/reciprocal_symmetry.py b/torchref/symmetry/reciprocal_symmetry.py index 15ad80c2..f595ee70 100644 --- a/torchref/symmetry/reciprocal_symmetry.py +++ b/torchref/symmetry/reciprocal_symmetry.py @@ -1,630 +1,36 @@ -"""Reciprocal space symmetry operations for structure factor grids. - -The reciprocal-space counterpart to ``map_symmetry.py``: :func:`ReciprocalSymmetry` -(grid operator), :func:`expand_hkl` / :func:`expand_reflections` / -:func:`expand_reciprocal_grid` (ASU -> P1), :func:`reduce_hkl` (P1 -> ASU), -:func:`complete_hkl` (find reflections missing from a dataset, same space group) -and :func:`canonicalize_hkl` (CCP4 ASU representative). Space groups accept -strings, ints 1-230 or ``gemmi.SpaceGroup``. - -Miller indices transform as h' = h @ R = R^T @ h with R the *real-space* rotation; -``reciprocal_matrices`` already holds the transpose, so do not transpose again. -Translations become phase shifts of -2π h·t; the sign is load-bearing and wrong -signs are invisible in P21/P212121/C2 (see :func:`expand_hkl`). +"""Asymmetric-unit conventions for Miller indices. + +The algorithms behind :class:`~torchref.symmetry.spacegroup.SpaceGroup`'s HKL verbs: +``expand_hkl`` (ASU -> P1), ``reduce_hkl`` (P1 -> ASU), ``complete_hkl`` (reflections +missing from a dataset, same space group) and ``canonicalize_hkl`` (CCP4 ASU +representative). All private -- call them through the space group, which is the only +public entry point. + +What makes these crystallographic rather than general symmetry is the choice of +asymmetric unit: the CCP4 convention, read off gemmi's ``ReciprocalAsu`` and keyed by +Laue class. That is why they hang off +:class:`~torchref.symmetry.spacegroup.SpaceGroup` and not +:class:`~torchref.symmetry.symmetry.Symmetry`. + +Miller indices transform as ``h' = h @ R = R^T @ h`` with R the *real-space* rotation; +:attr:`~torchref.symmetry.symmetry.Symmetry.reciprocal` already holds the transpose. +Translations enter as phase shifts of ``-2 pi h.t``. That sign is load-bearing and a +wrong one is invisible in P21/P212121/C2 -- see :func:`_expand_hkl` and +``tests/unit/symmetry/test_phase_convention.py``. """ -from typing import TYPE_CHECKING, Optional, Tuple +from typing import Optional, Tuple import numpy as np import torch -import torch.nn as nn -from torchref.config import get_float_dtype, normalize_device -from torchref.symmetry.spacegroup import SpaceGroup, SpaceGroupLike -from torchref.utils.device_mixin import DeviceMixin +from torchref.config import get_float_dtype -if TYPE_CHECKING: - from torchref.io.datasets.reflection_data import ReflectionData -def ReciprocalSymmetry( - space_group: SpaceGroupLike, - grid_shape, - dtype_float=None, - verbose=1, - device=None, -): - """ - Factory function to create the appropriate ReciprocalSymmetry implementation. - - Parameters - ---------- - space_group : str, int, or gemmi.SpaceGroup - Space group specification (e.g., 'P21', 4, gemmi.SpaceGroup('P 21')). - grid_shape : tuple of int - Shape of the reciprocal space grid (nh, nk, nl). - The grid spans from -n//2 to n//2 for each dimension. - dtype_float : torch.dtype, default: configured dtypes.float - Floating point precision to use. - verbose : int, default 1 - Verbosity level (0=silent, 1=info, 2=debug). - device : torch.device, default: configured device.current - Device to use for computation. - - Returns - ------- - ReciprocalSymmetryGrid - Implementation for reciprocal space grid symmetry operations. - """ - return ReciprocalSymmetryGrid(space_group, grid_shape, dtype_float, verbose, device) - - -class ReciprocalSymmetryGrid(DeviceMixin, nn.Module): - """ - Reciprocal space symmetry operations for Miller index grids. - - Covers Miller-index transformation, systematic absences, centric reflections, - Friedel pairs and symmetry expansion/averaging of structure factors. The - per-operation index maps, phase shifts, absence mask and centric mask are all - precomputed in ``__init__`` for the fixed ``grid_shape``. - - Attributes - ---------- - space_group : str - Space group name. - grid_shape : tuple of int - Shape of the reciprocal space grid (nh, nk, nl). - symmetry : SpaceGroup - Base symmetry operations handler. - n_ops : int - Number of symmetry operations. - - Examples - -------- - :: - - recip_sym = ReciprocalSymmetry('P21', grid_shape=(64, 64, 64)) - F_expanded = recip_sym(F_asym) # Expand from asymmetric unit - F_avg = recip_sym.symmetry_average(F_full) # Average symmetry-related reflections - """ - - def __init__( - self, - space_group, - grid_shape, - dtype_float=None, - verbose=1, - device=None, - ): - """ - Initialize reciprocal space symmetry operator. - - Parameters - ---------- - space_group : str - Space group name. - grid_shape : tuple of int - Shape of the reciprocal space grid (nh, nk, nl). - dtype_float : torch.dtype, default: configured dtypes.float - Floating point precision. - verbose : int, default 1 - Verbosity level. - device : torch.device, default: configured device.current - Computation device. - """ - super().__init__() - if dtype_float is None: - dtype_float = get_float_dtype() - device = normalize_device(device) - self.dtype_float = dtype_float - self.space_group = space_group - self.grid_shape = tuple(grid_shape) - self.verbose = verbose - self.device = device - - self.symmetry = SpaceGroup(space_group, dtype=dtype_float, device=device) - self.n_ops = self.symmetry.matrices.shape[0] - - self._setup_reciprocal_matrices() - - if self.verbose > 0: - print(f"ReciprocalSymmetryGrid initialized for {space_group}") - print(f" Number of symmetry operations: {self.n_ops}") - print(f" Grid shape: {self.grid_shape}") - - self._setup_hkl_grid() - self._setup_symmetry_index_grids() - self._setup_systematic_absences() - self._setup_centric_reflections() - - if self.verbose > 0: - n_absent = self.systematic_absences.sum().item() - n_centric = self.centric_mask.sum().item() - n_total = np.prod(self.grid_shape) - print(f" Systematic absences: {n_absent} ({100*n_absent/n_total:.2f}%)") - print(f" Centric reflections: {n_centric} ({100*n_centric/n_total:.2f}%)") - - def _setup_reciprocal_matrices(self): - """Cache ``reciprocal_matrices`` = R^T, for h' = R^T @ h.""" - # Translations cause phase shifts, not index changes, so they are not folded in. - recip_matrices = self.symmetry.matrices.transpose(-2, -1).contiguous() - self.register_buffer("reciprocal_matrices", recip_matrices) - - def _setup_hkl_grid(self): - """Build ``hkl_grid`` in FFT index order (0...n//2, -n//2+1...-1).""" - nh, nk, nl = self.grid_shape - - h = torch.fft.fftfreq(nh, d=1.0) * nh # gives 0,1,2,...,n//2,-n//2+1,...,-1 - k = torch.fft.fftfreq(nk, d=1.0) * nk - l = torch.fft.fftfreq(nl, d=1.0) * nl - - h = h.to(dtype=torch.int64, device=self.device) - k = k.to(dtype=torch.int64, device=self.device) - l = l.to(dtype=torch.int64, device=self.device) - - grid_h, grid_k, grid_l = torch.meshgrid(h, k, l, indexing="ij") - hkl_grid = torch.stack([grid_h, grid_k, grid_l], dim=-1) - - self.register_buffer("hkl_grid", hkl_grid) - - # Float copy for the matmuls; the int copy stays authoritative for indexing. - hkl_grid_float = hkl_grid.to(dtype=self.dtype_float) - self.register_buffer("hkl_grid_float", hkl_grid_float) - - def _setup_symmetry_index_grids(self): - """Precompute ``index_grids`` (n_ops, nh, nk, nl, 3) and ``phase_shifts``. - - The grid path uses the +2π h·t convention, unlike :func:`expand_hkl`. - """ - nh, nk, nl = self.grid_shape - hkl_flat = self.hkl_grid_float.reshape(-1, 3) # (N, 3) - - index_grids_list = [] - phase_shift_grids_list = [] - - for i in range(self.n_ops): - transformed = torch.matmul(hkl_flat, self.reciprocal_matrices[i].T) - # Exact for valid symmetry ops; round only mops up float error. - transformed_int = torch.round(transformed).to(torch.int64) - - # F(h') = F(h) * exp(2πi h·t) - translation = self.symmetry.translations[i] - phase_shift = 2.0 * np.pi * torch.matmul(hkl_flat, translation) - phase_shift = phase_shift.reshape(nh, nk, nl) - phase_shift_grids_list.append(phase_shift) - - # Wrap with periodic boundary to get grid indices. - idx_h = transformed_int[:, 0] % nh - idx_k = transformed_int[:, 1] % nk - idx_l = transformed_int[:, 2] % nl - - index_grid = torch.stack([idx_h, idx_k, idx_l], dim=-1) - index_grid = index_grid.reshape(nh, nk, nl, 3) - index_grids_list.append(index_grid) - - index_grids = torch.stack(index_grids_list, dim=0) - self.register_buffer("index_grids", index_grids) - - phase_shifts = torch.stack(phase_shift_grids_list, dim=0) - self.register_buffer("phase_shifts", phase_shifts) - - def _setup_systematic_absences(self): - """Mask reflections with some op mapping h -> h at h·t not integral. - - Such a reflection is destroyed by interference from the translation. - """ - nh, nk, nl = self.grid_shape - absences = torch.zeros(self.grid_shape, dtype=torch.bool, device=self.device) - - hkl_flat = self.hkl_grid_float.reshape(-1, 3) - - for i in range(self.n_ops): - transformed = torch.matmul(hkl_flat, self.reciprocal_matrices[i].T) - transformed_int = torch.round(transformed).to(torch.int64) - - hkl_int = self.hkl_grid.reshape(-1, 3) - same_reflection = (transformed_int == hkl_int).all(dim=-1) - - translation = self.symmetry.translations[i] - h_dot_t = torch.matmul(hkl_flat, translation) - - phase_mod = torch.abs(h_dot_t - torch.round(h_dot_t)) - non_zero_phase = phase_mod > 1e-6 - - absent_mask = (same_reflection & non_zero_phase).reshape(self.grid_shape) - absences = absences | absent_mask - - self.register_buffer("systematic_absences", absences) - - def _setup_centric_reflections(self): - """Mask reflections with some op mapping h -> -h (phase restricted to 0/π).""" - nh, nk, nl = self.grid_shape - centric = torch.zeros(self.grid_shape, dtype=torch.bool, device=self.device) - - hkl_flat = self.hkl_grid_float.reshape(-1, 3) - - for i in range(self.n_ops): - transformed = torch.matmul(hkl_flat, self.reciprocal_matrices[i].T) - transformed_int = torch.round(transformed).to(torch.int64) - - hkl_int = self.hkl_grid.reshape(-1, 3) - maps_to_minus_h = (transformed_int == -hkl_int).all(dim=-1) - - centric_mask = maps_to_minus_h.reshape(self.grid_shape) - centric = centric | centric_mask - - self.register_buffer("centric_mask", centric) - - def apply_to_indices(self, hkl, operation_index=None): - """ - Apply symmetry operation(s) to Miller indices. - - Parameters - ---------- - hkl : torch.Tensor, shape (..., 3) - Miller indices (h, k, l). - operation_index : int, optional - If specified, apply only this operation. - If None, apply all operations. - - Returns - ------- - torch.Tensor - Transformed Miller indices. - If operation_index is None: shape (n_ops, ..., 3) - Otherwise: shape (..., 3) - """ - hkl = hkl.to(dtype=self.dtype_float, device=self.device) - original_shape = hkl.shape[:-1] - - if operation_index is not None: - # Apply single operation - R = self.reciprocal_matrices[operation_index] - transformed = torch.matmul(hkl, R.T) - return torch.round(transformed).to(torch.int64) - else: - # Apply all operations - hkl_flat = hkl.reshape(-1, 3) # (N, 3) - results = [] - for i in range(self.n_ops): - R = self.reciprocal_matrices[i] - transformed = torch.matmul(hkl_flat, R.T) - results.append(torch.round(transformed).to(torch.int64)) - - # Stack: (n_ops, N, 3) - stacked = torch.stack(results, dim=0) - return stacked.reshape(self.n_ops, *original_shape, 3) - - def get_phase_shift(self, hkl, operation_index): - """ - Get phase shift for a symmetry operation on given Miller indices. - - The phase shift is exp(2πi h·t) where t is the translation. - - Parameters - ---------- - hkl : torch.Tensor, shape (..., 3) - Miller indices. - operation_index : int - Symmetry operation index. - - Returns - ------- - torch.Tensor - Phase shift in radians, shape (...). - """ - hkl = hkl.to(dtype=self.dtype_float, device=self.device) - translation = self.symmetry.translations[operation_index] - phase = 2.0 * np.pi * torch.matmul(hkl, translation) - return phase - - def get_symmetry_mate(self, F_grid, operation_index): - """ - Apply a single symmetry operation to a structure factor grid. - - Parameters - ---------- - F_grid : torch.Tensor, shape (nh, nk, nl) - Complex structure factor grid. - operation_index : int - Index of the symmetry operation (0 to n_ops-1). - - Returns - ------- - torch.Tensor, shape (nh, nk, nl) - Structure factors after applying symmetry operation. - Includes phase shift from translation component. - """ - if operation_index < 0 or operation_index >= self.n_ops: - raise ValueError( - f"Operation index {operation_index} out of range [0, {self.n_ops-1}]" - ) - - if F_grid.shape != self.grid_shape: - raise ValueError( - f"Grid shape {F_grid.shape} doesn't match expected {self.grid_shape}" - ) - - # Get precomputed index grid for this operation - idx_grid = self.index_grids[operation_index] # (nh, nk, nl, 3) - - # Gather structure factors from transformed positions - F_transformed = F_grid[idx_grid[..., 0], idx_grid[..., 1], idx_grid[..., 2]] - - # Apply phase shift from translation: F(h') = F(h) * exp(2πi h·t) - phase = self.phase_shifts[operation_index] - if F_grid.is_complex(): - phase_factor = torch.exp(1j * phase.to(F_grid.dtype)) - F_transformed = F_transformed * phase_factor - # For real-valued grids (amplitudes), no phase shift needed - - return F_transformed - - def get_all_symmetry_mates(self, F_grid): - """ - Get all symmetry-related structure factor grids. - - Parameters - ---------- - F_grid : torch.Tensor, shape (nh, nk, nl) - Complex structure factor grid. - - Returns - ------- - list of torch.Tensor - List of symmetry-related grids. - """ - return [self.get_symmetry_mate(F_grid, i) for i in range(self.n_ops)] - - def symmetry_average(self, F_grid, weights=None): - """ - Average structure factors over all symmetry equivalents. - - This is useful for enforcing symmetry constraints on structure factors. - - Parameters - ---------- - F_grid : torch.Tensor, shape (nh, nk, nl) - Complex structure factor grid. - weights : torch.Tensor, optional - Weights for averaging, shape (nh, nk, nl). - If None, equal weights are used. - - Returns - ------- - torch.Tensor, shape (nh, nk, nl) - Symmetry-averaged structure factors. - """ - mates = self.get_all_symmetry_mates(F_grid) - stacked = torch.stack(mates, dim=0) # (n_ops, nh, nk, nl) - - if weights is not None: - weights = weights.unsqueeze(0) # (1, nh, nk, nl) - stacked = stacked * weights - return stacked.sum(dim=0) / (weights.sum() * self.n_ops) - else: - return stacked.mean(dim=0) - - def expand_to_p1(self, F_asym, asym_mask=None): - """ - Expand structure factors from asymmetric unit to full P1. - - Takes structure factors defined on the asymmetric unit and - generates the full reciprocal space by applying all symmetry - operations. - - Parameters - ---------- - F_asym : torch.Tensor, shape (nh, nk, nl) - Structure factors on asymmetric unit (other positions can be zero). - asym_mask : torch.Tensor, optional - Currently a no-op: this parameter is accepted but ignored by the - implementation. Regardless of its value, positions are filled using - a non-zero heuristic (entries with ``|F| < 1e-10`` are treated as - unset and filled from symmetry mates). - - Returns - ------- - torch.Tensor, shape (nh, nk, nl) - Full structure factor grid with all symmetry equivalents filled. - """ - F_full = F_asym.clone() - - for i in range(1, self.n_ops): # Skip identity - F_mate = self.get_symmetry_mate(F_asym, i) - - # Only fill in positions that are zero (not yet set) - if F_full.is_complex(): - mask = F_full.abs() < 1e-10 - else: - mask = F_full.abs() < 1e-10 - - F_full = torch.where(mask, F_mate, F_full) - - return F_full - - def apply_friedel(self, F_grid): - """ - Apply Friedel's law: F(-h,-k,-l) = F*(h,k,l). - - For normal (non-anomalous) scattering, the structure factor - at -h is the complex conjugate of F(h). - - Parameters - ---------- - F_grid : torch.Tensor, shape (nh, nk, nl) - Complex structure factor grid. - - Returns - ------- - torch.Tensor, shape (nh, nk, nl) - Structure factors averaged toward Friedel symmetry. The result is - ``0.5 * (F(h) + F*(-h))``, i.e. an average of each reflection with - its conjugated Friedel mate, not a hard replacement/enforcement. - """ - # Flip all indices: F(-h,-k,-l) - F_friedel = torch.flip(F_grid, dims=[0, 1, 2]) - - # Roll to handle the asymmetry at 0 index - nh, nk, nl = self.grid_shape - F_friedel = torch.roll(F_friedel, shifts=(1, 1, 1), dims=(0, 1, 2)) - - if F_grid.is_complex(): - F_friedel = F_friedel.conj() - - # Average F(h) and F*(-h) - return 0.5 * (F_grid + F_friedel) - - def is_systematic_absence(self, h, k, l): - """ - Check if a reflection is systematically absent. - - Parameters - ---------- - h, k, l : int - Miller indices. - - Returns - ------- - bool - True if the reflection is systematically absent. - """ - nh, nk, nl = self.grid_shape - idx_h = h % nh - idx_k = k % nk - idx_l = l % nl - return self.systematic_absences[idx_h, idx_k, idx_l].item() - - def is_centric(self, h, k, l): - """ - Check if a reflection is centric. - - Parameters - ---------- - h, k, l : int - Miller indices. - - Returns - ------- - bool - True if the reflection is centric (phase restricted to 0 or π). - """ - nh, nk, nl = self.grid_shape - idx_h = h % nh - idx_k = k % nk - idx_l = l % nl - return self.centric_mask[idx_h, idx_k, idx_l].item() - - def get_epsilon(self): - """ - Compute epsilon (multiplicity) factors for each reflection. - - Epsilon is the number of symmetry operations that map h to itself - (h -> h) or to its Friedel mate (h -> -h). This count is taken - unconditionally for every space group; there is no centric/acentric - branch. Note that folding in Friedel mates inflates the count relative - to the conventional ε (pure rotational multiplicity, h -> h only). - - Returns - ------- - torch.Tensor, shape (nh, nk, nl) - Epsilon factors for each reflection. - """ - epsilon = torch.zeros(self.grid_shape, dtype=torch.int32, device=self.device) - - hkl_flat = self.hkl_grid.reshape(-1, 3) - - for i in range(self.n_ops): - # Get transformed indices - idx_grid = self.index_grids[i] - transformed_flat = idx_grid.reshape(-1, 3) - - # Check if h' ≡ h or h' ≡ -h (same reflection or Friedel) - same = (transformed_flat == hkl_flat).all(dim=-1) - - nh, nk, nl = self.grid_shape - # Also check Friedel (-h, -k, -l) - hkl_neg = (-self.hkl_grid).reshape(-1, 3) - hkl_neg[:, 0] = hkl_neg[:, 0] % nh - hkl_neg[:, 1] = hkl_neg[:, 1] % nk - hkl_neg[:, 2] = hkl_neg[:, 2] % nl - - friedel = (transformed_flat == hkl_neg).all(dim=-1) - - contributes = (same | friedel).reshape(self.grid_shape) - epsilon += contributes.to(torch.int32) - - return epsilon - - def forward(self, F_grid, mode="average"): - """ - Apply symmetry to structure factor grid. - - Parameters - ---------- - F_grid : torch.Tensor, shape (nh, nk, nl) - Complex structure factor grid. - mode : str, default 'average' - Operation mode: - - 'average': Average over all symmetry equivalents - - 'expand': Expand from asymmetric unit to full grid - - 'sum': Sum all symmetry mates (for accumulation) - - Returns - ------- - torch.Tensor, shape (nh, nk, nl) - Processed structure factor grid. - """ - if mode == "average": - return self.symmetry_average(F_grid) - elif mode == "expand": - return self.expand_to_p1(F_grid) - elif mode == "sum": - mates = self.get_all_symmetry_mates(F_grid) - return torch.stack(mates, dim=0).sum(dim=0) - else: - raise ValueError( - f"Unknown mode: {mode}. Use 'average', 'expand', or 'sum'." - ) - - def __call__(self, F_grid, mode="average"): - """Make the class callable.""" - return self.forward(F_grid, mode=mode) - - def get_symmetry_info(self): - """ - Get information about reciprocal space symmetry. - - Returns - ------- - dict - Dictionary with symmetry information. - """ - return { - "space_group": self.space_group, - "n_operations": self.n_ops, - "reciprocal_matrices": self.reciprocal_matrices, - "translations": self.symmetry.translations, - "n_systematic_absences": self.systematic_absences.sum().item(), - "n_centric": self.centric_mask.sum().item(), - "grid_shape": self.grid_shape, - } - - def __repr__(self): - return ( - f"ReciprocalSymmetryGrid(space_group='{self.space_group}', " - f"n_ops={self.n_ops}, grid_shape={self.grid_shape})" - ) - - -# ============================================================================= -# Standalone functions for symmetry expansion -# ============================================================================= - - -def expand_hkl( +def _expand_hkl( + sym, hkl: torch.Tensor, - spacegroup: SpaceGroupLike, include_friedel: bool = True, remove_absences: bool = True, device: Optional[torch.device] = None, @@ -636,10 +42,10 @@ def expand_hkl( Parameters ---------- + sym : SpaceGroup + The space group whose asymmetric unit convention applies. hkl : torch.Tensor, shape (N, 3) Input Miller indices (asymmetric unit). - spacegroup : str, int, or gemmi.SpaceGroup - Space group specification. include_friedel : bool, default True Include Friedel mates (-h, -k, -l). remove_absences : bool, default True @@ -661,12 +67,9 @@ def expand_hkl( device = hkl.device # Get symmetry operations - symmetry = SpaceGroup(spacegroup, dtype=get_float_dtype(), device=device) - n_ops = symmetry.matrices.shape[0] - - # Reciprocal space matrices (transpose of real space) - recip_matrices = symmetry.matrices.transpose(-2, -1) - translations = symmetry.translations + n_ops = sym.n_ops + recip_matrices = sym.reciprocal.matrices.to(device=device) + translations = sym.translations.to(device=device) # Convert hkl to float for matrix operations hkl_float = hkl.to(dtype=get_float_dtype(), device=device) @@ -679,7 +82,7 @@ def expand_hkl( for i in range(n_ops): # h' = h @ R^T hkl_transformed = torch.round(torch.matmul(hkl_float, recip_matrices[i].T)).to( - torch.int32 + torch.int32 # dtype-ok: transformed Miller indices (hkl); fixed-width int32 representation ) # Phase shift from translation: -2π h·t, for h' = hR under the convention # F(h) = Σ_j f_j exp(+2πi h·x_j). Do NOT "simplify" the sign: the wrong sign @@ -721,20 +124,13 @@ def expand_hkl( # Build output tensors expanded_hkl = torch.tensor( - [list(k) for k in unique_dict.keys()], dtype=torch.int32, device=device + [list(k) for k in unique_dict.keys()], dtype=torch.int32, device=device # dtype-ok: unique Miller indices (hkl); fixed-width int32 representation ) phase_shifts = torch.tensor(unique_phases, dtype=get_float_dtype(), device=device) - orig_idx_tensor = torch.tensor(orig_indices, dtype=torch.int64, device=device) + orig_idx_tensor = torch.tensor(orig_indices, dtype=torch.int64, device=device) # dtype-ok: reflection index mapping; int64 index tensor required - # Remove systematic absences if requested - sg = SpaceGroup(spacegroup) - is_p1 = sg.number == 1 - - if remove_absences and not is_p1: - absence_mask = _check_systematic_absences( - expanded_hkl, symmetry.matrices, translations, device - ) - keep_mask = ~absence_mask + if remove_absences and sym.number != 1: + keep_mask = ~sym.is_absent(expanded_hkl) expanded_hkl = expanded_hkl[keep_mask] phase_shifts = phase_shifts[keep_mask] @@ -743,27 +139,27 @@ def expand_hkl( return expanded_hkl, orig_idx_tensor, phase_shifts -def complete_hkl( +def _complete_hkl( + sym, input_hkl: torch.Tensor, cell: torch.Tensor, - spacegroup: SpaceGroupLike, d_min: float, device: Optional[torch.device] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Complete a set of Miller indices by identifying missing reflections. - Generates every reflection within ``d_min`` for ``spacegroup`` (minus + Generates every reflection within ``d_min`` for ``sym`` (minus systematic absences), then maps the input onto that complete set. This does *not* expand symmetry -- the output stays in the input space group. Parameters ---------- + sym : SpaceGroup + The space group whose asymmetric unit convention applies. input_hkl : torch.Tensor, shape (N, 3) Input Miller indices (may be incomplete). cell : torch.Tensor, shape (6,) Unit cell parameters [a, b, c, alpha, beta, gamma]. - spacegroup : str, int, or gemmi.SpaceGroup - Space group specification. d_min : float High resolution limit in Angstroms. device : torch.device, optional @@ -788,18 +184,8 @@ def complete_hkl( all_hkl = generate_possible_hkl(cell, d_min, device=device) # Get symmetry operations for absence check - symmetry = SpaceGroup(spacegroup, dtype=get_float_dtype(), device=device) - translations = symmetry.translations - - # Remove systematic absences - sg = SpaceGroup(spacegroup) - is_p1 = sg.number == 1 - - if not is_p1: - absence_mask = _check_systematic_absences( - all_hkl, symmetry.matrices, translations, device - ) - all_hkl = all_hkl[~absence_mask] + if sym.number != 1: + all_hkl = all_hkl[~sym.is_absent(all_hkl)] # Build lookup dictionary from input hkl to indices input_hkl_np = input_hkl.cpu().numpy() @@ -812,7 +198,7 @@ def complete_hkl( all_hkl_np = all_hkl.cpu().numpy() n_complete = len(all_hkl) - input_indices = torch.full((n_complete,), -1, dtype=torch.int64, device=device) + input_indices = torch.full((n_complete,), -1, dtype=torch.int64, device=device) # dtype-ok: reflection index buffer (-1 sentinel); int64 index required missing_mask = torch.ones(n_complete, dtype=torch.bool, device=device) for i, hkl in enumerate(all_hkl_np): @@ -824,163 +210,15 @@ def complete_hkl( return all_hkl, input_indices, missing_mask -def expand_reflections( - reflection_data: "ReflectionData", - include_friedel: bool = True, - remove_absences: bool = True, - verbose: int = 1, -) -> "ReflectionData": - """Expand ReflectionData from the asymmetric unit to P1. - - A :func:`expand_hkl` wrapper that carries every ReflectionData field (F, - sigmas, I, phases, FOM, R-free flags) through the expansion. Phases receive - the translation shift; resolution is recomputed and ``bin_indices`` cleared, - since expansion invalidates them. - - Parameters - ---------- - reflection_data : ReflectionData - Input reflection data with hkl, F, F_sigma, etc. - include_friedel : bool, default True - If True, also include Friedel mates (-h, -k, -l). - remove_absences : bool, default True - If True, remove systematically absent reflections from output. - verbose : int, default 1 - Verbosity level. - - Returns - ------- - ReflectionData - New object holding the expanded reflections, spacegroup set to 'P1'. - - See Also - -------- - expand_hkl : Low-level function for HKL expansion without ReflectionData. - """ - from torchref.io.datasets.reflection_data import ReflectionData as RefData - - if reflection_data.hkl is None: - raise ValueError("ReflectionData has no Miller indices loaded") - - space_group = reflection_data.spacegroup or "P1" - device = reflection_data.device - n_orig = len(reflection_data.hkl) - - if verbose > 0: - symmetry = SpaceGroup(space_group, dtype=get_float_dtype(), device=device) - print(f"Expanding reflections for {space_group}") - print(f" Original reflections: {n_orig}") - print(f" Symmetry operations: {symmetry.n_ops}") - - hkl_expanded, orig_idx_tensor, phase_shifts = expand_hkl( - reflection_data.hkl, - space_group, - include_friedel=include_friedel, - remove_absences=remove_absences, - device=device, - ) - - if verbose > 0: - print(f" After expansion: {len(hkl_expanded)} unique reflections") - - expanded = RefData(verbose=reflection_data.verbose, device=device) - - expanded.hkl = hkl_expanded - expanded.cell = ( - reflection_data.cell.clone() if reflection_data.cell is not None else None - ) - expanded.spacegroup = SpaceGroup("P1") # Now in P1 since symmetry is expanded - - if reflection_data.F is not None: - expanded.F = reflection_data.F[orig_idx_tensor] - - if reflection_data.F_sigma is not None: - expanded.F_sigma = reflection_data.F_sigma[orig_idx_tensor] - - if reflection_data.I is not None: - expanded.I = reflection_data.I[orig_idx_tensor] - - if hasattr(reflection_data, "I_sigma") and reflection_data.I_sigma is not None: - expanded.I_sigma = reflection_data.I_sigma[orig_idx_tensor] - - # Phases must absorb the translation shift; with no phases present, stash the - # shifts so a later phase assignment can still apply them. - if hasattr(reflection_data, "phase") and reflection_data.phase is not None: - expanded.phase = reflection_data.phase[orig_idx_tensor] + phase_shifts - else: - expanded._expansion_phase_shifts = phase_shifts - - if hasattr(reflection_data, "fom") and reflection_data.fom is not None: - expanded.fom = reflection_data.fom[orig_idx_tensor] - - if reflection_data.rfree_flags is not None: - expanded.rfree_flags = reflection_data.rfree_flags[orig_idx_tensor] - - if expanded.cell is not None: - expanded._calculate_resolution() - - expanded.bin_indices = None # invalidated by expansion - - expanded.amplitude_source = reflection_data.amplitude_source - expanded.intensity_source = reflection_data.intensity_source - expanded.phase_source = reflection_data.phase_source - expanded.rfree_source = reflection_data.rfree_source - - expanded.source = reflection_data - expanded.last_op = f"expand_to_p1(include_friedel={include_friedel})" - - return expanded - - -def _check_systematic_absences( - hkl: torch.Tensor, - matrices: torch.Tensor, - translations: torch.Tensor, - device: torch.device, -) -> torch.Tensor: - """Bool mask over ``hkl``, True where some op maps h -> h at non-integer h·t. - - ``matrices`` are the *real-space* rotations (n_ops, 3, 3); the transpose is - taken here. - """ - n_refl = len(hkl) - n_ops = matrices.shape[0] - absent = torch.zeros(n_refl, dtype=torch.bool, device=device) - - hkl_float = hkl.to(dtype=get_float_dtype(), device=device) - recip_matrices = matrices.transpose(-2, -1).to( - dtype=get_float_dtype(), device=device - ) - translations = translations.to(dtype=get_float_dtype(), device=device) - - for i in range(n_ops): - R = recip_matrices[i] - t = translations[i] - - hkl_transformed = torch.matmul(hkl_float, R.T) - hkl_transformed_int = torch.round(hkl_transformed).to(torch.int32) - - same_reflection = (hkl_transformed_int == hkl).all(dim=-1) - - h_dot_t = torch.matmul(hkl_float, t) - - phase_mod = torch.abs(h_dot_t - torch.round(h_dot_t)) - non_integer_phase = phase_mod > 1e-6 - - absent = absent | (same_reflection & non_integer_phase) - - return absent - - -def reduce_hkl( +def _reduce_hkl( + sym, hkl_p1: torch.Tensor, - spacegroup: SpaceGroupLike, include_friedel: bool = True, device: Optional[torch.device] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Reduce P1 Miller indices to the asymmetric unit of a target space group. - The inverse of :func:`expand_hkl`: symmetry-equivalent P1 reflections merge + The inverse of :meth:`~torchref.symmetry.spacegroup.SpaceGroup.expand_hkl`: symmetry-equivalent P1 reflections merge into one ASU reflection. The index map has *constant multiplicity* -- its second dimension is always ``n_equiv = n_ops * (2 if include_friedel else 1)`` however many equivalents actually exist in ``hkl_p1`` -- so aggregation needs @@ -988,10 +226,10 @@ def reduce_hkl( Parameters ---------- + sym : SpaceGroup + The target space group. hkl_p1 : torch.Tensor, shape (N, 3) Input Miller indices in P1 (complete hemisphere). - spacegroup : str, int, or gemmi.SpaceGroup - Target space group specification. include_friedel : bool, default True If True, also consider Friedel mates when finding ASU representative. device : torch.device, optional @@ -1012,12 +250,9 @@ def reduce_hkl( device = hkl_p1.device # Get symmetry operations - symmetry = SpaceGroup(spacegroup, dtype=get_float_dtype(), device=device) - n_ops = symmetry.matrices.shape[0] - - # Reciprocal space matrices (transpose of real space) - recip_matrices = symmetry.matrices.transpose(-2, -1) - translations = symmetry.translations + n_ops = sym.n_ops + recip_matrices = sym.reciprocal.matrices.to(device=device) + translations = sym.translations.to(device=device) # Total number of equivalent positions per ASU reflection n_equiv = n_ops * (2 if include_friedel else 1) @@ -1039,7 +274,7 @@ def get_canonical_hkl(hkl_single): for i in range(n_ops): # h' = h @ R^T hkl_trans = torch.round(torch.matmul(hkl_single, recip_matrices[i].T)).to( - torch.int32 + torch.int32 # dtype-ok: transformed Miller indices (hkl); fixed-width int32 representation ) equivalents.append(hkl_trans) @@ -1076,7 +311,7 @@ def get_canonical_hkl(hkl_single): R = recip_matrices[equiv_idx] t = translations[equiv_idx] - hkl_trans = torch.round(torch.matmul(hkl_single, R.T)).to(torch.int32) + hkl_trans = torch.round(torch.matmul(hkl_single, R.T)).to(torch.int32) # dtype-ok: transformed Miller indices (hkl); fixed-width int32 representation # -2π h·t, same convention as expand_hkl (see the derivation there). phase_shift = -2.0 * np.pi * torch.matmul(hkl_single, t) @@ -1098,9 +333,9 @@ def get_canonical_hkl(hkl_single): asu_list = sorted(asu_reflections.keys()) n_asu = len(asu_list) - hkl_asu = torch.tensor(asu_list, dtype=torch.int32, device=device) + hkl_asu = torch.tensor(asu_list, dtype=torch.int32, device=device) # dtype-ok: ASU Miller indices (hkl); fixed-width int32 representation reduction_indices = torch.full( - (n_asu, n_equiv), -1, dtype=torch.int64, device=device + (n_asu, n_equiv), -1, dtype=torch.int64, device=device # dtype-ok: reduction index map (-1 sentinel); int64 index tensor required ) phase_shifts = torch.zeros((n_asu, n_equiv), dtype=get_float_dtype(), device=device) @@ -1159,9 +394,9 @@ def _asu_condition_vectorized(h, k, l, condition_key): return fn(h, k, l) -def canonicalize_hkl( +def _canonicalize_hkl( + sym, hkl: torch.Tensor, - spacegroup: SpaceGroupLike, include_friedel: bool = True, device: Optional[torch.device] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: @@ -1175,10 +410,10 @@ def canonicalize_hkl( Parameters ---------- + sym : SpaceGroup + The space group whose asymmetric unit convention applies. hkl : torch.Tensor, shape (N, 3), dtype int32 Input Miller indices. - spacegroup : str, int, or gemmi.SpaceGroup - Space group specification. include_friedel : bool, default True Whether Friedel mates are considered equivalent. device : torch.device, optional @@ -1211,16 +446,15 @@ def canonicalize_hkl( empty_hkl = torch.empty((0, 3), dtype=hkl_dtype, device=device) empty_f = torch.empty(0, dtype=get_float_dtype(), device=device) empty_b = torch.empty(0, dtype=torch.bool, device=device) - empty_i = torch.empty(0, dtype=torch.int64, device=device) + empty_i = torch.empty(0, dtype=torch.int64, device=device) # dtype-ok: empty index tensor; int64 index dtype required return empty_hkl, empty_f, empty_b, empty_i - # Normalize spacegroup (CPU-only: this branch builds numpy-backed lookup tables) - sg_obj = SpaceGroup(spacegroup, dtype=get_float_dtype(), device=torch.device("cpu")) - sg_gemmi = sg_obj._gemmi - asu = gemmi.ReciprocalAsu(sg_gemmi) + # The ASU lookup tables are numpy-backed, so the operations come across to CPU + # regardless of where ``sym`` lives; only the returned tensors honour ``device``. + asu = gemmi.ReciprocalAsu(sym._gemmi) condition_key = asu.condition_str() - recip_mats = sg_obj.matrices.transpose(-2, -1).numpy() # (n_ops, 3, 3) - translations_np = sg_obj.translations.numpy() # (n_ops, 3) + recip_mats = sym.reciprocal.matrices.detach().cpu().numpy() # (n_ops, 3, 3) + translations_np = sym.translations.detach().cpu().numpy() # (n_ops, 3) n_ops = len(recip_mats) hkl_np = hkl.cpu().numpy().astype(np.int32) # (N, 3) @@ -1288,7 +522,7 @@ def canonicalize_hkl( example = hkl_np[np.where(remaining)[0][0]].tolist() raise ValueError( f"canonicalize_hkl could not map {n_unmapped} reflection(s) to the " - f"reciprocal ASU of space group {sg_obj} " + f"reciprocal ASU of space group {sym} " f"(include_friedel={include_friedel}); e.g. hkl={example}. With " f"include_friedel=False the Friedel half of reciprocal space has no " f"pure-rotation representative in the Laue-based CCP4 ASU." @@ -1316,9 +550,9 @@ def canonicalize_hkl( h_max = int(canonical_hkl.abs().max().item()) + 1 base = 2 * h_max + 1 sort_key = ( - canonical_hkl[:, 0].to(torch.int64) * base * base - + canonical_hkl[:, 1].to(torch.int64) * base - + canonical_hkl[:, 2].to(torch.int64) + canonical_hkl[:, 0].to(torch.int64) * base * base # dtype-ok: linear HKL hash/key; int64 avoids overflow for indexing + + canonical_hkl[:, 1].to(torch.int64) * base # dtype-ok: linear HKL hash/key; int64 avoids overflow for indexing + + canonical_hkl[:, 2].to(torch.int64) # dtype-ok: linear HKL hash/key; int64 avoids overflow for indexing ) sort_indices = torch.argsort(sort_key) @@ -1328,58 +562,3 @@ def canonicalize_hkl( friedel_flags[sort_indices], sort_indices, ) - - -def expand_reciprocal_grid( - F_grid: torch.Tensor, - space_group: str, - mode: str = "average", - include_friedel: bool = True, - device: Optional[torch.device] = None, -) -> torch.Tensor: - """Expand or symmetrize a reciprocal space grid using crystallographic symmetry. - - Convenience wrapper that builds a :class:`ReciprocalSymmetryGrid` for - ``F_grid.shape`` and applies it, so it re-pays the whole precompute on every - call -- hold the class yourself inside a loop. - - Parameters - ---------- - F_grid : torch.Tensor, shape (nh, nk, nl) - Input structure factor grid (can be complex or real). - space_group : str, int, or gemmi.SpaceGroup - Space group specification (e.g., 'P21', 'P212121'). The type hint is - narrowed to ``str``, but any ``SpaceGroupLike`` value is accepted. - mode : {'average', 'expand', 'sum'}, default 'average' - Average over all symmetry equivalents (symmetrize), expand from the - asymmetric unit to the full grid, or sum all symmetry mates. - include_friedel : bool, default True - If True, also apply Friedel symmetry after space group symmetry. - device : torch.device, optional - Device for computation. If None, uses F_grid's device. - - Returns - ------- - torch.Tensor, shape (nh, nk, nl) - Symmetrized or expanded structure factor grid. - """ - if device is None: - device = F_grid.device - - grid_shape = F_grid.shape - dtype = get_float_dtype() if not F_grid.is_complex() else F_grid.real.dtype - - recip_sym = ReciprocalSymmetryGrid( - space_group=space_group, - grid_shape=grid_shape, - dtype_float=dtype, - verbose=0, - device=device, - ) - - F_result = recip_sym(F_grid, mode=mode) - - if include_friedel: - F_result = recip_sym.apply_friedel(F_result) - - return F_result diff --git a/torchref/symmetry/spacegroup.py b/torchref/symmetry/spacegroup.py index da0f7c47..e7c7d5a4 100644 --- a/torchref/symmetry/spacegroup.py +++ b/torchref/symmetry/spacegroup.py @@ -1,39 +1,61 @@ -"""Space group utilities using gemmi as the canonical representation. - -:class:`SpaceGroup` is the interface used throughout torchref: an ``nn.Module`` -that normalizes its input (string, int, ``gemmi.SpaceGroup``, another -``SpaceGroup``, or None for P1), holds the rotation matrices and translations as -registered buffers, and applies them. The module-level functions here are -stateless equivalents plus the FFT/symmetry grid-size helpers. - -Real and reciprocal space use *transposed* conventions -- ``x' = R·x + t`` versus -``h' = Rᵀ·h`` -- so :meth:`SpaceGroup.apply` and :meth:`SpaceGroup.apply_to_hkl` -are not interchangeable. +"""Crystallographic space groups, using gemmi as the canonical source of truth. + +:class:`SpaceGroup` is the interface used throughout torchref. It specialises +:class:`~torchref.symmetry.symmetry.Symmetry` -- which owns the operations and +everything derivable from them -- with the crystallographic identity (Hermann-Mauguin +naming, number, point group, crystal system) and the CCP4 asymmetric-unit +conventions, the two things a bare operation list cannot supply. + +Construction normalizes any ``SpaceGroupLike``: a Hermann-Mauguin string, a number +1-230, a ``gemmi.SpaceGroup``, another :class:`SpaceGroup`, or None for P1. Only the +derived metadata is retained; no persistent ``gemmi`` reference is held, because a +lasting reference to the C++ singleton produces nanobind leak warnings at shutdown. + +Real and reciprocal space use *transposed* conventions. That is handled once, in +:attr:`~torchref.symmetry.symmetry.Symmetry.reciprocal`, rather than being re-decided +per call site. """ from __future__ import annotations -from typing import Union +from typing import Optional, Union import gemmi import torch -import torch.nn as nn from torchref.config import get_float_dtype, normalize_device -from torchref.utils.debug_utils import DebugMixin -from torchref.utils.device_mixin import DeviceMovementMixin +from torchref.symmetry.symmetry import Symmetry # Type alias for space group input - includes SpaceGroup class itself SpaceGroupLike = Union[str, int, gemmi.SpaceGroup, "SpaceGroup", None] +# gemmi stores rotations and translations as integers scaled by 24. +_GEMMI_SCALE = 24.0 + def _normalize_spacegroup(spacegroup: SpaceGroupLike) -> gemmi.SpaceGroup: """Normalize any ``SpaceGroupLike`` to a ``gemmi.SpaceGroup``. - Accepts a Hermann-Mauguin string (spacing-insensitive, retried upper-cased), - a number 1-230, a ``gemmi.SpaceGroup`` (returned unchanged), a - :class:`SpaceGroup` (unwrapped), or None (P1). Raises ``ValueError`` for an - unrecognised name/number, ``TypeError`` for any other type. + Accepts a Hermann-Mauguin string (spacing-insensitive, retried upper-cased), a + number 1-230, a ``gemmi.SpaceGroup`` (returned unchanged), a :class:`SpaceGroup` + (unwrapped), or None (P1). + + Parameters + ---------- + spacegroup : SpaceGroupLike + Space group in any supported form. + + Returns + ------- + gemmi.SpaceGroup + The normalized space group. + + Raises + ------ + ValueError + For an unrecognised name or number. + TypeError + For any other type. """ if spacegroup is None: return gemmi.SpaceGroup("P 1") @@ -41,47 +63,31 @@ def _normalize_spacegroup(spacegroup: SpaceGroupLike) -> gemmi.SpaceGroup: if isinstance(spacegroup, gemmi.SpaceGroup): return spacegroup - # Handle SpaceGroup class instances (forward reference resolved at runtime) - if hasattr(spacegroup, "_sg_hm") and hasattr(spacegroup, "matrices"): - return gemmi.find_spacegroup_by_name(spacegroup._sg_hm) + # Duck-typed rather than an isinstance check against SpaceGroup, so this stays + # usable from module scope before the class below is defined. + if hasattr(spacegroup, "_sg_xhm") and hasattr(spacegroup, "matrices"): + # The extended symbol carries the setting; the plain H-M symbol does not + # (``R 3:R`` would come back as ``R 3:H``). + return gemmi.find_spacegroup_by_name(spacegroup._sg_xhm) if isinstance(spacegroup, int): - # Space group number try: return gemmi.SpaceGroup(spacegroup) except Exception as e: raise ValueError(f"Invalid space group number: {spacegroup}") from e if isinstance(spacegroup, str): - # Try to parse as string - # Clean up common variations sg_clean = spacegroup.strip() - - # Handle double spaces that sometimes appear while " " in sg_clean: sg_clean = sg_clean.replace(" ", " ") - - try: - return gemmi.SpaceGroup(sg_clean) - except Exception: - pass - - # Try without spaces sg_nospace = sg_clean.replace(" ", "") - try: - return gemmi.SpaceGroup(sg_nospace) - except Exception: - pass - - # Try common substitutions - substitutions = [ - (sg_clean, sg_clean), - (sg_nospace, sg_nospace), - (sg_clean.upper(), sg_clean.upper()), - (sg_nospace.upper(), sg_nospace.upper()), - ] - - for _, variant in substitutions: + + for variant in ( + sg_clean, + sg_nospace, + sg_clean.upper(), + sg_nospace.upper(), + ): try: return gemmi.SpaceGroup(variant) except Exception: @@ -99,452 +105,86 @@ def _normalize_spacegroup(spacegroup: SpaceGroupLike) -> gemmi.SpaceGroup: ) -def spacegroup_to_str(spacegroup: SpaceGroupLike, style: str = "short") -> str: - """ - Convert space group to string representation. - - Parameters - ---------- - spacegroup : SpaceGroupLike - Space group in any supported format. - style : str, default 'short' - Output style: - - 'short': No spaces (e.g., 'P212121') - - 'hm': Hermann-Mauguin with spaces (e.g., 'P 21 21 21') - - 'xhm': Extended Hermann-Mauguin, including the setting/cell-choice - token where applicable (e.g., 'P 1 21 1' for a unique-axis-b setting) - - Returns - ------- - str - Space group name in requested style. - """ - sg = _normalize_spacegroup(spacegroup) - - if style == "short": - return sg.short_name() - elif style == "hm": - return sg.hm - elif style == "xhm": - return sg.xhm() - else: - raise ValueError(f"Unknown style: {style}. Use 'short', 'hm', or 'xhm'.") - - -def get_symmetry_operations(spacegroup: SpaceGroupLike): - """ - Get symmetry operations from a space group. - - Parameters - ---------- - spacegroup : SpaceGroupLike - Space group in any supported format. - - Returns - ------- - list of gemmi.Op - List of symmetry operations. - """ - sg = _normalize_spacegroup(spacegroup) - return list(sg.operations()) - - -def get_operations_as_tensors( - spacegroup: SpaceGroupLike, - dtype: torch.dtype = None, - device: torch.device = None, -): - """ - Get symmetry operations as PyTorch tensors. +def _operations_as_tensors( + sg: gemmi.SpaceGroup, + dtype: torch.dtype, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Extract a gemmi space group's operations as tensors. Parameters ---------- - spacegroup : SpaceGroupLike - Space group in any supported format. - dtype : torch.dtype, optional - Data type for tensors. Defaults to the configured ``dtypes.float``. - device : torch.device, optional - Device for tensors. Defaults to the configured ``device.current``. + sg : gemmi.SpaceGroup + Normalized space group. + dtype : torch.dtype + Floating dtype for both tensors. + device : torch.device + Device for both tensors. Returns ------- - matrices : torch.Tensor, shape (n_ops, 3, 3) - Rotation matrices. - translations : torch.Tensor, shape (n_ops, 3) - Translation vectors (in fractional coordinates). + matrices : torch.Tensor + Rotation matrices, shape ``(n_ops, 3, 3)``. + translations : torch.Tensor + Fractional translations, shape ``(n_ops, 3)``. """ - if dtype is None: - dtype = get_float_dtype() - device = normalize_device(device) - sg = _normalize_spacegroup(spacegroup) - - # Extract rotation matrices and translations from gemmi operations - # gemmi stores values as integers multiplied by 24, divide to get actual values - gemmi_ops = [ + ops = [ ( - torch.tensor(op.rot, dtype=dtype, device=device) / 24.0, - torch.tensor(op.tran, dtype=dtype, device=device) / 24.0, + torch.tensor(op.rot, dtype=dtype, device=device) / _GEMMI_SCALE, + torch.tensor(op.tran, dtype=dtype, device=device) / _GEMMI_SCALE, ) for op in sg.operations() ] - matrices, translations = zip(*gemmi_ops) - + matrices, translations = zip(*ops) return torch.stack(matrices), torch.stack(translations) -def is_same_spacegroup(sg1: SpaceGroupLike, sg2: SpaceGroupLike) -> bool: - """ - Check if two space groups are the same. - - Parameters - ---------- - sg1, sg2 : SpaceGroupLike - Space groups to compare. - - Returns - ------- - bool - True if the space groups are identical. - """ - return _normalize_spacegroup(sg1).number == _normalize_spacegroup(sg2).number +class SpaceGroup(Symmetry): + """A crystallographic space group: symmetry operations plus their identity. - -def get_point_group(spacegroup: SpaceGroupLike) -> str: - """ - Get the point group symbol for a space group. - - Parameters - ---------- - spacegroup : SpaceGroupLike - Space group in any supported format. - - Returns - ------- - str - Point group symbol (e.g., '222', 'mmm', '4/mmm'). - """ - sg = _normalize_spacegroup(spacegroup) - return sg.point_group_hm() - - -def get_crystal_system(spacegroup: SpaceGroupLike) -> str: - """ - Get the crystal system for a space group. - - Parameters - ---------- - spacegroup : SpaceGroupLike - Space group in any supported format. - - Returns - ------- - str - Crystal system name (triclinic, monoclinic, orthorhombic, - tetragonal, trigonal, hexagonal, or cubic). - """ - sg = _normalize_spacegroup(spacegroup) - return sg.crystal_system_str() - - -def is_centrosymmetric(spacegroup: SpaceGroupLike) -> bool: - """ - Check if a space group is centrosymmetric. - - Parameters - ---------- - spacegroup : SpaceGroupLike - Space group in any supported format. - - Returns - ------- - bool - True if the space group has an inversion center. - """ - sg = _normalize_spacegroup(spacegroup) - return sg.is_centrosymmetric() - - -def n_operations(spacegroup: SpaceGroupLike) -> int: - """ - Get the number of symmetry operations in a space group. - - Parameters - ---------- - spacegroup : SpaceGroupLike - Space group in any supported format. - - Returns - ------- - int - Number of symmetry operations. - """ - sg = _normalize_spacegroup(spacegroup) - return len(list(sg.operations())) - - -# ============================================================================= -# Grid size utilities (combined FFT-friendly and symmetry-friendly) -# ============================================================================= - - -def is_fft_friendly(n: int) -> bool: - """True if ``n`` factors into 2, 3 and 5 only (radix-2,3,5 FFT sizes). - - ``n <= 0`` is False; 128 and 135 are True, 131 is not. - """ - if n <= 0: - return False - - # Remove all factors of 2, 3, 5 - while n % 2 == 0: - n //= 2 - while n % 3 == 0: - n //= 3 - while n % 5 == 0: - n //= 5 - - # If we're left with 1, the number is FFT-friendly - return n == 1 - - -def find_fft_friendly_size(n: int, divisibility: int = 1) -> int: - """Smallest size >= ``n`` that is FFT-friendly and divisible by ``divisibility``. - - Parameters - ---------- - n : int - Minimum grid size. - divisibility : int, default 1 - Required divisibility (e.g. 2 for a screw axis). - - Returns - ------- - int - Optimal grid size (131 -> 135; 131 with divisibility 2 -> 160). - """ - candidate = n - - # Make sure it satisfies divisibility - if candidate % divisibility != 0: - candidate = ((candidate // divisibility) + 1) * divisibility - - # Now find nearest FFT-friendly size - while not is_fft_friendly(candidate): - candidate += divisibility - - return candidate - - -def get_grid_requirements(spacegroup: SpaceGroupLike) -> dict: - """Per-axis grid divisibility required for interpolation-free symmetry expansion. - - Derived from the denominators of the fractional translations, so a grid meeting - them indexes symmetry mates at exact integers. - - Returns - ------- - dict - ``{'nx_mod': int, 'ny_mod': int, 'nz_mod': int}`` -- e.g. P21 gives - ``(1, 2, 1)``, P212121 gives ``(2, 2, 2)``. - """ - import math - from fractions import Fraction - - sg = _normalize_spacegroup(spacegroup) - - # Start with no requirements - nx_lcm = 1 - ny_lcm = 1 - nz_lcm = 1 - - # Analyze each symmetry operation - for op in sg.operations(): - # gemmi stores translations as integers multiplied by 24 - trans = [t / 24.0 for t in op.tran] - - # For each axis, check if translation has fractional component - for axis_idx, t in enumerate(trans): - if abs(t) > 1e-9: - # Convert to fraction and get denominator - frac = Fraction(t).limit_denominator(24) - denom = frac.denominator - - if axis_idx == 0: - nx_lcm = math.lcm(nx_lcm, denom) - elif axis_idx == 1: - ny_lcm = math.lcm(ny_lcm, denom) - else: - nz_lcm = math.lcm(nz_lcm, denom) - - return {"nx_mod": nx_lcm, "ny_mod": ny_lcm, "nz_mod": nz_lcm} - - -def check_grid_compatibility(grid_shape: tuple, spacegroup: SpaceGroupLike) -> dict: - """Check a grid against both space-group divisibility and FFT-friendliness. - - Parameters - ---------- - grid_shape : tuple of int - Grid dimensions (nx, ny, nz). - spacegroup : SpaceGroupLike - Space group in any supported format. - - Returns - ------- - dict - ``compatible`` (both tests pass), ``symmetry_compatible``, - ``fft_friendly``, ``can_use_direct_indexing`` (interpolation-free - expansion possible -- equal to ``symmetry_compatible``), ``issues`` - (per-axis descriptions, empty when compatible) and ``requirements`` - (from :func:`get_grid_requirements`). - """ - nx, ny, nz = grid_shape - sg = _normalize_spacegroup(spacegroup) - requirements = get_grid_requirements(sg) - - issues = [] - sg_name = sg.short_name() - - # Check symmetry requirements - if nx % requirements["nx_mod"] != 0: - issues.append( - f"nx={nx} not divisible by {requirements['nx_mod']} " - f"(required for {sg_name} symmetry)" - ) - - if ny % requirements["ny_mod"] != 0: - issues.append( - f"ny={ny} not divisible by {requirements['ny_mod']} " - f"(required for {sg_name} symmetry)" - ) - - if nz % requirements["nz_mod"] != 0: - issues.append( - f"nz={nz} not divisible by {requirements['nz_mod']} " - f"(required for {sg_name} symmetry)" - ) - - symmetry_compatible = len(issues) == 0 - - # Check FFT-friendly - fft_x = is_fft_friendly(nx) - fft_y = is_fft_friendly(ny) - fft_z = is_fft_friendly(nz) - fft_friendly = fft_x and fft_y and fft_z - - if not fft_x: - issues.append(f"nx={nx} is not FFT-friendly (not a product of 2, 3, 5)") - if not fft_y: - issues.append(f"ny={ny} is not FFT-friendly (not a product of 2, 3, 5)") - if not fft_z: - issues.append(f"nz={nz} is not FFT-friendly (not a product of 2, 3, 5)") - - return { - "compatible": symmetry_compatible and fft_friendly, - "symmetry_compatible": symmetry_compatible, - "fft_friendly": fft_friendly, - "can_use_direct_indexing": symmetry_compatible, - "issues": issues, - "requirements": requirements, - } - - -def suggest_grid_size( - min_grid_shape: tuple, - spacegroup: SpaceGroupLike, - make_fft_friendly: bool = True, -) -> tuple: - """Smallest grid >= ``min_grid_shape`` meeting the symmetry divisibility. - - Parameters - ---------- - min_grid_shape : tuple of int - Minimum (nx, ny, nz) grid dimensions. - spacegroup : SpaceGroupLike - Space group in any supported format. - make_fft_friendly : bool, default True - If True, the result also factors into 2, 3, 5 only. - - Returns - ------- - tuple of int - Suggested grid dimensions (nx, ny, nz). - """ - requirements = get_grid_requirements(spacegroup) - - def find_next_valid(n, divisibility): - """Find next number >= n that satisfies divisibility and FFT constraints.""" - if n % divisibility == 0: - candidate = n - else: - candidate = ((n // divisibility) + 1) * divisibility - - if not make_fft_friendly: - return candidate - - # Find FFT-friendly size that also satisfies divisibility - while not is_fft_friendly(candidate): - candidate += divisibility - - return candidate - - nx = find_next_valid(min_grid_shape[0], requirements["nx_mod"]) - ny = find_next_valid(min_grid_shape[1], requirements["ny_mod"]) - nz = find_next_valid(min_grid_shape[2], requirements["nz_mod"]) - - return (nx, ny, nz) - - -# ============================================================================= -# SpaceGroup class - unified interface combining normalization and operations -# ============================================================================= - - -class SpaceGroup(DeviceMovementMixin, DebugMixin, nn.Module): - """ - Unified space group handler for crystallographic symmetry operations. - - Normalizes its input, holds the operations as buffers, applies them to - fractional coordinates (``__call__``) or Miller indices, and exposes the - grid-size helpers as methods. Only the derived metadata is retained -- no - persistent ``gemmi`` reference is held (see :attr:`_gemmi`). + Inherits the whole operation-derived surface from + :class:`~torchref.symmetry.symmetry.Symmetry` -- expansion, phases, reflection + predicates, grid sizing, map symmetrization -- and adds the crystallographic + naming and the CCP4 asymmetric-unit conventions. Parameters ---------- space_group : str, int, gemmi.SpaceGroup, SpaceGroup, or None - Hermann-Mauguin symbol, number 1-230, gemmi object, another instance, - or None for P1. + Hermann-Mauguin symbol, number 1-230, gemmi object, another instance, or None + for P1. dtype : torch.dtype, optional - Data type for matrices and translations. Defaults to the configured - ``dtypes.float`` (float32 unless ``TORCHREF_DTYPE_FLOAT=float64``). - device : torch.device, default: configured device.current - Device for computation. + Dtype for the operations. Defaults to the configured ``dtypes.float``. + device : torch.device, optional + Device for the operations. Defaults to the configured ``device.current``. Attributes ---------- - matrices : torch.Tensor, shape (n_ops, 3, 3) - Rotation matrices for all symmetry operations (registered buffer). - translations : torch.Tensor, shape (n_ops, 3) - Translation vectors for all symmetry operations (registered buffer). - n_ops : int - Number of symmetry operations. + matrices : torch.Tensor + Rotation matrices, shape ``(n_ops, 3, 3)``. + translations : torch.Tensor + Fractional translations, shape ``(n_ops, 3)``. + + Examples + -------- + >>> sg = SpaceGroup("P212121") + >>> sg.n_ops + 4 + >>> sg.crystal_system + 'orthorhombic' """ def __init__( self, space_group: SpaceGroupLike = None, - dtype: torch.dtype = None, - device: torch.device = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, ): - super(SpaceGroup, self).__init__() if dtype is None: dtype = get_float_dtype() device = normalize_device(device) - self._device = device - self._dtype = dtype - # Normalize to gemmi.SpaceGroup, extract metadata, then release gemmi_sg = _normalize_spacegroup(space_group) + self._sg_number: int = gemmi_sg.number self._sg_hm: str = gemmi_sg.hm self._sg_short_name: str = gemmi_sg.short_name() @@ -553,254 +193,315 @@ def __init__( self._sg_crystal_system: str = gemmi_sg.crystal_system_str() self._sg_centrosymmetric: bool = gemmi_sg.is_centrosymmetric() - # Get symmetry operations as tensors - matrices, translations = get_operations_as_tensors( - gemmi_sg, dtype=dtype, device=device - ) + matrices, translations = _operations_as_tensors(gemmi_sg, dtype, device) + # gemmi_sg goes out of scope here -- no persistent gemmi reference. - self.register_buffer("matrices", matrices) - self.register_buffer("translations", translations) - # gemmi_sg goes out of scope here — no persistent gemmi reference + super().__init__(matrices=matrices, translations=translations) # ========================================================================= - # Core properties + # Crystallographic identity # ========================================================================= - @property - def n_ops(self) -> int: - """Number of symmetry operations.""" - return self.matrices.shape[0] - @property def _gemmi(self) -> gemmi.SpaceGroup: - """Fresh gemmi.SpaceGroup each access; never cache it -- a persistent - reference to the C++ singleton produces nanobind leak warnings at shutdown.""" - return gemmi.find_spacegroup_by_name(self._sg_hm) + """A fresh ``gemmi.SpaceGroup`` on each access. + + Never cached: a persistent reference to the C++ singleton produces nanobind + leak warnings at interpreter shutdown. + """ + return gemmi.find_spacegroup_by_name(self._sg_xhm) @property def name(self) -> str: - """Short space group name (e.g., 'P21').""" + """Short space group name, e.g. ``'P21'``.""" return self._sg_short_name @property def hm(self) -> str: - """Hermann-Mauguin notation with spaces (e.g., 'P 21').""" + """Hermann-Mauguin notation with spaces, e.g. ``'P 21'``.""" return self._sg_hm @property def xhm(self) -> str: - """Extended Hermann-Mauguin notation.""" + """Extended Hermann-Mauguin notation, including the setting token.""" + return self._sg_xhm + + @property + def key(self) -> str: + """Value identity of this space group: the extended H-M symbol. + + Two groups with the same key have the same operations in the same + setting. The group number alone would not do -- ``P 1 21 1`` and + ``P 1 1 21`` share number 4 but place the screw axis differently. + """ return self._sg_xhm @property def number(self) -> int: - """Space group number (1-230).""" + """Space group number, 1-230.""" return self._sg_number @property def gemmi(self) -> gemmi.SpaceGroup: - """Access a gemmi.SpaceGroup object (created on demand, not stored).""" + """A ``gemmi.SpaceGroup``, created on demand and not stored.""" return self._gemmi @property def point_group(self) -> str: - """Point group symbol (e.g., '222', 'mmm').""" + """Point group symbol, e.g. ``'222'`` or ``'mmm'``.""" return self._sg_point_group @property def crystal_system(self) -> str: - """Crystal system name.""" + """Crystal system name, e.g. ``'orthorhombic'``.""" return self._sg_crystal_system @property def centrosymmetric(self) -> bool: - """True if space group has inversion center.""" + """Whether the group has an inversion centre.""" return self._sg_centrosymmetric - @property - def dtype(self) -> torch.dtype: - """Data type used for matrices.""" - return self._dtype + def short_name(self) -> str: + """Short space group name; the callable form of :attr:`name`.""" + return self._sg_short_name - @property - def device(self) -> torch.device: - """Device for matrices.""" - return self._device + def operations(self): + """The gemmi operations, from a temporary gemmi object. + + Returns + ------- + gemmi.GroupOps + The operation list. + """ + return self._gemmi.operations() # ========================================================================= - # Backward compatibility aliases + # Aliases retained for existing callers # ========================================================================= @property def spacegroup(self) -> gemmi.SpaceGroup: - """Alias for gemmi property (backward compatibility).""" + """Alias for :attr:`gemmi`.""" return self._gemmi @property def space_group(self) -> gemmi.SpaceGroup: - """Alias for gemmi property (backward compatibility).""" + """Alias for :attr:`gemmi`.""" return self._gemmi @property def space_group_name(self) -> str: - """Alias for name property (backward compatibility).""" + """Alias for :attr:`name`.""" return self.name @property def space_group_number(self) -> int: - """Alias for number property (backward compatibility).""" + """Alias for :attr:`number`.""" return self.number # ========================================================================= - # Gemmi method delegation for backward compatibility + # Asymmetric-unit conventions # ========================================================================= + # + # These need the CCP4 asymmetric unit, which is keyed by Laue class, so they live + # here rather than on ``Symmetry`` -- "the canonical ASU" is meaningless for a bare + # operation list. The algorithms are in ``reciprocal_symmetry``; these methods are + # the only public way in. - def short_name(self) -> str: - """Get short space group name.""" - return self._sg_short_name - - def operations(self): - """Get symmetry operations (creates temporary gemmi object on demand).""" - return self._gemmi.operations() + def expand_hkl( + self, + hkl: torch.Tensor, + include_friedel: bool = True, + remove_absences: bool = True, + device: Optional[torch.device] = None, + ): + """Expand Miller indices from the asymmetric unit to P1. - # ========================================================================= - # Symmetry operation methods - # ========================================================================= + Parameters + ---------- + hkl : torch.Tensor + Input Miller indices, shape ``(N, 3)``. + include_friedel : bool, default True + Include Friedel mates ``(-h, -k, -l)``. + remove_absences : bool, default True + Drop systematically absent reflections. + device : torch.device, optional + Computation device. Defaults to ``hkl``'s. - def apply( - self, xyz_fractional: torch.Tensor, apply_translation: bool = True - ) -> torch.Tensor: + Returns + ------- + expanded_hkl : torch.Tensor + Expanded indices, shape ``(M, 3)``, dtype ``int32``. + orig_indices : torch.Tensor + Map expanded -> original, shape ``(M,)``: ``F_exp = F_orig[orig_indices]``. + phase_shifts : torch.Tensor + Translation phase offsets in radians, shape ``(M,)``: + ``phase_exp = phase_orig[orig_indices] + phase_shifts``. """ - Apply symmetry operations to fractional coordinates (rotation + translation). + from torchref.symmetry.reciprocal_symmetry import _expand_hkl + + return _expand_hkl( + self, + hkl, + include_friedel=include_friedel, + remove_absences=remove_absences, + device=device, + ) + + def reduce_hkl( + self, + hkl_p1: torch.Tensor, + include_friedel: bool = True, + device: Optional[torch.device] = None, + ): + """Reduce P1 Miller indices to this group's asymmetric unit. - For real space coordinates, applies the full symmetry operation: x' = R·x + t + The inverse of :meth:`expand_hkl`. Parameters ---------- - xyz_fractional : torch.Tensor - Input tensor of shape (N, 3) representing fractional coordinates. - apply_translation : bool, default True - If True, apply the full operation x' = R·x + t. If False, apply - the rotational part only (x' = R·x), as used for Miller indices. + hkl_p1 : torch.Tensor + P1 Miller indices, shape ``(N, 3)``. + include_friedel : bool, default True + Consider Friedel mates when picking the ASU representative. + device : torch.device, optional + Computation device. Defaults to ``hkl_p1``'s. Returns ------- - torch.Tensor - Transformed coordinates of shape (N, 3, ops) where ops is the - number of symmetry operations. - - See Also - -------- - apply_to_hkl : For reciprocal space (Miller indices), rotation only. - """ - coords = xyz_fractional.to(self.matrices.device).to(self.matrices.dtype) - # coords: (N, 3), matrices: (ops, 3, 3) - # Apply rotation: result[n, i, o] = sum_j(matrices[o, i, j] * coords[n, j]) - transformed = torch.einsum("oij,nj->nio", self.matrices, coords) - # transformed: (N, 3, ops) - # Add translations: translations (ops, 3) -> (1, 3, ops) for broadcasting - if apply_translation: - transformed = transformed + self.translations.T.unsqueeze(0) - return transformed # (N, 3, ops) - - def apply_to_hkl(self, hkl: torch.Tensor) -> torch.Tensor: + hkl_asu : torch.Tensor + Unique ASU indices, shape ``(M, 3)``, dtype ``int32``. + reduction_indices : torch.Tensor + Indices into ``hkl_p1`` per equivalent, shape ``(M, n_equiv)``, **-1 where + no P1 reflection exists** -- mask or clamp before gathering, or a -1 + silently reads the last row. + phase_shifts : torch.Tensor + Phase shifts to apply before aggregation, shape ``(M, n_equiv)``. """ - Apply symmetry operations to Miller indices (rotation only, no translation). + from torchref.symmetry.reciprocal_symmetry import _reduce_hkl - Reciprocal space uses the *transpose*: ``h' = h·R = Rᵀ·h``. Substituting - ``R·h`` gives the wrong equivalents wherever the fractional rotation is - non-symmetric (trigonal, hexagonal, permutation-type cubic ops), silently - corrupting centric flags and epsilon multiplicities. Translations shift - structure-factor phases, not indices, so they are not applied here. + return _reduce_hkl( + self, hkl_p1, include_friedel=include_friedel, device=device + ) + + def complete_hkl( + self, + input_hkl: torch.Tensor, + cell: torch.Tensor, + d_min: float, + device: Optional[torch.device] = None, + ): + """Identify reflections missing from a dataset, without expanding symmetry. Parameters ---------- - hkl : torch.Tensor - Input tensor of shape (N, 3) representing Miller indices. + input_hkl : torch.Tensor + Possibly incomplete Miller indices, shape ``(N, 3)``. + cell : torch.Tensor + Unit cell parameters ``[a, b, c, alpha, beta, gamma]``, shape ``(6,)``. + d_min : float + High-resolution limit in Angstroms. + device : torch.device, optional + Computation device. Defaults to ``input_hkl``'s. Returns ------- - torch.Tensor - Transformed Miller indices of shape (N, 3, ops). - - See Also - -------- - apply : Real-space coordinates (``R·x + t``), the transposed convention. + complete_hkl : torch.Tensor + Every index within ``d_min`` minus systematic absences, shape ``(M, 3)``. + input_indices : torch.Tensor + Map complete -> input, shape ``(M,)``, ``-1`` where missing. + missing_mask : torch.Tensor + Boolean, shape ``(M,)``, True where absent from the input. """ - coords = hkl.to(self.matrices.device).to(self.matrices.dtype) - # result[n, i, o] = sum_j matrices[o, j, i] * coords[n, j] = (Rᵀ·h)_i - return torch.einsum("oji,nj->nio", self.matrices, coords) + from torchref.symmetry.reciprocal_symmetry import _complete_hkl - def expand_coords_to_P1(self, xyz_fractional: torch.Tensor) -> torch.Tensor: - """ - Expand fractional coordinates by applying all symmetry operations. + return _complete_hkl(self, input_hkl, cell, d_min, device=device) + + def canonicalize_hkl( + self, + hkl: torch.Tensor, + include_friedel: bool = True, + device: Optional[torch.device] = None, + ): + """Map Miller indices onto their canonical CCP4 ASU representatives. Parameters ---------- - xyz_fractional : torch.Tensor - Input tensor of shape (N, 3) representing fractional coordinates. + hkl : torch.Tensor + Input Miller indices, shape ``(N, 3)``. + include_friedel : bool, default True + Treat Friedel mates as equivalent. With ``False`` the Friedel half of + reciprocal space has no pure-rotation representative in the Laue-based + CCP4 ASU, and unmappable reflections raise. + device : torch.device, optional + Output device. Defaults to ``hkl``'s. The lookup itself runs on CPU + whatever device this group is on, because the ASU tables are numpy-backed. Returns ------- - torch.Tensor - Expanded coordinates of shape (N * ops, 3). + canonical_hkl : torch.Tensor + Remapped indices sorted lexicographically, shape ``(N, 3)``. + phase_shifts : torch.Tensor + Additive phase correction in radians, shape ``(N,)``. + friedel_flags : torch.Tensor + Boolean, shape ``(N,)``, True where Friedel conjugation was applied. + sort_indices : torch.Tensor + Permutation from original to sorted order, shape ``(N,)``. + + Notes + ----- + ``phase_shifts`` assumes the caller conjugates first: the contract is + ``phi_new = torch.where(friedel_flags, -phi_old, phi_old) + phase_shifts``. """ - transformed = self.apply(xyz_fractional) # (N, 3, ops) - N = xyz_fractional.shape[0] - ops = self.n_ops - # (N, 3, ops) -> (N, ops, 3) -> (N * ops, 3) - expanded = transformed.permute(0, 2, 1).reshape(N * ops, 3) - return expanded + from torchref.symmetry.reciprocal_symmetry import _canonicalize_hkl - def forward(self, xyz_fractional: torch.Tensor) -> torch.Tensor: - """Forward pass applies symmetry operations.""" - return self.apply(xyz_fractional) + return _canonicalize_hkl( + self, hkl, include_friedel=include_friedel, device=device + ) # ========================================================================= - # Grid utilities + # Copy and dunder # ========================================================================= - def get_grid_requirements(self) -> dict: - """Per-axis grid divisibility; see :func:`get_grid_requirements`.""" - return get_grid_requirements(self) - - def check_grid_compatibility(self, grid_shape: tuple) -> dict: - """Report whether ``(nx, ny, nz)`` suits this group and the FFT. + def copy(self) -> "SpaceGroup": + """An independent copy with cloned operations and an empty cache. - Returns the report dict documented in :func:`check_grid_compatibility`. + Returns + ------- + SpaceGroup + New instance carrying the same symmetry, dtype and device. """ - return check_grid_compatibility(grid_shape, self) - - def suggest_grid_size( - self, min_grid_shape: tuple, make_fft_friendly: bool = True - ) -> tuple: - """Smallest valid grid >= ``min_grid_shape``; see :func:`suggest_grid_size`.""" - return suggest_grid_size(min_grid_shape, self, make_fft_friendly) - - # ========================================================================= - # Dunder methods - # ========================================================================= - - def __repr__(self) -> str: - return f"SpaceGroup('{self.name}', number={self.number}, n_ops={self.n_ops})" + new = SpaceGroup.__new__(SpaceGroup) + new._sg_number = self._sg_number + new._sg_hm = self._sg_hm + new._sg_short_name = self._sg_short_name + new._sg_xhm = self._sg_xhm + new._sg_point_group = self._sg_point_group + new._sg_crystal_system = self._sg_crystal_system + new._sg_centrosymmetric = self._sg_centrosymmetric + # Through Symmetry's own initializer, so the operand-consistency checks and the + # device/dtype reconciliation in ``__post_init__`` run on the copy too. + Symmetry.__init__( + new, + matrices=self.matrices.clone(), + translations=self.translations.clone(), + ) + return new def __hash__(self) -> int: - """Hash based on space group number.""" - return hash(self._sg_number) + """Hash on :attr:`key` (the extended H-M symbol).""" + return hash(self._sg_xhm) def __eq__(self, other) -> bool: - """Equality based on space group number.""" + """Equality on :attr:`key`; also compares to a ``gemmi.SpaceGroup``.""" if isinstance(other, SpaceGroup): - return self._sg_number == other._sg_number + return self._sg_xhm == other._sg_xhm if isinstance(other, gemmi.SpaceGroup): - return self._sg_number == other.number + return self._sg_xhm == other.xhm() return False - # ========================================================================= - # Device movement - # ========================================================================= + def __repr__(self) -> str: + return f"SpaceGroup('{self.name}', number={self.number}, n_ops={self.n_ops})" - def copy(self) -> "SpaceGroup": - """A new SpaceGroup with the same symmetry, dtype and device (fresh buffers).""" - new_sg = SpaceGroup(self._sg_hm, dtype=self._dtype, device=self._device) - return new_sg + +__all__ = ["SpaceGroup", "SpaceGroupLike"] diff --git a/torchref/symmetry/symmetry.py b/torchref/symmetry/symmetry.py index 5213a130..4cb767d8 100644 --- a/torchref/symmetry/symmetry.py +++ b/torchref/symmetry/symmetry.py @@ -1,16 +1,801 @@ -"""DEPRECATED: ``Symmetry`` is a bare alias for :class:`SpaceGroup`. +"""Symmetry groups and everything derivable from their operations alone. -``Symmetry = SpaceGroup`` literally, so ``isinstance`` checks against either name -succeed and existing calls keep working. No ``DeprecationWarning`` is emitted -- -nothing tells a caller to migrate. Prefer ``SpaceGroup`` in new code. +:class:`Symmetry` holds a group as rotation matrices and fractional translations and +derives what needs nothing else: expansion of positions and Miller indices, +translation phases, the reflection predicates (centric, systematically absent, +multiplicity), symmetry-compatible grid sizes, and real-space map symmetrization. +Nothing here knows about crystals, so a group assembled from a raw operation list +serves non-crystallographic symmetry equally well. +:class:`~torchref.symmetry.spacegroup.SpaceGroup` specialises it with the +crystallographic identity and the CCP4 asymmetric-unit conventions. + +Rotation composes; translation does not. Translation acts *additively* on real-space +positions and as a *phase* ``exp(2 pi i h.t)`` in reciprocal space, so the two are +separate primitives rather than one method behind a flag. And ``h' = R^T h`` is not a +third law: it is :meth:`Symmetry.apply_rotations` on :attr:`Symmetry.reciprocal`, the +same group carrying the transposed rotations. Routing every caller through those +primitives is what keeps the real/reciprocal transpose from being re-decided, and got +wrong, at each site. + +Every expansion returns operations on the *leading* axis -- ``(n_ops, ...)``. """ -import warnings +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from fractions import Fraction +from typing import TYPE_CHECKING + +import torch + +from torchref.config import get_float_dtype +from torchref.utils.device_mixin import DeviceMixin + +if TYPE_CHECKING: + from torchref.symmetry.cell import Cell + +# Fractional translations are exact multiples of 1/24 (the denominator gemmi stores +# them over), which is what lets :meth:`Symmetry.grid_requirements` recover exact +# denominators instead of guessing at a float tolerance. +_TRANSLATION_DENOMINATOR = 24 + +# Tolerance for "this dot product is an integer" when testing a translation against a +# reflection. Miller indices are small and translations are k/24, so the true values +# are either integral or off by at least 1/24 -- far outside float32 noise. +_PHASE_TOL = 1e-6 + + +def is_fft_friendly(n: int) -> bool: + """Whether ``n`` factors into 2, 3 and 5 only, as radix-2,3,5 FFTs want. + + Parameters + ---------- + n : int + Candidate grid length. + + Returns + ------- + bool + True for 128 and 135, False for 131 and for any ``n <= 0``. + """ + if n <= 0: + return False + + for factor in (2, 3, 5): + while n % factor == 0: + n //= factor + + return n == 1 + + +def find_fft_friendly_size(n: int, divisibility: int = 1) -> int: + """Smallest FFT-friendly size at or above ``n`` that ``divisibility`` divides. + + Parameters + ---------- + n : int + Minimum grid length. + divisibility : int, default 1 + Required divisor, e.g. 2 for a screw axis. + + Returns + ------- + int + 131 gives 135; 131 at divisibility 2 gives 160. + """ + candidate = n + if candidate % divisibility != 0: + candidate = ((candidate // divisibility) + 1) * divisibility + + while not is_fft_friendly(candidate): + candidate += divisibility + + return candidate + + +@dataclass(eq=False, repr=False) +class Symmetry(DeviceMixin): + """A symmetry group as operations, plus everything they imply. + + Mutable by design; prefer :meth:`copy` over editing in place. Derived quantities + live in one cache that :meth:`reset_cache` clears -- and that ``.to()`` clears for + you, since :class:`~torchref.utils.device_mixin.DeviceMixin` invalidates caches on + every move. + + Parameters + ---------- + matrices : torch.Tensor + Rotation matrices, shape ``(n_ops, 3, 3)``, in the fractional basis. + translations : torch.Tensor + Fractional translations, shape ``(n_ops, 3)``. Coerced onto ``matrices``' + device and dtype, so the two can never end up split. + + Attributes + ---------- + matrices, translations : torch.Tensor + The operations, as above. + + Notes + ----- + Holds no refinable parameters -- symmetry operations are fixed constants -- so it + is a plain dataclass rather than an ``nn.Module``, and there is no gradient path + through it. + + Writing into :attr:`matrices` or :attr:`translations` in place does **not** + invalidate the cache, so the reciprocal stack and any built operator would keep + answering for the old operations. Nothing here mutates them, and :meth:`copy` is + the intended way to vary a group; if you must edit in place, call + :meth:`reset_cache` afterwards. + """ + + matrices: torch.Tensor + translations: torch.Tensor + _cache: dict = field(default_factory=dict, repr=False) + + def __post_init__(self) -> None: + """Validate the operation shapes and put both tensors on one device/dtype.""" + if self.matrices.ndim != 3 or self.matrices.shape[-2:] != (3, 3): + raise ValueError( + f"matrices must have shape (n_ops, 3, 3), got " + f"{tuple(self.matrices.shape)}" + ) + if self.translations.ndim != 2 or self.translations.shape[-1] != 3: + raise ValueError( + f"translations must have shape (n_ops, 3), got " + f"{tuple(self.translations.shape)}" + ) + if self.translations.shape[0] != self.matrices.shape[0]: + raise ValueError( + f"matrices and translations disagree on n_ops: " + f"{self.matrices.shape[0]} vs {self.translations.shape[0]}" + ) + + # One device and dtype for the pair. A split here would surface much later, as + # a device mismatch inside whichever expansion happened to touch both. + self.translations = self.translations.to( + device=self.matrices.device, dtype=self.matrices.dtype + ) + + # ========================================================================= + # Identity + # ========================================================================= + + @property + def n_ops(self) -> int: + """Number of symmetry operations.""" + return int(self.matrices.shape[0]) + + @property + def device(self) -> torch.device: + """Device the operations live on.""" + return self.matrices.device + + @property + def dtype(self) -> torch.dtype: + """Floating dtype of the operations.""" + return self.matrices.dtype + + # ========================================================================= + # The primitives + # ========================================================================= + + @property + def reciprocal(self) -> "Symmetry": + """The same group carrying the transposed rotations, for reciprocal space. + + ``h' = R^T h``, so Miller-index expansion is :meth:`apply_rotations` on this + object rather than a law of its own. Cached. + + Returns + ------- + Symmetry + Group with ``matrices = R^T`` and this group's translations. + + Notes + ----- + The translations come along because :meth:`phase_factors` needs them, but in + reciprocal space a translation is a *phase*, not a shift: calling + :meth:`apply_translations` on this object is meaningless. + """ + cached = self._cache.get("reciprocal") + if cached is None: + cached = Symmetry( + matrices=self.matrices.transpose(-2, -1).contiguous(), + translations=self.translations, + ) + self._cache["reciprocal"] = cached + return cached + + def apply_rotations(self, v: torch.Tensor) -> torch.Tensor: + """Rotate ``v`` by every operation: ``R v``. + + The one batched matmul behind every expansion. Translation is not applied -- + compose with :meth:`apply_translations` for real-space positions, or use + :attr:`reciprocal` for Miller indices. + + Parameters + ---------- + v : torch.Tensor + Vectors of shape ``(N, 3)``, in the fractional basis. + + Returns + ------- + torch.Tensor + Shape ``(n_ops, N, 3)``, rotated by each operation in turn. + """ + v = v.to(device=self.device, dtype=self.dtype) + # result[o, n, i] = sum_j matrices[o, i, j] * v[n, j] + return torch.einsum("oij,nj->oni", self.matrices, v) + + def apply_translations(self, v: torch.Tensor) -> torch.Tensor: + """Add each operation's fractional translation: ``v + t``. + + Real space only; in reciprocal space a translation is a phase, see + :meth:`phase_factors`. + + Parameters + ---------- + v : torch.Tensor + Shape ``(n_ops, N, 3)`` -- typically straight out of + :meth:`apply_rotations` -- or ``(N, 3)`` to broadcast one set of vectors + across all operations. + + Returns + ------- + torch.Tensor + Shape ``(n_ops, N, 3)``. + """ + v = v.to(device=self.device, dtype=self.dtype) + if v.ndim == 2: + v = v.unsqueeze(0).expand(self.n_ops, -1, -1) + elif v.shape[0] != self.n_ops: + raise ValueError( + f"expected a leading axis of n_ops={self.n_ops} or a bare (N, 3), " + f"got {tuple(v.shape)}" + ) + return v + self.translations.unsqueeze(1) + + def phase_factors(self, hkl: torch.Tensor) -> torch.Tensor: + """Structure-factor phase shift per operation: ``exp(2 pi i h.t)``. + + The reciprocal-space action of a translation. Pairs with + ``self.reciprocal.apply_rotations(hkl)`` to combine symmetry-equivalent + structure factors. + + Parameters + ---------- + hkl : torch.Tensor + Miller indices, shape ``(N, 3)``, integer or float. + + Returns + ------- + torch.Tensor + Complex phase factors of shape ``(n_ops, N)``. + + Notes + ----- + The phase stays at this group's floating dtype rather than being forced to + float32, so a float64 configuration yields complex128 instead of silently + narrowing to complex64. + + This is the *complex factor* form, ``exp(+2 pi i h.t)``, for combining + structure factors. Expanding the *phases* of reflection data instead needs a + signed offset in radians, ``-2 pi h.t``, which + :meth:`~torchref.symmetry.spacegroup.SpaceGroup.expand_hkl` computes itself. + The two are not interchangeable and the sign difference is invisible in + P21/P212121/C2 -- see ``tests/unit/symmetry/test_phase_convention.py``. + """ + hkl = hkl.to(device=self.device, dtype=self.dtype) + h_dot_t = torch.matmul(hkl, self.translations.T).T # (n_ops, N) + return torch.exp(1j * (2.0 * math.pi * h_dot_t)) + + # ========================================================================= + # Named expansions + # ========================================================================= + + def expand_positions(self, xyz_fractional: torch.Tensor) -> torch.Tensor: + """Expand fractional positions by every operation: ``R x + t``. + + Fixes the composition order in one place -- ``R x + t``, not ``R (x + t)``. + + Parameters + ---------- + xyz_fractional : torch.Tensor + Fractional coordinates, shape ``(N, 3)``. + + Returns + ------- + torch.Tensor + Shape ``(n_ops, N, 3)``, unwrapped (values may fall outside ``[0, 1)``). + """ + return self.apply_translations(self.apply_rotations(xyz_fractional)) + + def expand_directions(self, v_fractional: torch.Tensor) -> torch.Tensor: + """Expand fractional directions or displacements: ``R v``, no translation. + + Parameters + ---------- + v_fractional : torch.Tensor + Fractional vectors, shape ``(N, 3)``. + + Returns + ------- + torch.Tensor + Shape ``(n_ops, N, 3)``. + """ + return self.apply_rotations(v_fractional) + + def expand_to_P1(self, xyz_fractional: torch.Tensor) -> torch.Tensor: + """Flatten :meth:`expand_positions` into one P1 coordinate list. + + Parameters + ---------- + xyz_fractional : torch.Tensor + Fractional coordinates, shape ``(N, 3)``. + + Returns + ------- + torch.Tensor + Shape ``(n_ops * N, 3)``, operation-major. + """ + return self.expand_positions(xyz_fractional).reshape(-1, 3) + + def expand_reciprocal(self, hkl: torch.Tensor) -> torch.Tensor: + """Symmetry-equivalent Miller indices: ``h' = R^T h``. + + Parameters + ---------- + hkl : torch.Tensor + Miller indices, shape ``(N, 3)``. + + Returns + ------- + torch.Tensor + Shape ``(n_ops, N, 3)``, rounded to ``int64``. Rounding is exact for valid + operations on integer indices and only mops up float error. + """ + equivalents = self.reciprocal.apply_rotations(hkl) + return torch.round(equivalents).to(torch.int64) # dtype-ok: rounded Miller equivalents; int64 for exact integer compare/index + + # ========================================================================= + # Reflection predicates + # ========================================================================= + + def is_centric(self, hkl: torch.Tensor) -> torch.Tensor: + """Whether each reflection is centric, i.e. some operation maps ``h -> -h``. + + A centric reflection has its phase restricted to 0 or pi. + + Parameters + ---------- + hkl : torch.Tensor + Miller indices, shape ``(..., 3)``. + + Returns + ------- + torch.Tensor + Boolean mask of shape ``(...)``, on ``hkl``'s device. + """ + original_shape = hkl.shape[:-1] + with torch.no_grad(): + flat = hkl.reshape(-1, 3) + equivalents = self.expand_reciprocal(flat) # (n_ops, N, 3) + target = -flat.to(device=equivalents.device, dtype=torch.int64) # dtype-ok: compare target for int64 equivalents; dtype must match + centric = (equivalents == target).all(dim=-1).any(dim=0) + return centric.reshape(original_shape).to(hkl.device) + + def is_absent(self, hkl: torch.Tensor) -> torch.Tensor: + """Whether each reflection is systematically absent. + + Absent when some operation maps ``h -> h`` while ``h.t`` is non-integral: the + reflection is destroyed by interference from that translation. + + Parameters + ---------- + hkl : torch.Tensor + Miller indices, shape ``(..., 3)``. + + Returns + ------- + torch.Tensor + Boolean mask of shape ``(...)``, on ``hkl``'s device. + """ + original_shape = hkl.shape[:-1] + with torch.no_grad(): + flat = hkl.reshape(-1, 3) + equivalents = self.expand_reciprocal(flat) # (n_ops, N, 3) + target = flat.to(device=equivalents.device, dtype=torch.int64) # dtype-ok: compare target for int64 equivalents; dtype must match + maps_to_self = (equivalents == target).all(dim=-1) # (n_ops, N) + + h_dot_t = torch.matmul( + flat.to(device=self.device, dtype=self.dtype), self.translations.T + ).T # (n_ops, N) + non_integral = (h_dot_t - torch.round(h_dot_t)).abs() > _PHASE_TOL + + absent = (maps_to_self & non_integral).any(dim=0) + return absent.reshape(original_shape).to(hkl.device) + + def epsilon(self, hkl: torch.Tensor, *, friedel: bool = True) -> torch.Tensor: + """Reflection multiplicity: operations mapping ``h`` to ``h``, or also to ``-h``. + + Parameters + ---------- + hkl : torch.Tensor + Miller indices, shape ``(N, 3)``. + friedel : bool, default True + Whether operations mapping ``h -> -h`` count alongside ``h -> h``. + + Returns + ------- + torch.Tensor + Multiplicities of shape ``(N,)`` at the configured float dtype, floored at + 1, returned on ``hkl``'s device so it can weight data sitting beside it. + + Notes + ----- + The two settings answer different questions and both are wanted. + + Operations mapping ``h -> h`` add coherently and set the **mean**, + ``<|F|^2> = eps * Sigma``. That is the conventional crystallographic epsilon, + and what a Wilson normalisation or a likelihood's variance budget asks for. + Operations mapping ``h -> -h`` leave the mean alone and instead make ``F`` + real, changing the **distribution** from exponential to chi2_1 -- that is + centricity, and :meth:`is_centric` already carries it. Folding Friedel into + epsilon therefore mixes a mean effect with a distribution effect. + + The default keeps Friedel folded in because downstream sigma_A estimation is + calibrated against that convention; flipping it would silently decalibrate the + refinement path. Pass ``friedel=False`` for the conventional count, as the + molecular-replacement likelihood does -- counting Friedel there doubles + epsilon on exactly the reflections whose distribution the Woolfson branch is + already handling, inflating their ``V = eps - sigma_A**2``. + + The two differ on centric reflections and *only* there: measured across the + ten benchmark structures every disagreement was centric, and the counts are + not small -- 12360 reflections on 2DQ6, 7555 on 4BX9, 6680 on 3K7M. + + Both settings count lattice-centring cosets, so on a centred lattice every + reflection carries the centring order as a factor: C2 gives 2 for general + reflections where a primitive lattice gives 1. That is a separate axis from + this switch. Being uniform per lattice it is absorbed into ``Sigma`` wherever + epsilon is a factor -- which is why the refinement path never saw it -- and + bites only where epsilon is a term. + """ + float_dtype = get_float_dtype() + with torch.no_grad(): + equivalents = self.expand_reciprocal(hkl) # (n_ops, N, 3) + target = hkl.to(device=equivalents.device, dtype=torch.int64) # dtype-ok: compare target for int64 equivalents; dtype must match + fixes = (equivalents == target).all(dim=-1) + if friedel: + fixes = fixes | (equivalents == -target).all(dim=-1) + eps = fixes.sum(dim=0).clamp(min=1).to(float_dtype) + return eps.to(hkl.device) + + # ========================================================================= + # Grid compatibility + # ========================================================================= + + def grid_requirements(self) -> dict: + """Per-axis grid divisibility for interpolation-free symmetry expansion. + + Read off the denominators of the fractional translations, so a grid meeting + them indexes every symmetry mate at an exact integer. + + Returns + ------- + dict + ``{'nx_mod': int, 'ny_mod': int, 'nz_mod': int}`` -- P21 gives + ``(1, 2, 1)``, P212121 gives ``(2, 2, 2)``. + """ + mods = [1, 1, 1] + # Recover each denominator from the integer numerator over 1/24 rather than + # from the float directly: 1/3 is not representable in float32, so + # ``Fraction(float)`` would need a tolerance where this is exact. + numerators = torch.round( + self.translations.detach().cpu().double() * _TRANSLATION_DENOMINATOR + ).to(torch.int64) # dtype-ok: integer translation numerators for exact Fraction recovery + + for op_numerators in numerators.tolist(): + for axis, numerator in enumerate(op_numerators): + if numerator % _TRANSLATION_DENOMINATOR == 0: + continue + denominator = Fraction( + int(numerator), _TRANSLATION_DENOMINATOR + ).denominator + mods[axis] = math.lcm(mods[axis], denominator) + + return {"nx_mod": mods[0], "ny_mod": mods[1], "nz_mod": mods[2]} + + def check_grid_compatibility(self, grid_shape: tuple) -> dict: + """Check a grid against this group's divisibility and against the FFT. + + Parameters + ---------- + grid_shape : tuple of int + Grid dimensions ``(nx, ny, nz)``. + + Returns + ------- + dict + ``compatible`` (both tests pass), ``symmetry_compatible``, + ``fft_friendly``, ``can_use_direct_indexing`` (interpolation-free expansion + possible; equal to ``symmetry_compatible``), ``issues`` (per-axis + descriptions, empty when compatible) and ``requirements`` (from + :meth:`grid_requirements`). + """ + requirements = self.grid_requirements() + issues = [] + + for axis, name in enumerate(("nx", "ny", "nz")): + modulus = requirements[f"{name}_mod"] + length = int(grid_shape[axis]) + if length % modulus != 0: + issues.append( + f"{name}={length} not divisible by {modulus} " + f"(required by the symmetry)" + ) + + symmetry_compatible = len(issues) == 0 + + fft_friendly = True + for axis, name in enumerate(("nx", "ny", "nz")): + length = int(grid_shape[axis]) + if not is_fft_friendly(length): + fft_friendly = False + issues.append( + f"{name}={length} is not FFT-friendly (not a product of 2, 3, 5)" + ) + + return { + "compatible": symmetry_compatible and fft_friendly, + "symmetry_compatible": symmetry_compatible, + "fft_friendly": fft_friendly, + "can_use_direct_indexing": symmetry_compatible, + "issues": issues, + "requirements": requirements, + } + + def can_index_directly(self, grid_shape: tuple) -> bool: + """Whether ``grid_shape`` admits interpolation-free symmetry expansion. + + The question :meth:`symmetrize_map` answers internally when choosing an + implementation, exposed so callers can ask it without building an operator. + + Parameters + ---------- + grid_shape : tuple of int + Grid dimensions ``(nx, ny, nz)``. + + Returns + ------- + bool + True when every symmetry mate lands on an exact grid point. + """ + return bool(self.check_grid_compatibility(grid_shape)["symmetry_compatible"]) + + def suggest_grid_size( + self, min_grid_shape: tuple, make_fft_friendly: bool = True + ) -> tuple: + """Smallest grid at or above ``min_grid_shape`` meeting the divisibility. + + Parameters + ---------- + min_grid_shape : tuple of int + Minimum dimensions ``(nx, ny, nz)``. + make_fft_friendly : bool, default True + Also require factors of 2, 3 and 5 only. + + Returns + ------- + tuple of int + Suggested ``(nx, ny, nz)``. + """ + requirements = self.grid_requirements() + + def next_valid(length: int, divisibility: int) -> int: + if make_fft_friendly: + return find_fft_friendly_size(length, divisibility) + if length % divisibility == 0: + return length + return ((length // divisibility) + 1) * divisibility + + return tuple( + next_valid(int(min_grid_shape[axis]), requirements[f"{name}_mod"]) + for axis, name in enumerate(("nx", "ny", "nz")) + ) + + def optimal_grid_size( + self, cell: "Cell", max_res: float, make_fft_friendly: bool = True + ) -> tuple: + """Smallest grid that samples ``cell`` to ``max_res`` and suits this group. + + Composes the cell's Shannon-Nyquist minimum with + :meth:`suggest_grid_size`; the oversampling factor is the cell's, so every + grid-sizing path shares one setting. + + Parameters + ---------- + cell : Cell + Unit cell. + max_res : float + Maximum resolution in Angstroms. + make_fft_friendly : bool, default True + Also require factors of 2, 3 and 5 only. + + Returns + ------- + tuple of int + Grid dimensions ``(nx, ny, nz)``. + """ + return self.suggest_grid_size( + cell.compute_grid_size(max_res), make_fft_friendly=make_fft_friendly + ) + + # ========================================================================= + # Real-space maps + # ========================================================================= + + def map_operator(self, map_shape): + """Cached operator that applies this group to maps of ``map_shape``. + + Two implementations: exact integer indexing when the grid permits it, and + ``grid_sample`` interpolation otherwise. Which one you get depends on the grid, + so a mis-sized grid costs accuracy -- :meth:`can_index_directly` reports the + distinction, and :meth:`suggest_grid_size` fixes it. + + Parameters + ---------- + map_shape : tuple of int + Density map dimensions ``(nx, ny, nz)``. + + Returns + ------- + _MapSymmetryDirect or _MapSymmetryInterpolation + The operator, memoized for the most recent shape only. + + Notes + ----- + Only the last shape is kept. The interpolating operator holds sampling grids of + shape ``(n_ops, nx, ny, nz, 3)`` -- hundreds of megabytes at production grid + sizes -- so a dictionary keyed on shape would quietly make this object + expensive to hold and to move between devices. + """ + from torchref.symmetry.map_symmetry import build_map_operator + + key = tuple(int(n) for n in map_shape) + cached = self._cache.get("map_operator") + if cached is not None and cached[0] == key: + return cached[1] + + operator = build_map_operator(self, key) + self._cache["map_operator"] = (key, operator) + return operator + + def symmetrize_map( + self, density_map: torch.Tensor, combine: str = "sum" + ) -> torch.Tensor: + """Apply every operation to a density map and combine the mates. + + Parameters + ---------- + density_map : torch.Tensor + Asymmetric-unit density, shape ``(nx, ny, nz)``. + combine : {'sum', 'max'}, default 'sum' + ``'sum'`` for electron density, ``'max'`` for masks and boolean data. + + Returns + ------- + torch.Tensor + Symmetrized map, same shape as the input. Returned unchanged for a + one-operation group. + """ + if self.n_ops == 1: + return density_map + return self.map_operator(density_map.shape).symmetrize(density_map, combine) + + def expand_map_to_P1(self, density_map: torch.Tensor) -> torch.Tensor: + """Every symmetry mate of a density map, stacked. + + Parameters + ---------- + density_map : torch.Tensor + Asymmetric-unit density, shape ``(nx, ny, nz)``. + + Returns + ------- + torch.Tensor + Shape ``(n_ops, nx, ny, nz)``. + """ + return self.map_operator(density_map.shape).all_mates(density_map) + + # ========================================================================= + # Reciprocal-space extraction + # ========================================================================= + + def reciprocal_extractor(self, hkl: torch.Tensor, grid_shape: tuple): + """Cached extractor pulling symmetrized structure factors off a grid. + + Precomputes the equivalent indices, phases and flat gather indices for a fixed + ``hkl`` and ``grid_shape``, so each later call is one gather, multiply and sum. + + Parameters + ---------- + hkl : torch.Tensor + Target Miller indices, shape ``(N, 3)``. + grid_shape : tuple of int + Reciprocal grid dimensions ``(nx, ny, nz)``. + + Returns + ------- + ReciprocalSymmetryExtractor + Memoized against ``hkl``'s identity and ``grid_shape``; a different tensor + or shape rebuilds it. + """ + from torchref.base.reciprocal.symmetry import ReciprocalSymmetryExtractor + from torchref.utils.caching import ParameterFingerprint + + key = tuple(int(n) for n in grid_shape) + cached = self._cache.get("reciprocal_extractor") + if cached is not None: + cached_key, fingerprint, extractor = cached + if cached_key == key and fingerprint.matches([hkl]): + return extractor + + extractor = ReciprocalSymmetryExtractor(hkl, self, key) + self._cache["reciprocal_extractor"] = ( + key, + ParameterFingerprint([hkl]), + extractor, + ) + return extractor + + # ========================================================================= + # Cache and copy + # ========================================================================= + + def reset_cache(self) -> None: + """Drop every derived quantity. + + Called for you by :class:`~torchref.utils.device_mixin.DeviceMixin` on any + ``.to()``, including one targeting the current device. + """ + self._cache = {} + + def _apply(self, fn, recurse: bool = True): + """Clear the cache *before* the traversal moves anything. + + The base traversal walks ``__dict__`` first and invalidates caches afterwards, + which would transfer the cached sampling grids to the new device only to + discard them -- hundreds of megabytes of pointless copying at production grid + sizes. + """ + self.reset_cache() + return super()._apply(fn, recurse) + + def copy(self) -> "Symmetry": + """An independent copy with cloned operations and an empty cache. + + Returns + ------- + Symmetry + New instance; mutating its tensors cannot affect this one. + """ + return type(self)( + matrices=self.matrices.clone(), + translations=self.translations.clone(), + ) + + # ========================================================================= + # Dunder + # ========================================================================= + + def __len__(self) -> int: + """Number of symmetry operations.""" + return self.n_ops -from torchref.symmetry.spacegroup import SpaceGroup, SpaceGroupLike + def __repr__(self) -> str: + return f"Symmetry(n_ops={self.n_ops})" -# Backward compatibility alias - Symmetry is now SpaceGroup -# Using the class directly so isinstance() checks work. -Symmetry = SpaceGroup -__all__ = ["Symmetry", "SpaceGroupLike"] +__all__ = ["Symmetry", "find_fft_friendly_size", "is_fft_friendly"] diff --git a/torchref/topology/__init__.py b/torchref/topology/__init__.py new file mode 100644 index 00000000..9fc6c49c --- /dev/null +++ b/torchref/topology/__init__.py @@ -0,0 +1,63 @@ +"""Model topology as a graph: residues over atoms, connectivity over restraints. + +:class:`Topology` holds two levels. :class:`ResidueGraph` is the sequence -- residues as +template instances, inter-residue links as edges. :class:`AtomGraph` is the expansion -- +atoms as nodes, typed :class:`EdgeBlock` sets over them, and a CSR bond adjacency that +answers ``neighbors(i)``. + +The topology is **target-free**: it says what is connected, not what the ideal geometry +is. Ideal values and sigmas belong to a restraint layer keyed to the same edges, so one +connectivity can carry monomer-library targets, force-field parameters, or +ADP-similarity sigmas without duplicating the edges. + +Build one with :func:`build_topology`. + +Hydrogens come in two forms over the same graph. :func:`plan_hydrogens` instantiates +the monomer templates to add them as real atoms, which is the default. For a model +loaded heavy-only, :mod:`torchref.topology.riding` reconstructs them from their parents +at each non-bonded evaluation instead, so their sterics still count. Only one applies at +a time. +""" + +from .atom_graph import AtomGraph +from .build import build_topology, build_topology_with_values +from .edges import ORIGIN_ORDER, EdgeBlock +from .hydrogens import ( + HydrogenPlan, + augment_atom_table, + optimise_free_torsions, + plan_hydrogens, +) +from .residue_graph import ResidueGraph +from .restraints import Restraints +from .riding import ( + HydrogenTopology, + build_h_candidate_pairs, + build_hydrogen_topology, + place_riding_hydrogens, +) +from .restraint_sets import assemble_entries, max_period +from .templates import resolve_template_keys +from .topology import Topology + +__all__ = [ + "Topology", + "Restraints", + "ResidueGraph", + "AtomGraph", + "EdgeBlock", + "ORIGIN_ORDER", + "build_topology", + "build_topology_with_values", + "assemble_entries", + "max_period", + "HydrogenPlan", + "plan_hydrogens", + "optimise_free_torsions", + "augment_atom_table", + "HydrogenTopology", + "build_hydrogen_topology", + "build_h_candidate_pairs", + "place_riding_hydrogens", + "resolve_template_keys", +] diff --git a/torchref/topology/atom_graph.py b/torchref/topology/atom_graph.py new file mode 100644 index 00000000..05204315 --- /dev/null +++ b/torchref/topology/atom_graph.py @@ -0,0 +1,382 @@ +"""The atom level of a topology: atoms as nodes, typed edge blocks over them. + +Bonds are promoted to a real adjacency structure -- a CSR pair built once from the bond +block -- so :meth:`AtomGraph.neighbors` answers "what is atom *i* bonded to" without +inferring it from restraint index lists. Angles, torsions, chirals and planes stay typed +hyperedge sets read from the monomer library, because the library deliberately does not +restrain every path the bond graph implies. + +Every indexing structure here is a tensor, so it moves with ``.to(device)`` +alongside the edge blocks. Only the per-atom identifiers are NumPy, because they are +strings. +""" + +from dataclasses import dataclass, field +from typing import Dict, Optional, Set, Tuple + +import numpy as np +import torch + +from torchref.topology.edges import EdgeBlock +from torchref.utils.device_mixin import DeviceMixin + + +def _build_csr(bonds: torch.Tensor, n_atoms: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Symmetric CSR adjacency from an ``(E, 2)`` bond list. + + Parameters + ---------- + bonds : torch.Tensor + Bond atom indices, shape ``(E, 2)``, dtype ``int64``. + n_atoms : int + Number of atoms, so isolated trailing atoms still get an entry. + + Returns + ------- + indptr, indices : torch.Tensor + ``indices[indptr[i]:indptr[i + 1]]`` are atom ``i``'s bonded neighbours, + ascending, each partner listed once. A bond row repeated in the edge list -- + once per altloc conformer for a bond between two shared atoms, or from a LINK + record that appears twice -- therefore does not inflate an atom's degree. + """ + device = bonds.device + if bonds.numel() == 0: + return ( + torch.zeros(n_atoms + 1, dtype=torch.int64, device=device), # dtype-ok: CSR indptr offset array; int64 index required + torch.zeros(0, dtype=torch.int64, device=device), # dtype-ok: empty CSR neighbor index array; int64 index required + ) + + src = torch.cat([bonds[:, 0], bonds[:, 1]]) + dst = torch.cat([bonds[:, 1], bonds[:, 0]]) + + # Unique directed pairs, which ``torch.unique`` returns in lexicographic + # (src, dst) order -- the CSR layout wanted below. + pairs = torch.unique(torch.stack([src, dst], dim=1), dim=0) + src, dst = pairs[:, 0], pairs[:, 1] + + counts = torch.bincount(src, minlength=n_atoms) + indptr = torch.zeros(n_atoms + 1, dtype=torch.int64, device=device) # dtype-ok: CSR indptr offset array; int64 index required + torch.cumsum(counts, dim=0, out=indptr[1:]) + return indptr, dst.to(torch.int64) # dtype-ok: CSR neighbor (dst) index array; int64 index required + + +def _extend_paths( + indptr: torch.Tensor, indices: torch.Tensor, paths: torch.Tensor +) -> torch.Tensor: + """Extend each bonded path by one bonded step, without doubling back. + + Parameters + ---------- + indptr, indices : torch.Tensor + CSR adjacency. + paths : torch.Tensor + Existing paths, shape ``(P, L)`` with ``L >= 2``, each row a chain of bonded + atoms. + + Returns + ------- + torch.Tensor + Shape ``(P', L + 1)``. A path is extended by every bonded neighbour of its last + atom except the one it just came from, so ``(i, j, k)`` never yields ``k = i``. + """ + device = paths.device + if paths.numel() == 0: + return torch.zeros((0, paths.shape[1] + 1), dtype=torch.int64, device=device) # dtype-ok: empty BFS path index array; int64 index required + + last, prev = paths[:, -1], paths[:, -2] + counts = indptr[last + 1] - indptr[last] + total = int(counts.sum()) + if total == 0: + return torch.zeros((0, paths.shape[1] + 1), dtype=torch.int64, device=device) # dtype-ok: empty BFS path index array; int64 index required + + row = torch.repeat_interleave(torch.arange(len(paths), device=device), counts) + # Offset of each slot within its own neighbour list. + exclusive = torch.cumsum(counts, dim=0) - counts + pos = torch.arange(total, device=device) - torch.repeat_interleave( + exclusive, counts + ) + nxt = indices[torch.repeat_interleave(indptr[last], counts) + pos] + + keep = nxt != prev[row] + return torch.cat([paths[row][keep], nxt[keep, None]], dim=1) + + +@dataclass(eq=False, repr=False) +class AtomGraph(DeviceMixin): + """Atoms as nodes, typed edges over them, with bond adjacency. + + Parameters + ---------- + name, element, altloc : numpy.ndarray + Per-atom identifiers, shape ``(N,)``. Strings, so NumPy rather than tensors; + residue-level identity is reached through ``residue_of`` rather than duplicated + here. + residue_of : torch.Tensor + Residue index per atom, shape ``(N,)``, dtype ``int64``. + bonds, angles, torsions, chirals : EdgeBlock + Typed edge blocks. ``bonds`` also backs the adjacency. + planes : dict + ``{n_atoms_in_plane: EdgeBlock}`` -- planes are ragged, so they are grouped by + atom count the way the plane restraints already are. + energy_type : numpy.ndarray, optional + CCP4 energy type per atom (``NH1``, ``OC``, ``CH3``, ...), shape ``(N,)``, + ``''`` where the template does not say. Keys the contact radii and the + hydrogen-bond roles. + template_h_count : torch.Tensor, optional + How many hydrogens the atom carries in its template, shape ``(N,)``, + ``int8``; ``-1`` where unknown. Together with the bonded hydrogens actually + present this gives :meth:`implicit_h_count`. + hb_type : torch.Tensor, optional + Hydrogen-bond role code per atom, shape ``(N,)``, ``int8``; see the contact + policy for the enumeration. None until assigned. + + Notes + ----- + Holds no refinable parameters, so this is a dataclass rather than an ``nn.Module``. + The adjacency is derived from ``bonds`` at construction and rebuilt by + :meth:`rebuild_adjacency` if the bond block is replaced. + """ + + name: np.ndarray + element: np.ndarray + altloc: np.ndarray + residue_of: torch.Tensor + bonds: EdgeBlock + angles: EdgeBlock + torsions: EdgeBlock + chirals: EdgeBlock + planes: Dict[int, EdgeBlock] = field(default_factory=dict) + energy_type: Optional[np.ndarray] = None + template_h_count: Optional[torch.Tensor] = None + hb_type: Optional[torch.Tensor] = None + + _adj_indptr: Optional[torch.Tensor] = field(default=None, repr=False) + _adj_indices: Optional[torch.Tensor] = field(default=None, repr=False) + + def __post_init__(self) -> None: + if self._adj_indptr is None: + self.rebuild_adjacency() + + @property + def device(self) -> torch.device: + """Where the indexing tensors live. Derived from the bond block.""" + return self.bonds.indices.device + + @property + def n_atoms(self) -> int: + """Number of atom nodes.""" + return len(self.name) + + @property + def is_hydrogen(self) -> torch.Tensor: + """Boolean mask of hydrogen atoms, shape ``(N,)``.""" + flags = np.char.upper(np.char.strip(self.element.astype(str))) == "H" + return torch.as_tensor(flags, device=self.bonds.indices.device) + + def copy(self) -> "AtomGraph": + """An independent copy sharing no storage with this one.""" + return AtomGraph( + name=self.name.copy(), + element=self.element.copy(), + altloc=self.altloc.copy(), + residue_of=self.residue_of.clone(), + bonds=self.bonds.copy(), + angles=self.angles.copy(), + torsions=self.torsions.copy(), + chirals=self.chirals.copy(), + planes={size: block.copy() for size, block in self.planes.items()}, + energy_type=None if self.energy_type is None else self.energy_type.copy(), + template_h_count=( + None if self.template_h_count is None else self.template_h_count.clone() + ), + hb_type=None if self.hb_type is None else self.hb_type.clone(), + ) + + def implicit_h_count(self) -> Optional[torch.Tensor]: + """Hydrogens each atom should carry but the table does not hold, ``(N,)``. + + ``template_h_count`` minus the bonded hydrogens actually present, floored at + zero; ``0`` where the template count is unknown. None when the graph carries + no template counts. What decides whether an atom takes its with-hydrogen + contact radius. + """ + if self.template_h_count is None: + return None + is_h = self.is_hydrogen + bonds = self.bonds.indices + present = torch.zeros(self.n_atoms, dtype=torch.int64, device=bonds.device) # dtype-ok: bincount output; int64 + if bonds.numel(): + heavy_of_h = torch.cat( + [bonds[is_h[bonds[:, 1]] & ~is_h[bonds[:, 0]], 0], + bonds[is_h[bonds[:, 0]] & ~is_h[bonds[:, 1]], 1]] + ) + if heavy_of_h.numel(): + present = torch.bincount(heavy_of_h, minlength=self.n_atoms) + known = self.template_h_count >= 0 + missing = self.template_h_count.to(torch.int64) - present + return torch.where(known, missing.clamp(min=0), torch.zeros_like(missing)) + + def subset(self, remap: torch.Tensor, residue_remap: torch.Tensor) -> "AtomGraph": + """The atoms ``remap`` keeps, with every edge set reindexed. + + Parameters + ---------- + remap : torch.Tensor + Old atom index to new, shape ``(N_old,)``, ``-1`` where dropped. + residue_remap : torch.Tensor + Old residue index to new, shape ``(R_old,)``, ``-1`` where dropped. + + Returns + ------- + AtomGraph + Adjacency is rebuilt from the surviving bond block rather than subsetted: + CSR row offsets are not meaningful once the atoms are renumbered. + """ + keep = (remap >= 0).cpu().numpy() + planes = {} + for size, block in self.planes.items(): + reduced = block.subset(remap) + if reduced.n_edges: + planes[size] = reduced + + keep_t = torch.as_tensor(keep, device=self.residue_of.device) + return AtomGraph( + name=self.name[keep], + element=self.element[keep], + altloc=self.altloc[keep], + residue_of=residue_remap[self.residue_of[keep_t]], + bonds=self.bonds.subset(remap), + angles=self.angles.subset(remap), + torsions=self.torsions.subset(remap), + chirals=self.chirals.subset(remap), + planes=planes, + energy_type=None if self.energy_type is None else self.energy_type[keep], + template_h_count=( + None if self.template_h_count is None else self.template_h_count[keep_t] + ), + hb_type=None if self.hb_type is None else self.hb_type[keep_t], + ) + + def rebuild_adjacency(self) -> None: + """Rebuild the CSR adjacency from the current bond block.""" + self._adj_indptr, self._adj_indices = _build_csr( + self.bonds.indices, self.n_atoms + ) + + def neighbors(self, i: int) -> torch.Tensor: + """Atoms bonded to atom ``i``, ascending. + + Returns + ------- + torch.Tensor + Neighbour indices, a view into the adjacency, on the graph's device. + """ + return self._adj_indices[self._adj_indptr[i] : self._adj_indptr[i + 1]] + + def degree(self, i: int = None) -> torch.Tensor: + """Bonded-neighbour count, for atom ``i`` or for every atom.""" + deg = self._adj_indptr[1:] - self._adj_indptr[:-1] + return deg if i is None else deg[i] + + def _directed_bonds(self) -> torch.Tensor: + """Bonds as ``(2E, 2)`` directed pairs.""" + b = self.bonds.indices + return torch.cat([b, b.flip(1)], dim=0) + + # ------------------------------------------------------------------ + # Non-bonded exclusions + # ------------------------------------------------------------------ + + @staticmethod + def _pair_set(pairs: torch.Tensor) -> Set[Tuple[int, int]]: + """``(low, high)`` tuples of a ``(P, 2)`` index tensor, self-pairs dropped.""" + if pairs.numel() == 0: + return set() + lo = torch.minimum(pairs[:, 0], pairs[:, 1]) + hi = torch.maximum(pairs[:, 0], pairs[:, 1]) + keep = lo != hi + return set(zip(lo[keep].cpu().tolist(), hi[keep].cpu().tolist())) + + def exclusions_from_restraint_edges(self) -> Set[Tuple[int, int]]: + """1-2, 1-3 and 1-4 pairs taken from the bond, angle and torsion **edges**. + + 1-2 from every bond, 1-3 from each angle's outer pair, 1-4 from each torsion's + outer pair. Reproduces exactly the set the non-bonded term has always been + given. + + This is *not* the same as :meth:`exclusions_12_13_14`: a pair that is 1-3 bonded + but whose angle the monomer library does not restrain appears there and not + here, and so takes a repulsion it should not. Kept because switching the + non-bonded term to the connectivity-derived set changes its value and wants its + own measurement. + + Returns + ------- + set of tuple of int + ``(low, high)`` atom index pairs. + """ + excl: Set[Tuple[int, int]] = set() + for block, cols in ( + (self.bonds, (0, 1)), + (self.angles, (0, 2)), + (self.torsions, (0, 3)), + ): + if block.n_edges: + excl |= self._pair_set(block.indices[:, cols]) + return excl + + def exclusions_12_13_14(self) -> Set[Tuple[int, int]]: + """1-2, 1-3 and 1-4 pairs derived from bond **connectivity** alone. + + Walks the adjacency two and three steps out, so the result does not depend on + which angles and torsions the monomer library happens to restrain. This is the + physically correct exclusion set; :meth:`exclusions_from_restraint_edges` is the + one currently wired into the non-bonded term. + + Returns + ------- + set of tuple of int + ``(low, high)`` atom index pairs. + """ + p2 = self._directed_bonds() + p3 = _extend_paths(self._adj_indptr, self._adj_indices, p2) + p4 = _extend_paths(self._adj_indptr, self._adj_indices, p3) + return ( + self._pair_set(p2) + | self._pair_set(p3[:, (0, 2)]) + | self._pair_set(p4[:, (0, 3)]) + ) + + def hydrogen_parents(self) -> Dict[int, torch.Tensor]: + """``{hydrogen atom: heavy neighbours of its bonded parent}``. + + Taken from bond connectivity, so it does not depend on current coordinates the + way a distance criterion does. + + Returns + ------- + dict + Empty when the graph carries no hydrogens. + """ + is_h = self.is_hydrogen + out: Dict[int, torch.Tensor] = {} + for h in torch.nonzero(is_h, as_tuple=False).flatten().tolist(): + nb = self.neighbors(h) + heavy = nb[~is_h[nb]] + if heavy.numel() == 0: + continue + parent = int(heavy[0]) + parent_nb = self.neighbors(parent) + out[h] = parent_nb[~is_h[parent_nb] & (parent_nb != h)] + return out + + def __repr__(self) -> str: + return ( + f"AtomGraph(n_atoms={self.n_atoms}, bonds={self.bonds.n_edges}, " + f"angles={self.angles.n_edges}, torsions={self.torsions.n_edges}, " + f"chirals={self.chirals.n_edges}, " + f"planes={sum(b.n_edges for b in self.planes.values())})" + ) + + +__all__ = ["AtomGraph"] diff --git a/torchref/topology/build.py b/torchref/topology/build.py new file mode 100644 index 00000000..87970297 --- /dev/null +++ b/torchref/topology/build.py @@ -0,0 +1,1025 @@ +"""Assemble a :class:`~torchref.topology.topology.Topology` from an atom table. + +Intra-residue edges are matched here, template by template, through the Numba matchers +in :mod:`torchref.topology.builders_numba`. Inter-residue edges come from the +``InterResidue*Builder`` classes, which already encode the link geometry and are reused +rather than reimplemented. +""" + +from typing import Dict, List, Optional, Sequence, Tuple + +import numpy as np +import pandas as pd +import torch + +from torchref.topology.builders import ( + InterResidueAngleBuilder, + InterResidueBondBuilder, + InterResiduePlaneBuilder, + InterResidueTorsionBuilder, + PreprocessedCIF, +) +from torchref.topology.builders_numba import ( + match_angles_numba, + match_bonds_numba, + match_chirals_numba, + match_torsions_numba, +) +from torchref.topology.atom_graph import AtomGraph +from torchref.topology.edges import EdgeBlock, assemble_origins +from torchref.topology.residue_graph import ( + ResidueGraph, + build_residue_nodes, + find_disulfide_links, + find_peptide_links, +) +from torchref.topology.restraint_sets import to_tensor +from torchref.topology.templates import resolve_template_keys +from torchref.topology.topology import Topology + +#: Initial size of the matcher work arrays, grown on demand. +_WORK = 64 + + +def _atom_columns(pdb: pd.DataFrame) -> Dict[str, np.ndarray]: + """Per-atom identity arrays, with altlocs normalised so blank reads as ``' '``.""" + altloc = pdb["altloc"].values.astype(str) if "altloc" in pdb.columns else None + if altloc is None: + altloc = np.full(len(pdb), " ", dtype=" List[Tuple[np.ndarray, np.ndarray]]: + """Atom name/index arrays per alternative conformation of one residue. + + A residue without altlocs yields one conformation holding all its atoms. Otherwise + one per altloc, each holding the blank-altloc atoms plus that altloc's own -- so a + restraint spanning a shared backbone and a branching side chain is emitted once per + conformer. + """ + names = cols["name"][start:end] + indices = cols["index"][start:end] + altlocs = cols["altloc"][start:end] + unique = np.unique(altlocs) + + if len(unique) == 1 and unique[0] == " ": + return [(names, indices)] + if " " in unique: + common = altlocs == " " + out = [] + for alt in unique: + if alt == " ": + continue + m = altlocs == alt + out.append( + ( + np.concatenate([names[common], names[m]]), + np.concatenate([indices[common], indices[m]]), + ) + ) + return out + return [(names[altlocs == a], indices[altlocs == a]) for a in unique] + + +def _atom_types( + cols: Dict[str, np.ndarray], + nodes: Dict[str, np.ndarray], + template_key: np.ndarray, + comp_dict: Dict, +) -> Tuple[np.ndarray, np.ndarray]: + """Per-atom energy type and template hydrogen count, by name in the patched template. + + Returns + ------- + energy_type : numpy.ndarray + Shape ``(N,)``, ``''`` where the residue has no template or the atom is not + in it. + template_h_count : numpy.ndarray + Shape ``(N,)``, ``int8``; hydrogens the atom carries in its template, ``0`` + for template atoms with none (hydrogens included), ``-1`` where unknown. + """ + from torchref.topology.hydrogens import template_atom_types + + n_atoms = len(cols["name"]) + energy = np.full(n_atoms, "", dtype=" Tuple[Dict[str, np.ndarray], Dict[str, Dict[str, np.ndarray]]]: + """Intra-residue edges and the ideal values that belong to them. + + Keyed ``bonds`` / ``angles`` / ``torsions`` / ``chirals``. + + Emitted only where every named atom of a library restraint is present in the + conformation, which is the condition the matchers apply. + + Returns + ------- + indices : dict + ``{kind: (E, k) array}``. + values : dict + ``{kind: {property: (E,) array}}``, accumulated row-for-row with the indices. + """ + acc: Dict[str, List[np.ndarray]] = { + "bonds": [], + "angles": [], + "torsions": [], + "chirals": [], + } + val: Dict[str, Dict[str, List[np.ndarray]]] = { + "bonds": {"references": [], "sigmas": []}, + "angles": {"references": [], "sigmas": []}, + "torsions": {"references": [], "sigmas": [], "periods": []}, + "chirals": {"ideal_volumes": [], "sigmas": []}, + } + work = {k: np.zeros(_WORK, dtype=np.int64) for k in ("i1", "i2", "i3", "i4", "per")} + work["f1"] = np.zeros(_WORK, dtype=np.float64) + work["f2"] = np.zeros(_WORK, dtype=np.float64) + size = _WORK + + for r in range(len(nodes["chain"])): + key = str(template_key[r]) + start, end = int(nodes["atom_start"][r]), int(nodes["atom_end"][r]) + needed = max( + len(pp_cif.bonds.get(key, {}).get("atom1", ())), + len(pp_cif.angles.get(key, {}).get("atom1", ())), + len(pp_cif.torsions.get(key, {}).get("atom1", ())), + len(pp_cif.chirals.get(key, {}).get("atom1", ())), + ) + if needed == 0: + continue + if needed > size: + size = needed * 2 + work = {k: np.zeros(size, dtype=v.dtype) for k, v in work.items()} + + for names, indices in _conformers(cols, start, end): + if key in pp_cif.bonds: + b = pp_cif.bonds[key] + n = match_bonds_numba( + names, + indices, + b["atom1"], + b["atom2"], + b["value"], + b["sigma"], + work["i1"], + work["i2"], + work["f1"], + work["f2"], + ) + if n: + acc["bonds"].append( + np.column_stack([work["i1"][:n].copy(), work["i2"][:n].copy()]) + ) + val["bonds"]["references"].append(work["f1"][:n].copy()) + val["bonds"]["sigmas"].append(work["f2"][:n].copy()) + if key in pp_cif.angles: + a = pp_cif.angles[key] + n = match_angles_numba( + names, + indices, + a["atom1"], + a["atom2"], + a["atom3"], + a["value"], + a["sigma"], + work["i1"], + work["i2"], + work["i3"], + work["f1"], + work["f2"], + ) + if n: + acc["angles"].append( + np.column_stack( + [ + work["i1"][:n].copy(), + work["i2"][:n].copy(), + work["i3"][:n].copy(), + ] + ) + ) + val["angles"]["references"].append(work["f1"][:n].copy()) + val["angles"]["sigmas"].append(work["f2"][:n].copy()) + if key in pp_cif.torsions: + t = pp_cif.torsions[key] + n = match_torsions_numba( + names, + indices, + t["atom1"], + t["atom2"], + t["atom3"], + t["atom4"], + t["value"], + t["sigma"], + t["period"], + work["i1"], + work["i2"], + work["i3"], + work["i4"], + work["f1"], + work["f2"], + work["per"], + ) + if n: + acc["torsions"].append( + np.column_stack( + [ + work["i1"][:n].copy(), + work["i2"][:n].copy(), + work["i3"][:n].copy(), + work["i4"][:n].copy(), + ] + ) + ) + val["torsions"]["references"].append(work["f1"][:n].copy()) + val["torsions"]["sigmas"].append(work["f2"][:n].copy()) + val["torsions"]["periods"].append(work["per"][:n].copy()) + if key in pp_cif.chirals: + c = pp_cif.chirals[key] + n = match_chirals_numba( + names, + indices, + c["center"], + c["atom1"], + c["atom2"], + c["atom3"], + c["volume_sign"], + c["sigma"], + work["i1"], + work["i2"], + work["i3"], + work["i4"], + work["f1"], + work["f2"], + ) + if n: + acc["chirals"].append( + np.column_stack( + [ + work["i1"][:n].copy(), + work["i2"][:n].copy(), + work["i3"][:n].copy(), + work["i4"][:n].copy(), + ] + ) + ) + # Ideal volume is the sign times a typical tetrahedral volume. A + # sign of 0 ('both' / 'either') stays exactly 0, which the chiral + # target reads as an achiral centre and restrains |volume| instead. + val["chirals"]["ideal_volumes"].append(work["f1"][:n].copy() * 2.5) + val["chirals"]["sigmas"].append(work["f2"][:n].copy()) + + arity = {"bonds": 2, "angles": 3, "torsions": 4, "chirals": 4} + indices = { + k: (np.concatenate(v, axis=0) if v else np.zeros((0, arity[k]), dtype=np.int64)) + for k, v in acc.items() + } + values: Dict[str, Dict[str, np.ndarray]] = {} + for kind, properties in val.items(): + joined: Dict[str, np.ndarray] = {} + for prop, chunks in properties.items(): + if not chunks: + continue + array = np.concatenate(chunks) + if prop == "sigmas": + # A zero sigma divides by zero in the loss, so it is floored. + array = np.where(array == 0, 1e-4, array) + joined[prop] = array + values[kind] = joined + return indices, values + + +def _match_intra_planes( + cols: Dict[str, np.ndarray], + nodes: Dict[str, np.ndarray], + template_key: np.ndarray, + pp_cif: PreprocessedCIF, +) -> Tuple[Dict[int, np.ndarray], Dict[int, Dict[str, np.ndarray]]]: + """Intra-residue planes grouped by how many atoms survived matching. + + Missing atoms are dropped rather than voiding the plane; a plane is kept once at + least three of its atoms are present, so its arity depends on the model. Sigmas are + per atom, not per plane, so they carry the same ``(E, k)`` shape as the indices. + """ + by_size: Dict[int, List[np.ndarray]] = {} + sigmas_by_size: Dict[int, List[np.ndarray]] = {} + for r in range(len(nodes["chain"])): + key = str(template_key[r]) + if key not in pp_cif.planes: + continue + start, end = int(nodes["atom_start"][r]), int(nodes["atom_end"][r]) + for names, indices in _conformers(cols, start, end): + # Last-wins on a duplicate name, matching PlaneRestraintBuilder. + name_to_idx = dict(zip(names, indices)) + for plane in pp_cif.planes[key]: + present = [] + present_sigmas = [] + for position, atom_name in enumerate(plane["atoms"]): + if atom_name in name_to_idx: + present.append(name_to_idx[atom_name]) + present_sigmas.append(plane["sigmas"][position]) + if len(present) >= 3: + by_size.setdefault(len(present), []).append( + np.asarray(present, dtype=np.int64) + ) + sigmas_by_size.setdefault(len(present), []).append( + np.asarray(present_sigmas, dtype=np.float64) + ) + + indices_out = {n: np.stack(rows, axis=0) for n, rows in by_size.items()} + values_out = { + n: { + "sigmas": np.where( + np.stack(rows, axis=0) == 0, 1e-4, np.stack(rows, axis=0) + ) + } + for n, rows in sigmas_by_size.items() + } + return indices_out, values_out + + +def _inter_residue_edges( + pdb: pd.DataFrame, + link_dict: Optional[Dict], + verbose: int, +) -> Tuple[Dict[str, Dict[str, np.ndarray]], Dict[str, Dict], Dict[str, Dict]]: + """Peptide edges, their values, and the Ramachandran pairing, from the builders. + + Reuses ``InterResidue*Builder`` rather than reimplementing the link geometry. + + Returns + ------- + indices : dict + ``{edge type: {origin: (E, k) array}}``. + values : dict + ``{edge type: {origin: {property: array}}}`` -- every property a builder + returned besides the indices, so ``omega``'s ``is_proline`` comes along without + being named here. + extras : dict + Non-edge products of the same pass, currently the ``ramachandran`` phi/psi + pairing and its surface types. + """ + indices: Dict[str, Dict[str, np.ndarray]] = { + "bond": {}, + "angle": {}, + "torsion": {}, + "plane": {}, + } + values: Dict[str, Dict] = {"bond": {}, "angle": {}, "torsion": {}, "plane": {}} + extras: Dict[str, Dict] = {} + if not link_dict or "TRANS" not in link_dict: + return indices, values, extras + + cpu = torch.device("cpu") + trans = link_dict["TRANS"] + ptrans = link_dict.get("PTRANS") + + def split(group): + """A builder group as ``(indices array, {property: array})``.""" + rows = group["indices"].cpu().numpy() + rest = { + prop: tensor.cpu().numpy() + for prop, tensor in group.items() + if prop != "indices" and tensor is not None + } + return rows, rest + + bond = InterResidueBondBuilder(verbose=verbose).build( + pdb, trans, cpu, filter_atom_type="ATOM" + ) + if bond: + indices["bond"]["peptide"], values["bond"]["peptide"] = split(bond) + + ab = InterResidueAngleBuilder(verbose=verbose) + if ptrans is not None: + # PTRANS carries the extra C(i-1)-N-CD angle, so proline pairs are built from it + # and excluded from the TRANS pass to avoid two restraints on the same atoms. + groups = [ + ab.build( + pdb, trans, cpu, filter_atom_type="ATOM", exclude_next_resname="PRO" + ), + ab.build( + pdb, ptrans, cpu, filter_atom_type="ATOM", next_resname_filter="PRO" + ), + ] + else: + groups = [ab.build(pdb, trans, cpu, filter_atom_type="ATOM")] + parts = [split(g) for g in groups if g] + if parts: + indices["angle"]["peptide"] = np.concatenate([p[0] for p in parts], axis=0) + shared = set.intersection(*(set(p[1]) for p in parts)) + values["angle"]["peptide"] = { + prop: np.concatenate([p[1][prop] for p in parts]) for prop in shared + } + + tors = InterResidueTorsionBuilder(verbose=verbose).build( + pdb, trans, cpu, filter_atom_type="ATOM" + ) + if tors: + for origin in ("phi", "psi", "omega"): + if origin in tors: + indices["torsion"][origin], values["torsion"][origin] = split( + tors[origin] + ) + if "ramachandran" in tors: + extras["ramachandran"] = tors["ramachandran"] + + planes = InterResiduePlaneBuilder(verbose=verbose).build( + pdb, trans, cpu, filter_atom_type="ATOM" + ) + if planes: + for key, group in planes.items(): + indices["plane"][key], values["plane"][key] = split(group) + return indices, values, extras + + +def _origins( + intra: np.ndarray, + inter: Dict[str, np.ndarray], + disulfide: Optional[np.ndarray], +) -> Dict[str, np.ndarray]: + """Collect one edge type's per-origin arrays, dropping the empty ones.""" + per_origin: Dict[str, np.ndarray] = {} + if intra is not None and len(intra): + per_origin["intra"] = intra + for origin, rows in inter.items(): + if rows is not None and len(rows): + per_origin[origin] = rows + if disulfide is not None and len(disulfide): + per_origin["disulfide"] = disulfide + return per_origin + + +def _disulfide_edges( + pdb: pd.DataFrame, + nodes: Dict[str, np.ndarray], + cols: Dict[str, np.ndarray], + residue_of_row: Dict[int, int], + pairs: Sequence[Tuple[int, int]], + link_dict: Optional[Dict], + verbose: int, +) -> Tuple[Dict[str, np.ndarray], Dict[str, Dict[str, np.ndarray]]]: + """Bond, angle and torsion edges for the detected disulfide links, with values. + + Drives the ``InterResidue*Builder`` disulfide paths from the residue graph's + ``disulf`` edges, so the link geometry comes from the ``disulf`` dictionary entry + rather than being restated here. + + Returns + ------- + indices : dict + ``{'bond'|'angle'|'torsion': (E, k) array}``, omitting types with no edges. + values : dict + ``{edge type: {property: array}}`` for the same edges. + """ + out: Dict[str, np.ndarray] = {} + if not pairs or not link_dict or "disulf" not in link_dict: + return out, {} + + disulf = link_dict["disulf"] + bonds = disulf.get("bonds") + if bonds is None: + return out, {} + sg_sg = bonds[(bonds["atom1"] == "SG") & (bonds["atom2"] == "SG")] + if len(sg_sg) == 0: + return out, {} + length = float(sg_sg["value"].values[0]) + sigma = float(sg_sg["sigma"].values[0]) + + cpu = torch.device("cpu") + bond_builder = InterResidueBondBuilder(verbose=verbose) + angle_builder = InterResidueAngleBuilder(verbose=verbose) + torsion_builder = InterResidueTorsionBuilder(verbose=verbose) + + for row_a, row_b in pairs: + # The edge indices are the atom table's ``index`` column, not its row number. + bond_builder.process_disulfide_bond( + int(cols["index"][row_a]), int(cols["index"][row_b]), length, sigma + ) + res_a, res_b = residue_of_row[row_a], residue_of_row[row_b] + atoms_a = pdb.iloc[ + int(nodes["atom_start"][res_a]) : int(nodes["atom_end"][res_a]) + ] + atoms_b = pdb.iloc[ + int(nodes["atom_start"][res_b]) : int(nodes["atom_end"][res_b]) + ] + if disulf.get("angles") is not None: + angle_builder.process_disulfide_angles(atoms_a, atoms_b, disulf["angles"]) + if disulf.get("torsions") is not None: + torsion_builder.process_disulfide_torsions( + atoms_a, atoms_b, disulf["torsions"] + ) + + values: Dict[str, Dict[str, np.ndarray]] = {} + for edge_type, group in ( + ("bond", bond_builder.finalize(cpu)), + ("angle", angle_builder.finalize(cpu)), + ("torsion", torsion_builder.finalize_disulfide(cpu)), + ): + if not group: + continue + out[edge_type] = group["indices"].cpu().numpy() + values[edge_type] = { + prop: tensor.cpu().numpy() + for prop, tensor in group.items() + if prop != "indices" and tensor is not None + } + return out, values + + +def _lookup_link_atom( + pdb: pd.DataFrame, + chainid: str, + resseq: int, + icode: str, + resname: str, + name: str, + altloc: str, +): + """Resolve one ``LINK`` record's atom to a row of the atom table, or None. + + Matches on ``(chainid, resseq, icode, name)`` with ``resname`` as a tie-breaker. + Where a residue has alternative conformations the requested altloc wins, then the + blank one, then ``'A'``, then whatever is left -- a LINK naming a specific conformer + should reach that conformer, but one naming none should still resolve. + """ + sel = pdb[ + (pdb["chainid"].astype(str) == str(chainid)) + & (pdb["resseq"].astype(int) == int(resseq)) + & (pdb["icode"].astype(str) == str(icode)) + & (pdb["name"].astype(str).str.strip() == str(name).strip()) + ] + if len(sel) == 0: + return None + if resname: + tied = sel[sel["resname"].astype(str).str.strip() == str(resname).strip()] + if len(tied) > 0: + sel = tied + + if altloc: + for candidate in (altloc, ""): + hit = sel[sel["altloc"].astype(str) == candidate] + if len(hit) > 0: + return int(hit.iloc[0]["index"]) + for candidate in ("", "A"): + hit = sel[sel["altloc"].astype(str) == candidate] + if len(hit) > 0: + return int(hit.iloc[0]["index"]) + return int(sel.iloc[0]["index"]) + + +def _link_record_edges( + pdb: pd.DataFrame, + links, + disulfide_bonds: Optional[np.ndarray], + verbose: int, +) -> Tuple[np.ndarray, List[Tuple[int, int]], Dict[str, np.ndarray]]: + """Bond edges for the accepted ``LINK`` records, and the atom pairs they join. + + A record duplicating an auto-detected disulfide is dropped, since that link already + contributed its bond, angles and torsions; so is a record repeating an earlier one, + which would otherwise add a second bond edge and a second restraint on the same pair. + + Returns + ------- + edges : numpy.ndarray + Shape ``(L, 2)``; empty when nothing resolved. + atom_pairs : list of tuple of int + The same pairs, for lifting to residue-level link edges. + values : dict + ``references`` from each record's ``length`` (1.5 A where blank or unusable) and + a fixed ``sigmas`` of 0.02 A. + """ + if links is None or len(links) == 0: + return np.zeros((0, 2), dtype=np.int64), [], {} + + existing = set() + if disulfide_bonds is not None: + for a, b in disulfide_bonds: + existing.add((min(int(a), int(b)), max(int(a), int(b)))) + + rows: List[Tuple[int, int]] = [] + lengths: List[float] = [] + n_unresolved = 0 + for _, link in links.iterrows(): + idx1 = _lookup_link_atom( + pdb, + chainid=link["chainid1"], + resseq=int(link["resseq1"]), + icode=link["icode1"], + resname=link["resname1"], + name=link["name1"], + altloc=link["altloc1"], + ) + idx2 = _lookup_link_atom( + pdb, + chainid=link["chainid2"], + resseq=int(link["resseq2"]), + icode=link["icode2"], + resname=link["resname2"], + name=link["name2"], + altloc=link["altloc2"], + ) + if idx1 is None or idx2 is None or idx1 == idx2: + n_unresolved += 1 + continue + pair = (min(idx1, idx2), max(idx1, idx2)) + if pair in existing: + continue + existing.add(pair) + rows.append((idx1, idx2)) + length = link["length"] + usable = isinstance(length, (int, float)) and length == length and length > 0 + lengths.append(float(length) if usable else 1.5) + + if verbose > 1 and n_unresolved: + print(f"{n_unresolved} LINK records did not resolve to a pair of atoms") + if not rows: + return np.zeros((0, 2), dtype=np.int64), [], {} + values = { + "references": np.asarray(lengths, dtype=np.float64), + "sigmas": np.full(len(rows), 0.02, dtype=np.float64), + } + return np.asarray(rows, dtype=np.int64), rows, values + + +def _block_with_values( + per_origin: Dict[str, np.ndarray], + payload: Dict[str, Dict[str, np.ndarray]], + arity: int, + edge_type: str, + device, +) -> Tuple[EdgeBlock, Dict[str, Dict[str, torch.Tensor]]]: + """One canonical edge block plus its per-origin value tensors. + + The block and the values come out of a single :func:`assemble_origins` call, so the + same permutation is applied to both -- which is the only thing keeping a sigma + attached to the edge it belongs to. + """ + indices, bounds, sorted_payload = assemble_origins( + per_origin, arity, edge_type, payload + ) + block = EdgeBlock( + indices=torch.as_tensor(indices, dtype=torch.int64, device=device), # dtype-ok: atom index tensor for restraint edges; int64 index required + origin_bounds=bounds, + ) + values = { + origin: { + prop: to_tensor(array, prop, device=device) + for prop, array in properties.items() + } + for origin, properties in sorted_payload.items() + } + return block, values + + +def build_topology( + pdb: pd.DataFrame, + cif_dict: Dict, + link_dict: Optional[Dict] = None, + link_list=None, + links=None, + xyz: Optional[torch.Tensor] = None, + device=None, + verbose: int = 0, +) -> Topology: + """Build a topology, discarding the restraint values built along the way. + + See :func:`build_topology_with_values` for the parameters; this is the connectivity + half on its own, for callers that need the graph and no ideal geometry. + """ + topology, _, _ = build_topology_with_values( + pdb, + cif_dict, + link_dict=link_dict, + link_list=link_list, + links=links, + xyz=xyz, + device=device, + verbose=verbose, + ) + return topology + + +def build_topology_with_values( + pdb: pd.DataFrame, + cif_dict: Dict, + link_dict: Optional[Dict] = None, + link_list=None, + links=None, + xyz: Optional[torch.Tensor] = None, + device=None, + verbose: int = 0, +) -> Tuple[Topology, Dict[str, Dict], Dict[str, Dict]]: + """Build a topology from an atom table and the restraint dictionaries. + + Parameters + ---------- + pdb : pandas.DataFrame + Atom table, with ``name``, ``element``, ``altloc``, ``chainid``, ``resseq``, + ``icode``, ``resname``, ``ATOM`` and ``index`` columns. + cif_dict : dict + Restraint dictionary keyed by residue name. + link_dict : dict, optional + Link-type definitions. Without it no inter-residue edges are built. + link_list : pandas.DataFrame, optional + Link table used to resolve which modifications a peptide link applies. + links : pandas.DataFrame, optional + Parsed PDB ``LINK`` records. Each record that resolves to two distinct atoms and + does not duplicate an auto-detected disulfide contributes one bond edge. + xyz : torch.Tensor, optional + Coordinates, shape ``(N, 3)``. Needed only to detect disulfide links, which are + found by SG-SG distance. + device : torch.device, optional + Where to place the edge blocks. + verbose : int, default 0 + Verbosity level. + + Returns + ------- + topology : Topology + The connectivity. + values : dict + ``{edge_type: {origin: {property: tensor}}}`` for bonds, angles and torsions; + ``{'chiral': {property: tensor}}`` and ``{'plane': {size: {property: tensor}}}`` + for the two types that carry no origin. Row-aligned to the edge blocks. + extras : dict + Products of the same pass that are not edges -- currently ``ramachandran``. + """ + cols = _atom_columns(pdb) + nodes = build_residue_nodes( + cols["chain"], cols["resseq"], cols["icode"], cols["resname"] + ) + n_res = len(nodes["chain"]) + + names_by_residue = [ + set(cols["name"][int(nodes["atom_start"][r]) : int(nodes["atom_end"][r])]) + for r in range(n_res) + ] + is_polymer = np.array( + [cols["record"][int(nodes["atom_start"][r])] == "ATOM" for r in range(n_res)], + dtype=bool, + ) + + polymer_nodes = {k: v[is_polymer] for k, v in nodes.items()} + polymer_map = np.nonzero(is_polymer)[0] + polymer_names = [names_by_residue[r] for r in polymer_map] + peptide_local = find_peptide_links(polymer_nodes, polymer_names) + peptide_pairs = [ + (int(polymer_map[a]), int(polymer_map[b])) for a, b in peptide_local + ] + + comp_dict, template_key = resolve_template_keys( + nodes["resname"], peptide_pairs, cif_dict, link_list, verbose=verbose + ) + pp_cif = PreprocessedCIF(comp_dict) + match_cols = dict(cols) + match_cols["name"] = cols["name"].copy() + # PDB terminal H1 is the monomer dictionary's H. Resolve the alias only + # for matching, preserving the model's atom names and row identities. + for r in range(n_res): + start, end = int(nodes["atom_start"][r]), int(nodes["atom_end"][r]) + names = match_cols["name"][start:end] + if "H1" not in names or "H" in names: + continue + component = comp_dict.get(str(template_key[r]), {}) + atom_table = component.get("atoms") + if atom_table is None: + continue + template_names = set(atom_table["atom_id"].astype(str).str.strip()) + if "H" in template_names and "H1" not in template_names: + names[names == "H1"] = "H" + energy_type, template_h_count = _atom_types( + match_cols, nodes, template_key, comp_dict + ) + + intra, intra_values = _match_intra(match_cols, nodes, template_key, pp_cif) + intra_planes, intra_plane_values = _match_intra_planes( + match_cols, nodes, template_key, pp_cif + ) + inter, inter_values, extras = _inter_residue_edges(pdb, link_dict, verbose) + + residue_of_row = {} + for r in range(n_res): + for row in range(int(nodes["atom_start"][r]), int(nodes["atom_end"][r])): + residue_of_row[row] = r + + disulfide_pairs: List[Tuple[int, int]] = [] + disulfide: Dict[str, np.ndarray] = {} + disulfide_values: Dict[str, Dict[str, np.ndarray]] = {} + if xyz is not None: + sg_rows = [ + row + for row in range(len(cols["name"])) + if cols["name"][row] == "SG" and cols["record"][row] == "ATOM" + ] + disulfide_pairs = find_disulfide_links(sg_rows, residue_of_row, xyz) + disulfide, disulfide_values = _disulfide_edges( + pdb, nodes, cols, residue_of_row, disulfide_pairs, link_dict, verbose + ) + + link_edges, link_atom_pairs, link_values = _link_record_edges( + pdb, links, disulfide.get("bond"), verbose + ) + + # LINK edges carry ``index`` values, so lift them through that column. + index_to_residue = {int(cols["index"][row]): r for row, r in residue_of_row.items()} + + link_pairs = [(a, b, "TRANS") for a, b in peptide_pairs] + disulf_residue_pairs = sorted( + { + ( + min(residue_of_row[a], residue_of_row[b]), + max(residue_of_row[a], residue_of_row[b]), + ) + for a, b in disulfide_pairs + } + ) + link_pairs += [(a, b, "disulf") for a, b in disulf_residue_pairs] + for a, b in link_atom_pairs: + ra, rb = index_to_residue.get(a), index_to_residue.get(b) + if ra is not None and rb is not None and ra != rb: + link_pairs.append((ra, rb, "LINK")) + + residues = ResidueGraph( + chain=nodes["chain"], + resseq=nodes["resseq"], + icode=nodes["icode"], + resname=nodes["resname"], + template_key=template_key, + atom_start=nodes["atom_start"], + atom_end=nodes["atom_end"], + link_pairs=( + np.array([(a, b) for a, b, _ in link_pairs], dtype=np.int64) + if link_pairs + else np.zeros((0, 2), dtype=np.int64) + ), + link_kind=np.array([k for _, _, k in link_pairs], dtype=" np.ndarray: + """Row order that sorts ``rows`` lexicographically, left column most significant. + + Parameters + ---------- + rows : numpy.ndarray + Integer array of shape ``(E, k)``. + + Returns + ------- + numpy.ndarray + Permutation of ``arange(E)``. Total on the row values, so it does not depend + on the incoming order the way a single-column ``argsort`` does. + """ + if rows.size == 0: + return np.zeros(0, dtype=np.int64) + return np.lexsort(tuple(rows[:, c] for c in reversed(range(rows.shape[1])))) + + +def assemble_origins( + per_origin: Dict[str, np.ndarray], + arity: int, + edge_type: str, + payload: Dict[str, Dict[str, np.ndarray]] = None, +) -> Tuple[np.ndarray, Dict[str, Tuple[int, int]], Dict[str, Dict[str, np.ndarray]]]: + """Lay origins out in canonical order, carrying per-edge values through the sort. + + Origins follow :data:`ORIGIN_ORDER` for ``edge_type``, and rows within an origin + are sorted lexicographically. Anything in ``payload`` is permuted by the same order, + so a value array stays aligned row-for-row with the indices it belongs to, which + is the whole reason values cannot be concatenated separately. + + Parameters + ---------- + per_origin : dict + ``{origin: (E_o, k) integer array}``. Empty entries are skipped. + arity : int + Atoms per edge. + edge_type : str + Key into :data:`ORIGIN_ORDER`. + payload : dict, optional + ``{origin: {property: array}}``, each array indexed by row on axis 0. A property + need not be present for every origin -- ``phi`` and ``psi`` carry no reference + value or sigma, and must not acquire one here. + + Returns + ------- + indices : numpy.ndarray + Shape ``(E, k)``, canonical order. + bounds : dict + ``{origin: (start, end)}``, contiguous and covering the block. + sorted_payload : dict + ``{origin: {property: array}}``, permuted to match ``indices``. + """ + order = ORIGIN_ORDER.get(edge_type, tuple(sorted(per_origin))) + unknown = set(per_origin) - set(order) + if unknown: + raise ValueError( + f"{edge_type}: origins {sorted(unknown)} are not in ORIGIN_ORDER" + f"[{edge_type!r}] = {order}. Add them there so the layout stays " + f"deterministic." + ) + + payload = payload or {} + chunks: List[np.ndarray] = [] + bounds: Dict[str, Tuple[int, int]] = {} + sorted_payload: Dict[str, Dict[str, np.ndarray]] = {} + cursor = 0 + + for origin in order: + rows = per_origin.get(origin) + if rows is None or len(rows) == 0: + continue + rows = np.asarray(rows, dtype=np.int64).reshape(-1, arity) + permutation = _lexsort_rows(rows) + chunks.append(rows[permutation]) + bounds[origin] = (cursor, cursor + len(rows)) + cursor += len(rows) + + origin_payload = payload.get(origin) or {} + if origin_payload: + sorted_payload[origin] = { + prop: np.asarray(values)[permutation] + for prop, values in origin_payload.items() + if values is not None + } + + indices = ( + np.concatenate(chunks, axis=0) + if chunks + else np.zeros((0, arity), dtype=np.int64) + ) + return indices, bounds, sorted_payload + + +@dataclass(eq=False, repr=False) +class EdgeBlock(DeviceMixin): + """One edge type's index block plus its per-origin bounds. + + Parameters + ---------- + indices : torch.Tensor + Atom indices, shape ``(E, k)``, dtype ``int64``, in canonical order. + origin_bounds : dict + ``{origin: (start, end)}`` half-open row ranges into ``indices``. Ranges are + contiguous and cover the block. + + Notes + ----- + Holds no refinable parameters, so this is a dataclass rather than an + ``nn.Module``; ``DeviceMixin`` still moves ``indices`` with ``.to(device)``. + """ + + indices: torch.Tensor + origin_bounds: Dict[str, Tuple[int, int]] = field(default_factory=dict) + + @classmethod + def empty(cls, arity: int, device=None) -> "EdgeBlock": + """An edge-free block of the given arity.""" + return cls( + indices=torch.zeros((0, arity), dtype=torch.int64, device=device), # dtype-ok: empty edge index tensor (0,arity); int64 index required + origin_bounds={}, + ) + + @classmethod + def from_origins( + cls, + per_origin: Dict[str, np.ndarray], + arity: int, + edge_type: str, + device=None, + ) -> "EdgeBlock": + """Assemble a canonical block from ``{origin: (E_o, k) index array}``. + + Origins are laid out in :data:`ORIGIN_ORDER` for ``edge_type``, and rows + within an origin are sorted lexicographically. Origins with no rows are + omitted from ``origin_bounds`` rather than recorded as empty ranges. + + Parameters + ---------- + per_origin : dict + Integer index arrays keyed by origin. Empty arrays are skipped. + arity : int + Atoms per edge (2 for bonds, 3 for angles, ...). + edge_type : str + Key into :data:`ORIGIN_ORDER`. + device : torch.device, optional + Where to place the block. + + Returns + ------- + EdgeBlock + """ + indices, bounds, _ = assemble_origins(per_origin, arity, edge_type) + if len(indices) == 0: + return cls.empty(arity, device=device) + return cls( + indices=torch.as_tensor(indices, dtype=torch.int64, device=device), # dtype-ok: edge atom index tensor; int64 index required + origin_bounds=bounds, + ) + + @property + def device(self) -> torch.device: + """Where the block lives. Derived, so it cannot fall out of step.""" + return self.indices.device + + @property + def n_edges(self) -> int: + """Number of rows in the block.""" + return int(self.indices.shape[0]) + + @property + def arity(self) -> int: + """Atoms per edge.""" + return int(self.indices.shape[1]) + + def origins(self) -> List[str]: + """Origins present, in block layout order.""" + return sorted(self.origin_bounds, key=lambda o: self.origin_bounds[o][0]) + + def origin(self, name: str) -> torch.Tensor: + """Rows contributed by one origin, as a **view** into the block. + + Parameters + ---------- + name : str + Origin key. + + Returns + ------- + torch.Tensor + Shape ``(E_o, k)``. Shares storage with :attr:`indices`, so an in-place + edit to either is visible through the other. + + Raises + ------ + KeyError + If the origin contributed no rows. + """ + start, end = self.origin_bounds[name] + return self.indices[start:end] + + def copy(self) -> "EdgeBlock": + """An independent copy sharing no storage with this one.""" + return EdgeBlock( + indices=self.indices.clone(), + origin_bounds=dict(self.origin_bounds), + ) + + def subset(self, remap: torch.Tensor) -> "EdgeBlock": + """Edges whose every atom survives, reindexed by ``remap``. + + Parameters + ---------- + remap : torch.Tensor + Old atom index to new, shape ``(N_old,)``, with ``-1`` where the atom is + being dropped. An edge is kept only if none of its atoms maps to ``-1``: a + bond to a removed atom is not a bond, and an angle missing its apex is not + an angle. + + Returns + ------- + EdgeBlock + Canonically ordered, with ``origin_bounds`` recomputed over the survivors. + + Notes + ----- + No re-sort is needed. ``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 surviving rows stay sorted among + themselves. + """ + if self.n_edges == 0: + return EdgeBlock.empty(self.arity, device=self.indices.device) + + mapped = remap[self.indices] + keep = (mapped >= 0).all(dim=1) + + chunks = [] + bounds: Dict[str, Tuple[int, int]] = {} + cursor = 0 + for origin in self.origins(): + start, end = self.origin_bounds[origin] + surviving = mapped[start:end][keep[start:end]] + if surviving.shape[0] == 0: + continue + chunks.append(surviving) + bounds[origin] = (cursor, cursor + surviving.shape[0]) + cursor += surviving.shape[0] + + if not chunks: + return EdgeBlock.empty(self.arity, device=self.indices.device) + return EdgeBlock(indices=torch.cat(chunks, dim=0), origin_bounds=bounds) + + def tuple_set(self, origin: str = None) -> set: + """Edges as a set of index tuples, for order-free comparison. + + Parameters + ---------- + origin : str, optional + Restrict to one origin. None means the whole block. + + Returns + ------- + set of tuple of int + """ + rows = self.indices if origin is None else self.origin(origin) + return {tuple(int(v) for v in row) for row in rows.cpu().numpy()} + + def __repr__(self) -> str: + return ( + f"EdgeBlock(arity={self.arity}, n_edges={self.n_edges}, " + f"origins={self.origins()})" + ) + + +__all__ = ["EdgeBlock", "ORIGIN_ORDER", "assemble_origins"] diff --git a/torchref/topology/hydrogens.py b/torchref/topology/hydrogens.py new file mode 100644 index 00000000..5ec85303 --- /dev/null +++ b/torchref/topology/hydrogens.py @@ -0,0 +1,1289 @@ +"""Hydrogen generation as a graph operation: expand the template, map it on. + +A monomer 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 that are present, +read the hydrogen positions off it, and correct each to its ideal bond length. + +Two things the bond graph decides that a distance criterion previously guessed at: + +* **How many hydrogens a parent can carry.** The smaller of two budgets: the parent's + valence (including tetrahedral ammonium nitrogen) minus the heavy atoms actually + bonded to it in the graph, and the + template's own hydrogen count minus every graph bond the template does not know about + (a peptide bond, a LINK record, a metal contact). The first budget handles the + template's own chemistry; the second is what stops an acetyl cap's aldehyde hydrogen + or a metal-bound histidine NE2 hydrogen from being generated when the graph degree + sits below the nominal valence only because a double bond counts as one edge. Both + read the graph rather than a distance sweep, which gets a distorted or predicted model + wrong and cannot tell a bond from two atoms that merely sit close. +* **Which hydrogens have a free torsion.** A hydrogen whose parent has exactly one heavy + neighbour -- hydroxyl, thiol, amine, methyl -- can rotate about the parent-neighbour + axis, and the template's angle for it is arbitrary. Those get scanned; the rest are + fully determined by the template and are left alone. +""" + +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import numpy as np +import torch +from torchref.config import get_float_dtype + +#: Standard heavy-atom valences, one of the two budgets that cap how many hydrogens a +#: parent may take. Elements not listed fall back to 4 and are then bounded only by the +#: template's own hydrogen count. +STANDARD_VALENCE = {"C": 4, "N": 3, "O": 2, "S": 2} +_DEFAULT_VALENCE = 4 + +#: A template hydrogen further than this from its parent after alignment is discarded +#: rather than corrected: the alignment for that centre is too poor to trust. +MAX_PLACEMENT_DISTANCE = 1.5 + +#: Angles tried when scanning a free torsion. 24 gives 15-degree resolution, which is +#: finer than the placement error the alignment itself carries. +TORSION_SCAN_STEPS = 24 + +#: Heavy atoms beyond this distance cannot clash with a hydrogen, so the scan ignores +#: them. +CLASH_CUTOFF = 4.0 + + +@dataclass +class HydrogenPlan: + """Hydrogens to add, and where to put them. + + Parameters + ---------- + name, element, altloc : numpy.ndarray + Per-hydrogen identity, shape ``(H,)``. + residue : numpy.ndarray + Residue index each hydrogen belongs to, shape ``(H,)``. + parent : numpy.ndarray + Atom-table index of each hydrogen's parent, shape ``(H,)``. + position : numpy.ndarray + Cartesian coordinates, shape ``(H, 3)``. + bond_length : numpy.ndarray + Ideal parent-hydrogen distance, shape ``(H,)``. + group : numpy.ndarray + Free-torsion group id, shape ``(H,)``; ``-1`` for a hydrogen whose position the + template determines. Hydrogens on one parent that rotate together share an id. + """ + + name: np.ndarray + element: np.ndarray + altloc: np.ndarray + residue: np.ndarray + parent: np.ndarray + position: np.ndarray + bond_length: np.ndarray + group: np.ndarray + + @property + def n_hydrogens(self) -> int: + """How many hydrogens the plan adds.""" + return len(self.name) + + @property + def rotatable(self) -> np.ndarray: + """Mask of hydrogens whose torsion is free.""" + return self.group >= 0 + + def __repr__(self) -> str: + return ( + f"HydrogenPlan(n_hydrogens={self.n_hydrogens}, " + f"free_torsions={len(set(self.group[self.rotatable].tolist()))})" + ) + + +def _kabsch(source: np.ndarray, target: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Rotation and translation carrying ``source`` onto ``target``. + + Reflections are excluded, so the template's chirality survives the alignment. + + Returns + ------- + rotation, translation : numpy.ndarray + Shapes ``(3, 3)`` and ``(3,)``; ``rotation @ p + translation`` maps a source + point. + """ + source_centre, target_centre = source.mean(0), target.mean(0) + covariance = (source - source_centre).T @ (target - target_centre) + u, _, vt = np.linalg.svd(covariance) + flip = np.diag([1.0, 1.0, 1.0 if np.linalg.det(vt.T @ u.T) > 0 else -1.0]) + rotation = vt.T @ flip @ u.T + return rotation, target_centre - rotation @ source_centre + + +def template_atom_types(component: Dict) -> Tuple[Dict[str, str], Dict[str, int]]: + """Energy type of every template atom, and the hydrogen count of every heavy one. + + Parameters + ---------- + component : dict + One residue's restraint sections as the CIF reader returns them; needs an + ``atoms`` section, and ``bonds`` for the counts. + + Returns + ------- + energy_type : dict + ``{atom name: CCP4 energy type}``; ``''`` where the dictionary carries none. + h_count : dict + ``{heavy atom name: number of hydrogens bonded to it in the template}``. Heavy + atoms with none are absent. + """ + atoms = component.get("atoms") + if atoms is None or len(atoms) == 0: + return {}, {} + ids = atoms["atom_id"].astype(str).str.strip().values.astype(str) + elements = np.char.upper( + atoms["type_symbol"].astype(str).str.strip().values.astype(str) + ) + if "type_energy" in atoms.columns: + types = atoms["type_energy"].astype(str).str.strip().values.astype(str) + types = np.where(np.isin(types, ["nan", ".", "?", "", "None"]), "", types) + else: + types = np.full(len(ids), "", dtype=" 0: + first = bonds["atom1"].astype(str).str.strip().values + second = bonds["atom2"].astype(str).str.strip().values + for a, b in zip(first, second): + if a not in is_h or b not in is_h: + continue + if is_h[a] and not is_h[b]: + h_count[b] = h_count.get(b, 0) + 1 + elif is_h[b] and not is_h[a]: + h_count[a] = h_count.get(a, 0) + 1 + return energy_type, h_count + + +def _template(cif_dict: Dict, resname: str) -> Optional[Dict]: + """Template atoms, hydrogen parents, ideal bond lengths and heavy adjacency. + + Read from the restraint dictionary the caller already loaded. The link modifications + are deliberately not consulted: they rewrite restraint sections, not atom lists, so + they cannot say which hydrogens a linked residue keeps. The valence cap answers that + from the bond graph instead. + + Returns + ------- + dict or None + None when the component is absent or its atoms carry no coordinates. + """ + component = cif_dict.get(resname) + if component is None: + return None + atoms = component.get("atoms") + if atoms is None or len(atoms) == 0: + return None + if not all(column in atoms.columns for column in ("x", "y", "z")): + return None + + ids = atoms["atom_id"].astype(str).str.strip().values + elements = atoms["type_symbol"].astype(str).str.strip().values.astype(" 0: + import pandas as pd + + first = bonds["atom1"].astype(str).str.strip().values + second = bonds["atom2"].astype(str).str.strip().values + values = pd.to_numeric(bonds["value"], errors="coerce").values + for i in range(len(first)): + a, b = first[i], second[i] + ia, ib = id_to_index.get(a), id_to_index.get(b) + if ia is None or ib is None: + continue + if is_h[ia] and not is_h[ib]: + parent_of[a] = b + if np.isfinite(values[i]): + ideal_length[a] = float(values[i]) + elif is_h[ib] and not is_h[ia]: + parent_of[b] = a + if np.isfinite(values[i]): + ideal_length[b] = float(values[i]) + elif not is_h[ia] and not is_h[ib]: + heavy_adjacency.setdefault(a, []).append(b) + heavy_adjacency.setdefault(b, []).append(a) + + return { + "ids": ids, + "elements": elements, + "coords": coords, + "is_h": is_h, + "id_to_index": id_to_index, + "heavy_names": ids[~is_h], + "heavy_coords": coords[~is_h], + "h_names": ids[is_h], + "parent_of": parent_of, + "h_count": h_count, + "ideal_length": ideal_length, + "heavy_adjacency": heavy_adjacency, + } + + +def _orthonormal_frame(axis: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Two unit vectors completing ``axis`` into a right-handed frame.""" + seed = np.array([1.0, 0.0, 0.0]) + if abs(float(axis @ seed)) > 0.9: + seed = np.array([0.0, 1.0, 0.0]) + first = seed - axis * float(axis @ seed) + first = first / np.linalg.norm(first) + return first, np.cross(axis, first) + + +def _axis_frame_placement( + template: Dict, + parent_name: str, + neighbour_name: str, + parent_position: np.ndarray, + neighbour_position: np.ndarray, + h_names: List[str], +) -> Optional[np.ndarray]: + """Hydrogen positions for a centre with a single heavy neighbour. + + Maps the template's local geometry onto the model by carrying the parent-neighbour + axis across and completing the frame arbitrarily. Every bond angle at the parent is + preserved exactly; only the rotation about the axis is arbitrary, which is correct: + that is the degree of freedom the template cannot know, and + :func:`optimise_free_torsions` chooses it. + """ + index = template["id_to_index"] + if parent_name not in index or neighbour_name not in index: + return None + + template_axis = ( + template["coords"][index[parent_name]] + - template["coords"][index[neighbour_name]] + ) + model_axis = parent_position - neighbour_position + for vector in (template_axis, model_axis): + if np.linalg.norm(vector) < 1e-8: + return None + template_axis = template_axis / np.linalg.norm(template_axis) + model_axis = model_axis / np.linalg.norm(model_axis) + + t_first, t_second = _orthonormal_frame(template_axis) + m_first, m_second = _orthonormal_frame(model_axis) + + positions = [] + for name in h_names: + if name not in index: + return None + offset = ( + template["coords"][index[name]] - template["coords"][index[parent_name]] + ) + positions.append( + parent_position + + float(offset @ template_axis) * model_axis + + float(offset @ t_first) * m_first + + float(offset @ t_second) * m_second + ) + return np.array(positions) + + +def _half_hydrogen_angle(template: Dict, parent_name: str, h_names: List[str]) -> float: + """Half the hydrogen-parent-hydrogen angle, from the template where it has one.""" + index = template["id_to_index"] + if len(h_names) == 2 and all(n in index for n in h_names): + parent = template["coords"][index[parent_name]] + first = template["coords"][index[h_names[0]]] - parent + second = template["coords"][index[h_names[1]]] - parent + norms = np.linalg.norm(first) * np.linalg.norm(second) + if norms > 1e-12: + cosine = float(np.clip((first @ second) / norms, -1.0, 1.0)) + return 0.5 * float(np.arccos(cosine)) + # Tetrahedral, as a fallback for a template that does not carry both hydrogens. + return 0.5 * np.arccos(-1.0 / 3.0) + + +def _split_neighbours( + topology, atom_index: int, altloc: str = "" +) -> Tuple[np.ndarray, int]: + """Heavy neighbour rows of ``atom_index``, and how many hydrogens it already has. + + Coordinate-independent, unlike a distance sweep: a stretched bond in a predicted or + mid-refinement model still counts, and two atoms that merely sit close do not. + + Restricted to the conformer being hydrogenated when ``altloc`` is given: a shared + backbone atom is bonded to every altloc copy of a split neighbour, and counting them + all makes a CA with two CB copies look saturated and lose its HA, or an N with two CA + copies lose its H. Blank-altloc neighbours always count. + + The hydrogen count is what makes generation idempotent and makes a partially + hydrogenated structure top up correctly. Both consume the parent's valence, so + subtracting only the heavy neighbours leaves budget for a hydrogen the parent + already carries -- which is how a second pass came to add the free-amino-acid ``H2`` + to every backbone nitrogen that already had its ``H``. + """ + neighbours = topology.atoms.neighbors(atom_index) + if neighbours.numel() == 0: + return np.zeros(0, dtype=np.int64), 0 + if altloc: + alts = np.char.strip(topology.atoms.altloc[neighbours.cpu().numpy()].astype(str)) + keep = torch.as_tensor((alts == "") | (alts == altloc), device=neighbours.device) + neighbours = neighbours[keep] + if neighbours.numel() == 0: + return np.zeros(0, dtype=np.int64), 0 + is_h = topology.atoms.is_hydrogen[neighbours] + return neighbours[~is_h].cpu().numpy(), int(is_h.sum()) + + +def _template_bond_length(template: Dict, parent_name: str, h_name: str) -> float: + """Parent-hydrogen distance in the template, or NaN if either atom is missing.""" + index = template["id_to_index"] + if parent_name not in index or h_name not in index: + return float("nan") + return float( + np.linalg.norm( + template["coords"][index[h_name]] - template["coords"][index[parent_name]] + ) + ) + + +def _place_group( + template: Dict, + parent_name: str, + parent_position: np.ndarray, + neighbour_positions: np.ndarray, + heavy_bonded: int, + h_names: List[str], + lengths: np.ndarray, + name_to_row: Dict[str, int], + coords: np.ndarray, + template_names_of: List[str], +) -> Optional[np.ndarray]: + """Positions for the hydrogens on one parent, by the first strategy that applies. + + In order: + + 1. **Template frame.** A Kabsch fit over the parent and its immediate heavy + neighbours, used only when the template knows every heavy atom actually bonded to + the parent. Reproduces the library geometry exactly, torsions included. + 2. **Geometric construction.** Directions from the bonded neighbours alone, for a + centre whose template is missing a real substituent -- a peptide-linked backbone + nitrogen being the common case. + 3. **Axis frame.** For a single-neighbour centre, the template geometry carried over + about the one bond, leaving the rotation about it for the scan to choose. + + Returns None when none applies, so the caller can count the hydrogen as undetermined + rather than putting it somewhere arbitrary. + """ + if ( + heavy_bonded == 0 + and len(template["heavy_names"]) == 1 + and str(template["elements"][template["id_to_index"][parent_name]]).upper() + == "O" + ): + index = template["id_to_index"] + origin = template["coords"][index[parent_name]] + # Randomness is confined to initialization and follows TorchRef's torch seed. + rotation, r = np.linalg.qr( + torch.randn(3, 3, dtype=get_float_dtype(), device="cpu").numpy() + ) + rotation = rotation * np.sign(np.diag(r))[None, :] + rotation[:, -1] *= np.linalg.det(rotation) + present_h = [h for h in template["h_names"] if h in name_to_row] + if present_h: + h = present_h[0] + source = template["coords"][index[h]] - origin + target = coords[name_to_row[h]] - parent_position + if np.linalg.norm(target) < 1e-8: + return None + source /= np.linalg.norm(source) + target /= np.linalg.norm(target) + t1, t2 = _orthonormal_frame(source) + m1 = rotation[:, 0] - target * (rotation[:, 0] @ target) + if np.linalg.norm(m1) < 1e-8: + m1, _ = _orthonormal_frame(target) + m1 /= np.linalg.norm(m1) + rotation = ( + np.column_stack([target, m1, np.cross(target, m1)]) + @ np.column_stack([source, t1, t2]).T + ) + offsets = np.array([template["coords"][index[h]] - origin for h in h_names]) + offsets = offsets @ rotation.T + return ( + parent_position + + offsets * (lengths / np.linalg.norm(offsets, axis=1))[:, None] + ) + + covered = _template_covers_neighbours(template, parent_name, heavy_bonded) + + if covered: + alignment = _alignment_for(template, parent_name, name_to_row, coords) + if alignment is not None: + matrix, offset = alignment + index = template["id_to_index"] + positions = [] + for name, length in zip(h_names, lengths): + direction = ( + matrix @ template["coords"][index[name]] + offset + ) - parent_position + distance = float(np.linalg.norm(direction)) + if distance < 1e-6 or distance > MAX_PLACEMENT_DISTANCE: + positions = None + break + positions.append(parent_position + direction * (length / distance)) + if positions is not None: + return np.array(positions) + + if heavy_bonded >= 2: + directions = _construct_directions( + parent_position, + neighbour_positions, + len(h_names), + _half_hydrogen_angle(template, parent_name, h_names), + ) + if directions is not None: + return parent_position + directions * lengths[:, None] + + if heavy_bonded == 1 and template_names_of: + positions = _axis_frame_placement( + template, + parent_name, + template_names_of[0], + parent_position, + neighbour_positions[0], + h_names, + ) + if positions is None: + return None + # Rescale along each direction so the bond length is the library value rather + # than the template's own geometry, which differs from it by ~0.002 A. Scaling + # along the direction leaves every bond angle untouched. + offsets = positions - parent_position + norms = np.linalg.norm(offsets, axis=1) + if (norms < 1e-8).any(): + return None + return parent_position + offsets * (lengths / norms)[:, None] + + return None + + +def plan_hydrogens(topology, cif_dict: Dict, xyz, verbose: int = 0) -> HydrogenPlan: + """Decide which hydrogens to add and place them from the template. + + Parameters + ---------- + topology : Topology + Supplies the residue partition, the per-residue template key and the bond graph + the valence cap and the free-torsion test read. + cif_dict : dict + Restraint dictionary, keyed by residue name; must carry an ``atoms`` section + with coordinates for a residue to be hydrogenated. + xyz : torch.Tensor + Current coordinates, shape ``(N, 3)``. + verbose : int, default 0 + Verbosity level. + + Returns + ------- + HydrogenPlan + Missing hydrogen rows with Cartesian positions in Å. + + Notes + ----- + HOH uses the bundled water dictionary when no water dictionary is supplied. + Waters without an existing hydrogen get a random reference orientation; + ``torch.manual_seed`` controls reproducibility. An existing O–H direction is + preserved when completing a partially hydrogenated water. + """ + coords = np.asarray(xyz.detach().cpu(), dtype=np.float64) + residues = topology.residues + atoms = topology.atoms + names = atoms.name.astype(str) + altlocs = atoms.altloc.astype(str) + + out: Dict[str, list] = { + k: [] + for k in ( + "name", + "altloc", + "residue", + "parent", + "position", + "bond_length", + "group", + ) + } + if "HOH" in set(residues.resname.astype(str)) and "HOH" not in cif_dict: + from pathlib import Path + + from torchref import PATH_TORCHREF_DATA + from torchref.topology.monomer.cif import read_cif + + cif_dict = dict(cif_dict) + cif_dict.update( + read_cif(str(Path(PATH_TORCHREF_DATA) / "monomer_library/h/HOH.cif")) + ) + next_group = 0 + n_unplaceable = 0 + n_no_template = 0 + + for residue in range(residues.n_residues): + resname = str(residues.resname[residue]).strip() + template = _template(cif_dict, resname) + if template is None: + # Atoms without a dictionary cannot supply either bond geometry or + # hydrogen identities; leave those residues unchanged. + n_no_template += 1 + continue + + rows = np.arange( + int(residues.atom_start[residue]), int(residues.atom_end[residue]) + ) + present = set(names[rows]) + h1_alias = "H" in template["h_names"] and "H1" not in template["h_names"] + if h1_alias and "H1" in present: + present.add("H") + candidates = [h for h in template["h_names"] if h not in present] + if not candidates: + continue + + for altloc, conformer in _conformer_rows(rows, altlocs): + name_to_row = {} + for row in conformer: + name = "H" if h1_alias and names[row] == "H1" else names[row] + name_to_row.setdefault(name, row) + + # Hydrogens grouped by the parent they hang off, in name order so the cap + # below takes a deterministic subset. + by_parent: Dict[str, List[str]] = {} + for h in sorted(candidates): + parent = template["parent_of"].get(h) + if parent is not None and parent in name_to_row: + by_parent.setdefault(parent, []).append(h) + + for parent_name, group in by_parent.items(): + parent_row = name_to_row[parent_name] + parent_position = coords[parent_row] + + heavy_rows, existing_h = _split_neighbours(topology, parent_row, altloc) + # A water-metal contact does not replace an O-H covalent bond or + # provide the water's orientational reference. + if resname == "HOH": + heavy_rows = np.zeros(0, dtype=np.int64) + heavy_bonded = len(heavy_rows) + element = str( + template["elements"][template["id_to_index"][parent_name]] + ).upper() + valence = STANDARD_VALENCE.get(element, _DEFAULT_VALENCE) + # Two budgets; see the module docstring. ``extra_bonds`` are graph + # bonds the template has never heard of -- a peptide link, a LINK + # record, a metal -- each of which displaces one template hydrogen. + template_h = template["h_count"].get(parent_name, len(group)) + template_heavy = len(template["heavy_adjacency"].get(parent_name, [])) + extra_bonds = max(0, heavy_bonded - template_heavy) + energy_types = topology.atoms.energy_type + # Free NT* amines carry four neighbours. Peptide modifications + # retype N as NH1; explicit LINK/cap bonds also displace the free + # amine protonation even when no type modification is available. + if element == "N" and extra_bonds == 0 and energy_types is not None: + if str(energy_types[parent_row]).strip() in { + "NT", + "NT1", + "NT2", + "NT3", + "NT4", + }: + valence = 4 + allowed = max( + 0, + min(valence - heavy_bonded, template_h - extra_bonds) - existing_h, + ) + group = group[:allowed] + if not group: + continue + + lengths = np.array( + [ + template["ideal_length"].get( + h, _template_bond_length(template, parent_name, h) + ) + for h in group + ] + ) + if not np.isfinite(lengths).all(): + n_unplaceable += len(group) + continue + + placed = _place_group( + template, + parent_name, + parent_position, + coords[heavy_rows] if heavy_bonded else np.zeros((0, 3)), + heavy_bonded, + group, + lengths, + name_to_row, + coords, + template_names_of=[ + n + for n in template["heavy_adjacency"].get(parent_name, []) + if n in name_to_row + ], + ) + if placed is None: + n_unplaceable += len(group) + continue + + # One free-torsion group per parent with a single heavy neighbour: its + # hydrogens rotate together about that one bond. + group_id = -1 + if heavy_bonded == 1: + group_id = next_group + next_group += 1 + + for h, position, length in zip(group, placed, lengths): + out["name"].append(h) + out["altloc"].append(altloc) + out["residue"].append(residue) + out["parent"].append(int(parent_row)) + out["position"].append(position) + out["bond_length"].append(float(length)) + out["group"].append(group_id) + + if verbose > 1: + print( + f"Planned {len(out['name'])} hydrogens " + f"({n_no_template} residues with no usable template, " + f"{n_unplaceable} hydrogens whose direction was not determined)" + ) + + return HydrogenPlan( + name=np.array(out["name"], dtype=" Optional[Tuple[np.ndarray, np.ndarray]]: + """Template-to-model transform for one hydrogen-bearing centre. + + Fitted over the parent and its **immediate** heavy neighbours only. That set is the + rigid unit which fixes the hydrogen directions: the bond lengths and angles at the + parent are library constants, while anything further out sits across a rotatable + torsion whose value is the model's, not the template's. + + Reaching one bond further -- as a whole-residue or two-shell fit does -- makes the + rotation compromise between the real local geometry and a torsion the model does not + share, which lands hydrogens well off their parent. Measured on 7L84 the two-shell + fit around ``CB`` aligned to 0.75 A RMSD and put 12% of side-chain hydrogens beyond + 1.5 A of the atom they belong to. + + Returns None when fewer than three neighbours match, which leaves the rotation + undetermined; the caller then constructs the direction from the bond graph instead. + """ + local = [parent_name] + [ + n for n in sorted(template["heavy_adjacency"].get(parent_name, [])) + ] + matched = [n for n in local if n in name_to_row and n in template["id_to_index"]] + if len(matched) < 3: + return None + source = np.array([template["coords"][template["id_to_index"][n]] for n in matched]) + target = np.array([coords[name_to_row[n]] for n in matched]) + return _kabsch(source, target) + + +def _template_covers_neighbours( + template: Dict, parent_name: str, graph_heavy_count: int +) -> bool: + """Whether the template knows every heavy atom actually bonded to the parent. + + A peptide-linked backbone nitrogen is bonded to the preceding residue's carbon, + which the free-amino-acid template has never heard of. Placing its hydrogen from the + template frame would ignore that substituent and drop the hydrogen on top of it, so + those centres are built from the graph instead. + """ + return len(template["heavy_adjacency"].get(parent_name, [])) >= graph_heavy_count + + +def _construct_directions( + parent: np.ndarray, + neighbours: np.ndarray, + n_hydrogens: int, + half_angle: float, +) -> Optional[np.ndarray]: + """Unit directions for hydrogens on a centre, from its bonded neighbours alone. + + Two rules cover every centre whose orientation the neighbours determine: + + * one hydrogen, any number of neighbours -- it opposes the sum of the bond unit + vectors, which is where the remaining valence points. This is the backbone amide + and alpha hydrogens. + * two hydrogens on a two-neighbour centre -- they straddle that same direction, + opened out by ``half_angle`` in the plane perpendicular to the neighbour pair. + + Anything else (three hydrogens, or two on a one-neighbour centre) has a free + torsion and is handled by the scan, not here. + + Returns + ------- + numpy.ndarray or None + Shape ``(n_hydrogens, 3)`` unit vectors, or None when the rules do not apply or + the geometry is degenerate. + """ + bonds = neighbours - parent + lengths = np.linalg.norm(bonds, axis=1) + if (lengths < 1e-8).any(): + return None + units = bonds / lengths[:, None] + + total = units.sum(0) + norm = np.linalg.norm(total) + if norm < 1e-6: + return None + opposed = -total / norm + + if n_hydrogens == 1: + return opposed[None, :] + + if n_hydrogens == 2 and len(units) == 2: + perpendicular = np.cross(units[0], units[1]) + perpendicular_norm = np.linalg.norm(perpendicular) + if perpendicular_norm < 1e-6: + return None + perpendicular = perpendicular / perpendicular_norm + cos, sin = np.cos(half_angle), np.sin(half_angle) + return np.array( + [ + opposed * cos + perpendicular * sin, + opposed * cos - perpendicular * sin, + ] + ) + + return None + + +def optimise_free_torsions( + plan: HydrogenPlan, + topology, + xyz, + steps: int = TORSION_SCAN_STEPS, +) -> HydrogenPlan: + """Rotate each free-torsion hydrogen group to the least-clashing angle. + + The template's dihedral for a hydroxyl, thiol, amine or methyl hydrogen is whatever + the library happened to deposit, so it carries no information. Each group is scanned + about its parent-neighbour axis and scored by repulsion against nearby heavy atoms; + the best angle wins. + + Repulsion only: this removes clashes but does not seek hydrogen bonds, so a hydroxyl + is placed out of the way rather than donated to an acceptor. Scoring donors properly + is a separate piece of physics. + + Returns + ------- + HydrogenPlan + The same plan with ``position`` updated in place for the scanned groups. + """ + if plan.n_hydrogens == 0 or not plan.rotatable.any(): + return plan + + coords = np.asarray(xyz.detach().cpu(), dtype=np.float64) + is_h = topology.atoms.is_hydrogen.cpu().numpy() + heavy_rows = np.nonzero(~is_h)[0] + heavy_coords = coords[heavy_rows] + + angles = np.linspace(0.0, 2.0 * np.pi, steps, endpoint=False) + + for group_id in sorted(set(plan.group[plan.rotatable].tolist())): + members = np.nonzero(plan.group == group_id)[0] + parent_row = int(plan.parent[members[0]]) + parent = coords[parent_row] + + neighbours = topology.atoms.neighbors(parent_row).cpu().numpy() + heavy_neighbours = neighbours[~is_h[neighbours]] + if len(heavy_neighbours) != 1: + continue + axis = parent - coords[int(heavy_neighbours[0])] + norm = np.linalg.norm(axis) + if norm < 1e-8: + continue + axis = axis / norm + + # Heavy atoms that could clash, excluding the parent and its own neighbour. + near = np.nonzero( + (np.linalg.norm(heavy_coords - parent, axis=1) < CLASH_CUTOFF) + )[0] + exclude = {parent_row, int(heavy_neighbours[0])} + near_coords = np.array( + [heavy_coords[i] for i in near if heavy_rows[i] not in exclude] + ) + + offsets = plan.position[members] - parent + best_angle, best_score = 0.0, np.inf + for angle in angles: + rotated = _rotate_about(offsets, axis, angle) + if len(near_coords) == 0: + best_angle = 0.0 + break + trial = parent + rotated + separation = np.linalg.norm( + trial[:, None, :] - near_coords[None, :, :], axis=-1 + ) + score = float((1.0 / np.maximum(separation, 0.5) ** 2).sum()) + if score < best_score: + best_score, best_angle = score, float(angle) + + if best_angle != 0.0: + plan.position[members] = parent + _rotate_about(offsets, axis, best_angle) + + return plan + + +def _rotate_about(vectors: np.ndarray, axis: np.ndarray, angle: float) -> np.ndarray: + """Rotate ``vectors`` about a unit ``axis`` through the origin, Rodrigues form.""" + cos, sin = np.cos(angle), np.sin(angle) + return ( + vectors * cos + + np.cross(axis, vectors) * sin + + np.outer(vectors @ axis, axis) * (1.0 - cos) + ) + + +__all__ = [ + "HydrogenPlan", + "HydrogenFrames", + "plan_hydrogens", + "optimise_free_torsions", + "augment_atom_table", + "augment_atom_table_with_maps", + "hydrogen_frames", + "template_atom_types", + "STANDARD_VALENCE", + "MAX_PLACEMENT_DISTANCE", + "TORSION_SCAN_STEPS", +] + + +def augment_atom_table(pdb, plan: HydrogenPlan, topology): + """Insert a plan's hydrogens into an atom table. + + Each hydrogen is inserted immediately after the residue it belongs to, not appended + at the end: the residue partition is built from contiguous runs of + ``(chain, resseq, icode)``, so appending would split every hydrogenated residue into + two nodes. + + Rows are copied from the parent, then the hydrogen's own name, element and position + are written over them. Everything else -- chain, residue, altloc, occupancy, + B-factor, record type -- is inherited, so a hydrogen starts from its parent's + displacement parameter and refines from there. + + Parameters + ---------- + pdb : pandas.DataFrame + Atom table to extend. + plan : HydrogenPlan + topology : Topology + Supplies the residue partition the insertion points come from. + + Returns + ------- + pandas.DataFrame + A new table with ``serial`` and ``index`` renumbered. + """ + return augment_atom_table_with_maps(pdb, plan, topology)[0] + + +def augment_atom_table_with_maps(pdb, plan: HydrogenPlan, topology): + """:func:`augment_atom_table` plus the row maps the insertion implies. + + Parameters + ---------- + pdb : pandas.DataFrame + Atom table to extend. + plan : HydrogenPlan + topology : Topology + Supplies the residue partition the insertion points come from. + + Returns + ------- + augmented : pandas.DataFrame + The extended table, ``serial`` and ``index`` renumbered. + old_to_new : numpy.ndarray + New row of every old row, shape ``(N_old,)``. Existing rows are never dropped, + so every entry is valid. + plan_to_new : numpy.ndarray + New row of every planned hydrogen, shape ``(plan.n_hydrogens,)``. + """ + import pandas as pd + + n_old = len(pdb) + if plan.n_hydrogens == 0: + return pdb.copy(), np.arange(n_old, dtype=np.int64), np.zeros(0, dtype=np.int64) + + by_residue: Dict[int, List[int]] = {} + for i, residue in enumerate(plan.residue.tolist()): + by_residue.setdefault(residue, []).append(i) + + old_to_new = np.full(n_old, -1, dtype=np.int64) + plan_to_new = np.full(plan.n_hydrogens, -1, dtype=np.int64) + pieces = [] + offset = 0 + for residue in range(topology.n_residues): + start = int(topology.residues.atom_start[residue]) + end = int(topology.residues.atom_end[residue]) + pieces.append(pdb.iloc[start:end]) + old_to_new[start:end] = offset + np.arange(end - start) + offset += end - start + + members = by_residue.get(residue) + if not members: + continue + rows = pdb.loc[pdb.index[plan.parent[members]]].copy() + rows["name"] = plan.name[members] + rows["element"] = plan.element[members] + rows["altloc"] = plan.altloc[members] + rows[["x", "y", "z"]] = plan.position[members] + if "anisou_flag" in rows.columns: + rows["anisou_flag"] = False + for column in ("u11", "u22", "u33", "u12", "u13", "u23"): + if column in rows.columns: + rows[column] = float("nan") + pieces.append(rows) + plan_to_new[members] = offset + np.arange(len(members)) + offset += len(members) + + augmented = pd.concat(pieces, ignore_index=True) + augmented["index"] = augmented.index.to_numpy(dtype=int) + if "serial" in augmented.columns: + augmented["serial"] = augmented.index.to_numpy(dtype=int) + 1 + augmented.attrs = dict(pdb.attrs) + return augmented, old_to_new, plan_to_new + + +@dataclass +class HydrogenFrames: + """Which atom-table rows are riding hydrogens, and the frame each one rides in. + + Row indices are into the atom table the frames were built for; ``-1`` marks an + absent atom. A hydrogen whose ``parent_row`` is ``-1`` is not a riding hydrogen at + all and is dropped by :meth:`remap`; one whose ``n1_row`` or ``n2_row`` is ``-1`` + keeps riding but with ``frame_valid`` False, so it translates rigidly with its + parent instead of turning with the frame. + + The frame is ``(parent, n1, n2)``: ``n1`` is the parent's first heavy neighbour, + ``n2`` its second, or -- for a parent with a single heavy neighbour, i.e. every + hydroxyl, thiol, amine and methyl -- a heavy neighbour of ``n1`` other than the + parent, so the hydrogen turns with the torsion about the ``n1-parent`` bond. + + Parameters + ---------- + h_row, parent_row, n1_row, n2_row : numpy.ndarray + ``int64`` rows, shape ``(H,)``. ``h_row`` is ``-1`` for a planned hydrogen + that has not been inserted into a table yet; :meth:`fill_planned_rows` sets it. + frame_valid : numpy.ndarray + Boolean ``(H,)``; False where the frame is incomplete. + torsion_group, rotation_group : numpy.ndarray, optional + Group labels, shape ``(H,)``; ``-1`` means no independent orientation. + Torsion groups rotate about the parent-to-n1 bond; rotation groups have + three rotational degrees of freedom in Cartesian space. Labels need not + be contiguous. Hydrogens in one group share their parent and orientation. + """ + + h_row: np.ndarray + parent_row: np.ndarray + n1_row: np.ndarray + n2_row: np.ndarray + frame_valid: np.ndarray + torsion_group: Optional[np.ndarray] = None + rotation_group: Optional[np.ndarray] = None + + def __post_init__(self) -> None: + """Fill absent orientation groups with the fixed-orientation sentinel.""" + for name in ("torsion_group", "rotation_group"): + value = getattr(self, name) + if value is None: + value = np.full(len(self.h_row), -1, dtype=np.int64) + value = np.asarray(value, dtype=np.int64) + if value.shape != self.h_row.shape: + raise ValueError(f"{name} must have shape {self.h_row.shape}") + setattr(self, name, value) + if ((self.torsion_group >= 0) & (self.rotation_group >= 0)).any(): + raise ValueError("A hydrogen cannot belong to both orientation types") + + @classmethod + def empty(cls) -> "HydrogenFrames": + """Frames for a model with no riding hydrogens.""" + z = np.zeros(0, dtype=np.int64) + return cls(z, z.copy(), z.copy(), z.copy(), np.zeros(0, dtype=bool)) + + @property + def n_hydrogens(self) -> int: + """How many hydrogens ride.""" + return len(self.h_row) + + @property + def n_planned(self) -> int: + """How many entries still await a row from :meth:`fill_planned_rows`.""" + return int((self.h_row < 0).sum()) + + def remap(self, old_to_new: np.ndarray) -> "HydrogenFrames": + """The frames over a reindexed table. + + Parameters + ---------- + old_to_new : numpy.ndarray + New row of each old row, ``-1`` where the atom was dropped. + + Returns + ------- + HydrogenFrames + Entries whose hydrogen or parent was dropped are removed; a lost ``n1`` or + ``n2`` leaves the entry with ``frame_valid`` False. Planned entries + (``h_row == -1``) are kept as planned. + """ + table = np.asarray(old_to_new, dtype=np.int64) + + def follow(rows: np.ndarray) -> np.ndarray: + out = np.full(len(rows), -1, dtype=np.int64) + present = rows >= 0 + out[present] = table[rows[present]] + return out + + h = follow(self.h_row) + h[self.h_row < 0] = -1 + parent = follow(self.parent_row) + n1 = follow(self.n1_row) + n2 = follow(self.n2_row) + keep = (parent >= 0) & ((h >= 0) | (self.h_row < 0)) + return HydrogenFrames( + h_row=h[keep], + parent_row=parent[keep], + n1_row=n1[keep], + n2_row=n2[keep], + frame_valid=self.frame_valid[keep] & (n1[keep] >= 0) & (n2[keep] >= 0), + torsion_group=np.where(n1[keep] >= 0, self.torsion_group[keep], -1), + rotation_group=self.rotation_group[keep], + ) + + def fill_planned_rows(self, rows: np.ndarray) -> "HydrogenFrames": + """Give the planned entries their table rows, in plan order. + + Parameters + ---------- + rows : numpy.ndarray + New row of each planned hydrogen, shape ``(n_planned,)``. + """ + rows = np.asarray(rows, dtype=np.int64) + planned = self.h_row < 0 + if int(planned.sum()) != len(rows): + raise ValueError( + f"{int(planned.sum())} planned hydrogens but {len(rows)} rows given" + ) + h = self.h_row.copy() + h[planned] = rows + return HydrogenFrames( + h, + self.parent_row.copy(), + self.n1_row.copy(), + self.n2_row.copy(), + self.frame_valid.copy(), + self.torsion_group.copy(), + self.rotation_group.copy(), + ) + + def sorted_by_row(self) -> "HydrogenFrames": + """The same frames ordered by ``h_row``.""" + order = np.argsort(self.h_row, kind="stable") + return HydrogenFrames( + self.h_row[order], + self.parent_row[order], + self.n1_row[order], + self.n2_row[order], + self.frame_valid[order], + self.torsion_group[order], + self.rotation_group[order], + ) + + def to_tensors(self, device=None) -> Dict[str, torch.Tensor]: + """Return frame and orientation arrays as tensors, keyed by field name.""" + return { + "h_row": torch.as_tensor( + self.h_row, dtype=torch.int64, device=device + ), # dtype-ok: row index; int64 required + "parent_row": torch.as_tensor( + self.parent_row, dtype=torch.int64, device=device + ), # dtype-ok: row index; int64 required + "n1_row": torch.as_tensor( + self.n1_row, dtype=torch.int64, device=device + ), # dtype-ok: row index; int64 required + "n2_row": torch.as_tensor( + self.n2_row, dtype=torch.int64, device=device + ), # dtype-ok: row index; int64 required + "frame_valid": torch.as_tensor( + self.frame_valid, dtype=torch.bool, device=device + ), + "torsion_group": torch.as_tensor(self.torsion_group, device=device), + "rotation_group": torch.as_tensor(self.rotation_group, device=device), + } + + @classmethod + def from_tensors( + cls, + h_row: torch.Tensor, + parent_row: torch.Tensor, + n1_row: torch.Tensor, + n2_row: torch.Tensor, + frame_valid: torch.Tensor, + torsion_group: Optional[torch.Tensor] = None, + rotation_group: Optional[torch.Tensor] = None, + ) -> "HydrogenFrames": + """Rebuild from the tensors :meth:`to_tensors` produced.""" + as_np = lambda t: np.asarray(t.detach().cpu().numpy(), dtype=np.int64) + return cls( + as_np(h_row), + as_np(parent_row), + as_np(n1_row), + as_np(n2_row), + np.asarray(frame_valid.detach().cpu().numpy(), dtype=bool), + None if torsion_group is None else as_np(torsion_group), + None if rotation_group is None else as_np(rotation_group), + ) + + def __repr__(self) -> str: + return ( + f"HydrogenFrames(n_hydrogens={self.n_hydrogens}, " + f"planned={self.n_planned}, rigid={int((~self.frame_valid).sum())})" + ) + + +def _frame_atoms(topology, parent_row: int, altloc: str) -> Tuple[int, int]: + """``(n1, n2)`` rows for a frame anchored on ``parent_row``; ``-1`` where absent. + + ``n1`` is the parent's first heavy neighbour in the conformer, ``n2`` its second, + or a heavy neighbour of ``n1`` other than the parent when the parent has only one. + """ + heavy, _ = _split_neighbours(topology, parent_row, altloc) + if len(heavy) == 0: + return -1, -1 + n1 = int(heavy[0]) + if len(heavy) >= 2: + return n1, int(heavy[1]) + grand, _ = _split_neighbours(topology, n1, altloc) + grand = grand[grand != parent_row] + return n1, (int(grand[0]) if len(grand) else -1) + + +def hydrogen_frames(topology, plan: Optional[HydrogenPlan] = None) -> HydrogenFrames: + """Riding frames for every hydrogen the table has, plus the ones a plan adds. + + Read off the bond graph, not off distances, so a stretched or predicted model + still frames each hydrogen on its bonded parent. + + Parameters + ---------- + topology : Topology + Connectivity of the table the frames index into. + plan : HydrogenPlan, optional + Hydrogens about to be inserted. Their entries carry ``h_row == -1`` until + :meth:`HydrogenFrames.fill_planned_rows` is given the rows the insertion made; + their parent and frame atoms are rows of the *current* table, to be carried + through :meth:`HydrogenFrames.remap` with everything else. + + Returns + ------- + HydrogenFrames + Deposited hydrogens first, in row order, then planned ones in plan order. A + hydrogen bonded to no heavy atom is left out: nothing can carry it. + """ + atoms = topology.atoms + is_h = atoms.is_hydrogen.cpu().numpy() + altlocs = np.char.strip(atoms.altloc.astype(str)) + + rows: List[Tuple[int, int, int, int]] = [] + for h in np.nonzero(is_h)[0].tolist(): + neighbours = atoms.neighbors(h).cpu().numpy() + heavy = neighbours[~is_h[neighbours]] + if len(heavy) == 0: + continue + parent = int(heavy[0]) + altloc = str(altlocs[h]) + n1, n2 = _frame_atoms(topology, parent, altloc) + rows.append((h, parent, n1, n2)) + + if plan is not None: + for k in range(plan.n_hydrogens): + parent = int(plan.parent[k]) + n1, n2 = _frame_atoms(topology, parent, str(plan.altloc[k]).strip()) + rows.append((-1, parent, n1, n2)) + + if not rows: + return HydrogenFrames.empty() + arr = np.array(rows, dtype=np.int64) + torsion = np.full(len(rows), -1, dtype=np.int64) + rotation = np.full(len(rows), -1, dtype=np.int64) + elements = np.char.upper(np.char.strip(atoms.element.astype(str))) + groups = {} + n_existing = len(rows) - (0 if plan is None else plan.n_hydrogens) + for i, (h, parent, _, _) in enumerate(rows): + altloc = str(altlocs[h]) if h >= 0 else str(plan.altloc[i - n_existing]).strip() + groups.setdefault((parent, altloc), []).append(i) + for group, ((parent, altloc), members) in enumerate(groups.items()): + heavy, _ = _split_neighbours(topology, parent, altloc) + residue = int(atoms.residue_of[parent]) + is_water = str(topology.residues.resname[residue]).strip() == "HOH" + if len(heavy) == 0 or is_water: + rotation[members] = group + elif len(heavy) == 1 and ( + (elements[parent] == "C" and len(members) == 3) + or elements[parent] in ("O", "S") + ): + # Planar amide NH2 groups also have one heavy neighbour, but their + # orientation is constrained by conjugation rather than freely rotatable. + torsion[members] = group + return HydrogenFrames( + h_row=arr[:, 0], + parent_row=arr[:, 1], + n1_row=arr[:, 2], + n2_row=arr[:, 3], + frame_valid=(arr[:, 2] >= 0) & (arr[:, 3] >= 0), + torsion_group=torsion, + rotation_group=rotation, + ) diff --git a/torchref/topology/monomer/__init__.py b/torchref/topology/monomer/__init__.py new file mode 100644 index 00000000..b655740b --- /dev/null +++ b/torchref/topology/monomer/__init__.py @@ -0,0 +1,13 @@ +"""Monomer dictionaries: finding them, reading them, and patching them. + +Everything that turns files on disk into the ideal geometry a template carries. +:mod:`library` resolves and caches the CCP4 Monomer Library, fetching a component on +demand; :mod:`cif` reads a dictionary into DataFrames per section; :mod:`modifications` +applies the ``chem_mod`` records that change a template when a link forms. + +This is the data source. What is built from it -- the connectivity, the values over its +edges -- is the rest of :mod:`torchref.topology`. Import from the defining submodule +rather than from here. + +Reference: Long, F., et al. (2017). AceDRG. Acta Cryst. D73, 112-122. +""" diff --git a/torchref/restraints/restraints_helper.py b/torchref/topology/monomer/cif.py similarity index 96% rename from torchref/restraints/restraints_helper.py rename to torchref/topology/monomer/cif.py index b2d2276b..95731424 100644 --- a/torchref/restraints/restraints_helper.py +++ b/torchref/topology/monomer/cif.py @@ -94,7 +94,7 @@ def find_cif_file_in_library(resname): Delegates to :meth:`MonomerLibraryManager.get_cif_file`, whose last resort is an on-demand download. """ - from torchref.restraints.library import get_library_manager + from torchref.topology.monomer.library import get_library_manager return get_library_manager().get_cif_file(resname) @@ -128,14 +128,14 @@ def read_library_blocks(): """Return ``mon_lib_list.cif`` split into ``{block_name: block_text}``. Shared by :func:`read_link_definitions` and - :func:`~torchref.restraints.modifications.read_mod_definitions`, which read + :func:`~torchref.topology.monomer.modifications.read_mod_definitions`, which read disjoint parts of the same 4 MB file. Warnings -------- Cached process-wide; the returned dict is shared, so do not mutate it. """ - from torchref.restraints.library import get_library_manager + from torchref.topology.monomer.library import get_library_manager path = str(get_library_manager().get_link_definitions_path()) with open(path) as handle: @@ -155,7 +155,7 @@ def read_link_definitions(): link_list : DataFrame or None The ``chem_link`` table, or None if the file has no ``link_list`` block. Its ``mod_id_1``/``mod_id_2`` columns name the modifications each link - applies to its partners -- see :mod:`torchref.restraints.modifications`. + applies to its partners -- see :mod:`torchref.topology.monomer.modifications`. Warnings -------- diff --git a/torchref/restraints/library.py b/torchref/topology/monomer/library.py similarity index 96% rename from torchref/restraints/library.py rename to torchref/topology/monomer/library.py index e8375f95..0eea7309 100644 --- a/torchref/restraints/library.py +++ b/torchref/topology/monomer/library.py @@ -26,7 +26,12 @@ ) # Bundled package data location -_BUNDLED_PATH = Path(__file__).parent.parent / "data" / "monomer_library" +# Three levels up, not two: this module sits at torchref/topology/monomer/, so the +# package root is its great-grandparent. Computing it by depth is fragile, which is +# why the level is spelled out rather than left to be counted. +_BUNDLED_PATH = ( + Path(__file__).resolve().parents[2] / "data" / "monomer_library" +) # Legacy external monomer library path _LEGACY_PATH = ROOT_TORCHREF / "external_monomer_library" diff --git a/torchref/restraints/modifications.py b/torchref/topology/monomer/modifications.py similarity index 85% rename from torchref/restraints/modifications.py rename to torchref/topology/monomer/modifications.py index f9b10f02..ac70afac 100644 --- a/torchref/restraints/modifications.py +++ b/torchref/topology/monomer/modifications.py @@ -13,10 +13,13 @@ peptide carbonyl carbon the intra-residue ``CA-C-O`` plus the link's ``CA-C-N`` and ``O-C-N`` sum to 360 deg only once ``DEL-OXT`` has been applied. -``_chem_mod_atom`` and ``_chem_mod_tree`` are deliberately ignored. The restraint -builders match library restraints against the atoms actually present in the model -and silently skip any whose atoms are missing, so adding or deleting atom -*definitions* changes nothing downstream; only the restraint sections matter. +Of ``_chem_mod_atom`` only the ``change`` rows are applied, and only to the energy +type and charge: ``DEL-HN1`` retypes the backbone ``N`` from the free-amine ``NT3`` to +the amide ``NH1``, which is what the contact radii and hydrogen-bond roles read. +Adding or deleting atom *definitions* is ignored, as is ``_chem_mod_tree``: the +restraint builders match library restraints against the atoms actually present in +the model and silently skip any whose atoms are missing, so those rows change nothing +downstream. """ from functools import lru_cache @@ -24,10 +27,11 @@ import pandas as pd -from torchref.restraints.restraints_helper import read_library_blocks +from torchref.topology.monomer.cif import read_library_blocks #: CIF category -> section name, matching :func:`read_link_definitions`. _CATEGORY_MAP = { + "chem_mod_atom": "atoms", "chem_mod_bond": "bonds", "chem_mod_angle": "angles", "chem_mod_tor": "torsions", @@ -38,6 +42,7 @@ #: Per section: the columns identifying a restraint, and the columns a #: ``change``/``add`` row may overwrite. _ATOM_COLUMNS = { + "atoms": ("atom_id",), "bonds": ("atom1", "atom2"), "angles": ("atom1", "atom2", "atom3"), "torsions": ("atom1", "atom2", "atom3", "atom4"), @@ -45,12 +50,18 @@ "chirals": ("atom_centre", "atom1", "atom2", "atom3"), } _VALUE_COLUMNS = { + "atoms": ("type_energy", "charge"), "bonds": ("value", "sigma"), "angles": ("value", "sigma"), "torsions": ("value", "sigma", "periodicity"), "planes": ("sigma",), "chirals": ("volume_sign",), } +#: Value columns that hold text, so they are not coerced to numbers. +_TEXT_COLUMNS = {"volume_sign", "type_energy"} +#: Sections where only ``change`` rows are meaningful; ``add``/``delete`` rows are +#: skipped rather than editing the atom list (see the module docstring). +_CHANGE_ONLY = {"atoms"} #: Sections whose row is meaningless without a target value, so an ``add`` row #: carrying only ``.`` placeholders is dropped rather than appended as NaN. _REQUIRES_VALUE = {"bonds", "angles", "torsions"} @@ -63,6 +74,8 @@ def _restraint_key(section: str, row: Mapping) -> tuple: outer pair, torsions on the atom quadruple in either direction, chirals on the centre plus the unordered substituents. Planes key on ``(plane_id, atom)``. """ + if section == "atoms": + return (str(row["atom_id"]).strip(),) if section == "bonds": return tuple(sorted((row["atom1"], row["atom2"]))) if section == "angles": @@ -137,7 +150,10 @@ def _standardize_mod_columns(df: pd.DataFrame, section: str) -> pd.DataFrame: "atom_id_3": "atom3", "atom_id_4": "atom4", "atom_id_centre": "atom_centre", - "atom_id": "atom", + # Planes name their atom ``atom_id``; the atoms section keeps that name. + **({} if section == "atoms" else {"atom_id": "atom"}), + "new_type_energy": "type_energy", + "new_charge": "charge", "new_value_dist": "value", "new_value_dist_esd": "sigma", "new_value_angle": "value", @@ -152,7 +168,10 @@ def _standardize_mod_columns(df: pd.DataFrame, section: str) -> pd.DataFrame: for column in _VALUE_COLUMNS[section]: if column not in df.columns: df[column] = pd.NA - elif column != "volume_sign": + elif column in _TEXT_COLUMNS: + text = df[column].astype(str).str.strip() + df[column] = text.where(~text.isin(["", ".", "?", "nan"]), pd.NA) + else: df[column] = pd.to_numeric(df[column], errors="coerce") columns.append(column) @@ -167,7 +186,7 @@ def link_modifications(link_list: Optional[pd.DataFrame]) -> Dict[str, Tuple]: ---------- link_list : pandas.DataFrame or None The ``chem_link`` table returned by - :func:`~torchref.restraints.restraints_helper.read_link_definitions`. + :func:`~torchref.topology.monomer.cif.read_link_definitions`. Returns ------- @@ -200,7 +219,7 @@ def apply_modifications( ---------- comp : mapping of str to pandas.DataFrame One component's restraints, as produced by - :func:`~torchref.restraints.restraints_helper.read_cif`. + :func:`~torchref.topology.monomer.cif.read_cif`. mod_ids : sequence of str Modification IDs to apply, in order. Unknown IDs are ignored. mod_dict : mapping @@ -230,9 +249,14 @@ def apply_modifications( continue target = result.get(section) if target is None: + if section in _CHANGE_ONLY: + continue target = pd.DataFrame( columns=list(_ATOM_COLUMNS[section]) + list(_VALUE_COLUMNS[section]) ) + for column in _VALUE_COLUMNS[section]: + if column not in target.columns: + target[column] = pd.NA result[section] = _apply_section(target, mod_rows, section) return result @@ -253,6 +277,8 @@ def _apply_section( additions = [] for _, mod_row in mod_rows.iterrows(): function = mod_row["function"] + if section in _CHANGE_ONLY and function != "change": + continue key = _restraint_key(section, mod_row) positions = [p for p in by_key.get(key, []) if p not in dropped] diff --git a/torchref/restraints/neighbor_search.py b/torchref/topology/nonbonded.py similarity index 95% rename from torchref/restraints/neighbor_search.py rename to torchref/topology/nonbonded.py index 9204ca4f..404e7a6f 100644 --- a/torchref/restraints/neighbor_search.py +++ b/torchref/topology/nonbonded.py @@ -10,7 +10,7 @@ the input coordinates live on (CPU or GPU). """ -from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple import numpy as np import torch @@ -85,8 +85,8 @@ def prefilter_symop_offsets( valid_ops.append(op_idx) valid_offsets.append([dx, dy, dz]) - op_indices = torch.tensor(valid_ops, dtype=torch.long, device=device) - cell_offsets = torch.tensor(valid_offsets, dtype=torch.long, device=device) + op_indices = torch.tensor(valid_ops, dtype=torch.long, device=device) # dtype-ok: symmetry-operator index tensor; int64 + cell_offsets = torch.tensor(valid_offsets, dtype=torch.long, device=device) # dtype-ok: integer cell-offset lattice vectors; symmetry-image metadata return op_indices, cell_offsets @@ -146,7 +146,7 @@ def assign_to_grid( gd = grid_dims.to(device=device, dtype=fdtype) cell_ijk = (frac_wrapped * gd[None, None, :]).long() cell_ijk = cell_ijk.clamp( - min=torch.zeros(3, dtype=torch.long, device=device), + min=torch.zeros(3, dtype=torch.long, device=device), # dtype-ok: clamp min-bound for long grid-index tensor; matches int64 max=(grid_dims - 1).to(device), ) @@ -188,14 +188,14 @@ def build_cell_list( unique_cells, counts = torch.unique_consecutive( sorted_cells, return_counts=True ) - starts = torch.zeros(len(unique_cells) + 1, dtype=torch.long, device=device) + starts = torch.zeros(len(unique_cells) + 1, dtype=torch.long, device=device) # dtype-ok: CSR boundary/offset array; int64 required starts[1:] = counts.cumsum(0) cell_lookup = torch.full( - (n_grid_total,), -1, dtype=torch.long, device=device + (n_grid_total,), -1, dtype=torch.long, device=device # dtype-ok: grid-cell to index lookup table; used for indexing, int64 ) cell_lookup[unique_cells] = torch.arange( - len(unique_cells), dtype=torch.long, device=device + len(unique_cells), dtype=torch.long, device=device # dtype-ok: index values written into lookup table; int64 ) return sort_order, unique_cells, starts, cell_lookup @@ -233,7 +233,7 @@ def _get_canonical_offsets_14(device: torch.device) -> torch.Tensor: offsets.append([dx, dy, dz]) assert len(offsets) == 14, f"expected 14 canonical offsets, got {len(offsets)}" _NEIGHBOR_OFFSETS_14 = torch.tensor( - offsets, dtype=torch.long, device=device + offsets, dtype=torch.long, device=device # dtype-ok: grid neighbor-cell offset deltas used to compute index; int64 ) return _NEIGHBOR_OFFSETS_14 @@ -493,7 +493,7 @@ def find_pairs_periodic_grid_v2( all_pair_combo_j.append(cj) if not all_pair_atom_i: - empty = torch.tensor([], dtype=torch.long, device=device) + empty = torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty atom-pair index placeholder; int64 required return empty, empty, empty return ( @@ -517,11 +517,11 @@ def exclusion_set_to_hash( Hash: min(i,j) * max_idx + max(i,j), sorted for searchsorted. """ if not exclusion_set: - return torch.tensor([], dtype=torch.long, device=device) + return torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty exclusion-hash placeholder; int64 arr = np.array(list(exclusion_set), dtype=np.int64) hashes = arr[:, 0] * max_idx + arr[:, 1] # already (min, max) hashes.sort() - return torch.tensor(hashes, dtype=torch.long, device=device) + return torch.tensor(hashes, dtype=torch.long, device=device) # dtype-ok: packed pair-hash key for searchsorted; int64 avoids overflow def filter_pairs( @@ -629,11 +629,11 @@ def build_vdw_restraints_gpu( sg = SG(sg) empty_result = { - "indices": torch.zeros(0, 2, dtype=torch.long, device=device), + "indices": torch.zeros(0, 2, dtype=torch.long, device=device), # dtype-ok: atom-pair index tensor; torch indexing requires int64 "min_distances": torch.zeros(0, dtype=get_float_dtype(), device=device), "sigmas": torch.zeros(0, dtype=get_float_dtype(), device=device), - "symop_indices": torch.zeros(0, dtype=torch.long, device=device), - "cell_offsets": torch.zeros(0, 3, dtype=torch.long, device=device), + "symop_indices": torch.zeros(0, dtype=torch.long, device=device), # dtype-ok: symmetry-operator index tensor; int64 + "cell_offsets": torch.zeros(0, 3, dtype=torch.long, device=device), # dtype-ok: integer cell-offset lattice vectors; symmetry-image metadata } # Step 1: prefilter symop combos @@ -655,10 +655,10 @@ def build_vdw_restraints_gpu( if len(identity_indices) == 0: # Identity not in valid combos — should not happen, but add it op_indices = torch.cat([ - torch.zeros(1, dtype=torch.long, device=device), op_indices + torch.zeros(1, dtype=torch.long, device=device), op_indices # dtype-ok: identity prepended to symop-index tensor; int64 ]) cell_offsets_valid = torch.cat([ - torch.zeros(1, 3, dtype=torch.long, device=device), cell_offsets_valid + torch.zeros(1, 3, dtype=torch.long, device=device), cell_offsets_valid # dtype-ok: identity prepended to cell-offset tensor; int64 ]) identity_combo = 0 M = len(op_indices) @@ -774,7 +774,7 @@ def build_vdw_restraints_gpu( "valid_op_indices": op_indices, "valid_cell_offsets": cell_offsets_valid, "grid_dims": grid_dims, - "identity_combo": torch.tensor(identity_combo, dtype=torch.long, device=device), + "identity_combo": torch.tensor(identity_combo, dtype=torch.long, device=device), # dtype-ok: combo index scalar into symop/offset arrays; int64 } if verbose > 0: @@ -856,7 +856,7 @@ def find_h_vdw_pairs_gpu( xyz_all = torch.cat([xyz_heavy, xyz_h], dim=0) # (N_all, 3) n_all = xyz_all.shape[0] - empty = torch.tensor([], dtype=torch.long, device=device) + empty = torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty index placeholder tensor; int64 required if n_all == 0: return empty, empty, empty diff --git a/torchref/restraints/ramachandran.py b/torchref/topology/ramachandran.py similarity index 96% rename from torchref/restraints/ramachandran.py rename to torchref/topology/ramachandran.py index fc1cf691..9d45fc43 100644 --- a/torchref/restraints/ramachandran.py +++ b/torchref/topology/ramachandran.py @@ -33,7 +33,7 @@ #: Number of 1° bins per axis, covering [-180, +180). GRID_SIZE = 360 -_DATA_FILE = Path(__file__).resolve().parent.parent / "data" / "rama_nll_surfaces.pt" +_DATA_FILE = Path(__file__).resolve().parents[1] / "data" / "rama_nll_surfaces.pt" def load_nll_surfaces(device: torch.device) -> torch.Tensor: diff --git a/torchref/topology/residue_graph.py b/torchref/topology/residue_graph.py new file mode 100644 index 00000000..0a68a21c --- /dev/null +++ b/torchref/topology/residue_graph.py @@ -0,0 +1,296 @@ +"""The sequence level of a topology: residues as nodes, links as edges. + +A residue node is one instance of a monomer-library template. Its identity is +``(chain, resseq, icode)`` -- the insertion code included, so residues 100 and 100A are +distinct nodes. Edges are the inter-residue links: peptide bonds, disulfides, and +explicit ``LINK`` records. + +Holds no tensors, so this is a plain dataclass; the atom-level tensors live on +:class:`~torchref.topology.atom_graph.AtomGraph`. +""" + +from dataclasses import dataclass, field +from typing import Dict, List, Sequence, Tuple + +import numpy as np +import torch + +#: SG-SG separation below which two cysteines are taken to be disulfide-bonded. +DISULFIDE_MAX_DISTANCE = 2.5 + +#: Lower bound guarding against an atom being paired with itself through a +#: coordinate duplicate. +DISULFIDE_MIN_DISTANCE = 0.1 + + +def _residue_runs( + chain: np.ndarray, resseq: np.ndarray, icode: np.ndarray +) -> Tuple[np.ndarray, np.ndarray]: + """Start and end row of each contiguous ``(chain, resseq, icode)`` run. + + Contiguity is assumed rather than checked, matching the reader's guarantee that a + structure's atoms arrive grouped by residue. A residue split across two + non-adjacent runs would become two nodes. + + Returns + ------- + starts, ends : numpy.ndarray + Half-open row ranges, shape ``(R,)`` each. + """ + n = len(chain) + if n == 0: + return np.zeros(0, dtype=np.int64), np.zeros(0, dtype=np.int64) + changed = np.zeros(n, dtype=bool) + changed[0] = True + changed[1:] = ( + (chain[1:] != chain[:-1]) + | (resseq[1:] != resseq[:-1]) + | (icode[1:] != icode[:-1]) + ) + starts = np.nonzero(changed)[0].astype(np.int64) + ends = np.append(starts[1:], n).astype(np.int64) + return starts, ends + + +@dataclass +class ResidueGraph: + """Residues as nodes, inter-residue links as edges. + + Parameters + ---------- + chain, resseq, icode, resname : numpy.ndarray + Per-residue identity, shape ``(R,)``. + template_key : numpy.ndarray + Restraint-dictionary key per residue, shape ``(R,)``. Either the residue name + or a link-modified variant such as ``'ALA:DEL-HN1+DEL-OXT'``. + atom_start, atom_end : numpy.ndarray + Half-open row range of each residue's atoms, shape ``(R,)``. + link_pairs : numpy.ndarray + Residue index pairs, shape ``(L, 2)``. For a peptide link the first entry + donates its ``C`` and the second its ``N``. + link_kind : numpy.ndarray + Link type per edge, shape ``(L,)``: ``'TRANS'``, ``'PTRANS'``, ``'disulf'`` or + ``'LINK'``. + """ + + chain: np.ndarray + resseq: np.ndarray + icode: np.ndarray + resname: np.ndarray + template_key: np.ndarray + atom_start: np.ndarray + atom_end: np.ndarray + link_pairs: np.ndarray = field( + default_factory=lambda: np.zeros((0, 2), dtype=np.int64) + ) + link_kind: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=" int: + """Number of residue nodes.""" + return len(self.chain) + + def key(self, i: int) -> Tuple[str, int, str]: + """Identity of residue ``i`` as ``(chain, resseq, icode)``.""" + return (str(self.chain[i]), int(self.resseq[i]), str(self.icode[i])) + + def atom_rows(self, i: int) -> range: + """Row range of residue ``i``'s atoms.""" + return range(int(self.atom_start[i]), int(self.atom_end[i])) + + def copy(self) -> "ResidueGraph": + """An independent copy sharing no arrays with this one.""" + return ResidueGraph( + chain=self.chain.copy(), + resseq=self.resseq.copy(), + icode=self.icode.copy(), + resname=self.resname.copy(), + template_key=self.template_key.copy(), + atom_start=self.atom_start.copy(), + atom_end=self.atom_end.copy(), + link_pairs=self.link_pairs.copy(), + link_kind=self.link_kind.copy(), + ) + + def subset( + self, keep: np.ndarray, atom_start: np.ndarray, atom_end: np.ndarray + ) -> "ResidueGraph": + """The residues in ``keep``, with the atom ranges the caller recomputed. + + Parameters + ---------- + keep : numpy.ndarray + Boolean mask over residues, shape ``(R,)``. + atom_start, atom_end : numpy.ndarray + New half-open atom ranges for the surviving residues, in their order, shape + ``(R_kept,)``. Passed in rather than derived here because only the caller + knows how the atoms were renumbered. + + Returns + ------- + ResidueGraph + Link edges are kept only where **both** endpoints survive, and reindexed. A + peptide bond to a residue that is gone is not a peptide bond, and keeping it + would leave an edge pointing outside the graph. + """ + remap = np.full(self.n_residues, -1, dtype=np.int64) + remap[keep] = np.arange(int(keep.sum()), dtype=np.int64) + + if len(self.link_pairs): + mapped = remap[self.link_pairs] + survives = (mapped >= 0).all(axis=1) + link_pairs = mapped[survives] + link_kind = self.link_kind[survives] + else: + link_pairs = np.zeros((0, 2), dtype=np.int64) + link_kind = np.zeros(0, dtype=" np.ndarray: + """Link edges of one kind, shape ``(L_k, 2)``.""" + if len(self.link_kind) == 0: + return np.zeros((0, 2), dtype=np.int64) + return self.link_pairs[self.link_kind == kind] + + def __repr__(self) -> str: + kinds = ( + {k: int((self.link_kind == k).sum()) for k in np.unique(self.link_kind)} + if len(self.link_kind) + else {} + ) + return f"ResidueGraph(n_residues={self.n_residues}, links={kinds})" + + +def build_residue_nodes( + chain: np.ndarray, + resseq: np.ndarray, + icode: np.ndarray, + resname: np.ndarray, +) -> Dict[str, np.ndarray]: + """Per-residue identity arrays and atom ranges from per-atom columns. + + Returns + ------- + dict + ``chain``, ``resseq``, ``icode``, ``resname``, ``atom_start``, ``atom_end``, + each shape ``(R,)``. + """ + starts, ends = _residue_runs(chain, resseq, icode) + return { + "chain": chain[starts], + "resseq": resseq[starts], + "icode": icode[starts], + "resname": resname[starts], + "atom_start": starts, + "atom_end": ends, + } + + +def find_peptide_links( + nodes: Dict[str, np.ndarray], + names_by_residue: List[set], +) -> List[Tuple[int, int]]: + """Sequence-adjacent residue pairs carrying a C-N peptide bond. + + Two residues are sequence-adjacent when they are neighbours in their chain's + ``(resseq, icode)`` ordering **and** either share a ``resseq`` -- an insertion-code + step such as 100 to 100A -- or differ by exactly one. The second condition is what + stops a chain break being bridged: residues 49 and 56 are neighbours in the ordering + but not in the sequence. + + The C/N test then narrows to pairs that actually carry the bond, which is the + condition the link builders apply implicitly when they look the two atoms up. + + Parameters + ---------- + nodes : dict + Output of :func:`build_residue_nodes`. + names_by_residue : list of set + Atom names present in each residue. + + Returns + ------- + list of tuple of int + ``(residue donating C, residue donating N)`` pairs. + """ + by_chain: Dict[str, List[int]] = {} + for i in range(len(nodes["chain"])): + by_chain.setdefault(str(nodes["chain"][i]), []).append(i) + + pairs = [] + for members in by_chain.values(): + ordered = sorted( + members, key=lambda i: (int(nodes["resseq"][i]), str(nodes["icode"][i])) + ) + for a, b in zip(ordered, ordered[1:]): + step = int(nodes["resseq"][b]) - int(nodes["resseq"][a]) + if step not in (0, 1): + continue + if "C" in names_by_residue[a] and "N" in names_by_residue[b]: + pairs.append((a, b)) + return pairs + + +def find_disulfide_links( + sg_rows: Sequence[int], + residue_of_row: Dict[int, int], + xyz: torch.Tensor, +) -> List[Tuple[int, int]]: + """``SG`` atom pairs within :data:`DISULFIDE_MAX_DISTANCE` in different residues. + + Pairing is per **atom**, not per residue, because a cysteine modelled in two + alternative conformations carries two ``SG`` atoms and each can form its own bond. + Pairing per residue would keep only one of them. Two ``SG`` atoms of the same + residue -- its own alternative conformers -- are within bonding distance of each + other and are excluded by the differing-residue test. + + Parameters + ---------- + sg_rows : sequence of int + Atom rows of every ``SG`` under consideration. + residue_of_row : dict + ``{atom row: residue index}`` for those rows. + xyz : torch.Tensor + Cartesian coordinates, shape ``(N, 3)``. + + Returns + ------- + list of tuple of int + Atom-row pairs, lower row first, ascending. + """ + rows = list(sg_rows) + if len(rows) < 2: + return [] + idx = torch.as_tensor(rows, dtype=torch.int64, device=xyz.device) # dtype-ok: residue-atom index tensor; int64 index required + dist = torch.cdist(xyz[idx], xyz[idx]) + close = (dist > DISULFIDE_MIN_DISTANCE) & (dist < DISULFIDE_MAX_DISTANCE) + + a, b = torch.triu_indices(len(rows), len(rows), offset=1, device=xyz.device) + hit = close[a, b] + pairs = [] + for i, j in zip(a[hit].cpu().tolist(), b[hit].cpu().tolist()): + row_i, row_j = rows[i], rows[j] + if residue_of_row[row_i] != residue_of_row[row_j]: + pairs.append((row_i, row_j)) + return pairs + + +__all__ = [ + "ResidueGraph", + "build_residue_nodes", + "find_peptide_links", + "find_disulfide_links", + "DISULFIDE_MAX_DISTANCE", + "DISULFIDE_MIN_DISTANCE", +] diff --git a/torchref/topology/restraint_sets.py b/torchref/topology/restraint_sets.py new file mode 100644 index 00000000..45c079ae --- /dev/null +++ b/torchref/topology/restraint_sets.py @@ -0,0 +1,180 @@ +"""Ideal values and sigmas, layered over a topology's edges. + +The topology says what is connected. What the ideal geometry *is* lives here, keyed to +the same edges, so one connectivity can carry monomer-library targets, force-field +parameters or ADP-similarity sigmas without a second copy of the edges. + +:func:`assemble_entries` turns a topology plus its values into the nested mapping the +restraint consumers read, ``entries[edge_type][origin][property]``. Indices are +**views** into the contiguous edge blocks, so taking a per-origin subset costs nothing +and an in-place edit to a block is visible through every view of it. Access is three +dict lookups with no allocation, which is the point: it sits on the geometry targets' +hot path. +""" + +from typing import Dict, Optional, Sequence, Tuple + +import numpy as np +import torch + +from torchref.config import get_float_dtype + +#: Origins making up each edge type's ``all`` group -- what the geometry targets read. +#: ``None`` means every origin present. ``phi`` and ``psi`` are conformationally free +#: and carry no target, and ``omega`` has its own von Mises target, so the torsion group +#: holds only the two origins that are ordinary restrained torsions. +ALL_MEMBERS: Dict[str, Optional[Tuple[str, ...]]] = { + "bond": None, + "angle": None, + "torsion": ("intra", "disulfide"), +} + +#: Integer-valued edge properties, kept as ``int64`` rather than the float dtype. +_INTEGER_PROPERTIES = frozenset({"periods", "symop_indices", "cell_offsets"}) + +#: Boolean edge properties. +_BOOL_PROPERTIES = frozenset({"is_proline"}) + + +def to_tensor(values, prop: str, device=None) -> torch.Tensor: + """A per-edge value array as a tensor of the dtype that property calls for.""" + if isinstance(values, torch.Tensor): + return values.to(device=device) if device is not None else values + if prop in _INTEGER_PROPERTIES: + dtype = torch.int64 # dtype-ok: dtype var for index tensors; int64 index required + elif prop in _BOOL_PROPERTIES: + dtype = torch.bool + else: + dtype = get_float_dtype() + return torch.as_tensor(np.asarray(values), dtype=dtype, device=device) + + +def _contiguous_span( + bounds: Dict[str, Tuple[int, int]], members: Sequence[str] +) -> Optional[Tuple[int, int]]: + """The single row range covering ``members``, or None if they are not adjacent. + + Adjacency is what makes the ``all`` group a view instead of a copy, so the block + layout in :data:`~torchref.topology.edges.ORIGIN_ORDER` deliberately keeps each edge + type's ``all`` members together. + """ + present = [m for m in members if m in bounds] + if not present: + return None + spans = sorted(bounds[m] for m in present) + for (_, end), (start, _) in zip(spans, spans[1:]): + if end != start: + return None + return spans[0][0], spans[-1][1] + + +def _group( + indices: torch.Tensor, values: Dict[str, torch.Tensor] +) -> Dict[str, torch.Tensor]: + """One restraint group: its indices plus whatever properties it carries.""" + group = {"indices": indices} + group.update(values) + return group + + +def _all_group( + block, per_origin_values: Dict[str, Dict[str, torch.Tensor]], members +) -> Optional[Dict[str, torch.Tensor]]: + """The combined group a geometry target reads, or None when it would be empty. + + Carries only properties present in **every** member origin, so a property one member + lacks does not silently become a partial array. + """ + origins = ( + block.origins() + if members is None + else [m for m in members if m in block.origin_bounds] + ) + if not origins: + return None + + span = _contiguous_span(block.origin_bounds, origins) + if span is not None: + start, end = span + indices = block.indices[start:end] + else: + # Not reachable with the shipped layouts; kept so a future origin order that + # separates the members degrades to a copy rather than silently misaligning. + indices = torch.cat([block.origin(o) for o in origins], dim=0) + + shared = set(per_origin_values.get(origins[0], {})) + for origin in origins[1:]: + shared &= set(per_origin_values.get(origin, {})) + + values = { + prop: torch.cat([per_origin_values[o][prop] for o in origins], dim=0) + for prop in sorted(shared) + } + return _group(indices, values) + + +def assemble_entries( + topology, + values: Dict[str, Dict[str, Dict[str, torch.Tensor]]], +) -> Dict[str, Dict]: + """Build the nested mapping the restraint consumers read. + + Parameters + ---------- + topology : Topology + Supplies the edge blocks; per-origin indices come out as views into them. + values : dict + ``{edge_type: {origin: {property: tensor}}}`` for the keyed types, plus + ``{'chiral': {property: tensor}}`` and ``{'plane': {size: {property: tensor}}}`` + for the two that carry no origin. + + Returns + ------- + dict + ``entries[edge_type][origin][property]`` for bonds, angles and torsions; + ``entries['plane']['4_atoms'][property]``; ``entries['chiral'][property]``. + ``entries['vdw']`` starts empty and is filled when the pair list is built. + """ + entries: Dict[str, Dict] = {} + + for edge_type in ("bond", "angle", "torsion"): + block = topology.edge_block(edge_type) + per_origin = values.get(edge_type, {}) + group: Dict[str, Dict[str, torch.Tensor]] = {} + for origin in block.origins(): + group[origin] = _group(block.origin(origin), per_origin.get(origin, {})) + combined = _all_group(block, per_origin, ALL_MEMBERS[edge_type]) + if combined is not None: + group["all"] = combined + entries[edge_type] = group + + chirals = topology.atoms.chirals + entries["chiral"] = ( + _group(chirals.indices, values.get("chiral", {})) if chirals.n_edges else {} + ) + + entries["plane"] = { + f"{size}_atoms": _group(block.indices, values.get("plane", {}).get(size, {})) + for size, block in sorted(topology.atoms.planes.items()) + } + + entries["vdw"] = {} + return entries + + +def max_period(entries: Dict[str, Dict]) -> int: + """Largest torsion period in the ``all`` group. + + Read once at build time so the torsion target does not pay a device sync for it on + every iteration. + """ + group = entries.get("torsion", {}).get("all") + if not group: + return 1 + periods = group.get("periods") + if periods is None or periods.numel() == 0: + return 1 + return int(periods.max().item()) + + +__all__ = ["assemble_entries", "max_period", "to_tensor", "ALL_MEMBERS"] diff --git a/torchref/restraints/restraints.py b/torchref/topology/restraints.py similarity index 53% rename from torchref/restraints/restraints.py rename to torchref/topology/restraints.py index 7c76df5f..0c49868f 100644 --- a/torchref/restraints/restraints.py +++ b/torchref/topology/restraints.py @@ -1,183 +1,43 @@ -"""Restraints handler for crystallographic model refinement. +"""The restraint layer over a topology, and what it takes to build one. -Provides :class:`RestraintsNew`, which builds geometry restraints (bonds, -angles, torsions, planes, chirals, VDW) using dedicated builder classes. -It is decoupled from :class:`~torchref.model.Model`: it accepts a pdb -DataFrame and callable functions for accessing coordinates and ADPs. +:class:`Restraints` is the orchestrator. Given an atom table it resolves the monomer +dictionaries, builds the :class:`~torchref.topology.topology.Topology`, layers the ideal +values over its edges, derives the non-bonded pair list, and exposes the whole thing as +``restraints[edge_type][origin][property]`` -- three dict lookups into a mapping +assembled once, because the geometry targets read it on every iteration. + +Three kinds of thing live here, and only the first is really connectivity: + +* the topology and the values keyed to its edges, which are constants for the lifetime + of an atom set; +* the non-bonded pair list, which is distance-derived and rebuilt as the model moves, so + it is held apart from the rest; +* the Ramachandran map, a residue-level product of the same build. + +Deliberately decoupled from :class:`~torchref.model.Model`: it takes an atom table plus +callables for coordinates, ADPs and van der Waals radii, so it can be built and tested +without one. """ -from typing import Callable, Optional +from typing import Callable import numpy as np import pandas as pd import torch from torch.nn import Module -from torchref.restraints.builders_fast import ( - AngleRestraintBuilder, - BondRestraintBuilder, - ChiralRestraintBuilder, - InterResidueAngleBuilder, - InterResidueBondBuilder, - InterResiduePlaneBuilder, - InterResidueTorsionBuilder, - PlaneRestraintBuilder, - PreprocessedPDB, - TorsionRestraintBuilder, - find_peptide_link_pairs, -) -from torchref.restraints.modifications import ( - apply_modifications, - link_modifications, - read_mod_definitions, -) -from torchref.restraints.restraints_helper import ( +from torchref.topology.monomer.cif import ( find_cif_file_in_library, read_cif, read_link_definitions, ) from torchref.config import get_float_dtype from torchref.utils.debug_utils import DebugMixin -from torchref.utils.utils import TensorDict from torchref.utils.device_mixin import DeviceMixin -class _RestraintsAccessor: - """ - Provides backward-compatible dict-like access to restraints stored in TensorDict. - - This class mimics the old nested dict interface: - restraints["bond"]["intra"]["indices"] - - While actually accessing the TensorDict with flattened keys: - _tensor_storage["bond_intra_indices"] - """ - - # Types that don't have origin level (assigned directly as dicts) - _FLAT_TYPES = {"vdw", "chiral"} - - def __init__(self, parent: "RestraintsNew"): - self._parent = parent - - def __getitem__(self, rtype: str) -> "_RestraintTypeAccessor": - return _RestraintTypeAccessor(self._parent, rtype) - - def __setitem__(self, rtype: str, value): - """Handle direct assignment for flat types like vdw and chiral.""" - if rtype in self._FLAT_TYPES and isinstance(value, dict): - # Store all tensors with empty origin - self._parent._set_restraint_group(rtype, "", value) - else: - raise TypeError( - f"Cannot assign directly to restraints['{rtype}']. " - f"Use restraints['{rtype}'][origin] = data for nested types." - ) - - def __contains__(self, rtype: str) -> bool: - return len(self._parent._restraint_groups.get(rtype, set())) > 0 or \ - rtype in self._FLAT_TYPES and self._parent._has_restraint(rtype, "") - - def get(self, rtype: str, default=None): - if rtype in self: - return self[rtype] - return default - - def keys(self): - """Return all restraint types that have data.""" - result = [] - for rtype in ["bond", "angle", "torsion", "plane"]: - if len(self._parent._restraint_groups.get(rtype, set())) > 0: - result.append(rtype) - # Check for special types (vdw, chiral) which don't have origins - for rtype in self._FLAT_TYPES: - if self._parent._has_restraint(rtype, ""): - result.append(rtype) - return result - - -class _RestraintTypeAccessor: - """ - Provides access to origins within a restraint type. - - For regular types (bond, angle, torsion, plane): - restraints["bond"]["intra"] -> dict with indices, references, sigmas - - For special types (vdw, chiral), this class acts as the dict itself: - restraints["vdw"]["indices"] -> tensor - restraints["vdw"] = {"indices": ..., "sigmas": ...} - """ - - # Types that don't have origin level (accessed directly as dicts) - _FLAT_TYPES = {"vdw", "chiral"} - - def __init__(self, parent: "RestraintsNew", rtype: str): - self._parent = parent - self._rtype = rtype - - def __getitem__(self, key: str): - if self._rtype in self._FLAT_TYPES: - # For vdw/chiral, key is a property name (indices, sigmas, etc.) - tensor = self._parent._get_restraint_tensor(self._rtype, "", key) - if tensor is None: - raise KeyError(f"No {key} for {self._rtype}") - return tensor - else: - # For bond/angle/torsion/plane, key is an origin name - result = self._parent._get_restraint_group(self._rtype, key) - if result is None: - raise KeyError(f"No restraints for {self._rtype}/{key}") - return result - - def __setitem__(self, key: str, value): - if self._rtype in self._FLAT_TYPES: - # For vdw/chiral, if value is a tensor, store it directly - # If value is a dict, store all tensors - if isinstance(value, torch.Tensor): - self._parent._set_restraint_tensor(self._rtype, "", key, value) - elif isinstance(value, dict): - # This handles: restraints["vdw"] = {"indices": ..., "sigmas": ...} - # But this is called as restraints["vdw"][key] = value, so it won't work - # We need special handling in the parent accessor - pass - else: - # For bond/angle/torsion/plane, key is origin, value is dict - self._parent._set_restraint_group(self._rtype, key, value) - - def __contains__(self, key: str) -> bool: - if self._rtype in self._FLAT_TYPES: - return self._parent._get_restraint_tensor(self._rtype, "", key) is not None - return self._parent._has_restraint(self._rtype, key) - - def get(self, key: str, default=None): - try: - return self[key] - except KeyError: - return default - - def keys(self): - if self._rtype in self._FLAT_TYPES: - # Return property names for flat types - result = [] - for prop in ["indices", "references", "sigmas", "periods", "min_distances", - "symop_indices", "cell_offsets"]: - if self._parent._get_restraint_tensor(self._rtype, "", prop) is not None: - result.append(prop) - return result - return self._parent._get_origins_for_type(self._rtype) - - def items(self): - if self._rtype in self._FLAT_TYPES: - for prop in self.keys(): - yield prop, self._parent._get_restraint_tensor(self._rtype, "", prop) - else: - for origin in self.keys(): - yield origin, self._parent._get_restraint_group(self._rtype, origin) - - def __iter__(self): - return iter(self.keys()) - -class RestraintsNew(DeviceMixin, DebugMixin, Module): +class Restraints(DeviceMixin, DebugMixin, Module): """ Restraints handler for crystallographic model refinement. @@ -211,14 +71,17 @@ class RestraintsNew(DeviceMixin, DebugMixin, Module): Attributes ---------- - restraints : _RestraintsAccessor - Dict-*like* accessor over the flat TensorDict: it emulates - ``restraints["bond"]["intra"]["indices"]`` but is not a plain dict. + restraints : dict + Restraint groups as ``restraints["bond"]["intra"]["indices"]``. A plain nested + dict; the per-origin indices are views into ``topology``'s edge blocks. + topology : Topology + The connectivity the geometry restraints are defined over. cif_dict : dict Parsed CIF restraints keyed by residue type; ``missing_residues`` lists the types that could not be resolved. - h_topo, h_excl_hash - Riding-hydrogen topology and its exclusion hash, populated on demand. + h_topo + Riding-hydrogen map, built only when the model carries no hydrogens of its + own. Empty otherwise; see :mod:`torchref.topology.riding`. link_dict, link_list Link-type definitions from the monomer library, set only when ``pdb`` was provided. @@ -251,11 +114,15 @@ def __init__( self._cell = cell self._spacegroup = spacegroup - # Initialize TensorDict for restraint storage (registered as submodule) - self._tensor_storage = TensorDict() - - # Track which restraint groups exist (for iteration) - self._restraint_groups = {"bond": set(), "angle": set(), "torsion": set(), "plane": set()} + # Connectivity, the values layered over it, and the non-bonded pair list, which + # is rebuilt on displacement and so is kept apart from the rest. + self.topology = None + self._values = {} + self._vdw = {} + # Derived: per-origin views into the topology's edge blocks. Rebuilt by + # _rebuild_entries, which runs at build time and after any device move. + self._entries = {} + self._torsion_max_period = 1 # Empty initialization if pdb is None: @@ -353,68 +220,58 @@ def get_vdw_radii(self) -> torch.Tensor: return self._vdw_radii_fn() # ========================================================================= - # TensorDict Helper Methods for Restraint Storage + # Restraint storage # ========================================================================= - def _make_key(self, rtype: str, origin: str, prop: str) -> str: - """Create flattened key for TensorDict storage.""" - if origin: - return f"{rtype}_{origin}_{prop}" - else: - # For flat types (vdw, chiral) with no origin - return f"{rtype}_{prop}" - def _set_restraint_tensor( - self, rtype: str, origin: str, prop: str, tensor: torch.Tensor - ): - """Store a restraint tensor with flattened key.""" - key = self._make_key(rtype, origin, prop) - self._tensor_storage[key] = tensor - # Track that this origin exists for this restraint type - if rtype in self._restraint_groups: - self._restraint_groups[rtype].add(origin) - - def _get_restraint_tensor( - self, rtype: str, origin: str, prop: str - ) -> Optional[torch.Tensor]: - """Get a restraint tensor by type, origin, and property.""" - key = self._make_key(rtype, origin, prop) - if key in self._tensor_storage: - return self._tensor_storage[key] - return None - - def _has_restraint(self, rtype: str, origin: str) -> bool: - """Check if a restraint group exists.""" - key = self._make_key(rtype, origin, "indices") - return key in self._tensor_storage - - def _set_restraint_group(self, rtype: str, origin: str, data: dict): - """Store all tensors from a restraint data dict.""" - for prop, tensor in data.items(): - if tensor is not None and isinstance(tensor, torch.Tensor): - self._set_restraint_tensor(rtype, origin, prop, tensor) - - def _get_restraint_group(self, rtype: str, origin: str) -> Optional[dict]: - """Get all tensors for a restraint group as a dict.""" - if not self._has_restraint(rtype, origin): - return None - result = {} - # Common properties for different restraint types - for prop in ["indices", "references", "sigmas", "periods", "min_distances", - "is_proline"]: - tensor = self._get_restraint_tensor(rtype, origin, prop) - if tensor is not None: - result[prop] = tensor - return result if result else None - - def _get_origins_for_type(self, rtype: str) -> list: - """Get all origins (e.g., 'intra', 'peptide') for a restraint type.""" - return list(self._restraint_groups.get(rtype, set())) + + + + + @property - def restraints(self) -> "_RestraintsAccessor": - """Nested-dict-*like* accessor over the flat TensorDict (not a real dict).""" - return _RestraintsAccessor(self) + def restraints(self) -> dict: + """Restraint groups as ``[edge type][origin][property]``. + + A plain nested dict of tensors, assembled once at build time. Reading it costs + three dict lookups and no allocation, which matters because the geometry targets + do it on every iteration. Per-origin indices are **views** into the topology's + contiguous edge blocks, so an in-place edit to a block is visible here at once, + and taking a subset costs nothing. + """ + return self._entries + + def _rebuild_entries(self) -> None: + """Re-derive the entry views from the topology and its values. + + Cheap -- a handful of slices -- and idempotent. Runs at the end of a build and + again after any device or dtype move, because moving a tensor rebinds it and + leaves the old views pointing at freed storage. + """ + if self.topology is None: + return + from torchref.topology import assemble_entries, max_period + + self._entries = assemble_entries(self.topology, self._values) + if self._vdw: + self._entries["vdw"] = self._vdw + self._torsion_max_period = max_period(self._entries) + + def _apply(self, fn, recurse: bool = True): + """Drop the derived views before the traversal, re-slice them after. + + ``DeviceMixin``'s ``__dict__`` walk recurses into dicts, so leaving the entries + in place would move each slice on its own and quietly turn every view into an + independent tensor -- doubling the memory and breaking the aliasing the design + rests on. Rebuilding unconditionally rather than in ``_after_device_apply``, + because that hook only fires when the device or dtype actually changed, and a + ``.to()`` onto the current device must not leave the entries empty. + """ + self._entries = {} + result = super()._apply(fn, recurse) + self._rebuild_entries() + return result def _load_cif_dictionaries(self, cif_path): """Load CIF dictionaries from provided paths and monomer library.""" @@ -448,8 +305,16 @@ def _load_cif_dictionaries(self, cif_path): self.missing_residues = [ res for res in self.unique_residues if res not in self.cif_dict ] + from pathlib import Path + from torchref import PATH_TORCHREF_DATA + additional_files = [ - find_cif_file_in_library(res) for res in self.missing_residues + ( + Path(PATH_TORCHREF_DATA) / "monomer_library/h/HOH.cif" + if res == "HOH" + else find_cif_file_in_library(res) + ) + for res in self.missing_residues ] for cif_file in additional_files: @@ -467,94 +332,50 @@ def _load_cif_dictionaries(self, cif_path): res for res in self.unique_residues if res not in self.cif_dict ] - if len(self.missing_residues) > 1: + if len(self.missing_residues) >= 1: if self.verbose > 0: print( f"Warning: The following residues are missing from the CIF dictionary " f"and will have no restraints applied: {self.missing_residues}" ) - def expand_altloc(self, residue): - """ - Expand residue with alternative conformations into separate conformations. - - Yields one DataFrame per altloc (with common atoms included in each). - """ - residue = residue.copy() - residue.loc[residue["altloc"].isin(["", " "]), "altloc"] = " " - - alt_conf = residue["altloc"].unique() - if " " in alt_conf: - residue_no_alt = residue.loc[residue["altloc"] == " "] - for alt in alt_conf: - if alt == " ": - continue - residue_alt = residue.loc[residue["altloc"] == alt] - residue_combined = pd.concat( - [residue_no_alt, residue_alt], ignore_index=True - ) - yield residue_combined - else: - for alt_loc in alt_conf: - residue_alt = residue.loc[residue["altloc"] == alt_loc] - yield residue_alt def _load_rama_surfaces(self, device: torch.device): """Load pre-computed Ramachandran NLL surfaces as a buffer.""" - from torchref.restraints.ramachandran import load_nll_surfaces + from torchref.topology.ramachandran import load_nll_surfaces surfaces = load_nll_surfaces(device) self.register_buffer("_rama_surfaces", surfaces) def build_restraints(self): - """Build every restraint group; each builder handles all residues at once. + """Build the topology, the values over it, and the non-bonded pair list. Builds on CPU and moves the result to the ``xyz()`` device at the end. """ try: target_device = self.xyz().device device = torch.device("cpu") - pdb = self.pdb - # Must precede the intra-residue builders: it decides which residues - # draw their restraints from a link-modified component instead of the - # bare one. - comp_dict, res_keys = self._build_residue_variants() + from torchref.topology import build_topology_with_values - bond_result = BondRestraintBuilder(verbose=self.verbose).build( - pdb, comp_dict, device, residue_keys=res_keys - ) - if bond_result: - self.restraints["bond"]["intra"] = bond_result - - angle_result = AngleRestraintBuilder(verbose=self.verbose).build( - pdb, comp_dict, device, residue_keys=res_keys - ) - if angle_result: - self.restraints["angle"]["intra"] = angle_result - - torsion_result = TorsionRestraintBuilder(verbose=self.verbose).build( - pdb, comp_dict, device, residue_keys=res_keys - ) - if torsion_result: - self.restraints["torsion"]["intra"] = torsion_result - - plane_result = PlaneRestraintBuilder(verbose=self.verbose).build( - pdb, comp_dict, device, residue_keys=res_keys - ) - if plane_result: - for key, data in plane_result.items(): - self.restraints["plane"][key] = data - - chiral_result = ChiralRestraintBuilder(verbose=self.verbose).build( - pdb, comp_dict, device, residue_keys=res_keys + self.topology, self._values, extras = build_topology_with_values( + self.pdb, + self.cif_dict, + link_dict=self.link_dict, + link_list=self.link_list, + links=self.links, + xyz=self.xyz().detach().to(device), + device=device, + verbose=self.verbose, ) - if chiral_result: - self.restraints["chiral"] = chiral_result + self._rebuild_entries() - self._build_peptide_restraints(device) - self._build_disulfide_restraints(device) - self._build_link_restraints(device) + rama = extras.get("ramachandran") + if rama is not None: + self.register_buffer("_rama_phi_indices", rama["phi_indices"]) + self.register_buffer("_rama_psi_indices", rama["psi_indices"]) + self.register_buffer("_rama_surface_type", rama["surface_type"]) + self._load_rama_surfaces(device) # cutoff sits ~1 Å beyond the largest heavy-atom VDW sum (~3.6 Å) plus # expected drift, so a displacement-triggered rebuild stays inside the @@ -563,484 +384,18 @@ def build_restraints(self): cutoff=6.0, sigma=0.05, inter_residue_only=False, use_spatial_hash=True ) - # Register the concatenated 'all' buffers here, not lazily in forward: - # register_buffer() during a forward pass breaks CUDA-graph capture, and - # only registered buffers are moved by model.to(device). - self.cat_dict() - if target_device.type != "cpu": self.to(target_device) except Exception as e: - self.debug_on_error(e, context="RestraintsNew.build_restraints") + self.debug_on_error(e, context="Restraints.build_restraints") raise - def _build_residue_variants(self): - """Point peptide-linked residues at link-modified copies of their component. - The monomer library defines each amino acid free: ``ALA`` carries ``OXT`` - and a protonated ``N``, with carboxylate and ammonium geometry. Forming a - peptide bond applies the modifications the ``chem_link`` table names -- - ``DEL-OXT`` to the residue donating its C, ``DEL-HN1`` (``DEL-HNP`` for - proline) to the residue donating its N -- which delete the restraints the - link makes meaningless and overwrite the targets that change, notably - ``CA-C-O`` and ``CA-N-H``. Without them the intra-residue restraints fight - the link's own: around a peptide carbonyl carbon ``CA-C-O`` + ``CA-C-N`` + - ``O-C-N`` only sums to 360 degrees once ``DEL-OXT`` has been applied. - Chain termini are deliberately left unmodified -- a real C-terminus keeps - its ``OXT`` and carboxylate geometry, a real N-terminus its ammonium. - Returns - ------- - comp_dict : dict - :attr:`cif_dict` plus one entry per ``(residue type, modification - set)`` in use, keyed ``'ALA:DEL-HN1+DEL-OXT'``. :attr:`cif_dict` - itself is left keyed by residue type alone. - residue_keys : dict - ``{(chain_id, resseq): comp_dict key}`` for the modified residues, as - :meth:`~torchref.restraints.builders_fast.RestraintBuilder.build` - takes it. Residues absent from it use their residue name. - """ - comp_dict = dict(self.cif_dict) - residue_keys = {} - if not self.cif_dict or getattr(self, "link_list", None) is None: - return comp_dict, residue_keys - - modifications = link_modifications(self.link_list) - if "TRANS" not in modifications: - return comp_dict, residue_keys - trans_mods = modifications["TRANS"] - proline_mods = modifications.get("PTRANS", trans_mods) - - polymer = self.pdb[self.pdb["ATOM"] == "ATOM"] - if len(polymer) == 0: - return comp_dict, residue_keys - pp_pdb = PreprocessedPDB(polymer) - - mods_by_residue = {} - for res_i, res_next in find_peptide_link_pairs(pp_pdb): - donor_mod, acceptor_mod = ( - proline_mods - if pp_pdb.residue_resnames[res_next] == "PRO" - else trans_mods - ) - for res_idx, mod_id in ((res_i, donor_mod), (res_next, acceptor_mod)): - if mod_id is None: - continue - key = ( - str(pp_pdb.residue_chain_ids[res_idx]), - int(pp_pdb.residue_resseqs[res_idx]), - ) - mods_by_residue.setdefault(key, set()).add(mod_id) - if not mods_by_residue: - return comp_dict, residue_keys - mod_dict = read_mod_definitions() - for res_idx in range(pp_pdb.n_residues): - key = ( - str(pp_pdb.residue_chain_ids[res_idx]), - int(pp_pdb.residue_resseqs[res_idx]), - ) - mods = mods_by_residue.get(key) - resname = pp_pdb.residue_resnames[res_idx] - if not mods or resname not in comp_dict: - continue - mods = sorted(mods) - variant = f"{resname}:{'+'.join(mods)}" - if variant not in comp_dict: - comp_dict[variant] = apply_modifications( - comp_dict[resname], mods, mod_dict - ) - residue_keys[key] = variant - - if self.verbose > 1: - print( - f"Applied link modifications to {len(residue_keys)} residues " - f"({len(comp_dict) - len(self.cif_dict)} modified components)" - ) - return comp_dict, residue_keys - - def _build_peptide_restraints(self, device: torch.device): - """Build peptide bond/angle/torsion/plane restraints. - - TRANS/CIS links for standard peptide bonds; PTRANS/PCIS for bonds to - proline, which add the C(i-1)-N-CD angle and proline-specific targets. - """ - if "TRANS" not in self.link_dict: - if self.verbose > 0: - print( - "Warning: TRANS link not found in link dictionary, skipping peptide bonds" - ) - return - - trans_link = self.link_dict["TRANS"] - ptrans_link = self.link_dict.get("PTRANS") - pdb = self.pdb - - # Build peptide bonds using fast builder - bond_result = InterResidueBondBuilder(verbose=self.verbose).build( - pdb, trans_link, device, filter_atom_type="ATOM" - ) - if bond_result: - self.restraints["bond"]["peptide"] = bond_result - if self.verbose > 0: - print( - f"Built {bond_result['indices'].shape[0]} peptide bond restraints" - ) - - # Build peptide angles. - # If PTRANS is available, use it for proline pairs (excludes PRO - # from TRANS to avoid duplicate/conflicting restraints) and TRANS - # for non-proline pairs. Otherwise fall back to TRANS for all. - angle_builder = InterResidueAngleBuilder(verbose=self.verbose) - if ptrans_link is not None: - # Non-proline pairs: TRANS angles - angle_result = angle_builder.build( - pdb, trans_link, device, filter_atom_type="ATOM", - exclude_next_resname="PRO", - ) - # Proline pairs: PTRANS angles (includes C-N-CD) - pro_angle_result = angle_builder.build( - pdb, ptrans_link, device, filter_atom_type="ATOM", - next_resname_filter="PRO", - ) - # Merge results - if angle_result and pro_angle_result: - angle_result = { - "indices": torch.cat([angle_result["indices"], pro_angle_result["indices"]]), - "references": torch.cat([angle_result["references"], pro_angle_result["references"]]), - "sigmas": torch.cat([angle_result["sigmas"], pro_angle_result["sigmas"]]), - } - elif pro_angle_result: - angle_result = pro_angle_result - else: - angle_result = angle_builder.build( - pdb, trans_link, device, filter_atom_type="ATOM" - ) - - if angle_result: - self.restraints["angle"]["peptide"] = angle_result - if self.verbose > 0: - print( - f"Built {angle_result['indices'].shape[0]} peptide angle restraints" - ) - - # Build backbone torsions (phi, psi, omega) - torsion_result = InterResidueTorsionBuilder(verbose=self.verbose).build( - pdb, trans_link, device, filter_atom_type="ATOM" - ) - if torsion_result: - if "phi" in torsion_result: - self.restraints["torsion"]["phi"] = torsion_result["phi"] - if "psi" in torsion_result: - self.restraints["torsion"]["psi"] = torsion_result["psi"] - if "omega" in torsion_result: - self.restraints["torsion"]["omega"] = torsion_result["omega"] - if "ramachandran" in torsion_result: - rama = torsion_result["ramachandran"] - self.register_buffer("_rama_phi_indices", rama["phi_indices"]) - self.register_buffer("_rama_psi_indices", rama["psi_indices"]) - self.register_buffer("_rama_surface_type", rama["surface_type"]) - self._load_rama_surfaces(device) - - # Build peptide planes - plane_result = InterResiduePlaneBuilder(verbose=self.verbose).build( - pdb, trans_link, device, filter_atom_type="ATOM" - ) - if plane_result: - n_planes = 0 - for key, data in plane_result.items(): - n_planes += data["indices"].shape[0] - if self._has_restraint("plane", key): - # Append to existing planes of same atom count - existing = self.restraints["plane"][key] - self.restraints["plane"][key] = { - "indices": torch.cat( - [existing["indices"], data["indices"]], dim=0 - ), - "sigmas": torch.cat( - [existing["sigmas"], data["sigmas"]], dim=0 - ), - } - else: - self.restraints["plane"][key] = data - if self.verbose > 0: - print(f"Built {n_planes} peptide plane restraints") - - def _build_disulfide_restraints(self, device: torch.device): - """Build disulfide bond restraints.""" - if "disulf" not in self.link_dict: - if self.verbose > 1: - print( - "Warning: disulf link not found in link dictionary, skipping disulfide bonds" - ) - return - - disulf_link = self.link_dict["disulf"] - disulf_bonds = disulf_link.get("bonds") - disulf_angles = disulf_link.get("angles") - disulf_torsions = disulf_link.get("torsions") - - if disulf_bonds is None: - return - - # Get SG-SG bond parameters - sg_sg_bond = disulf_bonds[ - (disulf_bonds["atom1"] == "SG") & (disulf_bonds["atom2"] == "SG") - ] - - if len(sg_sg_bond) == 0: - return - - bond_length = float(sg_sg_bond["value"].values[0]) - bond_sigma = float(sg_sg_bond["sigma"].values[0]) - - # Find all SG atoms - pdb = self.pdb - sg_atoms = pdb[(pdb["name"] == "SG") & (pdb["ATOM"] == "ATOM")] - - if len(sg_atoms) == 0: - return - - # Get coordinates and find close pairs - xyz = self.xyz() - sg_indices = sg_atoms["index"].values - sg_coords = xyz[sg_indices] - sg_residues = ( - sg_atoms["chainid"].astype(str) + "_" + sg_atoms["resseq"].astype(str) - ).values - - distances = torch.cdist(sg_coords, sg_coords) - threshold = 2.5 - close_pairs = torch.where((distances < threshold) & (distances > 0.1)) - - valid_pairs = [] - for i, j in zip(close_pairs[0].cpu().numpy(), close_pairs[1].cpu().numpy()): - if i < j and sg_residues[i] != sg_residues[j]: - valid_pairs.append((i, j)) - - if len(valid_pairs) == 0: - return - - # Create builders - bond_builder = InterResidueBondBuilder(verbose=self.verbose) - angle_builder = InterResidueAngleBuilder(verbose=self.verbose) - torsion_builder = InterResidueTorsionBuilder(verbose=self.verbose) - - # Process each disulfide bond - for i_local, j_local in valid_pairs: - sg1_idx = int(sg_indices[i_local]) - sg2_idx = int(sg_indices[j_local]) - - # Add bond - bond_builder.process_disulfide_bond( - sg1_idx, sg2_idx, bond_length, bond_sigma - ) - - # Get residues for angle/torsion restraints - residue1 = pdb[pdb["index"] == sg1_idx].iloc[0] - residue2 = pdb[pdb["index"] == sg2_idx].iloc[0] - - res1_atoms = pdb[ - (pdb["chainid"] == residue1["chainid"]) - & (pdb["resseq"] == residue1["resseq"]) - ] - res2_atoms = pdb[ - (pdb["chainid"] == residue2["chainid"]) - & (pdb["resseq"] == residue2["resseq"]) - ] - - if disulf_angles is not None: - angle_builder.process_disulfide_angles( - res1_atoms, res2_atoms, disulf_angles - ) - - if disulf_torsions is not None: - torsion_builder.process_disulfide_torsions( - res1_atoms, res2_atoms, disulf_torsions - ) - - # Finalize - bond_result = bond_builder.finalize(device) - if bond_result: - self.restraints["bond"]["disulfide"] = bond_result - if self.verbose > 0: - print( - f"Built {bond_result['indices'].shape[0]} disulfide bond restraints" - ) - - angle_result = angle_builder.finalize(device) - if angle_result: - self.restraints["angle"]["disulfide"] = angle_result - if self.verbose > 0: - print( - f"Built {angle_result['indices'].shape[0]} disulfide angle restraints" - ) - - torsion_result = torsion_builder.finalize_disulfide(device) - if torsion_result: - self.restraints["torsion"]["disulfide"] = torsion_result - if self.verbose > 0: - print( - f"Built {torsion_result['indices'].shape[0]} disulfide torsion restraints" - ) - - def _build_link_restraints(self, device: torch.device): - """Build one bond restraint per accepted PDB LINK record. - - Target is the record's ``length`` at sigma=0.02 Å, falling back to 1.5 Å - when blank. LINKs duplicating an auto-detected CYS SG-SG disulfide are - skipped (that builder already added bond + angles + torsions); symmetry-mate - links were dropped earlier in ``extract_link_records``. Each bond joins the - VDW exclusion set via ``_build_exclusion_set``, so the non-bonded term does - not push linked atoms apart. - """ - if self.links is None or len(self.links) == 0: - return - - pdb = self.pdb - - # Already-bonded SG-SG pairs from auto-disulfide detection. - disulf = self.restraints.get("bond", {}).get("disulfide") - existing_disulf_pairs = set() - if disulf is not None and "indices" in disulf: - for i, j in disulf["indices"].cpu().numpy(): - existing_disulf_pairs.add((int(min(i, j)), int(max(i, j)))) - - bond_builder = InterResidueBondBuilder(verbose=self.verbose) - n_skipped_unresolved = 0 - n_skipped_dedup = 0 - - for _, link in self.links.iterrows(): - idx1 = self._lookup_link_atom( - pdb, - chainid=link["chainid1"], - resseq=int(link["resseq1"]), - icode=link["icode1"], - resname=link["resname1"], - name=link["name1"], - altloc=link["altloc1"], - ) - idx2 = self._lookup_link_atom( - pdb, - chainid=link["chainid2"], - resseq=int(link["resseq2"]), - icode=link["icode2"], - resname=link["resname2"], - name=link["name2"], - altloc=link["altloc2"], - ) - - if idx1 is None or idx2 is None: - n_skipped_unresolved += 1 - if self.verbose > 1: - print( - f"Warning: LINK atom not found " - f"({link['chainid1']}/{link['resname1']}{link['resseq1']}/" - f"{link['name1']} -- " - f"{link['chainid2']}/{link['resname2']}{link['resseq2']}/" - f"{link['name2']}); skipping." - ) - continue - - if idx1 == idx2: - n_skipped_unresolved += 1 - continue - - pair = (min(idx1, idx2), max(idx1, idx2)) - if pair in existing_disulf_pairs: - n_skipped_dedup += 1 - continue - - length = link["length"] - if not (isinstance(length, (int, float)) and length == length and length > 0): - length = 1.5 - bond_builder.process_disulfide_bond(idx1, idx2, float(length), 0.02) - - bond_result = bond_builder.finalize(device) - if bond_result: - self.restraints["bond"]["link"] = bond_result - if self.verbose > 0: - print( - f"Built {bond_result['indices'].shape[0]} LINK bond restraints" - + ( - f" (skipped {n_skipped_dedup} disulfide-dup," - f" {n_skipped_unresolved} unresolved)" - if (n_skipped_dedup or n_skipped_unresolved) - else "" - ) - ) - - @staticmethod - def _lookup_link_atom( - pdb: pd.DataFrame, - chainid: str, - resseq: int, - icode: str, - resname: str, - name: str, - altloc: str, - ): - """Resolve a LINK atom record to a row index in the model pdb, or None. - - Matches (chainid, resseq, icode, name), resname as tie-breaker; altloc - preference is requested, then blank, then 'A', then any. - """ - sel = pdb[ - (pdb["chainid"].astype(str) == str(chainid)) - & (pdb["resseq"].astype(int) == int(resseq)) - & (pdb["icode"].astype(str) == str(icode)) - & (pdb["name"].astype(str).str.strip() == str(name).strip()) - ] - if len(sel) == 0: - return None - if resname: - tied = sel[sel["resname"].astype(str).str.strip() == str(resname).strip()] - if len(tied) > 0: - sel = tied - - if altloc: - for cand in (altloc, ""): - hit = sel[sel["altloc"].astype(str) == cand] - if len(hit) > 0: - return int(hit.iloc[0]["index"]) - for cand in ("", "A"): - hit = sel[sel["altloc"].astype(str) == cand] - if len(hit) > 0: - return int(hit.iloc[0]["index"]) - return int(sel.iloc[0]["index"]) - - def _build_exclusion_set(self): - """Build set of atom pairs to exclude from VDW calculations.""" - exclusions = set() - - # 1-2: Direct bonds - for origin in self.restraints.get("bond", {}).keys(): - indices = self.restraints["bond"][origin].get("indices") - if indices is not None and len(indices) > 0: - idx_np = indices.cpu().numpy() - for i1, i2 in idx_np: - exclusions.add((int(min(i1, i2)), int(max(i1, i2)))) - - # 1-3: Angles - for origin in self.restraints.get("angle", {}).keys(): - indices = self.restraints["angle"][origin].get("indices") - if indices is not None and len(indices) > 0: - idx_np = indices.cpu().numpy() - for i1, i2, i3 in idx_np: - exclusions.add((int(min(i1, i3)), int(max(i1, i3)))) - - # 1-4: Torsions - for origin in self.restraints.get("torsion", {}).keys(): - indices = self.restraints["torsion"][origin].get("indices") - if indices is not None and len(indices) > 0: - idx_np = indices.cpu().numpy() - for i1, i2, i3, i4 in idx_np: - exclusions.add((int(min(i1, i4)), int(max(i1, i4)))) - - return exclusions def _find_nearby_pairs_spatial_hash(self, xyz, cutoff=6.0): """Atom pairs within ``cutoff`` of each other, as (M, 2) rows with i < j. @@ -1053,7 +408,7 @@ def _find_nearby_pairs_spatial_hash(self, xyz, cutoff=6.0): n_atoms = xyz.shape[0] if n_atoms == 0: - return torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) + return torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) # dtype-ok: empty atom-pair index tensor; int64 index required # Work on CPU to avoid per-iteration GPU kernel launch overhead coords = xyz.detach().cpu() @@ -1078,12 +433,12 @@ def _find_nearby_pairs_spatial_hash(self, xyz, cutoff=6.0): sorted_flat, return_counts=True ) n_unique = len(unique_cells) - starts = torch.zeros(n_unique + 1, dtype=torch.long) + starts = torch.zeros(n_unique + 1, dtype=torch.long) # dtype-ok: grid-cell CSR start offsets; int64 index required starts[1:] = counts.cumsum(0) # Lookup: flat_cell -> index in unique_cells (-1 if empty) n_grid = gx * gyz - cell_lookup = torch.full((n_grid,), -1, dtype=torch.long) + cell_lookup = torch.full((n_grid,), -1, dtype=torch.long) # dtype-ok: cell lookup table (-1 sentinel); int64 index required cell_lookup[unique_cells] = torch.arange(n_unique) # 14 unique neighbour offsets: self (0,0,0) + 13 forward neighbours. @@ -1169,9 +524,9 @@ def _find_nearby_pairs_spatial_hash(self, xyz, cutoff=6.0): if pair_chunks: all_pairs = np.concatenate(pair_chunks, axis=0) - return torch.from_numpy(all_pairs).to(dtype=torch.long, device=device) + return torch.from_numpy(all_pairs).to(dtype=torch.long, device=device) # dtype-ok: atom-pair index array from numpy; int64 index required else: - return torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) + return torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) # dtype-ok: empty atom-pair index tensor; int64 index required def _expand_with_symmetry_mates(self, xyz, cutoff): """Append symmetry-mate positions to ASU ``xyz`` for neighbour search. @@ -1296,10 +651,6 @@ def h_topo(self): """Access riding hydrogen topology (None if not built).""" return getattr(self, "_h_topo", None) - @property - def h_excl_hash(self): - """Access H-specific exclusion hash tensor (None if not built).""" - return getattr(self, "_h_excl_hash", None) def _build_h_exclusion_hash(self, h_topo, device): """Sorted 1-D hash tensor of H-specific 1-2 and 1-3 exclusions. @@ -1308,7 +659,7 @@ def _build_h_exclusion_hash(self, h_topo, device): ``torch.searchsorted`` lookup. """ if h_topo is None or h_topo.n_hydrogens == 0: - return torch.tensor([], dtype=torch.long, device=device) + return torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty index tensor; int64 index required n_heavy = len(self.pdb) n_h = h_topo.n_hydrogens @@ -1333,13 +684,13 @@ def _build_h_exclusion_hash(self, h_topo, device): exclusions.add((min(h_combined, nb), max(h_combined, nb))) if not exclusions: - return torch.tensor([], dtype=torch.long, device=device) + return torch.tensor([], dtype=torch.long, device=device) # dtype-ok: empty index tensor; int64 index required arr = np.array(list(exclusions), dtype=np.int64) max_idx = max(n_heavy + n_h, int(arr.max()) + 1) hashes = arr[:, 0] * max_idx + arr[:, 1] hashes.sort() - return torch.tensor(hashes, dtype=torch.long, device=device) + return torch.tensor(hashes, dtype=torch.long, device=device) # dtype-ok: grid-cell hash values used as keys/index; int64 required def _build_vdw_restraints( self, cutoff=6.0, sigma=0.2, inter_residue_only=True, use_spatial_hash=True @@ -1405,17 +756,15 @@ def vdw_radii_cpu(): else: cell_cpu = None if self._spacegroup is not None: - from torchref.symmetry.spacegroup import SpaceGroup - sg_cpu = SpaceGroup(self._spacegroup, device=cpu, - dtype=self._spacegroup._dtype) + sg_cpu = self._spacegroup.copy().to(cpu) else: sg_cpu = None if has_symmetry: - from torchref.restraints.neighbor_search import build_vdw_restraints_gpu + from torchref.topology.nonbonded import build_vdw_restraints_gpu - exclusions = self._build_exclusion_set() - self.restraints["vdw"] = build_vdw_restraints_gpu( + exclusions = self.topology.atoms.exclusions_from_restraint_edges() + self._vdw = build_vdw_restraints_gpu( xyz_fn=xyz_cpu, vdw_radii_fn=vdw_radii_cpu, cell=cell_cpu, @@ -1434,17 +783,33 @@ def vdw_radii_cpu(): use_spatial_hash=use_spatial_hash, ) - # Build riding hydrogen topology and precompute candidate pairs - from torchref.restraints.hydrogen_topology import ( - build_hydrogen_topology, + # Publish the new pair list before anything reads it back below. Unlike the + # geometry edges it is not derived from the topology, so it is held separately + # and re-inserted here and by _rebuild_entries. + self._entries["vdw"] = self._vdw + + # Riding hydrogens stand in for the sterics of hydrogens the model does not + # carry. Once it carries them they are ordinary atoms in the pair list above, and + # placing riding ones as well would put phantom hydrogens in the structure that + # push real atoms around. The two also disagree about how many belong on a + # parent -- the riding builder counts bonded neighbours by distance, the + # generator reads them off the bond graph -- so the leftovers are not even the + # hydrogens the generator declined to add. + from torchref.topology.riding import ( + HydrogenTopology, build_h_candidate_pairs, + build_hydrogen_topology, ) - self._h_topo = build_hydrogen_topology( - pdb=self.pdb, - device=cpu, - verbose=self.verbose, - ) + elements = self.pdb["element"].astype(str).str.strip().values + if (elements == "H").any(): + self._h_topo = HydrogenTopology(device=cpu) + else: + self._h_topo = build_hydrogen_topology( + pdb=self.pdb, + device=cpu, + verbose=self.verbose, + ) self._h_excl_hash = self._build_h_exclusion_hash(self._h_topo, cpu) # Precompute H candidate pairs from heavy-atom VDW pair list @@ -1497,7 +862,7 @@ def _build_vdw_restraints_legacy( ): """Legacy VDW restraint builder (no symmetry or CPU fallback).""" - exclusions = self._build_exclusion_set() + exclusions = self.topology.atoms.exclusions_from_restraint_edges() vdw_radii = self.get_vdw_radii() xyz = self.xyz() device = xyz.device @@ -1528,21 +893,21 @@ def _build_vdw_restraints_legacy( if dist_sq < cutoff_sq: pairs_list.append([i, j]) nearby_pairs = ( - torch.tensor(pairs_list, dtype=torch.long, device=device) + torch.tensor(pairs_list, dtype=torch.long, device=device) # dtype-ok: atom-pair index tensor; int64 index required if pairs_list - else torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) + else torch.tensor([], dtype=torch.long, device=device).reshape(0, 2) # dtype-ok: empty atom-pair index tensor; int64 index required ) empty_result = { - "indices": torch.tensor([], dtype=torch.long, device=device).reshape(0, 2), + "indices": torch.tensor([], dtype=torch.long, device=device).reshape(0, 2), # dtype-ok: empty atom-pair index tensor; int64 index required "min_distances": torch.tensor([], dtype=get_float_dtype(), device=device), "sigmas": torch.tensor([], dtype=get_float_dtype(), device=device), - "symop_indices": torch.tensor([], dtype=torch.long, device=device), - "cell_offsets": torch.tensor([], dtype=torch.long, device=device).reshape(0, 3), + "symop_indices": torch.tensor([], dtype=torch.long, device=device), # dtype-ok: empty symop index tensor; int64 index required + "cell_offsets": torch.tensor([], dtype=torch.long, device=device).reshape(0, 3), # dtype-ok: empty cell-offset index tensor; int64 index required } if len(nearby_pairs) == 0: - self.restraints["vdw"] = empty_result + self._vdw = empty_result return pairs_np = nearby_pairs.cpu().numpy() @@ -1660,7 +1025,7 @@ def _build_vdw_restraints_legacy( final_offsets = final_offsets[keep_mask] if len(final_i1) == 0: - self.restraints["vdw"] = empty_result + self._vdw = empty_result return # Compute min distances using VDW radii of ASU source atoms. @@ -1669,8 +1034,8 @@ def _build_vdw_restraints_legacy( # Store results final_pairs = np.stack([final_i1, final_i2], axis=1) - self.restraints["vdw"] = { - "indices": torch.tensor(final_pairs, dtype=torch.long, device=device), + self._vdw = { + "indices": torch.tensor(final_pairs, dtype=torch.long, device=device), # dtype-ok: final atom-pair index tensor; int64 index required "min_distances": torch.tensor( min_distances, dtype=get_float_dtype(), device=device ), @@ -1678,10 +1043,10 @@ def _build_vdw_restraints_legacy( (len(final_pairs),), sigma, dtype=get_float_dtype(), device=device ), "symop_indices": torch.tensor( - final_symop, dtype=torch.long, device=device + final_symop, dtype=torch.long, device=device # dtype-ok: symop index tensor; int64 index required ), "cell_offsets": torch.tensor( - final_offsets, dtype=torch.long, device=device + final_offsets, dtype=torch.long, device=device # dtype-ok: cell-offset index tensor; int64 index required ), } @@ -1694,8 +1059,8 @@ def _build_vdw_restraints_legacy( msg += f", {n_sym_count} symmetry contacts" print(msg) - # Device movement is handled automatically by TensorDict (registered as _tensor_storage) - # through PyTorch's Module.to(), cuda(), and cpu() methods + # Device movement goes through DeviceMixin: the topology and the value tensors are + # walked and moved, and _apply re-slices the derived entry views afterwards. def summary(self): """Print a detailed summary of all restraints.""" @@ -1793,50 +1158,11 @@ def get_count(rtype, origin): n_bonds_peptide = get_count("bond", "peptide") return ( - f"RestraintsNew(bonds={n_bonds}, angles={n_angles}, " + f"Restraints(bonds={n_bonds}, angles={n_angles}, " f"torsions={n_torsions}, peptide_bonds={n_bonds_peptide})" ) - def _get_all_indices(self, restraint_type, keys_to_merge=None): - """Concatenate the indices of one restraint type across origins, or None. - - ``keys_to_merge`` restricts to those origins; None means all. Rows follow - origin iteration order, so pair this only with - :meth:`_get_all_property` calls made with the same ``keys_to_merge``. - """ - indices_list = [] - for origin, data in self.restraints.get(restraint_type, {}).items(): - indices = data.get("indices") - if indices is not None: - if keys_to_merge is None: - indices_list.append(indices) - elif origin in keys_to_merge: - indices_list.append(indices) - - if not indices_list: - return None - - return torch.cat(indices_list, dim=0) - - def _get_all_property(self, restraint_type, property_name, keys_to_merge=None): - """Concatenate one property ('references'/'sigmas'/'periods') across origins. - - Returns None if no origin carries it. Row order matches - :meth:`_get_all_indices` for the same ``keys_to_merge``. - """ - values_list = [] - for origin, data in self.restraints.get(restraint_type, {}).items(): - values = data.get(property_name) - if values is not None: - if keys_to_merge is None: - values_list.append(values) - elif origin in keys_to_merge: - values_list.append(values) - - if not values_list: - return None - return torch.cat(values_list, dim=0) def bond_lengths(self, idx, xyz: torch.Tensor = None): """ @@ -1863,17 +1189,37 @@ def bond_lengths(self, idx, xyz: torch.Tensor = None): return torch.linalg.norm(pos2 - pos1, dim=-1) def copy(self): - """ - Create a deep copy of the Restraints object. + """An independent copy, sharing no state with this one. + + The entry views are re-sliced afterwards rather than left as deep-copied + tensors: ``deepcopy`` duplicates a view and the block it points into as two + unrelated tensors, so the copy would still hold the right values but would no + longer alias, and an in-place edit to one would stop being visible through the + other. Returns ------- Restraints - A deep copy of this Restraints instance. """ import copy - return copy.deepcopy(self) + # The coordinate and ADP accessors are borrowed from the model, not owned: + # duplicating them would hand the copy a third, orphaned parameter set (and + # deep-copying a wrapper with a cached forward fails on its graph tensor). + # They are carried across by reference; the owning model re-points them. + borrowed = ("_xyz_fn", "_adp_fn", "_vdw_radii_fn") + saved = {name: getattr(self, name, None) for name in borrowed} + for name in borrowed: + setattr(self, name, None) + try: + duplicate = copy.deepcopy(self) + finally: + for name, value in saved.items(): + setattr(self, name, value) + for name, value in saved.items(): + setattr(duplicate, name, value) + duplicate._rebuild_entries() + return duplicate def bond_deviations(self, xyz: torch.Tensor = None): """ @@ -2025,48 +1371,20 @@ def nll_angles(self, xyz: torch.Tensor = None): return gaussian_nll(deviations, sigmas) def cat_dict(self): - """ - Concatenate restraint origins into combined 'all' keys. + """Ensure the combined ``all`` groups are present. Idempotent. + + They are assembled with everything else at build time, so this normally has + nothing to do; it exists because the geometry targets guard their reads with + ``if "all" not in ...`` and call it when the guard trips. - Creates restraints['bond']['all'], restraints['angle']['all'], and - restraints['torsion']['all']. Bond and angle 'all' include every - origin; torsion 'all' includes only the 'intra' and 'disulfide' - origins (phi/psi have no reference values or sigmas, and omega is - handled by a dedicated OmegaTarget). + The previous implementation concatenated the origins on each call, and because + writing ``restraints['bond']['all']`` also registered ``'all'`` as an origin, a + second call folded the combined group into itself and doubled every restraint. + Deriving the group from the topology instead makes that impossible: ``all`` is a + span of the edge block, never an origin in its own right. """ - self.restraints["bond"]["all"] = { - "indices": self._get_all_indices("bond"), - "references": self._get_all_property("bond", "references"), - "sigmas": self._get_all_property("bond", "sigmas"), - } - self.restraints["angle"]["all"] = { - "indices": self._get_all_indices("angle"), - "references": self._get_all_property("angle", "references"), - "sigmas": self._get_all_property("angle", "sigmas"), - } - # Note: phi/psi origins are excluded because they have no reference - # values or sigmas (conformationally free). Omega is excluded here - # because it is handled by a dedicated OmegaTarget that uses a - # cis/trans von Mises mixture model. - _torsion_origins = ["intra", "disulfide"] - self.restraints["torsion"]["all"] = { - "indices": self._get_all_indices("torsion", _torsion_origins), - "references": self._get_all_property( - "torsion", "references", _torsion_origins - ), - "sigmas": self._get_all_property( - "torsion", "sigmas", _torsion_origins - ), - "periods": self._get_all_property( - "torsion", "periods", _torsion_origins - ), - } - # Cache max period to avoid .item() GPU sync every iteration - periods = self.restraints["torsion"]["all"]["periods"] - if periods is not None and periods.numel() > 0: - self._torsion_max_period = int(periods.max().item()) - else: - self._torsion_max_period = 1 + if self.topology is not None and "all" not in self._entries.get("bond", {}): + self._rebuild_entries() def torsions(self, idx, xyz: torch.Tensor = None): """ @@ -2187,49 +1505,6 @@ def _wrap_torsion_periodicity(self, diff_rad, periods): # All periods are 0 or 1, simple wrapping return torch.remainder(diff_rad + torch.pi, 2.0 * torch.pi) - torch.pi - def torsion_deviations(self, xyz: torch.Tensor = None, wrapped=True): - """ - Compute deviations between calculated and expected torsion angles. - - Parameters - ---------- - xyz : torch.Tensor, optional - Coordinates tensor. If None, uses the stored xyz_fn callable. - wrapped : bool, default True - If True, wrap deviations accounting for periodicity. - If False, return raw deviations (calculated - expected). - - Returns - ------- - torch.Tensor - Tensor of shape (n_torsions,) with deviations in degrees. - For wrapped=True, deviations are in range appropriate for the period. - - Notes - ----- - Expected values from the CIF library are discrete (typically -60°, 0°, - 60°, 90°, 180°) while calculated values from the structure are - continuous. Use wrapped=True for meaningful comparison and - visualization. - """ - if "all" not in self.restraints["torsion"]: - self.cat_dict() - - idx = self.restraints["torsion"]["all"]["indices"] - expected = self.restraints["torsion"]["all"]["references"] - periods = self.restraints["torsion"]["all"]["periods"] - calculated = self.torsions(idx, xyz) - - if not wrapped: - # Simple difference - return calculated - expected - else: - # Use the helper function for periodicity handling - diff_rad = (calculated - expected) * torch.pi / 180.0 - diff_wrapped_rad = self._wrap_torsion_periodicity(diff_rad, periods) - - # Convert back to degrees - return torch.rad2deg(diff_wrapped_rad) def torsion_deviations_with_sigmas(self, xyz: torch.Tensor = None): """ @@ -2263,143 +1538,8 @@ def torsion_deviations_with_sigmas(self, xyz: torch.Tensor = None): return deviations_rad, sigmas_deg - def nll_torsions(self, xyz: torch.Tensor = None): - """ - Compute negative log-likelihood for torsion angle restraints. - - von Mises: NLL = -κ·cos(θ-μ) + log(I₀(κ)) + log(2π), with κ = 1/σ². This is - the true NLL, so exp(-NLL) is a probability density. Deviations are folded - by the restraint period first (see :meth:`_wrap_torsion_periodicity`). - Parameters - ---------- - xyz : torch.Tensor, optional - Coordinates tensor. If None, uses the stored xyz_fn callable. - - Returns - ------- - torch.Tensor - Tensor of shape (n_torsions,) with negative log-likelihood values. - """ - from torchref.refinement.targets import von_mises_nll - - deviations_rad, sigmas_deg = self.torsion_deviations_with_sigmas(xyz) - return von_mises_nll(deviations_rad, sigmas_deg) - - def nll_planes(self, xyz: torch.Tensor = None): - """ - Compute negative log-likelihood for plane restraints. - - For each plane, computes the RMSD of atom deviations from the best-fit plane. - Uses Gaussian NLL: NLL = 0.5 * (deviation / σ)² + log(σ) + 0.5 * log(2π) - - Parameters - ---------- - xyz : torch.Tensor, optional - Coordinates tensor. If None, uses the stored xyz_fn callable. - - Returns - ------- - torch.Tensor - Tensor of shape (n_planes,) with negative log-likelihood values. - """ - from torchref.refinement.targets import gaussian_nll - - xyz = self.xyz(xyz) - device = xyz.device - - all_nlls = [] - - if "plane" in self.restraints: - for key, plane_data in self.restraints["plane"].items(): - indices = plane_data.get("indices") - sigmas = plane_data.get("sigmas") - - if indices is None or len(indices) == 0: - continue - - # indices shape: (n_planes, n_atoms_per_plane) - # sigmas shape: (n_planes, n_atoms_per_plane) - n_planes, n_atoms = indices.shape - - for i in range(n_planes): - plane_indices = indices[i] - plane_sigmas = sigmas[i] - - # Get positions of atoms in this plane - positions = xyz[plane_indices] # (n_atoms, 3) - - # Compute centroid - centroid = positions.mean(dim=0) - centered = positions - centroid - - # SVD to find best-fit plane normal - # The plane normal is the singular vector with smallest singular value - U, S, Vh = torch.linalg.svd(centered) - normal = Vh[-1] # Normal to best-fit plane - - # Compute deviations from plane (distance to plane) - deviations = torch.abs(centered @ normal) - - # Compute NLL for each atom - nll = gaussian_nll(deviations, plane_sigmas) - all_nlls.append(nll) - - if all_nlls: - return torch.cat(all_nlls) - return torch.tensor([0.0], device=device) - def nll_vdw(self, xyz: torch.Tensor = None): - """ - Compute negative log-likelihood for VDW (non-bonded) restraints. - - Uses a soft-repulsive potential based on distance violations. - NLL = 0.5 * (max(0, min_dist - actual_dist) / σ)² + log(σ) + 0.5 * log(2π) - - Only violations (distances shorter than minimum) contribute to the loss. - - Parameters - ---------- - xyz : torch.Tensor, optional - Coordinates tensor. If None, uses the stored xyz_fn callable. - - Returns - ------- - torch.Tensor - Tensor of shape (n_pairs,) with negative log-likelihood values. - """ - from torchref.refinement.targets import gaussian_nll - - xyz = self.xyz(xyz) - device = xyz.device - - if "vdw" not in self.restraints: - return torch.tensor([0.0], device=device) - - vdw_data = self.restraints["vdw"] - indices = vdw_data.get("indices") - - if indices is None or len(indices) == 0: - return torch.tensor([0.0], device=device) - - min_distances = vdw_data["min_distances"] - sigmas = vdw_data["sigmas"] - - # Get current positions - pos1 = xyz[indices[:, 0]] - pos2 = xyz[indices[:, 1]] - - # Compute actual distances - actual_distances = torch.norm(pos2 - pos1, dim=-1) - - # Violations: where actual distance is less than minimum - # Deviation = max(0, min_dist - actual_dist) - deviations = torch.clamp(min_distances - actual_distances, min=0.0) - - # Compute NLL (only non-zero for violations) - nll = gaussian_nll(deviations, sigmas) - - return nll def adp_b_differences(self, adp: torch.Tensor = None): """ @@ -2432,28 +1572,3 @@ def adp_b_differences(self, adp: torch.Tensor = None): return torch.cat(diffs_list, dim=0) return torch.tensor([], device=b_factors.device) - def adp_similarity_loss(self, adp: torch.Tensor = None, sigma: float = 2.0): - """ - Compute ADP similarity loss (SIMU in Phenix/SHELX). - - This restrains the B-factors of bonded atoms to be similar. - Loss = Σ ((B_i - B_j) / sigma)^2 - - Parameters - ---------- - adp : torch.Tensor, optional - ADP values. If None, uses the stored adp_fn callable. - sigma : float, default 2.0 - Target standard deviation for B-factor differences in Ų. - - Returns - ------- - torch.Tensor - Mean similarity loss. - """ - from torchref.refinement.targets import adp_similarity_nll - - b_diffs = self.adp_b_differences(adp) - if len(b_diffs) == 0: - return torch.tensor(0.0, device=self.xyz().device) - return adp_similarity_nll(b_diffs, sigma).mean() diff --git a/torchref/restraints/hydrogen_topology.py b/torchref/topology/riding.py similarity index 74% rename from torchref/restraints/hydrogen_topology.py rename to torchref/topology/riding.py index 6dd2632b..f36a681c 100644 --- a/torchref/restraints/hydrogen_topology.py +++ b/torchref/topology/riding.py @@ -1,21 +1,26 @@ -""" -Riding hydrogen topology and vectorized placement for VDW restraints. - -Builds a static topology map at restraints-construction time that describes -how to generate transient hydrogen atom positions from heavy-atom coordinates. -At each VDW evaluation the ``place_riding_hydrogens`` function produces H -positions in a single vectorized pass (no Python loops over atoms). - -Hydrogen positions are fully determined by the parent heavy atom and its -bonded heavy-atom neighbours, so gradients flow from the VDW loss through -the H positions back to the heavy-atom coordinates via standard autograd. +"""Riding hydrogens: the sterics of hydrogens a model does not carry. + +For a model loaded with ``strip_H=True``, whose atoms are heavy only. A static map +built once at restraint-construction time says how to reconstruct each absent hydrogen +from its parent and the parent's bonded neighbours; ``place_riding_hydrogens`` then +produces those positions in one vectorized pass at every non-bonded evaluation and +throws them away again. The positions are a function of the heavy atoms, so gradients +reach the heavy coordinates through them by ordinary autograd. + +Contrast :mod:`torchref.topology.hydrogens`, which *adds* hydrogens to the model as +real atoms with their own parameters. That is the default, and where both apply it is +the better answer: the hydrogen has a refinable position instead of one reconstructed +each step, and it contributes to the structure factors. Riding hydrogens are what is +left for the heavy-atom-only mode, and the two must not run together -- riding +placement alongside real hydrogens puts phantom atoms in the structure that push the +real ones around. """ +from dataclasses import dataclass, field from typing import Dict, Optional import numpy as np import torch -from torch import nn from torchref.config import dtypes, normalize_device from torchref.utils.device_resolution import resolve_device @@ -48,60 +53,114 @@ # --------------------------------------------------------------------------- -class HydrogenTopology(DeviceMixin, nn.Module): +@dataclass(eq=False, repr=False) +class HydrogenTopology(DeviceMixin): """Static topology describing riding hydrogens for VDW evaluation. - All data are stored as registered buffers so they move automatically - with ``.to(device)`` and appear in ``state_dict``. + Every tensor field starts as ``None``; :func:`build_hydrogen_topology` and + :func:`build_h_candidate_pairs` fill them, independently -- a topology can carry + hydrogens and no candidate pairs. Test :attr:`n_hydrogens` and + :attr:`has_candidates` rather than the fields. + + Holds no refinable parameters, so this is a dataclass rather than an + ``nn.Module``; ``DeviceMixin`` still moves every tensor with ``.to(device)``. + + Parameters + ---------- + device : torch.device, optional + Where the builders should allocate. Tracked from construction so a + ``resolve_device(h_topo, ...)`` called before anything is attached still + answers truthfully. Attributes ---------- - h_parent_idx : (N_h,) long - Index into heavy-atom array for each riding H. - h_bond_length : (N_h,) float - Ideal H–parent bond length (Å). - h_vdw_radius : (N_h,) float - Van der Waals radius for each H (1.20 Å). - h_placement_type : (N_h,) long - Placement-geometry enum (see module-level constants). - h_slot_in_parent : (N_h,) long - Ordinal within sibling H atoms on the same parent (0, 1, 2). - parent_neighbor_idx : (N_h, MAX_HEAVY_NB) long - Heavy-atom neighbour indices of the parent (-1 = padding). - parent_neighbor_count : (N_h,) long - Actual number of heavy-atom neighbours for the parent. - h_chainid_enc : (N_h,) long - Encoded chain ID (for same-residue filtering). - h_resseq : (N_h,) long - Residue sequence number (for same-residue filtering). - - Notes - ----- - The ``cand_*``/``n_asu_candidates``/``type_bounds`` buffers are added later by - ``build_h_candidate_pairs``, not by ``build_hydrogen_topology`` -- test - :attr:`has_candidates` before touching them. + h_parent_idx : torch.Tensor + Heavy-atom index of each riding H's parent, ``(N_h,)`` long. + h_bond_length : torch.Tensor + Ideal H-parent bond length in Angstroms, ``(N_h,)``. + h_vdw_radius : torch.Tensor + Van der Waals radius per H (1.20 A), ``(N_h,)``. + h_placement_type : torch.Tensor + Placement-geometry enum, ``(N_h,)`` long; see the module-level constants. + h_slot_in_parent : torch.Tensor + Ordinal among sibling H atoms on the same parent (0, 1, 2), ``(N_h,)`` long. + parent_neighbor_idx : torch.Tensor + Heavy-atom neighbours of the parent, ``(N_h, MAX_HEAVY_NB)`` long, ``-1`` + padded. + parent_neighbor_count : torch.Tensor + Heavy-atom neighbour count per parent, ``(N_h,)`` long. + h_chainid_enc : torch.Tensor + Encoded chain ID, ``(N_h,)`` long, for same-residue filtering. + h_resseq : torch.Tensor + Residue sequence number, ``(N_h,)`` long, for same-residue filtering. + type_bounds : dict + ``{placement_type: (start, end)}`` bounds into the type-sorted arrays. + cand_idx_i, cand_idx_j, cand_symop_idx, cand_cell_offset : torch.Tensor + Precomputed H candidate pairs, sorted so the asymmetric-unit ones come first. + cand_min_dist : torch.Tensor + Per-pair minimum-distance scratch buffer, ``(P,)``. + n_asu_candidates : int + How many leading candidate pairs lie inside the asymmetric unit. """ - def __init__(self, device=None): - super().__init__() - # Buffers are registered later by build_hydrogen_topology(), so this - # object is tensor-free at construction. The tracker still has to exist: - # ``DeviceMixin._refresh_device_trackers`` only maintains attributes - # already present in ``__dict__``, and callers reconcile against - # ``h_topo.device`` before attaching buffers to it. - self.device = normalize_device(device) + device: Optional[torch.device] = None + + h_parent_idx: Optional[torch.Tensor] = None + h_bond_length: Optional[torch.Tensor] = None + h_vdw_radius: Optional[torch.Tensor] = None + h_placement_type: Optional[torch.Tensor] = None + h_slot_in_parent: Optional[torch.Tensor] = None + parent_neighbor_idx: Optional[torch.Tensor] = None + parent_neighbor_count: Optional[torch.Tensor] = None + h_chainid_enc: Optional[torch.Tensor] = None + h_resseq: Optional[torch.Tensor] = None + type_bounds: Dict[int, tuple] = field(default_factory=dict) + + cand_idx_i: Optional[torch.Tensor] = None + cand_idx_j: Optional[torch.Tensor] = None + cand_symop_idx: Optional[torch.Tensor] = None + cand_cell_offset: Optional[torch.Tensor] = None + cand_min_dist: Optional[torch.Tensor] = None + n_asu_candidates: int = 0 + + # Derived at first placement and reused across steps; see reset_cache. + _dir_coeffs: Optional[torch.Tensor] = field(default=None, repr=False) + _nb_idx_clamped: Optional[torch.Tensor] = field(default=None, repr=False) + _nb_valid: Optional[torch.Tensor] = field(default=None, repr=False) + _bond_len_col: Optional[torch.Tensor] = field(default=None, repr=False) + + def __post_init__(self) -> None: + """Resolve the device tracker the builders allocate against.""" + self.device = normalize_device(self.device) @property def n_hydrogens(self) -> int: - """Number of riding hydrogens, or 0 before the buffers are attached.""" - if hasattr(self, "h_parent_idx"): - return self.h_parent_idx.shape[0] - return 0 + """Number of riding hydrogens, or 0 before the builders have run.""" + if self.h_parent_idx is None: + return 0 + return int(self.h_parent_idx.shape[0]) @property def has_candidates(self) -> bool: """Whether precomputed H candidate pairs are available.""" - return hasattr(self, "cand_idx_i") and self.cand_idx_i.shape[0] > 0 + return self.cand_idx_i is not None and self.cand_idx_i.shape[0] > 0 + + def reset_cache(self) -> None: + """Drop the derived placement tensors; rebuilt on the next placement call. + + Called by ``DeviceMixin`` on every ``.to()``, which is what keeps the clamped + neighbour indices and bond-length column from surviving a device move. + """ + self._dir_coeffs = None + self._nb_idx_clamped = None + self._nb_valid = None + self._bond_len_col = None + + def __repr__(self) -> str: + return ( + f"HydrogenTopology(n_hydrogens={self.n_hydrogens}, " + f"has_candidates={self.has_candidates})" + ) # --------------------------------------------------------------------------- @@ -109,20 +168,23 @@ def has_candidates(self) -> bool: # --------------------------------------------------------------------------- +#: Parsed monomer templates, keyed by residue name, shared across calls. Values are +#: None where the CIF is missing or carries no usable atom coordinates. +_TEMPLATE_CACHE: Dict = {} + + def _load_cif_hydrogen_info(pdb, verbose: int = 0) -> Dict: """``{resname: entry | None}`` H topology, ``None`` where the CIF is unusable. - Populates and returns the shared ``Model._hydrogenate_cif_cache``, so entries - from an earlier ``Model.hydrogenate()`` are reused. Each entry carries ``ids``, + Populates and returns :data:`_TEMPLATE_CACHE`. Each entry carries ``ids``, ``elems``, ``coords``, ``is_h``, ``id_to_idx``, ``heavy_names``, ``heavy_coords``, ``h_names``, ``h_coords``, ``parent_map``, ``ideal_bl`` and ``heavy_neighbor_map``. """ - from torchref.model.model import Model - from torchref.restraints.library import MonomerLibraryManager + from torchref.topology.monomer.library import MonomerLibraryManager lib = MonomerLibraryManager(verbose=0) - cache = Model._hydrogenate_cif_cache + cache = _TEMPLATE_CACHE for rn in pdb["resname"].unique(): rn_str = str(rn).strip() @@ -369,34 +431,17 @@ def build_hydrogen_topology( fdtype = dtypes.float if n_h_total == 0: - topo.register_buffer( - "h_parent_idx", torch.zeros(0, dtype=torch.long, device=device) - ) - topo.register_buffer( - "h_bond_length", torch.zeros(0, dtype=fdtype, device=device) - ) - topo.register_buffer( - "h_vdw_radius", torch.zeros(0, dtype=fdtype, device=device) - ) - topo.register_buffer( - "h_placement_type", torch.zeros(0, dtype=torch.long, device=device) - ) - topo.register_buffer( - "h_slot_in_parent", torch.zeros(0, dtype=torch.long, device=device) - ) - topo.register_buffer( - "parent_neighbor_idx", - torch.zeros(0, MAX_HEAVY_NB, dtype=torch.long, device=device), - ) - topo.register_buffer( - "parent_neighbor_count", torch.zeros(0, dtype=torch.long, device=device) - ) - topo.register_buffer( - "h_chainid_enc", torch.zeros(0, dtype=torch.long, device=device) - ) - topo.register_buffer( - "h_resseq", torch.zeros(0, dtype=torch.long, device=device) + topo.h_parent_idx = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: parent atom-index tensor (empty); int64 required + topo.h_bond_length = torch.zeros(0, dtype=fdtype, device=device) + topo.h_vdw_radius = torch.zeros(0, dtype=fdtype, device=device) + topo.h_placement_type = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: categorical H placement-type code (empty) + topo.h_slot_in_parent = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: slot index into parent (empty); int64 + topo.parent_neighbor_idx = torch.zeros( + 0, MAX_HEAVY_NB, dtype=torch.long, device=device # dtype-ok: parent neighbor atom-index tensor (empty); int64 required ) + topo.parent_neighbor_count = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: per-parent neighbor count (empty); structural int + topo.h_chainid_enc = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: categorical chain-id encoding (empty) + topo.h_resseq = torch.zeros(0, dtype=torch.long, device=device) # dtype-ok: residue sequence id (empty); categorical return topo # Sort all topology arrays by placement type for contiguous slicing @@ -421,43 +466,22 @@ def build_hydrogen_topology( idxs = np.where(mask)[0] type_bounds[t] = (int(idxs[0]), int(idxs[-1]) + 1) - topo.register_buffer( - "h_parent_idx", - torch.tensor(acc_parent_idx, dtype=torch.long, device=device), - ) - topo.register_buffer( - "h_bond_length", - torch.tensor(acc_bond_length, dtype=fdtype, device=device), - ) - topo.register_buffer( - "h_vdw_radius", - torch.full((n_h_total,), 1.20, dtype=fdtype, device=device), + topo.h_parent_idx = torch.tensor(acc_parent_idx, dtype=torch.long, device=device) # dtype-ok: parent atom-index tensor; torch indexing requires int64 + topo.h_bond_length = torch.tensor(acc_bond_length, dtype=fdtype, device=device) + topo.h_vdw_radius = torch.full((n_h_total,), 1.20, dtype=fdtype, device=device) + topo.h_placement_type = torch.tensor( + acc_placement_type, dtype=torch.long, device=device # dtype-ok: categorical H placement-type code; used for sort/slice ) - topo.register_buffer( - "h_placement_type", - torch.tensor(acc_placement_type, dtype=torch.long, device=device), + topo.h_slot_in_parent = torch.tensor(acc_slot, dtype=torch.long, device=device) # dtype-ok: slot index into parent neighbor slots; int64 + topo.parent_neighbor_idx = torch.tensor( + np.stack(acc_nb_idx), dtype=torch.long, device=device # dtype-ok: parent neighbor atom-index tensor; int64 required ) - topo.register_buffer( - "h_slot_in_parent", - torch.tensor(acc_slot, dtype=torch.long, device=device), - ) - topo.register_buffer( - "parent_neighbor_idx", - torch.tensor(np.stack(acc_nb_idx), dtype=torch.long, device=device), - ) - topo.register_buffer( - "parent_neighbor_count", - torch.tensor(acc_nb_count, dtype=torch.long, device=device), + topo.parent_neighbor_count = torch.tensor( + acc_nb_count, dtype=torch.long, device=device # dtype-ok: per-parent neighbor count; structural int metadata ) topo.type_bounds = type_bounds # dict: type_code -> (start, end) - topo.register_buffer( - "h_chainid_enc", - torch.tensor(acc_chainid_enc, dtype=torch.long, device=device), - ) - topo.register_buffer( - "h_resseq", - torch.tensor(acc_resseq, dtype=torch.long, device=device), - ) + topo.h_chainid_enc = torch.tensor(acc_chainid_enc, dtype=torch.long, device=device) # dtype-ok: categorical chain-id encoding + topo.h_resseq = torch.tensor(acc_resseq, dtype=torch.long, device=device) # dtype-ok: residue sequence id; categorical if verbose > 0: print(f" Hydrogen topology: {n_h_total} riding H atoms") @@ -612,7 +636,7 @@ def place_riding_hydrogens( xyz_h : (N_h, 3) float tensor, differentiable w.r.t. xyz_heavy """ # Function-local, matching every other target dispatch site: importing the gate at - # module scope would pull ``torchref.base.targets`` into ``torchref.restraints``. + # module scope would pull ``torchref.base.targets`` into ``torchref.topology``. from torchref.base.targets._dispatch import use_triton N_h = topo.h_parent_idx.shape[0] @@ -620,11 +644,11 @@ def place_riding_hydrogens( return torch.zeros(0, 3, dtype=xyz_heavy.dtype, device=xyz_heavy.device) # Precompute direction coefficients on first call - if not hasattr(topo, "_dir_coeffs") or topo._dir_coeffs is None: + if topo._dir_coeffs is None: topo._dir_coeffs = _precompute_direction_coefficients(topo) # Precompute static tensors on first call (avoid recomputing every step) - if not hasattr(topo, "_nb_idx_clamped"): + if topo._nb_idx_clamped is None: topo._nb_idx_clamped = topo.parent_neighbor_idx.clamp(min=0) topo._nb_valid = ( (topo.parent_neighbor_idx >= 0).unsqueeze(-1).to(topo.h_bond_length.dtype) @@ -712,15 +736,9 @@ def build_h_candidate_pairs( if n_h == 0: for name in ("cand_idx_i", "cand_idx_j", "cand_symop_idx"): - h_topo.register_buffer( - name, torch.zeros(0, dtype=torch.long, device=device) - ) - h_topo.register_buffer( - "cand_cell_offset", torch.zeros(0, 3, dtype=torch.long, device=device) - ) - h_topo.register_buffer( - "cand_min_dist", torch.zeros(0, dtype=dtypes.float, device=device) - ) + setattr(h_topo, name, torch.zeros(0, dtype=torch.long, device=device)) # dtype-ok: candidate atom/symop index tensors (empty); int64 required + h_topo.cand_cell_offset = torch.zeros(0, 3, dtype=torch.long, device=device) # dtype-ok: integer cell-offset lattice vectors (empty); symmetry metadata + h_topo.cand_min_dist = torch.zeros(0, dtype=dtypes.float, device=device) return heavy_indices = vdw_data["indices"] # (P, 2) @@ -818,21 +836,15 @@ def _same_res(chain_a, resseq_a, chain_b, resseq_b): if not acc_idx_i: for name in ("cand_idx_i", "cand_idx_j", "cand_symop_idx"): - h_topo.register_buffer( - name, torch.zeros(0, dtype=torch.long, device=device) - ) - h_topo.register_buffer( - "cand_cell_offset", torch.zeros(0, 3, dtype=torch.long, device=device) - ) - h_topo.register_buffer( - "cand_min_dist", torch.zeros(0, dtype=dtypes.float, device=device) - ) + setattr(h_topo, name, torch.zeros(0, dtype=torch.long, device=device)) # dtype-ok: candidate atom/symop index tensors (empty); int64 required + h_topo.cand_cell_offset = torch.zeros(0, 3, dtype=torch.long, device=device) # dtype-ok: integer cell-offset lattice vectors (empty); symmetry metadata + h_topo.cand_min_dist = torch.zeros(0, dtype=dtypes.float, device=device) return - cand_i = torch.tensor(acc_idx_i, dtype=torch.long, device=device) - cand_j = torch.tensor(acc_idx_j, dtype=torch.long, device=device) - cand_sym = torch.tensor(acc_symop, dtype=torch.long, device=device) - cand_off = torch.tensor(np.stack(acc_offset), dtype=torch.long, device=device) + cand_i = torch.tensor(acc_idx_i, dtype=torch.long, device=device) # dtype-ok: combined atom-index tensor; torch indexing requires int64 + cand_j = torch.tensor(acc_idx_j, dtype=torch.long, device=device) # dtype-ok: combined atom-index tensor; torch indexing requires int64 + cand_sym = torch.tensor(acc_symop, dtype=torch.long, device=device) # dtype-ok: symmetry-operator index; int64 + cand_off = torch.tensor(np.stack(acc_offset), dtype=torch.long, device=device) # dtype-ok: integer cell-offset lattice vectors; symmetry-image metadata # Apply 1-2 / 1-3 exclusions for intra-ASU candidates if h_excl_hash is not None and len(h_excl_hash) > 0: @@ -889,16 +901,13 @@ def _same_res(chain_a, resseq_a, chain_b, resseq_b): cand_off = cand_off[sort_order] n_asu_cand = is_asu.sum().item() - h_topo.register_buffer("cand_idx_i", cand_i) - h_topo.register_buffer("cand_idx_j", cand_j) - h_topo.register_buffer("cand_symop_idx", cand_sym) - h_topo.register_buffer("cand_cell_offset", cand_off) + h_topo.cand_idx_i = cand_i + h_topo.cand_idx_j = cand_j + h_topo.cand_symop_idx = cand_sym + h_topo.cand_cell_offset = cand_off h_topo.n_asu_candidates = n_asu_cand - h_topo.register_buffer( - "cand_min_dist", - torch.zeros(len(cand_i), dtype=dtypes.float, device=device), - ) + h_topo.cand_min_dist = torch.zeros(len(cand_i), dtype=dtypes.float, device=device) if verbose > 0: n_hh = ((cand_i >= n_heavy) & (cand_j >= n_heavy)).sum().item() diff --git a/torchref/topology/templates.py b/torchref/topology/templates.py new file mode 100644 index 00000000..afb44aea --- /dev/null +++ b/torchref/topology/templates.py @@ -0,0 +1,105 @@ +"""Monomer templates and the link modifications that patch them. + +The monomer library describes each residue in its **free** form. Forming a peptide bond +applies the modifications the ``chem_link`` table names -- ``DEL-OXT`` to the residue +donating its C, ``DEL-HN1`` (``DEL-HNP`` for proline) to the residue donating its N -- +which delete the restraints the link makes meaningless and overwrite the targets that +change. A residue therefore draws its restraints from a *patched* template, identified +by a key such as ``'ALA:DEL-HN1+DEL-OXT'``. + +Chain termini are deliberately left unpatched: a real C-terminus keeps its ``OXT`` and +carboxylate geometry, a real N-terminus its ammonium. +""" + +from typing import Dict, Sequence, Tuple + +import numpy as np + +from torchref.topology.monomer.modifications import ( + apply_modifications, + link_modifications, + read_mod_definitions, +) + + +def resolve_template_keys( + resnames: Sequence[str], + peptide_pairs: Sequence[Tuple[int, int]], + cif_dict: Dict, + link_list, + verbose: int = 0, +) -> Tuple[Dict, np.ndarray]: + """Assign each residue the template it should draw restraints from. + + Parameters + ---------- + resnames : sequence of str + Residue name per residue index. + peptide_pairs : sequence of tuple of int + ``(donates C, donates N)`` residue index pairs. + cif_dict : dict + Restraint dictionary keyed by residue name. + link_list : pandas.DataFrame or None + Link-type definitions, as + :func:`~torchref.topology.monomer.cif.read_link_definitions` returns + them. None disables patching. + verbose : int, default 0 + Verbosity level. + + Returns + ------- + comp_dict : dict + ``cif_dict`` plus one entry per patched ``(residue type, modification set)`` in + use. ``cif_dict`` itself is left keyed by residue name alone. + template_key : numpy.ndarray + Key per residue, shape ``(R,)``. Residues that are not patched carry their own + residue name. + """ + comp_dict = dict(cif_dict) + keys = np.array([str(r) for r in resnames], dtype=object) + + if not cif_dict or link_list is None or len(peptide_pairs) == 0: + return comp_dict, keys + + modifications = link_modifications(link_list) + if "TRANS" not in modifications: + return comp_dict, keys + trans_mods = modifications["TRANS"] + proline_mods = modifications.get("PTRANS", trans_mods) + + mods_by_residue: Dict[int, set] = {} + for res_c, res_n in peptide_pairs: + donor_mod, acceptor_mod = ( + proline_mods if str(resnames[res_n]) == "PRO" else trans_mods + ) + for res_idx, mod_id in ((res_c, donor_mod), (res_n, acceptor_mod)): + if mod_id is None: + continue + mods_by_residue.setdefault(res_idx, set()).add(mod_id) + + if not mods_by_residue: + return comp_dict, keys + + mod_dict = read_mod_definitions() + n_patched = 0 + for res_idx, mods in mods_by_residue.items(): + resname = str(resnames[res_idx]) + if resname not in comp_dict: + continue + variant = f"{resname}:{'+'.join(sorted(mods))}" + if variant not in comp_dict: + comp_dict[variant] = apply_modifications( + comp_dict[resname], sorted(mods), mod_dict + ) + keys[res_idx] = variant + n_patched += 1 + + if verbose > 1: + print( + f"Patched {n_patched} residues " + f"({len(comp_dict) - len(cif_dict)} template variants)" + ) + return comp_dict, keys + + +__all__ = ["resolve_template_keys"] diff --git a/torchref/topology/topology.py b/torchref/topology/topology.py new file mode 100644 index 00000000..51419991 --- /dev/null +++ b/torchref/topology/topology.py @@ -0,0 +1,178 @@ +"""The topology container: a residue graph over an atom graph. + +:class:`Topology` is what the model's connectivity lives in. The residue level carries +the sequence and the inter-residue links; the atom level carries the atoms, the typed +edge blocks and the bond adjacency. Per-atom residue identity is reached through +``atoms.residue_of`` rather than duplicated per atom. + +Mutable by design; prefer :meth:`Topology.copy` over editing in place. +""" + +from dataclasses import dataclass +from typing import Dict, Set, Tuple + +import numpy as np +import torch + +from torchref.topology.atom_graph import AtomGraph +from torchref.topology.residue_graph import ResidueGraph +from torchref.utils.device_mixin import DeviceMixin + + +@dataclass(eq=False, repr=False) +class Topology(DeviceMixin): + """Connectivity of one model, at both the residue and the atom level. + + Parameters + ---------- + residues : ResidueGraph + Sequence level -- residues as nodes, links as edges. + atoms : AtomGraph + Atom level -- atoms as nodes, typed edge blocks, bond adjacency. + + Notes + ----- + Holds no refinable parameters, so this is a dataclass rather than an ``nn.Module``. + Edge indices are ``int64`` constants and no gradient reaches them; gradients reach + the coordinates that the indices gather. + """ + + residues: ResidueGraph + atoms: AtomGraph + + @property + def device(self) -> torch.device: + """Where the indexing tensors live. Derived from the atom graph.""" + return self.atoms.device + + @property + def n_atoms(self) -> int: + """Number of atom nodes.""" + return self.atoms.n_atoms + + @property + def n_residues(self) -> int: + """Number of residue nodes.""" + return self.residues.n_residues + + def copy(self) -> "Topology": + """An independent copy sharing no storage with this one.""" + return Topology(residues=self.residues.copy(), atoms=self.atoms.copy()) + + def subset(self, keep) -> "Topology": + """The topology over a subset of the atoms. + + Reindexes what survives instead of rebuilding: no CIF is re-read and no template + is re-matched, which is what made ``Model.select`` expensive. + + Parameters + ---------- + keep : torch.Tensor or numpy.ndarray + Boolean mask over atoms, shape ``(N,)``, or integer atom indices. Indices + are taken as a set, not an order -- the result keeps the topology's own atom + order, because the edge blocks stay canonical only under a monotone + relabelling. + + Returns + ------- + Topology + Atoms in their original relative order. A residue with no surviving atoms is + dropped, and any link edge touching it goes with it. + + Notes + ----- + Selecting part of a residue leaves that residue's restraints partial: an edge + loses its whole restraint as soon as one of its atoms goes. That is the honest + outcome -- half a peptide plane is not a plane -- but it means a subset is a + weaker geometric model, not merely a smaller one. + """ + mask = torch.as_tensor(keep) + if mask.dtype != torch.bool: + selected = torch.zeros(self.n_atoms, dtype=torch.bool) + selected[mask.to(torch.int64)] = True # dtype-ok: boolean-mask->index cast for scatter select; int64 index required + mask = selected + mask = mask.to(device=self.atoms.residue_of.device) + + if int(mask.sum()) == 0: + raise ValueError("subset would keep no atoms") + + n_kept = int(mask.sum()) + remap = torch.full((self.n_atoms,), -1, dtype=torch.int64, device=mask.device) # dtype-ok: atom remap index array (-1 sentinel); int64 index required + remap[mask] = torch.arange(n_kept, dtype=torch.int64, device=mask.device) # dtype-ok: arange remap indices; int64 index required + + # A residue survives if any of its atoms does. Counting per residue also + # gives the new atom ranges, contiguous because the atom order is unchanged. + residue_of = self.atoms.residue_of + per_residue = ( + torch.bincount(residue_of[mask], minlength=self.n_residues).cpu().numpy() + ) + residue_keep = per_residue > 0 + counts = per_residue[residue_keep] + atom_end = np.cumsum(counts) + atom_start = atom_end - counts + + residue_remap = torch.full( + (self.n_residues,), -1, dtype=torch.int64, device=mask.device # dtype-ok: residue remap index array (-1 sentinel); int64 index required + ) + residue_remap[torch.as_tensor(residue_keep, device=mask.device)] = torch.arange( + int(residue_keep.sum()), dtype=torch.int64, device=mask.device # dtype-ok: arange residue remap indices; int64 index required + ) + + return Topology( + residues=self.residues.subset( + residue_keep, + atom_start.astype(np.int64), + atom_end.astype(np.int64), + ), + atoms=self.atoms.subset(remap, residue_remap), + ) + + def neighbors(self, i: int) -> torch.Tensor: + """Atoms bonded to atom ``i``. Delegates to :meth:`AtomGraph.neighbors`.""" + return self.atoms.neighbors(i) + + def residue_of_atom(self, i: int) -> int: + """Residue index of atom ``i``.""" + return int(self.atoms.residue_of[i]) + + def resname_of_atom(self, i: int) -> str: + """Residue name of atom ``i``, joined through the residue graph.""" + return str(self.residues.resname[self.residue_of_atom(i)]) + + def edge_block(self, edge_type: str): + """The :class:`~torchref.topology.edges.EdgeBlock` for one edge type. + + Parameters + ---------- + edge_type : str + ``'bond'``, ``'angle'``, ``'torsion'`` or ``'chiral'``. Planes are ragged + and reached through ``atoms.planes``. + """ + return { + "bond": self.atoms.bonds, + "angle": self.atoms.angles, + "torsion": self.atoms.torsions, + "chiral": self.atoms.chirals, + }[edge_type] + + def tuple_sets(self) -> Dict[str, Dict[str, Set[Tuple[int, ...]]]]: + """Every edge as ``{edge type: {origin: set of index tuples}}``. + + Order-free, so this is what an equivalence check against another builder should + compare. + """ + out: Dict[str, Dict[str, Set[Tuple[int, ...]]]] = {} + for name in ("bond", "angle", "torsion", "chiral"): + block = self.edge_block(name) + out[name] = {o: block.tuple_set(o) for o in block.origins()} + out["plane"] = {} + for size, block in self.atoms.planes.items(): + for origin in block.origins(): + out["plane"][f"{size}_atoms/{origin}"] = block.tuple_set(origin) + return out + + def __repr__(self) -> str: + return f"Topology({self.residues!r}, {self.atoms!r})" + + +__all__ = ["Topology"] diff --git a/torchref/utils/device_mixin.py b/torchref/utils/device_mixin.py index 6c8327f1..c4f4d2da 100644 --- a/torchref/utils/device_mixin.py +++ b/torchref/utils/device_mixin.py @@ -347,14 +347,14 @@ def run(device, dtype): # Each axis gets its own pair, varying only along the axis it measures. Sharing one # pair couples them: an accelerator scratch cannot be cast to float64 on MPS, so # probing dtype on the device pair makes ``.double()`` unprobeable there. - base = run(torch.device("cpu"), torch.float32) + base = run(torch.device("cpu"), torch.float32) # dtype-ok: fixed probe dtype is what the preservation test varies, not a config allocation if base is None: return None, None device = base.device if accel is not None: # Contrast pair for the device axis: same dtype, different device. - other = run(accel, torch.float32) + other = run(accel, torch.float32) # dtype-ok: fixed probe dtype (device-axis contrast) if other is None: device = None elif other.device != base.device: @@ -365,7 +365,7 @@ def run(device, dtype): # Contrast pair for the dtype axis: same device, different dtype. # float16 rather than float64 so this stays cheap and universally # supported; the CPU pin means ``.double()`` remains probeable. - other = run(torch.device("cpu"), torch.float16) + other = run(torch.device("cpu"), torch.float16) # dtype-ok: fixed probe dtype (dtype-axis contrast) if other is not None and other.dtype == base.dtype: dtype = base.dtype