From 5d91ac622e91fa151b57e9298fbb7d28abd5078c Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 21 Oct 2026 12:56:14 +0000 Subject: [PATCH 1/3] Answer CLI --help without importing torch; tighten difference-map I/O import torchref no longer imports torch: the top-level re-exports and the subpackages load on first access through a module __getattr__. Every CLI module defers its torch-dependent imports until after parse_args, and the choice lists --help shows come from torchref.cli._choices, a torch-free mirror pinned to its sources by a test. --help drops from ~4 s to ~0.1 s. torchref.difference-map: - refuses --all-columns without -lm, and a non-default --ded-weight unless -lm --all-columns writes the model-phased columns it weights - requires -o to name a .mtz file, checked before any data is read - writes _summary.json beside the MTZ (inputs, dataset scaling, R-factors, q-weight and shrinkage diagnostics), as difference-refine does - sizes the model FFT grid from the data's d_min when --dmin is not given - names itself, not difference-refine, in the MTZ history Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014aviwqFcZpxP3HWNy1ZruY --- docs/changelog.rst | 4 + docs/user_guide/cli.rst | 10 +- tests/integration/test_cli_ded_weights.py | 12 ++ tests/unit/cli/test_choices.py | 29 +++ tests/unit/cli/test_difference_map_args.py | 67 +++++++ tests/unit/cli/test_entry_points.py | 53 +++++- torchref/__init__.py | 109 ++++++----- torchref/cli/_choices.py | 52 ++++++ torchref/cli/_common.py | 5 +- torchref/cli/collection_difference_refine.py | 67 ++++--- torchref/cli/difference_map.py | 179 +++++++++++++------ torchref/cli/mtz2map.py | 9 +- torchref/cli/refine.py | 30 ++-- torchref/cli/simulate_noisy_data.py | 20 ++- torchref/cli/strip_altlocs.py | 4 +- torchref/cli/validate_ded.py | 28 +-- 16 files changed, 518 insertions(+), 160 deletions(-) create mode 100644 tests/unit/cli/test_choices.py create mode 100644 tests/unit/cli/test_difference_map_args.py create mode 100644 torchref/cli/_choices.py diff --git a/docs/changelog.rst b/docs/changelog.rst index 615ec9aa..5be35f86 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`` refuses options it would ignore (``--all-columns`` or a non-default ``--ded-weight`` without ``-lm``, a ``-o`` that is not a ``.mtz`` file) and writes ``_summary.json`` beside the MTZ - 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 +- ``torchref.difference-map`` without ``--dmin`` sizes the model's FFT grid from the data's resolution 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..e38f323a 100644 --- a/docs/user_guide/cli.rst +++ b/docs/user_guide/cli.rst @@ -267,8 +267,16 @@ 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. +``-o`` names the MTZ file (it must end in ``.mtz``), and ``_summary.json`` is +written beside it with the inputs, the dataset scaling, the R-factors and the fits +behind the ``q`` weights and the extrapolation shrinkage. Options that would change +nothing in the output are refused rather than ignored: ``--all-columns`` needs ``-lm``, +and ``--ded-weight`` only weights the model-phased columns that ``-lm`` with +``--all-columns`` adds -- the difference map itself carries every scheme's weight as a +column, chosen with ``mtz2map -cw``. + **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 +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; diff --git a/tests/integration/test_cli_ded_weights.py b/tests/integration/test_cli_ded_weights.py index 3c99c36b..938f8ebc 100644 --- a/tests/integration/test_cli_ded_weights.py +++ b/tests/integration/test_cli_ded_weights.py @@ -151,6 +151,18 @@ def test_difference_map_writes_df_weights_and_scale(diff_mtz): assert (w_q > 0).all() +def test_difference_map_writes_a_summary_beside_the_mtz(diff_mtz): + import gemmi + + summary = json.loads(diff_mtz.with_name("diff_summary.json").read_text()) + assert summary["output_files"]["mtz"] == str(diff_mtz) + assert set(summary["results"]["r_factors"]) == {"dark"} + assert summary["results"]["ded_weights"]["scheme"] == "q" + assert summary["input"]["model_grid_dmin"] == 2.2 + history = gemmi.read_mtz_file(str(diff_mtz)).history + assert history[0].startswith("torchref difference-map ") + + def test_mtz2map_builds_the_weighted_and_electron_maps(project_root, pair, diff_mtz): import gemmi diff --git a/tests/unit/cli/test_choices.py b/tests/unit/cli/test_choices.py new file mode 100644 index 00000000..52d39af1 --- /dev/null +++ b/tests/unit/cli/test_choices.py @@ -0,0 +1,29 @@ +"""The torch-free choice lists in ``torchref.cli._choices`` match their sources.""" + +import pytest + +from torchref.cli import _choices + +pytestmark = pytest.mark.unit + + +def test_ded_weight_schemes_match_the_registry(): + from torchref.maps.ded_weights import DEFAULT_SCHEME, SCHEMES + + assert _choices.DED_WEIGHT_SCHEMES == tuple(SCHEMES) + assert _choices.DEFAULT_DED_WEIGHT_SCHEME == DEFAULT_SCHEME + + +def test_scale_targets_match_the_scaler(): + from torchref.scaling.scaler_base import DEFAULT_SCALE_TARGET, SCALE_TARGETS + + assert _choices.SCALE_TARGETS == tuple(SCALE_TARGETS) + assert _choices.DEFAULT_SCALE_TARGET == DEFAULT_SCALE_TARGET + + +def test_xray_target_docs_match_the_table(): + from torchref.refinement.targets.xray._specs import XRAY_TARGETS + + assert _choices.XRAY_TARGET_DOCS == tuple( + (spec.name, spec.doc) for spec in XRAY_TARGETS.specs + ) diff --git a/tests/unit/cli/test_difference_map_args.py b/tests/unit/cli/test_difference_map_args.py new file mode 100644 index 00000000..a411b517 --- /dev/null +++ b/tests/unit/cli/test_difference_map_args.py @@ -0,0 +1,67 @@ +"""Argument checks that ``torchref.difference-map`` makes before reading any file.""" + +import sys + +import pytest + +from torchref.cli import difference_map + +pytestmark = pytest.mark.unit + + +def _run(monkeypatch, *extra, output="out.mtz"): + argv = [ + "torchref.difference-map", + "-dm", "missing_dark.pdb", + "-dsf", "missing_dark.mtz", "-lsf", "missing_light.mtz", + "-o", output, + *extra, + ] # fmt: skip + monkeypatch.setattr(sys, "argv", argv) + return difference_map.main() + + +def _refused(monkeypatch, capsys, *extra, **kwargs): + with pytest.raises(SystemExit) as exc: + _run(monkeypatch, *extra, **kwargs) + assert exc.value.code == 2 + return capsys.readouterr().err + + +def test_valid_arguments_reach_the_file_check(monkeypatch, capsys): + assert _run(monkeypatch) == 1 + assert "file not found" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "extra, message", + [ + (["--fraction", "0.3"], "--fraction needs -lm"), + (["-lm", "light.pdb"], "--fraction is required"), + (["--all-columns"], "--all-columns needs -lm"), + (["--ded-weight", "inverse_variance"], "--ded-weight only weights"), + ( + ["-lm", "light.pdb", "--fraction", "0.3", "--ded-weight", "none"], + "--ded-weight only weights", + ), + ], +) +def test_ignored_arguments_are_refused(monkeypatch, capsys, extra, message): + assert message in _refused(monkeypatch, capsys, *extra) + + +def test_ded_weight_is_accepted_where_it_weights_columns(monkeypatch, capsys): + extra = ["-lm", "light.pdb", "--fraction", "0.3", "--all-columns"] + assert _run(monkeypatch, *extra, "--ded-weight", "none") == 1 + assert "file not found" in capsys.readouterr().err + + +@pytest.mark.parametrize("output", ["out/", "out.map", "results"]) +def test_output_must_name_an_mtz_file(monkeypatch, capsys, output): + assert "-o/--output" in _refused(monkeypatch, capsys, output=output) + + +def test_output_refuses_an_existing_directory(monkeypatch, capsys, tmp_path): + (tmp_path / "out.mtz").mkdir() + err = _refused(monkeypatch, capsys, output=str(tmp_path / "out.mtz")) + assert "not a directory" in err diff --git a/tests/unit/cli/test_entry_points.py b/tests/unit/cli/test_entry_points.py index e812eb78..bd6eb98a 100644 --- a/tests/unit/cli/test_entry_points.py +++ b/tests/unit/cli/test_entry_points.py @@ -3,7 +3,8 @@ The settings only take effect if they are in place before torch loads its OpenMP runtime, so the launcher must not import torch, every console script must go through it, and a library import of torchref must leave the user's process alone. Each check that -depends on import order runs in a fresh interpreter. +depends on import order runs in a fresh interpreter. ``import torchref`` and every +``--help`` stay torch-free, so a typo on the command line costs no torch start-up. """ import importlib @@ -73,3 +74,53 @@ def test_every_console_script_goes_through_the_launcher(): entry = getattr(_torchref_cli, attr) cli_module = importlib.import_module(f"torchref.cli.{entry.__name__}") assert callable(cli_module.main), f"{name}: torchref.cli.{entry.__name__}" + + +def test_library_import_does_not_import_torch(): + out = _fresh( + "import sys, torchref\n" + "print('torch' in sys.modules)\n" + "from torchref import ReflectionData\n" + "print('torch' in sys.modules, ReflectionData.__module__)\n" + "print('Model' in dir(torchref), torchref.config.get_float_dtype())" + ) + assert out.splitlines()[-3:] == [ + "False", + "True torchref.io.datasets.reflection_data", + "True torch.float32", + ] + + +def test_unknown_attribute_raises_attribute_error(): + out = _fresh( + "import torchref\n" + "try:\n" + " torchref.no_such_name\n" + "except AttributeError as e:\n" + " print(type(e).__name__)" + ) + assert out.splitlines()[-1] == "AttributeError" + + +def _script_names(): + try: + import tomllib + except ImportError: # Python 3.10 + return [] + with open(ROOT / "pyproject.toml", "rb") as f: + return sorted(tomllib.load(f)["project"]["scripts"].items()) + + +@pytest.mark.parametrize("name, target", _script_names()) +def test_help_does_not_import_torch(name, target): + attr = target.split(":")[1] + out = _fresh( + "import sys, _torchref_cli\n" + f"sys.argv = [{name!r}, '--help']\n" + "try:\n" + f" _torchref_cli.{attr}()\n" + "except SystemExit as e:\n" + " assert e.code == 0, e.code\n" + "print('torch' in sys.modules)" + ) + assert out.splitlines()[-1] == "False", f"{name} --help imported torch" diff --git a/torchref/__init__.py b/torchref/__init__.py index c9ba9ddc..73793422 100644 --- a/torchref/__init__.py +++ b/torchref/__init__.py @@ -11,6 +11,12 @@ ) refinement.refine(macro_cycles=10) +``import torchref`` sets the thread count and ``PYTORCH_ENABLE_MPS_FALLBACK`` but does +not import torch: the re-exported names below (``Model``, ``ReflectionData``, +``dtypes``, ...) and the subpackages are imported on first access. The environment +variables :mod:`torchref.config` reads are therefore read on that first access, not at +``import torchref``. + Modules ------- io @@ -42,13 +48,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 +76,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,42 +89,55 @@ 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. +# Public name -> defining module. Resolved on first access by ``__getattr__`` so that +# ``import torchref`` (and with it every ``torchref.cli`` module, whose ``--help`` must +# not wait for torch) does not import torch. ``Restraints`` is deliberately absent: +# constructing it can trigger a monomer-library download. +_LAZY_EXPORTS = { + "dtypes": "torchref.config", + "device": "torchref.config", + "sigma_cutoff_ed": "torchref.config", + "DatasetCollection": "torchref.io", + "ReflectionData": "torchref.io", + "ScaledDataset": "torchref.io", + "read_cif": "torchref.io", + "read_mtz": "torchref.io", + "read_pdb": "torchref.io", + "DifferenceMap": "torchref.maps", + "Map": "torchref.maps", + "Model": "torchref.model", + "ModelFT": "torchref.model", + "RigidXYZTensor": "torchref.model.rigid_xyz", + "LBFGSRefinement": "torchref.refinement", + "Refinement": "torchref.refinement", + "RigidBodyRefinementStep": "torchref.refinement.rigid_body_refinement", + "Scaler": "torchref.scaling", + "ScalerBase": "torchref.scaling", + "SolventModel": "torchref.scaling", + "Cell": "torchref.symmetry", + "SpaceGroup": "torchref.symmetry", + "Symmetry": "torchref.symmetry", + "DeviceMixin": "torchref.utils.device_mixin", +} + + +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 ( + (PATH_TORCHREF / name / "__init__.py").exists() + or (PATH_TORCHREF / f"{name}.py").exists() + ): + 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__ = [ diff --git a/torchref/cli/_choices.py b/torchref/cli/_choices.py new file mode 100644 index 00000000..ee553b85 --- /dev/null +++ b/torchref/cli/_choices.py @@ -0,0 +1,52 @@ +"""Choice lists the CLI parsers show in ``--help``, held without importing torch. + +Each name mirrors a table defined beside the code that uses it, and that table stays +the source of truth: importing it imports torch, which a ``--help`` must not wait for. +``tests/unit/cli/test_choices.py`` pins every mirror to its source. +""" + +#: :data:`torchref.maps.ded_weights.SCHEMES`. +DED_WEIGHT_SCHEMES = ("none", "inverse_variance", "q") +#: :data:`torchref.maps.ded_weights.DEFAULT_SCHEME`. +DEFAULT_DED_WEIGHT_SCHEME = "q" + +#: :data:`torchref.scaling.scaler_base.SCALE_TARGETS`. +SCALE_TARGETS = ("nll", "ml_noalpha", "ls") +#: :data:`torchref.scaling.scaler_base.DEFAULT_SCALE_TARGET`. +DEFAULT_SCALE_TARGET = "ls" + +#: ``(name, doc)`` of each row of +#: :data:`torchref.refinement.targets.xray._specs.XRAY_TARGETS`, in table order. +XRAY_TARGET_DOCS = ( + ( + "ml", + "Read MLF, variance epsilon*beta, conditional mean alpha*|F_calc| (default).", + ), + ("ml_noalpha", "As 'ml' with the Luzzati mean coupling fixed at 1."), + ( + "ml_full", + "Full-form MLF: marginalises the unknown error-free amplitude, so the " + "observation error enters as an amplitude-only Gaussian. ~4x the cost.", + ), + ( + "nll_beta", + "Gaussian amplitude NLL on ml's model-error variance -- the large-signal " + "limit of 'ml'. Diagnostic: isolates the variance model from the shape.", + ), + ( + "nll", + "Gaussian amplitude NLL weighted by the experimental sigma only. No " + "model-error term, so it does not control overfitting.", + ), + ( + "nll_i", + "Gaussian NLL on the observed INTENSITIES weighted by sigma(I). As 'nll' " + "but skips the French-Wilson conversion, which reshapes the weak tail.", + ), + ("ls", "Least squares with unit weights; the scaler owns the overall scale."), + ( + "ls_wunit_k1", + "Phenix-style least squares: unit weights and a single global scale " + "recomputed every gradient call (bypasses the scaler).", + ), +) diff --git a/torchref/cli/_common.py b/torchref/cli/_common.py index 4f946203..03e01aea 100644 --- a/torchref/cli/_common.py +++ b/torchref/cli/_common.py @@ -422,7 +422,10 @@ def add_ded_weight_args(parser: argparse.ArgumentParser) -> None: 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 + from torchref.cli._choices import ( + DED_WEIGHT_SCHEMES as SCHEMES, + DEFAULT_DED_WEIGHT_SCHEME as DEFAULT_SCHEME, + ) parser.add_argument( "--ded-weight", diff --git a/torchref/cli/collection_difference_refine.py b/torchref/cli/collection_difference_refine.py index 34a8ef21..13996c62 100644 --- a/torchref/cli/collection_difference_refine.py +++ b/torchref/cli/collection_difference_refine.py @@ -20,15 +20,19 @@ --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, @@ -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() @@ -192,6 +190,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 +344,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 @@ -405,6 +407,9 @@ def _two_moment_columns(mc, dc, mask, fcalc_dark_full, fcalc_mixed_full, the type beside the value preserves the crystallographic column type. """ import numpy as np + import torch + + from torchref.io.datasets.french_wilson import french_wilson_auto empty = ({}, {}) if float(mc.sigma_alpha_sq) == 0.0: @@ -609,6 +614,8 @@ def _phasing_columns(mc, scaler, hkl_all, mask, *, fcalc_dark, Fobs_dark_vals, Returns ``(columns, types, ctx)``. ``ctx`` carries the intermediates the extrapolation and two-moment layers need, so nothing is computed twice. """ + import torch + mixed_model = mc["light"] with torch.no_grad(): @@ -701,6 +708,8 @@ def _extrapolation_columns( 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. """ + import torch + from torchref import ReflectionData, Scaler from torchref.base.metrics.rfactor import rfactor_work_free @@ -812,10 +821,6 @@ def _np(t): return columns, types, 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 = { @@ -831,7 +836,7 @@ def _np(t): } -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 +852,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 +880,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) @@ -888,8 +895,9 @@ def write_results_mtz( mc=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. @@ -926,6 +934,8 @@ def write_results_mtz( 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,6 +946,14 @@ 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_light = dc["light"] @@ -1058,7 +1076,9 @@ def write_results_mtz( weight_columns=weight_columns, kscale=kscale, ) - datasets = {name: "difference" for name in _DIFFERENCE_DATASET_COLUMNS} + datasets = dict.fromkeys( + ["dF", "SIGdF", "PHDELWT", "KSCALE", *weight_columns], "difference" + ) if mc is not None: phase_cols, phase_types, ctx = _phasing_columns( @@ -1118,7 +1138,7 @@ def write_results_mtz( 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 +1162,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(): @@ -1269,6 +1291,11 @@ def main(): add_general_args(parser) args = parser.parse_args() + + import torch + + from torchref.utils.serialization import convert_to_serializable + register_timing() # --- Parse fractions --- diff --git a/torchref/cli/difference_map.py b/torchref/cli/difference_map.py index 612f2b7d..3d8d0849 100644 --- a/torchref/cli/difference_map.py +++ b/torchref/cli/difference_map.py @@ -18,6 +18,9 @@ Given ``-lm``, the light state's amplitude and phase and the extrapolated map follow. ``--all-columns`` adds the alternative constructions of both. +Beside the MTZ, ``_summary.json`` records the inputs, the dataset scaling, the +R-factors and the fits behind the weights and the extrapolation shrinkage. + Examples -------- :: @@ -32,23 +35,23 @@ """ import argparse +import json import sys from pathlib import Path -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_dual_model_args, add_dmin_arg, + add_dual_model_args, add_general_args, add_output_arg, build_dual_column_names, configure_unbuffered_output, - register_timing, - parse_device_str, difference_config_from_args, + parse_device_str, + register_timing, validate_cif_files, validate_files, ) @@ -56,13 +59,64 @@ configure_unbuffered_output() +def _summary_path(out_path: Path) -> Path: + """The summary JSON beside the MTZ: ``results.mtz`` -> ``results_summary.json``.""" + return out_path.with_name(f"{out_path.stem}_summary.json") + + +def _check_arguments(parser: argparse.ArgumentParser, args: argparse.Namespace): + """Refuse argument combinations that would be silently ignored. + + Runs before any data is read, so a mistake costs no scale fit. + """ + has_light_model = args.light_model is not None + if has_light_model and args.fraction is None: + parser.error( + "--fraction is required with -lm/--light-model: the extrapolation " + "divides by the light-state occupancy." + ) + if not has_light_model and 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 not has_light_model and args.all_columns: + parser.error( + "--all-columns needs -lm/--light-model: every column it adds is built from " + "the light state's model. The difference map alone has no alternatives." + ) + if args.ded_weight != DEFAULT_DED_WEIGHT_SCHEME and not ( + has_light_model and args.all_columns + ): + parser.error( + "--ded-weight only weights the model-phased difference columns that " + "-lm/--light-model with --all-columns adds. The difference map carries " + "every scheme's weight as its own column (W_Q, W_InVa): choose one when " + "building the map, e.g. torchref.mtz2map -csf dF -cw W_InVa -cphi PHDELWT." + ) + + out_path = Path(args.output) + if out_path.is_dir() or args.output.endswith(("/", "\\")): + parser.error( + f"-o/--output names the MTZ file to write, not a directory: {args.output}" + ) + if out_path.suffix.lower() != ".mtz": + parser.error(f"-o/--output must end in .mtz, got {args.output}") + + 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, epilog=""" +Writes the MTZ named by -o and, beside it, _summary.json with the inputs, +the dataset scaling, the R-factors and the difference-weight fit. + Examples: # difference map only -- needs no light model and no fraction torchref.difference-map \\ @@ -84,7 +138,11 @@ def main(): 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_output_arg( + output, + help="Output MTZ file (e.g. results.mtz); the JSON summary is written beside " + "it as _summary.json", + ) add_all_columns_arg(output) add_ded_weight_args(output) @@ -94,27 +152,9 @@ def main(): add_general_args(parser) args = parser.parse_args() - - register_timing() + _check_arguments(parser, args) 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"), @@ -128,10 +168,28 @@ def main(): if validate_cif_files(args.cif): return 1 + col_dark, col_light = build_dual_column_names(args) + + import torch + + from torchref.cli.collection_difference_refine import ( + compute_rfactors, + setup_dark_only, + setup_dataset_collection, + setup_model_collection, + setup_scaler, + write_results_mtz, + ) + from torchref.utils.serialization import convert_to_serializable + + register_timing() + out_path = Path(args.output) + summary_path = _summary_path(out_path) out_path.parent.mkdir(parents=True, exist_ok=True) 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,7 +204,8 @@ 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: {out_path}") + print(f"Summary: {summary_path}") print(f"Device: {device}") if args.dmin: print(f"Resolution cutoff: {args.dmin:.2f} A") @@ -158,28 +217,21 @@ def main(): 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, ) + # The model's FFT grid has to reach the data it is compared with, so without + # --dmin it is sized from the data's own high-resolution limit. + d_min = args.dmin + if d_min is None: + d_min = min(float(dc[key].resolution.min()) for key in ("dark", "light")) + if args.verbose > 0: print("Setting up models...") sys.stdout.flush() @@ -196,31 +248,28 @@ def main(): 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() + fits = {"dark": (dark_model, dc["dark"])} + if has_light_model: + fits["mixed"] = (mc["light"], dc["light"]) + r_factors = {} + for name, (model, data) in fits.items(): + r_work, r_free = compute_rfactors(model, data, scaler) + r_factors[name] = {"r_work": r_work, "r_free": r_free} + if args.verbose > 0: + print(f" R-factor ({name}): R_work={r_work:.4f} R_free={r_free:.4f}") if args.verbose > 0: + print() print("Computing map coefficients...") sys.stdout.flush() with torch.no_grad(): - write_results_mtz( + map_diagnostics = write_results_mtz( dc, dark_model, scaler, @@ -230,9 +279,33 @@ def main(): 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, + "model_grid_dmin": d_min, + "all_columns": args.all_columns, + "ded_weight": args.ded_weight, + "difference_gamma": args.difference_gamma, + "difference_sigma_scale": args.difference_sigma_scale, + }, + "dataset_scaling": dc.scaling_metrics, + "results": {"r_factors": r_factors, **map_diagnostics}, + "output_files": {"mtz": str(out_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(f" Summary written to {summary_path}") print() 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..77352a8d 100644 --- a/torchref/cli/refine.py +++ b/torchref/cli/refine.py @@ -16,8 +16,11 @@ 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 +40,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,14 +148,11 @@ 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. + choices=[name for name, _ in XRAY_TARGET_DOCS], + # Choices and help both come from the mirror of XRAY_TARGETS, which a test pins + # to the table, so the two cannot drift from each other or from the dispatch. help="X-ray target function. " + " ".join( - f"'{spec.name}': {spec.doc}" for spec in XRAY_TARGETS.specs + f"'{name}': {doc}" for name, doc in XRAY_TARGET_DOCS ), ) refine_group.add_argument( @@ -235,6 +229,12 @@ def main(): args = parser.parse_args() + import torch + + # Imported before any json output: it makes StatEntry serialisable process-wide. + import torchref.utils.stats # noqa: F401 + 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) From 32b5a9c3a614efce7512c4c6eed31befd863f6d8 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 21 Oct 2026 13:33:24 +0000 Subject: [PATCH 2/3] difference-map: dark-only phases, DELFWT, all extrapolations, -o directory - The difference map (DELFWT/PHDELWT, dF, weights, KSCALE, Fc_dark) takes its phases from a scale fit of the dark model alone, in difference-map and difference-refine, so it no longer changes when a light model is given - DELFWT = dF times the --ded-weight weight, so Coot opens the difference map - The phase-aware and scalar extrapolations are written by default beside the shrunk one; --all-columns now gates only the model-phased difference and intensity columns - FC/PHIC/PHIC_diff are renamed Fc_light/PHFc_light/PHdFc, matching Fc_dark - difference-map takes -o as a directory (difference_map.mtz, summary.json) and both difference CLIs gain --dmax; difference-refine sizes the model grid from the data without --dmin - -lm/--fraction sit in an "optional" help group - Column builders return {label: (values, MTZ type)}; the torch-free choice mirror and lazy torchref exports are compacted, and the new tests dropped Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014aviwqFcZpxP3HWNy1ZruY --- docs/changelog.rst | 4 +- docs/user_guide/cli.rst | 61 ++- tests/integration/test_cli_ded_weights.py | 17 +- tests/integration/test_cli_two_moment_mtz.py | 43 +- tests/unit/cli/test_choices.py | 29 -- tests/unit/cli/test_difference_map_args.py | 67 --- tests/unit/cli/test_entry_points.py | 53 +- torchref/__init__.py | 98 +--- torchref/cli/_choices.py | 63 +-- torchref/cli/_common.py | 49 +- torchref/cli/collection_difference_refine.py | 494 +++++++------------ torchref/cli/difference_map.py | 194 +++----- torchref/cli/refine.py | 18 +- 13 files changed, 380 insertions(+), 810 deletions(-) delete mode 100644 tests/unit/cli/test_choices.py delete mode 100644 tests/unit/cli/test_difference_map_args.py diff --git a/docs/changelog.rst b/docs/changelog.rst index 5be35f86..6cb09d5c 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -14,7 +14,7 @@ Breaking changes - 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`` refuses options it would ignore (``--all-columns`` or a non-default ``--ded-weight`` without ``-lm``, a ``-o`` that is not a ``.mtz`` file) and writes ``_summary.json`` beside the MTZ +- ``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 @@ -48,7 +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 -- ``torchref.difference-map`` without ``--dmin`` sizes the model's FFT grid from the data's resolution instead of a fixed 1.0 Å +- 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 e38f323a..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,26 +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. -``-o`` names the MTZ file (it must end in ``.mtz``), and ``_summary.json`` is -written beside it with the inputs, the dataset scaling, the R-factors and the fits -behind the ``q`` weights and the extrapolation shrinkage. Options that would change -nothing in the output are refused rather than ignored: ``--all-columns`` needs ``-lm``, -and ``--ded-weight`` only weights the model-phased columns that ``-lm`` with -``--all-columns`` adds -- the difference map itself carries every scheme's weight as a -column, chosen with ``mtz2map -cw``. - -**Key options:** ``--ded-weight {q,inverse_variance,none}`` selects the scheme 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`` 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 938f8ebc..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): @@ -151,18 +152,6 @@ def test_difference_map_writes_df_weights_and_scale(diff_mtz): assert (w_q > 0).all() -def test_difference_map_writes_a_summary_beside_the_mtz(diff_mtz): - import gemmi - - summary = json.loads(diff_mtz.with_name("diff_summary.json").read_text()) - assert summary["output_files"]["mtz"] == str(diff_mtz) - assert set(summary["results"]["r_factors"]) == {"dark"} - assert summary["results"]["ded_weights"]["scheme"] == "q" - assert summary["input"]["model_grid_dmin"] == 2.2 - history = gemmi.read_mtz_file(str(diff_mtz)).history - assert history[0].startswith("torchref difference-map ") - - def test_mtz2map_builds_the_weighted_and_electron_maps(project_root, pair, diff_mtz): import gemmi 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/tests/unit/cli/test_choices.py b/tests/unit/cli/test_choices.py deleted file mode 100644 index 52d39af1..00000000 --- a/tests/unit/cli/test_choices.py +++ /dev/null @@ -1,29 +0,0 @@ -"""The torch-free choice lists in ``torchref.cli._choices`` match their sources.""" - -import pytest - -from torchref.cli import _choices - -pytestmark = pytest.mark.unit - - -def test_ded_weight_schemes_match_the_registry(): - from torchref.maps.ded_weights import DEFAULT_SCHEME, SCHEMES - - assert _choices.DED_WEIGHT_SCHEMES == tuple(SCHEMES) - assert _choices.DEFAULT_DED_WEIGHT_SCHEME == DEFAULT_SCHEME - - -def test_scale_targets_match_the_scaler(): - from torchref.scaling.scaler_base import DEFAULT_SCALE_TARGET, SCALE_TARGETS - - assert _choices.SCALE_TARGETS == tuple(SCALE_TARGETS) - assert _choices.DEFAULT_SCALE_TARGET == DEFAULT_SCALE_TARGET - - -def test_xray_target_docs_match_the_table(): - from torchref.refinement.targets.xray._specs import XRAY_TARGETS - - assert _choices.XRAY_TARGET_DOCS == tuple( - (spec.name, spec.doc) for spec in XRAY_TARGETS.specs - ) diff --git a/tests/unit/cli/test_difference_map_args.py b/tests/unit/cli/test_difference_map_args.py deleted file mode 100644 index a411b517..00000000 --- a/tests/unit/cli/test_difference_map_args.py +++ /dev/null @@ -1,67 +0,0 @@ -"""Argument checks that ``torchref.difference-map`` makes before reading any file.""" - -import sys - -import pytest - -from torchref.cli import difference_map - -pytestmark = pytest.mark.unit - - -def _run(monkeypatch, *extra, output="out.mtz"): - argv = [ - "torchref.difference-map", - "-dm", "missing_dark.pdb", - "-dsf", "missing_dark.mtz", "-lsf", "missing_light.mtz", - "-o", output, - *extra, - ] # fmt: skip - monkeypatch.setattr(sys, "argv", argv) - return difference_map.main() - - -def _refused(monkeypatch, capsys, *extra, **kwargs): - with pytest.raises(SystemExit) as exc: - _run(monkeypatch, *extra, **kwargs) - assert exc.value.code == 2 - return capsys.readouterr().err - - -def test_valid_arguments_reach_the_file_check(monkeypatch, capsys): - assert _run(monkeypatch) == 1 - assert "file not found" in capsys.readouterr().err - - -@pytest.mark.parametrize( - "extra, message", - [ - (["--fraction", "0.3"], "--fraction needs -lm"), - (["-lm", "light.pdb"], "--fraction is required"), - (["--all-columns"], "--all-columns needs -lm"), - (["--ded-weight", "inverse_variance"], "--ded-weight only weights"), - ( - ["-lm", "light.pdb", "--fraction", "0.3", "--ded-weight", "none"], - "--ded-weight only weights", - ), - ], -) -def test_ignored_arguments_are_refused(monkeypatch, capsys, extra, message): - assert message in _refused(monkeypatch, capsys, *extra) - - -def test_ded_weight_is_accepted_where_it_weights_columns(monkeypatch, capsys): - extra = ["-lm", "light.pdb", "--fraction", "0.3", "--all-columns"] - assert _run(monkeypatch, *extra, "--ded-weight", "none") == 1 - assert "file not found" in capsys.readouterr().err - - -@pytest.mark.parametrize("output", ["out/", "out.map", "results"]) -def test_output_must_name_an_mtz_file(monkeypatch, capsys, output): - assert "-o/--output" in _refused(monkeypatch, capsys, output=output) - - -def test_output_refuses_an_existing_directory(monkeypatch, capsys, tmp_path): - (tmp_path / "out.mtz").mkdir() - err = _refused(monkeypatch, capsys, output=str(tmp_path / "out.mtz")) - assert "not a directory" in err diff --git a/tests/unit/cli/test_entry_points.py b/tests/unit/cli/test_entry_points.py index bd6eb98a..e812eb78 100644 --- a/tests/unit/cli/test_entry_points.py +++ b/tests/unit/cli/test_entry_points.py @@ -3,8 +3,7 @@ The settings only take effect if they are in place before torch loads its OpenMP runtime, so the launcher must not import torch, every console script must go through it, and a library import of torchref must leave the user's process alone. Each check that -depends on import order runs in a fresh interpreter. ``import torchref`` and every -``--help`` stay torch-free, so a typo on the command line costs no torch start-up. +depends on import order runs in a fresh interpreter. """ import importlib @@ -74,53 +73,3 @@ def test_every_console_script_goes_through_the_launcher(): entry = getattr(_torchref_cli, attr) cli_module = importlib.import_module(f"torchref.cli.{entry.__name__}") assert callable(cli_module.main), f"{name}: torchref.cli.{entry.__name__}" - - -def test_library_import_does_not_import_torch(): - out = _fresh( - "import sys, torchref\n" - "print('torch' in sys.modules)\n" - "from torchref import ReflectionData\n" - "print('torch' in sys.modules, ReflectionData.__module__)\n" - "print('Model' in dir(torchref), torchref.config.get_float_dtype())" - ) - assert out.splitlines()[-3:] == [ - "False", - "True torchref.io.datasets.reflection_data", - "True torch.float32", - ] - - -def test_unknown_attribute_raises_attribute_error(): - out = _fresh( - "import torchref\n" - "try:\n" - " torchref.no_such_name\n" - "except AttributeError as e:\n" - " print(type(e).__name__)" - ) - assert out.splitlines()[-1] == "AttributeError" - - -def _script_names(): - try: - import tomllib - except ImportError: # Python 3.10 - return [] - with open(ROOT / "pyproject.toml", "rb") as f: - return sorted(tomllib.load(f)["project"]["scripts"].items()) - - -@pytest.mark.parametrize("name, target", _script_names()) -def test_help_does_not_import_torch(name, target): - attr = target.split(":")[1] - out = _fresh( - "import sys, _torchref_cli\n" - f"sys.argv = [{name!r}, '--help']\n" - "try:\n" - f" _torchref_cli.{attr}()\n" - "except SystemExit as e:\n" - " assert e.code == 0, e.code\n" - "print('torch' in sys.modules)" - ) - assert out.splitlines()[-1] == "False", f"{name} --help imported torch" diff --git a/torchref/__init__.py b/torchref/__init__.py index 73793422..bb0f7c72 100644 --- a/torchref/__init__.py +++ b/torchref/__init__.py @@ -11,11 +11,8 @@ ) refinement.refine(macro_cycles=10) -``import torchref`` sets the thread count and ``PYTORCH_ENABLE_MPS_FALLBACK`` but does -not import torch: the re-exported names below (``Model``, ``ReflectionData``, -``dtypes``, ...) and the subpackages are imported on first access. The environment -variables :mod:`torchref.config` reads are therefore read on that first access, not at -``import torchref``. +``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 ------- @@ -89,35 +86,25 @@ PATH_TORCHREF_DATA = PATH_TORCHREF / "data" -# Public name -> defining module. Resolved on first access by ``__getattr__`` so that -# ``import torchref`` (and with it every ``torchref.cli`` module, whose ``--help`` must -# not wait for torch) does not import torch. ``Restraints`` is deliberately absent: -# constructing it can trigger a monomer-library download. +# 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 = { - "dtypes": "torchref.config", - "device": "torchref.config", - "sigma_cutoff_ed": "torchref.config", - "DatasetCollection": "torchref.io", - "ReflectionData": "torchref.io", - "ScaledDataset": "torchref.io", - "read_cif": "torchref.io", - "read_mtz": "torchref.io", - "read_pdb": "torchref.io", - "DifferenceMap": "torchref.maps", - "Map": "torchref.maps", - "Model": "torchref.model", - "ModelFT": "torchref.model", - "RigidXYZTensor": "torchref.model.rigid_xyz", - "LBFGSRefinement": "torchref.refinement", - "Refinement": "torchref.refinement", - "RigidBodyRefinementStep": "torchref.refinement.rigid_body_refinement", - "Scaler": "torchref.scaling", - "ScalerBase": "torchref.scaling", - "SolventModel": "torchref.scaling", - "Cell": "torchref.symmetry", - "SpaceGroup": "torchref.symmetry", - "Symmetry": "torchref.symmetry", - "DeviceMixin": "torchref.utils.device_mixin", + 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() } @@ -125,9 +112,8 @@ 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 ( - (PATH_TORCHREF / name / "__init__.py").exists() - or (PATH_TORCHREF / f"{name}.py").exists() + 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: @@ -140,42 +126,4 @@ def __dir__(): return sorted(set(globals()) | set(_LAZY_EXPORTS)) -__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", -] +__all__ = ["__version__", "ROOT_TORCHREF", "PATH_TORCHREF", "N_CPUS", *_LAZY_EXPORTS] diff --git a/torchref/cli/_choices.py b/torchref/cli/_choices.py index ee553b85..5b73d265 100644 --- a/torchref/cli/_choices.py +++ b/torchref/cli/_choices.py @@ -1,52 +1,23 @@ """Choice lists the CLI parsers show in ``--help``, held without importing torch. -Each name mirrors a table defined beside the code that uses it, and that table stays -the source of truth: importing it imports torch, which a ``--help`` must not wait for. -``tests/unit/cli/test_choices.py`` pins every mirror to its source. +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`. -DED_WEIGHT_SCHEMES = ("none", "inverse_variance", "q") -#: :data:`torchref.maps.ded_weights.DEFAULT_SCHEME`. -DEFAULT_DED_WEIGHT_SCHEME = "q" +#: :data:`torchref.maps.ded_weights.SCHEMES` and ``DEFAULT_SCHEME``. +DED_WEIGHT_SCHEMES, DEFAULT_DED_WEIGHT_SCHEME = ("none", "inverse_variance", "q"), "q" -#: :data:`torchref.scaling.scaler_base.SCALE_TARGETS`. -SCALE_TARGETS = ("nll", "ml_noalpha", "ls") -#: :data:`torchref.scaling.scaler_base.DEFAULT_SCALE_TARGET`. -DEFAULT_SCALE_TARGET = "ls" +#: :data:`torchref.scaling.scaler_base.SCALE_TARGETS` and ``DEFAULT_SCALE_TARGET``. +SCALE_TARGETS, DEFAULT_SCALE_TARGET = ("nll", "ml_noalpha", "ls"), "ls" -#: ``(name, doc)`` of each row of -#: :data:`torchref.refinement.targets.xray._specs.XRAY_TARGETS`, in table order. -XRAY_TARGET_DOCS = ( - ( - "ml", - "Read MLF, variance epsilon*beta, conditional mean alpha*|F_calc| (default).", - ), - ("ml_noalpha", "As 'ml' with the Luzzati mean coupling fixed at 1."), - ( - "ml_full", - "Full-form MLF: marginalises the unknown error-free amplitude, so the " - "observation error enters as an amplitude-only Gaussian. ~4x the cost.", - ), - ( - "nll_beta", - "Gaussian amplitude NLL on ml's model-error variance -- the large-signal " - "limit of 'ml'. Diagnostic: isolates the variance model from the shape.", - ), - ( - "nll", - "Gaussian amplitude NLL weighted by the experimental sigma only. No " - "model-error term, so it does not control overfitting.", - ), - ( - "nll_i", - "Gaussian NLL on the observed INTENSITIES weighted by sigma(I). As 'nll' " - "but skips the French-Wilson conversion, which reshapes the weak tail.", - ), - ("ls", "Least squares with unit weights; the scaler owns the overall scale."), - ( - "ls_wunit_k1", - "Phenix-style least squares: unit weights and a single global scale " - "recomputed every gradient call (bypasses the scaler).", - ), -) +#: 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 03e01aea..28b26ff1 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,16 +391,14 @@ 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).", ) @@ -419,23 +421,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.cli._choices import ( - DED_WEIGHT_SCHEMES as SCHEMES, - DEFAULT_DED_WEIGHT_SCHEME as DEFAULT_SCHEME, - ) + 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 13996c62..7faa1b2c 100644 --- a/torchref/cli/collection_difference_refine.py +++ b/torchref/cli/collection_difference_refine.py @@ -36,12 +36,12 @@ 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, @@ -111,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( @@ -122,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) @@ -149,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): @@ -371,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 @@ -394,30 +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 from torchref.io.datasets.french_wilson import french_wilson_auto - empty = ({}, {}) - if float(mc.sigma_alpha_sq) == 0.0: - return empty - 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. @@ -460,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() ) @@ -474,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, + "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"), }) - 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 + return columns def _difference_columns( @@ -527,6 +490,7 @@ def _difference_columns( phases_dark, diff_Fobs, sig_diff, + weights, weight_columns, kscale, ): @@ -535,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. @@ -546,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 @@ -556,115 +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), + "FreeR_flag_dark": (_flags(data_dark), "I"), + "FreeR_flag_light": (_flags(data_light), "I"), } - 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. +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. """ import torch - mixed_model = mc["light"] - + 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( @@ -676,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) @@ -691,22 +632,19 @@ 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. + ``FEXT_SCALAR``, map ``2FEXT_SCALAR-Fc``/``PHFEXT_SCALAR`` + The classic ``(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. + 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 @@ -728,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 @@ -745,92 +680,64 @@ 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 # 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", } @@ -893,6 +800,7 @@ def write_results_mtz( filename, *, mc=None, + joint_scaler=None, all_columns=False, verbose=1, ded_weight=DEFAULT_DED_WEIGHT_SCHEME, @@ -902,10 +810,10 @@ def write_results_mtz( """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``. @@ -914,22 +822,23 @@ 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 @@ -955,7 +864,7 @@ def write_results_mtz( 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 @@ -972,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() @@ -992,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) @@ -1060,7 +960,7 @@ def write_results_mtz( f"before normalisation" ) - columns, types = _difference_columns( + columns = _difference_columns( data_dark, data_light, mask, @@ -1073,25 +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 = dict.fromkeys( - ["dF", "SIGdF", "PHDELWT", "KSCALE", *weight_columns], "difference" + ["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, @@ -1099,41 +998,30 @@ 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) @@ -1286,7 +1174,7 @@ def main(): ) res = parser.add_argument_group("Resolution") - add_dmin_arg(res) + add_resolution_args(res) add_general_args(parser) @@ -1372,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) @@ -1406,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...") @@ -1425,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: @@ -1756,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, @@ -1776,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 3d8d0849..068a42f5 100644 --- a/torchref/cli/difference_map.py +++ b/torchref/cli/difference_map.py @@ -3,23 +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. - -Beside the MTZ, ``_summary.json`` records the inputs, the dataset scaling, the -R-factors and the fits behind the weights and the extrapolation shrinkage. +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 -------- @@ -27,11 +27,11 @@ # 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 @@ -39,17 +39,17 @@ import sys from pathlib import Path -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_output_arg, + add_outdir_arg, + add_resolution_args, build_dual_column_names, configure_unbuffered_output, difference_config_from_args, + load_model, parse_device_str, register_timing, validate_cif_files, @@ -59,53 +59,6 @@ configure_unbuffered_output() -def _summary_path(out_path: Path) -> Path: - """The summary JSON beside the MTZ: ``results.mtz`` -> ``results_summary.json``.""" - return out_path.with_name(f"{out_path.stem}_summary.json") - - -def _check_arguments(parser: argparse.ArgumentParser, args: argparse.Namespace): - """Refuse argument combinations that would be silently ignored. - - Runs before any data is read, so a mistake costs no scale fit. - """ - has_light_model = args.light_model is not None - if has_light_model and args.fraction is None: - parser.error( - "--fraction is required with -lm/--light-model: the extrapolation " - "divides by the light-state occupancy." - ) - if not has_light_model and 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 not has_light_model and args.all_columns: - parser.error( - "--all-columns needs -lm/--light-model: every column it adds is built from " - "the light state's model. The difference map alone has no alternatives." - ) - if args.ded_weight != DEFAULT_DED_WEIGHT_SCHEME and not ( - has_light_model and args.all_columns - ): - parser.error( - "--ded-weight only weights the model-phased difference columns that " - "-lm/--light-model with --all-columns adds. The difference map carries " - "every scheme's weight as its own column (W_Q, W_InVa): choose one when " - "building the map, e.g. torchref.mtz2map -csf dF -cw W_InVa -cphi PHDELWT." - ) - - out_path = Path(args.output) - if out_path.is_dir() or args.output.endswith(("/", "\\")): - parser.error( - f"-o/--output names the MTZ file to write, not a directory: {args.output}" - ) - if out_path.suffix.lower() != ".mtz": - parser.error(f"-o/--output must end in .mtz, got {args.output}") - - def main(): """Entry point for ``torchref.difference-map``; returns the exit code.""" parser = argparse.ArgumentParser( @@ -114,47 +67,42 @@ def main(): "(no refinement).", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" -Writes the MTZ named by -o and, beside it, _summary.json with the inputs, -the dataset scaling, the R-factors and the difference-weight fit. - 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 + -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 (e.g. results.mtz); the JSON summary is written beside " - "it as _summary.json", + 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() - _check_arguments(parser, args) has_light_model = args.light_model is not None + 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"), (args.dark_structure_factor, "dark structure factor"), @@ -162,10 +110,7 @@ def main(): ] 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): + if validate_files(to_check) or validate_cif_files(args.cif): return 1 col_dark, col_light = build_dual_column_names(args) @@ -174,7 +119,7 @@ def main(): from torchref.cli.collection_difference_refine import ( compute_rfactors, - setup_dark_only, + setup_dark_scaler, setup_dataset_collection, setup_model_collection, setup_scaler, @@ -184,10 +129,9 @@ def main(): register_timing() - out_path = Path(args.output) - summary_path = _summary_path(out_path) - out_path.parent.mkdir(parents=True, exist_ok=True) - + 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 @@ -204,67 +148,54 @@ 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: {out_path}") - print(f"Summary: {summary_path}") + 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() - - if args.verbose > 0: print("Loading reflection data...") sys.stdout.flush() 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 model's FFT grid has to reach the data it is compared with, so without - # --dmin it is sized from the data's own high-resolution limit. - d_min = args.dmin - if d_min is None: - d_min = min(float(dc[key].resolution.min()) for key in ("dark", "light")) + # 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 + 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"])} + fits = {"dark": (dark_model, dc["dark"], scaler)} if has_light_model: - fits["mixed"] = (mc["light"], dc["light"]) + fits["mixed"] = (mc["light"], dc["light"], joint_scaler) r_factors = {} - for name, (model, data) in fits.items(): - r_work, r_free = compute_rfactors(model, data, scaler) + 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: print(f" R-factor ({name}): R_work={r_work:.4f} R_free={r_free:.4f}") + if args.verbose > 0: - print() print("Computing map coefficients...") sys.stdout.flush() @@ -273,8 +204,9 @@ def main(): 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, @@ -291,22 +223,18 @@ def main(): "fractions": fractions, "cif": args.cif, "dmin": args.dmin, + "dmax": args.dmax, "model_grid_dmin": d_min, - "all_columns": args.all_columns, - "ded_weight": args.ded_weight, - "difference_gamma": args.difference_gamma, - "difference_sigma_scale": args.difference_sigma_scale, }, "dataset_scaling": dc.scaling_metrics, "results": {"r_factors": r_factors, **map_diagnostics}, - "output_files": {"mtz": str(out_path), "summary": str(summary_path)}, + "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(f" Summary written to {summary_path}") - print() print("Done.") sys.stdout.flush() diff --git a/torchref/cli/refine.py b/torchref/cli/refine.py index 77352a8d..21e45db4 100644 --- a/torchref/cli/refine.py +++ b/torchref/cli/refine.py @@ -16,11 +16,7 @@ import warnings from pathlib import Path -from torchref.cli._choices import ( - DEFAULT_SCALE_TARGET, - SCALE_TARGETS, - XRAY_TARGET_DOCS, -) +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, @@ -148,12 +144,9 @@ def main(): "--xray-mode", type=str, default="ml", - choices=[name for name, _ in XRAY_TARGET_DOCS], - # Choices and help both come from the mirror of XRAY_TARGETS, which a test pins - # to the table, so the two cannot drift from each other or from the dispatch. - help="X-ray target function. " + " ".join( - f"'{name}': {doc}" for name, doc in XRAY_TARGET_DOCS - ), + 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", @@ -231,8 +224,7 @@ def main(): import torch - # Imported before any json output: it makes StatEntry serialisable process-wide. - import torchref.utils.stats # noqa: F401 + import torchref.utils.stats # noqa: F401 (makes json serialise StatEntry) from torchref.utils.serialization import convert_to_serializable register_timing() From af3bebc7e5e5f296e0c211f164c276cf5a2fb7e7 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 21 Oct 2026 13:47:23 +0000 Subject: [PATCH 3/3] Check --difference-sigma-scale bounds without importing torch The argparse converter imported difference_power for its bounds check, so a bad value paid torch start-up before the usage error. The bounds now come from the torch-free mirror; DifferencePowerConfig still validates at run time. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014aviwqFcZpxP3HWNy1ZruY --- torchref/cli/_choices.py | 3 +++ torchref/cli/_common.py | 12 +++++------- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/torchref/cli/_choices.py b/torchref/cli/_choices.py index 5b73d265..7add18f5 100644 --- a/torchref/cli/_choices.py +++ b/torchref/cli/_choices.py @@ -7,6 +7,9 @@ #: :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" diff --git a/torchref/cli/_common.py b/torchref/cli/_common.py index 28b26ff1..386aadb3 100644 --- a/torchref/cli/_common.py +++ b/torchref/cli/_common.py @@ -404,16 +404,14 @@ def add_all_columns_arg(parser: argparse.ArgumentParser) -> None: 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: