diffct 2.4.2__py3-none-any.whl → 2.4.3__py3-none-any.whl

diffct/__init__.py

@@ -64,7 +64,7 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.4.2'
+__version__ = '2.4.3'
 
 __all__ = [
     'Projector',

diffct/_solver_state.py

@@ -53,30 +53,32 @@ def owns_values(operator, kind):
     )
 
 
-def scalar_sum(operator, value, kind):
-    """Count physical rays/voxels once, including view/space replicas."""
+def scalar_sum(operator, value, kind, *, return_tensor=False):
+    """Count physical rays/voxels once, optionally retaining the reduction tensor."""
     world = getattr(operator, "world_size", 1)
+    result_device = value.device if isinstance(value, torch.Tensor) else torch.device("cpu")
     value = value if owns_values(operator, kind) 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.as_tensor(value, dtype=torch.float64, device=device).clone()
         dist.all_reduce(accumulator, group=group)
-        return accumulator.item()
+        return accumulator if return_tensor else accumulator.item()
+    if return_tensor:
+        return torch.as_tensor(value, dtype=torch.float64, device=result_device)
     return float(value)
 
 
-def squared_norm(operator, value, kind, elements=1048576):
+def squared_norm(operator, value, kind, elements=1048576, *, return_tensor=False):
     store = TensorStore(value) if isinstance(value, torch.Tensor) else value
-    total = 0.
+    total = torch.zeros((), dtype=torch.float64, device=getattr(store, "device", "cpu"))
     if owns_values(operator, kind):
-        total = torch.zeros((), dtype=torch.float64, device=getattr(store, "device", "cpu"))
         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()
             # Preserve the block reductions and their sequential FP64 order.
             total.add_(block.square_().sum().to(device=total.device))
-    return scalar_sum(operator, total, kind)
+    return scalar_sum(operator, total, kind, return_tensor=return_tensor)
 
 
 def copy_blocks(source, target, elements):
@@ -91,9 +93,10 @@ def fill_blocks(target, elements):
 
 
 def update_blocks(target, source, factor, elements, scale=1.):
+    device = getattr(target, "device", "cpu")
     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)
+        block = target.read(selected).to(device=device, dtype=target.dtype).clone()
+        block.mul_(scale).add_(source.read(selected).to(device=device, dtype=target.dtype), alpha=factor)
         target.write(selected, block)
 
 

diffct/_solver_tv.py

@@ -104,7 +104,8 @@ def total_variation(volume, *, operator=None, voxel_spacing=1., eps=1e-6, reduct
     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())
+        if getattr(operator, "world_size", 1) > 1:
+            global_value = scalar_sum(operator, local.detach(), "volume", return_tensor=True)
+            total = local + (global_value.to(local) - local.detach())
         elements = math.prod(operator.global_volume_shape)
     return total if reduction == "sum" else total / max(elements, 1)

diffct/analytical.py

@@ -523,8 +523,7 @@ def _fan_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec,
         d_sino, n_views, n_det, d_reco, Nx, Ny,
         _DTYPE(detector_spacing), d_src_pos, d_det_center, d_det_u_vec,
         cx, cy, _DTYPE(voxel_spacing),
-        _DTYPE(isocenter[0].item() / voxel_spacing),
-        _DTYPE(isocenter[1].item() / voxel_spacing),
+        *(_DTYPE(value / voxel_spacing) for value in isocenter.tolist()),
     )
 
     return reco
@@ -586,9 +585,7 @@ def _cone_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec, det
         _DTYPE(du), _DTYPE(dv),
         d_src_pos, d_det_center, d_det_u_vec, d_det_v_vec,
         cx, cy, cz, _DTYPE(voxel_spacing),
-        _DTYPE(isocenter[0].item() / voxel_spacing),
-        _DTYPE(isocenter[1].item() / voxel_spacing),
-        _DTYPE(isocenter[2].item() / voxel_spacing),
+        *(_DTYPE(value / voxel_spacing) for value in isocenter.tolist()),
     )
 
     # Convert back to (D, H, W) layout for the user.
@@ -640,7 +637,7 @@ def _analytical_transpose_raw(image, beam, geometry, parameters, isocenter, sino
         scale = (sdd_n / (2.0 * math.pi * sid_n)).view(
             (-1,) + (1,) * (len(sinogram_shape) - 1)
         )
-        iso = tuple(_DTYPE(value.item() / parameters[-1]) for value in isocenter)
+        iso = tuple(_DTYPE(value / parameters[-1]) for value in isocenter.tolist())
 
     if beam == "cone":
         D, H, W, du, dv, spacing = parameters

diffct/analytical_execution.py

@@ -53,6 +53,8 @@ def _sum(value, group, count):
 
 
 def _sum_blocks(value, group, count, block_elements):
+    if count == 1:
+        return value
     flat = value.reshape(-1)
     blocks = [_sum(flat.narrow(0, start, min(block_elements, flat.numel() - start)), group, count)
               for start in range(0, flat.numel(), block_elements)]
@@ -385,6 +387,21 @@ class AnalyticalExecution:
     def tensor_result(self):
         self._prepare_gradients()
         result = torch.zeros(self.volume_shape, device=self.output_device, dtype=self.dtype)
+        inputs = (self.source, self.projection_weights, self.angular_weights, self.filter_response)
+        active_gradients = torch.is_grad_enabled() and any(value is not None and value.requires_grad
+                                                         for value in inputs)
+        if self.partition.world_size == 1 and not active_gradients:
+            # Each tile still accumulates view batches in their original order.
+            # The numerical result itself holds the partial sums between batches.
+            for views, device in self._view_batches():
+                filtered, geometry = self._filtered(views, device)
+                for spatial, _ in _pixel_tiles(self.volume_shape, self.chunk_shape):
+                    partial = (self._native_gather(filtered, geometry, spatial) if self.backend == "numba"
+                               else self._portable_gather(filtered, geometry, views, spatial))
+                    result[spatial].add_(partial.to(self.output_device))
+                    del partial
+                del filtered, geometry
+            return result
         # Every owned tile's backward collective precedes the bundled input
         # transpose, even when this rank owns no data views. Empty space groups
         # have no tile collectives but still contribute zero input cotangents.

diffct/backends.py

@@ -61,7 +61,7 @@ def _rectangles(shape, chunks):
                     for start, step, size in zip(starts, chunks, shape))
 
 
-def _flat_admission_batches(owner, geometry, detector_chunk_shape):
+def _flat_admission_batches(owner, geometry, detector_chunk_shape, *, batched_corners=False):
     """Bound flat admission with corners and exact ambiguous-view fallback."""
     views = owner.view_slice
     if views.start == views.stop:
@@ -70,8 +70,13 @@ def _flat_admission_batches(owner, geometry, detector_chunk_shape):
     # coordinate; finite corner points and source differences bound the grid.
     endpoints = tuple((slice(0, 1),) if size == 1 else (slice(0, 1), slice(size - 1, size))
                       for size in owner.detector_shape)
-    for pixels in itertools.product(*endpoints):
-        yield views, pixels
+    if batched_corners:
+        # Native validation can compose the Cartesian corner set in one batch.
+        # Stepped pixel slices retain the exact global endpoint coordinates.
+        yield views, tuple(slice(0, size, max(1, size - 1)) for size in owner.detector_shape)
+    else:
+        for pixels in itertools.product(*endpoints):
+            yield views, pixels
     if owner.beam in ("parallel", "parallel3d"):
         return
     source, lower = (value[views].to(dtype=owner.dtype) for value in geometry[:2])

diffct/chunking.py

@@ -23,6 +23,7 @@ class _ChunkPlan:
     foreign_output: bool = False
     reservation_budget_bytes: int = (1 << 63) - 1
     estimated_reserved_bytes: int = 0
+    surface_batch: bool = True
 
 
 def _pixel_tiles(shape, chunk_shape):

diffct/kernels/surface_geometry.py

