diffct 2.2.1__tar.gz → 2.3.0__tar.gz

{diffct-2.2.1 → diffct-2.3.0}/CHANGELOG.md

@@ -5,7 +5,45 @@ All notable changes to this project are documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased]
+## [2.3.0] - 2026-10-14
+
+### Added
+
+- `backend="auto"`, `"numba"` or `"torch"` and `precision="fast"`, `"float32"` or `"float64"` on `Projector`, `fbp` and `fdk`. These settings are opt-in. Default calls keep native FP32 execution where compatible.
+- `backend="auto"` selects the PyTorch backend for `precision="float32"`, `precision="float64"`, per-axis `voxel_spacing` and hybrid partitions. An incompatible explicit `backend="numba"` request raises an error.
+- `backend="torch"` runs on CPU, NVIDIA CUDA and PyTorch ROCm. `precision="float64"` runs the data operations in FP64.
+- Portable FP32 ray setup clips rays and locates tile entries in FP64. Traversal measures distances from each entry, so distant endpoints keep voxel widths.
+- Per-axis `voxel_spacing`, such as `(dy, dx)` or `(dz, dy, dx)`, in tensor axis order.
+- Leading batch and channel dimensions for projection and adjoint, including empty batches.
+- Data gradients, JVP, `vmap`, `jacrev`, `jacfwd` and second data derivatives for the PyTorch projector. Dense Jacobians need a single-process projector.
+- `view_chunk_size` and `volume_chunk_shape` on `fbp` and `fdk` stream view batches and spatial tiles. `devices` schedules tiles across local devices.
+- `fbp_into` and `fdk_into` write into `TensorStore`, `NpyStore` or compatible block stores.
+- `partition="hybrid"` with `partition_grid=(view_ranks, space_ranks)` for projectors and analytical reconstruction. Uneven and empty shards are supported.
+- `cgls`, `sirt`, `tv_reconstruction` and `SolverResult`, with residual thresholds, residual history, callbacks and optional `x0`.
+- `cgls_into` keeps solver state in mapped files. It rejects non-finite measurements in bounded blocks before any output write.
+- `checkpoint_directory` and `resume` for `cgls`, `sirt` and `tv_reconstruction`. Fingerprints bind the acquisition, measurements, precision, ownership and settings.
+- `total_variation` with `reduction="mean"` or `"sum"`. Spatial and hybrid runs exchange neighbouring faces.
+- `docs/EXECUTION.md` documents precision, memory bounds, distributed loss replica divisors and checkpoint contracts.
+
+### Changed
+
+- Importing `diffct` and using `backend="torch"` no longer loads Numba CUDA or requires a CUDA runtime. Native execution initializes them on first use.
+- `pip install diffct` is enough for `backend="torch"`. Native Numba execution needs the CUDA extras `diffct[cu12]` or `diffct[cu13]`.
+- Distributed PyTorch-backend projectors require the same `volume_chunk_shape`, `view_chunk_size` and `detector_chunk_shape` on every rank.
+- The source distribution includes the four example helpers used by its shipped tests.
+
+### Fixed
+
+- Parker short-scan distance and detector-u offset checks bound position and
+  detector-axis rounding per component along the checked direction, not by the
+  absolute coordinate size.
+  A bfloat16 or float16 detector-v axis made a 50-unit change in detector distance
+  pass, and the Parker weights then used the first view's distance.
+- Fan and cone normalization in the analytical gather and its transpose uses each view's SID and SDD. Inactive views no longer rescale active views.
+- `angular_integration_weights` and `parker_weights` keep FP32 and FP64 input precision. FP16 and BF16 angles are computed in FP32.
+- Angular gaps come directly from sorted neighbouring angles. Float32 neighbours one ulp apart no longer get a zero weight, and Parker rows no longer depend on the input order.
+- Valid low-precision Parker short scans keep finite weights and angle gradients.
+- Output stores that overlap internally raise before any I/O. Disjoint interleaved `TensorStore` slices are accepted.
 
 ## [2.2.1] - 2026-10-13
 

{diffct-2.2.1 → diffct-2.3.0}/CITATION.cff

