diff --git a/docs/changelog.rst b/docs/changelog.rst index 615ec9aa..6cb09d5c 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -13,6 +13,8 @@ Breaking changes - Anomalous scattering is off unless ``--wavelength`` is given (it defaulted to 1.0 Å) - Bulk-solvent phase offset removed; ``optimize_phase`` is deprecated - ``phased-difference-map`` is now ``difference-map``: fewer default columns (``--all-columns``), renamed (``dF``, ``W_InVa``) +- ``import torchref`` no longer imports torch: the top-level names (``Model``, ``ReflectionData``, ``dtypes``, ...) load on first access, and ``TORCHREF_DTYPE_*``/``TORCHREF_DEVICE`` are read then +- ``torchref.difference-map`` takes ``-o`` as a directory (``difference_map.mtz``, ``summary.json``) and gains ``--dmax``; the difference MTZ adds ``DELFWT`` (``dF`` times the ``--ded-weight`` weight), writes the phase-aware and scalar extrapolations by default, and renames ``FC``/``PHIC``/``PHIC_diff`` to ``Fc_light``/``PHFc_light``/``PHdFc`` - Refinement checkpoints restore only through ``load_state``; ``Refinement.create_from_state_dict`` is removed - ``french_wilson_auto`` moved to ``torchref.io.datasets``; ``CollectionRiceTarget`` is now ``CollectionMLTarget`` - Refinement writes its own PDB/mmCIF header instead of copying the input's @@ -33,6 +35,7 @@ Performance ~~~~~~~~~~~ - CPU density kernels ship prebuilt in the new ``torchref-kernels`` package (Rust): nothing is compiled at runtime, ``ninja`` is no longer a dependency, and the CPU splat is multithreaded on macOS. Wheels cover Linux x86_64/aarch64, macOS arm64 and Windows x86_64; other platforms run the portable CPU kernels unless ``torchref-kernels`` is built from source - Faster CPU VDW pair search, hydrogen-mode switching and ADP-locality neighbours; ``SfFFT`` stores no real-space grid +- Every command-line tool answers ``--help`` and argument errors without importing torch - The command-line tools set ``OMP_WAIT_POLICY=PASSIVE``, so idle OpenMP workers no longer spin on the cores the CPU kernels need: multithreaded CPU refinement is faster. Library use is unchanged; see the CLI guide to opt in Fixes that change results @@ -45,6 +48,7 @@ Fixes that change results - Refinement: frozen parameters stay frozen, rigid-body translations refine, the scaler is refit warm each cycle - Density kernels on CPU, CUDA and MPS; long matrix products on MPS; REMARK 3 angle RMSD reported in degrees - CUDA ``ls``/``nll`` X-ray kernels score a signed ``F_calc`` as ``|F_calc|``, as the eager path does +- Difference maps take their phases from a fit of the dark model alone, so they no longer change with ``-lm``; without ``--dmin`` the models' FFT grid follows the data instead of a fixed 1.0 Å - ``WilsonNormaliser`` with ``s_lo``/``s_hi`` wider than the data: stable in float32 and no longer truncated by its log clamp Other fixes diff --git a/docs/user_guide/cli.rst b/docs/user_guide/cli.rst index b407dfa7..79a79fec 100644 --- a/docs/user_guide/cli.rst +++ b/docs/user_guide/cli.rst @@ -236,29 +236,36 @@ 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_Q``, the q-weight (the default), and ``W_InVa``, -the inverse variance ``1/sigma^2`` -- 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_Q -cphi PHDELWT``, adding -``--units electrons`` for e/A^3. It needs no light-state model, so ``-lm`` is optional: +The default output is the difference map ``DELFWT``/``PHDELWT``: the amplitude +difference ``dF`` times the ``--ded-weight`` weight, on the **dark** model's phases, which +Coot opens directly. The phases come from a fit of the dark model alone, so the +difference map is the same with or without a light model. Beside it sit ``dF``/``SIGdF``, +one mean-one weight column per registered scheme -- ``W_Q``, the q-weight (the default), +and ``W_InVa``, the inverse variance ``1/sigma^2`` -- and ``KSCALE``, the scaler's factor +from model to observed scale, so another weighting is +``torchref.mtz2map -csf dF -cw W_InVa -cphi PHDELWT`` (``--units electrons`` for e/A^3). +This is the construction ``torchref.validate-ded`` correlates against. ``-o`` is a +directory; it receives ``difference_map.mtz`` and ``summary.json`` (inputs, dataset +scaling, R-factors, the ``q``-weight fit and the extrapolation shrinkage). .. code-block:: bash torchref.difference-map \ -dm dark.pdb \ - -dsf dark.mtz -lsf light.mtz -o results.mtz + -dsf dark.mtz -lsf light.mtz -o results/ -Supplying ``-lm`` (with ``--fraction``) adds the light state's amplitude and -phase and the extrapolated map ``FWT``/``PHWT``: +The light-state model is optional. Supplying ``-lm`` (with ``--fraction``) adds the +light data's model ``Fc_light``/``PHFc_light`` and three extrapolated light-state maps: +``FWT``/``PHWT`` from the shrunk amplitudes ``FEXT``, ``2FEXT_PHASED-Fc``/``PHFEXT_PHASED`` +from the phase-aware ones and ``2FEXT_SCALAR-Fc``/``PHFEXT_SCALAR`` from the classic +``(Fo_light - w_dark Fo_dark) / w_light``, each on its own scale fit: .. code-block:: bash torchref.difference-map \ -dm dark.pdb -lm light.pdb \ -dsf dark.mtz -lsf light.mtz \ - --fraction 0.37 -o results.mtz + --fraction 0.37 -o results/ ``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 @@ -267,18 +274,18 @@ records this: the columns sit in named MTZ datasets -- ``observed``, ``differenc 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 {q,inverse_variance,none}`` selects the scheme the -model-phased and two-moment difference columns carry (default ``q``; its fit uses the -intensity differences when the data carry ``I``/``SIGI``, and falls back to inverse -variance with a warning when too few reflections exist to fit); ``--difference-gamma`` -fixes the dark-amplitude exponent of the difference power law instead of fitting it; -``--difference-sigma-scale`` scales the reported sigmas the ``q`` weight and the -extrapolation take as the noise (default 1; ``fit`` estimates the scale from the -differences, which reads noise as signal when the sigmas' calibration changes with -resolution, as French-Wilson and other posterior amplitude sigmas' does); -``--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. +**Key options:** ``--dmin``/``--dmax`` resolution limits (without ``--dmin`` the +models' grid follows the data); ``--ded-weight {q,inverse_variance,none}`` selects the +scheme ``DELFWT`` and the model-phased difference columns carry (default ``q``; its fit +uses the intensity differences when the data carry ``I``/``SIGI``, and falls back to +inverse variance with a warning when too few reflections exist to fit); +``--difference-gamma`` fixes the dark-amplitude exponent of the difference power law +instead of fitting it; ``--difference-sigma-scale`` scales the reported sigmas the ``q`` +weight and the extrapolation take as the noise (default 1; ``fit`` estimates the scale +from the differences, which reads noise as signal when the sigmas' calibration changes +with resolution, as French-Wilson and other posterior amplitude sigmas' does); +``--all-columns`` (needs ``-lm``) adds the model-phased difference and the intensity +block. :API: :mod:`torchref.cli.difference_map` diff --git a/tests/integration/test_cli_ded_weights.py b/tests/integration/test_cli_ded_weights.py index 3c99c36b..94c6a41a 100644 --- a/tests/integration/test_cli_ded_weights.py +++ b/tests/integration/test_cli_ded_weights.py @@ -24,6 +24,7 @@ "SIGFo_light": "Stddev", "dF": "SFAmplitude", "SIGdF": "Stddev", + "DELFWT": "SFAmplitude", "PHDELWT": "Phase", "W_InVa": "Weight", "W_Q": "Weight", @@ -105,7 +106,7 @@ def _run(project_root, module, *argv, timeout=1800): @pytest.fixture(scope="module") def diff_mtz(project_root, pair): - out = pair["dir"] / "diff.mtz" + out = pair["dir"] / "diff" _run( project_root, "difference_map.py", @@ -126,7 +127,7 @@ def diff_mtz(project_root, pair): "-o", out, ) - return out + return out / "difference_map.mtz" def _read(path): diff --git a/tests/integration/test_cli_two_moment_mtz.py b/tests/integration/test_cli_two_moment_mtz.py index 446c631d..1bc6e2cb 100644 --- a/tests/integration/test_cli_two_moment_mtz.py +++ b/tests/integration/test_cli_two_moment_mtz.py @@ -19,20 +19,29 @@ "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. + # The difference map: weighted dF on the dark phases, one weight column per + # registered scheme, and the observed-to-model scale. + "DELFWT": "SFAmplitude", "PHDELWT": "Phase", "W_InVa": "Weight", "W_Q": "Weight", "KSCALE": "MTZReal", "Fc_dark": "SFAmplitude", - # The mixed model, and the extrapolated map to refine against. - "FC": "SFAmplitude", - "PHIC": "Phase", + # The light data's (mixed) model, and the three extrapolated maps. + "Fc_light": "SFAmplitude", + "PHFc_light": "Phase", "FEXT": "SFAmplitude", "SIGFEXT": "Stddev", "FWT": "SFAmplitude", "PHWT": "Phase", + "FEXT_PHASED": "SFAmplitude", + "SIGFEXT_PHASED": "Stddev", + "2FEXT_PHASED-Fc": "SFAmplitude", + "PHFEXT_PHASED": "Phase", + "FEXT_SCALAR": "SFAmplitude", + "SIGFEXT_SCALAR": "Stddev", + "2FEXT_SCALAR-Fc": "SFAmplitude", + "PHFEXT_SCALAR": "Phase", } FLAG_COLUMNS = {"FreeR_flag_dark", "FreeR_flag_light"} @@ -49,19 +58,9 @@ ALL_COLUMNS_EXTRA = { "2mDFop-DFc": "SFAmplitude", "mDFop-DFc": "SFAmplitude", - "PHIC_diff": "Phase", + "PHdFc": "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. @@ -229,8 +228,8 @@ def test_map_coefficients_are_grouped_into_described_datasets(two_moment_all_mtz 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_Q"] == "difference" - assert where["FC"] == where["PHIC"] == "light_model" + assert where["DELFWT"] == where["PHDELWT"] == where["W_Q"] == "difference" + assert where["Fc_light"] == where["PHFc_light"] == "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) @@ -246,8 +245,8 @@ def test_the_difference_columns_carry_df_and_the_registered_weights( self, baseline_mtz ): """``dF`` must be ``Fo_light - Fo_dark``, ``W_InVa`` the mean-normalised inverse - variance and ``W_Q`` a mean-one weight -- the constructions - ``torchref.validate-ded`` correlates against. If these ever diverge, the map + variance, ``W_Q`` a mean-one weight and ``DELFWT`` ``dF * W_Q`` -- 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 @@ -270,12 +269,14 @@ def test_the_difference_columns_carry_df_and_the_registered_weights( assert np.abs(got_w - w).max() / max(float(np.abs(w).max()), 1e-30) < 1e-4 w_q = df["W_Q"].to_numpy().astype(float) assert np.isfinite(w_q).all() and abs(w_q.mean() - 1.0) < 1e-4 + got_delfwt = df["DELFWT"].to_numpy().astype(float) + assert np.abs(got_delfwt - got_df * w_q).max() / scale < 1e-5 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), + df["PHFc_light"].to_numpy().astype(float), ) diff --git a/torchref/__init__.py b/torchref/__init__.py index c9ba9ddc..bb0f7c72 100644 --- a/torchref/__init__.py +++ b/torchref/__init__.py @@ -11,6 +11,9 @@ ) refinement.refine(macro_cycles=10) +``import torchref`` does not import torch: the names re-exported here and the +subpackages load on first access, and :mod:`torchref.config` reads its environment then. + Modules ------- io @@ -42,13 +45,15 @@ __version__ = "0.7.0" +import importlib import os - -# Must be set before torch is imported below, or it has no effect. -os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") +import sys import warnings from pathlib import Path +# Must be set before torch is first imported, or it has no effect. +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") + from torchref._bootstrap import configure_threading, detect_available_cpus # Threading must be configured before torch is imported. @@ -68,12 +73,10 @@ configure_threading(N_CPUS) -import torch - -torch.set_num_threads(N_CPUS) - -# Must come after torch: reads its dtype/device state at import. -from torchref.config import device, dtypes, sigma_cutoff_ed +# A torch already imported has read its thread count; one imported later reads it from +# the OMP_NUM_THREADS just set. +if "torch" in sys.modules: + sys.modules["torch"].set_num_threads(N_CPUS) # Project root path for referencing package files ROOT_TORCHREF = Path(__file__).parent.parent.resolve() @@ -83,80 +86,44 @@ PATH_TORCHREF_DATA = PATH_TORCHREF / "data" -# ============================================================================= -# Convenience imports for common classes -# ============================================================================= - - -# Data I/O -from torchref.io import ( - DatasetCollection, - 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 - -# Refinement -from torchref.refinement import LBFGSRefinement, Refinement -from torchref.refinement.rigid_body_refinement import RigidBodyRefinementStep - -# Scaling -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__", - "ROOT_TORCHREF", - "PATH_TORCHREF", - "N_CPUS", - # Dtype and device configuration - "dtypes", - "device", - "sigma_cutoff_ed", - # Data I/O - "ReflectionData", - "ScaledDataset", - "DatasetCollection", - "read_mtz", - "read_cif", - "read_pdb", - # Model - "Model", - "ModelFT", - "RigidXYZTensor", - # Refinement - "Refinement", - "LBFGSRefinement", - "RigidBodyRefinementStep", - # Scaling - "Scaler", - "ScalerBase", - "SolventModel", - # Symmetry - "Cell", - "SpaceGroup", - "Symmetry", - # Maps - "Map", - "DifferenceMap", - # Device mixin - "DeviceMixin", -] +# Defining module -> public names, imported on first access by ``__getattr__`` so that +# ``import torchref`` (and every CLI ``--help``) doesn't import torch. ``Restraints`` is +# deliberately absent: constructing it can trigger a monomer-library download. +_LAZY_EXPORTS = { + name: module + for module, names in { + "torchref.config": "dtypes device sigma_cutoff_ed", + "torchref.io": "ReflectionData ScaledDataset DatasetCollection" + " read_mtz read_cif read_pdb", + "torchref.model": "Model ModelFT", + "torchref.model.rigid_xyz": "RigidXYZTensor", + "torchref.refinement": "Refinement LBFGSRefinement", + "torchref.refinement.rigid_body_refinement": "RigidBodyRefinementStep", + "torchref.scaling": "Scaler ScalerBase SolventModel", + "torchref.symmetry": "Cell SpaceGroup Symmetry", + "torchref.maps": "Map DifferenceMap", + "torchref.utils.device_mixin": "DeviceMixin", + }.items() + for name in names.split() +} + + +def __getattr__(name: str): + """Import a re-exported name, or a subpackage, on first access.""" + if name in _LAZY_EXPORTS: + value = getattr(importlib.import_module(_LAZY_EXPORTS[name]), name) + elif not name.startswith("_") and any( + (PATH_TORCHREF / n).exists() for n in (name, f"{name}.py") + ): + value = importlib.import_module(f"{__name__}.{name}") + else: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + globals()[name] = value + return value + + +def __dir__(): + return sorted(set(globals()) | set(_LAZY_EXPORTS)) + + +__all__ = ["__version__", "ROOT_TORCHREF", "PATH_TORCHREF", "N_CPUS", *_LAZY_EXPORTS] diff --git a/torchref/cli/_choices.py b/torchref/cli/_choices.py new file mode 100644 index 00000000..7add18f5 --- /dev/null +++ b/torchref/cli/_choices.py @@ -0,0 +1,26 @@ +"""Choice lists the CLI parsers show in ``--help``, held without importing torch. + +Each mirrors the table named beside it, which stays the source of truth and validates +the parsed value; importing that table imports torch, which ``--help`` must not await. +""" + +#: :data:`torchref.maps.ded_weights.SCHEMES` and ``DEFAULT_SCHEME``. +DED_WEIGHT_SCHEMES, DEFAULT_DED_WEIGHT_SCHEME = ("none", "inverse_variance", "q"), "q" + +#: ``difference_power.SIGMA_SCALE_BOUNDS`` (``refinement.model_error_estimation``). +SIGMA_SCALE_BOUNDS = (0.1, 10.0) + +#: :data:`torchref.scaling.scaler_base.SCALE_TARGETS` and ``DEFAULT_SCALE_TARGET``. +SCALE_TARGETS, DEFAULT_SCALE_TARGET = ("nll", "ml_noalpha", "ls"), "ls" + +#: The rows of :data:`torchref.refinement.targets.xray._specs.XRAY_TARGETS`, summarised. +XRAY_TARGET_DOCS = { + "ml": "Read MLF, variance epsilon*beta, mean alpha*|F_calc| (default).", + "ml_noalpha": "As 'ml' with the Luzzati mean coupling fixed at 1.", + "ml_full": "Full-form MLF, marginalising the error-free amplitude. ~4x the cost.", + "nll_beta": "Gaussian amplitude NLL on ml's model-error variance (diagnostic).", + "nll": "Gaussian amplitude NLL on sigma_obs only; does not control overfitting.", + "nll_i": "Gaussian NLL on the observed intensities, skipping French-Wilson.", + "ls": "Least squares with unit weights; the scaler owns the overall scale.", + "ls_wunit_k1": "Phenix-style least squares, one global scale per gradient call.", +} diff --git a/torchref/cli/_common.py b/torchref/cli/_common.py index 4f946203..386aadb3 100644 --- a/torchref/cli/_common.py +++ b/torchref/cli/_common.py @@ -335,7 +335,8 @@ def add_dual_model_args( ``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. + the dark state's phases. Refinement cannot: it refines the light model. ``-lm`` and + ``--fraction`` then sit in their own *optional* group, so ``--help`` says so. """ inp = parser.add_argument_group("Input files") inp.add_argument( @@ -345,14 +346,17 @@ 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=light_model_required, type=str, help=light_help + light = inp if light_model_required else parser.add_argument_group( + "Light-state model (optional)", + "Without these only the difference map, on the dark model's phases, is\n" + "written; with them the light-model and extrapolated maps are added.", + ) + light.add_argument( + "-lm", + "--light-model", + required=light_model_required, + type=str, + help="Light / triggered state model file (PDB or CIF)", ) inp.add_argument( "-dsf", @@ -377,7 +381,7 @@ def add_dual_model_args( frac_kwargs["required"] = True else: frac_kwargs["default"] = fraction_default - inp.add_argument("--fraction", **frac_kwargs) + light.add_argument("--fraction", **frac_kwargs) add_cif_arg(inp) col = parser.add_argument_group("Column selection") @@ -387,31 +391,27 @@ def add_dual_model_args( 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 + Off by default so the output file holds the maps a reader wants and can identify. + The gated columns -- the model-phased difference and the intensity block -- 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.", + help="Also write the model-phased difference and diagnostic columns, beyond " + "the default difference and extrapolated maps (needs a light model).", ) def _sigma_scale_arg(value: str) -> float | None: """``--difference-sigma-scale``: ``fit`` (``None``) or a scale within bounds.""" - from torchref.refinement.model_error_estimation.difference_power import ( - _check_sigma_scale, - ) + from torchref.cli._choices import SIGMA_SCALE_BOUNDS if value.strip().lower() == "fit": return None - try: - return _check_sigma_scale(float(value)) - except ValueError as err: - raise argparse.ArgumentTypeError(str(err)) from None + lo, hi = SIGMA_SCALE_BOUNDS + if not lo <= float(value) <= hi: + raise argparse.ArgumentTypeError(f"must lie in {lo}..{hi}, got {value}") + return float(value) def add_ded_weight_args(parser: argparse.ArgumentParser) -> None: @@ -419,20 +419,20 @@ def add_ded_weight_args(parser: argparse.ArgumentParser) -> None: 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. + choice here decides which one the headline products (``DELFWT``, validate-ded + correlations, model-phased difference columns) carry. """ - from torchref.maps.ded_weights import DEFAULT_SCHEME, SCHEMES + from torchref.cli._choices import DED_WEIGHT_SCHEMES, DEFAULT_DED_WEIGHT_SCHEME parser.add_argument( "--ded-weight", - choices=list(SCHEMES), - default=DEFAULT_SCHEME, + choices=list(DED_WEIGHT_SCHEMES), + default=DEFAULT_DED_WEIGHT_SCHEME, help="Per-reflection weight for difference coefficients: 'q' is the " "q-weight, a Wiener weight from a shell-free fit of the expected difference " "power that down-weights noisy reflections to no less than a tenth, " "'inverse_variance' is 1/sigma^2, 'none' is flat (default: " - f"{DEFAULT_SCHEME}). All weights are written as columns.", + f"{DEFAULT_DED_WEIGHT_SCHEME}). All weights are written as columns.", ) parser.add_argument( "--difference-gamma", diff --git a/torchref/cli/collection_difference_refine.py b/torchref/cli/collection_difference_refine.py index 34a8ef21..7faa1b2c 100644 --- a/torchref/cli/collection_difference_refine.py +++ b/torchref/cli/collection_difference_refine.py @@ -20,24 +20,28 @@ --fraction 0.37 -o output/ """ +# torch is imported inside the functions, after parse_args, so that --help does not +# wait for it; the annotations that name it stay strings. +from __future__ import annotations + import argparse import itertools import json import sys import warnings from pathlib import Path +from typing import TYPE_CHECKING -import torch - +from torchref.cli._choices import DEFAULT_DED_WEIGHT_SCHEME from torchref.cli._common import ( add_all_columns_arg, add_ded_weight_args, - add_dmin_arg, add_dual_model_args, add_general_args, add_metadata_args, add_outdir_arg, add_output_format_args, + add_resolution_args, add_weights_arg, apply_metadata_args, build_dual_column_names, @@ -52,15 +56,9 @@ validate_cif_files, validate_files, ) -from torchref.io.datasets.french_wilson import french_wilson_auto -from torchref.maps.ded_weights import ( - DEFAULT_SCHEME, - WEIGHT_COLUMNS, - DedWeightFallbackWarning, - all_ded_weights, - difference_snr, -) -from torchref.utils.serialization import convert_to_serializable + +if TYPE_CHECKING: + import torch configure_unbuffered_output() @@ -113,9 +111,9 @@ def setup_model_collection(pdb_dark, pdb_light, fractions, cif, d_min, device, v return mc -def setup_dataset_collection(sf_dark, sf_light, d_min, device, - column_names_dark=None, column_names_light=None): - """Load reflection data and create a DatasetCollection.""" +def setup_dataset_collection(sf_dark, sf_light, d_min, device, column_names_dark=None, + column_names_light=None, d_max=None): + """Load reflection data, cut to ``d_min``..``d_max`` (Å), and scale it jointly.""" from torchref import DatasetCollection data_dark = load_reflection_data( @@ -124,9 +122,9 @@ def setup_dataset_collection(sf_dark, sf_light, d_min, device, data_light = load_reflection_data( sf_light, device=device, column_names=column_names_light, ) - if d_min is not None: - data_dark.filter_by_resolution(d_min=d_min) - data_light.filter_by_resolution(d_min=d_min) + if d_min is not None or d_max is not None: + data_dark.filter_by_resolution(d_min=d_min, d_max=d_max) + data_light.filter_by_resolution(d_min=d_min, d_max=d_max) dc = DatasetCollection(device=device) dc.add_dataset("dark", data_dark) @@ -151,36 +149,21 @@ def setup_scaler(dataset_collection, model_collection, device, verbose=1): return scaler -def setup_dark_only(pdb_dark, dc, cif, d_min, device, verbose): - """Load the dark model alone and scale it against the dark data. +def setup_dark_scaler(model_dark, data_dark, device, verbose): + """Scale the dark model alone against the dark data; the difference map's phases. - 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``. + The difference amplitude ``|Fo_light| - |Fo_dark|`` is already on one scale + (:meth:`DatasetCollection.scale`) without reference to any model, so its phases need + the dark state only. Fitting them without the light model keeps the difference map + the same whether or not one is given. 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. """ from torchref import Scaler - model_dark = load_model( - pdb_dark, max_res=d_min, device=device, verbose=verbose, cif=cif, - ) - scaler = Scaler( - model_dark, dc["dark"], device=device, verbose=max(-1, verbose - 1), - ) + scaler = Scaler(model_dark, data_dark, device=device, verbose=max(-1, verbose - 1)) scaler.initialize().refine_lbfgs() - return model_dark, scaler + return scaler def compute_rfactors(model, data, scaler): @@ -192,6 +175,8 @@ def compute_rfactors(model, data, scaler): Takes either scaler kind; see the branch below. """ + import torch + from torchref.base.metrics.rfactor import rfactor_work_free with torch.no_grad(): @@ -344,6 +329,8 @@ def compute_bayes_extrapolated_amplitudes( per reflection. The variance holds the model phases and the fitted weight fixed, so it excludes uncertainty in the fit, phases and occupancy. """ + import torch + F_dark_phased = Fobs_dark * torch.exp(1j * phi_dark) F_light_phased = Fobs_light * torch.exp(1j * phi_mixed) z = F_dark_phased + (F_light_phased - F_dark_phased) / f @@ -369,10 +356,9 @@ def compute_bayes_extrapolated_amplitudes( return F_ext_bayes, var_ext_bayes, w -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 _two_moment_columns(mc, dc, mask, ctx, *, weights, diff_Fobs, Fobs_dark, sig_dark, + all_columns=False): + """Two-moment diagnostic columns, or an empty dict when the model is off. The observed light intensity carries a positive, phase-blind contamination ``sigma_alpha^2 |dF|^2`` from the spread of activation across crystals. Subtracting @@ -392,27 +378,27 @@ def _two_moment_columns(mc, dc, mask, fcalc_dark_full, fcalc_mixed_full, 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 + ctx : dict + The light-model intermediates of :func:`_phasing_columns`. + weights, diff_Fobs, 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. + dict + ``{label: (values, MTZ type letter)}``. """ import numpy as np + import torch - empty = ({}, {}) - if float(mc.sigma_alpha_sq) == 0.0: - return empty + from torchref.io.datasets.french_wilson import french_wilson_auto data_light = dc["light"] - if data_light.I is None: - return empty + if float(mc.sigma_alpha_sq) == 0.0 or data_light.I is None: + return {} + fcalc_dark_full, fcalc_mixed_full = ctx["fcalc_dark_full"], ctx["fcalc_mixed_full"] + F_obs_dark_phased, Fcalc_diff_amp = ctx["F_obs_dark_phased"], ctx["Fcalc_diff_amp"] with torch.no_grad(): # Full-size, so French-Wilson fits its prior to every working reflection. @@ -455,7 +441,7 @@ def _np(t): # 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) + ) * torch.exp(1j * ctx["phi_mixed"]) Fobs_diff_phased_corr = ( torch.abs(F_corr_phased - F_obs_dark_phased).detach().cpu().numpy() ) @@ -469,43 +455,25 @@ def _np(t): columns = { # dF_corr times the selected weight, on the dark phases PHDELWT. - "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", + "DELFWT_corr": (dF_corr * weights, "F"), + "Fo_light_corr": (F_corr, "F"), + "SIGFo_light_corr": (sig_F_corr, "Q"), + "dF_corr": (dF_corr, "F"), + "SIGdF_corr": (sig_dF_corr, "Q"), + "ddF": (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", + "Io_light": (I_light, "J"), + "SIGIo_light": (sig_I_light, "Q"), + "Ic_light_coh": (I_coherent, "J"), + "Ic_light_2mom": (I_two_moment, "J"), + "IVAR_ALPHA": (variance, "J"), + "W_2MOM": (w_two_moment, "W"), + "2mDFop-DFc_corr": (amp_2_corr, "F"), + "mDFop-DFc_corr": (amp_1_corr, "F"), }) - return columns, types + return columns def _difference_columns( @@ -522,6 +490,7 @@ def _difference_columns( phases_dark, diff_Fobs, sig_diff, + weights, weight_columns, kscale, ): @@ -530,9 +499,11 @@ def _difference_columns( ``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_Q``, ``W_InVa``; - 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_Q -cphi PHDELWT``. + correlates against. ``DELFWT`` is ``dF`` times ``weights``, the ``--ded-weight`` + scheme, so Coot opens ``DELFWT``/``PHDELWT`` as the difference map. One weight + column per registered scheme (``W_Q``, ``W_InVa``; MTZ type ``W``, mean one) sits + beside it, so any other weighting is reproducible from the file: + ``torchref.mtz2map -csf dF -cw W_InVa -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. @@ -541,6 +512,8 @@ def _difference_columns( 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. + + ``Fc_dark`` is the dark model's amplitude on the same fit as ``PHDELWT``. """ import numpy as np @@ -551,113 +524,90 @@ def _flags(data): 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, + return { + **{label: (hkl_np[:, i], "H") for i, label in enumerate("HKL")}, + "Fo_dark": (Fobs_dark, "F"), + "SIGFo_dark": (sig_dark, "Q"), + "Fo_light": (Fobs_light, "F"), + "SIGFo_light": (sig_light, "Q"), + "DELFWT": (diff_Fobs * weights, "F"), + "PHDELWT": (phases_dark, "P"), + "dF": (diff_Fobs, "F"), + "SIGdF": (sig_diff, "Q"), + **{name: (w, "W") for name, w in weight_columns.items()}, + "KSCALE": (kscale, "R"), + "Fc_dark": (Fcalc_dark, "F"), # 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", + "FreeR_flag_dark": (_flags(data_dark), "I"), + "FreeR_flag_light": (_flags(data_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. +def _phasing_columns(mc, scaler, hkl_all, mask, *, Fobs_dark_vals, Fobs_light_vals, + weights, all_columns=False): + """The light data's 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. + Everything here needs the light state's model and comes from the joint ``scaler``. + ``Fc_light``/``PHFc_light`` are the scaled amplitude and phase of the mixed model, + the model of the light data as ``Fc_dark`` is of the dark data. 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 + ``PHdFc``. These are a *different object* from the plain difference Fourier ``dF`` on ``PHDELWT``, not a refinement of it: the light state's model phases enter the observed amplitude, so they are model-biased where ``dF`` 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 + Returns ``(columns, ctx)``. ``ctx`` carries the joint fit's intermediates the extrapolation and two-moment layers need, so nothing is computed twice. """ - mixed_model = mc["light"] + import torch + dark_model, mixed_model = mc.dark_model, mc["light"] with torch.no_grad(): - fcalc_mixed_full = scaler.forward_mixed( - mixed_model(hkl_all), mixed_model.fractions + fcalc_dark_full, fcalc_mixed_full = ( + scaler.forward_mixed(model(hkl_all), model.fractions) + for model in (dark_model, mixed_model) ) - fcalc_mixed = fcalc_mixed_full[mask] + fcalc_dark, fcalc_mixed = fcalc_dark_full[mask], fcalc_mixed_full[mask] fcalc_diff = fcalc_mixed - fcalc_dark - phi_mixed = torch.angle(fcalc_mixed) + phi_dark, phi_mixed = torch.angle(fcalc_dark), 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) - Fcalc_light = torch.abs(fcalc_mixed).detach().cpu().numpy() - Fcalc_diff_amp = torch.abs(fcalc_diff).detach().cpu().numpy() + def _np(t): + return t.detach().cpu().numpy() + Fcalc_light = _np(torch.abs(fcalc_mixed)) + Fcalc_diff_amp = _np(torch.abs(fcalc_diff)) columns = { - "FC": Fcalc_light, - "PHIC": phi_mixed.detach().rad2deg().cpu().numpy(), + "Fc_light": (Fcalc_light, "F"), + "PHFc_light": (_np(phi_mixed.rad2deg()), "P"), } - 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", + Fobs_diff_phased = _np(torch.abs(F_obs_light_phased - F_obs_dark_phased)) + columns.update({ + "2mDFop-DFc": ((2 * Fobs_diff_phased - Fcalc_diff_amp) * weights, "F"), + "mDFop-DFc": ((Fobs_diff_phased - Fcalc_diff_amp) * weights, "F"), + "PHdFc": (_np(torch.angle(fcalc_diff).rad2deg()), "P"), + "dFc": (Fcalc_light - _np(torch.abs(fcalc_dark)), "F"), + # The real modulus of the complex vector difference. + "dFc_phased": (Fcalc_diff_amp, "F"), }) ctx = { + "fcalc_dark_full": fcalc_dark_full, "fcalc_mixed_full": fcalc_mixed_full, + "phi_dark": phi_dark, "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 + return columns, ctx def _extrapolation_columns( @@ -669,14 +619,12 @@ def _extrapolation_columns( Fobs_light_vals, sig_dark_vals, sig_light_vals, - phi_dark, ctx, rfree_flags_masked, snr_est, - all_columns=False, verbose=1, ): - """Extrapolated light-state amplitudes and the map to view them in. + """Extrapolated light-state amplitudes and the maps to view them in. ``snr_est`` is the :class:`~torchref.maps.ded_weights.DifferenceSNR` of the light-minus-dark differences on the same reflections; ``None`` (the fit failed) @@ -684,23 +632,22 @@ def _extrapolation_columns( Three constructions of the same quantity, all needing the light model: - ``FEXT`` (default, shrunk) + ``FEXT`` (shrunk), map ``FWT``/``PHWT`` The phase-aware amplitude shrunk toward ``Fo_dark`` by the signal fraction of each reflection's difference (:func:`compute_bayes_extrapolated_amplitudes`), which quiets the reflections where the extrapolation is noisiest. - ``FEXT_PHASED`` (``all_columns``) + ``FEXT_PHASED``, map ``2FEXT_PHASED-Fc``/``PHFEXT_PHASED`` 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. + ``FEXT_SCALAR``, map ``2FEXT_SCALAR-Fc``/``PHFEXT_SCALAR`` + The classic ``(Fo_light - w_dark * Fo_dark) / w_light`` on amplitudes only. - ``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. + Each map is ``2 * FEXT - Fc`` with ``F_calc`` rescaled against *its own* amplitudes + -- the three sets differ in overall scale -- so each costs one LBFGS scale fit, and + each takes the phase of its own fit: the scaler contributes one through ``f_sol``, + so pairing these coefficients with another fit's phase would be wrong. """ + import torch + from torchref import ReflectionData, Scaler from torchref.base.metrics.rfactor import rfactor_work_free @@ -719,10 +666,7 @@ def _fit(amp, sig): 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 + # The propagated sigma of the unshrunk extrapolations. sig_light_extra = torch.sqrt( sig_light_vals**2 + w_dark**2 * sig_dark_vals**2 ) / w_light @@ -736,102 +680,70 @@ def _fit(amp, sig): Fobs_dark_vals, Fobs_light_vals, sig_dark_vals, - phi_dark, + ctx["phi_dark"], ctx["phi_mixed"], w_light, snr=snr, sig_light=sig_light_vals, ) - sig_ext_bayes = torch.sqrt(var_ext_bayes) - - data_bayes, F_calc_bayes = _fit(F_ext_bayes_amp, sig_ext_bayes) - amp_calc_bayes = torch.abs(F_calc_bayes) + F_light_extra = ( + ctx["F_obs_light_phased"] - w_dark * ctx["F_obs_dark_phased"] + ) / w_light + F_light_scalar = (Fobs_light_vals - w_dark * Fobs_dark_vals) / w_light + constructions = ( + ("", "Bayes", F_ext_bayes_amp, torch.sqrt(var_ext_bayes)), + ("_PHASED", "Phase-aware", torch.abs(F_light_extra), sig_light_extra), + ("_SCALAR", "Scalar", F_light_scalar, sig_light_extra), + ) def _np(t): return t.detach().cpu().numpy() - 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), + columns = {} + for suffix, name, amp, sig in constructions: + data, F_calc = _fit(amp, sig) + amp_calc = torch.abs(F_calc) + # The shrunk map keeps the labels Coot opens by default. + map_label, phase_label = ( + (f"2FEXT{suffix}-Fc", f"PHFEXT{suffix}") if suffix else ("FWT", "PHWT") ) + columns.update({ + f"FEXT{suffix}": (_np(amp), "F"), + f"SIGFEXT{suffix}": (_np(sig), "Q"), + map_label: (_np(2 * amp - amp_calc), "F"), + phase_label: (_np(torch.angle(F_calc).rad2deg()), "P"), + }) + if verbose > 0: + print(f" {name} extrapolation R:", rfactor_work_free(data, amp_calc)) + if verbose > 0: print( f" Shrinkage: SNR from {source} differences, " f"mean w(h) = {w_shrinkage.mean().item():.3f}, " f"min w(h) = {w_shrinkage.min().item():.3g}" ) - 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 = { "shrinkage_source": source, "w_shrinkage_mean": float(w_shrinkage.mean().item()), "w_shrinkage_min": float(w_shrinkage.min().item()), } - return columns, types, diagnostics - + return columns, diagnostics -_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_Q, W_InVa", - "light_model": ( - "light_model: FC/PHIC, amplitude and phase of the mixed dark+light model" - ), + "difference": "difference: DELFWT/PHDELWT = weighted dF; dF, SIGdF, W_Q, W_InVa", + "light_model": "light_model: Fc_light/PHFc_light, the mixed dark+light model", "extrapolated_light": ( - "extrapolated_light: FWT/PHWT = 2*FEXT - Fc, the extrapolated light map" + "extrapolated_light: FWT/PHWT = 2*FEXT-Fc (shrunk), FEXT_PHASED, FEXT_SCALAR" ), "two_moment": "two_moment: difference columns with the two-moment correction", } -def _annotate_mtz(filename, datasets): +def _annotate_mtz(filename, datasets, program="difference-refine"): """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 @@ -847,6 +759,8 @@ def _annotate_mtz(filename, datasets): datasets : dict[str, str] Column label to dataset name. Unlisted columns, Miller indices included, stay in ``observed``. + program : str + The ``torchref.`` named in the first history line. """ import gemmi @@ -873,7 +787,7 @@ def _annotate_mtz(filename, datasets): col.dataset_id = ids[datasets[col.label]] mtz.history = [ - "torchref difference-refine map coefficients, grouped by dataset:", + f"torchref {program} map coefficients, grouped by dataset:", *(_MTZ_DATASET_HISTORY[name] for name in ids), ] mtz.write_to_file(filename) @@ -886,18 +800,20 @@ def write_results_mtz( filename, *, mc=None, + joint_scaler=None, all_columns=False, verbose=1, - ded_weight=DEFAULT_SCHEME, + ded_weight=DEFAULT_DED_WEIGHT_SCHEME, difference_config=None, + program="difference-refine", ): """Write the difference map, and map coefficients when a light model is given. The default output is the difference map on the dark model's phases - (:func:`_difference_columns`). It 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 runs no scale fit beyond the one that produced ``scaler``. Given ``mc``, the - light-state, extrapolated and two-moment layers follow (:func:`_phasing_columns`, + (:func:`_difference_columns`). It needs no light-state model, and takes its phases + from the dark-only ``scaler`` either way, so the difference map does not depend on + whether a light model is given. Given ``mc`` and ``joint_scaler``, the light-state, + extrapolated and two-moment layers follow (:func:`_phasing_columns`, :func:`_extrapolation_columns`), grouped into MTZ datasets by :func:`_annotate_mtz`. Coot auto-opens ``FWT``/``PHWT``, which here is the *extrapolated light-state* map, not a ``2mFo-DFc``. @@ -906,26 +822,29 @@ def write_results_mtz( ---------- dc : DatasetCollection Dark and light data, already inter-scaled by :meth:`DatasetCollection.scale`. - dark_model : Model or _SharedMixedModel + dark_model : Model 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. + scaler : Scaler + Scales ``dark_model`` against the dark data alone (:func:`setup_dark_scaler`). filename : str Output MTZ path. mc : ModelCollection, optional The dark+light collection. Absent means difference map only. + joint_scaler : CollectionScaler, optional + Scales ``mc`` against both datasets; required with ``mc``. all_columns : bool Also write the alternatives within each layer. verbose : int Verbosity level. ded_weight : str, optional - Weight scheme for the model-phased and two-moment difference columns; one of - :data:`torchref.maps.ded_weights.SCHEMES`. + Weight scheme ``DELFWT`` and the model-phased and two-moment difference columns + carry; one of :data:`torchref.maps.ded_weights.SCHEMES`. difference_config : DifferencePowerConfig, optional Its ``gamma`` fixes the ``F_dark`` exponent of the ``q`` scheme and its ``sigma_scale`` the scale on the reported sigmas of the SNR the ``q`` weights and the extrapolation share. + program : str + The tool named in the MTZ history, ``torchref.``. Returns ------- @@ -936,8 +855,16 @@ def write_results_mtz( extrapolated map is over-shrunk. """ import reciprocalspaceship as rs + import torch + + from torchref.maps.ded_weights import ( + WEIGHT_COLUMNS, + DedWeightFallbackWarning, + all_ded_weights, + difference_snr, + ) - data_dark = dc[mc.dark_key] if mc is not None else dc["dark"] + data_dark = dc["dark"] data_light = dc["light"] hkl_all = data_dark.hkl @@ -954,19 +881,10 @@ def write_results_mtz( 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. + # Fcalc on the full HKL list then masked, because the scaler was fitted on the + # full dataset. 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] - - phi_dark = torch.angle(fcalc_dark) + fcalc_dark = scaler(dark_model(hkl_all))[mask] hkl_np = hkl.cpu().numpy() Fobs_dark = Fobs_dark_vals.cpu().numpy() @@ -974,7 +892,7 @@ def write_results_mtz( sig_dark = sig_dark_vals.cpu().numpy() sig_light = sig_light_vals.cpu().numpy() Fcalc_dark = torch.abs(fcalc_dark).detach().cpu().numpy() - phases_dark = phi_dark.detach().rad2deg().cpu().numpy() + phases_dark = torch.angle(fcalc_dark).rad2deg().cpu().numpy() diff_t = Fobs_light_vals - Fobs_dark_vals sig_diff_t = torch.sqrt(sig_dark_vals**2 + sig_light_vals**2) @@ -1042,7 +960,7 @@ def write_results_mtz( f"before normalisation" ) - columns, types = _difference_columns( + columns = _difference_columns( data_dark, data_light, mask, @@ -1055,23 +973,24 @@ def write_results_mtz( phases_dark=phases_dark, diff_Fobs=diff_Fobs, sig_diff=sig_diff, + weights=weights, weight_columns=weight_columns, kscale=kscale, ) - datasets = {name: "difference" for name in _DIFFERENCE_DATASET_COLUMNS} + datasets = dict.fromkeys( + ["DELFWT", "PHDELWT", "dF", "SIGdF", *weight_columns, "KSCALE"], "difference" + ) 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, + phase_cols, ctx = _phasing_columns( + mc, joint_scaler, hkl_all, mask, + Fobs_dark_vals=Fobs_dark_vals, Fobs_light_vals=Fobs_light_vals, + 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( + ext_cols, ext_diagnostics = _extrapolation_columns( mc, dc, hkl, @@ -1079,46 +998,35 @@ def write_results_mtz( 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, snr_est=None if isinstance(snr_est, ValueError) else snr_est, - all_columns=all_columns, verbose=verbose, ) 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, + + tm_cols = _two_moment_columns( + mc, dc, mask, ctx, weights=weights, diff_Fobs=diff_Fobs, + Fobs_dark=Fobs_dark, sig_dark=sig_dark, all_columns=all_columns, ) columns.update(tm_cols) datasets.update(dict.fromkeys(tm_cols, "two_moment")) - types.update(tm_types) df = rs.DataSet( - columns, + {name: values for name, (values, _) in columns.items()}, cell=data_dark.cell.data.cpu().tolist(), spacegroup=data_dark.spacegroup.hm, ) # 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(): + for name, (_, letter) in columns.items(): df[name] = df[name].astype(letter) df = df.infer_mtz_dtypes() df.set_index(["H", "K", "L"], inplace=True) df.write_mtz(filename) - _annotate_mtz(filename, datasets) + _annotate_mtz(filename, datasets, program=program) if verbose > 0: print(f" Results MTZ written to {filename} ({len(columns)} columns)") @@ -1142,6 +1050,8 @@ def optimize_lbfgs(state, parameters, max_iter, nsteps, n_clean, verbose): optimizer's intent. The dark model effectively becomes a frozen reference at the autograd level for the duration of each step. """ + import torch + parameters = list(parameters) def _make_optimizer(): @@ -1264,11 +1174,16 @@ def main(): ) res = parser.add_argument_group("Resolution") - add_dmin_arg(res) + add_resolution_args(res) add_general_args(parser) args = parser.parse_args() + + import torch + + from torchref.utils.serialization import convert_to_serializable + register_timing() # --- Parse fractions --- @@ -1345,8 +1260,6 @@ def main(): outdir.mkdir(parents=True, exist_ok=True) device = parse_device_str(args.device) - d_min = args.dmin if args.dmin is not None else 1.0 - # --- Header --- if args.verbose > 0: print("=" * 72) @@ -1379,6 +1292,23 @@ def main(): print() sys.stdout.flush() + # --- Setup 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, d_max=args.dmax, + ) + data_dark = dc["dark"] + data_light = dc["light"] + # The models' FFT grid has to reach the data, so without --dmin it follows them. + d_min = args.dmin or min(data_dark.d_min, data_light.d_min) + # --- Setup models --- if args.verbose > 0: print("Setting up models...") @@ -1398,21 +1328,6 @@ def main(): else: mixed.freeze_fractions() - # --- Setup 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, - ) - data_dark = dc["dark"] - data_light = dc["light"] - if args.verbose > 0: # Compare R-free flags between datasets if data_dark.rfree_flags is not None and data_light.rfree_flags is not None: @@ -1729,10 +1644,11 @@ def _mtz_to_cif(mtz_path, cif_path): map_diagnostics = write_results_mtz( dc, - mc.dark_model, - scaler, + model_dark, + setup_dark_scaler(model_dark, data_dark, device, args.verbose), diff_mtz_out, mc=mc, + joint_scaler=scaler, all_columns=args.all_columns, verbose=args.verbose, ded_weight=args.ded_weight, @@ -1749,6 +1665,7 @@ def _mtz_to_cif(mtz_path, cif_path): "fractions": fractions, "cif": args.cif, "dmin": args.dmin, + "dmax": args.dmax, }, "dataset_scaling": dc.scaling_metrics, "parameters": { diff --git a/torchref/cli/difference_map.py b/torchref/cli/difference_map.py index 612f2b7d..068a42f5 100644 --- a/torchref/cli/difference_map.py +++ b/torchref/cli/difference_map.py @@ -3,20 +3,23 @@ """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_Q`` (q-weight, the default) and ``W_InVa`` (inverse variance), and the -observed-to-model scale ``KSCALE``. Build the map with -``torchref.mtz2map -csf dF -cw W_Q -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. +models are used as-is. Writes ``difference_map.mtz`` and ``summary.json`` (inputs, +dataset scaling, R-factors, weight and shrinkage fits) to the ``-o`` directory. + +The default output is the difference map: ``DELFWT``/``PHDELWT``, the amplitude +difference ``dF = |Fo_light| - |Fo_dark|`` times the ``--ded-weight`` weight on the +**dark** model's phases, which Coot opens directly. ``dF``/``SIGdF``, the weights of +every registered scheme (``W_Q``, ``W_InVa``) and the observed-to-model scale ``KSCALE`` +sit beside it, so another weighting is ``torchref.mtz2map -csf dF -cw W_InVa -cphi +PHDELWT`` (``--units electrons`` for e/A^3). The phases come from a fit of the dark +model alone, so the difference map is the same with or without ``-lm``. It is +deliberately not a *phased* difference map: putting the light state's model phases +into the observed amplitude biases the map toward the model the experiment is testing. + +Optionally, ``-lm`` with ``--fraction`` adds the light data's model ``Fc_light``/ +``PHFc_light`` and three extrapolated light-state maps: the shrunk ``FWT``/``PHWT``, the +phase-aware ``2FEXT_PHASED-Fc`` and the classic ``2FEXT_SCALAR-Fc``. ``--all-columns`` +adds the model-phased difference columns. Examples -------- @@ -24,31 +27,31 @@ # difference map only -- no light model, no fraction needed torchref.difference-map \\ - -dm dark.pdb -dsf dark.mtz -lsf light.mtz -o results.mtz + -dm dark.pdb -dsf dark.mtz -lsf light.mtz -o results/ 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 + --fraction 0.37 --dmin 1.7 --cif ligand.cif -o results/ """ import argparse +import json 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, + add_outdir_arg, + add_resolution_args, build_dual_column_names, configure_unbuffered_output, - register_timing, - parse_device_str, difference_config_from_args, + load_model, + parse_device_str, + register_timing, validate_cif_files, validate_files, ) @@ -59,6 +62,7 @@ def main(): """Entry point for ``torchref.difference-map``; returns the exit code.""" parser = argparse.ArgumentParser( + prog="torchref.difference-map", description="Compute difference and extrapolated map coefficients " "(no refinement).", formatter_class=argparse.RawDescriptionHelpFormatter, @@ -67,53 +71,37 @@ def main(): # 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 + -dsf dark.mtz -lsf light.mtz -o results/ + # adds the light data's model and the extrapolated light-state maps 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 + --fraction 0.37 --dmin 1.7 --dmax 20 -o results/ """, ) 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_outdir_arg( + output, help="Output directory for difference_map.mtz and summary.json" + ) add_all_columns_arg(output) add_ded_weight_args(output) res = parser.add_argument_group("Resolution") - add_dmin_arg(res) + add_resolution_args(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." - ) + if has_light_model != (args.fraction is not None): + parser.error("-lm/--light-model and --fraction go together, or not at all") + if args.all_columns and not has_light_model: + parser.error("--all-columns needs -lm/--light-model: its columns use the model") to_check = [ (args.dark_model, "dark model"), @@ -122,16 +110,30 @@ def main(): ] if has_light_model: to_check.insert(1, (args.light_model, "light model")) - if validate_files(to_check): + if validate_files(to_check) or validate_cif_files(args.cif): return 1 - if validate_cif_files(args.cif): - return 1 + col_dark, col_light = build_dual_column_names(args) - out_path = Path(args.output) - out_path.parent.mkdir(parents=True, exist_ok=True) + import torch + from torchref.cli.collection_difference_refine import ( + compute_rfactors, + setup_dark_scaler, + setup_dataset_collection, + setup_model_collection, + setup_scaler, + write_results_mtz, + ) + from torchref.utils.serialization import convert_to_serializable + + register_timing() + + outdir = Path(args.outdir) + outdir.mkdir(parents=True, exist_ok=True) + mtz_path, summary_path = outdir / "difference_map.mtz", outdir / "summary.json" device = parse_device_str(args.device) + fractions = [1.0 - args.fraction, args.fraction] if has_light_model else None if args.verbose > 0: print("=" * 72) @@ -146,94 +148,93 @@ def main(): 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"Output: {outdir}") print(f"Device: {device}") - if args.dmin: - print(f"Resolution cutoff: {args.dmin:.2f} A") + if args.dmin or args.dmax: + print(f"Resolution: {args.dmax or 'all'} - {args.dmin or 'all'} 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, + column_names_dark=col_dark, column_names_light=col_light, d_max=args.dmax, ) + # The models' FFT grid has to reach the data, so without --dmin it follows them. + d_min = args.dmin or min(dc["dark"].d_min, dc["light"].d_min) if args.verbose > 0: - print("Setting up models...") + print("Setting up models and scalers...") sys.stdout.flush() + mc = joint_scaler = None 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() + joint_scaler = setup_scaler(dc, mc, device, args.verbose) + dark_model = mc.base_models[0] else: - mc = None - dark_model, scaler = setup_dark_only( - args.dark_model, dc, args.cif, d_min, device, args.verbose, + dark_model = load_model( + args.dark_model, max_res=d_min, device=device, verbose=args.verbose, + cif=args.cif, ) + scaler = setup_dark_scaler(dark_model, dc["dark"], device, args.verbose) + + fits = {"dark": (dark_model, dc["dark"], scaler)} + if has_light_model: + fits["mixed"] = (mc["light"], dc["light"], joint_scaler) + r_factors = {} + for name, (model, data, fit) in fits.items(): + r_work, r_free = compute_rfactors(model, data, fit) + r_factors[name] = {"r_work": r_work, "r_free": r_free} 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() + print(f" R-factor ({name}): R_work={r_work:.4f} R_free={r_free:.4f}") if args.verbose > 0: print("Computing map coefficients...") sys.stdout.flush() with torch.no_grad(): - write_results_mtz( + map_diagnostics = write_results_mtz( dc, dark_model, scaler, - str(out_path), + str(mtz_path), mc=mc, + joint_scaler=joint_scaler, all_columns=args.all_columns, verbose=args.verbose, ded_weight=args.ded_weight, difference_config=difference_config_from_args(args), + program="difference-map", ) + summary = { + "input": { + "dark_model": args.dark_model, + "light_model": args.light_model, + "dark_structure_factor": args.dark_structure_factor, + "light_structure_factor": args.light_structure_factor, + "fractions": fractions, + "cif": args.cif, + "dmin": args.dmin, + "dmax": args.dmax, + "model_grid_dmin": d_min, + }, + "dataset_scaling": dc.scaling_metrics, + "results": {"r_factors": r_factors, **map_diagnostics}, + "output_files": {"mtz": str(mtz_path), "summary": str(summary_path)}, + } + with open(summary_path, "w") as f: + json.dump(convert_to_serializable(summary), f, indent=2) + if args.verbose > 0: - print() + print(f" Summary written to {summary_path}") print("Done.") sys.stdout.flush() diff --git a/torchref/cli/mtz2map.py b/torchref/cli/mtz2map.py index 713fa940..30a24865 100644 --- a/torchref/cli/mtz2map.py +++ b/torchref/cli/mtz2map.py @@ -15,16 +15,12 @@ import argparse import sys -import numpy as np -import torch - from torchref.cli._common import ( add_general_args, add_resolution_args, register_timing, parse_device_str, ) -from torchref.config import get_float_dtype, get_int_dtype def main(): @@ -120,6 +116,11 @@ def main(): args = parser.parse_args() + import numpy as np + import torch + + from torchref.config import get_float_dtype, get_int_dtype + register_timing() # --- Device --- diff --git a/torchref/cli/refine.py b/torchref/cli/refine.py index 0b775150..21e45db4 100644 --- a/torchref/cli/refine.py +++ b/torchref/cli/refine.py @@ -16,8 +16,7 @@ import warnings from pathlib import Path -import torch - +from torchref.cli._choices import DEFAULT_SCALE_TARGET, SCALE_TARGETS, XRAY_TARGET_DOCS from torchref.cli._common import ( add_adp_mode_arg, add_dmin_arg, @@ -37,15 +36,9 @@ validate_files, write_refinement_outputs, ) -from torchref.refinement.targets.xray._specs import XRAY_TARGETS -from torchref.scaling.scaler_base import DEFAULT_SCALE_TARGET, SCALE_TARGETS -from torchref.utils.serialization import convert_to_serializable configure_unbuffered_output() -# Import stats module early to patch json with StatEntry encoder -import torchref.utils.stats # noqa: F401,E402 - def _sigma_a_kwargs(args) -> dict: """Only the σ_A estimator knobs the user actually set. @@ -151,15 +144,9 @@ def main(): "--xray-mode", type=str, default="ml", - choices=list(XRAY_TARGETS.names), - # Generated from the table, not hand-written. The hand-written version drifted - # badly: it advertised 'bhattacharyya' (not in `choices`, so argparse rejected it) - # and 'gaussian' (an alias), while never mentioning ml_noalpha, nll_beta or nll -- - # three of the modes it actually accepted. `spec.doc` promised it was "surfaced in - # --help" and had no reader at all. - help="X-ray target function. " + " ".join( - f"'{spec.name}': {spec.doc}" for spec in XRAY_TARGETS.specs - ), + choices=list(XRAY_TARGET_DOCS), + help="X-ray target function. " + + " ".join(f"'{name}': {doc}" for name, doc in XRAY_TARGET_DOCS.items()), ) refine_group.add_argument( "--scale-target", @@ -235,6 +222,11 @@ def main(): args = parser.parse_args() + import torch + + import torchref.utils.stats # noqa: F401 (makes json serialise StatEntry) + from torchref.utils.serialization import convert_to_serializable + register_timing() # Parse weights diff --git a/torchref/cli/simulate_noisy_data.py b/torchref/cli/simulate_noisy_data.py index d25f15d9..ab583774 100644 --- a/torchref/cli/simulate_noisy_data.py +++ b/torchref/cli/simulate_noisy_data.py @@ -40,9 +40,7 @@ import argparse import sys from pathlib import Path - -import pandas as pd -import torch +from typing import TYPE_CHECKING from torchref.cli._common import ( add_device_arg, @@ -51,8 +49,9 @@ load_model, parse_device_str, ) -from torchref.io import mtz -from torchref.io.datasets import FcalcDataset + +if TYPE_CHECKING: + from torchref.io.datasets import FcalcDataset def _build_parser() -> argparse.ArgumentParser: @@ -158,8 +157,11 @@ def main(): def _run_reference_mode(args, model, device) -> int: """Reference-driven: scale model to CrystFEL hkl and graft sigmas.""" + import torch + from torchref import ReflectionData from torchref.base.reciprocal import get_d_spacing + from torchref.io.datasets import FcalcDataset from torchref.scaling import Scaler if args.verbose: @@ -219,6 +221,8 @@ def _run_reference_mode(args, model, device) -> int: def _run_parametric_mode(args, model, device) -> int: """No reference: build HKL from cell+resolution, use three-term model.""" + from torchref.io.datasets import FcalcDataset + if args.verbose: print( f"Generating HKL to d_min={args.d_min} A" @@ -248,8 +252,12 @@ def _run_parametric_mode(args, model, device) -> int: return 0 -def _write_output(args, noisy: FcalcDataset) -> None: +def _write_output(args, noisy: "FcalcDataset") -> None: """Write noisy amplitudes or intensities and their uncertainties to MTZ.""" + import pandas as pd + + from torchref.io import mtz + hkl_np = noisy.hkl.cpu().numpy() columns = { "H": hkl_np[:, 0], diff --git a/torchref/cli/strip_altlocs.py b/torchref/cli/strip_altlocs.py index 878000d2..1faa27dd 100644 --- a/torchref/cli/strip_altlocs.py +++ b/torchref/cli/strip_altlocs.py @@ -18,8 +18,6 @@ import argparse import sys -from torchref import Model - def main(): """Entry point for ``torchref.strip-altlocs``; returns the exit code.""" @@ -32,6 +30,8 @@ def main(): parser.add_argument("output", help="Output PDB file") args = parser.parse_args() + from torchref import Model + # Pure bookkeeping on the atom table: no reason to start an accelerator. model = Model(device="cpu", verbose=0).load_pdb(args.input) result = model.strip_altlocs() diff --git a/torchref/cli/validate_ded.py b/torchref/cli/validate_ded.py index 57be77b8..d92c9b0f 100644 --- a/torchref/cli/validate_ded.py +++ b/torchref/cli/validate_ded.py @@ -27,9 +27,7 @@ import warnings from pathlib import Path -import numpy as np -import torch - +from torchref.cli._choices import DEFAULT_DED_WEIGHT_SCHEME from torchref.cli._common import ( add_ded_weight_args, add_dmin_arg, @@ -47,14 +45,6 @@ validate_cif_files, validate_files, ) -from torchref.config import get_int_dtype -from torchref.maps.ded_weights import ( - DEFAULT_SCHEME, - DedWeightFallbackWarning, - all_ded_weights, -) -from torchref.utils.serialization import convert_to_serializable - configure_unbuffered_output() @@ -70,6 +60,8 @@ def build_atom_mask(selection_xyz, real_space_grid, cell, mask_radius, device): ``real_space_grid`` comes from ``get_real_grid()`` and ``cell`` may be a ``Cell`` or the raw parameter tensor. """ + import torch + from torchref.base.coordinates.transforms_torch import ( get_fractional_matrix, ) @@ -78,6 +70,7 @@ def build_atom_mask(selection_xyz, real_space_grid, cell, mask_radius, device): ) from torchref.base.electron_density.solvent_mask import add_to_solvent_mask from torchref.base.electron_density.voxel_utils import find_relevant_voxels + from torchref.config import get_int_dtype grid_shape = real_space_grid.shape[:3] cell_t = cell.data if hasattr(cell, "data") else cell @@ -106,6 +99,8 @@ def build_atom_mask(selection_xyz, real_space_grid, cell, mask_radius, device): def compute_correlation(map1, map2, mask): """Pearson correlation coefficient between two maps within a mask.""" + import torch + v1 = map1[mask] v2 = map2[mask] v1 = v1 - v1.mean() @@ -120,6 +115,8 @@ def compute_map_from_coefficients(amplitudes, phases_rad, hkl_p1, gridsize): Takes ``amplitudes`` ``(N,)`` (which may be signed), ``phases_rad`` ``(N,)`` and P1-expanded ``hkl_p1`` ``(N, 3)``. """ + import torch + from torchref.base.reciprocal.grid_operations import place_on_grid coefficients = amplitudes * torch.exp(1j * phases_rad) @@ -206,7 +203,7 @@ def setup_ded_context( col_light=None, n_bins=20, verbose=0, - ded_weight=DEFAULT_SCHEME, + ded_weight=DEFAULT_DED_WEIGHT_SCHEME, difference_config=None, ): """Load reflection data and prepare the model-independent side of DED validation. @@ -239,9 +236,12 @@ def setup_ded_context( orig_idx, phase_shifts, w_dfo_p1, weights_p1, work_mask, free_mask, n_bins. """ import gemmi + import numpy as np + import torch from torchref import DatasetCollection from torchref.config import get_float_dtype, normalize_device + from torchref.maps.ded_weights import all_ded_weights from torchref.symmetry import Cell, SpaceGroup device = normalize_device(device) @@ -419,6 +419,8 @@ def compute_ded_maps( resolution_bins, reciprocal_cc_overall, reciprocal_cc_work, reciprocal_cc_free, w_delta_fcalc_asu. """ + import torch + from torchref.base.fourier.grid import get_real_grid from torchref.cli.collection_difference_refine import ( setup_scaler as setup_collection_scaler, @@ -626,7 +628,9 @@ def run_validation(args): from torchref.cli.collection_difference_refine import ( setup_scaler as setup_collection_scaler, ) + from torchref.maps.ded_weights import DedWeightFallbackWarning from torchref.model.model_collection import ModelCollection + from torchref.utils.serialization import convert_to_serializable device = parse_device_str(args.device) outdir = Path(args.outdir)