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

diffct/__init__.py

@@ -64,7 +64,7 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.4.2'
+__version__ = '2.4.4'
 
 __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

@@ -27,7 +27,7 @@ import math
 
 import torch
 
-from .constants import _DTYPE
+from .constants import _DTYPE, _native_spacing
 from .utils import (
     DeviceManager,
     TorchCUDABridge,
@@ -431,7 +431,7 @@ def _parallel_weighted_backproject_raw(sinogram, ray_dir, det_origin, det_u_vec,
     _resolve_kernel('_parallel_2d_fbp_backproject_kernel')[grid, tpb, numba_stream](
         d_sino, n_views, n_det, d_reco, Nx, Ny,
         _DTYPE(detector_spacing), d_ray_dir, d_det_origin, d_det_u_vec,
-        cx, cy, _DTYPE(voxel_spacing),
+        cx, cy, _native_spacing(voxel_spacing),
     )
 
     return reco * (1.0 / (2.0 * math.pi))
@@ -522,9 +522,9 @@ def _fan_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec,
     _resolve_kernel('_fan_2d_fbp_backproject_kernel')[grid, tpb, numba_stream](
         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),
+        cx, cy, _native_spacing(voxel_spacing),
+        *(_DTYPE(value if isinstance(voxel_spacing, (tuple, list)) else value / voxel_spacing)
+          for value in isocenter.tolist()),
     )
 
     return reco
@@ -585,10 +585,9 @@ def _cone_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec, det
         d_sino, n_views, n_u, n_v, d_reco, Nx, Ny, Nz,
         _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),
+        cx, cy, cz, _native_spacing(voxel_spacing),
+        *(_DTYPE(value if isinstance(voxel_spacing, (tuple, list)) else value / voxel_spacing)
+          for value in isocenter.tolist()),
     )
 
     # Convert back to (D, H, W) layout for the user.
@@ -640,7 +639,8 @@ 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 if isinstance(parameters[-1], (tuple, list)) else value / parameters[-1])
+                    for value in isocenter.tolist())
 
     if beam == "cone":
         D, H, W, du, dv, spacing = parameters
@@ -650,7 +650,7 @@ def _analytical_transpose_raw(image, beam, geometry, parameters, isocenter, sino
         _resolve_kernel('_cone_3d_fdk_transpose_kernel')[grid, tpb, stream](
             d_sino, *sinogram_shape, d_image, W, H, D, _DTYPE(du), _DTYPE(dv),
             *d_geometry, _DTYPE(W * 0.5), _DTYPE(H * 0.5), _DTYPE(D * 0.5),
-            _DTYPE(spacing), *iso,
+            _native_spacing(spacing), *iso,
         )
     else:
         pitch, H, W, spacing = parameters
@@ -660,7 +660,7 @@ def _analytical_transpose_raw(image, beam, geometry, parameters, isocenter, sino
                                  else '_fan_2d_fbp_transpose_kernel')
         kernel[grid, tpb, stream](
             d_sino, *sinogram_shape, d_image, W, H, _DTYPE(pitch), *d_geometry,
-            _DTYPE(W * 0.5), _DTYPE(H * 0.5), _DTYPE(spacing), *iso,
+            _DTYPE(W * 0.5), _DTYPE(H * 0.5), _native_spacing(spacing), *iso,
         )
     return sinogram * scale
 

diffct/analytical_execution.py

@@ -16,6 +16,7 @@ from .analytical import (
     fan_weighted_backproject, parallel_weighted_backproject,
 )
 from .backends import _zero_dependency, select_backend, voxel_spacing_tuple
+from .constants import _native_spacing
 from .chunking import _memory_budget, _pixel_tiles
 from .distribution import Partition, preflight
 from .operators import _detector_spacing, _normalize_devices, _positive_int, _shape, _validate_trajectory
@@ -53,6 +54,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)]
@@ -141,6 +144,9 @@ class AnalyticalExecution:
         self.rank = 3 if beam == "cone" else 2
         self.global_volume_shape = _shape(volume_shape, self.rank, "volume_shape")
         self.spacing = voxel_spacing_tuple(voxel_spacing, self.rank)
+        self.voxel_spacing = self.spacing if isinstance(voxel_spacing, (tuple, list)) else self.spacing[0]
+        if self.backend == "numba" and isinstance(self.voxel_spacing, tuple):
+            _native_spacing(self.voxel_spacing)
         self.pitch = _detector_spacing(detector_spacing, beam)
         self.devices = _normalize_devices(devices)
         if not isinstance(source, torch.Tensor):
@@ -334,14 +340,14 @@ class AnalyticalExecution:
                               device=data.device, dtype=self.dtype)
         if self.beam == "parallel":
             geometry = (geometry[0], geometry[1] - centre, geometry[2])
-            return parallel_weighted_backproject(data, *geometry, self.pitch[0], *shape, self.spacing[0])
+            return parallel_weighted_backproject(data, *geometry, self.pitch[0], *shape, self.voxel_spacing)
         geometry = (geometry[0] - centre, geometry[1] - centre, *geometry[2:])
         isocenter = self.isocenter.to(data.device) - centre
         if self.beam == "fan":
             return fan_weighted_backproject(data, *geometry, self.pitch[0], *shape,
-                                            self.spacing[0], isocenter=isocenter)
+                                            self.voxel_spacing, isocenter=isocenter)
         return cone_weighted_backproject(data, *geometry, *shape, *self.pitch,
-                                         self.spacing[0], isocenter=isocenter)
+                                         self.voxel_spacing, isocenter=isocenter)
 
     def _tile(self, spatial):
         shape = tuple(part.stop - part.start for part in spatial)
@@ -385,6 +391,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.