@@ -13,8 +13,8 @@ authors:
     given-names: Yipeng
     email: yipeng.sun@fau.de
     alias: sypsyp97
-version: 2.2.1
-date-released: "2026-10-13"
+version: 2.3.0
+date-released: "2026-10-14"
 license: Apache-2.0
 repository-code: "https://github.com/sypsyp97/diffct"
 url: "https://sypsyp97.github.io/diffct/"

{diffct-2.2.1 → diffct-2.3.0}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.5
 Name: diffct
-Version: 2.2.1
-Summary: Differentiable CUDA projectors for CT: arbitrary trajectories, multi-GPU and multi-node execution, geometry gradients
+Version: 2.3.0
+Summary: Differentiable CT projection and reconstruction: precision modes, streaming, distributed execution and geometry gradients
 Project-URL: Homepage, https://github.com/sypsyp97/diffct
 Project-URL: Documentation, https://sypsyp97.github.io/diffct/
 Project-URL: Changelog, https://github.com/sypsyp97/diffct/blob/main/CHANGELOG.md
@@ -10,7 +10,7 @@ Author-email: Yipeng Sun <yipeng.sun@fau.de>
 License-Expression: Apache-2.0
 License-File: LICENSE
 License-File: NOTICE
-Keywords: CT reconstruction,CUDA,PyTorch,computed tomography,differentiable programming
+Keywords: CT reconstruction,CUDA,PyTorch,ROCm,computed tomography,differentiable programming
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: GPU :: NVIDIA CUDA
 Classifier: Intended Audience :: Science/Research
@@ -47,11 +47,11 @@ Description-Content-Type: text/markdown
 </p>
 
 <p align="center">
-  <img src="https://raw.githubusercontent.com/sypsyp97/diffct/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
+  <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
 </p>
 
 <p align="center">
-  <a href="https://github.com/sypsyp97/diffct/blob/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
+  <a href="https://github.com/sypsyp97/diffct/blob/main/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
   <a href="https://www.preprints.org/manuscript/202605.1446/v1">Technical report</a> ·
   <a href="https://doi.org/10.20944/preprints202605.1446.v1">DOI</a>
 </p>
@@ -76,21 +76,23 @@ To our knowledge, diffct is the only open-source GPU CT library that combines ar
 - **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
 - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
 - **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
-- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Default `partition="views"` shards projections and replicates volumes. `partition="space"` gives each rank an owned volume slab and combines ray contributions. Speedup depends on workload and communication costs.
+- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Views, space and hybrid partitions control replication of volumes and projections. Speedup depends on workload and communication costs.
+- **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts and validation limits](docs/EXECUTION.md).
+- **Solver API.** `cgls`, `sirt` and `tv_reconstruction` report residual history, tolerances and callbacks, with scan/data fingerprints and complete checkpoint state. `cgls_into` keeps state in bounded mapped blocks.
 - **Differentiable FBP and FDK.** `fbp`, `fdk` and approximate `helical_fdk` compose the public filters, weights and analytical backprojectors. Their backward kernels transpose the actual interpolation stencil; custom filters and weights can be learned. See the [analytical guide](docs/source/analytical_learning.rst).
 
 Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
 
 ## Install
 
-You need an NVIDIA GPU. Install PyTorch for your CUDA version first. Then install diffct with the extra for the same CUDA major version:
+For the native path, install NVIDIA CUDA PyTorch first, then diffct with the matching CUDA major extra:
 
 ```bash
 pip install "diffct[cu12]"    # PyTorch built for CUDA 12
 pip install "diffct[cu13]"    # PyTorch built for CUDA 13
 ```
 
-The extra installs the CUDA compiler libraries that Numba CUDA needs to compile the kernels. A plain `pip install diffct` does not install them. Use it only if a CUDA Toolkit is already installed on the system.
+The extra installs the CUDA compiler libraries Numba needs. For CPU or PyTorch ROCm execution, install the appropriate PyTorch build and plain `pip install diffct`, then select `backend="torch"`. Importing diffct and using this backend require no CUDA runtime; Numba CUDA initializes when native execution is used. Plain installation also works for native execution when a compatible CUDA Toolkit is already present.
 
 From source, with the examples:
 