@@ -52,9 +52,10 @@ def compose_surface(frames, offsets, source, points, view_index, pixel_start,
     ``offset32`` gives the offset arithmetic dtype. Values are stored exactly
     in float64 staging buffers. Results use caller-owned float32 buffers.
 
-    A launch may compose several views with the same offsets: grid row ``y``
+    A launch may compose several views: grid row ``y``
     writes view ``view_index + y`` from the ``y``-th of ``gridDim.y`` equal
-    frame packets in ``frames``.
+    frame packets in ``frames``. A shared offset buffer holds ``rank * capacity``
+    values plus status; a larger buffer holds one such packet per output view.
     """
     pixel = cuda.grid(1)
     view_index += cuda.blockIdx.y
@@ -68,15 +69,16 @@ def compose_surface(frames, offsets, source, points, view_index, pixel_start,
     if pixel >= count:
         return
 
-    tangent = offsets[pixel]
-    depth = offsets[(rank - 1) * capacity + pixel]
+    offset_base = view_index * rank * capacity if offsets.size > rank * capacity + 1 else 0
+    tangent = offsets[offset_base + pixel]
+    depth = offsets[offset_base + (rank - 1) * capacity + pixel]
     tangent32 = offset32 and frame32[2]
     point32 = frame32[1] and tangent32
     normal32 = frame32[2]
     vertical = np.float64(0.0)
     vertical32 = False
     if rank == 3:
-        vertical = offsets[capacity + pixel]
+        vertical = offsets[offset_base + capacity + pixel]
         vertical32 = offset32 and frame32[3]
         point32 = point32 and vertical32
         normal32 = normal32 and frame32[3]

diffct/operators.py

@@ -16,7 +16,7 @@ from .storage import TensorStore
 from .streaming import _validate_stores
 from .surfaces import ParameterizedSurface, _GeometrySpec, _sample_geometry, _sample_offsets
 from .surface_gradients import surface_batch_vjp
-from .utils import TorchCUDABridge, _get_numba_external_stream_for
+from .utils import TorchCUDABridge, _get_numba_external_stream_for, _raise_geometry_errors
 from .backends import ReferenceProjector, _flat_admission_batches, select_backend, voxel_spacing_tuple
 from .distribution import Partition
 from .projectors import (
@@ -77,22 +77,27 @@ def _validate_trajectory(trajectory, beam, *, detector_surface=False, clone=True
     if not isinstance(trajectory, (tuple, list)) or len(trajectory) != count:
         raise TypeError(f"{beam} trajectory must contain {count} tensors")
 
-    geometry = []
+    geometry, checks = [], []
     n_views = None
     for component in trajectory:
-        if not isinstance(component, torch.Tensor):
-            raise TypeError("trajectory components must be tensors")
-        if component.ndim != 2 or component.shape[1] != dimensions:
-            raise ValueError(f"trajectory tensors must have shape (views, {dimensions})")
-        if not torch.is_floating_point(component):
-            raise TypeError("trajectory tensors must have floating-point dtype")
-        if component.shape[0] == 0:
-            raise ValueError("trajectory must contain at least one view")
-        if n_views is not None and component.shape[0] != n_views:
-            raise ValueError("trajectory tensors must have the same number of views")
+        try:
+            if not isinstance(component, torch.Tensor):
+                raise TypeError("trajectory components must be tensors")
+            if component.ndim != 2 or component.shape[1] != dimensions:
+                raise ValueError(f"trajectory tensors must have shape (views, {dimensions})")
+            if not torch.is_floating_point(component):
+                raise TypeError("trajectory tensors must have floating-point dtype")
+            if component.shape[0] == 0:
+                raise ValueError("trajectory must contain at least one view")
+            if n_views is not None and component.shape[0] != n_views:
+                raise ValueError("trajectory tensors must have the same number of views")
+        except (TypeError, ValueError):
+            # A deferred value error of an earlier component takes precedence.
+            if checks:
+                _raise_geometry_errors(checks)
+            raise
         n_views = component.shape[0]
-        if not torch.isfinite(component).all().item():
-            raise ValueError("trajectory tensors must contain only finite values")
+        checks.append((~torch.isfinite(component).all(), "trajectory tensors must contain only finite values"))
         geometry.append(component.detach().clone() if clone else component.detach())
 
     device = next((value.device for value in geometry if value.is_cuda), geometry[0].device)
@@ -103,28 +108,29 @@ def _validate_trajectory(trajectory, beam, *, detector_surface=False, clone=True
     )
     for axis in axes:
         norms = torch.linalg.vector_norm(axis, dim=1)
-        if not torch.allclose(norms, torch.ones_like(norms), rtol=1e-4, atol=1e-4):
-            raise ValueError("trajectory direction axes must be unit vectors")
+        checks.append((~torch.isclose(norms, torch.ones_like(norms), rtol=1e-4, atol=1e-4).all(),
+                       "trajectory direction axes must be unit vectors"))
     if beam in ("parallel", "cone", "parallel3d"):
         first, second = (geometry[0], geometry[2]) if beam == "parallel" else (geometry[2], geometry[3])
         dots = torch.sum(first.to(device) * second.to(device), dim=1)
-        if torch.any(torch.abs(dots) > 1e-4).item():
-            raise ValueError("trajectory direction axes must be orthogonal")
+        checks.append((torch.any(torch.abs(dots) > 1e-4), "trajectory direction axes must be orthogonal"))
     if beam == "parallel3d":
         direction = geometry[0].to(device=device, dtype=torch.float64)
         frame = torch.stack((direction, geometry[2].to(direction.device, torch.float64),
                              geometry[3].to(direction.device, torch.float64)), dim=-1)
         # Numerical rank rejects a tangent frame at machine roundoff without
         # imposing an application-level minimum angle of incidence.
-        if torch.any(torch.linalg.matrix_rank(frame) < 3).item():
-            raise ValueError("parallel3d rays must not be tangent to the detector plane")
+        # Bad entries already have a finite-value error above; keep the rank
+        # solver well-defined until the combined error flags reach the host.
+        frame = torch.where(torch.isfinite(frame), frame, 0.)
+        checks.append((torch.any(torch.linalg.matrix_rank(frame) < 3),
+                       "parallel3d rays must not be tangent to the detector plane"))
     if beam in ("fan", "cone") and not detector_surface:
         # float64 avoids overflow of half-precision input and dtype mismatches in cross().
         source = geometry[0].to(device=device, dtype=torch.float64)
         principal = geometry[1].to(source.device, torch.float64) - source
         lengths = torch.linalg.vector_norm(principal, dim=1)
-        if torch.any(lengths == 0).item():
-            raise ValueError("source and detector center must differ in every view")
+        checks.append((torch.any(lengths == 0), "source and detector center must differ in every view"))
         principal = principal / lengths[:, None]
         det_u = geometry[2].to(source.device, torch.float64)
         if beam == "fan":
@@ -132,9 +138,10 @@ def _validate_trajectory(trajectory, beam, *, detector_surface=False, clone=True
         else:
             normal = torch.linalg.cross(det_u, geometry[3].to(source.device, torch.float64), dim=1)
             facing = torch.sum(principal * normal, dim=1)
-        if torch.any(torch.abs(facing) < 1e-4).item():
-            raise ValueError("the detector must not be edge-on to the source in any view")
+        checks.append((torch.any(torch.abs(facing) < 1e-4),
+                       "the detector must not be edge-on to the source in any view"))
 
+    _raise_geometry_errors(checks)
     return tuple(geometry), n_views
 
 
@@ -544,16 +551,14 @@ class Projector:
                 )
                 for component in self._trajectory[:2]
             ])
-            if torch.any(distances.min(dim=0).values > 1e6).item():
-                raise ValueError(
+            _raise_geometry_errors([
+                (torch.any(distances.min(dim=0).values > 1e6),
                     "the source or the detector center must lie within 1e6 voxels "
-                    "of the volume center in every view"
-                )
-            if torch.any(distances > 1e15).item():
-                raise ValueError(
+                    "of the volume center in every view"),
+                (torch.any(distances > 1e15),
                     "the source and the detector center must lie within 1e15 voxels "
-                    "of the volume center"
-                )
+                    "of the volume center"),
+            ])
 
     def _agreement_revision(self):
         surface = self.detector_surface
@@ -768,6 +773,7 @@ class Projector:
         peers = len(active_devices) > 1 or (self._distributed and self.world_size > 1)
         extra_rays = (2 * int(self.partition == "space") + int(is_project and foreign_output)
                       + int(is_project and self.partition == "space" and peers))
+        surface_batch = True
 
         def fits(shape):
             points = surface or pixels != self.detector_shape
@@ -776,17 +782,21 @@ class Projector:
             buffers = _pipeline_buffers(self.beam, shape, batch, pixels, points,
                                         geometry_grad=self._geometry_trainable,
                                         ray_accumulators=extra_rays, peers=peers or foreign_output)
-            buffers += self._surface_work_buffers(batch, pixels, spec=geometry_spec, streamed=True)
+            workspace = self._surface_work_buffers(batch, pixels, spec=geometry_spec,
+                                                   streamed=True, batched=surface_batch)
             return (
-                sum(buffers) <= budget
-                and _reservation_bytes(buffers) <= reservation_budget
+                sum(buffers) + sum(workspace) <= budget
+                and _reservation_bytes(buffers + self._surface_reserved_buffers(
+                    batch, pixels, spec=geometry_spec, streamed=True, batched=surface_batch)) <= reservation_budget
                 and _largest_buffer_bytes(self.beam, shape, batch, pixels, points) <= allocation_limit
-                and max(self._surface_work_buffers(batch, pixels, spec=geometry_spec, streamed=True), default=0) <= allocation_limit
+                and max(workspace, default=0) <= allocation_limit
                 and _host_geometry_bytes(self.beam, batch, pixels, host_points) <= _pipeline._GEOMETRY_WORK_BYTES
             )
 
         while not fits(minimum_shape):
-            if self.view_chunk_size is None and batch > 1:
+            if self.detector_surface is not None and surface_batch:
+                surface_batch = False
+            elif self.view_chunk_size is None and batch > 1:
                 batch = max(1, batch // 2)
             elif self.detector_chunk_shape is None and max(pixels) > 1:
                 axis = max(range(len(pixels)), key=pixels.__getitem__)
@@ -796,25 +806,36 @@ class Projector:
             else:
                 raise RuntimeError("CUDA or host geometry memory budget cannot fit the minimum tile/view/pixel batch; reduce explicit limits")
         while not fits(chunk):
+            if self.detector_surface is not None and surface_batch:
+                surface_batch = False
+                continue
             if self.volume_chunk_shape is not None or max(chunk) == 1:
                 raise RuntimeError("CUDA memory budget cannot fit the requested tile/view batch")
             axis = max(range(len(chunk)), key=chunk.__getitem__)
             smaller = list(chunk)
             smaller[axis] = max(1, smaller[axis] // 2)
             chunk = tuple(smaller)
+        if not surface_batch:
+            # Smaller automatic tiles may now fit the complete workspace.
+            surface_batch = True
+            if not fits(chunk):
+                surface_batch = False
         estimate = _pipeline_bytes(self.beam, chunk, batch, pixels,
                                    surface or pixels != self.detector_shape,
                                    geometry_grad=self._geometry_trainable,
                                    ray_accumulators=extra_rays, peers=peers or foreign_output)
-        estimate += sum(self._surface_work_buffers(batch, pixels, spec=geometry_spec, streamed=True))
+        estimate += sum(self._surface_work_buffers(batch, pixels, spec=geometry_spec,
+                                                  streamed=True, batched=surface_batch))
         return _ChunkPlan(chunk, batch, devices, pixels, budget, allocation_limit,
                           estimated_gpu_bytes=estimate, foreign_output=foreign_output,
+                          surface_batch=surface_batch,
                           reservation_budget_bytes=reservation_budget,
                           estimated_reserved_bytes=_reservation_bytes(_pipeline_buffers(
                               self.beam, chunk, batch, pixels, surface or pixels != self.detector_shape,
                               geometry_grad=self._geometry_trainable,
                               ray_accumulators=extra_rays, peers=peers or foreign_output)
-                              + self._surface_work_buffers(batch, pixels, spec=geometry_spec, streamed=True)))
+                              + self._surface_reserved_buffers(batch, pixels, spec=geometry_spec,
+                                                               streamed=True, batched=surface_batch)))
 
     def _surface_capacity(self, views, pixels):
         rank, columns = len(self.volume_shape), 3 if self.beam == "cone" else 2
@@ -822,37 +843,80 @@ class Projector:
         maximum = max(views * math.prod(pixels) * rank, views * rank * frames)
         return max(1, min(math.prod(pixels), 21, (maximum - 1) // columns))
 
-    def _surface_work_buffers(self, views, pixels, slots=1, spec=None, *, streamed=False):
-        """Bound device-local offset staging and world-point expression storage."""
+    def _surface_snapshot_buffers(self):
+        """Allocator-rounded FP64 bounds for independent geometry snapshots."""
+        return ((_allocation_bytes(2 * self._n_views * math.prod(self.detector_shape) * 3),)
+                + (_allocation_bytes(2 * self._n_views * len(self.volume_shape)),) * len(self._trajectory))
+
+    def _surface_work_buffers(self, *args, **kwargs):
+        """One inventory bounding every surface workspace phase.
+
+        Its total is the largest phase total and its largest entry is the
+        largest allocation of any phase. Phases do not overlap in time.
+        """
+        phases = self._surface_work_phases(*args, **kwargs)
+        buffers = list(max(phases, key=sum))
+        peak = max(max(phase, default=0) for phase in phases)
+        if peak > max(buffers, default=0):
+            # The largest total covers the peak allocation: move the excess
+            # from the smallest entries into the largest one.
+            largest = buffers.index(max(buffers))
+            excess, buffers[largest] = peak - buffers[largest], peak
+            for index in sorted((i for i in range(len(buffers)) if i != largest), key=buffers.__getitem__):
+                taken = min(buffers[index], excess)
+                buffers[index] -= taken
+                excess -= taken
+            buffers = [size for size in buffers if size]
+        return tuple(buffers)
+
+    def _surface_reserved_buffers(self, *args, **kwargs):
+        """Every surface workspace allocation of a call, for reservation bounds.
+
+        Phases never hold live tensors together, but blocks of one phase can
+        leave allocator segments pinned for the next, so all of them count.
+        """
+        return tuple(size for phase in self._surface_work_phases(*args, **kwargs) for size in phase)
+
+    def _surface_work_phases(self, views, pixels, slots=1, spec=None, *, streamed=False, batched=True):
+        """Bound device-local offset staging and world-point expression storage.
+
+        Returns the buffer inventories of phases that never coexist.
+        """
         if self.detector_surface is None:
-            return ()
+            return ((),)
         if streamed:
             rank, frames = len(self.volume_shape), 4 if self.beam == "cone" else 3
             capacity = self._surface_capacity(views, pixels)
             columns = 3 if self.beam == "cone" else 2
             shared = spec is not None and spec.kind == "offsets" and spec.shared
-            if shared:
+            if shared or batched:
                 capacity = math.prod(pixels)
-            # Shared offsets pack the frame packets of every view of a job.
-            buffers = (_allocation_bytes(2 * (views if shared else 1) * (frames + (3 if rank == 3 else 0)) * rank),
-                       _allocation_bytes(2 * (columns * capacity + 1)))
+            # Batched jobs retain every frame and, unless shared, every offset.
+            buffers = (_allocation_bytes(2 * (views if shared or batched else 1) * (frames + (3 if rank == 3 else 0)) * rank),
+                       _allocation_bytes(2 * (columns * capacity * (views if batched and not shared else 1) + 1)))
             if self._geometry_trainable:
                 # Replay one view and a bounded detector row at a time.
-                replay = rank * min(capacity, pixels[-1])
+                replay = rank * min(self._surface_capacity(views, pixels), pixels[-1])
                 replay_buffers = ((_allocation_bytes(2 * replay),) * 8
                                   + (_allocation_bytes(2 * rank),) * (frames + 3))
                 if isinstance(self.detector_surface, ParameterizedSurface):
                     # A CUDA sampler needs a device cotangent for its selected output.
                     replay_buffers += (_allocation_bytes(6 * views * math.prod(pixels)),)
-                if sum(replay_buffers) > sum(buffers):
-                    return replay_buffers
-            return buffers
+                # Composition frees its workspace before replay starts.
+                return buffers, replay_buffers
+            return (buffers,)
         elements = views * math.prod(pixels) * 3
         dtype = torch.float64 if spec is None or spec.dtype is None else spec.dtype
         buffers = (torch.empty((), dtype=dtype).element_size() * elements,) * (4 * slots)
         if dtype != torch.float64:
             buffers += (8 * elements,) * slots
-        return buffers + (views * 3 * 8 * (4 if self.beam == "cone" else 3),) * slots
+        buffers += (views * 3 * 8 * (4 if self.beam == "cone" else 3),) * slots
+        if spec is None and not isinstance(self.detector_surface, ParameterizedSurface):
+            # The preliminary full plan precedes the callback. Admit independent
+            # snapshots plus one bounded arena for the callback's saved tensors.
+            snapshots = self._surface_snapshot_buffers()
+            buffers += (*snapshots, sum(snapshots))
+        return (buffers,)
 
     def _surface_spec(self, frames, offsets):
         dtype = torch.float32
@@ -873,9 +937,10 @@ class Projector:
     def _surface_validation_inputs(self):
         """Stage frame metadata once without expanding surface world points."""
         with torch.enable_grad():
-            frames = tuple(value.to(device="cpu") for value in (self._learnable or self._trajectory))
+            frames = self._learnable or self._trajectory
             surface = self.detector_surface
             if isinstance(surface, ParameterizedSurface):
+                frames = tuple(value.to(device="cpu") for value in frames)
                 parameters = tuple(value.to(device="cpu") for value in surface.parameters)
                 geometry = (*frames, *parameters)
                 return (geometry, _GeometrySpec("sampled", surface.sampler, torch.float64),
@@ -1154,6 +1219,7 @@ class Projector:
         return dict(operation="project" if is_project else "backproject", schedule=plan.schedule,
                     chunk_shape=tuple(plan.chunk_shape), view_chunk_size=plan.view_chunk_size,
                     detector_chunk_shape=tuple(plan.detector_chunk_shape or self.detector_shape),
+                    surface_batch=plan.surface_batch,
                     window_size=plan.window_size, gpu_budget_bytes=plan.gpu_budget_bytes,
                     estimated_gpu_bytes=plan.estimated_gpu_bytes,
                     reservation_budget_bytes=plan.reservation_budget_bytes,
@@ -1172,14 +1238,16 @@ class Projector:
             return _pipeline_bytes(self.beam, shape, plan.view_chunk_size, pixels, points,
                                    slots=slots, retained=retained, geometry_grad=self._geometry_trainable,
                                    peers=peers, ray_accumulators=ray_accumulators) + sum(
-                                       self._surface_work_buffers(plan.view_chunk_size, pixels, slots, spec, streamed=True))
+                                       self._surface_work_buffers(plan.view_chunk_size, pixels, slots, spec,
+                                                                  streamed=True, batched=plan.surface_batch))
 
         def reserved(shape, slots=1, retained=1, ray_accumulators=0):
             return _reservation_bytes(_pipeline_buffers(
                 self.beam, shape, plan.view_chunk_size, pixels, points,
                 slots=slots, retained=retained, geometry_grad=self._geometry_trainable,
                 peers=peers, ray_accumulators=ray_accumulators)
-                + self._surface_work_buffers(plan.view_chunk_size, pixels, slots, spec, streamed=True))
+                + self._surface_reserved_buffers(plan.view_chunk_size, pixels, slots, spec,
+                                                 streamed=True, batched=plan.surface_batch))
 
         shapes = [plan.chunk_shape]
         if self.volume_chunk_shape is None and len(volume_shape) == 3:
@@ -1255,8 +1323,8 @@ class Projector:
         return _pipeline_buffers(
             self.beam, plan.chunk_shape, plan.view_chunk_size, pixels, points,
             slots=plan.slots, retained=plan.window_size, geometry_grad=self._geometry_trainable,
-            peers=peers, ray_accumulators=extra_rays) + self._surface_work_buffers(
-                plan.view_chunk_size, pixels, plan.slots, spec, streamed=True)
+            peers=peers, ray_accumulators=extra_rays) + self._surface_reserved_buffers(
+                plan.view_chunk_size, pixels, plan.slots, spec, streamed=True, batched=plan.surface_batch)
 
     def _choose_schedule(self, source, sink, is_project, geometry, plan, spec):
         candidates = self._candidate_plans(plan, is_project, spec)
@@ -1266,7 +1334,8 @@ class Projector:
                tuple(device.index for device in plan.devices), str(source.dtype), plan.gpu_budget_bytes,
                plan.allocation_limit, plan.reservation_budget_bytes,
                str(getattr(source, "device", "cpu")), str(getattr(sink, "device", "cpu")),
-               spec.kind if spec is not None else "flat", self.schedule)
+               spec.kind if spec is not None else "flat", self.schedule, plan.surface_batch,
+               spec.shared if spec is not None else False)
         collective = self._distributed and self.world_size > 1
         if not collective and key in self._schedule_cache:
             return self._schedule_cache[key], [], True
@@ -1337,6 +1406,7 @@ class Projector:
                                       geometry, spec, views, pixels)
         origin = torch.cuda.current_stream(offsets.device) if offsets.is_cuda else None
         frames = geometry[:frame_count]
+        frame_origins = tuple(torch.cuda.current_stream(value.device) for value in frames if value.is_cuda)
         common = torch.float32
         for value in (*frames, offsets):
             common = torch.promote_types(common, value.dtype)
@@ -1349,25 +1419,33 @@ class Projector:
         # Offsets shared by all views are staged once per call and composed
         # with one launch for all views instead of per view and row segment.
         shared = spec.kind == "offsets" and offsets.ndim == len(pixel_shape) + 1
+        batched = shared or groups.get("_surface_batch", True)
         pixel_count = math.prod(pixel_shape)
         # One packed frame packet per view: source, center, axes, normal and
         # float32 cross-product scratch. Shared offsets pack every view of a job.
         stride = (frame_count + (3 if rank == 3 else 0)) * rank
         if "_surface_workspace" not in groups:
-            capacity = pixel_count if shared else self._surface_capacity(count, pixel_shape)
+            maximum_views = groups.get("_max_geometry", (count * rank,))[0] // rank
+            capacity = (groups.get("_max_rays", count * pixel_count) // maximum_views
+                        if batched else self._surface_capacity(count, pixel_shape))
             groups["_surface_workspace"] = (
-                pipe.empty((count if shared else 1) * stride, dtype=torch.float64),
-                pipe.empty((3 if rank == 3 else 2) * capacity + 1, dtype=torch.float64), capacity,
+                pipe.empty((maximum_views if batched else 1) * stride, dtype=torch.float64),
+                pipe.empty(rank * capacity * (maximum_views if batched and not shared else 1) + 1,
+                           dtype=torch.float64), capacity,
             )
         packed_frames, packed_offsets, capacity = groups["_surface_workspace"]
         columns = (0, 1, 2) if rank == 3 else (0, 2)
         source, points = targets[0].view(count, rank), targets[1].view(count, -1, rank)
         # A workspace sized for an earlier, smaller tile keeps the segment path.
         shared = shared and capacity >= pixel_count
-        if shared and packed_frames.numel() < count * stride:
+        batched = batched and capacity >= pixel_count
+        if batched and packed_frames.numel() < count * stride:
             # Earlier launches on the upload stream may still read the smaller buffer.
             pipe.retain(pipe.event(pipe.upload), packed_frames, None)
             packed_frames = pipe.empty(count * stride, dtype=torch.float64)
+            if not shared:
+                pipe.retain(pipe.event(pipe.upload), packed_offsets, None)
+                packed_offsets = pipe.empty(rank * capacity * count + 1, dtype=torch.float64)
             groups["_surface_workspace"] = (packed_frames, packed_offsets, capacity)
         # Every job of a call slices the same offset snapshot: stage it once per
         # call, pixel tile and workspace. The key holds the snapshot and the
@@ -1380,39 +1458,55 @@ class Projector:
                            and groups["_surface_shared_staged"][1] is staged_key[1]
                            and groups["_surface_shared_staged"][2] == staged_key[2])
         else:
-            restage = shared
+            restage = batched
             groups.pop("_surface_shared_staged", None)
         with torch.cuda.device(pipe.device), torch.cuda.stream(pipe.upload), cuda.gpus[pipe.device.index]:
             if origin is not None:
                 pipe.upload.wait_stream(origin)
+            for frame_origin in frame_origins:
+                pipe.upload.wait_stream(frame_origin)
             packed_offsets[-1].zero_()
             stream = _get_numba_external_stream_for(pipe.upload)
             arrays = tuple(TorchCUDABridge.tensor_to_cuda_array(value)
                            for value in (packed_frames, packed_offsets, source, points))
             if restage:
-                # Pack the float64 columns where the offsets live, then copy each
-                # column into the workspace without a device staging buffer.
-                flat = offsets.detach().reshape(pixel_count, 3)
+                # Preserve strided sampler outputs: copy directly into admitted
+                # column storage without flattening the input into a temporary.
+                offset_views = 1 if shared else count
+                packed = packed_offsets[:offset_views * rank * capacity].view(offset_views, rank, capacity)
+                selected_offsets = offsets.detach()
+                if selected_offsets.ndim == len(pixel_shape) + 1:
+                    selected_offsets = selected_offsets.unsqueeze(0)
                 if offsets.is_cuda:
                     for column, original in enumerate(columns):
-                        packed_offsets[column * capacity:column * capacity + pixel_count].copy_(flat[:, original])
+                        target = packed[:, column, :pixel_count].view(offset_views, *pixel_shape)
+                        value = selected_offsets[..., original].expand(offset_views, *pixel_shape)
+                        _pipeline._copy(target, value, pipe.stats)
+                        pipe.retain(pipe.event(pipe.upload), value, packed_offsets)
                 else:
-                    host = flat[:, list(columns)].T.to(dtype=torch.float64).contiguous()
-                    for column in range(len(columns)):
-                        _pipeline._copy(packed_offsets[column * capacity:column * capacity + pixel_count],
-                                        host[column], pipe.stats)
+                    host = torch.empty((offset_views, rank, capacity), dtype=torch.float64)
+                    for column, original in enumerate(columns):
+                        host[:, column, :pixel_count].view(offset_views, *pixel_shape).copy_(
+                            selected_offsets[..., original])
+                    _pipeline._copy(packed, host, pipe.stats)
                     pipe.retain(pipe.event(pipe.upload), host, packed_offsets)
-                if cacheable:
+                if shared and cacheable:
                     groups["_surface_shared_staged"] = staged_key
-            if shared:
+            if batched:
                 # Stage the frame packets of every view in one transfer and
                 # compose all views with one launch; grid row y composes view y.
                 block = packed_frames[:count * stride].view(count, stride)
-                packet = torch.zeros((count, stride), dtype=torch.float64)
-                for index, value in enumerate(frames):
-                    packet[:, index * rank:(index + 1) * rank].copy_(value[views].detach())
-                _pipeline._copy(block, packet, pipe.stats)
-                pipe.retain(pipe.event(pipe.upload), packet, packed_frames)
+                if frame_origins:
+                    for index, value in enumerate(frames):
+                        selected = value[views].detach()
+                        _pipeline._copy(block[:, index * rank:(index + 1) * rank], selected, pipe.stats)
+                        pipe.retain(pipe.event(pipe.upload), selected, packed_frames)
+                else:
+                    packet = torch.zeros((count, stride), dtype=torch.float64)
+                    for index, value in enumerate(frames):
+                        packet[:, index * rank:(index + 1) * rank].copy_(value[views].detach())
+                    _pipeline._copy(block, packet, pipe.stats)
+                    pipe.retain(pipe.event(pipe.upload), packet, packed_frames)
                 if rank == 3:
                     normal = block[:, 4 * rank:5 * rank]
                     axis_u, axis_v = block[:, 2 * rank:3 * rank], block[:, 3 * rank:4 * rank]
@@ -1433,7 +1527,7 @@ class Projector:
                         self.beam == "parallel", frame32, offset32,
                         tuple(centre) + (0.,) * (3 - rank), self.voxel_spacing,
                     )
-            for local in range(0 if shared else count):
+            for local in range(0 if batched else count):
                 packet = torch.cat([value[views.start + local].detach().to(device="cpu", dtype=torch.float64)
                                     for value in frames])
                 _pipeline._copy(packed_frames[:frame_count * rank], packet, pipe.stats)
@@ -1556,7 +1650,7 @@ class Projector:
         points = spec is not None or self.detector_chunk_shape is not None or pixels != self.detector_shape
         sizes = ((plan.view_chunk_size * rank, rays * rank) if points
                  else (plan.view_chunk_size * rank,) * (4 if self.beam == "cone" else 3))
-        return {"_max_rays": rays, "_max_geometry": sizes}
+        return {"_max_rays": rays, "_max_geometry": sizes, "_surface_batch": plan.surface_batch}
 
     def _prefetch_job(self, job, source, geometry, spec, transfer, groups, index, is_project,
                       native, reduced, active_rays, next_epoch=False):
@@ -2127,7 +2221,7 @@ class Projector:
         spacing = self.detector_spacing if self.beam == "cone" else self.detector_spacing[0]
         sampled = geometry_spec is not None and geometry_spec.kind in ("sampled", "offsets")
         grads = [torch.zeros(component.shape, dtype=torch.float32,
-                             device="cpu" if sampled else volume.device) for component in geometry]
+                             device=component.device if sampled else volume.device) for component in geometry]
         for device_rank, device in enumerate(devices):
             local_slice = _balanced_slice(local_views, device_rank, len(devices))
             if local_slice.start == local_slice.stop:
@@ -2287,7 +2381,8 @@ class Projector:
             self.last_execution_stats = stats
             return tuple(grad.to(device=component.device, dtype=component.dtype)
                          for grad, component in zip(grads, geometry))
-        groups = {device: {} for device in plan.devices}
+        # Replay composition follows the admitted batched/fallback workspace.
+        groups = {device: {"_surface_batch": plan.surface_batch} for device in plan.devices}
         with _pipeline._TransferPipeline(plan.devices, self._pool_elements(plan, geometry_spec), stats,
                                         slots=plan.slots) as transfer:
             for spatial_slice, shape, centre in self._volume_tiles(plan.chunk_shape):
@@ -2514,22 +2609,28 @@ class Projector:
                             staged = tuple(value[owned].to(device=device, dtype=owner.dtype) for value in geometry)
                             admission_owner = copy(owner)
                             admission_owner.view_slice = slice(0, shard.stop - shard.start)
-                            for views, selected in _flat_admission_batches(admission_owner, staged, pixels):
+                            nonfinite = torch.zeros((), device=device, dtype=torch.bool)
+                            coincident = torch.zeros_like(nonfinite)
+                            for views, selected in _flat_admission_batches(
+                                admission_owner, staged, pixels, batched_corners=True,
+                            ):
                                 source, points = _sample_geometry(
                                     owner.beam, owner.detector_shape, owner.detector_spacing, staged,
                                     None, views, selected, owner.voxel_spacing, device=device,
                                     dtype=owner.dtype, validate=False,
                                 )
-                                if not torch.isfinite(source).all().item() or not torch.isfinite(points).all().item():
-                                    raise ValueError("detector surface world coordinates must be finite")
+                                nonfinite = nonfinite | ~torch.isfinite(source).all() | ~torch.isfinite(points).all()
                                 sources = source.reshape(source.shape[0], *([1] * len(selected)), -1)
-                                if owner.beam != "parallel" and torch.any(torch.all(points == sources, dim=-1)).item():
-                                    raise ValueError("source and detector surface pixels must differ")
+                                if owner.beam != "parallel":
+                                    coincident = coincident | torch.any(torch.all(points == sources, dim=-1))
                                 directions = (source if owner.beam == "parallel" else
                                               points - sources)
-                                if not torch.isfinite(directions).all().item():
-                                    raise ValueError("detector surface world coordinates must be finite")
+                                nonfinite = nonfinite | ~torch.isfinite(directions).all()
                                 del source, points, sources, directions
+                            _raise_geometry_errors([
+                                (nonfinite, "detector surface world coordinates must be finite"),
+                                (coincident, "source and detector surface pixels must differ"),
+                            ])
                             del staged
             except (TypeError, ValueError, RuntimeError) as failure:
                 error = failure
@@ -2591,7 +2692,25 @@ class Projector:
     def _prepare_execution(self, tensor, is_project, output_store=None):
         device = torch.device(getattr(tensor, "device", "cpu"))
         if self.detector_surface is not None:
-            spec, geometry = self._effective_geometry(None, device)
+            # Only eligible full calls need a preliminary snapshot reservation.
+            # Explicit streamed jobs first obtain exact shared-offset metadata.
+            devices = self._devices or (device,)
+            active_devices = devices[:min(len(devices), self.projection_shape[0])]
+            full_candidate = (not isinstance(self.detector_surface, ParameterizedSurface)
+                              and isinstance(tensor, torch.Tensor) and device.type == "cuda"
+                              and self.world_size == 1 and device in active_devices
+                              and output_store is None and self.volume_chunk_shape is None
+                              and self.view_chunk_size is None and self.detector_chunk_shape is None
+                              and self.schedule == "auto" and self.partition == "views")
+            if full_candidate:
+                plan = self._execution_plan(tensor, is_project, output_store)
+                spec, geometry = self._effective_geometry(plan, device)
+                if plan is None:
+                    return plan, spec, geometry
+            else:
+                spec, geometry = self._effective_geometry(None, torch.device("cpu"))
+            # Exact dtype/shared metadata can admit a borderline full call;
+            # reuse the one host snapshot if the preliminary plan streamed.
             plan = self._execution_plan(tensor, is_project, output_store, spec)
         else:
             plan = self._execution_plan(tensor, is_project, output_store)
@@ -2606,11 +2725,28 @@ class Projector:
             return None, tuple(component.clone() if plan is None or component.device.type == "cpu"
                                else component.to(device="cpu", copy=True)
                                for component in self._learnable)
-        # Callback graphs can save mutable captured parameters. Snapshot their
-        # values on the host, restoring their original device only when the
-        # callback's own backward needs them. Native tile graphs are not saved.
+        # Full CUDA execution has admitted its snapshots before the callback.
+        # Streamed execution and explicit CPU samplers retain host snapshots.
+        local = plan is None and input_device.type == "cuda"
+        capacity = sum(self._surface_snapshot_buffers()) if local else 0
+        arena, cursor = None, 0
+
         def pack(tensor):
-            return tensor.detach().to(device="cpu").clone(), tensor.device
+            nonlocal arena, cursor
+            if (local and tensor.device == input_device and tensor.layout == torch.strided
+                    and not tensor.is_nested and not tensor.is_quantized):
+                itemsize = tensor.element_size()
+                start = (cursor + itemsize - 1) // itemsize * itemsize
+                end = start + tensor.numel() * itemsize
+                if end <= capacity:
+                    with torch.no_grad():
+                        if arena is None:
+                            arena = torch.empty(capacity, dtype=torch.uint8, device=input_device)
+                        saved = arena[start:end].view(tensor.dtype).reshape(tensor.shape)
+                        saved.copy_(tensor.detach())
+                    cursor = end
+                    return saved, tensor.device
+            return tensor.detach().to(device="cpu", copy=True), tensor.device
 
         def unpack(saved):
             tensor, device = saved
@@ -2630,12 +2766,14 @@ class Projector:
                 with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
                     offsets = self._surface_offsets()
                     values = (*frames, offsets)
-                    snapshots = tuple(value.clone() if value.device.type == "cpu" else value.to(device="cpu", copy=True)
+                    snapshots = tuple(value.clone() if (local and value.device == input_device)
+                                      or value.device.type == "cpu" else value.to(device="cpu", copy=True)
                                       for value in values)
                 spec = self._surface_spec(frames, offsets)
                 metadata = (tuple(offsets.shape), str(offsets.dtype))
             self._snapshot_d2h_bytes += sum(value.numel() * value.element_size()
-                                            for value in values if value.is_cuda)
+                                            for value, snapshot in zip(values, snapshots)
+                                            if value.is_cuda and not snapshot.is_cuda)
         except (TypeError, ValueError, RuntimeError) as failure:
             error = failure
         self._partition.preflight(error, metadata, cache_key="surface snapshot",

diffct/solvers.py

@@ -96,7 +96,7 @@ def cgls(operator, measurements, iterations=30, *, x0=None, rtol=1e-6, atol=0.,
         x = _initial(operator, data, x0)
         residual = data.clone() if x0 is None else data - operator.project(x)
         norm = math.sqrt(squared_norm(operator, data, "rays"))
-        history = [math.sqrt(squared_norm(operator, residual, "rays"))]
+        history = [norm if x0 is None else math.sqrt(squared_norm(operator, residual, "rays"))]
         completed = 0
         adjoint = (operator.backproject(residual) if iterations and history[0] > atol + rtol * norm
                    else torch.zeros_like(x))
@@ -117,11 +117,15 @@ def cgls(operator, measurements, iterations=30, *, x0=None, rtol=1e-6, atol=0.,
         states["x"].add_(states["p"], alpha=alpha)
         states["r"].sub_(states["q"], alpha=alpha)
         states["s"] = operator.backproject(states["r"])
-        next_gamma = squared_norm(operator, states["s"], "volume")
+        next_gamma = squared_norm(operator, states["s"], "volume", return_tensor=True)
+        residual_norm = squared_norm(operator, states["r"], "rays", return_tensor=True)
+        # Both reductions retain the same FP64 block/collective order. Read
+        # their results together; denominator breakdown is still checked first.
+        next_gamma, residual_norm = torch.stack((next_gamma, residual_norm.to(next_gamma.device))).cpu().tolist()
         states["p"].mul_(next_gamma / gamma).add_(states["s"])
         gamma = next_gamma
         completed += 1
-        history.append(math.sqrt(squared_norm(operator, states["r"], "rays")))
+        history.append(math.sqrt(residual_norm))
         if checkpoint_directory is not None:
             save_checkpoint(operator, checkpoint_directory, states, metadata, completed, history,
                             norm, 1048576, gamma=gamma)
@@ -148,11 +152,12 @@ def sirt(operator, measurements, iterations=100, *, x0=None, rtol=1e-6, atol=0.,
         saved, stores = load_checkpoint(operator, resume, metadata, iterations, 1048576, ("x", "row", "column"))
         states = {name: _read_tensor(store, data.device) for name, store in stores.items()}
         completed, history, norm = saved["iteration"], list(saved["residual_history"]), saved["measurement_norm"]
+        residual = None
     else:
         states = dict(x=_initial(operator, data, x0))
         completed, norm = 0, math.sqrt(squared_norm(operator, data, "rays"))
-        initial = data if x0 is None else data - operator.project(states["x"])
-        history = [math.sqrt(squared_norm(operator, initial, "rays"))]
+        residual = data if x0 is None else data - operator.project(states["x"])
+        history = [norm if x0 is None else math.sqrt(squared_norm(operator, residual, "rays"))]
         if iterations and history[0] > atol + rtol * norm:
             row = operator.project(torch.ones_like(states["x"]))
             column = operator.backproject(torch.ones_like(data))
@@ -160,12 +165,15 @@ def sirt(operator, measurements, iterations=100, *, x0=None, rtol=1e-6, atol=0.,
             states["column"] = torch.where(column > 0, column.reciprocal(), 0.)
     threshold, stopped = atol + rtol * norm, False
     while completed < iterations and history[-1] > threshold:
-        residual = data - operator.project(states["x"])
+        # A callback may change either the iterate or the operator after history.
+        if callbacks or residual is None:
+            residual = data - operator.project(states["x"])
         states["x"].add_(states["column"] * operator.backproject(states["row"] * residual), alpha=relaxation)
         if nonnegative:
             states["x"].clamp_(min=0)
         completed += 1
-        history.append(math.sqrt(squared_norm(operator, data - operator.project(states["x"]), "rays")))
+        residual = data - operator.project(states["x"])
+        history.append(math.sqrt(squared_norm(operator, residual, "rays")))
         if checkpoint_directory is not None:
             save_checkpoint(operator, checkpoint_directory, states, metadata, completed, history, norm, 1048576)
         if callbacks:
@@ -195,14 +203,16 @@ def tv_reconstruction(operator, measurements, iterations=100, weight=1., lr=.1,
                               for name in ("exp_avg", "exp_avg_sq", "step")}
         optimizer.state[x]["step"] = optimizer.state[x]["step"].reshape(())
         completed, history, norm = saved["iteration"], list(saved["residual_history"]), saved["measurement_norm"]
+        residual = None
     else:
         completed, norm = 0, math.sqrt(squared_norm(operator, data, "rays"))
-        initial = data if x0 is None else data - operator.project(x)
-        history = [math.sqrt(squared_norm(operator, initial, "rays"))]
+        residual = -data if x0 is None else operator.project(x) - data
+        history = [norm if x0 is None else math.sqrt(squared_norm(operator, residual, "rays"))]
     threshold, stopped = atol + rtol * norm, False
     ray_count = scalar_sum(operator, data.numel(), "rays")
     while completed < iterations and history[-1] > threshold:
-        residual = operator.project(x) - data
+        if callbacks or residual is None:
+            residual = operator.project(x) - data
         gradient = operator.backproject(residual) / max(ray_count, 1)
         if weight:
             with torch.enable_grad():
@@ -214,7 +224,8 @@ def tv_reconstruction(operator, measurements, iterations=100, weight=1., lr=.1,
         optimizer.step()
         x.clamp_(min=0)
         completed += 1
-        history.append(math.sqrt(squared_norm(operator, data - operator.project(x), "rays")))
+        residual = operator.project(x) - data
+        history.append(math.sqrt(squared_norm(operator, residual, "rays")))
         if checkpoint_directory is not None:
             states = dict(x=x, **optimizer.state[x])
             states["step"] = states["step"].reshape(1)

diffct/surfaces.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass
 
 import torch
 
+from .utils import _raise_geometry_errors
 
 @dataclass(frozen=True)
 class _GeometrySpec:
@@ -41,7 +42,8 @@ class ParameterizedSurface:
 
 
 def _detector_grids(detector_shape, spacing, pixels):
-    axes = [(torch.arange(selected.start, selected.stop, dtype=torch.float64, device="cpu") + .5 - size / 2) * pitch
+    axes = [(torch.arange(selected.start, selected.stop, selected.step or 1,
+                          dtype=torch.float64, device="cpu") + .5 - size / 2) * pitch
             for selected, size, pitch in zip(pixels, detector_shape, spacing)]
     if len(axes) == 2:
         return torch.meshgrid(*axes, indexing="ij")
@@ -106,9 +108,7 @@ def _sample_geometry(beam, detector_shape, spacing, geometry, spec, views, pixel
         depth = offsets[..., 2].to(device=device, dtype=offset_dtype)
         if offsets.ndim == dimensions + 1:
             tangent, depth = tangent.unsqueeze(0), depth.unsqueeze(0)
-    if validate:
-        if not torch.isfinite(offsets).all().item():
-            raise ValueError("detector surface coordinates must be finite")
+    checks = [(~torch.isfinite(offsets).all(), "detector surface coordinates must be finite")] if validate else []
 
     def frame(component):
         selected = component if local_geometry else component[views]
@@ -132,17 +132,18 @@ def _sample_geometry(beam, detector_shape, spacing, geometry, spec, views, pixel
     source = (frames[0] if local_geometry else frames[0][views]).to(device=device, dtype=source_dtype)
     if validate:
         for component in (source, points):
-            if not torch.isfinite(component).all().item() or not torch.isfinite(component.float()).all().item():
-                raise ValueError("surface world coordinates must be finite in float32")
+            checks.append((~torch.isfinite(component).all() | ~torch.isfinite(component.float()).all(),
+                           "surface world coordinates must be finite in float32"))
         if beam not in ("parallel", "parallel3d") and view_count:
             sources = source.unsqueeze(1)
             endpoints = points.reshape(view_count, -1, source.shape[-1])
-            if torch.any(torch.all(sources.float() == endpoints.float(), dim=-1)).item():
-                raise ValueError("source and detector surface pixels must differ")
+            checks.append((torch.any(torch.all(sources.float() == endpoints.float(), dim=-1)),
+                           "source and detector surface pixels must differ"))
             source_distance = torch.linalg.vector_norm(sources.double(), dim=-1) / voxel_spacing
             pixel_distance = torch.linalg.vector_norm(endpoints.double(), dim=-1) / voxel_spacing
-            if torch.any(torch.minimum(source_distance, pixel_distance) > 1e6).item():
-                raise ValueError("surface source or pixel must lie within 1e6 voxels")
-            if torch.any(source_distance > 1e15).item() or torch.any(pixel_distance > 1e15).item():
-                raise ValueError("surface coordinates must lie within 1e15 voxels")
+            checks.append((torch.any(torch.minimum(source_distance, pixel_distance) > 1e6),
+                           "surface source or pixel must lie within 1e6 voxels"))
+            checks.append((torch.any(source_distance > 1e15) | torch.any(pixel_distance > 1e15),
+                           "surface coordinates must lie within 1e15 voxels"))
+        _raise_geometry_errors(checks)
     return source, points

diffct/utils.py

@@ -18,6 +18,16 @@ import torch
 from .constants import _DTYPE, _TPB_2D, _TPB_3D, _INF, _EPSILON
 
 
+def _raise_geometry_errors(checks):
+    """Read one validity scalar; fetch individual error flags only on failure."""
+    device = next((flag.device for flag, _ in checks if flag.is_cuda), checks[0][0].device)
+    failed = torch.stack([flag.to(device=device) for flag, _ in checks])
+    if failed.any().item():
+        for invalid, (_, message) in zip(failed.tolist(), checks):
+            if invalid:
+                raise ValueError(message)
+
+
 # ============================================================================
 # Device Management Utilities
 # ============================================================================

{diffct-2.4.2.dist-info → diffct-2.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffct
-Version: 2.4.2
+Version: 2.4.3
 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/
@@ -68,7 +68,7 @@ Description-Content-Type: text/markdown
 > **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
 > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
 
-This README describes the source checkout, including [unreleased APIs](CHANGELOG.md#unreleased). Per-call trajectories, angle-tensor circular helpers, `transform_trajectory`, `parallel3d` and `native_volume_layout` require installation from source.
+This README describes the source checkout. Features listed under [Unreleased](CHANGELOG.md#unreleased) require installation from source.
 
 ## Why diffct
 
@@ -177,18 +177,18 @@ volume slabs use `--partition space` under torchrun. Run
 
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.4.1">
-    <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.4.1" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.4.3">
+    <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.4.3" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
   </picture>
 </p>
 
 | 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
 |---|---:|---:|---:|
-| Forward projection | 569.99 ms | 151.02 ms (3.8×) | 81.50 ms (7.0×) |
-| Adjoint (backprojection) | 1061.36 ms | 278.75 ms (3.8×) | 220.34 ms (4.8×) |
-| CGLS iteration | 1756.03 ms | 468.05 ms (3.8×) | 333.08 ms (5.3×) |
+| Forward projection | 566.79 ms | 150.76 ms (3.8×) | 81.58 ms (6.9×) |
+| Adjoint (backprojection) | 1060.77 ms | 279.37 ms (3.8×) | 224.04 ms (4.7×) |
+| CGLS iteration | 1753.73 ms | 468.09 ms (3.7×) | 337.39 ms (5.2×) |
 
-diffct 2.4.1, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
+diffct 2.4.3, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
 
 Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
 
@@ -228,7 +228,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
 
 ## Examples
 
-`quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `benchmark_trajectory_admission.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
+`quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `benchmark_trajectory_admission.py`, `benchmark_surface_geometry.py`, `benchmark_gpu_overheads.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
 
 ## Documentation
 

{diffct-2.4.2.dist-info → diffct-2.4.3.dist-info}/RECORD

@@ -1,33 +1,33 @@
-diffct/__init__.py,sha256=saZamKDQ9CLB1HoVwIeZG6HMWNnbZbdNXtaT_MIpDLw,2965
-diffct/_solver_state.py,sha256=DKfrgrKcSpx3YetbF8hjA2yzMzJag-iHVpMZGup_dLw,16258
-diffct/_solver_tv.py,sha256=9nE1DVRzdUlcfd-OQQ86GpntX6CcaVvsSVWK-L4uR_I,5783
-diffct/analytical.py,sha256=u2PUp8GTNHoqwP2S4s0y9xNjiVBLfnYMULuuThcWmdU,33245
-diffct/analytical_execution.py,sha256=ZuOmt45-TQg3wb2J0erJe2fXIq0ZQr-uXrrbYiAWg7E,25660
-diffct/backends.py,sha256=VaALT_iE-aL4jAbzJ5xGvINfGLrHqqvuPozz0R8h048,47960
-diffct/chunking.py,sha256=EzozqIuG8NH_vuBRPqaN5zL0rvRFp5ftwReAhZXBP44,11342
+diffct/__init__.py,sha256=m5Oy9KZOD9hqS_vMAk-FTd2BoMUqACyyitYrDAEVhFA,2965
+diffct/_solver_state.py,sha256=ExWTTGftGOsPo9MQwxv-ylmYvhiPYXUo3Tmhi6VEvg0,16606
+diffct/_solver_tv.py,sha256=2x16Bq5Bx9WCBszgPy3K-wicsH9AeuzzJS82uyiHBkE,5847
+diffct/analytical.py,sha256=YTyI_lwgnhHdoglO9fboshfNQtrZsacK-gra63ja-xE,33130
+diffct/analytical_execution.py,sha256=Al452qBRevGcEpwcdf7LpY8XbgZzgJ4YoupyRD_1LWA,26788
+diffct/backends.py,sha256=zNNjO_sDJ8KzstTmP0WqPeUYIJDBeYg7m5KD5zQ8ZFw,48278
+diffct/chunking.py,sha256=jJo7JlR8GrQz-_Rj44_HQLsNJeNgL1SpleJsL4QKJvE,11373
 diffct/constants.py,sha256=oy3QNC6an43Ogdn1m81znz1-EMp9GeinVf2wkGYiXfE,3160
 diffct/differentiable.py,sha256=cFirhqXRpUemL7y7Sl_0uqZKGJNj__ptQEcTwYNKzys,1879
 diffct/distribution.py,sha256=iLC3BgAQ9lOkndbEOvNB_j_UaA2Ug9w1KiTUHkOzG8o,9982
 diffct/geometry.py,sha256=iwWXQksIJGOPSuNWJ8tImWj6yPZCC71G3RBh2VDVylY,42581
-diffct/operators.py,sha256=QH7TedxIc35616bn0PwgaLDllH3Xt3aadFzIpbe8y0M,158446
+diffct/operators.py,sha256=yGDie8Kvjq8wkBU3hnIKUuKWR9VxgkYFFMINE03a0LY,167009
 diffct/pipeline.py,sha256=jp7AvDeQQOME3OPkO9XTcjWFLvOLrv8wczhAf1jcklM,10891
 diffct/projectors.py,sha256=0oAFSA5mycN59dayTAWK2W7o96MP8oLuVp2V9X-1388,45334
 diffct/reconstruction.py,sha256=KdDrpxQ31ufyUlL_bkzvDzaTcrK4BB97tnJ04ZfVfos,29368
-diffct/solvers.py,sha256=ZroVUlhHOE2IAXB0clw9hUCvomrFSgYdGRQWwnSFZxs,17297
+diffct/solvers.py,sha256=B7M4akFUTuwb_MBpQbpZLdlMv5Eyxy3WczExDAs613s,18000
 diffct/storage.py,sha256=EdcFTzRtEAPRvAZhRPrSf8gRNF5o_JVbPLE8eKPhNdY,4153
 diffct/streaming.py,sha256=2LB5xpl6vIJE1kXMBtQbG0OVxlutEEqXLBrcvlzBHl8,6086
 diffct/surface_gradients.py,sha256=LeM_eDY1Alv604SdiPUx2CWFzr-J9NqCEcJ_ovVFGa8,6290
-diffct/surfaces.py,sha256=z_0O98H7eMdSd9qrgUtcg5hBaE643s6sG_AjnH0Dyhw,7935
-diffct/utils.py,sha256=cih3GyZ33gnOrTW7O8r2_aXIP2a-rFjOMuAPuvHeQ4I,14794
+diffct/surfaces.py,sha256=JQ-PyMnKTTtK8VUykJ3kvIcquQHeDThVHQ18umlOYQg,8009
+diffct/utils.py,sha256=ZOZpR_qpVSxv_uSQJ6--x-RXAqUrZJjWXYJoPKhu5nY,15240
 diffct/kernels/__init__.py,sha256=YNdsqAkexN25pI7QaHy4phAnQNICQ7tO8K8S1zt7TnI,1899
 diffct/kernels/analytical_transpose.py,sha256=D4vr_ZFPrjFeQyhoDqTyijbkTY-g-sHJet52ho-cbkQ,9187
 diffct/kernels/cone_beam.py,sha256=jKlbqQPHhrng-pTbyMFx2b2l0pdQIMiaaB5hj6ayPQI,26753
 diffct/kernels/fan_beam.py,sha256=NPha0Wl-_ayOdS6kAoXJngQffsUczp3rRnojXZ-F1y8,17556
 diffct/kernels/geometry_vjp.py,sha256=YiPMdsuXINLu1oFhPcSbG1LW2P3WSfeJeEvHNLrmQzU,22008
 diffct/kernels/parallel_beam.py,sha256=tCiSx9W8nO4mn8DVyRuK4E0Ry8q38g3pAUACab0UYVU,17143
-diffct/kernels/surface_geometry.py,sha256=eXD_aqW0e6vBMsS0R5oZ-cyedK_wrX7RDjFfiM0HIg4,5602
-diffct-2.4.2.dist-info/METADATA,sha256=0NsA9YkbCMTwARCVFc1RCqY53W8-zwGjK6ZNOlbcTgU,17661
-diffct-2.4.2.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
-diffct-2.4.2.dist-info/licenses/LICENSE,sha256=a1fjzGs44LpDnV_LmA2Tzp8z-eJKoNVVdLOnh-F6J6E,11344
-diffct-2.4.2.dist-info/licenses/NOTICE,sha256=du9I2sDF5IDfiY9vF64WWzMJh-MarpVXyLfWQmhbvmI,714
-diffct-2.4.2.dist-info/RECORD,,
+diffct/kernels/surface_geometry.py,sha256=W0G0ilIJJ_EotU4SELeppeWrzokzFWJqpm6qqKaWN7Y,5842
+diffct-2.4.3.dist-info/METADATA,sha256=PeKQAUhUPwodcv-NYQahBD1mceTa5w7WPQDvBDXmolE,17612
+diffct-2.4.3.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+diffct-2.4.3.dist-info/licenses/LICENSE,sha256=a1fjzGs44LpDnV_LmA2Tzp8z-eJKoNVVdLOnh-F6J6E,11344
+diffct-2.4.3.dist-info/licenses/NOTICE,sha256=du9I2sDF5IDfiY9vF64WWzMJh-MarpVXyLfWQmhbvmI,714
+diffct-2.4.3.dist-info/RECORD,,