diffct 2.4.1__tar.gz → 2.4.3__tar.gz

{diffct-2.4.1 → diffct-2.4.3}/CHANGELOG.md

@@ -7,6 +7,30 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [2.4.3] - 2026-10-22
+
+### Changed
+
+- Reduce GPU transfers and host synchronization in native geometry admission: validity flags are reduced on the device and flat detector corners are admitted in one batch. Invalid-geometry errors and their order are unchanged (#42).
+- Compose streamed sampled and per-view detector surfaces for a complete admitted view/pixel batch in one launch, with a bounded fallback workspace. Eligible full legacy-surface calls keep their geometry snapshots on the input GPU (#42).
+- Read the two CGLS scalar norms together, keep tensor-store updates and TV reductions on their device, and reuse SIRT/TV residuals when no callback can change the volume or operator (#42).
+- Reuse filtered view batches across volume tiles in single-rank analytical reconstruction (#42).
+- Add `examples/benchmark_gpu_overheads.py`, a synthetic benchmark for launch and transfer overheads (#42).
+
+### Fixed
+
+- Geometry-gradient replay uses the surface workspace that admission selected. Surface admission charges the larger of the composition and replay totals, the largest allocation of either phase, and a reservation bound that counts both phases (#42).
+- Spawned multi-process test workers exit after a successful teardown, which removes intermittent `terminate called without an active exception` aborts in CI (#43).
+- Add a formal check that the union reservation bounds every modelled composition/replay allocation order (#44).
+
+## [2.4.2] - 2026-10-21
+
+### Fixed
+
+- Count host world points in the streamed geometry budget only for surfaces that expand them on the host. Shared, non-trainable legacy offsets no longer shrink view batches (#40).
+- Compose shared legacy offsets with one launch per job: the frames of all views are staged in one transfer and `compose_surface` runs on a 2-D grid whose rows are views. The offset snapshot is staged once per call and pixel tile; offsets supplied by the caller, such as in the geometry VJP replay, are staged on every call (#40).
+- Count the batched frame staging in the streamed surface memory budget, and pack shared offsets into the workspace without a device staging buffer (#40).
+
 ## [2.4.1] - 2026-10-20
 
 ### Fixed

{diffct-2.4.1 → diffct-2.4.3}/CITATION.cff

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

{diffct-2.4.1 → diffct-2.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffct
-Version: 2.4.1
+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.1 → diffct-2.4.3}/README.md

@@ -37,7 +37,7 @@
 > **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
 
@@ -146,18 +146,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).
 
@@ -197,7 +197,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.1 → diffct-2.4.3}/diffct/__init__.py

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

{diffct-2.4.1 → diffct-2.4.3}/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-2.4.1 → diffct-2.4.3}/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-2.4.1 → diffct-2.4.3}/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-2.4.1 → diffct-2.4.3}/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-2.4.1 → diffct-2.4.3}/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-2.4.1 → diffct-2.4.3}/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-2.4.1 → diffct-2.4.3}/diffct/kernels/surface_geometry.py

@@ -51,26 +51,34 @@ def compose_surface(frames, offsets, source, points, view_index, pixel_start,
     ``frame32`` gives the arithmetic dtype of each packed frame component.
     ``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: grid row ``y``
+    writes view ``view_index + y`` from the ``y``-th of ``gridDim.y`` equal
+    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
+    base = cuda.blockIdx.y * (frames.size // cuda.gridDim.y)
     if pixel == 0:
         for axis in range(rank):
-            value = frames[axis]
+            value = frames[base + axis]
             if not parallel:
                 value = _surface_add(value, -centre[axis], frame32[0])
             source[view_index, axis] = np.float32(value)
     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]
@@ -85,20 +93,20 @@ def compose_surface(frames, offsets, source, points, view_index, pixel_start,
     point_squared = np.float64(0.0)
     for axis in range(rank):
         point = _surface_add(
-            frames[rank + axis],
-            _surface_multiply(tangent, frames[2 * rank + axis], tangent32),
+            frames[base + rank + axis],
+            _surface_multiply(tangent, frames[base + 2 * rank + axis], tangent32),
             frame32[1] and tangent32,
         )
         if rank == 3:
             point = _surface_add(
-                point, _surface_multiply(vertical, frames[3 * rank + axis], vertical32),
+                point, _surface_multiply(vertical, frames[base + 3 * rank + axis], vertical32),
                 point32,
             )
-            normal = frames[4 * rank + axis]
+            normal = frames[base + 4 * rank + axis]
         else:
-            normal = frames[2 * rank + 1] if axis == 0 else -frames[2 * rank]
+            normal = frames[base + 2 * rank + 1] if axis == 0 else -frames[base + 2 * rank]
         point = _surface_add(point, _surface_multiply(depth, normal, depth32), final32)
-        source_value = frames[axis]
+        source_value = frames[base + axis]
         source_native = np.float32(source_value)
         point_native = np.float32(point)
         finite = (finite and math.isfinite(source_value) and math.isfinite(point)