@@ -226,6 +228,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
 
 | Document | Contents |
 |---|---|
+| [Precision, training and reconstruction](docs/EXECUTION.md) | Backend modes, batch/spacing contracts, analytical storage, hybrid loss scaling and solver checkpoints |
 | [Trajectory guide](https://github.com/sypsyp97/diffct/blob/main/docs/source/trajectories.rst) | Trajectory tuples and `Projector` usage |
 | [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md) | Capabilities and limits, custom trajectories, geometry gradients, multi-GPU details, geometry and units |
 | [docs/DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md) | Execution/memory choices, distributed loss rules, Slurm and cross-node checks |

{diffct-2.2.1 → diffct-2.3.0}/README.md

@@ -16,11 +16,11 @@
 </p>
 
 <p align="center">
-  <img src="https://raw.githubusercontent.com/sypsyp97/diffct/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
+  <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
 </p>
 
 <p align="center">
-  <a href="https://github.com/sypsyp97/diffct/blob/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
+  <a href="https://github.com/sypsyp97/diffct/blob/main/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
   <a href="https://www.preprints.org/manuscript/202605.1446/v1">Technical report</a> ·
   <a href="https://doi.org/10.20944/preprints202605.1446.v1">DOI</a>
 </p>
@@ -45,21 +45,23 @@ To our knowledge, diffct is the only open-source GPU CT library that combines ar
 - **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
 - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
 - **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
-- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Default `partition="views"` shards projections and replicates volumes. `partition="space"` gives each rank an owned volume slab and combines ray contributions. Speedup depends on workload and communication costs.
+- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Views, space and hybrid partitions control replication of volumes and projections. Speedup depends on workload and communication costs.
+- **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts and validation limits](docs/EXECUTION.md).
+- **Solver API.** `cgls`, `sirt` and `tv_reconstruction` report residual history, tolerances and callbacks, with scan/data fingerprints and complete checkpoint state. `cgls_into` keeps state in bounded mapped blocks.
 - **Differentiable FBP and FDK.** `fbp`, `fdk` and approximate `helical_fdk` compose the public filters, weights and analytical backprojectors. Their backward kernels transpose the actual interpolation stencil; custom filters and weights can be learned. See the [analytical guide](docs/source/analytical_learning.rst).
 
 Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
 
 ## Install
 
-You need an NVIDIA GPU. Install PyTorch for your CUDA version first. Then install diffct with the extra for the same CUDA major version:
+For the native path, install NVIDIA CUDA PyTorch first, then diffct with the matching CUDA major extra:
 
 ```bash
 pip install "diffct[cu12]"    # PyTorch built for CUDA 12
 pip install "diffct[cu13]"    # PyTorch built for CUDA 13
 ```
 
-The extra installs the CUDA compiler libraries that Numba CUDA needs to compile the kernels. A plain `pip install diffct` does not install them. Use it only if a CUDA Toolkit is already installed on the system.
+The extra installs the CUDA compiler libraries Numba needs. For CPU or PyTorch ROCm execution, install the appropriate PyTorch build and plain `pip install diffct`, then select `backend="torch"`. Importing diffct and using this backend require no CUDA runtime; Numba CUDA initializes when native execution is used. Plain installation also works for native execution when a compatible CUDA Toolkit is already present.
 
 From source, with the examples:
 
@@ -195,6 +197,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
 
 | Document | Contents |
 |---|---|
+| [Precision, training and reconstruction](docs/EXECUTION.md) | Backend modes, batch/spacing contracts, analytical storage, hybrid loss scaling and solver checkpoints |
 | [Trajectory guide](https://github.com/sypsyp97/diffct/blob/main/docs/source/trajectories.rst) | Trajectory tuples and `Projector` usage |
 | [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md) | Capabilities and limits, custom trajectories, geometry gradients, multi-GPU details, geometry and units |
 | [docs/DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md) | Execution/memory choices, distributed loss rules, Slurm and cross-node checks |

{diffct-2.2.1 → diffct-2.3.0}/diffct/__init__.py

@@ -20,11 +20,15 @@ from .surfaces import ParameterizedSurface
 from .reconstruction import (
     fbp,
     fdk,
+    fbp_into,
+    fdk_into,
     helical_fdk,
     helical_view_weights,
     helical_weighted_backproject,
 )
 
+from .solvers import SolverResult, cgls, sirt, tv_reconstruction, cgls_into, total_variation
+
 from .geometry import (
     circular_trajectory_3d,
     random_trajectory_3d,
@@ -58,7 +62,7 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.2.1'
+__version__ = '2.3.0'
 
 __all__ = [
     'Projector',
@@ -67,6 +71,14 @@ __all__ = [
     'ParameterizedSurface',
     'fbp',
     'fdk',
+    'fbp_into',
+    'fdk_into',
+    'SolverResult',
+    'cgls',
+    'sirt',
+    'tv_reconstruction',
+    'cgls_into',
+    'total_variation',
     'helical_fdk',
     'helical_view_weights',
     'helical_weighted_backproject',

diffct-2.3.0/diffct/_solver_state.py

@@ -0,0 +1,284 @@
+"""Bounded solver state, owned reductions and complete checkpoint publication."""
+
+import hashlib
+import json
+import math
+import numbers
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+
+from .pipeline import _fragments
+from .storage import NpyStore, TensorStore
+
+
+def count(value, name, positive=False):
+    if isinstance(value, bool) or not isinstance(value, numbers.Integral) or value < int(positive):
+        raise ValueError(f"{name} must be a {'positive' if positive else 'nonnegative'} integer")
+    return int(value)
+
+
+def finite(value, name, positive=False):
+    if isinstance(value, bool) or not isinstance(value, numbers.Real) or not math.isfinite(value) or value < 0 or (positive and value == 0):
+        raise ValueError(f"{name} must be finite and {'positive' if positive else 'nonnegative'}")
+    return float(value)
+
+
+def agreement(operator, error=None, token=None):
+    """Every rank rejects a bad request before numerical execution starts."""
+    world = getattr(operator, "world_size", 1)
+    if world == 1:
+        if error is not None:
+            raise error
+        return
+    records = [None] * world
+    dist.all_gather_object(records, (None if error is None else str(error), token),
+                           group=operator.process_group)
+    failures = [record[0] for record in records if record[0] is not None]
+    if failures or any(record[1] != records[0][1] for record in records):
+        raise ValueError("distributed solver agreement failed: " + (failures[0] if failures else "inconsistent checkpoint iteration or controls"))
+
+
+def scalar_sum(operator, value, kind):
+    """Count physical rays/voxels once, including view/space replicas."""
+    world = getattr(operator, "world_size", 1)
+    partition = getattr(operator, "partition", "views")
+    rank = getattr(operator, "rank", 0)
+    owns = (getattr(operator, "view_rank", 0) == 0 if kind == "volume"
+            else getattr(operator, "space_rank", 0) == 0) if partition == "hybrid" else (
+        rank == 0 if (kind == "volume") == (partition == "views") else True
+    )
+    value = float(value) if owns else 0.
+    if world > 1:
+        group = operator.process_group
+        device = torch.device("cuda", torch.cuda.current_device()) if dist.get_backend(group) == "nccl" else torch.device("cpu")
+        accumulator = torch.tensor(value, dtype=torch.float64, device=device)
+        dist.all_reduce(accumulator, group=group)
+        value = accumulator.item()
+    return value
+
+
+def squared_norm(operator, value, kind, elements=1048576):
+    store = TensorStore(value) if isinstance(value, torch.Tensor) else value
+    total = 0.
+    for selected in _fragments(tuple(store.shape), elements):
+        # Clone even an FP64 view: norms must never alter caller measurements.
+        block = store.read(selected).to(dtype=torch.float64).clone()
+        total += block.square_().sum().item()
+    return scalar_sum(operator, total, kind)
+
+
+def copy_blocks(source, target, elements):
+    for selected in _fragments(tuple(target.shape), elements):
+        target.write(selected, source.read(selected).to(device=getattr(target, "device", "cpu"), dtype=target.dtype))
+
+
+def fill_blocks(target, elements):
+    for selected in _fragments(tuple(target.shape), elements):
+        shape = tuple(part.stop - part.start for part in selected)
+        target.write(selected, torch.zeros(shape, dtype=target.dtype))
+
+
+def update_blocks(target, source, factor, elements, scale=1.):
+    for selected in _fragments(tuple(target.shape), elements):
+        block = target.read(selected).to(device="cpu", dtype=target.dtype).clone()
+        block.mul_(scale).add_(source.read(selected).to(device="cpu", dtype=target.dtype), alpha=factor)
+        target.write(selected, block)
+
+
+def fingerprint(value, elements):
+    store = TensorStore(value.detach()) if isinstance(value, torch.Tensor) else value
+    digest = hashlib.sha256(json.dumps((tuple(store.shape), str(store.dtype))).encode())
+    for selected in _fragments(tuple(store.shape), elements):
+        block = store.read(selected).detach().to(device="cpu").contiguous()
+        # Hash storage bytes without requiring NumPy support for the original
+        # floating dtype (for example bfloat16). This view allocates no payload.
+        digest.update(block.reshape(-1).view(torch.uint8).numpy().tobytes())
+    return digest.hexdigest()
+
+
+def _effective_surface(operator, frames, elements):
+    """Hash the actual physical pixel rays, without materializing a volume."""
+    surface = getattr(operator, "detector_surface", None)
+    if surface is None:
+        return None
+    from .backends import ReferenceProjector
+    from .surfaces import ParameterizedSurface, _detector_grids
+
+    dimensions = len(operator.global_volume_shape)
+    if elements < max(3, dimensions):
+        raise ValueError("block_elements cannot hold one physical detector pixel")
+    if isinstance(surface, ParameterizedSurface):
+        if any(value.device.type != "cpu" and value.numel() > elements for value in surface.parameters):
+            raise ValueError("surface parameters cannot be staged within block_elements")
+        geometry = (*frames, *surface.parameters)
+    else:
+        # Legacy callbacks require complete grids. Their documented output can
+        # contain every view, so admit that snapshot only when it fits the
+        # bound; a smaller arbitrary callback invocation may change its meaning.
+        if operator._n_views * math.prod(operator.detector_shape) * 3 > elements:
+            raise ValueError("legacy surface fingerprint exceeds block_elements; use ParameterizedSurface")
+        pixels = tuple(slice(0, size) for size in operator.detector_shape)
+        u, v = _detector_grids(operator.detector_shape, operator.detector_spacing, pixels)
+        offsets = surface(u, v)
+        if not isinstance(offsets, torch.Tensor) or not torch.is_floating_point(offsets):
+            raise TypeError("detector_surface must return floating-point offsets")
+        if tuple(offsets.shape) not in ((*operator.detector_shape, 3),
+                                       (operator._n_views, *operator.detector_shape, 3)):
+            raise ValueError("legacy detector surface must match the complete pixel grids")
+        geometry = (*frames, offsets)
+    reference = getattr(operator, "_reference", None)
+    if reference is None:
+        reference = ReferenceProjector(operator, validate=False)
+    description = json.dumps((operator._n_views, operator.detector_shape, str(operator.dtype))).encode()
+    origins_digest, directions_digest = hashlib.sha256(description), hashlib.sha256(description)
+    for view in range(operator._n_views):
+        for pixels in _fragments(operator.detector_shape, elements // max(3, dimensions)):
+            origins, directions, _ = reference._rays(slice(view, view + 1), pixels,
+                                                      torch.device("cpu"), validate=True, geometry=geometry)
+            # Separate streams preserve row-major ray ordering across different
+            # block sizes; changing a memory ceiling does not change the hash.
+            origins_digest.update(origins.detach().contiguous().numpy().tobytes())
+            directions_digest.update(directions.detach().contiguous().numpy().tobytes())
+    return [origins_digest.hexdigest(), directions_digest.hexdigest()]
+
+
+def _signature_local(operator, source, algorithm, settings, elements, compute_dtype):
+    frames = getattr(operator, "_learnable", ()) or getattr(operator, "_trajectory", ())
+    geometry = tuple(fingerprint(value, elements) for value in frames)
+    surface = getattr(operator, "detector_surface", None)
+    surface_values = (() if surface is None else getattr(surface, "parameters", ()))
+    if isinstance(surface, torch.Tensor):
+        surface_values = (surface,)
+    return dict(
+        schema_version=1, algorithm=algorithm, settings=settings,
+        rank=getattr(operator, "rank", 0), world_size=getattr(operator, "world_size", 1),
+        partition=getattr(operator, "partition", "views"),
+        partition_grid=getattr(operator, "partition_grid", None),
+        global_volume_shape=list(operator.global_volume_shape), volume_shape=list(operator.volume_shape),
+        volume_slice=[[part.start, part.stop] for part in operator.volume_slice],
+        projection_shape=list(operator.projection_shape),
+        view_slice=[operator.view_slice.start, operator.view_slice.stop],
+        beam=operator.beam, voxel_spacing=operator.voxel_spacing,
+        detector_spacing=list(operator.detector_spacing), detector_shape=list(operator.detector_shape),
+        input_dtype=str(source.dtype), geometry=list(geometry),
+        surface=[fingerprint(value, elements) for value in surface_values],
+        effective_surface=_effective_surface(operator, frames, elements),
+        compute_dtype=str(compute_dtype), precision=getattr(operator, "precision", None),
+        backend=getattr(operator, "backend", None),
+        measurements=fingerprint(source, elements),
+    )
+
+
+def signature(operator, source, algorithm, settings, elements, *, compute_dtype=None):
+    error, metadata = None, None
+    try:
+        dtype = getattr(operator, "dtype", source.dtype) if compute_dtype is None else compute_dtype
+        metadata = _signature_local(operator, source, algorithm, settings, elements, dtype)
+    except Exception as caught:
+        error = caught
+    agreement(operator, error)
+    if getattr(operator, "world_size", 1) > 1:
+        records = [None] * operator.world_size
+        dist.all_gather_object(records, metadata, group=operator.process_group)
+        local_fields = {"rank", "volume_shape", "volume_slice", "projection_shape", "view_slice", "measurements"}
+        scans = [{key: value for key, value in record.items() if key not in local_fields} for record in records]
+        if any(scan != scans[0] for scan in scans):
+            raise ValueError("distributed checkpoint scan and compute configuration must agree")
+        replicas = {}
+        for record in records:
+            views = tuple(record["view_slice"])
+            digest = record["measurements"]
+            if views in replicas and replicas[views] != digest:
+                raise ValueError("distributed checkpoint measurement replicas must agree")
+            replicas[views] = digest
+    return metadata
+
+
+def load_checkpoint(operator, directory, expected, iterations, elements, names):
+    error, saved, states = None, None, None
+    try:
+        directory = Path(directory).resolve()
+        saved = json.loads((directory / "metadata.json").read_text(encoding="utf-8"))
+        # JSON canonicalization also normalizes tuple-valued physical spacing.
+        canonical = json.loads(json.dumps(expected))
+        if any(saved.get(key) != value for key, value in canonical.items()):
+            raise ValueError("checkpoint scan, data, ownership, dtype or algorithm settings do not match")
+        completed = count(saved.get("iteration"), "checkpoint iteration")
+        history = saved.get("residual_history")
+        if completed > iterations or not isinstance(history, list) or len(history) != completed + 1:
+            raise ValueError("checkpoint iteration or residual history is invalid")
+        for value in history:
+            finite(value, "checkpoint residual")
+        finite(saved.get("measurement_norm"), "checkpoint measurement norm")
+        if expected["algorithm"] in ("cgls", "cgls_into"):
+            finite(saved.get("gamma"), "checkpoint gamma")
+        manifest = saved.get("states", {})
+        if set(manifest) != set(names):
+            raise ValueError("checkpoint does not contain the complete solver state")
+        states = {}
+        for name in names:
+            store = NpyStore(directory / f"{name}.npy", mode="r")
+            specification = manifest[name]
+            shape = operator.projection_shape if name in ("r", "q", "row") else operator.volume_shape
+            if name == "step":
+                shape = (1,)
+            if tuple(store.shape) != tuple(shape) or list(store.shape) != specification.get("shape") or str(store.dtype) != specification.get("dtype"):
+                raise ValueError("checkpoint solver state shape or dtype does not match")
+            if fingerprint(store, elements) != specification.get("fingerprint"):
+                raise ValueError("checkpoint solver state fingerprint does not match")
+            states[name] = store
+    except (ValueError, TypeError, OSError, EOFError, KeyError) as caught:
+        error = caught
+    token = None if error is not None else (saved["iteration"], saved["residual_history"], saved.get("gamma"), saved["measurement_norm"])
+    agreement(operator, error, token)
+    return saved, states
+
+
+def save_checkpoint(operator, directory, states, metadata, iteration, history, measurement_norm, elements, **scalars):
+    error, manifest, rank = None, {}, None
+    try:
+        step = Path(directory).resolve() / f"step-{iteration:06d}"
+        step.mkdir(parents=True, exist_ok=True)
+        rank = step / f"rank-{getattr(operator, 'rank', 0):05d}"
+        rank.mkdir()
+        for name, source in states.items():
+            source = TensorStore(source.detach()) if isinstance(source, torch.Tensor) else source
+            target = NpyStore.create(rank / f"{name}.npy", source.shape, dtype=source.dtype)
+            copy_blocks(source, target, elements)
+            target.flush()
+            manifest[name] = dict(shape=list(source.shape), dtype=str(source.dtype), fingerprint=fingerprint(target, elements))
+    except (OSError, ValueError, TypeError) as caught:
+        error = caught
+    # All owned states on every rank are written and flushed before any rank
+    # atomically publishes the metadata marking a completed iteration.
+    agreement(operator, error, (iteration, history))
+    error = None
+    try:
+        document = dict(metadata, iteration=iteration, residual_history=history,
+                        measurement_norm=measurement_norm, states=manifest, **scalars)
+        temporary = rank / "metadata.tmp"
+        with temporary.open("x", encoding="utf-8") as file:
+            json.dump(document, file)
+            file.flush()
+        temporary.replace(rank / "metadata.json")
+    except Exception as caught:
+        error = caught
+    # A failed metadata write/rename must stop every rank before another
+    # projection or callback can enter a numerical collective.
+    agreement(operator, error, iteration)
+
+
+def callback_stop(operator, callback, iteration, residual, volume):
+    error, stop = None, False
+    try:
+        stop = callback is not None and callback(iteration, residual, volume) is False
+    except Exception as caught:
+        error = caught
+    agreement(operator, error)
+    if getattr(operator, "world_size", 1) > 1:
+        records = [None] * operator.world_size
+        dist.all_gather_object(records, stop, group=operator.process_group)
+        stop = any(records)
+    return stop

diffct-2.3.0/diffct/_solver_tv.py

@@ -0,0 +1,110 @@
+"""Physical forward-difference TV with differentiable slab boundary faces."""
+
+import math
+import numbers
+
+import torch
+import torch.distributed as dist
+
+from ._solver_state import finite, scalar_sum
+from .backends import _zero_dependency
+
+
+def _peers(operator, nonempty):
+    spaces = operator.space_world_size
+    current = operator.space_rank
+    occupied = [rank for rank in range(spaces) if rank < min(spaces, operator.global_volume_shape[0])]
+    previous = max((rank for rank in occupied if rank < current), default=None) if nonempty else None
+    following = min((rank for rank in occupied if rank > current), default=None) if nonempty else None
+    def global_rank(rank):
+        if rank is None:
+            return None
+        parent = operator.view_rank * spaces + rank
+        return operator._partition.global_rank(parent)
+    return global_rank(previous), global_rank(following)
+
+
+class _FollowingFace(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, volume, operator):
+        previous, following = _peers(operator, volume.shape[0] > 0)
+        ctx.operator, ctx.previous, ctx.following = operator, previous, following
+        ctx.shape, ctx.dtype, ctx.device = volume.shape, volume.dtype, volume.device
+        nccl = dist.get_backend(operator.space_group) == "nccl"
+        ctx.communication_device = (torch.device("cuda", torch.cuda.current_device())
+                                    if nccl else torch.device("cpu"))
+        if nccl:
+            # Initialize the subgroup communicator collectively, including
+            # empty slabs, before the first subset-only face exchange.
+            dist.barrier(group=operator.space_group)
+        face = torch.zeros(volume.shape[1:], dtype=volume.dtype, device=ctx.communication_device)
+        operations = []
+        if previous is not None:
+            first = volume[0].detach().to(device=ctx.communication_device).contiguous()
+            operations.append(dist.P2POp(dist.isend, first, previous, operator.space_group))
+        if following is not None:
+            operations.append(dist.P2POp(dist.irecv, face, following, operator.space_group))
+        if operations:
+            for work in dist.batch_isend_irecv(operations):
+                work.wait()
+        return face.to(device=volume.device)
+
+    @staticmethod
+    def backward(ctx, grad_face):
+        received = torch.zeros(ctx.shape[1:], dtype=ctx.dtype, device=ctx.communication_device)
+        operations = []
+        if ctx.following is not None:
+            outgoing = grad_face.to(device=ctx.communication_device).contiguous()
+            operations.append(dist.P2POp(dist.isend, outgoing, ctx.following, ctx.operator.space_group))
+        if ctx.previous is not None:
+            operations.append(dist.P2POp(dist.irecv, received, ctx.previous, ctx.operator.space_group))
+        if operations:
+            for work in dist.batch_isend_irecv(operations):
+                work.wait()
+        gradient = torch.zeros(ctx.shape, dtype=ctx.dtype, device=ctx.device)
+        if ctx.shape[0] > 0:
+            gradient[0] = received.to(device=ctx.device)
+        return gradient, None
+
+
+def total_variation(volume, *, operator=None, voxel_spacing=1., eps=1e-6, reduction="mean"):
+    """Isotropic TV over every voxel, with physical forward differences.
+
+    Differences vanish at the global outer boundary. Spatially partitioned
+    volumes exchange a single neighboring face in forward and backward; the
+    reported scalar counts replicated ownership once and gradients stay local.
+    """
+    if not isinstance(volume, torch.Tensor) or not torch.is_floating_point(volume) or volume.ndim < 1:
+        raise TypeError("volume must be a real floating-point tensor")
+    eps = finite(eps, "eps", positive=True)
+    if reduction not in ("mean", "sum"):
+        raise ValueError("reduction must be mean or sum")
+    spacing = (voxel_spacing,) * volume.ndim if isinstance(voxel_spacing, numbers.Real) else tuple(voxel_spacing)
+    if len(spacing) != volume.ndim:
+        raise ValueError("voxel_spacing must match the volume axes")
+    spacing = tuple(finite(pitch, "voxel_spacing", positive=True) for pitch in spacing)
+    if operator is not None and tuple(volume.shape) != tuple(operator.volume_shape):
+        raise ValueError("volume shape must match its operator ownership")
+    distributed_space = operator is not None and getattr(operator, "world_size", 1) > 1 and operator.partition in ("space", "hybrid")
+    halo = _FollowingFace.apply(volume, operator) if distributed_space else None
+    squared = torch.zeros_like(volume)
+    for axis, pitch in enumerate(spacing):
+        current, following = [slice(None)] * volume.ndim, [slice(None)] * volume.ndim
+        current[axis], following[axis] = slice(0, -1), slice(1, None)
+        difference = torch.zeros_like(volume)
+        difference[tuple(current)] = (volume[tuple(following)] - volume[tuple(current)]) / pitch
+        if axis == 0 and halo is not None and volume.shape[0] and operator.volume_slice[0].stop < operator.global_volume_shape[0]:
+            difference[-1] = (halo - volume[-1]) / pitch
+        squared = squared + difference.square()
+    local = torch.sqrt(squared + eps).sum()
+    if halo is not None:
+        # The final slab also needs backward to receive the previous slab's
+        # contribution, even though its own global outer difference is zero.
+        local = local + _zero_dependency(halo)
+    total = local
+    elements = volume.numel()
+    if operator is not None:
+        global_value = scalar_sum(operator, local.detach().item(), "volume")
+        total = local + (local.new_tensor(global_value) - local.detach())
+        elements = math.prod(operator.global_volume_shape)
+    return total if reduction == "sum" else total / max(elements, 1)