diffct 2.0.2__py3-none-any.whl → 2.2.0__py3-none-any.whl

diffct/__init__.py

@@ -15,6 +15,8 @@ from .projectors import (
 )
 
 from .operators import Projector
+from .storage import TensorStore, NpyStore
+from .surfaces import ParameterizedSurface
 
 from .geometry import (
     circular_trajectory_3d,
@@ -49,10 +51,13 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.0.2'
+__version__ = '2.2.0'
 
 __all__ = [
     'Projector',
+    'TensorStore',
+    'NpyStore',
+    'ParameterizedSurface',
     'ParallelProjectorFunction',
     'ParallelBackprojectorFunction',
     'FanProjectorFunction',

diffct/chunking.py

@@ -0,0 +1,201 @@
+"""Spatial slices and immutable execution limits for streamed projectors."""
+
+from dataclasses import dataclass
+from itertools import product
+import math
+
+import torch
+
+
+@dataclass(frozen=True)
+class _ChunkPlan:
+    chunk_shape: tuple
+    view_chunk_size: int
+    devices: tuple
+    detector_chunk_shape: tuple = None
+    gpu_budget_bytes: int = 0
+    allocation_limit: int = (1 << 63) - 1
+    schedule: str = "spatial"
+    window_size: int = 1
+    slots: int = 1
+    estimated_gpu_bytes: int = 0
+    fallback_reason: str = ""
+    foreign_output: bool = False
+    reservation_budget_bytes: int = (1 << 63) - 1
+    estimated_reserved_bytes: int = 0
+
+
+def _pixel_tiles(shape, chunk_shape):
+    for starts in product(*(range(0, size, part) for size, part in zip(shape, chunk_shape))):
+        clipped = tuple(min(part, size - start) for size, start, part in zip(shape, starts, chunk_shape))
+        yield tuple(slice(start, start + size) for start, size in zip(starts, clipped)), clipped
+
+
+def _memory_budget(device):
+    """Usable allocator bytes, retaining 25% headroom for allocation overhead.
+
+    Reusable reserved blocks count as available, but never bypass the process
+    fraction ceiling. Older PyTorch without its public fraction getter can
+    only observe driver/allocator memory; explicit limits remain available.
+    """
+    free, total = torch.cuda.mem_get_info(device)
+    allocated = torch.cuda.memory_allocated(device)
+    reusable = max(0, torch.cuda.memory_reserved(device) - allocated)
+    getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
+    fraction = getter(device) if getter is not None else 1.0
+    allocator_room = max(0, int(total * fraction) - allocated)
+    backend = getattr(torch.cuda, "get_allocator_backend", lambda: "native")()
+    foreign_splits = 0
+    if backend == "native":
+        caller = torch.cuda.current_stream(device).cuda_stream
+        for segment in _native_segments(device):
+            blocks = segment["blocks"]
+            if segment["stream"] != caller and any(block["state"] != "inactive" for block in blocks):
+                foreign_splits += sum(block["size"] for block in blocks if block["state"] == "inactive")
+    available = max(0, min(free + reusable, allocator_room) - foreign_splits)
+    budget = available * 3 // 4
+    allocation_limit = (1 << 63) - 1  # Unbounded limit in scalar int64 agreement.
+    # Native requests >1 MiB and <10 MiB reserve a 20 MiB segment. A smaller
+    # current available room must use the <=1 MiB small pool, whose segments are 2 MiB.
+    # Source: pytorch/c10/cuda/CUDACachingAllocator.cpp (Native allocator notes).
+    if backend == "native" and available < 20 * 1024 ** 2:
+        allocation_limit = 1024 ** 2
+    return budget, allocation_limit
+
+
+def _native_segments(device):
+    index = torch.device(device).index if not isinstance(device, int) else device
+    index = torch.cuda.current_device() if index is None else index
+    return [segment for segment in torch.cuda.memory_snapshot() if segment["device"] == index]
+
+
+def _reservation_budget(device):
+    """Room for fresh segments, conservatively excluding every pinned segment.
+
+    Free splits in a live segment may serve the caller, but cannot guarantee
+    placement of the complete new inventory. Wholly inactive segments can be
+    released by the native allocator, including those owned by other streams.
+    """
+    if torch.cuda.get_allocator_backend() != "native":
+        return (1 << 63) - 1
+    free, total = torch.cuda.mem_get_info(device)
+    segments = _native_segments(device)
+    pinned = sum(segment["total_size"] for segment in segments
+                 if any(block["state"] != "inactive" for block in segment["blocks"]))
+    reclaimable = sum(segment["total_size"] for segment in segments
+                      if all(block["state"] == "inactive" for block in segment["blocks"]))
+    getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
+    fraction = getter(device) if getter is not None else 1.0
+    return max(0, min(free + reclaimable, int(total * fraction) - pinned))
+
+
+def _working_set_bytes(beam, shape, views, detector_shape, surface):
+    """Conservative native staging/layout, cotangent, geometry and VJP scratch."""
+    rank = len(shape)
+    voxels = math.prod(shape)
+    rays = views * math.prod(detector_shape)
+    geometry = (rank * (views + rays) if surface
+                else rank * (4 if beam == "cone" else 3) * views)
+    # A resident native tile, staging and completed cone layout can coexist.
+    # Two ray buffers cover reusable output/cotangent and transfer scratch; two
+    # geometry buffers cover staged values and their first-order VJP/reduction.
+    tile_buffers = 3 if beam == "cone" else 2
+    return 4 * (tile_buffers * voxels + 2 * rays + 2 * geometry)
+
+
+def _largest_buffer_bytes(beam, shape, views, detector_shape, surface):
+    rank = len(shape)
+    rays = views * math.prod(detector_shape)
+    geometry = (rank * (views + rays) if surface
+                else rank * (4 if beam == "cone" else 3) * views)
+    return 4 * max(math.prod(shape), rays, geometry)
+
+
+def _allocation_bytes(elements):
+    size = 4 * elements
+    if torch.cuda.get_allocator_backend() == "native":
+        return (size + 511) // 512 * 512
+    return size
+
+
+def _reservation_bytes(buffer_sizes):
+    """Fresh, monotone fixed-buffer inventory bound for standard native pools.
+
+    Arbitrary free holes and allocator overrides are outside this model. Large
+    and middle requests share one pool, so mixed inventories use separate fresh
+    segments per request instead of assuming their placement is independent.
+    """
+    sizes = [size for size in buffer_sizes if size > 0]
+    if torch.cuda.get_allocator_backend() != "native":
+        return sum(sizes)
+    mib = 1024 ** 2
+    sizes = [(size + 511) // 512 * 512 for size in sizes]
+    large = [size for size in sizes if size >= 10 * mib]
+    reservation = sum((size + 2 * mib - 1) // (2 * mib) * (2 * mib) for size in large)
+    middle = [size for size in sizes if mib < size < 10 * mib]
+    if large:
+        reservation += len(middle) * 20 * mib
+        middle = []
+    for segment_size, requests in (
+        (2 * mib, [size for size in sizes if size <= mib]),
+        (20 * mib, middle),
+    ):
+        if not requests:
+            continue
+        total = sum(requests)
+        # At a new segment's birth every older segment holds more than
+        # segment_size - max(requests); the new one holds at least min(requests).
+        # This bounds online best-fit placement in any allocation order.
+        count = (1 if total <= segment_size else
+                 (total - min(requests)) // (segment_size - max(requests)) + 1)
+        reservation += count * segment_size
+    return reservation
+
+
+def _pipeline_buffers(beam, shape, views, pixels, points, *, slots=1, retained=1,
+                      geometry_grad=False, peers=False, ray_accumulators=0):
+    """Live typed allocation inventory, with native rounding per buffer."""
+    rank, rays = len(shape), views * math.prod(pixels)
+    native = _allocation_bytes(math.prod(shape))
+    # Two producer epochs permit the preceding native/layout/output storage
+    # to stay in flight while the next retained tile window queues work.
+    tiles = slots * (retained + int(beam == "cone")) + int(peers)
+    sizes = ((views * rank, rays * rank) if points
+             else (views * rank,) * (4 if beam == "cone" else 3))
+    geometry = tuple(_allocation_bytes(size) for size in sizes)
+    return ((native,) * tiles + (_allocation_bytes(rays),) * (slots + ray_accumulators)
+            + geometry * (slots * (1 + int(geometry_grad))))
+
+
+def _pipeline_bytes(beam, shape, views, pixels, points, *, slots=1, retained=1,
+                    geometry_grad=False, peers=False, ray_accumulators=0):
+    """Live typed bytes; segment reservation is admitted separately."""
+    rank, rays = len(shape), views * math.prod(pixels)
+    native = _allocation_bytes(math.prod(shape))
+    tile = native * (slots * (retained + int(beam == "cone")) + int(peers))
+    sizes = ((views * rank, rays * rank) if points
+             else (views * rank,) * (4 if beam == "cone" else 3))
+    geometry = sum(_allocation_bytes(size) for size in sizes)
+    return (tile + slots * (_allocation_bytes(rays) + geometry)
+            + ray_accumulators * _allocation_bytes(rays) + (slots * geometry if geometry_grad else 0))
+
+
+def _host_geometry_bytes(beam, views, pixels, points):
+    rank = 3 if beam == "cone" else 2
+    if not points:
+        return 8 * 2 * rank * views * (4 if beam == "cone" else 3)
+    # Grids/offsets, world expressions and first-order cotangent/graph work.
+    # Arbitrary sampler internals and explicit parameter snapshots are separate.
+    return 8 * rank * (math.prod(pixels) * (1 + 6 * views) + 4 * views)
+
+
+def _tiles(shape, chunk_shape, spacing):
+    """Yield tensor slices, clipped shapes and physical centres in xyz order."""
+    starts = (range(0, size, chunk) for size, chunk in zip(shape, chunk_shape))
+    for start in product(*starts):
+        tile_shape = tuple(min(chunk, size - first)
+                           for size, chunk, first in zip(shape, chunk_shape, start))
+        slices = tuple(slice(first, first + size) for first, size in zip(start, tile_shape))
+        centre = tuple((first + size / 2 - full / 2) * spacing
+                       for first, size, full in zip(start, tile_shape, shape))
+        yield slices, tile_shape, centre[::-1]

diffct/kernels/cone_beam.py

@@ -30,7 +30,7 @@ def _cone_3d_forward_kernel(
     d_vol, Nx, Ny, Nz,
     d_sino, n_views, n_u, n_v,
     du, dv, d_src_pos, d_det_center, d_det_u_vec, d_det_v_vec,
-    cx, cy, cz, voxel_spacing
+    cx, cy, cz, voxel_spacing, d_detector_positions=None,
 ):
     """Compute the 3D cone-beam forward projection with arbitrary source-detector trajectories.
 
@@ -93,27 +93,32 @@ def _cone_3d_forward_kernel(
     src_y = d_src_pos[iview, 1] / voxel_spacing
     src_z = d_src_pos[iview, 2] / voxel_spacing
 
-    # Read detector center and orientation vectors
-    det_cx = d_det_center[iview, 0] / voxel_spacing
-    det_cy = d_det_center[iview, 1] / voxel_spacing
-    det_cz = d_det_center[iview, 2] / voxel_spacing
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        # Read detector center and orientation vectors
+        det_cx = d_det_center[iview, 0] / voxel_spacing
+        det_cy = d_det_center[iview, 1] / voxel_spacing
+        det_cz = d_det_center[iview, 2] / voxel_spacing
 
-    u_vec_x = d_det_u_vec[iview, 0]
-    u_vec_y = d_det_u_vec[iview, 1]
-    u_vec_z = d_det_u_vec[iview, 2]
+        u_vec_x = d_det_u_vec[iview, 0]
+        u_vec_y = d_det_u_vec[iview, 1]
+        u_vec_z = d_det_u_vec[iview, 2]
 
-    v_vec_x = d_det_v_vec[iview, 0]
-    v_vec_y = d_det_v_vec[iview, 1]
-    v_vec_z = d_det_v_vec[iview, 2]
+        v_vec_x = d_det_v_vec[iview, 0]
+        v_vec_y = d_det_v_vec[iview, 1]
+        v_vec_z = d_det_v_vec[iview, 2]
 
-    # Calculate detector element offset from center
-    u_offset = (np.float32(iu) + _HALF - np.float32(n_u) * _HALF) * du / voxel_spacing
-    v_offset = (np.float32(iv) + _HALF - np.float32(n_v) * _HALF) * dv / voxel_spacing
+        # Calculate detector element offset from center
+        u_offset = (np.float32(iu) + _HALF - np.float32(n_u) * _HALF) * du / voxel_spacing
+        v_offset = (np.float32(iv) + _HALF - np.float32(n_v) * _HALF) * dv / voxel_spacing
 
-    # Calculate 3D detector element position using center + u*u_vec + v*v_vec
-    det_x = det_cx + u_offset * u_vec_x + v_offset * v_vec_x
-    det_y = det_cy + u_offset * u_vec_y + v_offset * v_vec_y
-    det_z = det_cz + u_offset * u_vec_z + v_offset * v_vec_z
+        # Calculate 3D detector element position using center + u*u_vec + v*v_vec
+        det_x = det_cx + u_offset * u_vec_x + v_offset * v_vec_x
+        det_y = det_cy + u_offset * u_vec_y + v_offset * v_vec_y
+        det_z = det_cz + u_offset * u_vec_z + v_offset * v_vec_z
+    else:
+        det_x = d_detector_positions[iview, iu, iv, 0] / voxel_spacing
+        det_y = d_detector_positions[iview, iu, iv, 1] / voxel_spacing
+        det_z = d_detector_positions[iview, iu, iv, 2] / voxel_spacing
 
     # === 3D RAY DIRECTION CALCULATION ===
     # Ray direction vector from source to detector element in 3D space
@@ -238,7 +243,7 @@ def _cone_3d_backward_kernel(
     d_sino, n_views, n_u, n_v,
     d_vol, Nx, Ny, Nz,
     du, dv, d_src_pos, d_det_center, d_det_u_vec, d_det_v_vec,
-    cx, cy, cz, voxel_spacing
+    cx, cy, cz, voxel_spacing, d_detector_positions=None,
 ):
     """Compute the 3D cone-beam backprojection with arbitrary source-detector trajectories.
 
@@ -305,27 +310,32 @@ def _cone_3d_backward_kernel(
     src_y = d_src_pos[iview, 1] / voxel_spacing
     src_z = d_src_pos[iview, 2] / voxel_spacing
 
-    # Read detector center and orientation vectors
-    det_cx = d_det_center[iview, 0] / voxel_spacing
-    det_cy = d_det_center[iview, 1] / voxel_spacing
-    det_cz = d_det_center[iview, 2] / voxel_spacing
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        # Read detector center and orientation vectors
+        det_cx = d_det_center[iview, 0] / voxel_spacing
+        det_cy = d_det_center[iview, 1] / voxel_spacing
+        det_cz = d_det_center[iview, 2] / voxel_spacing
 
-    u_vec_x = d_det_u_vec[iview, 0]
-    u_vec_y = d_det_u_vec[iview, 1]
-    u_vec_z = d_det_u_vec[iview, 2]
+        u_vec_x = d_det_u_vec[iview, 0]
+        u_vec_y = d_det_u_vec[iview, 1]
+        u_vec_z = d_det_u_vec[iview, 2]
 
-    v_vec_x = d_det_v_vec[iview, 0]
-    v_vec_y = d_det_v_vec[iview, 1]
-    v_vec_z = d_det_v_vec[iview, 2]
+        v_vec_x = d_det_v_vec[iview, 0]
+        v_vec_y = d_det_v_vec[iview, 1]
+        v_vec_z = d_det_v_vec[iview, 2]
 
-    # Calculate detector element offset from center
-    u_offset = (np.float32(iu) + _HALF - np.float32(n_u) * _HALF) * du / voxel_spacing
-    v_offset = (np.float32(iv) + _HALF - np.float32(n_v) * _HALF) * dv / voxel_spacing
+        # Calculate detector element offset from center
+        u_offset = (np.float32(iu) + _HALF - np.float32(n_u) * _HALF) * du / voxel_spacing
+        v_offset = (np.float32(iv) + _HALF - np.float32(n_v) * _HALF) * dv / voxel_spacing
 
-    # Calculate 3D detector element position using center + u*u_vec + v*v_vec
-    det_x = det_cx + u_offset * u_vec_x + v_offset * v_vec_x
-    det_y = det_cy + u_offset * u_vec_y + v_offset * v_vec_y
-    det_z = det_cz + u_offset * u_vec_z + v_offset * v_vec_z
+        # Calculate 3D detector element position using center + u*u_vec + v*v_vec
+        det_x = det_cx + u_offset * u_vec_x + v_offset * v_vec_x
+        det_y = det_cy + u_offset * u_vec_y + v_offset * v_vec_y
+        det_z = det_cz + u_offset * u_vec_z + v_offset * v_vec_z
+    else:
+        det_x = d_detector_positions[iview, iu, iv, 0] / voxel_spacing
+        det_y = d_detector_positions[iview, iu, iv, 1] / voxel_spacing
+        det_z = d_detector_positions[iview, iu, iv, 2] / voxel_spacing
 
     # === 3D RAY DIRECTION CALCULATION ===
     # Ray direction vector from source to detector element in 3D space

diffct/kernels/fan_beam.py

@@ -30,7 +30,7 @@ def _fan_2d_forward_kernel(
     d_image, Nx, Ny,
     d_sino, n_ang, n_det,
     det_spacing, d_src_pos, d_det_center, d_det_u_vec,
-    cx, cy, voxel_spacing
+    cx, cy, voxel_spacing, d_detector_positions=None,
 ):
     """Compute the 2D fan beam forward projection with arbitrary source-detector trajectories.
 
@@ -82,19 +82,23 @@ def _fan_2d_forward_kernel(
     src_x = d_src_pos[iang, 0] / voxel_spacing
     src_y = d_src_pos[iang, 1] / voxel_spacing
 
-    # Read detector center and orientation vector
-    det_cx = d_det_center[iang, 0] / voxel_spacing
-    det_cy = d_det_center[iang, 1] / voxel_spacing
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        # Read detector center and orientation vector
+        det_cx = d_det_center[iang, 0] / voxel_spacing
+        det_cy = d_det_center[iang, 1] / voxel_spacing
 
-    u_vec_x = d_det_u_vec[iang, 0]
-    u_vec_y = d_det_u_vec[iang, 1]
+        u_vec_x = d_det_u_vec[iang, 0]
+        u_vec_y = d_det_u_vec[iang, 1]
 
-    # Calculate detector element offset from center
-    u_offset = (np.float32(idet) + _HALF - np.float32(n_det) * _HALF) * det_spacing / voxel_spacing
+        # Calculate detector element offset from center
+        u_offset = (np.float32(idet) + _HALF - np.float32(n_det) * _HALF) * det_spacing / voxel_spacing
 
-    # Calculate 2D detector element position using center + u*u_vec
-    det_x = det_cx + u_offset * u_vec_x
-    det_y = det_cy + u_offset * u_vec_y
+        # Calculate 2D detector element position using center + u*u_vec
+        det_x = det_cx + u_offset * u_vec_x
+        det_y = det_cy + u_offset * u_vec_y
+    else:
+        det_x = d_detector_positions[iang, idet, 0] / voxel_spacing
+        det_y = d_detector_positions[iang, idet, 1] / voxel_spacing
 
     # === RAY DIRECTION CALCULATION ===
     # Ray direction vector from source to detector element
@@ -192,7 +196,7 @@ def _fan_2d_backward_kernel(
     d_sino, n_ang, n_det,
     d_image, Nx, Ny,
     det_spacing, d_src_pos, d_det_center, d_det_u_vec,
-    cx, cy, voxel_spacing
+    cx, cy, voxel_spacing, d_detector_positions=None,
 ):
     """Compute the 2D fan beam backprojection with arbitrary source-detector trajectories.
 
@@ -248,19 +252,23 @@ def _fan_2d_backward_kernel(
     src_x = d_src_pos[iang, 0] / voxel_spacing
     src_y = d_src_pos[iang, 1] / voxel_spacing
 
-    # Read detector center and orientation vector
-    det_cx = d_det_center[iang, 0] / voxel_spacing
-    det_cy = d_det_center[iang, 1] / voxel_spacing
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        # Read detector center and orientation vector
+        det_cx = d_det_center[iang, 0] / voxel_spacing
+        det_cy = d_det_center[iang, 1] / voxel_spacing
 
-    u_vec_x = d_det_u_vec[iang, 0]
-    u_vec_y = d_det_u_vec[iang, 1]
+        u_vec_x = d_det_u_vec[iang, 0]
+        u_vec_y = d_det_u_vec[iang, 1]
 
-    # Calculate detector element offset from center
-    u_offset = (np.float32(idet) + _HALF - np.float32(n_det) * _HALF) * det_spacing / voxel_spacing
+        # Calculate detector element offset from center
+        u_offset = (np.float32(idet) + _HALF - np.float32(n_det) * _HALF) * det_spacing / voxel_spacing
 
-    # Calculate 2D detector element position using center + u*u_vec
-    det_x = det_cx + u_offset * u_vec_x
-    det_y = det_cy + u_offset * u_vec_y
+        # Calculate 2D detector element position using center + u*u_vec
+        det_x = det_cx + u_offset * u_vec_x
+        det_y = det_cy + u_offset * u_vec_y
+    else:
+        det_x = d_detector_positions[iang, idet, 0] / voxel_spacing
+        det_y = d_detector_positions[iang, idet, 1] / voxel_spacing
 
     # === RAY DIRECTION CALCULATION ===
     # Ray direction vector from source to detector element

diffct/kernels/geometry_vjp.py

@@ -58,6 +58,7 @@ def _fan_2d_geometry_vjp_kernel(
     det_spacing, d_src_pos, d_det_center, d_det_u_vec,
     cx, cy, voxel_spacing,
     d_grad_src, d_grad_det_center, d_grad_det_u,
+    d_detector_positions=None, d_grad_detector_positions=None,
 ):
     """Accumulate fan-beam geometry gradients of ``<grad_sino, A x>`` per view."""
     iang, idet = cuda.grid(2)
@@ -71,14 +72,18 @@ def _fan_2d_geometry_vjp_kernel(
     # the entry and the exit (-1 = segment end at the source or detector).
     src_x = d_src_pos[iang, 0] / voxel_spacing
     src_y = d_src_pos[iang, 1] / voxel_spacing
-    det_cx = d_det_center[iang, 0] / voxel_spacing
-    det_cy = d_det_center[iang, 1] / voxel_spacing
-    u_vec_x = d_det_u_vec[iang, 0]
-    u_vec_y = d_det_u_vec[iang, 1]
-    u_phys = (np.float32(idet) + _HALF - np.float32(n_det) * _HALF) * det_spacing
-    u_offset = u_phys / voxel_spacing
-    det_x = det_cx + u_offset * u_vec_x
-    det_y = det_cy + u_offset * u_vec_y
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        det_cx = d_det_center[iang, 0] / voxel_spacing
+        det_cy = d_det_center[iang, 1] / voxel_spacing
+        u_vec_x = d_det_u_vec[iang, 0]
+        u_vec_y = d_det_u_vec[iang, 1]
+        u_phys = (np.float32(idet) + _HALF - np.float32(n_det) * _HALF) * det_spacing
+        u_offset = u_phys / voxel_spacing
+        det_x = det_cx + u_offset * u_vec_x
+        det_y = det_cy + u_offset * u_vec_y
+    else:
+        det_x = d_detector_positions[iang, idet, 0] / voxel_spacing
+        det_y = d_detector_positions[iang, idet, 1] / voxel_spacing
 
     dir_x, dir_y = det_x - src_x, det_y - src_y
     length = math.sqrt(dir_x * dir_x + dir_y * dir_y)
@@ -203,10 +208,14 @@ def _fan_2d_geometry_vjp_kernel(
 
     cuda.atomic.add(d_grad_src, (iang, 0), g * ga_x)
     cuda.atomic.add(d_grad_src, (iang, 1), g * ga_y)
-    cuda.atomic.add(d_grad_det_center, (iang, 0), g * gb_x)
-    cuda.atomic.add(d_grad_det_center, (iang, 1), g * gb_y)
-    cuda.atomic.add(d_grad_det_u, (iang, 0), g * u_phys * gb_x)
-    cuda.atomic.add(d_grad_det_u, (iang, 1), g * u_phys * gb_y)
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        cuda.atomic.add(d_grad_det_center, (iang, 0), g * gb_x)
+        cuda.atomic.add(d_grad_det_center, (iang, 1), g * gb_y)
+        cuda.atomic.add(d_grad_det_u, (iang, 0), g * u_phys * gb_x)
+        cuda.atomic.add(d_grad_det_u, (iang, 1), g * u_phys * gb_y)
+    else:
+        d_grad_detector_positions[iang, idet, 0] = g * gb_x
+        d_grad_detector_positions[iang, idet, 1] = g * gb_y
 
 
 @_FASTMATH_DECORATOR
@@ -216,6 +225,7 @@ def _cone_3d_geometry_vjp_kernel(
     du, dv, d_src_pos, d_det_center, d_det_u_vec, d_det_v_vec,
     cx, cy, cz, voxel_spacing,
     d_grad_src, d_grad_det_center, d_grad_det_u, d_grad_det_v,
+    d_detector_positions=None, d_grad_detector_positions=None,
 ):
     """Accumulate cone-beam geometry gradients of ``<grad_sino, A x>`` per view."""
     iv, iu, iview = cuda.grid(3)
@@ -228,22 +238,27 @@ def _cone_3d_geometry_vjp_kernel(
     src_x = d_src_pos[iview, 0] / voxel_spacing
     src_y = d_src_pos[iview, 1] / voxel_spacing
     src_z = d_src_pos[iview, 2] / voxel_spacing
-    det_cx = d_det_center[iview, 0] / voxel_spacing
-    det_cy = d_det_center[iview, 1] / voxel_spacing
-    det_cz = d_det_center[iview, 2] / voxel_spacing
-    u_vec_x = d_det_u_vec[iview, 0]
-    u_vec_y = d_det_u_vec[iview, 1]
-    u_vec_z = d_det_u_vec[iview, 2]
-    v_vec_x = d_det_v_vec[iview, 0]
-    v_vec_y = d_det_v_vec[iview, 1]
-    v_vec_z = d_det_v_vec[iview, 2]
-    u_phys = (np.float32(iu) + _HALF - np.float32(n_u) * _HALF) * du
-    v_phys = (np.float32(iv) + _HALF - np.float32(n_v) * _HALF) * dv
-    u_offset = u_phys / voxel_spacing
-    v_offset = v_phys / voxel_spacing
-    det_x = det_cx + u_offset * u_vec_x + v_offset * v_vec_x
-    det_y = det_cy + u_offset * u_vec_y + v_offset * v_vec_y
-    det_z = det_cz + u_offset * u_vec_z + v_offset * v_vec_z
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        det_cx = d_det_center[iview, 0] / voxel_spacing
+        det_cy = d_det_center[iview, 1] / voxel_spacing
+        det_cz = d_det_center[iview, 2] / voxel_spacing
+        u_vec_x = d_det_u_vec[iview, 0]
+        u_vec_y = d_det_u_vec[iview, 1]
+        u_vec_z = d_det_u_vec[iview, 2]
+        v_vec_x = d_det_v_vec[iview, 0]
+        v_vec_y = d_det_v_vec[iview, 1]
+        v_vec_z = d_det_v_vec[iview, 2]
+        u_phys = (np.float32(iu) + _HALF - np.float32(n_u) * _HALF) * du
+        v_phys = (np.float32(iv) + _HALF - np.float32(n_v) * _HALF) * dv
+        u_offset = u_phys / voxel_spacing
+        v_offset = v_phys / voxel_spacing
+        det_x = det_cx + u_offset * u_vec_x + v_offset * v_vec_x
+        det_y = det_cy + u_offset * u_vec_y + v_offset * v_vec_y
+        det_z = det_cz + u_offset * u_vec_z + v_offset * v_vec_z
+    else:
+        det_x = d_detector_positions[iview, iu, iv, 0] / voxel_spacing
+        det_y = d_detector_positions[iview, iu, iv, 1] / voxel_spacing
+        det_z = d_detector_positions[iview, iu, iv, 2] / voxel_spacing
 
     dir_x, dir_y, dir_z = det_x - src_x, det_y - src_y, det_z - src_z
     length = math.sqrt(dir_x * dir_x + dir_y * dir_y + dir_z * dir_z)
@@ -406,15 +421,20 @@ def _cone_3d_geometry_vjp_kernel(
     cuda.atomic.add(d_grad_src, (iview, 0), g * ga_x)
     cuda.atomic.add(d_grad_src, (iview, 1), g * ga_y)
     cuda.atomic.add(d_grad_src, (iview, 2), g * ga_z)
-    cuda.atomic.add(d_grad_det_center, (iview, 0), g * gb_x)
-    cuda.atomic.add(d_grad_det_center, (iview, 1), g * gb_y)
-    cuda.atomic.add(d_grad_det_center, (iview, 2), g * gb_z)
-    cuda.atomic.add(d_grad_det_u, (iview, 0), g * u_phys * gb_x)
-    cuda.atomic.add(d_grad_det_u, (iview, 1), g * u_phys * gb_y)
-    cuda.atomic.add(d_grad_det_u, (iview, 2), g * u_phys * gb_z)
-    cuda.atomic.add(d_grad_det_v, (iview, 0), g * v_phys * gb_x)
-    cuda.atomic.add(d_grad_det_v, (iview, 1), g * v_phys * gb_y)
-    cuda.atomic.add(d_grad_det_v, (iview, 2), g * v_phys * gb_z)
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        cuda.atomic.add(d_grad_det_center, (iview, 0), g * gb_x)
+        cuda.atomic.add(d_grad_det_center, (iview, 1), g * gb_y)
+        cuda.atomic.add(d_grad_det_center, (iview, 2), g * gb_z)
+        cuda.atomic.add(d_grad_det_u, (iview, 0), g * u_phys * gb_x)
+        cuda.atomic.add(d_grad_det_u, (iview, 1), g * u_phys * gb_y)
+        cuda.atomic.add(d_grad_det_u, (iview, 2), g * u_phys * gb_z)
+        cuda.atomic.add(d_grad_det_v, (iview, 0), g * v_phys * gb_x)
+        cuda.atomic.add(d_grad_det_v, (iview, 1), g * v_phys * gb_y)
+        cuda.atomic.add(d_grad_det_v, (iview, 2), g * v_phys * gb_z)
+    else:
+        d_grad_detector_positions[iview, iu, iv, 0] = g * gb_x
+        d_grad_detector_positions[iview, iu, iv, 1] = g * gb_y
+        d_grad_detector_positions[iview, iu, iv, 2] = g * gb_z
 
 
 @_FASTMATH_DECORATOR
@@ -424,6 +444,7 @@ def _parallel_2d_geometry_vjp_kernel(
     det_spacing, d_ray_dir, d_det_origin, d_det_u_vec,
     cx, cy, voxel_spacing,
     d_grad_ray_dir, d_grad_det_origin, d_grad_det_u,
+    d_detector_positions=None, d_grad_detector_positions=None,
 ):
     """Accumulate parallel-beam geometry gradients of ``<grad_sino, A x>`` per view.
 
@@ -442,14 +463,18 @@ def _parallel_2d_geometry_vjp_kernel(
     # entry and exit faces.
     dir_x = np.float64(d_ray_dir[iang, 0])
     dir_y = np.float64(d_ray_dir[iang, 1])
-    det_ox = np.float64(d_det_origin[iang, 0]) / voxel_spacing
-    det_oy = np.float64(d_det_origin[iang, 1]) / voxel_spacing
-    u_vec_x = np.float64(d_det_u_vec[iang, 0])
-    u_vec_y = np.float64(d_det_u_vec[iang, 1])
     u_phys = (np.float32(idet) + _HALF - np.float32(n_det) * _HALF) * det_spacing
-    u_offset = (np.float64(idet) + _HALF - np.float64(n_det) * _HALF) * det_spacing / voxel_spacing
-    pnt_x = det_ox + u_offset * u_vec_x
-    pnt_y = det_oy + u_offset * u_vec_y
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        det_ox = np.float64(d_det_origin[iang, 0]) / voxel_spacing
+        det_oy = np.float64(d_det_origin[iang, 1]) / voxel_spacing
+        u_vec_x = np.float64(d_det_u_vec[iang, 0])
+        u_vec_y = np.float64(d_det_u_vec[iang, 1])
+        u_offset = (np.float64(idet) + _HALF - np.float64(n_det) * _HALF) * det_spacing / voxel_spacing
+        pnt_x = det_ox + u_offset * u_vec_x
+        pnt_y = det_oy + u_offset * u_vec_y
+    else:
+        pnt_x = np.float64(d_detector_positions[iang, idet, 0]) / voxel_spacing
+        pnt_y = np.float64(d_detector_positions[iang, idet, 1]) / voxel_spacing
 
     t_min, t_max = -_BIG, _BIG
     ent_axis = -1
@@ -556,7 +581,11 @@ def _parallel_2d_geometry_vjp_kernel(
 
     cuda.atomic.add(d_grad_ray_dir, (iang, 0), g * gr_x)
     cuda.atomic.add(d_grad_ray_dir, (iang, 1), g * gr_y)
-    cuda.atomic.add(d_grad_det_origin, (iang, 0), g * gp_x)
-    cuda.atomic.add(d_grad_det_origin, (iang, 1), g * gp_y)
-    cuda.atomic.add(d_grad_det_u, (iang, 0), g * u_phys * gp_x)
-    cuda.atomic.add(d_grad_det_u, (iang, 1), g * u_phys * gp_y)
+    if d_detector_positions is None or d_detector_positions.size == 0:
+        cuda.atomic.add(d_grad_det_origin, (iang, 0), g * gp_x)
+        cuda.atomic.add(d_grad_det_origin, (iang, 1), g * gp_y)
+        cuda.atomic.add(d_grad_det_u, (iang, 0), g * u_phys * gp_x)
+        cuda.atomic.add(d_grad_det_u, (iang, 1), g * u_phys * gp_y)
+    else:
+        d_grad_detector_positions[iang, idet, 0] = g * gp_x
+        d_grad_detector_positions[iang, idet, 1] = g * gp_y