diff --git a/docs/changelog.rst b/docs/changelog.rst index 4e4d40f7..b56d0310 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -1,6 +1,14 @@ Changelog ========= +Version 0.6.2 +------------- +- Fixed crash in difference refinement with mismatched reflection files: HKL reindexing (``validate_hkl``/``remap``/``reduce_to_spacegroup``) now carries all per-reflection fields (including the anomalous bookkeeping read by ``hkl_for_sf()``) instead of a hardcoded subset. +- Added a metal shader for structure factor calculation +- Standardized structure factor calculation geometry +- Split gpu tests into cuda and mps +- Reworked VDW pair list creation + Version 0.6.1 ------------- - Fixed bug in beta estimation that caused instability in GPU refinement diff --git a/pyproject.toml b/pyproject.toml index 1ecb1db1..d7de55bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchref" -version = "0.6.1" +version = "0.6.2" description = "Pytorch based crystallographic refinement" readme = "README.md" requires-python = ">=3.10" @@ -98,8 +98,28 @@ profile = "black" line_length = 88 [tool.pytest.ini_options] +# Used when pytest is invoked from the repo root. ``tests/pytest.ini`` takes +# over when invoked from inside ``tests/``; keep the two in step so behaviour +# does not depend on the working directory. Accelerator tests are gated by +# ``tests/conftest.py``, which runs whatever the host supports -- do not add a +# default ``-m`` expression here, as it deselects at collection time and no +# flag can undo it. testpaths = ["tests"] python_files = "test_*.py" +python_classes = "Test*" +python_functions = "test_*" +pythonpath = ["."] +markers = [ + "unit: Unit tests (fast, no I/O, isolated)", + "integration: Integration tests (slower, real I/O, pipelines)", + "gpu: Needs any accelerator (CUDA or MPS); auto-skipped if the host has none", + "cuda: Needs CUDA specifically (e.g. Triton kernels); auto-skipped if absent", + "mps: Needs MPS specifically (Metal kernels); auto-skipped if absent", + "cuda_only: Deprecated alias for 'cuda'", + "slow: Long-running tests (skipped in quick runs)", + "openmm: Tests requiring OpenMM (the [amber] extra); skipped if absent", + "amber: Tests requiring OpenMM + AmberTools (antechamber/tleap); skipped if absent", +] [tool.ruff.lint] ignore = ["F841","E741","F841","E402","E722"] \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index 905a1953..d5afc48e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,7 @@ """ import importlib.util import shutil +import warnings import pytest import torchref @@ -21,13 +22,40 @@ _HAS_AMBERTOOLS = bool(shutil.which("antechamber") and shutil.which("tleap")) +def _cuda_available() -> bool: + return torch.cuda.is_available() + + +def _mps_available() -> bool: + return hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + + def pytest_addoption(parser): """Add custom command line options.""" + parser.addoption( + "--run-cuda", + action="store_true", + default=False, + help=( + "Require CUDA: fail instead of skipping if it is unavailable. " + "CUDA tests already run automatically when a CUDA device is " + "present; use this in CI to catch a runner that lost its GPU." + ), + ) + parser.addoption( + "--run-mps", + action="store_true", + default=False, + help=( + "Require MPS: fail instead of skipping if it is unavailable. " + "MPS tests already run automatically on Apple silicon." + ), + ) parser.addoption( "--run-gpu", action="store_true", default=False, - help="Run GPU tests" + help="Deprecated no-op: accelerator tests now run automatically.", ) parser.addoption( "--run-slow", @@ -41,26 +69,77 @@ def pytest_configure(config): """Configure pytest markers.""" config.addinivalue_line("markers", "unit: Unit tests (fast, no I/O)") config.addinivalue_line("markers", "integration: Integration tests (slower, real I/O)") - config.addinivalue_line("markers", "gpu: GPU-requiring tests (CUDA or MPS; skipped by default)") - config.addinivalue_line("markers", "cuda_only: Tests that specifically require CUDA (e.g. Triton)") + config.addinivalue_line("markers", "gpu: Needs any accelerator (CUDA or MPS); auto-skipped if none") + config.addinivalue_line("markers", "cuda: Needs CUDA specifically (e.g. Triton); auto-skipped if absent") + config.addinivalue_line("markers", "mps: Needs MPS specifically (Metal kernels); auto-skipped if absent") + config.addinivalue_line("markers", "cuda_only: Deprecated alias for 'cuda'") config.addinivalue_line("markers", "slow: Slow tests (skipped by default)") config.addinivalue_line("markers", "openmm: Needs OpenMM (the [amber] extra); skipped if absent") config.addinivalue_line("markers", "amber: Needs OpenMM + AmberTools (antechamber/tleap); skipped if absent") + if config.getoption("--run-gpu"): + # UserWarning, not DeprecationWarning: pytest.ini filters the latter, + # and a silently-swallowed notice is worse than none when the whole + # point is telling someone their flag no longer does anything. + warnings.warn( + "--run-gpu is deprecated and does nothing: accelerator tests now " + "run automatically wherever the backend is available. Use " + "--run-cuda / --run-mps to *require* a backend (fail rather than " + "skip when it is missing).", + UserWarning, + stacklevel=2, + ) + def pytest_collection_modifyitems(config, items): - """Skip GPU/slow tests unless explicitly requested.""" - skip_gpu = pytest.mark.skip(reason="Need --run-gpu option to run") + """Gate tests on what this host can actually do. + + Accelerator tests are **not** opt-in: a ``cuda``-marked test runs whenever + CUDA is present, an ``mps``-marked test whenever MPS is, and a ``gpu``-marked + (backend-agnostic) test whenever either is. Anything the host cannot run is + skipped with a reason naming the missing backend. + + ``--run-cuda`` / ``--run-mps`` invert the *absence* case from skip to + failure, for CI that expects a specific backend and would otherwise go + green on a runner that quietly lost its GPU. + """ + has_cuda = _cuda_available() + has_mps = _mps_available() + + require_cuda = config.getoption("--run-cuda") + require_mps = config.getoption("--run-mps") + if require_cuda and not has_cuda: + raise pytest.UsageError("--run-cuda given but no CUDA device is available") + if require_mps and not has_mps: + raise pytest.UsageError("--run-mps given but MPS is not available") + skip_slow = pytest.mark.skip(reason="Need --run-slow option to run") skip_openmm = pytest.mark.skip(reason="OpenMM not installed (pip install '.[amber]')") skip_amber = pytest.mark.skip( reason="AmberTools (antechamber/tleap) not on PATH (conda install ambertools)" ) + skip_cuda = pytest.mark.skip(reason="No CUDA device on this host") + skip_mps = pytest.mark.skip(reason="No MPS device on this host") + skip_gpu = pytest.mark.skip(reason="No accelerator (CUDA or MPS) on this host") for item in items: - if "gpu" in item.keywords and not config.getoption("--run-gpu"): + keywords = item.keywords + # ``cuda_only`` is the retired spelling of ``cuda``. + wants_cuda = "cuda" in keywords or "cuda_only" in keywords + wants_mps = "mps" in keywords + if wants_cuda and not has_cuda: + item.add_marker(skip_cuda) + if wants_mps and not has_mps: + item.add_marker(skip_mps) + # A bare ``gpu`` mark means "any accelerator"; a test that also names a + # specific backend has already been gated on the stricter condition. + if ( + "gpu" in keywords + and not (wants_cuda or wants_mps) + and not (has_cuda or has_mps) + ): item.add_marker(skip_gpu) - if "slow" in item.keywords and not config.getoption("--run-slow"): + if "slow" in keywords and not config.getoption("--run-slow"): item.add_marker(skip_slow) # Amber stack gates: "amber" needs OpenMM + AmberTools; "openmm" needs # just OpenMM. Skip with the most specific missing-dependency reason. @@ -135,26 +214,93 @@ def cpu_device() -> torch.device: return torch.device("cpu") -def _gpu_available() -> bool: - """Return True if any GPU backend (CUDA or MPS) is available.""" - if torch.cuda.is_available(): - return True - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return True - return False - - @pytest.fixture(scope="session") def gpu_device() -> torch.device: """GPU torch device (only use with @pytest.mark.gpu). Prefers CUDA, falls back to MPS; skips if neither is available. """ - if torch.cuda.is_available(): - return torch.device("cuda") - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return torch.device("mps") - pytest.skip("No GPU (CUDA or MPS) available") + accel = _accelerator() + if accel is None: + pytest.skip("No accelerator (CUDA or MPS) on this host") + return accel + + +def _accelerator() -> "torch.device | None": + """The canonical accelerator this host can actually use, or ``None``. + + Indices are filled in (``cuda:0`` / ``mps:0``) so the value compares equal + to a device read back off a real tensor -- ``torch.device('mps')`` and + ``torch.device('mps:0')`` are *not* equal even though they name the same + physical device. + """ + if _cuda_available(): + return torch.device("cuda", torch.cuda.current_device()) + if _mps_available(): + return torch.device("mps", 0) + return None + + +# Built at import time so the ``gpu`` mark is attached during *collection*. +# Adding it later (e.g. via ``request.node.add_marker`` inside the fixture) is +# too late for ``pytest_collection_modifyitems`` to gate on. +_DEVICE_PARAMS = [pytest.param(torch.device("cpu"), id="cpu")] +_ACCELERATOR = _accelerator() +if _ACCELERATOR is not None: + _DEVICE_PARAMS.append( + pytest.param( + _ACCELERATOR, + id=_ACCELERATOR.type, + # Backend-specific mark, so a CUDA-less host skips the cuda leg and + # a non-Mac skips the mps leg, each with an accurate reason. + marks=getattr(pytest.mark, _ACCELERATOR.type), + ) + ) + + +@pytest.fixture(params=_DEVICE_PARAMS) +def any_device(request) -> torch.device: + """Every device this host can actually use, one test run per device. + + The CPU leg always runs. The accelerator leg is ``gpu``-marked, so a plain + ``pytest`` run skips it and ``pytest --run-gpu`` picks up CUDA on a CUDA + box or MPS on a Mac. On a CPU-only host the accelerator parameter does not + exist at all, so there is no skip noise. + """ + return request.param + + +@pytest.fixture(scope="session") +def _device_model_cache() -> dict: + """``{device_str: ModelFT}`` built at most once per device, per session.""" + return {} + + +@pytest.fixture +def device_model_bundle(_device_model_cache, pdb_dir, any_device): + """A loaded model on ``any_device``, for target conformance tests. + + The existing ``loaded_model`` / ``model_and_data`` fixtures are + function-scoped and construct on the process default, so a + device-parametrized sweep over them would reload the structure once per + test per device. This caches one model per device instead. + + Shared mutable state: callers must treat the bundle as read-only. A test + that moves a *target* will drag the borrowed model with it, poisoning every + later test on that device -- see ``test_target_device_round_trip``, which + deliberately builds its own. + """ + key = str(any_device) + if key not in _device_model_cache: + pdb = pdb_dir / "1DAW.pdb" + if not pdb.exists(): + pytest.skip("1DAW.pdb fixture not present") + from torchref.model import ModelFT + + _device_model_cache[key] = ModelFT(device=any_device, verbose=0).load_pdb( + str(pdb) + ) + return {"model": _device_model_cache[key]} @pytest.fixture @@ -171,7 +317,7 @@ def device(request) -> torch.device: markers = {m.name for m in request.node.iter_markers()} if "cuda_only" in markers and not torch.cuda.is_available(): pytest.skip("Test requires CUDA") - if "gpu" in markers and not _gpu_available(): + if "gpu" in markers and not (_cuda_available() or _mps_available()): pytest.skip("No GPU (CUDA or MPS) available") return get_default_device() diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py new file mode 100644 index 00000000..52060c21 --- /dev/null +++ b/tests/helpers/__init__.py @@ -0,0 +1 @@ +"""Shared test helpers for the torchref suite.""" diff --git a/tests/helpers/device_asserts.py b/tests/helpers/device_asserts.py new file mode 100644 index 00000000..20fdd4bf --- /dev/null +++ b/tests/helpers/device_asserts.py @@ -0,0 +1,292 @@ +"""Device-consistency assertions for torchref objects. + +The walker here is written **independently** of +:func:`torchref.utils.device_mixin._apply_to_obj`. That is deliberate: an oracle +that mirrors the implementation it checks can only ever confirm the +implementation's own idea of what is reachable, so it would be blind to exactly +the tensors ``DeviceMixin`` forgets to move. This walker instead enumerates +everything it can reach by generic Python/PyTorch means -- ``__dict__``, +``__slots__``, module parameters/buffers/children, containers, dataclass +fields, and :class:`~torchref.utils.utils.ModuleReference` wrappers. + +Exports +------- +assert_device_consistent + Fail unless every reachable tensor *and* every device tracker sits on the + expected device. Reports all mismatches at once. +collect_device_map + The underlying survey, returned as ``{path: torch.device}``. Useful when a + test wants to assert something more specific than "all on one device". +HOST_SIDE + Paths that are intentionally CPU-resident regardless of the object's + device. Every entry needs a comment justifying it. +""" + +from __future__ import annotations + +import dataclasses +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple + +import torch +from torch import nn + +from torchref.config import canonical_device + +__all__ = [ + "assert_device_consistent", + "collect_device_map", + "HOST_SIDE", +] + + +# Attribute *leaf names* that are legitimately CPU-resident even when the +# owning object lives on an accelerator. Keep this list short and justified -- +# it is the escape hatch that makes the conformance suite meaningful rather +# than a rubber stamp. +# +# NOTE: a registered ``nn.Module`` buffer can never be host-side in practice, +# because ``module.to(device)`` relocates it regardless of what this set says. +# Anything that genuinely must stay on the host belongs in plain Python state +# (or ``get_extra_state``), not in a buffer. If you find yourself wanting to add +# a buffer name here, convert the buffer instead. +HOST_SIDE: Set[str] = set() + + +# Attributes on ``nn.Module`` that hold PyTorch's own bookkeeping. They are +# traversed explicitly via ``named_parameters``/``named_buffers``/ +# ``named_children``, so walking the raw dicts as well would only duplicate work +# and produce confusing double paths. +_MODULE_INTERNALS = frozenset( + { + "_parameters", + "_buffers", + "_modules", + "_non_persistent_buffers_set", + "_backward_hooks", + "_backward_pre_hooks", + "_forward_hooks", + "_forward_hooks_with_kwargs", + "_forward_pre_hooks", + "_forward_pre_hooks_with_kwargs", + "_state_dict_hooks", + "_state_dict_pre_hooks", + "_load_state_dict_pre_hooks", + "_load_state_dict_post_hooks", + "_is_full_backward_hook", + "training", + } +) + +_TRACKER_ATTRS = ("device", "_device") +_DTYPE_TRACKER_ATTRS = ("dtype_float", "_dtype") + +_MAX_DEPTH = 12 + + +def _slot_names(obj: Any) -> Iterable[str]: + """Every ``__slots__`` entry declared anywhere in ``type(obj)``'s MRO.""" + for klass in type(obj).__mro__: + for name in getattr(klass, "__slots__", ()) or (): + yield name + + +def _iter_attributes(obj: Any) -> Iterable[Tuple[str, Any]]: + """Yield ``(name, value)`` for the object's own state. + + Covers ``__dict__`` (skipping ``nn.Module`` bookkeeping), ``__slots__``, and + dataclass fields -- a dataclass with ``__slots__`` exposes neither through + ``__dict__``, which is how ``_ReflectionSubset``-style views hide tensors + from a naive walk. + """ + seen: Set[str] = set() + + for name, value in list(getattr(obj, "__dict__", {}).items()): + if name in _MODULE_INTERNALS: + continue + seen.add(name) + yield name, value + + for name in _slot_names(obj): + if name in seen or name.startswith("__"): + continue + seen.add(name) + try: + yield name, getattr(obj, name) + except AttributeError: + continue # declared slot, never assigned + + if dataclasses.is_dataclass(obj) and not isinstance(obj, type): + for field in dataclasses.fields(obj): + if field.name in seen: + continue + seen.add(field.name) + try: + yield field.name, getattr(obj, field.name) + except AttributeError: + continue + + +def _walk( + obj: Any, + path: str, + found: Dict[str, torch.device], + trackers: Dict[str, torch.device], + dtypes: Dict[str, torch.dtype], + dtype_trackers: Dict[str, torch.dtype], + visited: Set[int], + depth: int, +) -> None: + if obj is None or depth > _MAX_DEPTH: + return + if isinstance(obj, (str, bytes, int, float, bool, complex, torch.dtype, torch.device)): + return + + if isinstance(obj, torch.Tensor): + found[path] = obj.device + # Only floating/complex tensors carry the configured float dtype; + # integer index buffers and boolean masks have their own and must not + # be swept into the comparison. + if obj.is_floating_point() or obj.is_complex(): + dtypes[path] = obj.dtype + return + + if id(obj) in visited: + return + visited.add(id(obj)) + + args = (found, trackers, dtypes, dtype_trackers, visited) + + if isinstance(obj, nn.Module): + for name, tensor in list(obj.named_parameters(recurse=False)) + list( + obj.named_buffers(recurse=False) + ): + if tensor is None: + continue + found[f"{path}.{name}"] = tensor.device + if tensor.is_floating_point() or tensor.is_complex(): + dtypes[f"{path}.{name}"] = tensor.dtype + for name, child in obj.named_children(): + _walk(child, f"{path}.{name}", *args, depth + 1) + + # ``ModuleReference`` deliberately hides its payload from PyTorch's module + # tree, so a parameters/buffers sweep never sees it -- but the referenced + # object still has to agree on a device with the referrer. + wrapped = getattr(obj, "_wrapped_module", None) + if wrapped is not None and not isinstance(obj, nn.Module): + _walk(wrapped, f"{path}->ref", *args, depth + 1) + return + + for name, value in _iter_attributes(obj): + if name in _TRACKER_ATTRS and isinstance(value, (torch.device, str)): + try: + trackers[f"{path}.{name}"] = torch.device(value) + except (RuntimeError, TypeError): + pass # a ``device``-named attribute that isn't a device + continue + if name in _DTYPE_TRACKER_ATTRS and isinstance(value, torch.dtype): + dtype_trackers[f"{path}.{name}"] = value + continue + if name in HOST_SIDE: + continue + _walk(value, f"{path}.{name}", *args, depth + 1) + + if isinstance(obj, dict): + for key, value in list(obj.items()): + _walk(value, f"{path}[{key!r}]", *args, depth + 1) + elif isinstance(obj, (list, tuple, set, frozenset)): + for index, value in enumerate(obj): + _walk(value, f"{path}[{index}]", *args, depth + 1) + + +def collect_device_map(obj: Any, name: str = "obj") -> Tuple[Dict, Dict, Dict, Dict]: + """Survey ``obj`` along both the device and dtype axes. + + Returns ``(tensor_devices, tracker_devices, tensor_dtypes, tracker_dtypes)``, + each ``{path: value}``. Paths are dotted attribute chains, with ``[i]`` / + ``['k']`` for container elements and ``->ref`` where the walk crossed a + :class:`ModuleReference`. + + The dtype maps cover only floating/complex tensors and the + ``dtype_float`` / ``_dtype`` trackers, mirroring the mixin's rule that the + two axes are resolved independently -- an integer index buffer says nothing + about the configured float dtype. + """ + tensors: Dict[str, torch.device] = {} + trackers: Dict[str, torch.device] = {} + tensor_dtypes: Dict[str, torch.dtype] = {} + tracker_dtypes: Dict[str, torch.dtype] = {} + _walk(obj, name, tensors, trackers, tensor_dtypes, tracker_dtypes, set(), 0) + return tensors, trackers, tensor_dtypes, tracker_dtypes + + +def assert_device_consistent( + obj: Any, + expected: Any, + *, + name: str = "obj", + ignore: Optional[Iterable[str]] = None, + expected_dtype: Optional[torch.dtype] = None, +) -> None: + """Assert every reachable tensor and tracker sits on ``expected``. + + Parameters + ---------- + obj + Object to survey. + expected + Target device. Compared through + :func:`torchref.config.canonical_device` on both sides, so a bare + ``mps`` and an indexed ``mps:0`` are treated as equal -- the helper + tests placement, not spelling. + name + Root label used to build the reported paths. + ignore + Path substrings to exclude. Use for genuinely borrowed state, not to + paper over a split object. + expected_dtype + When given, also assert every floating/complex tensor and every + ``dtype_float`` / ``_dtype`` tracker matches. This is the only way to + catch the ``torch.tensor(float(x))`` class of bug, where a tunable is + hardcoded float32 under a float64 configuration. + + Raises + ------ + AssertionError + Listing **every** mismatch. A half-moved object usually has several, + and reporting only the first turns one debugging session into five. + """ + expected_dev = canonical_device(expected) + ignore = tuple(ignore or ()) + + tensors, trackers, tensor_dtypes, tracker_dtypes = collect_device_map(obj, name) + problems: List[str] = [] + + def skipped(path: str) -> bool: + return any(frag in path for frag in ignore) + + for path, dev in sorted(tensors.items()): + if not skipped(path) and canonical_device(dev) != expected_dev: + problems.append(f"{path}: tensor on {dev}, expected {expected_dev}") + + for path, dev in sorted(trackers.items()): + if not skipped(path) and canonical_device(dev) != expected_dev: + problems.append(f"{path}: tracker == {dev}, expected {expected_dev}") + + if expected_dtype is not None: + for path, dt in sorted(tensor_dtypes.items()): + if not skipped(path) and dt != expected_dtype: + problems.append(f"{path}: tensor dtype {dt}, expected {expected_dtype}") + for path, dt in sorted(tracker_dtypes.items()): + if not skipped(path) and dt != expected_dtype: + problems.append( + f"{path}: dtype tracker == {dt}, expected {expected_dtype}" + ) + + if problems: + target = f"{expected_dev}" + if expected_dtype is not None: + target += f" / {expected_dtype}" + raise AssertionError( + f"{type(obj).__name__} is not consistently on {target} " + f"({len(problems)} mismatch(es)):\n " + "\n ".join(problems) + ) diff --git a/tests/helpers/device_cases.py b/tests/helpers/device_cases.py new file mode 100644 index 00000000..321a982b --- /dev/null +++ b/tests/helpers/device_cases.py @@ -0,0 +1,337 @@ +"""Registry of cheaply-constructible device-bearing objects for conformance tests. + +Each :class:`DeviceCase` builds one object on a requested device. The suite in +``tests/unit/test_device_conformance.py`` then asserts that everything the +object owns actually landed there, and that it survives a device round trip. + +``UNCOVERED`` is the other half of the contract: every device-bearing class in +the source tree must appear either here or there, with a reason. That is what +keeps the registry from quietly rotting as the package grows -- see +``test_every_device_bearing_class_is_accounted_for``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List + +import torch + +__all__ = [ + "DeviceCase", + "CASES", + "TargetDeviceCase", + "TARGET_CASES", + "UNCOVERED", + "UNCOVERED_PREFIXES", +] + +_CELL = [50.0, 60.0, 70.0, 90.0, 90.0, 90.0] +_SG = "P 21 21 21" + + +@dataclass(frozen=True) +class DeviceCase: + """One constructible object under device test. + + Parameters + ---------- + name + Test id. + build + ``device -> object``. Must place everything on ``device``. + cls_name + Class this case covers, for the coverage guard. + tensor_free + True when the object legitimately owns no tensors (an empty shell). + Such objects can only be checked via their tracker. + ignore + Path fragments to exclude from the walk (borrowed state). + """ + + name: str + build: Callable[[torch.device], Any] + cls_name: str + tensor_free: bool = False + ignore: tuple = field(default_factory=tuple) + + +def _cell(device): + from torchref.symmetry import Cell + + return Cell(_CELL, device=device) + + +CASES: List[DeviceCase] = [ + DeviceCase("Cell", _cell, "Cell"), + DeviceCase( + "SpaceGroup", + lambda d: __import__( + "torchref.symmetry", fromlist=["SpaceGroup"] + ).SpaceGroup(_SG, device=d), + "SpaceGroup", + ), + # D4: device implied by the cell; the SpaceGroup used to be built from the + # raw (None) device argument and land on the process default instead. + DeviceCase( + "SfFFT_from_cell", + lambda d: __import__( + "torchref.model.sf_fft", fromlist=["SfFFT"] + ).SfFFT(cell=_cell(d), spacegroup=_SG, max_res=2.0), + "SfFFT", + ), + # D4: explicit device disagreeing with the supplied cell. + DeviceCase( + "SfFFT_explicit_device", + lambda d: __import__( + "torchref.model.sf_fft", fromlist=["SfFFT"] + ).SfFFT(cell=_cell("cpu"), spacegroup=_SG, max_res=2.0, device=d), + "SfFFT", + ), + DeviceCase( + "SfDS_from_cell", + lambda d: __import__( + "torchref.model.sf_ds", fromlist=["SfDS"] + ).SfDS(cell=_cell(d), spacegroup=_SG), + "SfDS", + ), + # D1: tensor-free shells, whose tracker is the only thing to check. + DeviceCase( + "ScalerBase_empty", + lambda d: __import__( + "torchref.scaling", fromlist=["ScalerBase"] + ).ScalerBase(device=d), + "ScalerBase", + tensor_free=True, + ), + DeviceCase( + "LossState", + lambda d: __import__( + "torchref.refinement.loss_state", fromlist=["LossState"] + ).LossState(device=d), + "LossState", + tensor_free=True, + ), + DeviceCase( + "SolventModel", + lambda d: __import__( + "torchref.scaling.solvent", fromlist=["SolventModel"] + ).SolventModel(device=d), + "SolventModel", + ), + # D9: empty-init wrappers used to drop the requested device entirely. + DeviceCase( + "MixedTensor_empty", + lambda d: __import__( + "torchref.model.parameter_wrappers", fromlist=["MixedTensor"] + ).MixedTensor(device=d), + "MixedTensor", + ), + DeviceCase( + "MixedTensor_populated", + lambda d: __import__( + "torchref.model.parameter_wrappers", fromlist=["MixedTensor"] + ).MixedTensor( + torch.arange(12, dtype=torch.float32).reshape(4, 3), + # A 2D value tensor takes a 1D per-row mask. + refinable_mask=torch.zeros(4, dtype=torch.bool), + device=d, + ), + "MixedTensor", + ), + DeviceCase( + "PositiveMixedTensor", + lambda d: __import__( + "torchref.model.parameter_wrappers", fromlist=["PositiveMixedTensor"] + ).PositiveMixedTensor( + torch.arange(1, 5, dtype=torch.float32), + refinable_mask=torch.zeros(4, dtype=torch.bool), + device=d, + ), + "PositiveMixedTensor", + ), + DeviceCase( + "OccupancyTensor_empty", + lambda d: __import__( + "torchref.model.parameter_wrappers", fromlist=["OccupancyTensor"] + ).OccupancyTensor(device=d), + "OccupancyTensor", + ), + DeviceCase( + "RigidXYZTensor_empty", + lambda d: __import__( + "torchref.model.rigid_xyz", fromlist=["RigidXYZTensor"] + ).RigidXYZTensor(device=d), + "RigidXYZTensor", + ), + DeviceCase( + "ReciprocalSymmetryGrid", + lambda d: __import__( + "torchref.symmetry", fromlist=["ReciprocalSymmetryGrid"] + ).ReciprocalSymmetryGrid(_SG, grid_shape=(16, 16, 16), device=d), + "ReciprocalSymmetryGrid", + ), + DeviceCase( + "MapSymmetryDirect", + lambda d: __import__( + "torchref.symmetry", fromlist=["MapSymmetryDirect"] + ).MapSymmetryDirect( + _SG, map_shape=(16, 16, 16), cell_params=_CELL, device=d + ), + "MapSymmetryDirect", + ), + DeviceCase( + "TensorMasks", + lambda d: __import__( + "torchref.utils", fromlist=["TensorMasks"] + ).TensorMasks(device=d), + "TensorMasks", + tensor_free=True, + ), + DeviceCase( + "ReflectionData_empty", + lambda d: __import__( + "torchref.io", fromlist=["ReflectionData"] + ).ReflectionData(device=d), + "ReflectionData", + ), + DeviceCase( + "ManualWeighting", + lambda d: __import__( + "torchref.refinement.weighting", fromlist=["ManualWeighting"] + ).ManualWeighting({"geometry": 1.0}, device=d), + "ManualWeighting", + tensor_free=True, + ), +] + + +@dataclass(frozen=True) +class TargetDeviceCase: + """A refinement target, which needs a real model/data/scaler to construct. + + Separate from :class:`DeviceCase` because ``build`` takes a fixture-supplied + bundle as well as a device. The bundle is already on ``device``; ``build`` + must not cache or mutate it. + + ``ignore`` defaults to the wrapped objects: a target *borrows* its model and + data, and the conformance walk would otherwise re-check the entire structure + through every target that points at it. + """ + + name: str + build: Callable[[Dict[str, Any], torch.device], Any] + cls_name: str + needs: tuple = ("model",) + ignore: tuple = ("_model", "_data", "_scaler", "_model_dark", "_model_light") + + +TARGET_CASES: List[TargetDeviceCase] = [ + TargetDeviceCase( + "NonBondedTarget", + lambda b, d: __import__( + "torchref.refinement.targets.geometry.non_bonded", + fromlist=["NonBondedTarget"], + ).NonBondedTarget(b["model"]), + "NonBondedTarget", + ), + TargetDeviceCase( + "ADPSimilarityTarget", + lambda b, d: __import__( + "torchref.refinement.targets.adp.similarity", + fromlist=["ADPSimilarityTarget"], + ).ADPSimilarityTarget(b["model"]), + "ADPSimilarityTarget", + ), + TargetDeviceCase( + "ADPLocalityTarget", + lambda b, d: __import__( + "torchref.refinement.targets.adp.locality", fromlist=["ADPLocalityTarget"] + ).ADPLocalityTarget(b["model"]), + "ADPLocalityTarget", + ), + # Owns no tensors at all -- the case that exercises the request-driven + # tracker path rather than the owned-tensor path. + TargetDeviceCase( + "BondTarget", + lambda b, d: __import__( + "torchref.refinement.targets.geometry.bonds", fromlist=["BondTarget"] + ).BondTarget(b["model"]), + "BondTarget", + ), +] + + +# Device-bearing classes deliberately not in CASES. Every entry needs a reason; +# "hard to build" is a reason, "didn't get to it" is not. +UNCOVERED: Dict[str, str] = { + # --- abstract / mixin bases: never instantiated directly ----------------- + "Target": "abstract base; covered through its concrete subclasses", + "ModelTarget": "abstract base; needs a loaded model", + "DataTarget": "abstract base; needs model + data + scaler", + "XrayTarget": "abstract base; needs model + data + scaler", + "GeometryTarget": "abstract base; needs a model with restraints", + "CrystalDataset": "abstract dataclass base; covered via ReflectionData", + "BaseWeighting": "abstract base; covered via ManualWeighting", + "Refinement": "abstract base; covered via LBFGSRefinement in integration", + "PassThroughTensor": "documented non-functional stub (parameter_wrappers.py)", + "ADPTarget": "abstract base; needs a model with ADPs", + "CombinedTargets": "composite container; needs its component targets", + "CombinedModelTargets": "composite container; needs a loaded model", + "CollectionXrayTarget": "abstract base; needs a dataset collection", + # --- need loaded structures/data: covered by integration tests ----------- + "Model": "needs a PDB; covered by tests/unit/model/test_model_state_dict_device.py", + "ModelFT": "needs a PDB; covered by tests/unit/model/test_model_state_dict_device.py", + "MixedModel": "needs two loaded models", + "ModelCollection": "needs several loaded models", + "_SharedMixedModel": "internal view owned by ModelCollection", + "Scaler": "needs a loaded model + data; covered in integration", + "RestraintsNew": "needs a model + monomer library", + "HydrogenTopology": "needs a built restraint topology", + "FrenchWilson": "needs loaded intensities", + "DatasetCollection": "needs several loaded datasets", + "FcalcDataset": "needs computed structure factors", + "Map": "needs data + model", + "DifferenceMap": "needs two datasets + a model", + "LBFGSRefinement": "full pipeline; covered in integration", + "MapSymmetry": "interpolation variant; needs a real map grid", + "CholeskyMixedTensor": "needs a valid ADP tensor; shares MixedTensor's paths", + "CollectionScaler": "needs a dataset collection", + "CollectionDifferenceTarget": "needs a dataset collection", + "CollectionMLTarget": "needs a dataset collection", + "CollectionRiceTarget": "needs a dataset collection", + "ADPEntropyTarget": "needs a model with ADPs", + "AngleTarget": "needs a model with restraints", + "ChiralTarget": "needs a model with restraints", + "BhattacharyyaXrayTarget": "needs a scaler with initialised bins; buffers audited", + "ReciprocalSymmetryExtractor": "needs an hkl tensor + SpaceGroup; device is " + "intentionally inherited from hkl, so the generic 'requested device' " + "assertion does not apply", + "CoordinateSimilarityTarget": "needs two loaded models", + "MultiModelADPTarget": "needs a model collection", + "MultiModelGeometryTarget": "needs a model collection", + "TotalGeometryTarget": "composite; needs a model with restraints", + "TotalADPTarget": "composite; needs a model with restraints", + "NonBondedHTarget": "needs a model with hydrogen topology", + "PlanarityTarget": "needs a model with restraints", + "RamachandranTarget": "needs a model with restraints", + "TorsionTarget": "needs a model with restraints", + "RigidBondTarget": "needs a model with restraints", + "ScalerLogScaleTrendTarget": "needs a scaler", + "ScalerURegularizationTarget": "needs a scaler", + "GaussianXrayTarget": "needs model + data + scaler", + "LeastSquaresXrayTarget": "needs model + data + scaler", + "MaximumLikelihoodXrayTarget": "needs model + data + scaler", + "RiceXrayTarget": "needs model + data + scaler", + "DifferenceXrayTarget": "needs two datasets", + "PhaseInformedDifferenceTarget": "needs two datasets + phases", + "RiceDifferenceTarget": "needs two datasets", + "TaylorCorrectedDifferenceTarget": "needs two datasets", + "RigidTransform": "alignment helper; needs a coordinate set", + "RigidBodyRefinement": "experimental; needs model + data", +} + +# Everything under torchref/experimental is out of scope for the conformance +# sweep: it is explicitly unstable, and several modules need optional +# dependencies (OpenMM, AmberTools, JAX) that are not installed in CI. +UNCOVERED_PREFIXES = ("torchref/experimental/",) diff --git a/tests/helpers/device_inventory.py b/tests/helpers/device_inventory.py new file mode 100644 index 00000000..a404ac8e --- /dev/null +++ b/tests/helpers/device_inventory.py @@ -0,0 +1,70 @@ +"""Static inventory of ``DeviceMixin`` subclasses in the torchref source tree. + +Why AST rather than ``DeviceMixin.__subclasses__()``: the runtime hook only +sees classes whose defining module has been imported, so a coverage guard built +on it silently shrinks whenever an import is removed or an optional dependency +is missing -- exactly when you most want to be told about a gap. Parsing the +source finds every class whether or not it is importable here. + +Inheritance is resolved **transitively**: ``PositiveMixedTensor(MixedTensor)`` +never names a mixin alias itself, but it is still a device-bearing class and +still needs conformance coverage. +""" + +from __future__ import annotations + +import ast +from pathlib import Path +from typing import Dict, Set, Tuple + +__all__ = ["MIXIN_ALIASES", "device_mixin_classes"] + +# Every spelling of the one mixin. ``DeviceMovementMixin`` and +# ``_NonModuleDeviceMixin`` are aliases kept for backward compatibility. +MIXIN_ALIASES = frozenset( + {"DeviceMixin", "DeviceMovementMixin", "_NonModuleDeviceMixin"} +) + + +def _base_names(node: ast.ClassDef) -> Set[str]: + """Base-class names for ``node``, ignoring subscripts and keywords.""" + names = set() + for base in node.bases: + if isinstance(base, ast.Name): + names.add(base.id) + elif isinstance(base, ast.Attribute): + names.add(base.attr) + return names + + +def device_mixin_classes(package_root: Path) -> Dict[str, str]: + """Return ``{class_name: "relative/path.py:lineno"}`` for device-bearing classes. + + A class qualifies if any of its bases is a mixin alias, or is another + qualifying class -- iterated to a fixed point so multi-level hierarchies + are captured. + """ + definitions: Dict[str, Tuple[Set[str], str]] = {} + + for path in sorted(package_root.rglob("*.py")): + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except (SyntaxError, UnicodeDecodeError): # pragma: no cover + continue + rel = path.relative_to(package_root.parent) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + definitions[node.name] = (_base_names(node), f"{rel}:{node.lineno}") + + qualifying: Dict[str, str] = {} + changed = True + while changed: + changed = False + for name, (bases, where) in definitions.items(): + if name in qualifying: + continue + if bases & MIXIN_ALIASES or bases & qualifying.keys(): + qualifying[name] = where + changed = True + + return qualifying diff --git a/tests/integration/test_cli_difference_refine_mismatched_mtz.py b/tests/integration/test_cli_difference_refine_mismatched_mtz.py new file mode 100644 index 00000000..a1a2723b --- /dev/null +++ b/tests/integration/test_cli_difference_refine_mismatched_mtz.py @@ -0,0 +1,110 @@ +"""CPU end-to-end test for ``torchref.difference-refine`` with MISMATCHED MTZ files. + +Difference refinement builds a ``DatasetCollection`` from two independent MTZ +files and aligns the light dataset onto the dark reference via +``ReflectionData.validate_hkl``. When the two files had different reflection +sets, alignment left ``hkl_anomalous`` (read by ``hkl_for_sf``) at the +pre-alignment length and the run crashed inside the collection difference target +(``RuntimeError: The size of tensor a (...) must match the size of tensor b``). + +This test generates two MTZ files with genuinely different reflection sets from +the smallest fixture (3GR5) and asserts the CLI now completes. It is the +end-to-end guard for the 0.6.2 fix; pre-fix it exited non-zero. + +Wall-clock budget: a couple of minutes on CPU (1 macro-cycle, 1 step). +""" + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + + +@pytest.fixture +def diff_refine_cli_script(project_root) -> Path: + script = project_root / "torchref" / "cli" / "collection_difference_refine.py" + if not script.exists(): + pytest.skip(f"difference-refine CLI script not found: {script}") + return script + + +@pytest.fixture +def mismatched_mtz_pair(test_files_dir, tmp_path): + """Two MTZ files sharing a cell/spacegroup but with different reflections. + + Built from 3GR5: ``dark`` drops the last 10% of reflections, ``light`` drops + the first 15%, so each contains reflections the other lacks and the two + counts differ -- exactly the condition that triggered the crash. + """ + pdb_file = test_files_dir / "pdb" / "3GR5.pdb" + mtz_file = test_files_dir / "mtz" / "3GR5.mtz" + if not pdb_file.exists() or not mtz_file.exists(): + pytest.skip("3GR5 test files not found") + + import torch + + from torchref.io.datasets.reflection_data import ReflectionData + + full = ReflectionData(device="cpu", verbose=0).load_mtz(str(mtz_file)) + n = len(full.hkl) + idx = torch.arange(n) + dark = full.__select__(idx < int(n * 0.90)) + light = full.__select__(idx >= int(n * 0.15)) + assert len(dark.hkl) != len(light.hkl), "subsets must differ in size" + + dark_mtz = tmp_path / "dark.mtz" + light_mtz = tmp_path / "light.mtz" + dark.write_mtz(str(dark_mtz)) + light.write_mtz(str(light_mtz)) + return {"pdb": pdb_file, "dark": dark_mtz, "light": light_mtz} + + +@pytest.mark.integration +def test_difference_refine_mismatched_mtz_cpu( + diff_refine_cli_script, mismatched_mtz_pair, tmp_path +): + outdir = tmp_path / "diff_out" + result = subprocess.run( + [ + sys.executable, + str(diff_refine_cli_script), + "-dm", str(mismatched_mtz_pair["pdb"]), + "-lm", str(mismatched_mtz_pair["pdb"]), + "-dsf", str(mismatched_mtz_pair["dark"]), + "-lsf", str(mismatched_mtz_pair["light"]), + "--fraction", "0.3", + "--n-cycles", "1", + "--n-steps", "1", + "--max-iter", "5", + "-o", str(outdir), + "--device", "cpu", + "--verbose", "1", + ], + capture_output=True, + text=True, + timeout=1800, + ) + + if result.returncode != 0: + print("STDOUT:", result.stdout[-3000:]) + print("STDERR:", result.stderr[-3000:]) + + assert result.returncode == 0, ( + f"difference-refine exited with {result.returncode} on mismatched MTZ " + f"files. stderr tail: {result.stderr[-800:]}" + ) + + # The exact stale-tensor shape mismatch must not reappear. + assert "must match the size of tensor" not in result.stderr + + prefix = "fractions_70_30" + summary = outdir / f"{prefix}_summary.json" + diff_mtz = outdir / f"{prefix}_difference_data.mtz" + assert summary.exists(), "summary JSON not written" + assert diff_mtz.exists(), "difference MTZ not written" + + with open(summary) as f: + data = json.load(f) + assert "results" in data and "r_factor_light" in data["results"] diff --git a/tests/integration/test_device_mixin.py b/tests/integration/test_device_mixin.py index 7e586442..788c1dc6 100644 --- a/tests/integration/test_device_mixin.py +++ b/tests/integration/test_device_mixin.py @@ -13,9 +13,6 @@ import pytest import torch -CUDA_AVAILABLE = torch.cuda.is_available() - - def _load_model_ft(pdb_file, mtz_file): """Helper: load a ModelFT and matching reflection data on CPU. @@ -36,6 +33,7 @@ def _load_model_ft(pdb_file, mtz_file): @pytest.mark.integration +@pytest.mark.cuda def test_modelft_cpu_gpu_cpu_sf_round_trip(sample_pdb_file, sample_mtz_file): """CPU -> GPU -> CPU structure-factor round-trip via the unified mixin. @@ -46,9 +44,6 @@ def test_modelft_cpu_gpu_cpu_sf_round_trip(sample_pdb_file, sample_mtz_file): 3. Move back to CPU, recompute Fcalc; verify CPU placement and values match the original CPU result. """ - if not CUDA_AVAILABLE: - pytest.skip("CUDA device not available") - model, data = _load_model_ft(sample_pdb_file, sample_mtz_file) hkl, *_ = data() @@ -66,9 +61,8 @@ def test_modelft_cpu_gpu_cpu_sf_round_trip(sample_pdb_file, sample_mtz_file): assert fcalc_cuda.device.type == "cuda" assert model.device.type == "cuda" assert model.cell.device.type == "cuda", "Cell did not migrate to GPU" - for buf in model.buffers(): - assert buf.device.type == "cuda", "buffer left behind on CPU" - break + for name, buf in model.named_buffers(): + assert buf.device.type == "cuda", f"buffer {name} left behind on CPU" # Values must agree within numerical tolerance after the round-trip. The CPU # (C++ box-separable splat) and GPU (Triton work-queue splat) paths differ diff --git a/tests/integration/test_ds_triton_vs_eager.py b/tests/integration/test_ds_triton_vs_eager.py index 6e68a8a7..373f2d40 100644 --- a/tests/integration/test_ds_triton_vs_eager.py +++ b/tests/integration/test_ds_triton_vs_eager.py @@ -4,14 +4,14 @@ eager backend evaluated in float64 (downcast for comparison), for both the isolated P1 kernels and the full ``SfDS`` symmetry loop. -Markers: ``@pytest.mark.gpu`` (skipped without ``--run-gpu``) and +Markers: ``@pytest.mark.cuda`` (auto-skipped without CUDA) and ``@pytest.mark.integration``. Requires a CUDA device. """ import pytest import torch -pytestmark = [pytest.mark.gpu, pytest.mark.integration] +pytestmark = [pytest.mark.cuda, pytest.mark.integration] _ATOL = 1e-2 _RTOL = 1e-3 diff --git a/tests/integration/test_refinement_pipeline.py b/tests/integration/test_refinement_pipeline.py index 78c1639a..4f385578 100644 --- a/tests/integration/test_refinement_pipeline.py +++ b/tests/integration/test_refinement_pipeline.py @@ -166,5 +166,5 @@ def test_gpu_refinement_setup(self, sample_structure_pair, gpu_device): data.load_mtz(sample_structure_pair["reflections"]) # Everything should be on GPU - assert model.xyz().device.type == 'cuda' - assert data.hkl.device.type == 'cuda' + assert model.xyz().device.type == gpu_device.type + assert data.hkl.device.type == gpu_device.type diff --git a/tests/integration/test_sfds_device.py b/tests/integration/test_sfds_device.py index e12d18ff..8599e793 100644 --- a/tests/integration/test_sfds_device.py +++ b/tests/integration/test_sfds_device.py @@ -38,7 +38,7 @@ def test_sfds_same_device_cpu(): assert torch.isfinite(F.real).all() and torch.isfinite(F.imag).all() -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration def test_sfds_hkl_on_different_device(): """hkl on CPU while the module + atoms are on CUDA must not crash.""" diff --git a/tests/integration/test_triton_vs_eager_targets.py b/tests/integration/test_triton_vs_eager_targets.py index 3267118d..0b795d88 100644 --- a/tests/integration/test_triton_vs_eager_targets.py +++ b/tests/integration/test_triton_vs_eager_targets.py @@ -12,7 +12,7 @@ Toggling is done with the shared ``torchref.utils.use_engine`` context manager (``Engine.EAGER`` vs ``Engine.TRITON``). No process restart needed. -Markers: ``@pytest.mark.gpu`` (skipped by default) and +Markers: ``@pytest.mark.cuda`` (skipped by default) and ``@pytest.mark.integration``. Run on a CUDA box with ``pytest tests/integration/test_triton_vs_eager_targets.py -m gpu``. """ @@ -177,7 +177,7 @@ def _target_names(state): return list(state.targets.keys()) -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration @pytest.mark.parametrize( "target_name", @@ -218,7 +218,7 @@ def test_triton_matches_eager_per_target(target_name, gpu_refinement, gpu_state) _assert_close(target_name, eager, triton, atol, rtol) -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration @pytest.mark.parametrize("target_mode", ["bhattacharyya", "rice", "ls", "gaussian"]) def test_triton_matches_eager_xray_modes(target_mode, _1daw_pair): @@ -255,7 +255,7 @@ def test_triton_matches_eager_xray_modes(target_mode, _1daw_pair): _assert_close(name, eager, triton, atol, rtol) -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration @pytest.mark.parametrize("n_atoms", [4, 5, 6]) def test_planarity_triton_per_atom_sigma(n_atoms): @@ -298,7 +298,7 @@ def test_planarity_triton_per_atom_sigma(n_atoms): assert torch.allclose(ge, gt, atol=1e-4, rtol=1e-3) -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration def test_geometry_degenerate_finite_grads(): """At degenerate geometry, BOTH eager and Triton give finite gradients. @@ -403,7 +403,7 @@ def grad_at(xv): return cos, rel -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration def test_eager_gpu_hessian_iso(tmp_path): """`use_engine(Engine.EAGER)` gives correct GPU second derivatives (iso). @@ -450,7 +450,7 @@ def test_eager_gpu_hessian_iso(tmp_path): dtypes.float, dtypes.complex = f0, c0 -@pytest.mark.gpu +@pytest.mark.cuda @pytest.mark.integration def test_eager_gpu_hessian_aniso(pdb_dir): """Same as above but for anisotropic ADPs (real ANISOU structure).""" diff --git a/tests/integration/test_variable_radius_gpu.py b/tests/integration/test_variable_radius_gpu.py index 23a2c07d..5f3ca855 100644 --- a/tests/integration/test_variable_radius_gpu.py +++ b/tests/integration/test_variable_radius_gpu.py @@ -5,7 +5,7 @@ forward maps must agree to float32 + analytic-vs-gathered-coord tolerance and the gradients (xyz / adp / u / occ) must be parallel. Requires CUDA + Triton. -Markers: ``@pytest.mark.gpu`` (skipped without ``--run-gpu``), ``integration``. +Markers: ``@pytest.mark.cuda`` (auto-skipped without CUDA), ``integration``. """ import math @@ -16,7 +16,7 @@ from torchref.base.electron_density.main import build_electron_density from torchref.utils import Engine, use_engine -pytestmark = [pytest.mark.gpu, pytest.mark.integration] +pytestmark = [pytest.mark.cuda, pytest.mark.integration] def _cell(): diff --git a/tests/integration/test_variable_radius_mps.py b/tests/integration/test_variable_radius_mps.py new file mode 100644 index 00000000..cba7553e --- /dev/null +++ b/tests/integration/test_variable_radius_mps.py @@ -0,0 +1,198 @@ +"""MPS variable-radius density path: native Metal kernels (Engine.AUTO) vs the +portable plain-scatter reference (Engine.EAGER), on the same MPS device. + +Both truncate each atom at its own per-atom radius, so the forward maps must +agree to float32 + truncation-shape tolerance and the gradients (xyz / adp / u / +occ) must be parallel. Requires an Apple-silicon GPU (MPS); skipped elsewhere. + +Markers: ``@pytest.mark.mps`` (auto-skipped without MPS), ``integration``. +""" + +import math + +import pytest +import torch + +from torchref.base.electron_density.kernels.mps import mps_kernels_available +from torchref.base.electron_density.main import build_electron_density +from torchref.utils import Engine, use_engine + +pytestmark = [pytest.mark.mps, pytest.mark.integration] + + +@pytest.fixture +def mps_device(gpu_device): + if gpu_device.type != "mps": + pytest.skip("MPS-specific test (Metal kernels)") + if not mps_kernels_available(): + pytest.skip("Metal splat kernels failed to compile") + return gpu_device + + +def _cell(): + a, b, c = 30.0, 25.0, 20.0 + nx, ny, nz = 60, 50, 40 + frac = torch.diag(torch.tensor([a, b, c])) + inv_frac = torch.diag(torch.tensor([1 / a, 1 / b, 1 / c])) + voxel = torch.tensor([a / nx, b / ny, c / nz]) + ii, jj, kk = torch.meshgrid( + torch.arange(nx), torch.arange(ny), torch.arange(nz), indexing="ij" + ) + rsg = (torch.stack([ii / nx, jj / ny, kk / nz], -1).to(torch.float32)) @ frac.T + return (a, b, c), (nx, ny, nz), frac.float(), inv_frac.float(), voxel.float(), rsg + + +def _cell_monoclinic(): + """Non-orthogonal (beta ~ 100 deg) cell: exercises the per-axis box and the + off-diagonal coordinate math (u_c gains an x-component).""" + a, b, c = 30.0, 25.0, 20.0 + nx, ny, nz = 60, 50, 40 + beta = math.radians(100.0) + frac = torch.tensor( + [[a, 0.0, c * math.cos(beta)], + [0.0, b, 0.0], + [0.0, 0.0, c * math.sin(beta)]], dtype=torch.float64) + inv_frac = torch.linalg.inv(frac) + voxel = (frac.norm(dim=0) / torch.tensor([nx, ny, nz], dtype=torch.float64)).float() + ii, jj, kk = torch.meshgrid( + torch.arange(nx), torch.arange(ny), torch.arange(nz), indexing="ij" + ) + fc = torch.stack([ii / nx, jj / ny, kk / nz], -1).double() + rsg = (fc @ frac.T).float() + return (a, b, c), (nx, ny, nz), frac.float(), inv_frac.float(), voxel, rsg + + +def _iso_atoms(cell, n=60, seed=0): + g = torch.Generator().manual_seed(seed) + a, b, c = cell + xyz = torch.rand(n, 3, generator=g) * torch.tensor([a, b, c]) + adp = torch.rand(n, generator=g) * 35 + 3 + occ = torch.ones(n) + A = torch.rand(n, 5, generator=g) * 5 + B = torch.rand(n, 5, generator=g) * 20 + 2 + return [t.float() for t in (xyz, adp, occ, A, B)] + + +def _aniso_atoms(cell, n=30, seed=1): + g = torch.Generator().manual_seed(seed) + a, b, c = cell + xyz = torch.rand(n, 3, generator=g) * torch.tensor([a, b, c]) + u = torch.zeros(n, 6) + u[:, :3] = torch.rand(n, 3, generator=g) * 0.12 + 0.02 + occ = torch.ones(n) + A = torch.rand(n, 5, generator=g) * 5 + B = torch.rand(n, 5, generator=g) * 20 + 2 + return [t.float() for t in (xyz, u, occ, A, B)] + + +def _cos(a, b): + a, b = a.reshape(-1).double(), b.reshape(-1).double() + return float((a @ b) / (a.norm() * b.norm() + 1e-12)) + + +def _rel_map(x, y): + return float((x - y).abs().max() / (y.abs().max() + 1e-8)) + + +def _to(dev, *ts): + return [t.to(dev) for t in ts] + + +def test_iso_metal_matches_plain(mps_device): + _, grid, frac, inv_frac, voxel, rsg = _cell() + xyz, adp, occ, A, B = _to(mps_device, *_iso_atoms(_cell()[0])) + rsg, frac, inv_frac, voxel = _to(mps_device, rsg, frac, inv_frac, voxel) + + def run(engine): + with use_engine(engine): + return build_electron_density(rsg, xyz, adp, occ, A, B, inv_frac, frac, voxel) + + ref = run(Engine.EAGER) # portable plain splat + got = run(Engine.AUTO) # Metal + assert _rel_map(got.cpu(), ref.cpu()) < 2e-2 + assert _cos(got.cpu(), ref.cpu()) > 0.9995 + + +def test_iso_metal_matches_plain_monoclinic(mps_device): + _, grid, frac, inv_frac, voxel, rsg = _cell_monoclinic() + xyz, adp, occ, A, B = _to(mps_device, *_iso_atoms(_cell_monoclinic()[0])) + rsg, frac, inv_frac, voxel = _to(mps_device, rsg, frac, inv_frac, voxel) + + def run(engine): + with use_engine(engine): + return build_electron_density(rsg, xyz, adp, occ, A, B, inv_frac, frac, voxel) + + ref = run(Engine.EAGER) + got = run(Engine.AUTO) + assert _rel_map(got.cpu(), ref.cpu()) < 2e-2 + assert _cos(got.cpu(), ref.cpu()) > 0.9995 + + +def test_aniso_metal_matches_plain(mps_device): + _, grid, frac, inv_frac, voxel, rsg = _cell() + xa, ua, oa, Aa, Ba = _to(mps_device, *_aniso_atoms(_cell()[0])) + rsg, frac, inv_frac, voxel = _to(mps_device, rsg, frac, inv_frac, voxel) + empty = torch.zeros(0, 3, device=mps_device) + z0 = torch.zeros(0, device=mps_device) + z05 = torch.zeros(0, 5, device=mps_device) + + def run(engine): + with use_engine(engine): + return build_electron_density( + rsg, empty, z0, z0, z05, z05, inv_frac, frac, voxel, + xyz_aniso=xa, u_aniso=ua, occ_aniso=oa, A_aniso=Aa, B_aniso=Ba, + ) + + ref = run(Engine.EAGER) + got = run(Engine.AUTO) + assert _rel_map(got.cpu(), ref.cpu()) < 2e-2 + assert _cos(got.cpu(), ref.cpu()) > 0.9995 + + +def test_iso_gradients_match_plain(mps_device): + _, grid, frac, inv_frac, voxel, rsg = _cell() + xyz0, adp0, occ0, A, B = _iso_atoms(_cell()[0]) + rsg, frac, inv_frac, voxel, A, B = _to(mps_device, rsg, frac, inv_frac, voxel, A, B) + w = torch.randn(grid, device=mps_device) + + def run(engine): + xx = xyz0.to(mps_device).clone().requires_grad_() + aa = adp0.to(mps_device).clone().requires_grad_() + oo = occ0.to(mps_device).clone().requires_grad_() + with use_engine(engine): + dm = build_electron_density(rsg, xx, aa, oo, A, B, inv_frac, frac, voxel) + (dm * w).sum().backward() + return xx.grad, aa.grad, oo.grad + + gx_r, ga_r, go_r = run(Engine.EAGER) + gx_m, ga_m, go_m = run(Engine.AUTO) + assert _cos(gx_m.cpu(), gx_r.cpu()) > 0.999 + assert _cos(ga_m.cpu(), ga_r.cpu()) > 0.999 + assert _cos(go_m.cpu(), go_r.cpu()) > 0.999 + + +def test_aniso_gradients_match_plain(mps_device): + _, grid, frac, inv_frac, voxel, rsg = _cell() + xa0, ua0, oa0, Aa, Ba = _aniso_atoms(_cell()[0]) + rsg, frac, inv_frac, voxel, Aa, Ba = _to(mps_device, rsg, frac, inv_frac, voxel, Aa, Ba) + empty = torch.zeros(0, 3, device=mps_device) + z0 = torch.zeros(0, device=mps_device) + z05 = torch.zeros(0, 5, device=mps_device) + w = torch.randn(grid, device=mps_device) + + def run(engine): + xx = xa0.to(mps_device).clone().requires_grad_() + uu = ua0.to(mps_device).clone().requires_grad_() + with use_engine(engine): + dm = build_electron_density( + rsg, empty, z0, z0, z05, z05, inv_frac, frac, voxel, + xyz_aniso=xx, u_aniso=uu, occ_aniso=oa0.to(mps_device), + A_aniso=Aa, B_aniso=Ba, + ) + (dm * w).sum().backward() + return xx.grad, uu.grad + + gx_r, gu_r = run(Engine.EAGER) + gx_m, gu_m = run(Engine.AUTO) + assert _cos(gx_m.cpu(), gx_r.cpu()) > 0.999 + assert _cos(gu_m.cpu(), gu_r.cpu()) > 0.999 diff --git a/tests/pytest.ini b/tests/pytest.ini index c8c990ef..8896f04f 100644 --- a/tests/pytest.ini +++ b/tests/pytest.ini @@ -12,17 +12,25 @@ pythonpath = .. markers = unit: Unit tests (fast, no I/O, isolated) integration: Integration tests (slower, real I/O, pipelines) - gpu: Tests requiring CUDA GPU (skipped by default) + gpu: Needs any accelerator (CUDA or MPS); auto-skipped if the host has none + cuda: Needs CUDA specifically (e.g. Triton kernels); auto-skipped if absent + mps: Needs MPS specifically (Metal kernels); auto-skipped if absent + cuda_only: Deprecated alias for 'cuda' slow: Long-running tests (skipped in quick runs) openmm: Tests requiring OpenMM (the [amber] extra); skipped if absent amber: Tests requiring OpenMM + AmberTools (antechamber/tleap); skipped if absent -# Default options (exclude GPU tests by default) +# Default options. +# +# Accelerator tests are gated by ``pytest_collection_modifyitems`` in +# conftest.py, which runs whatever the host supports and skips the rest with a +# reason naming the missing backend. Do NOT add a default ``-m`` expression +# here: a mark expression *deselects* at collection time, which no flag can +# undo, so the gating would silently stop working. addopts = -v --tb=short --strict-markers - -m "not gpu" -ra # Ignore warnings from dependencies diff --git a/tests/unit/base/test_kernel_import_shims.py b/tests/unit/base/test_kernel_import_shims.py index 6f1bb2c6..4bb8fc7f 100644 --- a/tests/unit/base/test_kernel_import_shims.py +++ b/tests/unit/base/test_kernel_import_shims.py @@ -11,6 +11,12 @@ import pytest +from torchref.utils.triton_dispatch import triton_available + +_needs_triton = pytest.mark.skipif( + not triton_available(), reason="CUDA/Triton backend requires the triton package" +) + def test_kernels_public_api_resolves(): kernels = importlib.import_module("torchref.base.kernels") @@ -90,8 +96,21 @@ def test_solvent_radius_offsets_import_path(): "torchref.base.electron_density.kernels.cpu.scatter_dispatch", "torchref.base.electron_density.kernels.cpu.jit_reference", "torchref.base.electron_density.kernels.cpu.variable_radius", - "torchref.base.electron_density.kernels.cuda.fused", - "torchref.base.electron_density.kernels.cuda.variable_radius", + # CUDA/Triton backend: importing pulls in `triton`, absent on non-CUDA + # hosts (e.g. macOS), so gate these on triton availability. + pytest.param( + "torchref.base.electron_density.kernels.cuda.fused", marks=_needs_triton + ), + pytest.param( + "torchref.base.electron_density.kernels.cuda.variable_radius", + marks=_needs_triton, + ), + # MPS Metal kernels: importing must not trigger shader compilation, so + # these resolve cleanly on every platform (compile is deferred to first use). + "torchref.base.electron_density.kernels.mps", + "torchref.base.electron_density.kernels.mps.compile", + "torchref.base.electron_density.kernels.mps.variable_radius", + "torchref.base.electron_density.kernels.mps._shaders", ], ) def test_new_kernel_modules_import(modname): diff --git a/tests/unit/base/test_math_torch.py b/tests/unit/base/test_math_torch.py index 9ac12c10..3bf19be0 100644 --- a/tests/unit/base/test_math_torch.py +++ b/tests/unit/base/test_math_torch.py @@ -92,8 +92,8 @@ def test_coordinate_transforms_gpu(self, mock_cell, random_coordinates, gpu_devi frac = cartesian_to_fractional_torch(coords, cell) cart_back = fractional_to_cartesian_torch(frac, cell) - assert frac.device.type == "cuda" - assert cart_back.device.type == "cuda" + assert frac.device.type == gpu_device.type + assert cart_back.device.type == gpu_device.type class TestGridFunctions: diff --git a/tests/unit/base/test_max_eig_sym3.py b/tests/unit/base/test_max_eig_sym3.py new file mode 100644 index 00000000..3bc1da5d --- /dev/null +++ b/tests/unit/base/test_max_eig_sym3.py @@ -0,0 +1,72 @@ +"""Closed-form largest eigenvalue of a symmetric 3x3 vs torch.linalg.eigvalsh. + +``_max_eig_sym3`` replaces ``eigvalsh`` in the anisotropic splat-radius policy so +it runs natively on MPS; it must agree with the reference eigendecomposition, +including on diagonal and near-degenerate matrices. +""" + +import pytest +import torch + +from torchref.base.electron_density.radius_policy import ( + _max_eig_sym3, + _u6_to_u3, + per_atom_radius_aniso, +) + +pytestmark = pytest.mark.unit + + +def _sym(n, seed, scale=1.0): + g = torch.Generator().manual_seed(seed) + M = torch.randn(n, 3, 3, generator=g) * scale + return (M + M.transpose(1, 2)) / 2 + + +def test_matches_eigvalsh_random(): + S = _sym(4000, 0) + ref = torch.linalg.eigvalsh(S).max(dim=1).values + got = _max_eig_sym3(S) + assert torch.allclose(got, ref, atol=1e-4, rtol=1e-4) + + +def test_matches_eigvalsh_diagonal(): + # purely diagonal (off-diagonals zero) -> largest diagonal entry + d = torch.tensor([[3.0, -1.0, 2.0], [5.0, 5.0, 5.0], [0.1, 0.2, 0.05]]) + S = torch.diag_embed(d) + ref = torch.linalg.eigvalsh(S).max(dim=1).values + got = _max_eig_sym3(S) + assert torch.allclose(got, ref, atol=1e-5) + assert torch.allclose(got, d.max(dim=1).values, atol=1e-5) + + +def test_matches_eigvalsh_near_degenerate(): + # near-isotropic (all eigenvalues nearly equal) stresses the p2->0 branch + S = torch.eye(3).expand(50, 3, 3).clone() + S += _sym(50, 7, scale=1e-4) + ref = torch.linalg.eigvalsh(S).max(dim=1).values + got = _max_eig_sym3(S) + assert torch.allclose(got, ref, atol=1e-4) + + +def test_radius_matches_eigvalsh_path(): + # The (quantized) aniso radius must be identical to the eigvalsh-based one. + g = torch.Generator().manual_seed(3) + n = 500 + u = torch.zeros(n, 6) + u[:, :3] = torch.rand(n, 3, generator=g) * 0.12 + 0.02 + u[:, 3:] = (torch.rand(n, 3, generator=g) - 0.5) * 0.02 + B = torch.rand(n, 5, generator=g) * 20 + 2 + + rad_closed = per_atom_radius_aniso(B, u, n_sigma=3.0) + + # Reference radius using eigvalsh directly. + from torchref.base.electron_density.radius_policy import ( + EIGHT_PI2, _ceil_round, R_LO, R_HI, + ) + b_form = B.max(dim=1).values + lam = torch.linalg.eigvalsh(_u6_to_u3(u)).max(dim=1).values + sigma = torch.sqrt((b_form / EIGHT_PI2 + lam).clamp(min=1e-6)) + rad_ref = _ceil_round(3.0 * sigma).clamp(min=R_LO, max=R_HI) + + assert torch.equal(rad_closed, rad_ref) diff --git a/tests/unit/io/test_data.py b/tests/unit/io/test_data.py index 2cbd4c51..a615bfba 100644 --- a/tests/unit/io/test_data.py +++ b/tests/unit/io/test_data.py @@ -80,14 +80,14 @@ def test_reflection_data_cpu(self): @pytest.mark.unit @pytest.mark.gpu - def test_reflection_data_cuda(self, gpu_device): - """Test CUDA movement.""" + def test_reflection_data_accelerator(self, gpu_device): + """Movement onto whichever accelerator this host has.""" from torchref.io import ReflectionData data = ReflectionData() - data = data.cuda() + data = data.to(gpu_device) - assert data.device.type == "cuda" + assert data.device.type == gpu_device.type class TestReflectionDataAttributes: diff --git a/tests/unit/io/test_reflection_data_reindex.py b/tests/unit/io/test_reflection_data_reindex.py new file mode 100644 index 00000000..b5a1243d --- /dev/null +++ b/tests/unit/io/test_reflection_data_reindex.py @@ -0,0 +1,148 @@ +"""Regression tests for per-reflection field reindexing. + +``validate_hkl`` / ``remap`` / ``reduce_to_spacegroup`` must carry EVERY +per-reflection field onto the new HKL grid, not a hand-maintained subset. The +historical bug left ``hkl_anomalous`` (read by ``hkl_for_sf``) at the +pre-alignment length, which crashed difference refinement whenever the dark and +light datasets had different reflection sets. See docs/changelog.rst 0.6.2. +""" + +import pytest +import torch + +from torchref.io.datasets.reflection_data import ReflectionData + + +def _base_grid(h=10, k=10, lmax=10): + """A block of Miller indices with strictly-positive l (already in the ASU).""" + hs = torch.arange(-h, h + 1) + ks = torch.arange(-k, k + 1) + ls = torch.arange(1, lmax + 1) + return ( + torch.stack(torch.meshgrid(hs, ks, ls, indexing="ij"), dim=-1) + .reshape(-1, 3) + .to(torch.int32) + ) + + +def _synthetic(hkl, seed=0, device="cpu"): + n = hkl.shape[0] + g = torch.Generator().manual_seed(seed) + F = torch.rand(n, generator=g) * 100.0 + 1.0 + return ReflectionData.from_tensors( + hkl, + F, + F * 0.1, + (60.0, 60.0, 60.0, 90.0, 90.0, 90.0), + "P 21 21 21", + device=device, + verbose=0, + ) + + +class TestValidateHklReindex: + """The reported crash lives in validate_hkl (collection HKL alignment).""" + + def test_carries_all_per_reflection_fields(self): + grid = _base_grid() + n = grid.shape[0] + # ``light`` lacks the last 10%; the reference grid lacks the first 10%, + # so the two sets genuinely differ (each has reflections the other lacks). + light = _synthetic(grid[: int(n * 0.9)], seed=1) + ref_hkl = grid[int(n * 0.1):].clone() + assert len(light.hkl_anomalous) == len(light.hkl) # sane before + + light.validate_hkl(ref_hkl) + + m = len(light.hkl) + assert m == len(ref_hkl) + # The field the old code left stale (the direct cause of the crash): + assert light.hkl_anomalous.shape[0] == m + assert light.hkl_for_sf().shape[0] == m + # Derived-from-HKL fields recompute for the new grid: + assert light.centric.shape[0] == m + assert light.friedel_flags.shape[0] == m + # Global invariant: no per-reflection tensor left at the old length. + light._assert_per_reflection_consistent() + + def test_identical_hkl_preserves_count(self): + grid = _base_grid(6, 6, 6) + d = _synthetic(grid, seed=2) + n0 = len(d.hkl) + d.validate_hkl(d.hkl.clone()) + assert len(d.hkl) == n0 + d._assert_per_reflection_consistent() + + +class TestP1RoundTripReindex: + """The same class of bug lived latently in remap/expand_to_p1 and + reduce_to_spacegroup (silent data loss rather than a crash).""" + + def test_expand_to_p1_carries_validation_flags(self): + grid = _base_grid(6, 6, 6) + d = _synthetic(grid, seed=3) + d.generate_validation_set(val_fraction_of_free=0.5, seed=0) + assert d.validation_flags is not None + + p1 = d.expand_to_p1() + # validation_flags used to be dropped to None by remap; now carried. + assert p1.validation_flags is not None + assert p1.validation_flags.shape[0] == len(p1.hkl) + p1._assert_per_reflection_consistent() + + def test_reduce_to_spacegroup_consistent(self): + grid = _base_grid(6, 6, 6) + d = _synthetic(grid, seed=4) + p1 = d.expand_to_p1() + back = p1.reduce_to_spacegroup("P 21 21 21") + back._assert_per_reflection_consistent() + + +@pytest.mark.integration +class TestCollectionDifferenceMismatchedHKL: + """In-process reproduction of the reported failure: a DatasetCollection + built from two different reflection sets, run through the collection + difference target's forward (the exact call site that crashed).""" + + def test_forward_finite_with_different_reflection_sets(self, pdb_dir, mtz_dir): + pdb = pdb_dir / "3GR5.pdb" + mtz = mtz_dir / "3GR5.mtz" + if not (pdb.exists() and mtz.exists()): + pytest.skip("3GR5 fixture not present") + + from torchref.cli._common import load_model + from torchref.config import get_default_device + from torchref.io.datasets.collection import DatasetCollection + from torchref.io.datasets.reflection_data import ReflectionData + from torchref.model.model_collection import ModelCollection + from torchref.refinement.targets.collection import CollectionDifferenceTarget + from torchref.scaling.collection_scaler import CollectionScaler + + dev = get_default_device() + full = ReflectionData(device=dev, verbose=0).load_mtz(str(mtz)) + n = len(full.hkl) + idx = torch.arange(n, device=full.hkl.device) + # Two DIFFERENT reflection sets with DIFFERENT counts (drops the last + # 10% vs the first 15%) -> reproduces the shape-mismatch crash pre-fix. + dark = full.__select__(idx < int(n * 0.90)) + light = full.__select__(idx >= int(n * 0.15)) + assert len(dark.hkl) != len(light.hkl) + + dc = DatasetCollection(device=dev, verbose=0) + dc.add_dataset("dark", dark, set_as_reference=True) + dc.add_dataset("light", light) # -> validate_hkl aligns onto dark grid + + # High-resolution limit for the model FFT grid (covers the data). + d_min = float(full.resolution.min()) + model_dark = load_model(str(pdb), max_res=d_min, device=dev, verbose=0) + model_light = load_model(str(pdb), max_res=d_min, device=dev, verbose=0) + mc = ModelCollection([model_dark, model_light], dark_key="dark", verbose=0) + mc.add_dark() + mc.add_timepoint("light", [0.7, 0.3]) + + scaler = CollectionScaler(dc, mc, verbose=0) + scaler.initialize() + + target = CollectionDifferenceTarget(dc, mc, scaler=scaler, verbose=0) + loss = target.forward() + assert torch.isfinite(loss) diff --git a/tests/unit/model/test_model.py b/tests/unit/model/test_model.py index 7f584546..1781a89a 100644 --- a/tests/unit/model/test_model.py +++ b/tests/unit/model/test_model.py @@ -101,7 +101,7 @@ def test_model_gpu_device(self, gpu_device): model = Model(device=gpu_device) - assert model.device.type == 'cuda' + assert model.device.type == gpu_device.type class TestModelGetSelectionMask: diff --git a/tests/unit/model/test_parameter_wrappers.py b/tests/unit/model/test_parameter_wrappers.py index bb1b6b07..5cfeb037 100644 --- a/tests/unit/model/test_parameter_wrappers.py +++ b/tests/unit/model/test_parameter_wrappers.py @@ -133,7 +133,7 @@ def test_mixed_tensor_gpu(self, random_coordinates, gpu_device): mixed = MixedTensor(values, device=gpu_device) result = mixed() - assert result.device.type == 'cuda' + assert result.device.type == gpu_device.type class TestMixedTensorSet: diff --git a/tests/unit/refinement/test_forcefield_target.py b/tests/unit/refinement/test_forcefield_target.py index dacebdf3..df5db224 100644 --- a/tests/unit/refinement/test_forcefield_target.py +++ b/tests/unit/refinement/test_forcefield_target.py @@ -35,6 +35,21 @@ def __init__(self, n_atoms=5, has_hydrogens=True): # Create coordinates as a parameter for gradient testing self._coords = nn.Parameter(torch.randn(len(self.Z), 3)) + @property + def device(self): + """Device of the model's coordinates. + + Real ``Model`` / ``ModelFT`` expose this, and ``ModelTarget`` now + reconciles its own device against the model's, so a stand-in has to + carry it too. + """ + return self._coords.device + + @property + def dtype_float(self): + """Float dtype of the model's coordinates (mirrors ``Model``).""" + return self._coords.dtype + def xyz(self): """Return current coordinates.""" return self._coords diff --git a/tests/unit/refinement/test_ml_sigmaa.py b/tests/unit/refinement/test_ml_sigmaa.py index a264d2da..405836d5 100644 --- a/tests/unit/refinement/test_ml_sigmaa.py +++ b/tests/unit/refinement/test_ml_sigmaa.py @@ -189,15 +189,20 @@ def test_beta_physical_and_falls_with_resolution(self, refinement): assert (beta[v] > 0).all() and torch.isfinite(beta[v]).all() # The falling-with-resolution trend is an estimator-math property, - # checked in float64 so it is device-independent. + # checked in float64 so it is device-independent. Move to CPU before + # casting -- MPS has no float64, so the double() must happen off-device. data = refinement.reflection_data with torch.no_grad(): - fc = torch.abs(t._scaled_F_calc_full()).double().reshape(-1) - fo = data.get_corrected_data()[0].double().reshape(-1) - epsd = epsilon_from_hkl(data.hkl, getattr(data, "spacegroup", None)).double() + fc = torch.abs(t._scaled_F_calc_full()).cpu().double().reshape(-1) + fo = data.get_corrected_data()[0].cpu().double().reshape(-1) + epsd = epsilon_from_hkl( + data.hkl, getattr(data, "spacegroup", None) + ).cpu().double() s = get_scattering_vectors(data.hkl, data.cell) - dss = (torch.norm(s, dim=1) ** 2).double() - _b, bbin, _ = estimate_beta(fo, fc, data.centric, epsd, dss, data.free.mask) + dss = (torch.norm(s, dim=1) ** 2).cpu().double() + _b, bbin, _ = estimate_beta( + fo, fc, data.centric.cpu(), epsd, dss, data.free.mask.cpu() + ) assert (bbin > 0).all() and torch.isfinite(bbin).all() # beta is an absolute model-error variance in F^2 units (~(1-sigma_A^2)* # Sigma_N); Sigma_N ~= decays steeply with resolution, so absolute diff --git a/tests/unit/refinement/test_targets.py b/tests/unit/refinement/test_targets.py index e8d87944..fe2728e7 100644 --- a/tests/unit/refinement/test_targets.py +++ b/tests/unit/refinement/test_targets.py @@ -160,7 +160,7 @@ def test_target_gpu_tensors(self, mock_F_obs, mock_F_sigma, gpu_device): loss = torch.mean((fobs / sigma) ** 2) - assert loss.device.type == 'cuda' + assert loss.device.type == gpu_device.type assert torch.isfinite(loss) diff --git a/tests/unit/restraints/test_restraints.py b/tests/unit/restraints/test_restraints.py index a4ba9d3a..1ec929cf 100644 --- a/tests/unit/restraints/test_restraints.py +++ b/tests/unit/restraints/test_restraints.py @@ -176,7 +176,7 @@ def test_bond_distance_gpu(self, random_coordinates, gpu_device): distance = torch.sqrt(torch.sum((coords[0] - coords[1]) ** 2)) - assert distance.device.type == 'cuda' + assert distance.device.type == gpu_device.type class TestRestraintNumericStability: diff --git a/tests/unit/scaling/test_scaler.py b/tests/unit/scaling/test_scaler.py index 77d5aa6a..44353856 100644 --- a/tests/unit/scaling/test_scaler.py +++ b/tests/unit/scaling/test_scaler.py @@ -81,7 +81,7 @@ def test_scaler_gpu_device(self, gpu_device): scaler = Scaler(device=gpu_device) - assert scaler.device.type == 'cuda' + assert scaler.device.type == gpu_device.type class TestScalingCalculations: diff --git a/tests/unit/symmetry/test_hkl_symmetry_gemmi.py b/tests/unit/symmetry/test_hkl_symmetry_gemmi.py new file mode 100644 index 00000000..8953b20e --- /dev/null +++ b/tests/unit/symmetry/test_hkl_symmetry_gemmi.py @@ -0,0 +1,158 @@ +"""Regression tests for reciprocal-space (Miller-index) symmetry against gemmi. + +These guard the ``h' = h·R = Rᵀ·h`` reciprocal-space convention used by +``SpaceGroup.apply_to_hkl``. A previous bug applied the real-space transform +``R·h`` instead, which is only correct when the fractional rotation matrix is +symmetric. For space groups whose rotation matrices are non-symmetric +(trigonal, hexagonal, and permutation-type cubic operations) ``R·h`` produced +the wrong set of symmetry equivalents, corrupting the centric flags +(``is_centric_from_hkl``) and epsilon multiplicities (``epsilon_from_hkl``) +that feed French-Wilson intensity conversion and ML sigma_A weighting. + +Ground truth comes from gemmi: +- transform: ``gemmi.Op.apply_to_hkl`` +- centric: ``GroupOps.is_reflection_centric`` +- epsilon: ``GroupOps.epsilon_factor_without_centering`` (Friedel-doubled + for centric reflections, matching the Friedel-aware count in + ``epsilon_from_hkl``) +""" + +import numpy as np +import pytest +import torch + +from torchref.config import get_default_device, get_float_dtype, get_int_dtype + +# Space groups whose rotation matrices are non-symmetric — these are the ones +# that regress if apply_to_hkl uses R·h instead of Rᵀ·h. Plus orthorhombic / +# tetragonal controls (symmetric matrices) that must remain correct either way. +_TRIGONAL_HEXAGONAL = ["P 32 2 1", "P 31 2 1", "P 61 2 2", "P 6 2 2", "P 3 1 2"] +_CONTROLS = ["P 21 21 21", "P 43 21 2", "P 1"] +_SPACE_GROUPS = _TRIGONAL_HEXAGONAL + _CONTROLS + +# A spread of reflections including negative indices and axial/general cases. +_HKLS = [ + (1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), (2, 1, 0), (1, -1, 0), + (1, 2, 3), (0, 1, 2), (3, 1, 0), (0, 0, 2), (1, 0, 1), (-1, 2, -3), + (2, 2, 1), (1, 1, 1), +] + + +def _hkl_tensor(): + """Miller indices on the configured default device and integer dtype. + + Kept at the config defaults rather than a hardcoded float64/CPU so the test + exercises the same precision and device production does, and stays + MPS-compatible (MPS has no float64). The arithmetic here is exact in + float32 regardless: fractional rotation matrices are integer-valued + (gemmi's ``op.rot / 24``) and the Miller indices are small, so every + product and sum is well inside the exactly-representable integer range. + """ + return torch.tensor(_HKLS, dtype=get_int_dtype(), device=get_default_device()) + + +@pytest.mark.unit +@pytest.mark.parametrize("sg_name", _SPACE_GROUPS) +def test_apply_to_hkl_matches_gemmi_transform(sg_name): + """apply_to_hkl must reproduce gemmi's per-operation reciprocal transform. + + Compared as the *set* of equivalents per reflection so the test is + insensitive to operation ordering between torchref and gemmi. + """ + import gemmi + + from torchref.symmetry import SpaceGroup + + sg = SpaceGroup(sg_name) + hkl = _hkl_tensor() + + # torchref: (N, 3, ops) -> per-reflection set of equivalents + out = sg.apply_to_hkl(hkl.to(get_float_dtype())) + out_int = torch.round(out).to(torch.int64).cpu().numpy() # (N, 3, ops) + + gemmi_ops = list(gemmi.SpaceGroup(sg_name).operations().sym_ops) + + for n, h in enumerate(_HKLS): + got = {tuple(out_int[n, :, o]) for o in range(out_int.shape[2])} + expected = {tuple(op.apply_to_hkl(h)) for op in gemmi_ops} + assert got == expected, ( + f"{sg_name} hkl={h}: apply_to_hkl equivalents {sorted(got)} " + f"!= gemmi {sorted(expected)}" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("sg_name", _SPACE_GROUPS) +def test_is_centric_from_hkl_matches_gemmi(sg_name): + """Centric flags must match gemmi's is_reflection_centric.""" + import gemmi + + from torchref.base.french_wilson import is_centric_from_hkl + + ops = gemmi.SpaceGroup(sg_name).operations() + hkl = _hkl_tensor() + + got = is_centric_from_hkl(hkl, sg_name).cpu().numpy().astype(bool) + expected = np.array([ops.is_reflection_centric(h) for h in _HKLS], dtype=bool) + + assert np.array_equal(got, expected), ( + f"{sg_name}: centric {got.astype(int)} != gemmi {expected.astype(int)}" + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("sg_name", _SPACE_GROUPS) +def test_epsilon_from_hkl_matches_gemmi(sg_name): + """Friedel-aware epsilon must match gemmi's epsilon (doubled for centrics).""" + import gemmi + + from torchref.base.targets.xray_ml_sigmaa import epsilon_from_hkl + from torchref.symmetry import SpaceGroup + + ops = gemmi.SpaceGroup(sg_name).operations() + sg = SpaceGroup(sg_name) + hkl = _hkl_tensor() + + got = epsilon_from_hkl(hkl, sg).cpu().numpy().round().astype(int) + expected = np.array( + [ + ops.epsilon_factor_without_centering(h) + * (2 if ops.is_reflection_centric(h) else 1) + for h in _HKLS + ], + dtype=int, + ) + + assert np.array_equal(got, expected), ( + f"{sg_name}: epsilon {got} != gemmi-derived {expected}" + ) + + +@pytest.mark.unit +def test_apply_to_hkl_is_transpose_not_plain_rotation(): + """Explicit guard on the exact bug: apply_to_hkl == Rᵀ·h, not R·h. + + Uses P3, whose 3-fold rotation matrix is non-symmetric, so R·h and Rᵀ·h + give genuinely different results. + """ + from torchref.symmetry import SpaceGroup + + sg = SpaceGroup("P 3") + # Use the SpaceGroup's own dtype (config default) -- no hardcoded float64, + # so this runs on MPS as well as CPU/CUDA. + mats = sg.matrices # (ops, 3, 3) + hkl = torch.tensor([[1, 2, 3]], dtype=mats.dtype, device=mats.device) + + out = sg.apply_to_hkl(hkl) # (1, 3, ops) + + # Correct reciprocal transform: Rᵀ·h + expected_t = torch.einsum("oji,nj->nio", mats, hkl) + # The old (buggy) real-space transform: R·h + plain_r = torch.einsum("oij,nj->nio", mats, hkl) + + assert torch.allclose(out, expected_t), "apply_to_hkl should compute Rᵀ·h" + # For P3 the two conventions must actually differ (non-symmetric matrix), + # otherwise this test would not detect a regression. + assert not torch.allclose(expected_t, plain_r), ( + "P3 rotation matrices should be non-symmetric; test cannot detect the bug" + ) diff --git a/tests/unit/symmetry/test_symmetry.py b/tests/unit/symmetry/test_symmetry.py index d04044e1..1ddab2bf 100644 --- a/tests/unit/symmetry/test_symmetry.py +++ b/tests/unit/symmetry/test_symmetry.py @@ -212,8 +212,8 @@ def test_spacegroup_gpu(self, gpu_device): sg = SpaceGroup("P21", device=gpu_device) - assert sg.matrices.device.type == 'cuda' - assert sg.translations.device.type == 'cuda' + assert sg.matrices.device.type == gpu_device.type + assert sg.translations.device.type == gpu_device.type @pytest.mark.unit def test_spacegroup_dtype(self): diff --git a/tests/unit/test_device_conformance.py b/tests/unit/test_device_conformance.py new file mode 100644 index 00000000..ca9ac4ba --- /dev/null +++ b/tests/unit/test_device_conformance.py @@ -0,0 +1,255 @@ +"""Device conformance: constructors place tensors where they were told, and +``.to()`` moves *everything* -- tensors and trackers alike. + +The package resolves a single default device at import and expects every object +to live on it, with ``self.device`` as the authority for where new tensors are +allocated (200+ call sites pass ``device=self.device``). Two failure modes +follow, and both used to be untested: + +* a constructor that builds a sub-object on the *global* default rather than + the device it was handed, producing a split object that only errors much + later, deep inside some unrelated op; +* a ``.to()`` that moves the tensors but leaves ``self.device`` stale, so the + object keeps allocating on the device it just left. + +Before this file, the only test that exercised a real device transition was +CUDA-gated -- meaning zero coverage on CPU-only CI and on Apple silicon. +""" + +from pathlib import Path + +import pytest +import torch + +from torchref.config import canonical_device + +# Absolute ``tests.`` imports, not bare ``helpers.``: the repo carries two +# pytest configs -- ``tests/pytest.ini`` (rootdir ``tests/``, ``pythonpath = ..``) +# and ``[tool.pytest.ini_options]`` in ``pyproject.toml`` (rootdir the repo +# root). Only the repo root is on ``sys.path`` under both, so a bare +# ``helpers.`` import works from ``tests/`` and fails from the repo root. +from tests.helpers.device_asserts import assert_device_consistent, collect_device_map +from tests.helpers.device_cases import ( + CASES, + TARGET_CASES, + UNCOVERED, + UNCOVERED_PREFIXES, +) +from tests.helpers.device_inventory import device_mixin_classes + +_IDS = [c.name for c in CASES] + + +@pytest.mark.unit +@pytest.mark.parametrize("case", CASES, ids=_IDS) +def test_constructed_on_requested_device(case, any_device): + """Everything a constructor builds lands on the device it was given.""" + obj = case.build(any_device) + assert_device_consistent(obj, any_device, name=case.name, ignore=case.ignore) + + +@pytest.mark.unit +@pytest.mark.parametrize("case", CASES, ids=_IDS) +def test_device_round_trip(case, any_device): + """cpu -> device -> cpu, with tensors and trackers following every leg.""" + obj = case.build(torch.device("cpu")) + assert_device_consistent(obj, "cpu", name=case.name, ignore=case.ignore) + + obj.to(any_device) + assert_device_consistent(obj, any_device, name=case.name, ignore=case.ignore) + + obj.to("cpu") + assert_device_consistent(obj, "cpu", name=case.name, ignore=case.ignore) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "case", [c for c in CASES if not c.tensor_free], ids=[c.name for c in CASES if not c.tensor_free] +) +def test_tracker_agrees_with_owned_tensor(case, any_device): + """``obj.device`` must equal a real tensor's device, index and all. + + ``torch.device('mps') != torch.device('mps:0')``, so a tracker carrying the + un-indexed spelling compares unequal to every tensor the object owns even + though nothing is actually misplaced. Callers do write + ``if t.device != self.device``, so the two spellings have to agree. + """ + obj = case.build(any_device) + tensors, _, _, _ = collect_device_map(obj, case.name) + if not tensors: + pytest.skip(f"{case.name} owns no tensors on this path") + tracker = getattr(obj, "device", None) + assert isinstance(tracker, torch.device), f"{case.name}.device is {tracker!r}" + for path, dev in tensors.items(): + assert tracker == dev, ( + f"{case.name}.device == {tracker!r} but {path} is on {dev!r}; " + "these must compare equal, not merely name the same backend" + ) + + +@pytest.mark.unit +def test_every_device_bearing_class_is_accounted_for(): + """A new device-bearing class must be given a case or an explicit excuse. + + Uses a static AST inventory rather than ``DeviceMixin.__subclasses__()``: + the runtime hook only sees classes whose module happens to be imported, so + it would silently under-report exactly when coverage regressed. + """ + package_root = Path(__file__).resolve().parents[2] / "torchref" + found = device_mixin_classes(package_root) + + covered = ( + {c.cls_name for c in CASES} + | {c.cls_name for c in TARGET_CASES} + | set(UNCOVERED) + ) + missing = { + name: where + for name, where in found.items() + if name not in covered + and not any(where.startswith(p) for p in UNCOVERED_PREFIXES) + } + + assert not missing, ( + "device-bearing classes with no conformance case and no entry in " + "helpers.device_cases.UNCOVERED:\n " + + "\n ".join(f"{n} ({w})" for n, w in sorted(missing.items())) + ) + + +@pytest.mark.unit +def test_uncovered_entries_still_exist(): + """Stop ``UNCOVERED`` accumulating excuses for classes that are long gone.""" + package_root = Path(__file__).resolve().parents[2] / "torchref" + found = set(device_mixin_classes(package_root)) + stale = sorted(set(UNCOVERED) - found) + assert not stale, f"UNCOVERED lists classes that no longer exist: {stale}" + + +# --------------------------------------------------------------------------- +# Refinement targets +# +# Targets were the largest hole in the device story: none of them carried a +# ``device`` tracker at all, so ``DeviceMixin``'s machinery was a no-op for +# every one, and their scalar tunables were allocated on CPU in float32 +# regardless of the model's device or the configured dtype. +# --------------------------------------------------------------------------- + +_TARGET_IDS = [c.name for c in TARGET_CASES] + + +@pytest.mark.unit +@pytest.mark.parametrize("case", TARGET_CASES, ids=_TARGET_IDS) +def test_target_follows_its_model_device(case, device_model_bundle, any_device): + """A target's own buffers land beside the model it acts on.""" + target = case.build(device_model_bundle, any_device) + assert_device_consistent(target, any_device, name=case.name, ignore=case.ignore) + + +@pytest.mark.unit +@pytest.mark.parametrize("case", TARGET_CASES, ids=_TARGET_IDS) +def test_target_tracker_is_set(case, device_model_bundle, any_device): + """``target.device`` exists and is truthful. + + Tensor-free targets (``BondTarget``) can only be checked this way -- and + they are exactly the ones that used to have no tracker at all. + """ + target = case.build(device_model_bundle, any_device) + assert isinstance(target.device, torch.device) + assert canonical_device(target.device) == canonical_device(any_device) + + +@pytest.mark.unit +@pytest.mark.parametrize("case", TARGET_CASES, ids=_TARGET_IDS) +def test_target_forward_does_not_reallocate_buffers(case, device_model_bundle, any_device): + """``forward()`` must not move or re-register buffers. + + Three targets used to repair their CPU-resident buffers lazily on the first + forward. That is what this pins shut: allocation inside forward breaks + CUDA-graph capture, and it silently costs a transfer on the hot path. The + old code fails this on the first call and passes on the second. + """ + target = case.build(device_model_bundle, any_device) + + def snapshot(): + return { + n: (b.data_ptr(), b.device, b.dtype) + for n, b in target.named_buffers(recurse=False) + if b is not None + } + + try: + target() + except Exception as exc: # pragma: no cover - target needs state we lack + pytest.skip(f"{case.name} cannot run forward() here: {exc}") + before = snapshot() + target() + assert snapshot() == before, f"{case.name} mutated its buffers during forward()" + + +@pytest.mark.unit +@pytest.mark.parametrize("case", TARGET_CASES, ids=_TARGET_IDS) +def test_empty_target_state_dict_round_trip(case): + """An empty target must round-trip its own state dict strictly. + + The empty-init path exists precisely so ``load_state_dict`` has something + to load into. ``CoordinateSimilarityTarget`` used to fail this: its index + buffers were registered only inside ``_build_atom_map``, which never ran + without models. + """ + cls = type(case.build({"model": None}, torch.device("cpu"))) + shell = cls() + other = cls() + other.load_state_dict(shell.state_dict(), strict=True) + + +@pytest.mark.unit +def test_target_scalar_buffers_follow_config_dtype(pdb_dir, monkeypatch): + """Scalar tunables honour ``dtypes.float``, not a hardcoded float32. + + CPU-only: MPS has no float64. Before the migration every one of these was + ``torch.tensor(float(x))`` -- silently float32 under a float64 config. + """ + pdb = pdb_dir / "1DAW.pdb" + if not pdb.exists(): + pytest.skip("1DAW.pdb fixture not present") + + from torchref.config import dtypes + from torchref.model import ModelFT + from torchref.refinement.targets.geometry.non_bonded import NonBondedTarget + + monkeypatch.setattr(dtypes, "_float", torch.float64) + model = ModelFT(device="cpu", dtype_float=torch.float64, verbose=0).load_pdb(str(pdb)) + target = NonBondedTarget(model) + + assert target.dtype_float == torch.float64 + for name in ("_sigma_vdw", "_r_exp", "_c_rep"): + assert getattr(target, name).dtype == torch.float64, name + + +@pytest.mark.unit +def test_register_target_warns_on_device_mismatch(pdb_dir): + """A mismatched target is reported, not silently relocated. + + Moving it would drag the model and data it borrows onto the state's + device -- registering a loss term must not have that side effect. + """ + pdb = pdb_dir / "1DAW.pdb" + if not pdb.exists(): + pytest.skip("1DAW.pdb fixture not present") + + from torchref.model import ModelFT + from torchref.refinement.loss_state import LossState + from torchref.refinement.targets.geometry.non_bonded import NonBondedTarget + + model = ModelFT(device="cpu", verbose=0).load_pdb(str(pdb)) + target = NonBondedTarget(model) + assert target.device.type == "cpu" + + state = LossState(device=torch.device("meta")) + with pytest.warns(UserWarning, match="is on cpu but the state is on"): + state.register_target("nb", target, probe=False) + + # ...and the target must not have been moved. + assert target.device.type == "cpu" + assert target._sigma_vdw.device.type == "cpu" diff --git a/tests/unit/test_device_mixin.py b/tests/unit/test_device_mixin.py index 5b52efc0..32009cab 100644 --- a/tests/unit/test_device_mixin.py +++ b/tests/unit/test_device_mixin.py @@ -248,3 +248,100 @@ def test_cell_cache_repopulates_on_target_dtype(): assert "volume" not in cell._cache, "cache must be cleared after .to()" new_volume = cell.volume assert new_volume.dtype == torch.float64 + + +# --------------------------------------------------------------------------- +# Probe-driven tracker inference (objects that own no tensors) +# --------------------------------------------------------------------------- + + +class _TensorFreeTracker(DeviceMixin, nn.Module): + """Carries device/dtype trackers but owns no tensor to read them from. + + This is the shape of every geometry target: the trackers say where the + object *will* allocate, and nothing it currently holds can confirm it. Such + objects are the only consumers of ``_probe_target``. + """ + + def __init__(self, dtype=torch.float64): + super().__init__() + self.device = torch.device("cpu") + self.dtype_float = dtype + + +@pytest.mark.unit +def test_device_move_does_not_clobber_dtype_tracker(any_device): + """A device-only move must leave ``dtype_float`` alone. + + Regression: the probe used a single float32 scratch for the dtype axis, so + a plain ``.to(device)`` reported ``float32`` and overwrote a float64 + tracker. Both axes need their own contrast pair. + """ + mod = _TensorFreeTracker(dtype=torch.float64) + mod.to(any_device) + assert mod.dtype_float == torch.float64, "device-only move changed dtype_float" + assert mod.device.type == any_device.type + + +@pytest.mark.unit +def test_dtype_only_move_does_not_clobber_device_tracker(): + """Mirror case: ``.float()`` must leave the device tracker alone.""" + mod = _TensorFreeTracker(dtype=torch.float64) + mod.device = torch.device("cpu") + mod.float() + assert mod.dtype_float == torch.float32 + assert mod.device == torch.device("cpu"), "dtype-only move changed device" + + +@pytest.mark.unit +@pytest.mark.parametrize( + "fn,expected_dtype", + [ + (lambda t: t.to("cpu"), None), + (lambda t: t.double(), torch.float64), + (lambda t: t.float(), torch.float32), + (lambda t: t.half(), torch.float16), + (lambda t: t, None), + ], + ids=["to_cpu", "double", "float", "half", "identity"], +) +def test_probe_target_dtype_axis(fn, expected_dtype): + """The dtype axis reports only what ``fn`` actually casts to. + + Host-independent: the dtype contrast pair is pinned to CPU, which is also + what keeps ``.double()`` probeable on Apple silicon (MPS has no float64, so + contrasting dtype on an accelerator scratch would fail outright). + """ + from torchref.utils.device_mixin import _probe_target + + _, dtype = _probe_target(fn) + assert dtype == expected_dtype + + +@pytest.mark.unit +@pytest.mark.gpu +@pytest.mark.parametrize( + "fn,preserves_device", + [ + (lambda t: t.to("cpu"), False), + (lambda t: t.double(), True), + (lambda t: t.float(), True), + (lambda t: t, True), + ], + ids=["to_cpu", "double", "float", "identity"], +) +def test_probe_target_device_axis(fn, preserves_device): + """The device axis distinguishes "moves to D" from "keeps its input's device". + + Requires an accelerator: telling those two apart needs two devices to + contrast. On a single-device host they are genuinely indistinguishable -- + and reporting ``cpu`` there is correct, since that is where everything is + anyway -- so there is nothing to assert. + """ + from torchref.utils.device_mixin import _probe_target + + device, _ = _probe_target(fn) + if preserves_device: + assert device is None, "device-preserving fn must not pin a device" + else: + assert device == torch.device("cpu") diff --git a/tests/unit/test_gradient_correctness.py b/tests/unit/test_gradient_correctness.py index e98eb081..364ad49d 100644 --- a/tests/unit/test_gradient_correctness.py +++ b/tests/unit/test_gradient_correctness.py @@ -29,8 +29,8 @@ validated with the same cosine/gradnorm metric against central finite differences of a scalar loss (:func:`test_model_forward_xyz_adp_*`). -CPU-only by default; the Triton comparisons are ``gpu``/``cuda_only`` and skip -unless ``--run-gpu`` is passed on a CUDA host. +CPU-only by default; the Triton comparisons are marked ``cuda`` and are +auto-skipped unless the host actually has a CUDA device. """ import itertools @@ -361,10 +361,9 @@ def loss_fn(): # ============================================================================= # 6. Triton kernels (CUDA float32, hand-written backward) vs eager autograd -# Metric: cosine similarity + gradnorm ratio. Skipped without --run-gpu. +# Metric: cosine similarity + gradnorm ratio. Auto-skipped without CUDA. # ============================================================================= -@pytest.mark.gpu -@pytest.mark.cuda_only +@pytest.mark.cuda @pytest.mark.skipif(not _HAS_CUDA, reason="Triton kernels require CUDA") def test_triton_sf_iso_matches_eager_cosine(): """DS isotropic Triton kernel gradient == eager autograd (CUDA float32).""" @@ -388,8 +387,7 @@ def test_triton_sf_iso_matches_eager_cosine(): assert_grads_agree(g_triton, g_eager, min_cos=0.999, ratio_tol=1e-2, ctx="iso ") -@pytest.mark.gpu -@pytest.mark.cuda_only +@pytest.mark.cuda @pytest.mark.skipif(not _HAS_CUDA, reason="Triton kernels require CUDA") def test_triton_sf_aniso_matches_eager_cosine(): """DS anisotropic Triton kernel gradient == eager autograd (CUDA float32).""" @@ -412,8 +410,7 @@ def test_triton_sf_aniso_matches_eager_cosine(): assert_grads_agree(g_triton, g_eager, min_cos=0.999, ratio_tol=1e-2, ctx="aniso ") -@pytest.mark.gpu -@pytest.mark.cuda_only +@pytest.mark.cuda @pytest.mark.skipif(not _HAS_CUDA, reason="Triton kernels require CUDA") def test_triton_gaussian_xray_matches_eager_cosine(): """Gaussian X-ray Triton kernel gradient == eager autograd (CUDA float32).""" @@ -435,8 +432,7 @@ def test_triton_gaussian_xray_matches_eager_cosine(): assert_grads_agree([g_t], [g_e], min_cos=0.999, ratio_tol=1e-2, ctx="xray ") -@pytest.mark.gpu -@pytest.mark.cuda_only +@pytest.mark.cuda @pytest.mark.skipif(not _HAS_CUDA, reason="Triton kernels require CUDA") def test_triton_bond_matches_eager_cosine(): """Bond restraint Triton kernel gradient == eager autograd (CUDA float32).""" diff --git a/tests/unit/utils/test_device_resolution.py b/tests/unit/utils/test_device_resolution.py index 5ef428aa..7e5baa28 100644 --- a/tests/unit/utils/test_device_resolution.py +++ b/tests/unit/utils/test_device_resolution.py @@ -96,7 +96,7 @@ def test_explicit_overrides_skips_nones(self): # GPU tests (cuda required for the cross-device branches) # --------------------------------------------------------------------------- -@pytest.mark.gpu +@pytest.mark.cuda class TestMixedDevices: def test_inconsistent_warns_and_moves_to_first(self): if not torch.cuda.is_available(): diff --git a/torchref/__init__.py b/torchref/__init__.py index 092bfef9..2ff89bbd 100644 --- a/torchref/__init__.py +++ b/torchref/__init__.py @@ -51,7 +51,7 @@ General utilities and debugging tools. """ -__version__ = "0.6.1" +__version__ = "0.6.2" import os diff --git a/torchref/base/alignment/normalization.py b/torchref/base/alignment/normalization.py index 2599a47e..e51e600b 100644 --- a/torchref/base/alignment/normalization.py +++ b/torchref/base/alignment/normalization.py @@ -14,7 +14,7 @@ import torch from torchref.base.math_torch import U_to_matrix -from torchref.config import get_default_device +from torchref.config import normalize_device def compute_radial_shells( @@ -46,8 +46,7 @@ def compute_radial_shells( shell_centers : torch.Tensor Shell centers in Angstroms^-1, shape (n_shells,). """ - if device is None: - device = get_default_device() + device = normalize_device(device) s_min = 1.0 / d_max # Low resolution end s_max = 1.0 / d_min # High resolution end diff --git a/torchref/base/alignment/rotation.py b/torchref/base/alignment/rotation.py index a07f8251..ad0d4399 100644 --- a/torchref/base/alignment/rotation.py +++ b/torchref/base/alignment/rotation.py @@ -8,7 +8,7 @@ import numpy as np import torch -from torchref.config import dtypes, get_default_device, get_float_dtype +from torchref.config import dtypes, get_float_dtype, normalize_device def rotate_coords_torch(coords, phi, rho): @@ -261,8 +261,7 @@ def random_rotation_uniform( torch.Tensor Rotation matrices with shape (n, 3, 3) or (3, 3) if n=1. """ - if device is None: - device = get_default_device() + device = normalize_device(device) if dtype is None: dtype = get_float_dtype() # Sample uniform random numbers diff --git a/torchref/base/electron_density/kernels/mps/__init__.py b/torchref/base/electron_density/kernels/mps/__init__.py new file mode 100644 index 00000000..d324da65 --- /dev/null +++ b/torchref/base/electron_density/kernels/mps/__init__.py @@ -0,0 +1,25 @@ +"""Metal (MPS) variable-radius electron-density splat kernels. + +Native Metal kernels (compiled at runtime via ``torch.mps.compile_shader``) for +the density splat on Apple-silicon GPUs, replacing the portable eager splat that +dominates fcalc time on MPS. Gated behind ``device.type == 'mps'`` in +``electron_density.main``; every other platform is unaffected. +""" + +from torchref.base.electron_density.kernels.mps.compile import ( + clear_cache, + mps_kernels_available, + warmup, +) +from torchref.base.electron_density.kernels.mps.variable_radius import ( + add_anisotropic_mps_var, + add_isotropic_mps_var, +) + +__all__ = [ + "add_isotropic_mps_var", + "add_anisotropic_mps_var", + "mps_kernels_available", + "warmup", + "clear_cache", +] diff --git a/torchref/base/electron_density/kernels/mps/_shaders.py b/torchref/base/electron_density/kernels/mps/_shaders.py new file mode 100644 index 00000000..facd694c --- /dev/null +++ b/torchref/base/electron_density/kernels/mps/_shaders.py @@ -0,0 +1,434 @@ +"""Metal Shading Language (MSL) source for the variable-radius density splat. + +Compiled at runtime via ``torch.mps.compile_shader`` (see ``compile.py``). One +GPU thread per atom (dispatched ``threads=[n_atoms]``): each thread sizes its +own cubic bounding box from ``r2cut`` and the inverse-cell row norms, iterates +the box, truncates to the per-atom sphere (``r2 <= r2cut``), evaluates the +5-term ITC92 Gaussian, and accumulates into the grid via a portable +compare-exchange float atomic-add (``atomic_add_f``) that works on every Metal +GPU family (Apple7 / M1 onward), not just Metal-3 native ``atomic_float``. + +The math mirrors the CUDA variable-radius kernels +(``kernels/cuda/variable_radius.py``) and is validated against the portable CPU +reference (``add_isotropic_plain_var`` / ``add_anisotropic_plain_var``). Because +each thread owns one atom, the backward kernels accumulate per-atom gradients +with no atomics. + +Coordinates: work in Cartesian offsets. ``frac`` is the fractional->Cartesian +matrix (its columns are the cell vectors a,b,c); ``inv_frac`` is its inverse +(Cartesian->fractional). For voxel offset ``o`` from the atom's grid anchor, +``w = frac @ (o/n - residual)`` is the Cartesian atom->voxel vector and +``r2 = w.w``. + +Constants match the reference bit-for-bit: PI_1P5 = pi^1.5, PI_SQ = pi^2. +""" + +# NOTE: kept as a single translation unit so one compile_shader call yields all +# kernels. Backward (iso/aniso) and anisotropic kernels are appended below as +# each is validated. +MSL_SOURCE = r""" +#include +#include +using namespace metal; + +constant float PI_1P5 = 5.568327996831708f; // pi^1.5 +constant float PI_SQ = 9.869604401089358f; // pi^2 + +// Portable float atomic-add via compare-exchange on uint. Works on EVERY Metal +// GPU family (incl. Apple7 / M1), not just Metal-3 native atomic_float. Measured +// identical to native on Apple8/M2 -- the splat's per-atom locality keeps +// contention low, so the CAS loop essentially never retries. +inline void atomic_add_f(device atomic_uint* addr, float val) { + uint old = atomic_load_explicit(addr, memory_order_relaxed), nxt; + do { nxt = as_type(as_type(old) + val); } + while (!atomic_compare_exchange_weak_explicit( + addr, &old, nxt, memory_order_relaxed, memory_order_relaxed)); +} + +// --------------------------------------------------------------------------- +// Isotropic forward: accumulate 5-Gaussian density into grid (atomic add). +// --------------------------------------------------------------------------- +kernel void iso_splat_fwd( + device atomic_uint* grid [[buffer(0)]], // (nx*ny*nz,) accumulator + device const float* xyz [[buffer(1)]], // (n,3) Cartesian + device const float* adp [[buffer(2)]], // (n,) isotropic B + device const float* occ [[buffer(3)]], // (n,) + device const float* A [[buffer(4)]], // (n,5) ITC92 amplitudes + device const float* B [[buffer(5)]], // (n,5) ITC92 widths + device const float* r2cut [[buffer(6)]], // (n,) squared radius + device const float* mask [[buffer(7)]], // (n,5) per-Gaussian mask + device const float* inv_frac [[buffer(8)]], // 9 (row-major 3x3) + device const float* frac [[buffer(9)]], // 9 (row-major 3x3) + constant int& n_atoms [[buffer(10)]], + constant int& nx [[buffer(11)]], + constant int& ny [[buffer(12)]], + constant int& nz [[buffer(13)]], + uint a [[thread_position_in_grid]]) +{ + if (a >= (uint)n_atoms) return; + + // Cell geometry (frac columns are cell vectors a,b,c). + float f0=frac[0],f1=frac[1],f2=frac[2], + f3=frac[3],f4=frac[4],f5=frac[5], + f6=frac[6],f7=frac[7],f8=frac[8]; + float i0=inv_frac[0],i1=inv_frac[1],i2=inv_frac[2], + i3=inv_frac[3],i4=inv_frac[4],i5=inv_frac[5], + i6=inv_frac[6],i7=inv_frac[7],i8=inv_frac[8]; + float fnx=(float)nx, fny=(float)ny, fnz=(float)nz; + + // Per-axis Cartesian voxel step vectors (cell vector / n). + float uax=f0/fnx, uay=f3/fnx, uaz=f6/fnx; + float ubx=f1/fny, uby=f4/fny, ubz=f7/fny; + float ucx=f2/fnz, ucy=f5/fnz, ucz=f8/fnz; + + float rc2 = r2cut[a]; + float r = sqrt(rc2); + // Triclinic-correct half-widths: max|off_axis| = n_axis * r * ||inv_frac row||. + int bhx = (int)ceil(r*fnx*sqrt(i0*i0+i1*i1+i2*i2)); + int bhy = (int)ceil(r*fny*sqrt(i3*i3+i4*i4+i5*i5)); + int bhz = (int)ceil(r*fnz*sqrt(i6*i6+i7*i7+i8*i8)); + + // Atom -> fractional -> wrap into [0,1) -> nearest grid node + residual. + float ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + float fx=ax*i0+ay*i1+az*i2; + float fy=ax*i3+ay*i4+az*i5; + float fz=ax*i6+ay*i7+az*i8; + fx-=floor(fx); fy-=floor(fy); fz-=floor(fz); + int cix=(int)rint(fx*fnx), ciy=(int)rint(fy*fny), ciz=(int)rint(fz*fnz); + float sx=fx-(float)cix/fnx, sy=fy-(float)ciy/fny, sz=fz-(float)ciz/fnz; + float w0x=f0*sx+f1*sy+f2*sz; // residual -> Cartesian + float w0y=f3*sx+f4*sy+f5*sz; + float w0z=f6*sx+f7*sy+f8*sz; + + // Per-Gaussian amplitude / width. + float occa=occ[a], b_iso=adp[a]; + float Bt[5], An[5]; + for (int g=0; g<5; ++g) { + float Bt_g = max((B[5*a+g]+b_iso)*0.25f, 0.1f); + Bt[g] = Bt_g; + An[g] = mask[5*a+g]*A[5*a+g]*occa*PI_1P5/(Bt_g*sqrt(Bt_g)); + } + + for (int ox=-bhx; ox<=bhx; ++ox) { + float fox=(float)ox; + for (int oy=-bhy; oy<=bhy; ++oy) { + float foy=(float)oy; + for (int oz=-bhz; oz<=bhz; ++oz) { + float foz=(float)oz; + float wx=fox*uax+foy*ubx+foz*ucx - w0x; + float wy=fox*uay+foy*uby+foz*ucy - w0y; + float wz=fox*uaz+foy*ubz+foz*ucz - w0z; + float r2=wx*wx+wy*wy+wz*wz; + if (r2 > rc2) continue; + float dens=0.0f; + for (int g=0; g<5; ++g) dens += An[g]*fast::exp(-PI_SQ*r2/Bt[g]); + int vix=cix+ox; vix=((vix%nx)+nx)%nx; + int viy=ciy+oy; viy=((viy%ny)+ny)%ny; + int viz=ciz+oz; viz=((viz%nz)+nz)%nz; + int idx=(vix*ny+viy)*nz+viz; + atomic_add_f(&grid[idx], dens); + } + } + } +} + +// --------------------------------------------------------------------------- +// Isotropic backward: analytic grads to xyz, adp, occ. One thread per atom, so +// each thread owns its atom's gradient slots -> no atomics. Recomputes the +// forward geometry and gathers the upstream grad from the grid. +// --------------------------------------------------------------------------- +kernel void iso_splat_bwd( + device float* grad_xyz [[buffer(0)]], // (n,3) out + device float* grad_adp [[buffer(1)]], // (n,) out + device float* grad_occ [[buffer(2)]], // (n,) out + device const float* grad_out [[buffer(3)]], // (nx*ny*nz,) upstream grad + device const float* xyz [[buffer(4)]], + device const float* adp [[buffer(5)]], + device const float* occ [[buffer(6)]], + device const float* A [[buffer(7)]], + device const float* B [[buffer(8)]], + device const float* r2cut [[buffer(9)]], + device const float* mask [[buffer(10)]], + device const float* inv_frac [[buffer(11)]], + device const float* frac [[buffer(12)]], + constant int& n_atoms [[buffer(13)]], + constant int& nx [[buffer(14)]], + constant int& ny [[buffer(15)]], + constant int& nz [[buffer(16)]], + uint a [[thread_position_in_grid]]) +{ + if (a >= (uint)n_atoms) return; + + float f0=frac[0],f1=frac[1],f2=frac[2], + f3=frac[3],f4=frac[4],f5=frac[5], + f6=frac[6],f7=frac[7],f8=frac[8]; + float i0=inv_frac[0],i1=inv_frac[1],i2=inv_frac[2], + i3=inv_frac[3],i4=inv_frac[4],i5=inv_frac[5], + i6=inv_frac[6],i7=inv_frac[7],i8=inv_frac[8]; + float fnx=(float)nx, fny=(float)ny, fnz=(float)nz; + float uax=f0/fnx, uay=f3/fnx, uaz=f6/fnx; + float ubx=f1/fny, uby=f4/fny, ubz=f7/fny; + float ucx=f2/fnz, ucy=f5/fnz, ucz=f8/fnz; + + float rc2 = r2cut[a]; + float r = sqrt(rc2); + int bhx = (int)ceil(r*fnx*sqrt(i0*i0+i1*i1+i2*i2)); + int bhy = (int)ceil(r*fny*sqrt(i3*i3+i4*i4+i5*i5)); + int bhz = (int)ceil(r*fnz*sqrt(i6*i6+i7*i7+i8*i8)); + + float ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + float fx=ax*i0+ay*i1+az*i2; + float fy=ax*i3+ay*i4+az*i5; + float fz=ax*i6+ay*i7+az*i8; + fx-=floor(fx); fy-=floor(fy); fz-=floor(fz); + int cix=(int)rint(fx*fnx), ciy=(int)rint(fy*fny), ciz=(int)rint(fz*fnz); + float sx=fx-(float)cix/fnx, sy=fy-(float)ciy/fny, sz=fz-(float)ciz/fnz; + float w0x=f0*sx+f1*sy+f2*sz; + float w0y=f3*sx+f4*sy+f5*sz; + float w0z=f6*sx+f7*sy+f8*sz; + + float occa=occ[a], b_iso=adp[a]; + float Bt[5], An[5], clampf[5]; + for (int g=0; g<5; ++g) { + float raw = (B[5*a+g]+b_iso)*0.25f; + float Bt_g = max(raw, 0.1f); + Bt[g] = Bt_g; + An[g] = mask[5*a+g]*A[5*a+g]*occa*PI_1P5/(Bt_g*sqrt(Bt_g)); + clampf[g] = (raw > 0.1f) ? 1.0f : 0.0f; // d(clamped Bt)/d(adp)=0 in clamp region + } + + float gx=0.0f, gy=0.0f, gz=0.0f, gb=0.0f, go=0.0f; + for (int ox=-bhx; ox<=bhx; ++ox) { + float fox=(float)ox; + for (int oy=-bhy; oy<=bhy; ++oy) { + float foy=(float)oy; + for (int oz=-bhz; oz<=bhz; ++oz) { + float foz=(float)oz; + float wx=fox*uax+foy*ubx+foz*ucx - w0x; + float wy=fox*uay+foy*uby+foz*ucy - w0y; + float wz=fox*uaz+foy*ubz+foz*ucz - w0z; + float r2=wx*wx+wy*wy+wz*wz; + if (r2 > rc2) continue; + int vix=cix+ox; vix=((vix%nx)+nx)%nx; + int viy=ciy+oy; viy=((viy%ny)+ny)%ny; + int viz=ciz+oz; viz=((viz%nz)+nz)%nz; + float g_out = grad_out[(vix*ny+viy)*nz+viz]; + float dens=0.0f, coeff_xyz=0.0f, db_sum=0.0f; + for (int g=0; g<5; ++g) { + float e = fast::exp(-PI_SQ*r2/Bt[g]); + float Ae = An[g]*e; + dens += Ae; + coeff_xyz += Ae/Bt[g]; + db_sum += Ae*(-1.5f/Bt[g] + PI_SQ*r2/(Bt[g]*Bt[g]))*clampf[g]; + } + float sxyz = g_out*2.0f*PI_SQ*coeff_xyz; + gx += sxyz*wx; gy += sxyz*wy; gz += sxyz*wz; + gb += g_out*0.25f*db_sum; + go += g_out*dens; + } + } + } + grad_xyz[3*a+0]=gx; grad_xyz[3*a+1]=gy; grad_xyz[3*a+2]=gz; + grad_adp[a]=gb; + grad_occ[a]=(occa!=0.0f) ? go/occa : 0.0f; +} + +constant float TWO_PI_SQ = 19.739208802178716f; // 2*pi^2 (= 8*pi^2 / 4) + +// --------------------------------------------------------------------------- +// Anisotropic forward. Per Gaussian g: M_g = (B_g*I + 8*pi^2*U)/4 (symmetric +// 3x3), inverted analytically; A_norm_g = mask*A*occ*pi^1.5/sqrt(det M_g); +// q_g = w^T Minv_g w; density = sum_g A_norm_g fast::exp(-pi^2 q_g). Truncation +// is a per-axis bounding box (from r2cut + inv-cell row norms) with a sphere +// cull (r2 <= r2cut) -- far tighter than a full cube on anisotropic high-res +// cells. Uses metal::fast::exp (accurate to ~1e-4, well under the grid floor). +// --------------------------------------------------------------------------- +kernel void aniso_splat_fwd( + device atomic_uint* grid [[buffer(0)]], + device const float* xyz [[buffer(1)]], // (n,3) + device const float* u [[buffer(2)]], // (n,6) U11,U22,U33,U12,U13,U23 + device const float* occ [[buffer(3)]], + device const float* A [[buffer(4)]], // (n,5) + device const float* B [[buffer(5)]], // (n,5) + device const float* r2cut [[buffer(6)]], // (n,) squared truncation radius + device const float* mask [[buffer(7)]], // (n,5) + device const float* inv_frac [[buffer(8)]], + device const float* frac [[buffer(9)]], + constant int& n_atoms [[buffer(10)]], + constant int& nx [[buffer(11)]], + constant int& ny [[buffer(12)]], + constant int& nz [[buffer(13)]], + uint a [[thread_position_in_grid]]) +{ + if (a >= (uint)n_atoms) return; + + float f0=frac[0],f1=frac[1],f2=frac[2], + f3=frac[3],f4=frac[4],f5=frac[5], + f6=frac[6],f7=frac[7],f8=frac[8]; + float i0=inv_frac[0],i1=inv_frac[1],i2=inv_frac[2], + i3=inv_frac[3],i4=inv_frac[4],i5=inv_frac[5], + i6=inv_frac[6],i7=inv_frac[7],i8=inv_frac[8]; + float fnx=(float)nx, fny=(float)ny, fnz=(float)nz; + float uax=f0/fnx, uay=f3/fnx, uaz=f6/fnx; + float ubx=f1/fny, uby=f4/fny, ubz=f7/fny; + float ucx=f2/fnz, ucy=f5/fnz, ucz=f8/fnz; + + float ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + float fx=ax*i0+ay*i1+az*i2, fy=ax*i3+ay*i4+az*i5, fz=ax*i6+ay*i7+az*i8; + fx-=floor(fx); fy-=floor(fy); fz-=floor(fz); + int cix=(int)rint(fx*fnx), ciy=(int)rint(fy*fny), ciz=(int)rint(fz*fnz); + float sx=fx-(float)cix/fnx, sy=fy-(float)ciy/fny, sz=fz-(float)ciz/fnz; + float w0x=f0*sx+f1*sy+f2*sz, w0y=f3*sx+f4*sy+f5*sz, w0z=f6*sx+f7*sy+f8*sz; + + float occa=occ[a]; + float u11=u[6*a+0],u22=u[6*a+1],u33=u[6*a+2],u12=u[6*a+3],u13=u[6*a+4],u23=u[6*a+5]; + float p00[5],p11[5],p22[5],p01[5],p02[5],p12[5],An[5]; + for (int g=0; g<5; ++g) { + float Bg=B[5*a+g]; + float ma=0.25f*Bg+TWO_PI_SQ*u11, mb=0.25f*Bg+TWO_PI_SQ*u22, mc=0.25f*Bg+TWO_PI_SQ*u33; + float md=TWO_PI_SQ*u12, me=TWO_PI_SQ*u13, mf=TWO_PI_SQ*u23; + float det=ma*(mb*mc-mf*mf)-md*(md*mc-me*mf)+me*(md*mf-me*mb); + float inv=1.0f/det; + p00[g]=(mb*mc-mf*mf)*inv; p11[g]=(ma*mc-me*me)*inv; p22[g]=(ma*mb-md*md)*inv; + p01[g]=(me*mf-md*mc)*inv; p02[g]=(md*mf-me*mb)*inv; p12[g]=(md*me-ma*mf)*inv; + An[g]=mask[5*a+g]*A[5*a+g]*occa*PI_1P5/sqrt(max(det,1e-10f)); + } + + float rc2=r2cut[a], r=sqrt(rc2); + int bhx=(int)ceil(r*fnx*sqrt(i0*i0+i1*i1+i2*i2)); + int bhy=(int)ceil(r*fny*sqrt(i3*i3+i4*i4+i5*i5)); + int bhz=(int)ceil(r*fnz*sqrt(i6*i6+i7*i7+i8*i8)); + for (int ox=-bhx; ox<=bhx; ++ox) { + float fox=(float)ox; + for (int oy=-bhy; oy<=bhy; ++oy) { + float foy=(float)oy; + for (int oz=-bhz; oz<=bhz; ++oz) { + float foz=(float)oz; + float wx=fox*uax+foy*ubx+foz*ucx - w0x; + float wy=fox*uay+foy*uby+foz*ucy - w0y; + float wz=fox*uaz+foy*ubz+foz*ucz - w0z; + if (wx*wx+wy*wy+wz*wz > rc2) continue; + float dens=0.0f; + for (int g=0; g<5; ++g) { + float q=p00[g]*wx*wx+p11[g]*wy*wy+p22[g]*wz*wz + +2.0f*(p01[g]*wx*wy+p02[g]*wx*wz+p12[g]*wy*wz); + dens += An[g]*fast::exp(-PI_SQ*q); + } + int vix=cix+ox; vix=((vix%nx)+nx)%nx; + int viy=ciy+oy; viy=((viy%ny)+ny)%ny; + int viz=ciz+oz; viz=((viz%nz)+nz)%nz; + atomic_add_f(&grid[(vix*ny+viy)*nz+viz], dens); + } + } + } +} + +// --------------------------------------------------------------------------- +// Anisotropic backward: grads to xyz, U (6), occ. v_g = Minv_g w; dg_g = +// A_norm_g exp(-pi^2 w.v_g). grad_xyz = go*2pi^2*sum dg v; grad_U diag = +// go*2pi^2*sum dg(-0.5 p_ii + pi^2 v_i^2); grad_U offdiag = go*4pi^2*sum +// dg(-0.5 p_ij + pi^2 v_i v_j); grad_occ = go*sum dg / occ. One thread/atom. +// --------------------------------------------------------------------------- +kernel void aniso_splat_bwd( + device float* grad_xyz [[buffer(0)]], // (n,3) + device float* grad_u [[buffer(1)]], // (n,6) + device float* grad_occ [[buffer(2)]], // (n,) + device const float* grad_out [[buffer(3)]], + device const float* xyz [[buffer(4)]], + device const float* u [[buffer(5)]], + device const float* occ [[buffer(6)]], + device const float* A [[buffer(7)]], + device const float* B [[buffer(8)]], + device const float* r2cut [[buffer(9)]], + device const float* mask [[buffer(10)]], + device const float* inv_frac [[buffer(11)]], + device const float* frac [[buffer(12)]], + constant int& n_atoms [[buffer(13)]], + constant int& nx [[buffer(14)]], + constant int& ny [[buffer(15)]], + constant int& nz [[buffer(16)]], + uint a [[thread_position_in_grid]]) +{ + if (a >= (uint)n_atoms) return; + + float f0=frac[0],f1=frac[1],f2=frac[2], + f3=frac[3],f4=frac[4],f5=frac[5], + f6=frac[6],f7=frac[7],f8=frac[8]; + float i0=inv_frac[0],i1=inv_frac[1],i2=inv_frac[2], + i3=inv_frac[3],i4=inv_frac[4],i5=inv_frac[5], + i6=inv_frac[6],i7=inv_frac[7],i8=inv_frac[8]; + float fnx=(float)nx, fny=(float)ny, fnz=(float)nz; + float uax=f0/fnx, uay=f3/fnx, uaz=f6/fnx; + float ubx=f1/fny, uby=f4/fny, ubz=f7/fny; + float ucx=f2/fnz, ucy=f5/fnz, ucz=f8/fnz; + + float ax=xyz[3*a+0], ay=xyz[3*a+1], az=xyz[3*a+2]; + float fx=ax*i0+ay*i1+az*i2, fy=ax*i3+ay*i4+az*i5, fz=ax*i6+ay*i7+az*i8; + fx-=floor(fx); fy-=floor(fy); fz-=floor(fz); + int cix=(int)rint(fx*fnx), ciy=(int)rint(fy*fny), ciz=(int)rint(fz*fnz); + float sx=fx-(float)cix/fnx, sy=fy-(float)ciy/fny, sz=fz-(float)ciz/fnz; + float w0x=f0*sx+f1*sy+f2*sz, w0y=f3*sx+f4*sy+f5*sz, w0z=f6*sx+f7*sy+f8*sz; + + float occa=occ[a]; + float u11=u[6*a+0],u22=u[6*a+1],u33=u[6*a+2],u12=u[6*a+3],u13=u[6*a+4],u23=u[6*a+5]; + float p00[5],p11[5],p22[5],p01[5],p02[5],p12[5],An[5]; + for (int g=0; g<5; ++g) { + float Bg=B[5*a+g]; + float ma=0.25f*Bg+TWO_PI_SQ*u11, mb=0.25f*Bg+TWO_PI_SQ*u22, mc=0.25f*Bg+TWO_PI_SQ*u33; + float md=TWO_PI_SQ*u12, me=TWO_PI_SQ*u13, mf=TWO_PI_SQ*u23; + float det=ma*(mb*mc-mf*mf)-md*(md*mc-me*mf)+me*(md*mf-me*mb); + float inv=1.0f/det; + p00[g]=(mb*mc-mf*mf)*inv; p11[g]=(ma*mc-me*me)*inv; p22[g]=(ma*mb-md*md)*inv; + p01[g]=(me*mf-md*mc)*inv; p02[g]=(md*mf-me*mb)*inv; p12[g]=(md*me-ma*mf)*inv; + An[g]=mask[5*a+g]*A[5*a+g]*occa*PI_1P5/sqrt(max(det,1e-10f)); + } + + float gx=0,gy=0,gz=0,gu0=0,gu1=0,gu2=0,gu3=0,gu4=0,gu5=0,go=0; + float rc2=r2cut[a], r=sqrt(rc2); + int bhx=(int)ceil(r*fnx*sqrt(i0*i0+i1*i1+i2*i2)); + int bhy=(int)ceil(r*fny*sqrt(i3*i3+i4*i4+i5*i5)); + int bhz=(int)ceil(r*fnz*sqrt(i6*i6+i7*i7+i8*i8)); + for (int ox=-bhx; ox<=bhx; ++ox) { + float fox=(float)ox; + for (int oy=-bhy; oy<=bhy; ++oy) { + float foy=(float)oy; + for (int oz=-bhz; oz<=bhz; ++oz) { + float foz=(float)oz; + float wx=fox*uax+foy*ubx+foz*ucx - w0x; + float wy=fox*uay+foy*uby+foz*ucy - w0y; + float wz=fox*uaz+foy*ubz+foz*ucz - w0z; + if (wx*wx+wy*wy+wz*wz > rc2) continue; + int vix=cix+ox; vix=((vix%nx)+nx)%nx; + int viy=ciy+oy; viy=((viy%ny)+ny)%ny; + int viz=ciz+oz; viz=((viz%nz)+nz)%nz; + float g_out=grad_out[(vix*ny+viy)*nz+viz]; + float sxx=0,syy=0,szz=0,dens=0; + float su0=0,su1=0,su2=0,su3=0,su4=0,su5=0; + for (int g=0; g<5; ++g) { + float vx=p00[g]*wx+p01[g]*wy+p02[g]*wz; + float vy=p01[g]*wx+p11[g]*wy+p12[g]*wz; + float vz=p02[g]*wx+p12[g]*wy+p22[g]*wz; + float q=wx*vx+wy*vy+wz*vz; + float dg=An[g]*fast::exp(-PI_SQ*q); + dens+=dg; sxx+=dg*vx; syy+=dg*vy; szz+=dg*vz; + su0+=dg*(-0.5f*p00[g]+PI_SQ*vx*vx); + su1+=dg*(-0.5f*p11[g]+PI_SQ*vy*vy); + su2+=dg*(-0.5f*p22[g]+PI_SQ*vz*vz); + su3+=dg*(-0.5f*p01[g]+PI_SQ*vx*vy); + su4+=dg*(-0.5f*p02[g]+PI_SQ*vx*vz); + su5+=dg*(-0.5f*p12[g]+PI_SQ*vy*vz); + } + float s=g_out*2.0f*PI_SQ; + gx+=s*sxx; gy+=s*syy; gz+=s*szz; + gu0+=s*su0; gu1+=s*su1; gu2+=s*su2; + gu3+=g_out*4.0f*PI_SQ*su3; gu4+=g_out*4.0f*PI_SQ*su4; gu5+=g_out*4.0f*PI_SQ*su5; + go+=g_out*dens; + } + } + } + grad_xyz[3*a+0]=gx; grad_xyz[3*a+1]=gy; grad_xyz[3*a+2]=gz; + grad_u[6*a+0]=gu0; grad_u[6*a+1]=gu1; grad_u[6*a+2]=gu2; + grad_u[6*a+3]=gu3; grad_u[6*a+4]=gu4; grad_u[6*a+5]=gu5; + grad_occ[a]=(occa!=0.0f) ? go/occa : 0.0f; +} +""" diff --git a/torchref/base/electron_density/kernels/mps/compile.py b/torchref/base/electron_density/kernels/mps/compile.py new file mode 100644 index 00000000..d61dbbe4 --- /dev/null +++ b/torchref/base/electron_density/kernels/mps/compile.py @@ -0,0 +1,87 @@ +"""Lazy compilation + availability probe for the Metal (MPS) splat kernels. + +Mirrors the memoize / permanent-failure / graceful-fallback structure of +``kernels/cpu/scatter.py`` (``_get_module``), but the build is a single +in-process ``torch.mps.compile_shader`` call -- no ninja, build directory, or +file locking, since PyTorch caches the compiled pipeline-state objects itself. + +``mps_kernels_available()`` is the gate the dispatcher checks; it returns False +(and the caller falls back to the portable plain splat) whenever MPS is absent, +``compile_shader`` is missing (torch < 2.9), or the shader fails to build. +""" + +from __future__ import annotations + +import traceback +from typing import Optional, Tuple + +import torch + +from torchref.base.electron_density.kernels.mps._shaders import MSL_SOURCE + +# Memoized compile state (attempted at most once per process). +_lib = None +_lib_failed = False +_lib_error: Optional[Tuple[str, str]] = None # (message, traceback) + + +def _mps_shader_supported() -> bool: + """True iff this torch build can compile+run Metal shaders on this host.""" + return ( + hasattr(torch, "mps") + and hasattr(torch.mps, "compile_shader") + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ) + + +def _get_lib(): + """Return the compiled shader library, or None if unavailable. + + Short-circuits both prior outcomes so compilation is attempted exactly once. + Any failure (unsupported torch, MPS off, MSL compile error) is recorded and + turned into a None return so callers fall back to the plain splat. + """ + global _lib, _lib_failed, _lib_error + if _lib is not None: + return _lib + if _lib_failed: + return None + if not _mps_shader_supported(): + _lib_failed = True + _lib_error = ( + "torch.mps.compile_shader unavailable or MPS not available", + "", + ) + return None + try: + _lib = torch.mps.compile_shader(MSL_SOURCE) + except Exception as e: # noqa: BLE001 - any build failure -> fall back + _lib_failed = True + _lib_error = (f"{type(e).__name__}: {e}", traceback.format_exc()) + return None + return _lib + + +def mps_kernels_available() -> bool: + """Whether the Metal splat kernels compiled and are ready to dispatch.""" + return _get_lib() is not None + + +def warmup() -> bool: + """Eagerly trigger compilation (e.g. to move the one-time cost off the + first refinement step). Returns availability.""" + return _get_lib() is not None + + +def clear_cache() -> None: + """Forget the compiled library and failure state (recompiled on next use).""" + global _lib, _lib_failed, _lib_error + _lib = None + _lib_failed = False + _lib_error = None + + +def last_error() -> Optional[Tuple[str, str]]: + """The (message, traceback) of the last compile failure, if any.""" + return _lib_error diff --git a/torchref/base/electron_density/kernels/mps/variable_radius.py b/torchref/base/electron_density/kernels/mps/variable_radius.py new file mode 100644 index 00000000..e11f2c24 --- /dev/null +++ b/torchref/base/electron_density/kernels/mps/variable_radius.py @@ -0,0 +1,192 @@ +"""Metal (MPS) variable-radius electron-density splat, autograd-wrapped. + +``add_isotropic_mps_var`` / ``add_anisotropic_mps_var`` mirror the signatures of +the portable CPU reference (``add_isotropic_plain_var`` / +``add_anisotropic_plain_var`` in ``kernels/cpu/variable_radius.py``) so the +dispatch site in ``main.py`` is a near-copy of the CPU branch. Each delegates to +a ``torch.autograd.Function`` that dispatches the compiled Metal kernels +(one thread per atom) for the forward and analytic backward. + +Gradients flow to ``xyz``, ``adp``/``u``, and ``occ`` (and identity to the input +``density_map``); ``A``/``B`` and the cell matrices receive no gradient -- the +same set as the CUDA kernels. Backward is first-order only (like CUDA); double +backward must use ``Engine.EAGER`` (the plain splat). +""" + +from __future__ import annotations + +import torch + +from torchref.base.electron_density.kernels.mps.compile import _get_lib + + +def _quantized_r2cut(radius_per_atom, voxel_size): + """Per-atom squared cutoff quantized to a voxel multiple, matching the plain + reference (``_box_radius_per_atom`` * ``min_voxel`` in + ``kernels/cpu/variable_radius.py``).""" + min_voxel = voxel_size.min() + r_eff = torch.ceil(radius_per_atom / min_voxel) * min_voxel + return r_eff * r_eff + + +class MetalGridDensity(torch.autograd.Function): + """Isotropic per-atom Metal splat: returns ``density_map + splat``.""" + + @staticmethod + def forward(ctx, density_map, xyz, adp, occ, A, B, r2cut, mask, inv_frac, frac): + lib = _get_lib() + if lib is None: + raise RuntimeError("Metal splat kernels unavailable") + nx, ny, nz = (int(s) for s in density_map.shape) + out = density_map.contiguous().clone() + n = xyz.shape[0] + if n > 0: + lib.iso_splat_fwd( + out.view(-1), + xyz.contiguous(), + adp.contiguous(), + occ.contiguous(), + A.contiguous(), + B.contiguous(), + r2cut.contiguous(), + mask.contiguous(), + inv_frac.contiguous().view(-1), + frac.contiguous().view(-1), + n, nx, ny, nz, + threads=[n], + ) + ctx.save_for_backward(xyz, adp, occ, A, B, r2cut, mask, inv_frac, frac) + ctx.grid_shape = (nx, ny, nz) + return out + + @staticmethod + def backward(ctx, grad_out): + xyz, adp, occ, A, B, r2cut, mask, inv_frac, frac = ctx.saved_tensors + nx, ny, nz = ctx.grid_shape + n = xyz.shape[0] + grad_xyz = torch.zeros_like(xyz) + grad_adp = torch.zeros_like(adp) + grad_occ = torch.zeros_like(occ) + if n > 0: + lib = _get_lib() + lib.iso_splat_bwd( + grad_xyz.view(-1), + grad_adp, + grad_occ, + grad_out.contiguous().view(-1), + xyz.contiguous(), + adp.contiguous(), + occ.contiguous(), + A.contiguous(), + B.contiguous(), + r2cut.contiguous(), + mask.contiguous(), + inv_frac.contiguous().view(-1), + frac.contiguous().view(-1), + n, nx, ny, nz, + threads=[n], + ) + # forward returned density_map + splat -> grad wrt density_map is identity. + # order: density_map, xyz, adp, occ, A, B, r2cut, mask, inv_frac, frac + return (grad_out, grad_xyz, grad_adp, grad_occ, + None, None, None, None, None, None) + + +def add_isotropic_mps_var( + density_map, xyz, adp, occ, A, B, + inv_frac_matrix, frac_matrix, grid_shape_tuple, voxel_size, radius_per_atom, +): + """Isotropic variable-radius Metal splat; adds into ``density_map``. + + Signature mirrors ``add_isotropic_plain_var`` (``grid_shape_tuple`` is + unused -- the grid shape comes from ``density_map``). + + The truncation radius is quantized up to a voxel multiple + (``ceil(radius/min_voxel)*min_voxel``) to exactly match the plain-splat + reference's per-atom cutoff, so the Metal and fallback paths agree. + """ + r2cut = _quantized_r2cut(radius_per_atom, voxel_size) + mask = torch.ones(xyz.shape[0], 5, dtype=xyz.dtype, device=xyz.device) + return MetalGridDensity.apply( + density_map, xyz, adp, occ, A, B, r2cut, mask, inv_frac_matrix, frac_matrix + ) + + +class MetalGridDensityAniso(torch.autograd.Function): + """Anisotropic per-atom Metal splat: returns ``density_map + splat``.""" + + @staticmethod + def forward(ctx, density_map, xyz, u, occ, A, B, r2cut, mask, inv_frac, frac): + lib = _get_lib() + if lib is None: + raise RuntimeError("Metal splat kernels unavailable") + nx, ny, nz = (int(s) for s in density_map.shape) + out = density_map.contiguous().clone() + n = xyz.shape[0] + if n > 0: + lib.aniso_splat_fwd( + out.view(-1), + xyz.contiguous(), + u.contiguous(), + occ.contiguous(), + A.contiguous(), + B.contiguous(), + r2cut.contiguous(), + mask.contiguous(), + inv_frac.contiguous().view(-1), + frac.contiguous().view(-1), + n, nx, ny, nz, + threads=[n], + ) + ctx.save_for_backward(xyz, u, occ, A, B, r2cut, mask, inv_frac, frac) + ctx.grid_shape = (nx, ny, nz) + return out + + @staticmethod + def backward(ctx, grad_out): + xyz, u, occ, A, B, r2cut, mask, inv_frac, frac = ctx.saved_tensors + nx, ny, nz = ctx.grid_shape + n = xyz.shape[0] + grad_xyz = torch.zeros_like(xyz) + grad_u = torch.zeros_like(u) + grad_occ = torch.zeros_like(occ) + if n > 0: + lib = _get_lib() + lib.aniso_splat_bwd( + grad_xyz.view(-1), + grad_u.view(-1), + grad_occ, + grad_out.contiguous().view(-1), + xyz.contiguous(), + u.contiguous(), + occ.contiguous(), + A.contiguous(), + B.contiguous(), + r2cut.contiguous(), + mask.contiguous(), + inv_frac.contiguous().view(-1), + frac.contiguous().view(-1), + n, nx, ny, nz, + threads=[n], + ) + # order: density_map, xyz, u, occ, A, B, r2cut, mask, inv_frac, frac + return (grad_out, grad_xyz, grad_u, grad_occ, + None, None, None, None, None, None) + + +def add_anisotropic_mps_var( + real_space_grid, density_map, xyz, u, occ, A, B, + inv_frac_matrix, frac_matrix, radius_per_atom, voxel_size, +): + """Anisotropic variable-radius Metal splat; adds into ``density_map``. + + Signature mirrors ``add_anisotropic_plain_var`` (``real_space_grid`` is used + only for its grid shape). Each atom is truncated at its per-axis bounding box + with a sphere cull at the (quantized) per-atom radius -- far tighter than the + plain reference's full cube on anisotropic high-resolution cells. + """ + r2cut = _quantized_r2cut(radius_per_atom, voxel_size) + mask = torch.ones(xyz.shape[0], 5, dtype=xyz.dtype, device=xyz.device) + return MetalGridDensityAniso.apply( + density_map, xyz, u, occ, A, B, r2cut, mask, inv_frac_matrix, frac_matrix + ) diff --git a/torchref/base/electron_density/main.py b/torchref/base/electron_density/main.py index 621ad98e..9c1b0d2f 100644 --- a/torchref/base/electron_density/main.py +++ b/torchref/base/electron_density/main.py @@ -10,13 +10,15 @@ - ``Engine.AUTO`` — fastest available per device, all variable-radius: CUDA+float32 -> the work-queue Triton kernels - (``WorkQueueGridDensity`` / ``WorkQueueGridDensityAniso``); CPU -> the + (``WorkQueueGridDensity`` / ``WorkQueueGridDensityAniso``); CPU+float32 -> the grouped-separable variable-radius splat (``add_isotropic_cpu_separable_var`` / ``add_anisotropic_cpu_var``); - everything else (CUDA+float64, MPS) -> the portable plain-scatter - variable-radius splat (``add_isotropic_plain_var`` / - ``add_anisotropic_plain_var``). On a Triton kernel failure under AUTO it falls - through to the plain-scatter variable-radius splat. + MPS+float32 -> the native Metal kernels + (``add_isotropic_mps_var`` / ``add_anisotropic_mps_var``, compiled via + ``torch.mps.compile_shader``); everything else (CUDA/MPS float64) -> the + portable plain-scatter variable-radius splat (``add_isotropic_plain_var`` / + ``add_anisotropic_plain_var``). On a Triton/Metal kernel failure or + unavailability under AUTO it falls through to the plain-scatter splat. - ``Engine.EAGER`` — the portable plain-scatter variable-radius splat on every device. Double-differentiable; use it for Hessians / debugging. Force it with ``with use_engine(Engine.EAGER): ...``. @@ -206,8 +208,10 @@ def _add_isotropic( variable-radius Triton kernel (``WorkQueueGridDensity``). On kernel failure under AUTO it falls through to the portable splat; under ``Engine.TRITON`` it raises (never silently degrade). - - CPU + AUTO -> the fast C++-scatter grouped-separable splat. - - Everything else (``Engine.EAGER`` on any device, CUDA float64, MPS) -> the + - CPU + float32 + AUTO -> the fast C++-scatter grouped-separable splat. + - MPS + float32 + AUTO -> the native Metal kernel (``add_isotropic_mps_var``); + falls through to the portable splat if the shader is unavailable. + - Everything else (``Engine.EAGER`` on any device, CUDA/MPS float64) -> the portable plain-``scatter_add`` grouped splat: identical per-atom radius, double-differentiable, float64-capable, device-agnostic. @@ -251,6 +255,26 @@ def _add_isotropic( density_map, xyz, adp, occ, A, B, inv_frac_matrix, frac_matrix, grid_shape_tuple, voxel_size, radius_per_atom, ) + # Native Metal splat on Apple-silicon GPUs (float32 only). Falls through to + # the portable plain splat if the shader is unavailable or fails. + if ( + get_engine() is Engine.AUTO + and density_map.device.type == "mps" + and density_map.dtype == torch.float32 + ): + from torchref.base.electron_density.kernels.mps import ( + add_isotropic_mps_var, + mps_kernels_available, + ) + + if mps_kernels_available(): + try: + return add_isotropic_mps_var( + density_map, xyz, adp, occ, A, B, + inv_frac_matrix, frac_matrix, grid_shape_tuple, voxel_size, radius_per_atom, + ) + except Exception: + pass # fall through to the portable splat return add_isotropic_plain_var( density_map, xyz, adp, occ, A, B, inv_frac_matrix, frac_matrix, grid_shape_tuple, voxel_size, radius_per_atom, @@ -274,12 +298,14 @@ def _add_anisotropic( The per-atom radius is the isotropic bounding radius of the ellipsoid (largest principal axis, ``per_atom_radius_aniso``). - CUDA+float32 (engine permitting) -> ``WorkQueueGridDensityAniso``. CPU + AUTO - -> the fast C++-scatter grouped box-splat ``add_anisotropic_cpu_var``. - Everything else (``Engine.EAGER`` on any device, CUDA float64, MPS) -> the - portable plain-``scatter_add`` box-splat ``add_anisotropic_plain_var`` - (double-diff, float64-capable, device-agnostic). All paths use the per-atom - radius (isotropic bounding sphere of the ellipsoid). + CUDA+float32 (engine permitting) -> ``WorkQueueGridDensityAniso``. CPU + + float32 + AUTO -> the fast C++-scatter grouped box-splat + ``add_anisotropic_cpu_var``. MPS + float32 + AUTO -> the native Metal kernel + ``add_anisotropic_mps_var`` (falls through if unavailable). Everything else + (``Engine.EAGER`` on any device, CUDA/MPS float64) -> the portable + plain-``scatter_add`` box-splat ``add_anisotropic_plain_var`` (double-diff, + float64-capable, device-agnostic). All paths use the per-atom radius + (isotropic bounding sphere of the ellipsoid). """ n_sigma = get_sigma_cutoff_ed() radius_per_atom = per_atom_radius_aniso(B, u, n_sigma=n_sigma) @@ -318,6 +344,26 @@ def _add_anisotropic( real_space_grid, density_map, xyz, u, occ, A, B, inv_frac_matrix, frac_matrix, radius_per_atom, voxel_size, ) + # Native Metal splat on Apple-silicon GPUs (float32 only). Falls through to + # the portable plain splat if the shader is unavailable or fails. + if ( + get_engine() is Engine.AUTO + and density_map.device.type == "mps" + and density_map.dtype == torch.float32 + ): + from torchref.base.electron_density.kernels.mps import ( + add_anisotropic_mps_var, + mps_kernels_available, + ) + + if mps_kernels_available(): + try: + return add_anisotropic_mps_var( + real_space_grid, density_map, xyz, u, occ, A, B, + inv_frac_matrix, frac_matrix, radius_per_atom, voxel_size, + ) + except Exception: + pass # fall through to the portable splat return add_anisotropic_plain_var( real_space_grid, density_map, xyz, u, occ, A, B, inv_frac_matrix, frac_matrix, radius_per_atom, voxel_size, diff --git a/torchref/base/electron_density/radius_policy.py b/torchref/base/electron_density/radius_policy.py index ffcb5af4..0b280559 100644 --- a/torchref/base/electron_density/radius_policy.py +++ b/torchref/base/electron_density/radius_policy.py @@ -48,6 +48,49 @@ def _u6_to_u3(u: torch.Tensor) -> torch.Tensor: return U3 +def _max_eig_sym3(A: torch.Tensor) -> torch.Tensor: + """Largest eigenvalue of a batch of symmetric 3x3 matrices, in closed form. + + Uses the analytic (trigonometric) solution of the characteristic cubic + (Smith 1961) so no ``torch.linalg.eigvalsh`` is needed -- that op is + unimplemented on MPS, and the eigendecomposition here feeds only the + (quantized) splat radius, so an exact decomposition is overkill. + + Parameters + ---------- + A : torch.Tensor + Symmetric matrices, shape (n, 3, 3). + + Returns + ------- + torch.Tensor + Largest eigenvalue per matrix, shape (n,). + """ + a00 = A[:, 0, 0]; a11 = A[:, 1, 1]; a22 = A[:, 2, 2] + a01 = A[:, 0, 1]; a02 = A[:, 0, 2]; a12 = A[:, 1, 2] + + p1 = a01 * a01 + a02 * a02 + a12 * a12 + q = (a00 + a11 + a22) / 3.0 + p2 = (a00 - q) ** 2 + (a11 - q) ** 2 + (a22 - q) ** 2 + 2.0 * p1 + p = torch.sqrt((p2 / 6.0).clamp(min=1e-30)) + + # B = (A - q I) / p ; r = det(B) / 2 in [-1, 1] + b00 = (a00 - q) / p; b11 = (a11 - q) / p; b22 = (a22 - q) / p + b01 = a01 / p; b02 = a02 / p; b12 = a12 / p + detB = ( + b00 * (b11 * b22 - b12 * b12) + - b01 * (b01 * b22 - b12 * b02) + + b02 * (b01 * b12 - b11 * b02) + ) + r = (detB / 2.0).clamp(-1.0, 1.0) + phi = torch.acos(r) / 3.0 + eig_max = q + 2.0 * p * torch.cos(phi) + + # Diagonal matrices (p1 == 0): eigenvalues are the diagonal entries. + diag_max = torch.maximum(torch.maximum(a00, a11), a22) + return torch.where(p1 <= 1e-30, diag_max, eig_max) + + def sigma_eff_iso(adp: torch.Tensor, B_widths: torch.Tensor) -> torch.Tensor: """``sigma_eff_i = sqrt((max_k B_widths[i,k] + adp_i) / 8pi^2)``, shape (n,). @@ -93,7 +136,8 @@ def sigma_eff_aniso(B_widths: torch.Tensor, u: torch.Tensor) -> torch.Tensor: Anisotropic U parameters [U11,U22,U33,U12,U13,U23], shape (n, 6). """ b_form = B_widths.detach().max(dim=1).values # broadest ITC92 width - lam_max = torch.linalg.eigvalsh(_u6_to_u3(u.detach())).max(dim=1).values + # Closed-form largest eigenvalue (no eigvalsh -> runs natively on MPS). + lam_max = _max_eig_sym3(_u6_to_u3(u.detach())) return torch.sqrt((b_form / EIGHT_PI2 + lam_max).clamp(min=1e-6)) diff --git a/torchref/base/reciprocal/symmetry.py b/torchref/base/reciprocal/symmetry.py index e99e0f03..902afdbe 100644 --- a/torchref/base/reciprocal/symmetry.py +++ b/torchref/base/reciprocal/symmetry.py @@ -34,7 +34,7 @@ import numpy as np import torch -from torchref.config import get_float_dtype +from torchref.config import canonical_device, get_float_dtype from torchref.utils.autograd_ops import gather_with_index_add from .grid_operations import extract_structure_factor_from_grid @@ -251,7 +251,12 @@ def __init__( grid_shape: tuple, device: Optional[torch.device] = None, ): - self.device = device or hkl.device + # ``is not None``, not ``or``: ``device=0`` means cuda:0/mps:0 and is + # falsy, so ``or`` silently discarded it. ``hkl`` is a bare tensor, so + # this reads its device rather than going through ``resolve_device``. + self.device = canonical_device( + device if device is not None else hkl.device + ) self.hkl = hkl.to(device=self.device) self.symmetry = symmetry self.n_ops = symmetry.n_ops diff --git a/torchref/base/targets/xray_ml_sigmaa.py b/torchref/base/targets/xray_ml_sigmaa.py index cb92f092..f08e3066 100644 --- a/torchref/base/targets/xray_ml_sigmaa.py +++ b/torchref/base/targets/xray_ml_sigmaa.py @@ -36,6 +36,8 @@ import numpy as np import torch +from torchref.config import get_float_dtype + # ===================================================================== # Loss math (mean = |Fc|, variance = epsilon*beta) # ===================================================================== @@ -112,19 +114,22 @@ def epsilon_from_hkl(hkl: torch.Tensor, spacegroup) -> torch.Tensor: ``apply_to_hkl``. """ n = hkl.shape[0] - ones = torch.ones(n, device=hkl.device) + float_dtype = get_float_dtype() if spacegroup is None or not hasattr(spacegroup, "apply_to_hkl"): - return ones - try: - with torch.no_grad(): - Hs = spacegroup.apply_to_hkl(hkl.to(torch.float64)) # (N,3,ops) - h0 = hkl.to(torch.float64).unsqueeze(-1) # (N,3,1) - same = (Hs == h0).all(dim=1) - friedel = (Hs == -h0).all(dim=1) - eps = (same | friedel).sum(dim=1).clamp(min=1).to(torch.get_default_dtype()) - return eps.to(hkl.device) - except Exception: - return ones + return torch.ones(n, device=hkl.device, dtype=float_dtype) + + with torch.no_grad(): + # Configured float dtype, not float64: MPS has no float64 and casting + # there raises. Symmetry arithmetic on Miller indices is exact in + # float32 (integer-valued rotation matrices, small indices), so the + # exact `==` comparisons below remain valid. + h = hkl.to(float_dtype) + Hs = spacegroup.apply_to_hkl(h) # (N,3,ops) + h0 = h.unsqueeze(-1) # (N,3,1) + same = (Hs == h0).all(dim=1) + friedel = (Hs == -h0).all(dim=1) + eps = (same | friedel).sum(dim=1).clamp(min=1).to(float_dtype) + return eps # ===================================================================== diff --git a/torchref/cli/_common.py b/torchref/cli/_common.py index f91b8542..67a90602 100644 --- a/torchref/cli/_common.py +++ b/torchref/cli/_common.py @@ -633,10 +633,9 @@ def load_model( ModelFT """ from torchref import ModelFT - from torchref.config import get_default_device + from torchref.config import normalize_device - if device is None: - device = get_default_device() + device = normalize_device(device) model = ModelFT(max_res=max_res, device=device, verbose=verbose) suffix = Path(path).suffix.lower() if suffix in (".cif", ".mmcif"): @@ -674,10 +673,9 @@ def load_reflection_data( ReflectionData """ from torchref import ReflectionData - from torchref.config import get_default_device + from torchref.config import normalize_device - if device is None: - device = get_default_device() + device = normalize_device(device) data = ReflectionData(device=str(device), verbose=verbose) suffix = Path(path).suffix.lower() if suffix in (".cif",): diff --git a/torchref/cli/validate_ded.py b/torchref/cli/validate_ded.py index 60156f6a..1f457df6 100644 --- a/torchref/cli/validate_ded.py +++ b/torchref/cli/validate_ded.py @@ -292,10 +292,9 @@ def setup_ded_context( from torchref.symmetry.grid_utils import calculate_optimal_grid_size from torchref.symmetry.reciprocal_symmetry import expand_hkl - if device is None: - from torchref.config import get_default_device + from torchref.config import normalize_device - device = get_default_device() + device = normalize_device(device) # Load and scale data_dark = load_reflection_data( diff --git a/torchref/config.py b/torchref/config.py index 94b253f2..33eceb62 100644 --- a/torchref/config.py +++ b/torchref/config.py @@ -358,6 +358,65 @@ def _auto_detect_device() -> torch.device: return torch.device("cpu") +def canonical_device(dev): + """Return ``dev`` with its default index filled in. + + ``torch.device('cuda') != torch.device('cuda:0')`` even though both name the + same physical device. Devices read back off a real tensor always carry an + index, so any comparison between a *requested* device and an *observed* one + needs a shared normal form -- otherwise ``obj.device == tensor.device`` is + False on a freshly constructed object and True after its first ``.to()``. + + ``cpu`` deliberately stays bare: ``torch.empty(0, device='cpu').device`` has + no index, so canonicalising it to ``cpu:0`` would recreate the very mismatch + this function exists to remove. + + This is the allocation-free form of ``torch.empty(0, device=d).device``, + which matters because it is called on constructor and comparison paths. + Deliberately *not* memoised: ``torch.cuda.current_device()`` is mutable via + ``torch.cuda.set_device``, so a cache would freeze a stale answer. + + Parameters + ---------- + dev : torch.device or str or int or None + Device to normalise. ``None`` passes through so callers can chain. + + Returns + ------- + torch.device or None + """ + if dev is None: + return None + d = dev if isinstance(dev, torch.device) else torch.device(dev) + if d.index is not None: + return d + if d.type == "cpu": + return d + if d.type == "cuda": + return torch.device("cuda", torch.cuda.current_device()) + if d.type == "mps": + return torch.device("mps", 0) + # Unknown/future backend: fall back to asking PyTorch directly. + return torch.empty(0, device=d).device + + +def normalize_device(dev=None) -> torch.device: + """Coerce a user-supplied ``device`` argument to a canonical device. + + ``None`` resolves to :func:`get_default_device`. This is the pure, + side-effect-free counterpart of :func:`torchref.utils.resolve_device`: + + * one device source (or none) -> ``normalize_device`` + * several device-bearing inputs to reconcile -> ``resolve_device`` + + ``resolve_device`` *moves* the objects it is given, so using it for a + single-input constructor with nothing to reconcile is overreach. + """ + if dev is None: + return get_default_device() + return canonical_device(dev) + + class DeviceConfig: """ Device configuration with property-based access. @@ -381,7 +440,10 @@ class DeviceConfig: def __init__(self): override = os.environ.get("TORCHREF_DEVICE", "auto").lower() if override == "auto": - self._device = _auto_detect_device() + # Canonicalise here too, not just in ``_coerce``: the ``auto`` + # branch bypasses ``_coerce`` entirely, and it is the branch almost + # every user takes. + self._device = canonical_device(_auto_detect_device()) else: self._device = self._coerce(override) self._warn_if_mps_dtype_mismatch() @@ -408,7 +470,7 @@ def _coerce(value) -> torch.device: hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ): raise RuntimeError("MPS requested but not available on this system.") - return dev + return canonical_device(dev) def _warn_if_mps_dtype_mismatch(self) -> None: if self._device.type == "mps" and dtypes.float == torch.float64: diff --git a/torchref/io/datasets/base.py b/torchref/io/datasets/base.py index dd7013b3..53271a74 100644 --- a/torchref/io/datasets/base.py +++ b/torchref/io/datasets/base.py @@ -20,7 +20,7 @@ import gemmi import torch -from torchref.config import get_default_device +from torchref.config import get_default_device, normalize_device from torchref.symmetry import Cell from torchref.utils.device_mixin import DeviceMovementMixin @@ -121,9 +121,10 @@ class CrystalDataset(DeviceMovementMixin): def __post_init__(self): """Initialize non-field attributes after dataclass init.""" - # Ensure device is a torch.device object - if isinstance(self.device, str): - object.__setattr__(self, "device", torch.device(self.device)) + # Canonicalise, don't merely coerce: ``torch.device("mps")`` has no + # index and so compares unequal to the ``mps:0`` every tensor reports, + # which makes ``resolve_device``'s equality checks misfire. + object.__setattr__(self, "device", normalize_device(self.device)) # Import here to avoid circular imports from torchref.utils.utils import TensorMasks @@ -204,8 +205,7 @@ def _from_state(cls, state: Dict[str, Any], device=None) -> "CrystalDataset": """ from torchref.utils.utils import TensorMasks - if device is None: - device = get_default_device() + device = normalize_device(device) # Extract masks before creating object masks_state = state.pop("masks", {}) diff --git a/torchref/io/datasets/fcalc_data.py b/torchref/io/datasets/fcalc_data.py index d5139792..2184b4e1 100644 --- a/torchref/io/datasets/fcalc_data.py +++ b/torchref/io/datasets/fcalc_data.py @@ -14,7 +14,7 @@ import pandas as pd import torch -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_default_device, get_float_dtype, normalize_device from torchref.symmetry import Cell, SpaceGroup, SpaceGroupLike from .base import CrystalDataset @@ -131,13 +131,21 @@ def from_cell_and_resolution( from torchref.base.reciprocal import get_d_spacing - if device is None: - device = get_default_device() + # ``Cell`` is the unit-cell carrier: when one is handed in and no device + # is requested, follow it rather than the global default -- otherwise a + # CPU-default host silently relocates the caller's cell. + if device is None and isinstance(cell, Cell): + device = cell.device + device = normalize_device(device) if dtype is None: dtype = get_float_dtype() # Handle Cell input - convert to Cell object if needed if isinstance(cell, Cell): + # ``Cell.to`` is in-place, so this moves the *caller's* object when + # an explicit device disagrees with it. That is the documented + # resolve_device contract; the branch above avoids it entirely when + # no device was requested. cell_obj = cell.to(device=device) cell_tensor = cell_obj.data else: @@ -376,8 +384,7 @@ def _from_state(cls, state: Dict[str, Any], device=None) -> "FcalcDataset": """ from torchref.utils.utils import TensorMasks - if device is None: - device = get_default_device() + device = normalize_device(device) # Extract masks before creating object masks_state = state.pop("masks", {}) diff --git a/torchref/io/datasets/reflection_data.py b/torchref/io/datasets/reflection_data.py index 4fb74f60..006f0daf 100644 --- a/torchref/io/datasets/reflection_data.py +++ b/torchref/io/datasets/reflection_data.py @@ -8,7 +8,8 @@ import warnings from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import numpy as np import pandas as pd @@ -17,7 +18,7 @@ from torchref.base import math_torch from torchref.base.french_wilson import FrenchWilson -from torchref.config import dtypes, get_default_device +from torchref.config import dtypes, get_default_device, normalize_device from torchref.io import cif, mtz from torchref.io.datasets.base import CrystalDataset from torchref.symmetry import Cell, SpaceGroup @@ -331,6 +332,153 @@ def _tv(t): self._corrected_fp = fp return self._corrected_cache + # ===================== per-reflection field reindexing ===================== + # + # Any routine that changes the HKL set (expand onto a reference grid, remap, + # symmetry reduction) must carry EVERY per-reflection field along, not a + # hand-maintained subset. Historically ``validate_hkl`` / ``remap`` / + # ``reduce_to_spacegroup`` each hardcoded their own field list and silently + # dropped fields added later (notably ``hkl_anomalous``, read by + # ``hkl_for_sf``), which broke difference refinement on mismatched reflection + # files. The single source of truth below is shared by all of them. + + # Fill value used for MISSING output rows (index == -1) per field. Missing + # rows are always masked out downstream (``hkl_present`` / ``missing``), so + # these fills only need to be shape-correct and non-poisonous. Fields not + # listed default to 0. NOTE: ``hkl_anomalous`` is special-cased (filled with + # the reference HKL row, never 0 which would be a spurious Miller index) and + # the derived-from-HKL fields below are recomputed, never gathered. + _REINDEX_FILL = { + "F": 0.0, + "I": 0.0, + "phase": 0.0, + "fom": 0.0, + "E": 0.0, + "E_squared": 0.0, + "F_squared_corrected": 0.0, + "radial_shell_indices": 0, + "F_sigma": 1.0, + "I_sigma": 1.0, + "rfree_flags": 1, # missing reflections default to the work set + "friedel_flags": False, + "validation_flags": False, + "outlier_flags": False, + } + + # Per-reflection fields that are pure functions of (hkl, cell, spacegroup): + # never gathered/aggregated, always recomputed or invalidated after an HKL + # change (``resolution`` recomputed; ``bin_indices`` / ``_centric_flags`` + # lazily rebuilt by ``get_bins`` / the ``centric`` property). + _REINDEX_DERIVED = ("resolution", "bin_indices", "_centric_flags") + + # Tensor dataclass fields that are NOT per-reflection (exempt from the + # length invariant): overall anisotropy parameters have shape (6,). + _NON_PER_REFLECTION_TENSORS = frozenset({"U_aniso"}) + + def _reindex_per_reflection( + self, + index_map: torch.Tensor, + new_hkl: torch.Tensor, + target: Optional["ReflectionData"] = None, + ) -> torch.Tensor: + """Gather every per-reflection field from ``self`` onto ``new_hkl``. + + Enumerates dataclass fields generically (same ``shape[0] == n`` guard as + :meth:`__select__` / :meth:`_canonicalize_in_place`) so any current or + future per-reflection field is carried automatically. + + Parameters + ---------- + index_map : torch.Tensor + LongTensor of length ``len(new_hkl)`` mapping each output row to a + source row index into ``self``, or ``-1`` for reflections absent + from the source (filled per :attr:`_REINDEX_FILL`). + new_hkl : torch.Tensor + Miller indices for the reindexed dataset, shape ``(M, 3)``. + target : ReflectionData, optional + Object to write the gathered fields onto. Defaults to ``self`` + (in-place). Pass a fresh instance for out-of-place reindexing + (``remap``); the caller is responsible for setting ``target.cell`` / + ``target.spacegroup`` before calling so resolution can be recomputed. + + Returns + ------- + torch.Tensor + Boolean presence mask (``index_map >= 0``), for the caller's use in + building ``hkl_present`` / ``missing`` masks. + """ + from dataclasses import fields as dc_fields + + if target is None: + target = self + + n_src = len(self.hkl) if self.hkl is not None else 0 + new_hkl = new_hkl.to(dtype=dtypes.int, device=self.device) + n_out = len(new_hkl) + index_map = index_map.to(device=self.device, dtype=torch.long) + present = index_map >= 0 + src_idx = index_map[present] + + derived = set(self._REINDEX_DERIVED) + for f in dc_fields(self): + name = f.name + if name == "hkl" or name in derived: + continue + val = getattr(self, name) + if not isinstance(val, torch.Tensor): + continue + if not (val.shape and val.shape[0] == n_src): + # Non-per-reflection tensor (e.g. U_aniso (6,)): leave target's + # own value untouched (fresh default when target is new). + continue + if name == "hkl_anomalous": + # Present rows keep their signed (anomalous) index; missing rows + # fall back to the canonical reference HKL (never a 0,0,0 row). + out = new_hkl.clone() + out[present] = val[src_idx] + else: + fill = self._REINDEX_FILL.get(name, 0) + out = torch.full( + (n_out,) + tuple(val.shape[1:]), + fill, + dtype=val.dtype, + device=self.device, + ) + out[present] = val[src_idx] + setattr(target, name, out) + + # Install the new HKL and recompute / invalidate derived-from-HKL fields. + target.hkl = new_hkl + target.bin_indices = None + target._centric_flags = None + if target.cell is not None: + target._calculate_resolution() + else: + target.resolution = None + return present + + def _assert_per_reflection_consistent(self) -> None: + """Invariant: every per-reflection tensor matches ``len(self.hkl)``. + + A cheap post-condition for the reindex routines. Turns any future + hardcoded-list regression (a field silently left at the old length) into + a loud failure instead of a downstream shape mismatch. + """ + from dataclasses import fields as dc_fields + + n = len(self.hkl) if self.hkl is not None else 0 + bad = [] + for f in dc_fields(self): + if f.name in self._NON_PER_REFLECTION_TENSORS: + continue + val = getattr(self, f.name) + if isinstance(val, torch.Tensor) and val.ndim >= 1 and val.shape[0] != n: + bad.append((f.name, tuple(val.shape))) + if bad: + raise RuntimeError( + f"ReflectionData per-reflection length mismatch (n_hkl={n}): {bad}" + ) + def _canonicalize_in_place(self) -> None: """Remap HKL to canonical CCP4 ASU form and reorder all data in-place.""" from dataclasses import fields as dc_fields @@ -455,7 +603,12 @@ def load(self, reader, french_wilson: bool = True): ) if spacegroup is not None: - self.spacegroup = SpaceGroup(spacegroup) + # Build on the dataset's device, not the process default: a + # ``ReflectionData(device='cpu')`` on an MPS host would otherwise + # end up with mps-resident symmetry matrices while every other + # tensor it owns is on CPU, and nothing downstream would notice + # until an op mixed the two. + self.spacegroup = SpaceGroup(spacegroup, device=self.device) self._calculate_resolution() # Prefer intensities (via French-Wilson) only when French-Wilson is @@ -640,9 +793,14 @@ def from_tensors( ReflectionData Fully initialized reflection data with all cleanup applied. """ - if device is None: - device = get_default_device() - data = cls(device=device, verbose=verbose) + # Follow the incoming tensors when no device is given, rather than + # the global default -- otherwise building from accelerator-resident + # arrays on a CPU-default host silently round-trips every one of them + # through host memory. ``hkl`` is a bare tensor, so read its device + # directly; ``resolve_device`` is for objects it can move in place. + if device is None and isinstance(hkl, torch.Tensor): + device = hkl.device + data = cls(device=normalize_device(device), verbose=verbose) def _prep(t: torch.Tensor) -> torch.Tensor: # Detach (constant data) unless the caller wants the graph preserved. @@ -657,7 +815,9 @@ def _prep(t: torch.Tensor) -> torch.Tensor: else Cell(cell, device=data.device) ) data.spacegroup = ( - spacegroup if isinstance(spacegroup, SpaceGroup) else SpaceGroup(spacegroup) + spacegroup + if isinstance(spacegroup, SpaceGroup) + else SpaceGroup(spacegroup, device=data.device) ) if rfree_flags is not None: @@ -691,7 +851,7 @@ def _prep(t: torch.Tensor) -> torch.Tensor: def load_mtz( self, - path: str, + path: Union[str, Path], column_names: Optional[dict] = None, french_wilson: bool = True, anomalous: Optional[bool] = None, @@ -723,14 +883,18 @@ def load_mtz( ReflectionData Self, for method chaining. """ + # ``str(path)``: gemmi's readers take a str, and the public + # ``torchref.io.read_mtz`` already advertises ``Union[str, Path]``, so + # accepting a Path here too keeps the two entry points consistent + # instead of failing deep inside gemmi with an argument-type error. reader = mtz.MTZReader( verbose=self.verbose, column_names=column_names, anomalous=anomalous - ).read(path) + ).read(str(path)) return self.load(reader, french_wilson=french_wilson) def load_cif( self, - path: str, + path: Union[str, Path], data_block: Optional[str] = None, anomalous: Optional[bool] = None, ) -> "ReflectionData": @@ -756,7 +920,7 @@ def load_cif( Self, for method chaining. """ self.reader = cif.ReflectionCIFReader( - path, verbose=self.verbose, data_block=data_block, anomalous=anomalous + str(path), verbose=self.verbose, data_block=data_block, anomalous=anomalous ) return self.load(self.reader) @@ -2223,63 +2387,15 @@ def validate_hkl(self, hkl_ref: torch.Tensor) -> "ReflectionData": [data_hkl_to_idx.get(tuple(hkl), -1) for hkl in hkl_ref_np], dtype=np.int64 ) - # Create presence mask: True where data exists - presence_mask = torch.from_numpy(ref_to_data_idx >= 0).to(device=self.device) + # Index map into this dataset for each reference row (-1 where missing). valid_indices = torch.from_numpy(ref_to_data_idx).to(device=self.device) - # Helper to expand a tensor to reference size - def expand_tensor(tensor, fill_value=0.0): - if tensor is None: - return None - expanded = torch.full( - (n_ref,) + tensor.shape[1:], - fill_value, - dtype=tensor.dtype, - device=self.device, - ) - # Copy existing data to correct positions - mask = valid_indices >= 0 - expanded[mask] = tensor[valid_indices[mask]] - return expanded - - # Expand all data arrays - old_F = self.F - old_F_sigma = self.F_sigma - old_I = self.I - old_I_sigma = getattr(self, "I_sigma", None) - old_rfree = self.rfree_flags - old_phase = getattr(self, "phase", None) - old_fom = getattr(self, "fom", None) - - # Replace HKL with reference - self.hkl = hkl_ref - - # Expand data tensors - self.F = expand_tensor(old_F, fill_value=0.0) - self.F_sigma = expand_tensor(old_F_sigma, fill_value=1.0) - - if old_I is not None: - self.I = expand_tensor(old_I, fill_value=0.0) - if old_I_sigma is not None: - self.I_sigma = expand_tensor(old_I_sigma, fill_value=1.0) - - # For rfree, default missing to work set (1) - if old_rfree is not None: - rfree_expanded = torch.ones( - n_ref, dtype=old_rfree.dtype, device=self.device - ) - mask = valid_indices >= 0 - rfree_expanded[mask] = old_rfree[valid_indices[mask]] - self.rfree_flags = rfree_expanded - - # Recalculate resolution for new HKL set - self._calculate_resolution() - - # Expand phase and fom if present - if old_phase is not None: - self.phase = expand_tensor(old_phase, fill_value=0.0) - if old_fom is not None: - self.fom = expand_tensor(old_fom, fill_value=0.0) + # Reindex EVERY per-reflection field onto the reference grid via the + # shared primitive (carries hkl_anomalous / friedel_flags / centric / + # validation / outlier that this routine used to silently drop, which + # broke difference refinement on mismatched reflection files). Masks are + # handled separately below because they are not dataclass fields. + presence_mask = self._reindex_per_reflection(valid_indices, hkl_ref) # Transfer existing masks to new indexing old_masks = dict(self.masks.items()) @@ -2308,6 +2424,7 @@ def expand_tensor(tensor, fill_value=0.0): print(f" Present in data: {n_present} ({100*n_present/n_ref:.1f}%)") print(f" Missing (masked): {n_missing} ({100*n_missing/n_ref:.1f}%)") + self._assert_per_reflection_consistent() return self def find_outliers( @@ -3046,8 +3163,9 @@ def remap( """ from torchref.symmetry.spacegroup import SpaceGroup - # Helper function for remapping tensors with missing handling - def _remap_tensor(tensor, fill_value): + # Mask remapper. Masks are not dataclass fields, so the shared + # per-reflection reindexer below does not touch them. + def _remap_mask(tensor, fill_value): if tensor is None: return None valid_mask = index_mapping >= 0 @@ -3060,56 +3178,36 @@ def _remap_tensor(tensor, fill_value): result[valid_mask] = tensor[index_mapping[valid_mask]] return result - # Create new ReflectionData + # Create new ReflectionData; set cell/spacegroup first so the shared + # reindexer can recompute resolution on the new grid. remapped = ReflectionData(verbose=self.verbose, device=self.device) + remapped.cell = self.cell.clone() if self.cell is not None else None + if spacegroup is not None: + remapped.spacegroup = SpaceGroup(spacegroup, device=remapped.device) + else: + remapped.spacegroup = self.spacegroup - # Set new HKL - remapped.hkl = new_hkl.to(device=self.device) - - # Remap amplitude and intensity fields - remapped.F = _remap_tensor(self.F, fill_value=0.0) - remapped.F_sigma = _remap_tensor(self.F_sigma, fill_value=1.0) - remapped.I = _remap_tensor(self.I, fill_value=0.0) - remapped.I_sigma = _remap_tensor(self.I_sigma, fill_value=1.0) - remapped.fom = _remap_tensor(self.fom, fill_value=0.0) - - # Carry forward prior mask if available - prior_mask = self.masks() - if prior_mask is not None: - remapped.masks["prior_flagged"] = _remap_tensor( - prior_mask.to(dtype=dtypes.int), fill_value=0 - ).to(torch.bool) - - # Handle rfree_flags (True = include in Rfree for missing) - if self.rfree_flags is not None: - remapped.rfree_flags = _remap_tensor( - self.rfree_flags.to(dtype=dtypes.int), fill_value=1 - ).to(torch.bool) + # Reindex ALL per-reflection dataclass fields onto new_hkl. This carries + # the anomalous / validation / centric / outlier fields the old hardcoded + # list silently dropped, and sets hkl / resolution while invalidating + # bin_indices and _centric_flags. Missing rows (index -1) get the + # per-field fills in _REINDEX_FILL. + self._reindex_per_reflection(index_mapping, new_hkl, target=remapped) - # Handle phases with optional shifts - if self.phase is not None: - remapped.phase = _remap_tensor(self.phase, fill_value=0.0) + # Apply optional phase shifts (e.g. from symmetry translations). + if remapped.phase is not None: if phase_shifts is not None: remapped.phase = remapped.phase + phase_shifts.to(device=self.device) elif phase_shifts is not None: - # Store phase shifts even if no original phases + # No original phases: store the shifts for later phase reconstruction. remapped._expansion_phase_shifts = phase_shifts.to(device=self.device) - # Clone cell - remapped.cell = self.cell.clone() if self.cell is not None else None - - # Set spacegroup - if spacegroup is not None: - remapped.spacegroup = SpaceGroup(spacegroup) - else: - remapped.spacegroup = self.spacegroup - - # Recalculate resolution - if remapped.cell is not None and remapped.hkl is not None: - remapped._calculate_resolution() - - # Invalidate dependent fields - remapped.bin_indices = None + # Carry forward prior combined mask if available. + prior_mask = self.masks() + if prior_mask is not None: + remapped.masks["prior_flagged"] = _remap_mask( + prior_mask.to(dtype=dtypes.int), fill_value=0 + ).to(torch.bool) # Copy metadata sources remapped.amplitude_source = self.amplitude_source @@ -3126,6 +3224,7 @@ def _remap_tensor(tensor, fill_value): if missing_mask.any(): remapped.masks["missing"] = ~missing_mask.to(device=self.device) + remapped._assert_per_reflection_consistent() return remapped def fill(self, d_min: Optional[float] = None) -> "ReflectionData": @@ -3440,18 +3539,69 @@ def _aggregate_sigma(tensor, agg_func="mean"): # 0 = free, non-zero = work. Take min to get free if any is free. reduced.rfree_flags = rfree_gathered.min(dim=1).values != 0 + # Boolean per-reflection flags: an equivalent's flag propagates to the + # merged reflection if ANY contributor has it set (validation/outlier + # are conservative "exclude if any"). + def _aggregate_any(tensor): + if tensor is None: + return None + gathered = tensor[reduction_indices.clamp(min=0)].to(torch.bool) + gathered = gathered & valid_mask + return gathered.any(dim=1) + + if self.validation_flags is not None: + reduced.validation_flags = _aggregate_any(self.validation_flags) + if self.outlier_flags is not None: + reduced.outlier_flags = _aggregate_any(self.outlier_flags) + + # Completeness pass: carry any remaining per-reflection dataclass tensor + # field not handled above so the merge never silently drops data (the + # historical bug). Derived-from-HKL fields are recomputed/invalidated + # below, not aggregated; 'first' is a safe representative for the rest. + from dataclasses import fields as dc_fields + + _already_set = { + "hkl", + "F", + "F_sigma", + "I", + "I_sigma", + "phase", + "fom", + "rfree_flags", + "validation_flags", + "outlier_flags", + } + _recomputed = set(self._REINDEX_DERIVED) | {"hkl_anomalous", "friedel_flags"} + n_src = len(self.hkl) + for f in dc_fields(self): + name = f.name + if name in _already_set or name in _recomputed: + continue + if name in self._NON_PER_REFLECTION_TENSORS: + continue + val = getattr(self, name) + if not isinstance(val, torch.Tensor): + continue + if not (val.shape and val.shape[0] == n_src): + continue + setattr(reduced, name, _aggregate_tensor(val, "first")) + # Clone cell reduced.cell = self.cell.clone() if self.cell is not None else None - # Set spacegroup - reduced.spacegroup = SpaceGroup(spacegroup) + # Set spacegroup (on the reduced dataset's device, not the global default) + reduced.spacegroup = SpaceGroup(spacegroup, device=reduced.device) # Recalculate resolution if reduced.cell is not None and reduced.hkl is not None: reduced._calculate_resolution() - # Invalidate dependent fields + # Invalidate derived-from-HKL fields (recomputed lazily for the new ASU). + # hkl_anomalous / friedel_flags are left unset: the merged ASU is + # Friedel-merged, so hkl_for_sf() correctly falls back to hkl. reduced.bin_indices = None + reduced._centric_flags = None # Copy metadata sources reduced.amplitude_source = self.amplitude_source @@ -3465,6 +3615,7 @@ def _aggregate_sigma(tensor, agg_func="mean"): f"reduce_to_spacegroup({spacegroup}, aggregation={aggregation})" ) + reduced._assert_per_reflection_consistent() return reduced def canonicalize(self, include_friedel: bool = True) -> "ReflectionData": diff --git a/torchref/io/ihm.py b/torchref/io/ihm.py index 5b1e858c..a054dc68 100644 --- a/torchref/io/ihm.py +++ b/torchref/io/ihm.py @@ -474,10 +474,9 @@ def build_model_collection( "Call read_atom_data() first." ) - if device is None: - from torchref.config import get_default_device + from torchref.config import normalize_device - device = get_default_device() + device = normalize_device(device) # Build one ModelFT per state base_models = [] diff --git a/torchref/maps/difference_map.py b/torchref/maps/difference_map.py index 44a19e1a..6051c5ed 100644 --- a/torchref/maps/difference_map.py +++ b/torchref/maps/difference_map.py @@ -12,7 +12,6 @@ import torch from torchref.base.reciprocal.grid_operations import place_on_grid -from torchref.config import get_default_device from torchref.io.datasets.collection import DatasetCollection from torchref.maps.map import Map from torchref.symmetry.reciprocal_symmetry import expand_hkl diff --git a/torchref/model/mixed_model.py b/torchref/model/mixed_model.py index 787cebe2..e24e443e 100644 --- a/torchref/model/mixed_model.py +++ b/torchref/model/mixed_model.py @@ -12,6 +12,7 @@ from torch import nn from torchref.utils.device_mixin import DeviceMovementMixin +from torchref.utils.device_resolution import resolve_device if TYPE_CHECKING: from torchref.model.model_ft import ModelFT @@ -104,12 +105,12 @@ def __init__( self.verbose = verbose - # Infer device from first model if not specified - if device is None: - device = models[0].device + # First-wins reconciliation across the supplied models, and the move + # itself: ``resolve_device`` warns when they disagree, which the manual + # loop this replaces did silently. + device = resolve_device(*models, device=device) # Store models as ModuleList for proper PyTorch handling - models = [model.to(device=device) for model in models] self.models = nn.ModuleList(models) # Validate model compatibility diff --git a/torchref/model/model.py b/torchref/model/model.py index 94d4c00d..6056a1df 100644 --- a/torchref/model/model.py +++ b/torchref/model/model.py @@ -19,7 +19,7 @@ import torch.nn as nn from torchref.base import math_torch -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.io import cif, pdb from torchref.model.parameter_wrappers import ( CholeskyMixedTensor, @@ -151,8 +151,7 @@ def __init__( # ``dtypes.float`` / ``device.current`` change is honored. if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) # Configuration self.dtype_float = dtype_float self.verbose = verbose @@ -287,7 +286,10 @@ def spacegroup(self, value): a space group name string, or a space group number. """ if value is not None: - self._spacegroup = SpaceGroup(value) + # ``device=self.device``: SpaceGroup falls back to the global + # default otherwise, so setting a spacegroup on a CPU-pinned Model + # would silently plant accelerator-resident matrices on it. + self._spacegroup = SpaceGroup(value, device=self.device) else: self._spacegroup = None @@ -1091,21 +1093,28 @@ def get_vdw_radii(self): ), f"vdW radii length mismatch with number of atoms {len(self.vdw_radii)} != {len(self.pdb)}" return self.vdw_radii - def to(self, *args, **kwargs): - """Move Model and rebuild device-specific SF indices. + def _after_device_apply( + self, old_device, new_device, old_dtype, new_dtype, *, + device_changed, dtype_changed, + ): + """Rebuild the precomputed SF indices after a real device/dtype change. - Delegates to :class:`~torchref.utils.device_mixin.DeviceMovementMixin`, which - walks ``self.__dict__`` (picking up ``self.cell``, ``self.altloc_pairs``, + :class:`~torchref.utils.device_mixin.DeviceMixin` walks + ``self.__dict__`` (picking up ``self.cell``, ``self.altloc_pairs``, ``self._restraints`` and all registered parameters / buffers), refreshes - the ``self.device`` tracker, and invalidates caches. Afterwards this - override rebuilds the precomputed SF indices on the new device. - """ - result = super().to(*args, **kwargs) - if hasattr(result, "aniso_flag") and result.aniso_flag is not None: - result._rebuild_sf_indices() - if result.verbose > 0: - print(f"Model moved to device: {result.device}") - return result + the ``self.device`` tracker and invalidates caches; this hook then + regenerates the iso/aniso index tensors on the new device. + + This lives on the movement hook rather than a ``to()`` override because + a parent module reaches ``Model`` through ``_apply``, which never calls + ``to()``. It is deliberately not ``reset_cache()``: that hook fires + after every optimizer step (see ``LossState.reset_caches``), and index + reconstruction does not belong on that path. + """ + if getattr(self, "aniso_flag", None) is not None: + self._rebuild_sf_indices() + if self.verbose > 0: + print(f"Model moved to device: {self.device}") def copy(self): """ @@ -2934,8 +2943,7 @@ def create_from_state_dict( """ # Resolve dtype/device at call time so the fallbacks below use the # current config, not the import-time default. - if device is None: - device = get_default_device() + device = normalize_device(device) if dtype_float is None: dtype_float = get_float_dtype() # Extract metadata (non-tensor data that we handle specially) diff --git a/torchref/model/model_collection.py b/torchref/model/model_collection.py index 34084c2f..c798872a 100644 --- a/torchref/model/model_collection.py +++ b/torchref/model/model_collection.py @@ -11,8 +11,8 @@ import torch from torch import nn -from torchref.config import get_default_device from torchref.utils.device_mixin import DeviceMovementMixin +from torchref.utils.device_resolution import resolve_device if TYPE_CHECKING: from torchref.model.model_ft import ModelFT @@ -69,8 +69,11 @@ def __init__( # Normalize to handle floating point drift initial_fractions = [f / total for f in initial_fractions] - if device is None: - device = base_models[0].device if hasattr(base_models[0], "device") else get_default_device() + # Reconcile across *all* base models, not just the first: reading + # ``base_models[0].device`` left a mixed-device list unreconciled, so + # ``fractions_tensor`` below could land on a device the later models + # were not on. + device = resolve_device(*base_models, device=device) # Match base models' float dtype (consistent under a float64 config). fractions_tensor = torch.tensor( diff --git a/torchref/model/model_ft.py b/torchref/model/model_ft.py index 214fc0c6..a4195edb 100644 --- a/torchref/model/model_ft.py +++ b/torchref/model/model_ft.py @@ -6,7 +6,7 @@ import torch from torchref.base.fourier import fft, ifft -from torchref.config import dtypes, get_default_device, get_float_dtype +from torchref.config import dtypes, get_float_dtype, normalize_device from torchref.model.model import Model from torchref.model.sf_fft import SfFFT from torchref.symmetry import SpaceGroup @@ -1151,8 +1151,7 @@ def create_from_state_dict( # Resolve dtype/device at call time so the fallback below uses the # current config rather than an import-time default. - if device is None: - device = get_default_device() + device = normalize_device(device) if dtype_float is None: dtype_float = get_float_dtype() diff --git a/torchref/model/parameter_wrappers.py b/torchref/model/parameter_wrappers.py index 50289f8d..550d68c6 100644 --- a/torchref/model/parameter_wrappers.py +++ b/torchref/model/parameter_wrappers.py @@ -8,6 +8,7 @@ import torch from torch import nn +from torchref.config import get_float_dtype, normalize_device from torchref.utils.caching import CachedForwardMixin from torchref.utils.device_mixin import DeviceMixin @@ -133,12 +134,19 @@ def __init__( # Empty initialization if initial_values is None: + # Honour the requested device/dtype even with no values yet: the + # empty shell is the documented ``load_state_dict`` entry point, and + # a caller that asked for a device should not get a CPU parameter + # (nor a ``.device`` property that raises because every buffer is + # ``None``). + device = normalize_device(device) + dtype = dtype if dtype is not None else get_float_dtype() self.register_buffer("refinable_mask", None) self.register_buffer("fixed_mask", None) self.register_buffer("fixed_values", None) - self.register_buffer("_shape", None) self.refinable_params = nn.Parameter( - torch.empty(0), requires_grad=requires_grad + torch.empty(0, device=device, dtype=dtype), + requires_grad=requires_grad, ) self._has_refinable = False self._refinable_indices = None @@ -200,8 +208,12 @@ def __init__( refinable_values, requires_grad=requires_grad ) - # Store shape for reconstruction - self.register_buffer("_shape", torch.tensor(initial_values.shape)) + # Shape is host-side metadata, deliberately *not* a registered buffer: + # a buffer would be dragged onto the accelerator by ``.to()`` and then + # read back with ``.tolist()``, i.e. a device sync on every ``.shape`` + # access, purely to recover numbers we already have. It is also fully + # derivable from ``fixed_values`` (a clone of ``initial_values``), so + # storing it separately only creates something that can disagree. # Pre-compute index cache to avoid boolean indexing at runtime self._build_index_cache() @@ -476,17 +488,30 @@ def set(self, values: torch.Tensor, mask: torch.Tensor) -> None: @property def shape(self): """Return the shape of the full tensor.""" - return tuple(self._shape.tolist()) + if self.fixed_values is None: + return () + return tuple(self.fixed_values.shape) + + # ``fixed_values`` is ``None`` on the empty-shell path, so both properties + # fall back to the (empty) parameter, which always exists and carries the + # device/dtype the constructor was given. Without the fallback the + # ``AttributeError`` raised by ``None.device`` is intercepted by + # ``nn.Module.__getattr__`` and re-raised as the thoroughly misleading + # "'MixedTensor' object has no attribute 'device'". @property def dtype(self): """Return the dtype of the tensor.""" - return self.fixed_values.dtype + if self.fixed_values is not None: + return self.fixed_values.dtype + return self.refinable_params.dtype @property def device(self): """Return the device of the tensor.""" - return self.fixed_values.device + if self.fixed_values is not None: + return self.fixed_values.device + return self.refinable_params.device def get_refinable_count(self) -> int: """Return the number of refinable parameters.""" @@ -517,6 +542,16 @@ def update_fixed_values(self, new_values: torch.Tensor): ) self.fixed_values = new_values.to(dtype=self.dtype, device=self.device).detach() + def _normalize_refinable_mask(self, new_mask: torch.Tensor) -> torch.Tensor: + """Coerce an incoming mask onto this wrapper's device as ``bool``. + + Must be applied *before* deriving ``fixed_mask = ~new_mask``: negating + the caller's un-migrated tensor leaves ``refinable_mask`` and + ``fixed_mask`` on different devices, which then fails far away from + here, in whichever op first uses both. + """ + return new_mask.to(device=self.device, dtype=torch.bool) + def update_refinable_mask( self, new_mask: torch.Tensor, reset_refinable: bool = False ): @@ -542,7 +577,8 @@ def update_refinable_mask( current_full = self.forward().detach() - self.refinable_mask = new_mask.to(device=self.device) + new_mask = self._normalize_refinable_mask(new_mask) + self.refinable_mask = new_mask self.fixed_mask = ~new_mask if reset_refinable: @@ -609,11 +645,27 @@ def clip(self, min_value=None, max_value=None) -> "MixedTensor": ) return new_mixed - def to(self, *args, **kwargs): - """Move via :class:`DeviceMixin` and rebuild the index cache.""" - result = super().to(*args, **kwargs) - result._build_index_cache() - return result + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Drop the legacy ``_shape`` buffer key before delegating. + + ``_shape`` used to be a registered buffer; it is now derived from + ``fixed_values``. Checkpoints written before that change still carry the + key, and a ``strict=True`` load would reject it as unexpected. + """ + state_dict.pop(prefix + "_shape", None) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def _after_device_apply(self, *args, device_changed, dtype_changed, **kwargs): + """Rebuild the index cache after a real device/dtype change. + + ``_refinable_indices`` / ``_refinable_idx_1d`` are plain attributes + holding tensors, so they must be regenerated on the new device. Using + the movement hook rather than a ``to()`` override matters twice over: + ``to()`` is bypassed when a parent module reaches this wrapper through + ``_apply``, and ``reset_cache()`` would put this rebuild on the + per-optimizer-step path. + """ + self._build_index_cache() def refine( self, selection: Union[slice, torch.Tensor, tuple], reset_values: bool = False @@ -676,6 +728,7 @@ def refine( new_mask |= temp_mask # Update the mask + new_mask = self._normalize_refinable_mask(new_mask) self.refinable_mask = new_mask self.fixed_mask = ~new_mask @@ -755,6 +808,7 @@ def fix( new_mask &= ~temp_mask # Update the mask + new_mask = self._normalize_refinable_mask(new_mask) self.refinable_mask = new_mask self.fixed_mask = ~new_mask @@ -1180,7 +1234,8 @@ def update_refinable_mask( # Convert to log space current_log = torch.log(current_normal.clamp(min=self.epsilon)) - self.refinable_mask = new_mask.to(device=self.device) + new_mask = self._normalize_refinable_mask(new_mask) + self.refinable_mask = new_mask self.fixed_mask = ~new_mask # Update fixed_values with log-space values @@ -1407,7 +1462,8 @@ def update_refinable_mask( ) with torch.no_grad(): current_raw = self._u6_to_raw6(self.forward()) - self.refinable_mask = new_mask.to(device=self.device) + new_mask = self._normalize_refinable_mask(new_mask) + self.refinable_mask = new_mask self.fixed_mask = ~new_mask self.fixed_values = current_raw.clone() new_refinable = current_raw[self.refinable_mask].clone() @@ -1556,16 +1612,22 @@ def __init__( self._name = name or "occupancy" # Empty initialization + # + # ``OccupancyTensor`` builds its own shell rather than delegating to + # ``MixedTensor``, so the device/dtype handling has to be repeated here + # -- fixing only the base class would leave this path on CPU. if initial_values is None: + device = normalize_device(device) + dtype = dtype if dtype is not None else get_float_dtype() self._full_shape = 0 self._collapsed_shape = 0 self.register_buffer("refinable_mask", None) self.register_buffer("fixed_mask", None) self.register_buffer("fixed_values", None) - self.register_buffer("_shape", None) self.register_buffer("expansion_mask", None) self.refinable_params = nn.Parameter( - torch.empty(0), requires_grad=requires_grad + torch.empty(0, device=device, dtype=dtype), + requires_grad=requires_grad, ) self._has_refinable = False self._refinable_indices = None @@ -1644,8 +1706,11 @@ def __init__( refinable_values, requires_grad=requires_grad ) - # Store collapsed shape - self.register_buffer("_shape", torch.tensor([self._collapsed_shape])) + # Collapsed shape is host-side metadata, not a buffer -- same reasoning + # as ``MixedTensor`` (a buffer gets dragged onto the accelerator by + # ``.to()`` and then read back with a sync, purely to recover a number + # we already have). ``_collapsed_shape`` already holds it as an int; + # the ``shape`` property below derives from it. # Pre-compute index cache to avoid boolean indexing at runtime self._build_index_cache() diff --git a/torchref/model/rigid_xyz.py b/torchref/model/rigid_xyz.py index e99f56ed..1c40feb4 100644 --- a/torchref/model/rigid_xyz.py +++ b/torchref/model/rigid_xyz.py @@ -18,6 +18,7 @@ from torch import nn from torchref.base.alignment.rotation import rotation_matrix_euler_xyz +from torchref.config import get_float_dtype, normalize_device from torchref.utils.caching import CachedForwardMixin from torchref.utils.device_mixin import DeviceMixin @@ -56,14 +57,22 @@ def __init__( self._name = name if original_xyz is None: - # Empty init for state_dict loading. - self.register_buffer("original_xyz", torch.empty(0, 3)) - self.register_buffer("chain_indices", torch.empty(0, dtype=torch.long)) - self.register_buffer("chain_centers", torch.empty(0, 3)) - self.register_buffer("mobile_mask", torch.empty(0, dtype=torch.bool)) - self.register_buffer("atom_weights", torch.empty(0)) - self.euler_angles = nn.Parameter(torch.empty(0, 3)) - self.translations = nn.Parameter(torch.empty(0, 3)) + # Empty init for state_dict loading. Honour the requested + # device/dtype: otherwise every buffer lands on CPU regardless of + # what the caller asked for, and only a later ``.to()`` repairs it. + device = normalize_device(device) + dtype = dtype if dtype is not None else get_float_dtype() + self.register_buffer("original_xyz", torch.empty(0, 3, device=device, dtype=dtype)) + self.register_buffer( + "chain_indices", torch.empty(0, dtype=torch.long, device=device) + ) + self.register_buffer("chain_centers", torch.empty(0, 3, device=device, dtype=dtype)) + self.register_buffer( + "mobile_mask", torch.empty(0, dtype=torch.bool, device=device) + ) + self.register_buffer("atom_weights", torch.empty(0, device=device, dtype=dtype)) + self.euler_angles = nn.Parameter(torch.empty(0, 3, device=device, dtype=dtype)) + self.translations = nn.Parameter(torch.empty(0, 3, device=device, dtype=dtype)) self._n_chains = 0 self._chain_id_order: list = [] return diff --git a/torchref/model/sf_ds.py b/torchref/model/sf_ds.py index e122c5fd..6deee280 100644 --- a/torchref/model/sf_ds.py +++ b/torchref/model/sf_ds.py @@ -129,10 +129,11 @@ def __init__( super().__init__() if dtype_float is None: dtype_float = dtypes.float - if device is None: - device = get_default_device() self.dtype_float = dtype_float - self.device = device + # Derive from ``cell`` when no device is given, instead of jumping to + # the global default and leaving a caller-supplied cell behind on + # another device. An explicit ``device`` still wins and moves the cell. + self.device = resolve_device(cell, device=device) self.verbose = verbose self.max_memory_gb = max_memory_gb self.engine = engine @@ -142,7 +143,9 @@ def __init__( self._spacegroup = None if spacegroup is not None or cell is not None: - self._spacegroup = SpaceGroup(spacegroup, dtype=dtype_float, device=device) + self._spacegroup = SpaceGroup( + spacegroup, dtype=dtype_float, device=self.device + ) # Cache reciprocal basis matrix self._recB: Optional[torch.Tensor] = None @@ -201,7 +204,13 @@ def set_cell_and_spacegroup(self, cell: Cell, spacegroup: SpaceGroupLike = None) Unit cell object. spacegroup : SpaceGroupLike, optional Space group specification. + + Notes + ----- + Receiver wins: an incoming cell on another device is moved to match + this module rather than the other way round. """ + self.device = resolve_device(self, cell) self._cell = cell self._recB = None # Invalidate cache self.spacegroup = spacegroup diff --git a/torchref/model/sf_fft.py b/torchref/model/sf_fft.py index 182f6de7..d2deb8fb 100644 --- a/torchref/model/sf_fft.py +++ b/torchref/model/sf_fft.py @@ -25,6 +25,7 @@ from torchref.symmetry.spacegroup import SpaceGroupLike from torchref.utils.caching import ParameterFingerprint from torchref.utils.device_mixin import DeviceMovementMixin +from torchref.utils.device_resolution import resolve_device class SfFFT(DeviceMovementMixin, nn.Module): @@ -127,11 +128,10 @@ def __init__( dtype_float = dtypes.float self.dtype_float = dtype_float - self.device = ( - device - if device is not None - else cell.device if cell is not None else get_default_device() - ) + # One device for the module and everything it builds. ``resolve_device`` + # also moves ``cell`` when an explicit ``device`` disagrees with it, so + # the cell and the SpaceGroup below cannot end up split. + self.device = resolve_device(cell, device=device) self.verbose = verbose self.use_late_symmetry = use_late_symmetry @@ -141,7 +141,12 @@ def __init__( self._spacegroup = None if spacegroup is not None or cell is not None: - self._spacegroup = SpaceGroup(spacegroup, dtype=dtype_float, device=device) + # ``self.device``, not the raw ``device`` argument: the latter is + # ``None`` on the derive-from-cell path, which would silently put + # the symmetry matrices on the global default instead. + self._spacegroup = SpaceGroup( + spacegroup, dtype=dtype_float, device=self.device + ) # Buffers (registered during setup_grid) self.register_buffer("gridsize", None) @@ -227,7 +232,14 @@ def set_cell_and_spacegroup(self, cell: Cell, spacegroup: SpaceGroupLike = None) Unit cell object. spacegroup : SpaceGroupLike, optional Space group specification. + + Notes + ----- + Receiver wins: this module may already own grid buffers, so an incoming + cell on another device is moved to match rather than dragging the + module after it. """ + self.device = resolve_device(self, cell) self._cell = cell self.spacegroup = spacegroup @@ -305,8 +317,10 @@ def compute_real_space_grid( torch.Tensor Real-space grid with shape (nx, ny, nz, 3). """ - if device is None: - device = get_default_device() + # Forward ``device`` as-is, including ``None``. ``get_real_grid`` + # already infers from ``fractional_matrix`` when no device is given; + # resolving the global default here meant it never saw ``None`` and + # its inference was dead code. return get_real_grid( fractional_matrix=fractional_matrix, gridsize=gridsize, device=device ) diff --git a/torchref/refinement/base_refinement.py b/torchref/refinement/base_refinement.py index 79c21ab5..bfff2886 100644 --- a/torchref/refinement/base_refinement.py +++ b/torchref/refinement/base_refinement.py @@ -7,7 +7,7 @@ import torch from torch.nn import Module as nnModule -from torchref.config import get_default_device +from torchref.config import normalize_device from torchref.io import ReflectionData from torchref.model.model_ft import ModelFT from torchref.refinement.logger import Logger @@ -1207,8 +1207,7 @@ def create_from_state_dict( print(f"Restored at R-work={rwork:.4f}, R-free={rfree:.4f}") """ - if device is None: - device = get_default_device() + device = normalize_device(device) # Helper to extract submodule state from flattened state_dict def extract_submodule_state(state_dict: dict, prefix: str) -> dict: diff --git a/torchref/refinement/loss_state.py b/torchref/refinement/loss_state.py index dad74447..4fc5308c 100644 --- a/torchref/refinement/loss_state.py +++ b/torchref/refinement/loss_state.py @@ -31,7 +31,7 @@ import torch from torch import nn -from torchref.config import get_default_device +from torchref.config import canonical_device, get_default_device from torchref.utils.autograd_introspection import collect_loss_leaves, _iter_roots from torchref.utils.device_mixin import DeviceMovementMixin from torchref.utils.loss_validation import validate_loss @@ -276,8 +276,33 @@ def register_target( if probe: self._probe_and_merge_leaves(target) self._collect_resettable_modules(target) + self._warn_on_device_mismatch(key, target) return self + def _warn_on_device_mismatch(self, key: str, target: Callable) -> None: + """Warn if ``target`` does not agree with this state's device. + + Deliberately a warning and not a move. ``target.to(self.device)`` would + drag the model, data and scaler the target *borrows* along with it, so + registering a loss term could silently relocate an entire + ``ReflectionData``. Targets resolve their own device at construction + (see ``Target._adopt_device``); a disagreement here means something + upstream built them inconsistently, which is worth surfacing rather + than papering over. + """ + target_device = getattr(target, "device", None) + if target_device is None or not isinstance(self.device, torch.device): + return + if canonical_device(target_device) == canonical_device(self.device): + return + warnings.warn( + f"LossState: target {key!r} is on {target_device} but the state is " + f"on {self.device}. Not moving it -- a target borrows its model and " + "data, so relocating it here would move those too. Construct the " + "target on the right device instead.", + stacklevel=3, + ) + def _collect_resettable_modules(self, target: Callable) -> None: """Walk ``target``'s submodules and collect any that expose a ``reset_cache`` method, deduplicating against modules already @@ -914,23 +939,12 @@ def summary(self) -> None: # ========================================================================= # Device Management # ========================================================================= - - def to(self, *args, **kwargs): - """Move via :class:`DeviceMixin`; honour an explicit device when no tensors exist yet.""" - result = super().to(*args, **kwargs) - # If no tensor was found to refresh ``self.device``, fall back to the - # explicit device argument so subsequent allocations land correctly. - if not isinstance(result.device, torch.device): - from torchref.utils.device_mixin import _parse_to_args - - device, _ = _parse_to_args(args, kwargs) - if device is not None: - result.device = ( - torch.device(device) - if not isinstance(device, torch.device) - else device - ) - return result + # + # ``LossState`` normally owns no tensors (targets are callables, weights are + # floats), so its ``device`` tracker used to need a bespoke ``to()`` + # override to survive a move. ``DeviceMixin`` now carries the parsed request + # down the traversal and updates tensor-free objects itself, so the override + # is gone; ``self.device`` follows ``.to()`` via the shared machinery. # ========================================================================= # Utility diff --git a/torchref/refinement/targets/adp/base.py b/torchref/refinement/targets/adp/base.py index 3e8940c6..68675a9a 100644 --- a/torchref/refinement/targets/adp/base.py +++ b/torchref/refinement/targets/adp/base.py @@ -39,6 +39,7 @@ def __init__( self, model: "Model" = None, verbose: int = 0, + device=None, **kwargs, ): - super().__init__(model, verbose) + super().__init__(model, verbose, device=device) diff --git a/torchref/refinement/targets/adp/locality.py b/torchref/refinement/targets/adp/locality.py index 9dbf41ad..b43ec182 100644 --- a/torchref/refinement/targets/adp/locality.py +++ b/torchref/refinement/targets/adp/locality.py @@ -60,18 +60,22 @@ def __init__( sigma_aniso: float = 0.5, exclude_bonded: bool = True, verbose: int = 0, + device=None, ): - super().__init__(model, verbose) - self.register_buffer( - "_k_neighbors", torch.tensor(k_neighbors, dtype=torch.int64) - ) - self.register_buffer("_correlation_length", torch.tensor(correlation_length)) - self.register_buffer("_scale", torch.tensor(scale)) + super().__init__(model, verbose, device=device) + # Host-side, deliberately not buffers: all three are only ever reached + # through the properties below, which immediately ``.item()`` them, and + # their consumers (k-NN sizing, the exp() falloff, stats reporting) are + # host-side. Keeping them on the device bought a sync per access. + self._k_neighbors = int(k_neighbors) + self._correlation_length = float(correlation_length) + self._scale = float(scale) # Sigma for the deviatoric (anisotropy) channel; only used when # anisotropic atoms are present. Dimensionless and on the same scale as # the magnitude channel's fixed 0.5 log-sigma (the channel restrains - # fractional anisotropy dev/B_eq, the analogue of log B_eq). - self.register_buffer("_sigma_aniso", torch.tensor(sigma_aniso)) + # fractional anisotropy dev/B_eq, the analogue of log B_eq). This one + # *is* a buffer: it is handed to adp_locality_aniso_math as a tensor. + self._register_scalar("_sigma_aniso", float(sigma_aniso)) self.exclude_bonded = exclude_bonded # Cache for neighbor indices and distances @@ -79,29 +83,48 @@ def __init__( self._neighbor_distances = None # (N, k_neighbors) self._last_xyz_hash = None + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Absorb the retired ``_k_neighbors`` / ``_correlation_length`` / + ``_scale`` buffers. + + All three are host-side Python scalars now (see ``__init__``). + Checkpoints predating that change still carry them, and a + ``strict=True`` load would reject them as unexpected keys -- so restore + their values rather than discarding a user's tuning. + """ + for legacy, cast in ( + ("_k_neighbors", int), + ("_correlation_length", float), + ("_scale", float), + ): + saved = state_dict.pop(prefix + legacy, None) + if saved is not None: + setattr(self, legacy, cast(saved.item())) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + @property def k_neighbors(self) -> int: - return self._k_neighbors.item() + return self._k_neighbors @k_neighbors.setter def k_neighbors(self, value: int): - self._k_neighbors.fill_(value) + self._k_neighbors = int(value) @property def correlation_length(self) -> float: - return self._correlation_length.item() + return self._correlation_length @correlation_length.setter def correlation_length(self, value: float): - self._correlation_length.fill_(value) + self._correlation_length = float(value) @property def scale(self) -> float: - return self._scale.item() + return self._scale @scale.setter def scale(self, value: float): - self._scale.fill_(value) + self._scale = float(value) @property def sigma_aniso(self) -> float: @@ -287,11 +310,6 @@ def forward(self, recompute_neighbors: bool = False) -> torch.Tensor: # deviatoric/anisotropy channel). Isotropic-only models keep the # original B-factor path (numerically unchanged, zero overhead). if not getattr(self.model, "_aniso_is_empty", True): - if (self._sigma_aniso.device != device - or self._sigma_aniso.dtype != adp.dtype): - self._sigma_aniso = self._sigma_aniso.to( - device=device, dtype=adp.dtype - ) u6 = self.model.adp_u6() return adp_locality_aniso_math( u6, indices, distances, self._sigma_aniso diff --git a/torchref/refinement/targets/adp/similarity.py b/torchref/refinement/targets/adp/similarity.py index 78335a1c..45dd914b 100644 --- a/torchref/refinement/targets/adp/similarity.py +++ b/torchref/refinement/targets/adp/similarity.py @@ -42,12 +42,15 @@ class ADPSimilarityTarget(ADPTarget): def __init__( self, model: "Model" = None, simu_sigma: float = 2.0, - simu_sigma_aniso: float = 1.0, verbose: int = 0 + simu_sigma_aniso: float = 1.0, verbose: int = 0, device=None ): - super().__init__(model, verbose) - # Register simu-specific sigma as buffer (separate from base sigma) - self.register_buffer("_simu_sigma", torch.tensor(simu_sigma)) - self.register_buffer("_simu_sigma_aniso", torch.tensor(simu_sigma_aniso)) + super().__init__(model, verbose, device=device) + # Both sigmas are handed to adp_simu_math / adp_simu_aniso_math, which + # dispatch to Triton on CUDA float32. Allocating them on the target's + # resolved device and dtype is what lets the lazy repair inside + # forward() go away. + self._register_scalar("_simu_sigma", float(simu_sigma)) + self._register_scalar("_simu_sigma_aniso", float(simu_sigma_aniso)) @property def simu_sigma(self) -> float: @@ -96,25 +99,10 @@ def forward(self) -> torch.Tensor: adp_t = self.model.adp() if pair_indices.shape[0] == 0: return torch.zeros((), device=adp_t.device, dtype=adp_t.dtype) - # Lazily move the ``_simu_sigma`` buffer onto the model's device - # the first time we reach here. Once moved, subsequent forwards - # (and CUDA-Graph captures) skip the device transfer — calling - # ``.to()`` on a CPU buffer inside a capture region triggers a - # ``cudaErrorStreamCaptureUnsupported``. - if (self._simu_sigma.device != adp_t.device - or self._simu_sigma.dtype != adp_t.dtype): - self._simu_sigma = self._simu_sigma.to( - device=adp_t.device, dtype=adp_t.dtype, - ) # Anisotropic atoms present: restrain the full U tensors. Isotropic-only # models keep the original (Triton-accelerated) B-factor path so they # pay nothing and are numerically unchanged. if not getattr(self.model, "_aniso_is_empty", True): - if (self._simu_sigma_aniso.device != adp_t.device - or self._simu_sigma_aniso.dtype != adp_t.dtype): - self._simu_sigma_aniso = self._simu_sigma_aniso.to( - device=adp_t.device, dtype=adp_t.dtype, - ) u6 = self.model.adp_u6() return adp_simu_aniso_math( u6, pair_indices, self._simu_sigma, self._simu_sigma_aniso diff --git a/torchref/refinement/targets/base.py b/torchref/refinement/targets/base.py index 237b53af..7ea97cd2 100644 --- a/torchref/refinement/targets/base.py +++ b/torchref/refinement/targets/base.py @@ -21,7 +21,9 @@ from torch import nn from torch.special import i0 +from torchref.config import get_float_dtype, normalize_device from torchref.utils.device_mixin import DeviceMixin +from torchref.utils.device_resolution import resolve_device from torchref.utils.stats import ( VERBOSITY_DEBUG, VERBOSITY_DETAILED, @@ -80,6 +82,7 @@ class Target(DeviceMixin, nn.Module): def __init__( self, verbose: int = 0, + device=None, **kwargs, ): """ @@ -89,9 +92,79 @@ def __init__( ---------- verbose : int, optional Verbosity level. Default is 0. + device : torch.device, optional + Where this target allocates. Defaults to the configured default; + :meth:`_adopt_device` refines it from whatever model / data / + scaler the subclass is given. """ super().__init__() self.verbose = verbose + # Seeded here for two distinct reasons. + # + # (1) Subclasses allocate their tunables in ``__init__``, so + # ``device=self.device`` / ``dtype=self.dtype_float`` must already + # have an answer at that point. Before this existed, every scalar + # tunable in the package landed on CPU float32 regardless of the + # model's device or the configured dtype. + # + # (2) ``DeviceMixin._refresh_device_trackers`` early-returns unless one + # of these names is already in ``obj.__dict__``. Merely defining + # them is what switches the mixin on for targets -- it was + # previously a complete no-op for every target in the package. + self.device = normalize_device(device) + self.dtype_float = get_float_dtype() + + # ---- device / dtype resolution -------------------------------------- + + def _adopt_device(self, *sources, device=None): + """Reconcile this target with the device-bearing objects it wraps. + + Call *after* the sources are attached and *before* allocating any + buffer. ``None`` sources are dropped, so the empty-init path used by + ``load_state_dict`` keeps the seeded default, and the method stays + re-runnable if a source is attached later. + + ``self`` is deliberately **not** passed to + :func:`~torchref.utils.resolve_device` -- unlike + ``ScalerBase.set_data``, which does pass it. A freshly constructed + target owns no tensors, so there is nothing of its own to reconcile, + and a target's state is a handful of scalars against a ``Model``'s + entire structure. Letting an empty shell's default device outvote a + model already on the accelerator would move megabytes to satisfy five + floats. + + Returns + ------- + torch.device + The resolved device (also stored on ``self.device``). + """ + present = [s for s in sources if s is not None] + if device is None and not present: + return self.device + self.device = resolve_device(*present, device=device) + for src in present: + src_dtype = getattr(src, "dtype_float", None) + if isinstance(src_dtype, torch.dtype): + self.dtype_float = src_dtype + break + return self.device + + def _register_scalar(self, name: str, value, dtype=None) -> None: + """Register a 0-dim tunable on this target's device and float dtype. + + Scalar tunables reach Triton kernels as tensors, where a CPU-resident + buffer is a raw pointer into host memory rather than a promotable + scalar. Allocating them correctly here is what lets the lazy ``.to()`` + repairs inside ``forward()`` go away. + """ + self.register_buffer( + name, + torch.tensor( + value, + device=self.device, + dtype=self.dtype_float if dtype is None else dtype, + ), + ) def forward(self) -> torch.Tensor: """Compute and return the loss. Override in subclasses.""" @@ -181,6 +254,7 @@ def __init__( self, model: "Model" = None, verbose: int = 0, + device=None, **kwargs, ): """ @@ -192,11 +266,16 @@ def __init__( Reference to the Model object (optional for empty init). verbose : int, optional Verbosity level. Default is 0. + device : torch.device, optional + Explicit device. When omitted, follows ``model``. """ - super().__init__(verbose=verbose) + super().__init__(verbose=verbose, device=device) # Register model as a proper submodule (not in state_dict but handles device) # Use add_module to allow None values self.add_module("_model", model) + # Follow the model rather than the global default, so subclass buffers + # allocated below this line land beside the coordinates they act on. + self._adopt_device(model, device=device) @property def model(self) -> "Model": @@ -268,6 +347,7 @@ def __init__( model: "Model" = None, scaler: "Scaler" = None, verbose: int = 0, + device=None, **kwargs, ): """ @@ -284,12 +364,21 @@ def __init__( Reference to the Scaler object. verbose : int, optional Verbosity level. Default is 0. + device : torch.device, optional + Explicit device. When omitted, model / data / scaler are + reconciled onto one device, the model winning on disagreement. """ - super().__init__(verbose=verbose) - # Register as proper submodules (allows None values) + super().__init__(verbose=verbose, device=device) + # Register as proper submodules (allows None values). Note ``_data`` is + # a plain attribute: ReflectionData is a dataclass, not an nn.Module, + # so add_module would reject it -- but DeviceMixin still walks it. self.add_module("_model", model) self._data = data self.add_module("_scaler", scaler) + # The forward path mixes tensors from all three, so pin them together + # here. Order is precedence: ``resolve_device`` is first-wins, and the + # model is what the rest of the pipeline follows. + self._adopt_device(model, data, scaler, device=device) @property def model(self) -> "Model": diff --git a/torchref/refinement/targets/geometry/base.py b/torchref/refinement/targets/geometry/base.py index d9908591..04c7c794 100644 --- a/torchref/refinement/targets/geometry/base.py +++ b/torchref/refinement/targets/geometry/base.py @@ -39,9 +39,10 @@ def __init__( self, model: "Model" = None, verbose: int = 0, + device=None, **kwargs, ): - super().__init__(model, verbose) + super().__init__(model, verbose, device=device) def stats(self) -> Dict[str, StatEntry]: """ diff --git a/torchref/refinement/targets/geometry/non_bonded.py b/torchref/refinement/targets/geometry/non_bonded.py index d4f457bb..1d61abc5 100644 --- a/torchref/refinement/targets/geometry/non_bonded.py +++ b/torchref/refinement/targets/geometry/non_bonded.py @@ -109,6 +109,7 @@ def __init__( rebuild_threshold: float = 1.0, verbose: int = 0, scale: float = 10.0, + device=None, ): """ Initialize non-bonded target. @@ -142,23 +143,50 @@ def __init__( Public attribute stored as ``self.scale``. Default is 10.0. Set but not consumed by ``forward()`` (informational only). """ - super().__init__(model, verbose) + super().__init__(model, verbose, device=device) self.mode = mode self.scale = scale - # Register sigma / r_exp / buffer as buffers so .to(device) moves them. - self.register_buffer("_sigma_vdw", torch.tensor(float(sigma))) - self.register_buffer("_r_exp", torch.tensor(float(r_exp))) - self.register_buffer("_buffer", torch.tensor(float(buffer))) - self.register_buffer( - "_rebuild_threshold", torch.tensor(float(rebuild_threshold)) - ) + # Tunables that reach the kernel live on the target's resolved device + # and float dtype: the prolsq branch hands these straight to a Triton + # kernel, where a CPU tensor is a raw pointer into host memory rather + # than a promotable scalar. + self._register_scalar("_sigma_vdw", float(sigma)) + self._register_scalar("_r_exp", float(r_exp)) # c_rep: back-door override; by default derived from sigma so that # PROLSQ shape term equals v^p / (p * sigma^p). if c_rep is None: c_rep_val = 1.0 / (float(r_exp) * float(sigma) ** float(r_exp)) else: c_rep_val = float(c_rep) - self.register_buffer("_c_rep", torch.tensor(c_rep_val)) + self._register_scalar("_c_rep", c_rep_val) + # Host-side, deliberately not buffers. + # + # ``buffer`` is consumed as a Python float -- forward() already called + # ``float(self._buffer)`` before handing it to the kernel, where it + # becomes a compile-time constant. Storing it on the device only bought + # a device->host sync on the hot path. + # + # ``rebuild_threshold`` is only ever compared against an already + # ``.item()``-ed displacement in maintenance(). + self._buffer = float(buffer) + self._rebuild_threshold = float(rebuild_threshold) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Absorb the retired ``_buffer`` / ``_rebuild_threshold`` buffers. + + Both are host-side Python floats now (see ``__init__``). Checkpoints + written before that change still carry them, and a ``strict=True`` load + would reject them as unexpected keys -- so restore their values rather + than discarding a user's tuning. + """ + for legacy, attr in ( + ("_buffer", "_buffer"), + ("_rebuild_threshold", "_rebuild_threshold"), + ): + saved = state_dict.pop(prefix + legacy, None) + if saved is not None: + setattr(self, attr, float(saved.item())) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) @property def c_rep(self) -> float: @@ -199,13 +227,13 @@ def r_exp(self, value: float): @property def buffer(self) -> float: - """Get distance buffer.""" - return self._buffer.item() + """Get distance buffer (host-side; see ``__init__``).""" + return self._buffer @buffer.setter def buffer(self, value: float): """Set distance buffer.""" - self._buffer.fill_(value) + self._buffer = float(value) def maintenance(self) -> None: """Rebuild the VDW pair list if any ASU atom drifted too far. @@ -245,12 +273,12 @@ def maintenance(self) -> None: max_disp_sq = (delta * delta).sum(dim=-1).max() thresh_sq = self._rebuild_threshold * self._rebuild_threshold - if max_disp_sq.item() <= thresh_sq.item(): + if max_disp_sq.item() <= thresh_sq: return # within slack — nothing to do if self.verbose > 0: max_disp = float(max_disp_sq.item()) ** 0.5 - thresh = float(self._rebuild_threshold.item()) + thresh = self._rebuild_threshold print( f" VDW rebuild: max drift {max_disp:.2f} Å > " f"threshold {thresh:.2f} Å" @@ -358,7 +386,7 @@ def forward(self) -> torch.Tensor: self.model.cell.fractional_matrix, self.model.cell.inv_fractional_matrix, self._c_rep, self._r_exp, - float(self._buffer), self._sigma_vdw, + self._buffer, self._sigma_vdw, ) # Compute positions (handles symmetry transparently) diff --git a/torchref/refinement/targets/similarity.py b/torchref/refinement/targets/similarity.py index 01872a92..83edc2c6 100644 --- a/torchref/refinement/targets/similarity.py +++ b/torchref/refinement/targets/similarity.py @@ -72,11 +72,26 @@ def __init__( model_light: "Model" = None, alpha: float = 2.0, verbose: int = 0, + device=None, ): - super().__init__(verbose=verbose) + super().__init__(verbose=verbose, device=device) self.add_module("_model_dark", model_dark) self.add_module("_model_light", model_light) - self.register_buffer("_alpha", torch.tensor(alpha)) + self._adopt_device(model_dark, model_light, device=device) + self._register_scalar("_alpha", float(alpha)) + # Registered unconditionally, before the map is built. They used to be + # created only inside ``_build_atom_map``, which runs only when both + # models are present -- so on the empty-init path (the one that exists + # for ``load_state_dict``) they never existed at all: ``forward()`` + # raised AttributeError and a strict load of a populated checkpoint + # failed on unexpected keys. ``_build_atom_map`` now overwrites rather + # than creates. + self.register_buffer( + "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) + ) + self.register_buffer( + "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) + ) if model_dark is not None and model_light is not None: self._build_atom_map() @@ -148,10 +163,10 @@ def _build_atom_map(self): "dark and light models" ) self.register_buffer( - "_idx_dark", torch.zeros(0, dtype=torch.long) + "_idx_dark", torch.zeros(0, dtype=torch.long, device=self.device) ) self.register_buffer( - "_idx_light", torch.zeros(0, dtype=torch.long) + "_idx_light", torch.zeros(0, dtype=torch.long, device=self.device) ) return @@ -168,13 +183,21 @@ def _build_atom_map(self): f"(dark={n_dark}, light={n_light})" ) + # On ``self.device``: these are 1-D index tensors, so unlike the 0-dim + # scalars they are not covered by PyTorch's scalar-promotion rule. A + # CPU-resident index against accelerator coordinates costs a host sync + # on every forward. self.register_buffer( "_idx_dark", - torch.tensor(merged["_idx_dark"].values, dtype=torch.long), + torch.tensor( + merged["_idx_dark"].values, dtype=torch.long, device=self.device + ), ) self.register_buffer( "_idx_light", - torch.tensor(merged["_idx_light"].values, dtype=torch.long), + torch.tensor( + merged["_idx_light"].values, dtype=torch.long, device=self.device + ), ) def forward(self) -> torch.Tensor: @@ -186,8 +209,9 @@ def forward(self) -> torch.Tensor: Scalar mean loss over all matched atom pairs. """ if len(self._idx_dark) == 0: - device = self._alpha.device - return torch.tensor(0.0, device=device) + # ``self.device``, not ``self._alpha.device``: an empty target must + # still hand back a loss on the refinement's device. + return torch.tensor(0.0, device=self.device, dtype=self.dtype_float) xyz_dark = self._model_dark.xyz() xyz_light = self._model_light.xyz() diff --git a/torchref/refinement/targets/xray/base.py b/torchref/refinement/targets/xray/base.py index c0ccb425..b74f606e 100644 --- a/torchref/refinement/targets/xray/base.py +++ b/torchref/refinement/targets/xray/base.py @@ -67,6 +67,7 @@ def __init__( use_work_set: bool = True, verbose: int = 0, use_set: str = None, + device=None, ): """ Initialize X-ray target. @@ -89,7 +90,9 @@ def __init__( verbose : int, optional Verbosity level. Default is 0. """ - super().__init__(data=data, model=model, scaler=scaler, verbose=verbose) + super().__init__( + data=data, model=model, scaler=scaler, verbose=verbose, device=device + ) # ``use_set`` (3-way: "work"/"free"/"val") is the canonical subset # selector; the legacy ``use_work_set`` bool maps onto it. When # ``use_set`` is not given explicitly, fall back to the bool so older diff --git a/torchref/refinement/targets/xray/bhattacharyya.py b/torchref/refinement/targets/xray/bhattacharyya.py index 04f3aeb9..fe3cb702 100644 --- a/torchref/refinement/targets/xray/bhattacharyya.py +++ b/torchref/refinement/targets/xray/bhattacharyya.py @@ -107,30 +107,49 @@ def __init__( verbose=verbose, use_set=use_set, ) - # log-spaced B grid + # log-spaced B grid. ``log_b_grid`` itself is a construction-time + # local, not a buffer: it is read twice below for the min/max and never + # touched again, so registering it only created something for ``.to()`` + # to carry around. log_b_grid = torch.linspace( - math.log(b_grid_min), math.log(b_grid_max), b_grid_n + math.log(b_grid_min), + math.log(b_grid_max), + b_grid_n, + device=self.device, + dtype=self.dtype_float, ) - self.register_buffer("log_b_grid", log_b_grid) self.register_buffer("b_grid", torch.exp(log_b_grid)) self._log_b_min = float(log_b_grid[0].item()) self._log_b_max = float(log_b_grid[-1].item()) self._log_b_step = (self._log_b_max - self._log_b_min) / (b_grid_n - 1) # Global σ_m scale (tunable, non-learnable) + self._register_scalar("sigma_m_scale", float(sigma_m_scale)) + # Populated by _initialize_cache() on first forward(). Empty, but still + # allocated on this target's device: a 0-element tensor reports a + # device, and a CPU one here reads as a split object. + empty = dict(device=self.device, dtype=self.dtype_float) + self.register_buffer("exp_table", torch.empty(0, **empty)) # (b_grid_n, N_refl) + self.register_buffer("s_sq_per_refl", torch.empty(0, **empty)) # (N_refl,) + self.register_buffer("s_4_per_refl", torch.empty(0, **empty)) # (N_refl,) + self.register_buffer("f_sq_kh", torch.empty(0, **empty)) # (K, N_refl) + self.register_buffer("g_w_table", torch.empty(0, **empty)) # (K, b_grid_n) + self.register_buffer("g_4_table", torch.empty(0, **empty)) # (K, b_grid_n) self.register_buffer( - "sigma_m_scale", torch.tensor(float(sigma_m_scale)) + "atom_to_element", torch.empty(0, dtype=torch.long, device=self.device) ) - # Populated by _initialize_cache() on first forward() - self.register_buffer("exp_table", torch.empty(0)) # (b_grid_n, N_refl) - self.register_buffer("s_sq_per_refl", torch.empty(0)) # (N_refl,) - self.register_buffer("s_4_per_refl", torch.empty(0)) # (N_refl,) - self.register_buffer("f_sq_kh", torch.empty(0)) # (K, N_refl) - self.register_buffer("g_w_table", torch.empty(0)) # (K, b_grid_n) - self.register_buffer("g_4_table", torch.empty(0)) # (K, b_grid_n) - self.register_buffer("atom_to_element", torch.empty(0, dtype=torch.long)) - self.register_buffer("sigma_d_mean", torch.tensor(0.0)) + self._register_scalar("sigma_d_mean", 0.0) self._initialized = False + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Drop the retired ``log_b_grid`` buffer key. + + It is a construction-time local now (derived from the same + ``b_grid_min``/``max``/``n`` the constructor already has), so a + ``strict=True`` load of an older checkpoint would reject it. + """ + state_dict.pop(prefix + "log_b_grid", None) + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + # ------------------------------------------------------------------ # Cache initialisation (called once, on first forward) # ------------------------------------------------------------------ diff --git a/torchref/refinement/targets/xray/least_squares.py b/torchref/refinement/targets/xray/least_squares.py index 4dbe340f..5fdff64e 100644 --- a/torchref/refinement/targets/xray/least_squares.py +++ b/torchref/refinement/targets/xray/least_squares.py @@ -46,6 +46,7 @@ def __init__( n_bins: int = 20, verbose: int = 0, use_set: str = None, + device=None, ): if scale_mode not in ("scaler", "binwise_optimal"): raise ValueError( @@ -64,6 +65,7 @@ def __init__( use_work_set=use_work_set, verbose=verbose, use_set=use_set, + device=device, ) self.weighting = weighting self.scale_mode = scale_mode diff --git a/torchref/refinement/targets/xray/maximum_likelihood.py b/torchref/refinement/targets/xray/maximum_likelihood.py index af36f00c..65878db7 100644 --- a/torchref/refinement/targets/xray/maximum_likelihood.py +++ b/torchref/refinement/targets/xray/maximum_likelihood.py @@ -8,7 +8,6 @@ epsilon_from_hkl, ml_xray_loss_beta_math, ) -from torchref.utils.device_resolution import resolve_device from .base import XrayTarget from .gaussian import GaussianXrayTarget @@ -178,10 +177,10 @@ def create_xray_target( ) mode = "ml" - # Pin model/data/scaler onto one device before constructing the - # target — its forward path mixes tensors from all three. - resolve_device(model, data, scaler, device=device) - + # Device reconciliation is ``DataTarget.__init__``'s job now (it calls + # ``_adopt_device(model, data, scaler)`` before allocating anything), so + # doing it again here would be a second copy of the same policy, free to + # drift. ``device`` is forwarded instead. kwargs = dict( data=data, model=model, @@ -189,6 +188,7 @@ def create_xray_target( use_work_set=use_work_set, verbose=verbose, use_set=use_set, + device=device, ) if mode == "gaussian": return GaussianXrayTarget(**kwargs) diff --git a/torchref/restraints/hydrogen_topology.py b/torchref/restraints/hydrogen_topology.py index b4b8e94f..9d0ad649 100644 --- a/torchref/restraints/hydrogen_topology.py +++ b/torchref/restraints/hydrogen_topology.py @@ -17,7 +17,8 @@ import torch from torch import nn -from torchref.config import dtypes, get_default_device +from torchref.config import dtypes, normalize_device +from torchref.utils.device_resolution import resolve_device from torchref.utils.device_mixin import DeviceMixin # --------------------------------------------------------------------------- @@ -83,9 +84,14 @@ class HydrogenTopology(DeviceMixin, nn.Module): ``build_h_candidate_pairs`` and checked via :attr:`has_candidates`. """ - def __init__(self): + def __init__(self, device=None): super().__init__() - # Buffers are registered by build_hydrogen_topology() + # Buffers are registered later by build_hydrogen_topology(), so this + # object is tensor-free at construction. The tracker still has to exist: + # ``DeviceMixin._refresh_device_trackers`` only maintains attributes + # already present in ``__dict__``, and callers reconcile against + # ``h_topo.device`` before attaching buffers to it. + self.device = normalize_device(device) @property def n_hydrogens(self) -> int: @@ -250,8 +256,7 @@ def build_hydrogen_topology( HydrogenTopology Module with registered buffer tensors. """ - if device is None: - device = get_default_device() + device = normalize_device(device) cache = _load_cif_hydrogen_info(pdb, verbose) model_names = pdb["name"].astype(str).str.strip().values @@ -366,7 +371,9 @@ def build_hydrogen_topology( acc_chainid_enc.append(chainid_enc) acc_resseq.append(resseq) - topo = HydrogenTopology() + # Seed the tracker with the device its buffers are about to be built on, + # so a later ``resolve_device(h_topo, ...)`` sees the truth. + topo = HydrogenTopology(device=device) n_h_total = len(acc_parent_idx) fdtype = dtypes.float @@ -719,8 +726,10 @@ def build_h_candidate_pairs( device : torch.device verbose : int """ - if device is None: - device = get_default_device() + # Buffers are registered onto ``h_topo`` below, so follow it rather than + # the global default -- otherwise they attach to a module living somewhere + # else. + device = resolve_device(h_topo, device=device) n_h = h_topo.n_hydrogens n_heavy = len(pdb) diff --git a/torchref/restraints/restraints.py b/torchref/restraints/restraints.py index ada6ef82..f3da1fa7 100644 --- a/torchref/restraints/restraints.py +++ b/torchref/restraints/restraints.py @@ -29,7 +29,7 @@ read_cif, read_link_definitions, ) -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype from torchref.utils.debug_utils import DebugMixin from torchref.utils.utils import TensorDict from torchref.utils.device_mixin import DeviceMixin diff --git a/torchref/scaling/collection_scaler.py b/torchref/scaling/collection_scaler.py index 2683ea4b..7b8a82e6 100644 --- a/torchref/scaling/collection_scaler.py +++ b/torchref/scaling/collection_scaler.py @@ -26,7 +26,7 @@ epsilon_from_hkl, ml_xray_loss_beta_math, ) -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype from torchref.scaling.scaler_base import ScalerBase from torchref.scaling.solvent import SolventModel from torchref.utils.utils import ModuleReference @@ -80,10 +80,12 @@ def __init__( verbose: int = 1, device: torch.device = None, ): - if device is None: - device = get_default_device() # Bind to the dark/reference dataset for bins and scattering vectors dark_data = dataset_collection[model_collection.dark_key] + # Forward ``device`` through rather than resolving the global default + # here: a resolved default arrives at ``ScalerBase`` as an *explicit* + # device, which then drags ``dark_data`` onto it. Passing ``None`` + # lets ScalerBase derive from the data, which is the point. super().__init__( data=dark_data, nbins=nbins, diff --git a/torchref/scaling/scaler.py b/torchref/scaling/scaler.py index 59a8b3b9..d29ff0f4 100644 --- a/torchref/scaling/scaler.py +++ b/torchref/scaling/scaler.py @@ -19,7 +19,6 @@ import torch import torch.nn as nn -from torchref.config import get_default_device from torchref.io import ReflectionData from torchref.base.reciprocal import get_scattering_vectors from torchref.scaling.scaler_base import ScalerBase @@ -157,7 +156,15 @@ def set_model_and_data(self, model: "Model", data: ReflectionData): Model object for structure factor calculation. data : ReflectionData ReflectionData object with observed data. + + Notes + ----- + Receiver wins: the scaler already owns buffers by this point, so + ``model`` and ``data`` are reconciled onto *its* device. This is the + state-dict restore path, where the three objects are built separately + and can easily disagree. """ + resolve_device(self, model, data) # Set _model_ref directly: nn.Module.__setattr__ intercepts assignments # of nn.Module instances (like `self.model = model`) and registers them # as submodules, bypassing the property setter entirely. diff --git a/torchref/scaling/scaler_base.py b/torchref/scaling/scaler_base.py index 18706063..b94f967e 100644 --- a/torchref/scaling/scaler_base.py +++ b/torchref/scaling/scaler_base.py @@ -30,7 +30,7 @@ epsilon_from_hkl, ml_xray_loss_beta_math, ) -from torchref.config import get_complex_dtype, get_default_device, get_float_dtype +from torchref.config import get_complex_dtype, get_float_dtype from torchref.utils.autograd_ops import gather_with_index_add from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMixin @@ -158,7 +158,14 @@ def set_data(self, data: "ReflectionData"): ---------- data : ReflectionData ReflectionData object with observed data. + + Notes + ----- + Receiver wins: this scaler may already hold buffers, so ``data`` is + moved onto the scaler's device rather than the buffers being registered + from whatever device ``data`` happened to be on. """ + self.device = resolve_device(self, data) self._data = ModuleReference(data) if data.cell is not None: self.cell = data.cell diff --git a/torchref/scaling/solvent.py b/torchref/scaling/solvent.py index 131d036f..90a6c65c 100644 --- a/torchref/scaling/solvent.py +++ b/torchref/scaling/solvent.py @@ -11,9 +11,10 @@ ifft, ) from torchref.base.electron_density.main import _get_radius_offsets -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMixin +from torchref.utils.device_resolution import resolve_device from torchref.utils.utils import ModuleReference, TensorDict @@ -113,8 +114,11 @@ def __init__( super(SolventModel, self).__init__() if float_type is None: float_type = get_float_dtype() - if device is None: - device = get_default_device() + # Follow the model when no device is given. ``realspace.py`` constructs + # a SolventModel with only a model, so falling back to the global + # default here put the solvent grids on a different device than the + # structure they are computed from. + device = resolve_device(model, device=device) self.device = device self.verbose = verbose self.float_type = float_type diff --git a/torchref/symmetry/cell.py b/torchref/symmetry/cell.py index 18d42f0c..88f5104f 100644 --- a/torchref/symmetry/cell.py +++ b/torchref/symmetry/cell.py @@ -14,8 +14,8 @@ from torchref.config import ( NYQUIST_OVERSAMPLING, - get_default_device, get_float_dtype, + normalize_device, ) from torchref.utils.device_mixin import _NonModuleDeviceMixin @@ -76,8 +76,7 @@ def __init__( """ if dtype is None: dtype = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) # Convert to tensor first to get shape if isinstance(data, torch.Tensor): tensor = data.to(dtype=dtype, device=device) diff --git a/torchref/symmetry/map_symmetry.py b/torchref/symmetry/map_symmetry.py index ae5dc613..457641e7 100644 --- a/torchref/symmetry/map_symmetry.py +++ b/torchref/symmetry/map_symmetry.py @@ -14,7 +14,7 @@ import torch import torch.nn as nn -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.symmetry.spacegroup import SpaceGroup, SpaceGroupLike from torchref.utils.device_mixin import DeviceMixin @@ -57,8 +57,12 @@ def MapSymmetry( """ if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + # ``cell_params`` is documented as a tensor, so follow it when no device is + # given rather than jumping to the global default and leaving the caller's + # cell behind. + if device is None and isinstance(cell_params, torch.Tensor): + device = cell_params.device + device = normalize_device(device) # Check grid compatibility symmetry = SpaceGroup(space_group, dtype=dtype_float, device=device) compat = symmetry.check_grid_compatibility(map_shape) @@ -117,14 +121,24 @@ def __init__( super().__init__() if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + if device is None and isinstance(cell_params, torch.Tensor): + device = cell_params.device self.dtype_float = dtype_float self.space_group = space_group self.map_shape = tuple(map_shape) - self.cell_params = cell_params self.verbose = verbose - self.device = device + self.device = normalize_device(device) + # Store ``cell_params`` *on* this module's device. It used to be kept + # exactly as handed in -- a plain attribute, so ``DeviceMixin`` could + # neither see it nor repair it, and ``self.device`` could disagree with + # it indefinitely. + if isinstance(cell_params, torch.Tensor): + cell_params = cell_params.to(device=self.device, dtype=self.dtype_float) + else: + cell_params = torch.as_tensor( + cell_params, device=self.device, dtype=self.dtype_float + ) + self.cell_params = cell_params self.symmetry = SpaceGroup( space_group, dtype=self.dtype_float, device=self.device diff --git a/torchref/symmetry/map_symmetry_interpolation.py b/torchref/symmetry/map_symmetry_interpolation.py index 3debb3a2..b11213a7 100644 --- a/torchref/symmetry/map_symmetry_interpolation.py +++ b/torchref/symmetry/map_symmetry_interpolation.py @@ -10,7 +10,7 @@ import torch.nn as nn import torch.nn.functional as F -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.symmetry.spacegroup import SpaceGroup from torchref.utils.device_mixin import DeviceMixin @@ -82,8 +82,7 @@ def __init__( super().__init__() if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) self.dtype_float = dtype_float self.space_group = space_group self.map_shape = tuple(map_shape) diff --git a/torchref/symmetry/reciprocal_symmetry.py b/torchref/symmetry/reciprocal_symmetry.py index ad7070b6..9745ca03 100644 --- a/torchref/symmetry/reciprocal_symmetry.py +++ b/torchref/symmetry/reciprocal_symmetry.py @@ -30,7 +30,7 @@ import torch import torch.nn as nn -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.symmetry.spacegroup import SpaceGroup, SpaceGroupLike from torchref.utils.device_mixin import DeviceMixin @@ -133,8 +133,7 @@ def __init__( super().__init__() if dtype_float is None: dtype_float = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) self.dtype_float = dtype_float self.space_group = space_group self.grid_shape = tuple(grid_shape) diff --git a/torchref/symmetry/spacegroup.py b/torchref/symmetry/spacegroup.py index 02f4ecda..dda5067f 100644 --- a/torchref/symmetry/spacegroup.py +++ b/torchref/symmetry/spacegroup.py @@ -36,7 +36,7 @@ import torch import torch.nn as nn -from torchref.config import get_default_device, get_float_dtype +from torchref.config import get_float_dtype, normalize_device from torchref.utils.debug_utils import DebugMixin from torchref.utils.device_mixin import DeviceMovementMixin @@ -212,8 +212,7 @@ def get_operations_as_tensors( """ if dtype is None: dtype = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) sg = _normalize_spacegroup(spacegroup) # Extract rotation matrices and translations from gemmi operations @@ -692,8 +691,7 @@ def __init__( super(SpaceGroup, self).__init__() if dtype is None: dtype = get_float_dtype() - if device is None: - device = get_default_device() + device = normalize_device(device) self._device = device self._dtype = dtype @@ -864,9 +862,15 @@ def apply_to_hkl(self, hkl: torch.Tensor) -> torch.Tensor: """ Apply symmetry operations to Miller indices (rotation only, no translation). - For reciprocal space, only the rotational part of symmetry operations - applies to Miller indices: h' = R·h. The translation vector affects the - phase of structure factors, not the indices themselves. + For reciprocal space, the symmetry-equivalent Miller index is obtained + by the *transpose* of the (real-space, fractional) rotation matrix: + ``h' = h·R = Rᵀ·h``. This is NOT the same as the real-space transform + ``x' = R·x`` unless ``R`` is symmetric: for space groups whose fractional + rotation matrices are non-symmetric (trigonal, hexagonal, and permutation- + type cubic operations) applying ``R·h`` instead of ``Rᵀ·h`` yields the + wrong set of equivalents, corrupting centric-flag and epsilon-multiplicity + determination. The translation vector affects the phase of structure + factors, not the indices themselves, so it is not applied here. Parameters ---------- @@ -881,9 +885,12 @@ def apply_to_hkl(self, hkl: torch.Tensor) -> torch.Tensor: See Also -------- - apply : For real space coordinates (rotation + translation). + apply : For real space coordinates (rotation + translation), which uses + ``R·x`` and is the transpose of the reciprocal-space convention here. """ - return self.apply(hkl, apply_translation=False) + coords = hkl.to(self.matrices.device).to(self.matrices.dtype) + # result[n, i, o] = sum_j matrices[o, j, i] * coords[n, j] = (Rᵀ·h)_i + return torch.einsum("oji,nj->nio", self.matrices, coords) def expand_coords_to_P1(self, xyz_fractional: torch.Tensor) -> torch.Tensor: """ diff --git a/torchref/utils/device_mixin.py b/torchref/utils/device_mixin.py index f3533305..72f8ecf4 100644 --- a/torchref/utils/device_mixin.py +++ b/torchref/utils/device_mixin.py @@ -35,11 +35,33 @@ class MyDataclass(DeviceMixin): from __future__ import annotations +import inspect import threading import torch from torch import nn +from torchref.config import canonical_device + +# ``torch._C._nn._parse_to`` is the same parser ``nn.Module.to`` uses, so it +# accepts every overload PyTorch does -- including ``.to(other_tensor)`` and +# ``.to(0)``, which the hand-rolled ``_parse_to_args`` fallback below cannot +# express. It is a private API, so probe once at import rather than catching +# per call: catching ``TypeError``/``RuntimeError`` around each invocation +# would turn a genuinely invalid user argument into a silent no-op. +try: + _PARSE_TO = torch._C._nn._parse_to +except AttributeError: # pragma: no cover - defensive against API drift + _PARSE_TO = None + +# Whether this torch exposes ``nn.Module._apply(fn, recurse=...)``. Detected +# once by signature rather than by catching ``TypeError`` at the call site: a +# blanket catch there would also swallow a genuine ``TypeError`` raised *inside* +# ``_apply`` after some tensors had already moved, and retry the move. +_MODULE_APPLY_TAKES_RECURSE = ( + "recurse" in inspect.signature(nn.Module._apply).parameters +) + # --------------------------------------------------------------------------- # Thread-local traversal state (cycle detection across one top-level .to()) # --------------------------------------------------------------------------- @@ -71,11 +93,50 @@ class MyDataclass(DeviceMixin): ) +class _ToRequest: + """The ``(device, dtype)`` a top-level ``.to()`` asked for. + + :meth:`DeviceMixin._apply` receives only an opaque tensor-to-tensor ``fn``, + so an object that owns no tensors cannot work out where it was just asked + to go. Recording the parsed request on the traversal thread-local lets + :func:`_refresh_device_trackers` answer that question at *every* node -- + including tensor-free children, which are reached through + ``child._apply(fn)`` and therefore never pass through their own ``to()``. + + ``probe`` caches a :func:`_probe_target` result for the traversal so a + graph with many tensor-free nodes probes at most once. + """ + + __slots__ = ("device", "dtype", "probe") + + def __init__(self, device=None, dtype=None): + self.device = canonical_device(device) + self.dtype = dtype + self.probe = None + + def _current_visited(): """Return the active visited set, or ``None`` if not in a traversal.""" return getattr(_traversal_state, "visited", None) +def _current_request(): + """Return the active :class:`_ToRequest`, or ``None``.""" + return getattr(_traversal_state, "request", None) + + +def _push_request(request): + """Install ``request`` for the current traversal, returning the previous.""" + prev = getattr(_traversal_state, "request", None) + _traversal_state.request = request + return prev + + +def _pop_request(prev): + """Restore the request saved by :func:`_push_request`.""" + _traversal_state.request = prev + + def _enter_traversal(): """Begin a top-level traversal. Returns a token for :func:`_exit_traversal`. @@ -85,6 +146,9 @@ def _enter_traversal(): prev = getattr(_traversal_state, "visited", None) if prev is None: _traversal_state.visited = set() + # Never inherit a request leaked by an earlier traversal that unwound + # abnormally: a stale target would silently misdirect this one. + _traversal_state.request = None return "owner" return None @@ -93,6 +157,7 @@ def _exit_traversal(token): """End a top-level traversal.""" if token == "owner": _traversal_state.visited = None + _traversal_state.request = None def _parse_to_args(args, kwargs): @@ -116,6 +181,15 @@ def _parse_to_args(args, kwargs): return device, dtype +def _accepts_single_arg(func) -> bool: + """Whether ``func`` can be called with exactly one positional argument.""" + try: + inspect.signature(func).bind(None) + except (TypeError, ValueError): + return False + return True + + def _apply_to_obj(val, fn, visited): """Apply ``fn`` to any tensors inside ``val``, recursing as needed. @@ -147,12 +221,13 @@ def _apply_to_obj(val, fn, visited): if callable(apply_method) and not isinstance(val, type): if id(val) in visited: return val - try: + # Check the signature up front instead of catching ``TypeError`` from + # the call: a catch there cannot distinguish "wrong signature" from a + # real ``TypeError`` raised mid-traversal, and would silently leave the + # object partially moved. + if _accepts_single_arg(apply_method): apply_method(fn) return val - except TypeError: - # Object has an _apply with a different signature; skip silently. - pass if isinstance(val, list): new_list = [_apply_to_obj(v, fn, visited) for v in val] @@ -180,68 +255,259 @@ def _apply_to_obj(val, fn, visited): def _invalidate_caches(obj): - """Call ``reset_forward_cache`` / ``reset_cache`` if present.""" - reset_fwd = getattr(obj, "reset_forward_cache", None) - if callable(reset_fwd): - try: - reset_fwd() - except Exception: - pass - reset_full = getattr(obj, "reset_cache", None) - if callable(reset_full): + """Call ``reset_forward_cache`` / ``reset_cache`` if present. + + Failures propagate rather than being swallowed. A cache that fails to clear + keeps tensors from the *previous* device, which produces either a + cross-device error much later or -- worse -- silently stale numbers; that + is precisely the corruption this mixin exists to prevent, so it must not be + hidden behind a bare ``except``. + + Raises + ------ + RuntimeError + Chained from the hook's own exception, naming the object and hook. + Note that movement is **not atomic**: by the time a cache hook runs, + some of the object's tensors have already been transformed, so the + object may be left partially moved. + """ + for hook_name in ("reset_forward_cache", "reset_cache"): + hook = getattr(obj, hook_name, None) + if not callable(hook): + continue try: - reset_full() - except Exception: - pass - - -def _representative_tensor(obj): - """Return one tensor that reflects the current device/dtype of *obj*. - - Prefers buffers (their dtype tracks ``dtype_float`` for crystallographic - code) and falls back to parameters, then to a plain ``torch.Tensor`` - attribute found in ``__dict__``. + hook() + except Exception as exc: + raise RuntimeError( + f"{type(obj).__name__}.{hook_name}() failed during device/dtype " + f"movement: {type(exc).__name__}: {exc}. The object may be " + "partially moved and its caches may still hold tensors on the " + "previous device." + ) from exc + + +def _owned_tensors(obj): + """Yield the tensors *obj* itself owns (never a child's). + + ``recurse=False`` is deliberate: the trackers describe where this object + allocates, so inheriting a device from a submodule would report a + half-moved graph as consistent. """ if isinstance(obj, nn.Module): - for buf in obj.buffers(): - return buf - for param in obj.parameters(): - return param + # Read the registration dicts rather than ``buffers()``/``parameters()``: + # several classes here override those accessors with a no-argument + # signature (``SolventModel.parameters``, ``MixedTensor.parameters``), + # and the dicts are in any case the literal "owned, not inherited" set. + for buf in obj._buffers.values(): + if buf is not None: + yield buf + for param in obj._parameters.values(): + if param is not None: + yield param for val in obj.__dict__.values(): if isinstance(val, torch.Tensor): - return val - if hasattr(val, "_data") and isinstance(getattr(val, "_data"), torch.Tensor): - return val._data + yield val + else: + data = getattr(val, "_data", None) + if isinstance(data, torch.Tensor): + yield data + # Container subclasses (``TensorMasks`` is a ``dict``) keep their tensors in + # the container's own storage, not in ``__dict__``. + if isinstance(obj, dict): + for val in obj.values(): + if isinstance(val, torch.Tensor): + yield val + elif isinstance(obj, (list, tuple)): + for val in obj: + if isinstance(val, torch.Tensor): + yield val + + +def _representative_tensor(obj): + """Return one tensor reflecting *obj*'s own device, or ``None``.""" + for tensor in _owned_tensors(obj): + return tensor return None -def _refresh_device_trackers(obj): +def _observed_state(obj): + """Return ``(device, floating_dtype)`` observed from *obj*'s own tensors. + + The two axes are resolved **independently**: the first owned tensor fixes + the device, but only a floating/complex tensor may fix ``dtype_float``. + Deciding both from a single representative tensor lets an integer buffer + that happens to be registered first (``hkl``, ``aniso_flag``) veto the + dtype answer. + """ + device = None + dtype = None + for tensor in _owned_tensors(obj): + if device is None: + device = tensor.device + if dtype is None and (tensor.is_floating_point() or tensor.is_complex()): + dtype = tensor.dtype + if device is not None and dtype is not None: + break + return device, dtype + + +def _probe_target(fn): + """Discover what ``fn`` does to a tensor's device and floating dtype. + + Needed only on the paths that bypass :meth:`DeviceMixin.to` and therefore + record no request: ``.float()`` / ``.double()`` / ``.half()``, and an + ``_apply`` driven by a plain (non-mixin) ``nn.Module`` parent. + + ``fn`` is probed on scratch tensors whose devices are **known to be + valid**, never on the object's own tracker -- a tensor-free object can be + carrying a stale tracker naming a backend this host does not have (say + ``cuda:0`` on a CPU-only machine), and allocating there would fail exactly + when the answer matters most. + + **Both** axes need two reference points, and the scratches must differ on + both. A transforming ``fn`` funnels its two inputs to a single value on the + axis it touches, while a preserving ``fn`` hands each input's own value + back: + + ================== =================== =================== + ``fn`` devices agree? dtypes agree? + ================== =================== =================== + ``.to(device=D)`` yes -> device is D no -> dtype is None + ``.half()`` no -> None yes -> dtype is f16 + ``.to(D, f32)`` yes -> D yes -> f32 + ================== =================== =================== + + Probing one axis with a single reference is what makes a device-only move + report the scratch's own dtype and clobber ``dtype_float``. + + Scratch devices are always ones **known to be valid**, never the object's + own tracker -- a tensor-free object can carry a stale tracker naming a + backend this host does not have (``cuda:0`` on a CPU-only machine), and + allocating there would fail exactly when the answer matters most. + + Returns + ------- + tuple + ``(device, dtype)``; either is ``None`` when ``fn`` leaves that axis + untouched, or when ``fn`` could not be probed at all. + """ + accel = None + if torch.cuda.is_available(): + accel = torch.device("cuda", torch.cuda.current_device()) + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + accel = torch.device("mps", 0) + + def run(device, dtype): + """Apply ``fn`` to one scratch tensor; ``None`` if that is not possible.""" + try: + scratch = torch.empty(0, device=device, dtype=dtype) + except (RuntimeError, TypeError): + return None # backend advertised but unusable for this dtype + # ``fn`` is arbitrary caller-supplied code, so a broad catch is correct + # here -- unlike the cache hooks above, a non-conforming ``fn`` is an + # expected outcome, not a corruption signal. The caller decides what an + # unprobeable ``fn`` means. + try: + out = fn(scratch) + except Exception: + return None + return out if isinstance(out, torch.Tensor) else None + + # The two axes get their own contrast pair, and each pair varies only along + # the axis it measures. Sharing one pair for both would couple them: an + # accelerator scratch cannot be cast to float64 on MPS, so probing dtype on + # the device pair makes ``.double()`` unprobeable on Apple silicon. + base = run(torch.device("cpu"), torch.float32) + if base is None: + return None, None + + device = base.device + if accel is not None: + # Contrast pair for the device axis: same dtype, different device. + other = run(accel, torch.float32) + if other is None: + device = None + elif other.device != base.device: + device = None # each kept its own device -> device-preserving + + dtype = None + if base.is_floating_point() or base.is_complex(): + # Contrast pair for the dtype axis: same device, different dtype. + # float16 rather than float64 so this stays cheap and universally + # supported; the CPU pin means ``.double()`` remains probeable. + other = run(torch.device("cpu"), torch.float16) + if other is not None and other.dtype == base.dtype: + dtype = base.dtype + + return device, dtype + + +_DEVICE_TRACKERS = ("device", "_device") +_DTYPE_TRACKERS = ("dtype_float", "_dtype") + + +def _refresh_device_trackers(obj, fn=None): """Refresh ``device`` / ``_device`` / ``dtype_float`` / ``_dtype`` trackers. - Only updates attributes whose current value is already a ``torch.device`` - / ``torch.dtype`` (or ``None`` / ``str`` / ``int``) — never overwrites - unrelated attributes that happen to share the same name. The tracker - controls where new tensors are subsequently allocated, so it must follow - ``.to()`` moves. + The tracker decides where the object allocates *next* -- over 200 call + sites in the package pass ``device=self.device`` -- so leaving it stale on + a tensor-free object silently misplaces every tensor that object later + creates. + + Resolution order, applied per axis: + + 1. a tensor the object itself owns (authoritative, already moved), + 2. the ``.to()`` request recorded on the traversal thread-local, which is + the only source available to an object holding no tensors, + 3. a probe of ``fn`` (see :func:`_probe_target`), for the ``.float()`` and + plain-parent paths that record no request. + + Only attributes already present in ``obj.__dict__`` and already holding a + device/dtype-shaped value are written, so an unrelated attribute that + happens to be named ``device`` is never clobbered. Written values are + canonical, so ``obj.device == some_tensor.device`` holds. """ - rep = _representative_tensor(obj) - if rep is None: + has_device_tracker = any(a in obj.__dict__ for a in _DEVICE_TRACKERS) + has_dtype_tracker = any(a in obj.__dict__ for a in _DTYPE_TRACKERS) + if not (has_device_tracker or has_dtype_tracker): return - for attr_name in ("device", "_device"): - if attr_name not in obj.__dict__: - continue - current = obj.__dict__[attr_name] - if current is None or isinstance(current, (torch.device, str, int)): - obj.__dict__[attr_name] = rep.device + device, dtype = _observed_state(obj) + + if device is None or dtype is None: + request = _current_request() + if request is not None: + if device is None: + device = request.device + if dtype is None: + dtype = request.dtype + if (device is None or dtype is None) and fn is not None: + if request.probe is None: + request.probe = _probe_target(fn) + probed_device, probed_dtype = request.probe + device = device if device is not None else probed_device + dtype = dtype if dtype is not None else probed_dtype + elif fn is not None: + probed_device, probed_dtype = _probe_target(fn) + device = device if device is not None else probed_device + dtype = dtype if dtype is not None else probed_dtype + + if device is not None: + device = canonical_device(device) + for attr_name in _DEVICE_TRACKERS: + if attr_name not in obj.__dict__: + continue + current = obj.__dict__[attr_name] + if current is None or isinstance(current, (torch.device, str, int)): + obj.__dict__[attr_name] = device - if rep.is_floating_point() or rep.is_complex(): - for attr_name in ("dtype_float", "_dtype"): + if dtype is not None: + for attr_name in _DTYPE_TRACKERS: if attr_name not in obj.__dict__: continue current = obj.__dict__[attr_name] if current is None or isinstance(current, torch.dtype): - obj.__dict__[attr_name] = rep.dtype + obj.__dict__[attr_name] = dtype def _safe_setattr(obj, name, value): @@ -293,19 +559,39 @@ def to(self, *args, **kwargs): # type: ignore[override] pair is parsed and applied via :meth:`_apply`. A call that resolves to neither a device nor a dtype is a no-op that returns ``self``. """ + if _PARSE_TO is not None: + # Let PyTorch's own parser raise on invalid arguments. + device, dtype = _PARSE_TO(*args, **kwargs)[:2] + else: # pragma: no cover - only on a torch build without the private API + device, dtype = _parse_to_args(args, kwargs) + + # Build the request BEFORE claiming the traversal. ``_ToRequest`` + # canonicalises the device, which raises for a backend this host does + # not have (``.to('cuda')`` on a CUDA-less machine). Constructing it + # after ``_enter_traversal`` but outside the ``try`` leaked the + # thread-local visited set on that raise, and every later ``.to()`` on + # the thread then short-circuited as "already visited" -- moving + # nothing, silently. + request = _ToRequest(device, dtype) + token = _enter_traversal() + prev_request = _push_request(request) try: if isinstance(self, nn.Module): return super().to(*args, **kwargs) - device, dtype = _parse_to_args(args, kwargs) if device is None and dtype is None: return self + # Forward the caller's original arguments rather than the parsed + # pair, so overloads and options the parser folds away -- + # ``.to(other_tensor)``, ``non_blocking=``, ``memory_format=`` -- + # reach the tensors intact. def fn(t): - return t.to(device=device, dtype=dtype) + return t.to(*args, **kwargs) return self._apply(fn) finally: + _pop_request(prev_request) _exit_traversal(token) def cuda(self, device=None): # type: ignore[override] @@ -356,11 +642,15 @@ def _apply(self, fn, recurse=True): # type: ignore[override] return self visited.add(id(self)) + # Snapshot before anything moves, so step 5 can tell a real + # transformation from a ``.to()`` that targets the current state. + old_device, old_dtype = _observed_state(self) + # 1. Standard nn.Module traversal (params, buffers, child modules). if isinstance(self, nn.Module): - try: + if _MODULE_APPLY_TAKES_RECURSE: super()._apply(fn, recurse=recurse) - except TypeError: + else: # pragma: no cover - older torch without the kwarg super()._apply(fn) # Mark registered children as visited so that other plain @@ -379,15 +669,72 @@ def _apply(self, fn, recurse=True): # type: ignore[override] # 3. Refresh ``device`` / ``_device`` / ``dtype`` trackers so # subsequent tensor allocations target the new device. - _refresh_device_trackers(self) + _refresh_device_trackers(self, fn) # 4. Invalidate caches. _invalidate_caches(self) + + # 5. Notify the object iff something actually changed, so index + # rebuilds do not fire on a no-op ``.to(current_device)``. + new_device, new_dtype = _observed_state(self) + device_changed = ( + old_device is not None + and new_device is not None + and canonical_device(old_device) != canonical_device(new_device) + ) + dtype_changed = ( + old_dtype is not None + and new_dtype is not None + and old_dtype != new_dtype + ) + if device_changed or dtype_changed: + self._after_device_apply( + old_device, + new_device, + old_dtype, + new_dtype, + device_changed=device_changed, + dtype_changed=dtype_changed, + ) return self finally: if token is not None: _exit_traversal(token) + def _after_device_apply( + self, + old_device, + new_device, + old_dtype, + new_dtype, + *, + device_changed, + dtype_changed, + ): + """Hook called once per object after a *real* device/dtype change. + + No-op by default. Override to rebuild state derived from tensor + placement -- precomputed index tensors, device-specific caches -- that + the generic walk cannot repair on its own. + + Use this rather than ``reset_cache()``: ``reset_cache`` is a + functional-cache hook that :meth:`LossState.reset_caches` fires after + *every* optimizer step, so rebuilding indices there would put index + reconstruction and GPU syncs on the hot path. This hook only fires when + movement actually happened. + + Parameters + ---------- + old_device, new_device : torch.device or None + Device before and after. ``None`` when the object owned no tensors + on that side of the move. + old_dtype, new_dtype : torch.dtype or None + Floating/complex dtype before and after, resolved independently of + the device. + device_changed, dtype_changed : bool + Which axes actually changed. At least one is always True. + """ + # --------------------------------------------------------------------------- # Backwards-compatibility aliases diff --git a/torchref/utils/device_resolution.py b/torchref/utils/device_resolution.py index f9680023..8029439d 100644 --- a/torchref/utils/device_resolution.py +++ b/torchref/utils/device_resolution.py @@ -18,19 +18,10 @@ import torch +from torchref.config import canonical_device as _canonical from torchref.config import get_default_device -def _canonical(device: torch.device) -> torch.device: - """Return ``device`` with its default index filled in. - - ``torch.device('cuda') != torch.device('cuda:0')`` even though both - point to the same physical device. Materialising an empty tensor - on the device is the cheapest way to get the canonical form. - """ - return torch.empty(0, device=device).device - - def resolve_device( *modules: Any, device: Optional[Union[torch.device, str]] = None, @@ -71,6 +62,17 @@ def resolve_device( torch.device The resolved device. + Raises + ------ + TypeError + If a bare ``torch.Tensor`` (or ``nn.Parameter``) is passed. Such an + object satisfies the ``.device`` / ``.to()`` precondition + *syntactically* but violates it semantically, because + ``Tensor.to()`` returns a new tensor rather than moving in place -- + so the move would be dropped without a word. Use + :func:`torchref.config.normalize_device` to read a device off a + tensor; use this function only to reconcile owning objects. + Examples -------- Empty call returns the configured default:: @@ -89,10 +91,29 @@ def resolve_device( >>> resolve_device(cuda_model, cpu_data) # doctest: +SKIP device(type='cuda') """ + for m in modules: + if isinstance(m, torch.Tensor): + raise TypeError( + "resolve_device() moves its inputs in place and returns only " + f"the resolved device, but a bare {type(m).__name__} was " + "passed. torch.Tensor.to() is out-of-place, so the move would " + "be silently discarded and the tensor left where it was. Pass " + "the object that owns the tensor, or move it yourself: " + "t = t.to(normalize_device(...))." + ) + if device is not None: - resolved = torch.device(device) if not isinstance(device, torch.device) else device + resolved = _canonical(device) for m in modules: - if m is not None: + # Skip modules already on target. ``.to()`` invalidates caches and + # fires ``_after_device_apply``, so a no-op move is not free -- + # ``SfDS.forward`` calls ``resolve_device`` on every evaluation. + # + # This is only safe because ``DeviceMixin`` now keeps ``m.device`` + # truthful; before that, the unconditional ``.to()`` here was + # accidentally repairing objects whose tracker lied about where + # their tensors were. + if m is not None and _canonical(m.device) != resolved: m.to(resolved) return resolved diff --git a/torchref/utils/utils.py b/torchref/utils/utils.py index 3e5d0366..06a4f85f 100644 --- a/torchref/utils/utils.py +++ b/torchref/utils/utils.py @@ -257,11 +257,12 @@ class TensorMasks(DeviceMovementMixin, dict): def __init__(self, data=None, device=None): super().__init__() - if device is None: - from torchref.config import get_default_device + from torchref.config import normalize_device - device = get_default_device() - self.device = torch.device(device) + # ``normalize_device`` rather than ``torch.device(...)``: the latter + # keeps an un-indexed spelling ("mps"), which compares unequal to the + # indexed device every real tensor reports. + self.device = normalize_device(device) self._cache = None self._updated = True @@ -314,12 +315,14 @@ def _apply(self, fn): self._cache = None self._updated = True - # Refresh the ``device`` tracker so future ``__setitem__`` calls - # (which migrate incoming tensors to ``self.device``) land correctly. - for v in self.values(): - if isinstance(v, torch.Tensor): - self.device = v.device - break + # Refresh the ``device`` tracker (future ``__setitem__`` calls migrate + # incoming tensors to it) through the shared helper rather than a local + # copy: it also handles the empty-``TensorMasks`` case, where there is + # no mask tensor to read a device from and the tracker has to come from + # the recorded ``.to()`` request. + from torchref.utils.device_mixin import _refresh_device_trackers + + _refresh_device_trackers(self, fn) return self def reset_cache(self) -> None: