diffct 2.4.0__tar.gz → 2.4.1__tar.gz

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

@@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [2.4.1] - 2026-10-20
+
+### Fixed
+
+- Let `ParameterizedSurface` use full native CUDA execution when the existing memory limits admit it. Full curved-detector geometry gradients now pull back one GPU composition graph instead of replaying tiny pixel fragments through host autograd; explicitly chunked and memory-limited calls retain bounded streaming (#37).
+- Revalidate updated `ParameterizedSurface` pixels on portable calls and native data transforms. Native `*_into` calls validate all sampled batches before changing caller output (#37).
+- Apply call-time trajectory admission to retained, trainable constructor frames, including geometry updated into degenerate or out-of-range values. Flat CUDA frame validation avoids bulk CPU round trips (#37).
+- Keep each portable distributed forward's surface sampler for backward when callers replace the sampler or parameters before differentiating (#37).
+- Preserve native fan/cone ray weights for distant endpoints by clipping and rebasing each ray in FP64 on the GPU, while retaining FP32 local traversal (#37).
+- Reject malformed checkpoint metadata structures collectively before solver work, keeping distributed process groups usable after a failed resume (#37).
+
 ## [2.4.0] - 2026-10-20
 
 ### Added

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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffct
-Version: 2.4.0
+Version: 2.4.1
 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/
@@ -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.0">
-    <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.4.0" 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.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">
   </picture>
 </p>
 
 | 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
 |---|---:|---:|---:|
-| Forward projection | 493.40 ms | 131.78 ms (3.7×) | 72.42 ms (6.8×) |
-| Adjoint (backprojection) | 1118.98 ms | 293.01 ms (3.8×) | 232.50 ms (4.8×) |
-| CGLS iteration | 1744.92 ms | 465.62 ms (3.7×) | 337.98 ms (5.2×) |
+| 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×) |
 
-diffct 2.4.0, 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.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.
 
 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).
 

{diffct-2.4.0 → diffct-2.4.1}/README.md

@@ -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.0">
-    <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.4.0" 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.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">
   </picture>
 </p>
 
 | 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
 |---|---:|---:|---:|
-| Forward projection | 493.40 ms | 131.78 ms (3.7×) | 72.42 ms (6.8×) |
-| Adjoint (backprojection) | 1118.98 ms | 293.01 ms (3.8×) | 232.50 ms (4.8×) |
-| CGLS iteration | 1744.92 ms | 465.62 ms (3.7×) | 337.98 ms (5.2×) |
+| 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×) |
 
-diffct 2.4.0, 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.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.
 
 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).
 

{diffct-2.4.0 → diffct-2.4.1}/diffct/__init__.py

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

{diffct-2.4.0 → diffct-2.4.1}/diffct/_solver_state.py

@@ -211,6 +211,8 @@ def load_checkpoint(operator, directory, expected, iterations, elements, names):
     try:
         directory = Path(directory).resolve()
         saved = json.loads((directory / "metadata.json").read_text(encoding="utf-8"))
+        if not isinstance(saved, dict):
+            raise ValueError("checkpoint metadata must be an object")
         # JSON canonicalization also normalizes tuple-valued physical spacing.
         canonical = json.loads(json.dumps(expected))
         if any(saved.get(key) != value for key, value in canonical.items()):
@@ -225,12 +227,16 @@ def load_checkpoint(operator, directory, expected, iterations, elements, names):
         if expected["algorithm"] in ("cgls", "cgls_into"):
             finite(saved.get("gamma"), "checkpoint gamma")
         manifest = saved.get("states", {})
+        if not isinstance(manifest, dict):
+            raise ValueError("checkpoint states must be an object")
         if set(manifest) != set(names):
             raise ValueError("checkpoint does not contain the complete solver state")
         states = {}
         for name in names:
-            store = NpyStore(directory / f"{name}.npy", mode="r")
             specification = manifest[name]
+            if not isinstance(specification, dict):
+                raise ValueError("checkpoint state specification must be an object")
+            store = NpyStore(directory / f"{name}.npy", mode="r")
             shape = operator.projection_shape if name in ("r", "q", "row") else operator.volume_shape
             if name == "step":
                 shape = (1,)

{diffct-2.4.0 → diffct-2.4.1}/diffct/backends.py

@@ -168,7 +168,8 @@ class ReferenceProjector:
         self.spacing = voxel_spacing_tuple(owner.voxel_spacing, len(owner.global_volume_shape))
         self.chunk_shape, self.view_chunk_size, self.detector_chunk_shape = self._work_chunks()
         if not validate:
-            # Native owners already validated geometry before transform dispatch.
+            # Call-time owners validate frames separately and sample surfaces
+            # from the invocation's single geometry snapshot below.
             return
         with torch.enable_grad():
             geometry = self._snapshot_geometry(synchronize=synchronize)
@@ -273,8 +274,7 @@ class ReferenceProjector:
     def _validate_geometry(self, geometry, fallback, *, admission=False, synchronize=True):
         """Finish every local geometry check before any rank starts numerics."""
         owner = self.owner
-        if not admission and (owner.detector_surface is None or
-                              (owner.world_size == 1 and isinstance(owner.detector_surface, ParameterizedSurface))):
+        if not admission and owner.detector_surface is None:
             return
         devices = tuple(torch.device(device) for device in (owner._devices or (fallback,)))
         error = None
@@ -294,14 +294,8 @@ class ReferenceProjector:
                     device = devices[batch % len(devices)]
                     origins, directions, _ = self._rays(
                         views, pixels, device, validate=True, geometry=staged[device],
-                        admission=admission,
                     )
-                    if (not torch.isfinite(origins.to(dtype=owner.dtype)).all().item()
-                            or not torch.isfinite(directions.to(dtype=owner.dtype)).all().item()):
-                        raise ValueError("detector surface world coordinates must be finite")
-                    if (owner.beam not in ("parallel", "parallel3d")
-                            and torch.any(torch.all(directions == 0, dim=-1)).item()):
-                        raise ValueError("source and detector surface pixels must differ")
+                    self._check_rays(origins, directions)
         except (TypeError, ValueError, RuntimeError) as failure:
             error = failure
         if synchronize:
@@ -309,10 +303,23 @@ class ReferenceProjector:
         elif error is not None:
             raise error
 
-    def _rays(self, views, pixels, device, validate=False, stats=None, geometry=None, *, admission=True):
+    def _check_rays(self, origins, directions):
+        owner = self.owner
+        if (not torch.isfinite(origins.to(dtype=owner.dtype)).all().item()
+                or not torch.isfinite(directions.to(dtype=owner.dtype)).all().item()):
+            raise ValueError("detector surface world coordinates must be finite")
+        if (owner.beam not in ("parallel", "parallel3d")
+                and torch.any(torch.all(directions == 0, dim=-1)).item()):
+            raise ValueError("source and detector surface pixels must differ")
+
+    def _rays(self, views, pixels, device, validate=False, stats=None, geometry=None, *, admit=False):
+        """Rays of one view/pixel batch.
+
+        ``admit`` applies the admission checks of ``_validate_geometry`` to the
+        same sampled offsets before composing the numerical rays, so a
+        parameterized sampler runs once for both.
+        """
         owner = self.owner
-        # Admission compares original coordinates before FP32 rounding.
-        dtype = torch.float64 if validate and admission else owner.dtype
         axes = [(torch.arange(pixel.start, pixel.stop, dtype=torch.float64, device="cpu")
                  + .5 - size / 2) * pitch
                 for pixel, size, pitch in zip(pixels, owner.detector_shape, owner.detector_spacing)]
@@ -338,99 +345,138 @@ class ReferenceProjector:
             offsets = offsets[views]
         elif tuple(offsets.shape) not in (shared_shape, (len(ids), *shared_shape)):
             raise ValueError("detector_surface output shape must match the selected view/pixel batch")
-        if validate:
+        if validate or admit:
             if not torch.isfinite(offsets).all().item():
                 raise ValueError("detector_surface offsets must be finite")
             if owner.beam not in ("cone", "parallel3d") and torch.any(offsets[..., 1] != 0).item():
                 raise ValueError("2D detector_surface middle offsets must be zero")
-        frames = (owner._learnable or owner._trajectory) if geometry is None else geometry[:frame_count]
-        native_surface = owner.backend == "numba" and surface is not None
-        native_legacy = native_surface and not isinstance(surface, ParameterizedSurface)
-        if native_surface and not native_legacy:
-            dtype = torch.float32
-            for component in (*frames, offsets):
-                dtype = torch.promote_types(dtype, component.dtype)
-        local_dtype = offsets.dtype if native_legacy else dtype
-        local = self._transfer(offsets.to(dtype=local_dtype), device, stats)
-        if local.ndim == u.ndim + 1:
-            local = local.unsqueeze(0)
-        dimensions = len(owner.global_volume_shape)
-        selected = tuple(self._transfer(value[views].to(
-            dtype=torch.promote_types(torch.float32, value.dtype) if native_legacy else dtype,
-        ), device, stats) for value in frames)
-
-        def frame(value):
-            return value.reshape(len(ids), *([1] * u.ndim), dimensions)
-
-        def world_points(offsets, components):
-            axis_u = frame(components[2])
-            points = frame(components[1]) + offsets[..., 0, None] * axis_u
-            if owner.beam in ("cone", "parallel3d"):
-                axis_v = frame(components[3])
-                points = points + offsets[..., 1, None] * axis_v
-                normal_dtype = torch.promote_types(axis_u.dtype, axis_v.dtype)
-                normal = torch.linalg.cross(axis_u.to(dtype=normal_dtype),
-                                            axis_v.to(dtype=normal_dtype), dim=-1)
+        def compose(validate, stats):
+            check_range = (validate and owner.beam not in ("parallel", "parallel3d")
+                           and owner.dtype == torch.float32 and surface is not None)
+            dtype = torch.float64 if check_range else owner.dtype
+            frames = (owner._learnable or owner._trajectory) if geometry is None else geometry[:frame_count]
+            native_surface = owner.backend == "numba" and surface is not None
+            native_legacy = native_surface and not isinstance(surface, ParameterizedSurface)
+            if native_surface and not native_legacy:
+                dtype = torch.float32
+                for component in (*frames, offsets):
+                    dtype = torch.promote_types(dtype, component.dtype)
+            local_dtype = offsets.dtype if native_legacy else dtype
+            local = self._transfer(offsets.to(dtype=local_dtype), device, stats)
+            if local.ndim == u.ndim + 1:
+                local = local.unsqueeze(0)
+            dimensions = len(owner.global_volume_shape)
+            selected = tuple(self._transfer(value[views].to(
+                dtype=torch.promote_types(torch.float32, value.dtype) if native_legacy else dtype,
+            ), device, stats) for value in frames)
+
+            def frame(value):
+                return value.reshape(len(ids), *([1] * u.ndim), dimensions)
+
+            def world_points(offsets, components):
+                axis_u = frame(components[2])
+                points = frame(components[1]) + offsets[..., 0, None] * axis_u
+                if owner.beam in ("cone", "parallel3d"):
+                    axis_v = frame(components[3])
+                    points = points + offsets[..., 1, None] * axis_v
+                    normal_dtype = torch.promote_types(axis_u.dtype, axis_v.dtype)
+                    normal = torch.linalg.cross(axis_u.to(dtype=normal_dtype),
+                                                axis_v.to(dtype=normal_dtype), dim=-1)
+                else:
+                    normal = torch.stack((axis_u[..., 1], -axis_u[..., 0]), dim=-1)
+                return points + offsets[..., 2, None] * normal
+
+            if check_range and not native_surface:
+                # Certify ordinary batches without expanding raw FP64 world points.
+                # L1 norms bound Euclidean distances without squaring tiny bases.
+                pitches = local.new_tensor(self.spacing[::-1])
+                maximum = local.abs().amax(dim=tuple(range(1, local.ndim - 1)))
+                axis_u = selected[2]
+                if owner.beam == "cone":
+                    basis = (axis_u, selected[3], torch.linalg.cross(axis_u, selected[3], dim=-1))
+                    components = (0, 1, 2)
+                else:
+                    basis = (axis_u, torch.stack((axis_u[:, 1], -axis_u[:, 0]), dim=-1))
+                    components = (0, 2)
+                source_bound = torch.linalg.vector_norm(selected[0] / pitches, ord=1, dim=-1)
+                pixel_bound = torch.linalg.vector_norm(selected[1] / pitches, ord=1, dim=-1)
+                for component, axis in zip(components, basis):
+                    pixel_bound = pixel_bound + maximum[:, component] * torch.linalg.vector_norm(
+                        axis / pitches, ord=1, dim=-1,
+                    )
+                # Cover rounding in the few frame/offset operations; nonfinite or
+                # boundary-adjacent bounds remain inconclusive and use exact points.
+                slack = 64 * torch.finfo(torch.float64).eps
+                source_bound = source_bound + slack * source_bound.clamp_min(1)
+                pixel_bound = pixel_bound + slack * pixel_bound.clamp_min(1)
+                certified = (torch.isfinite(source_bound) & torch.isfinite(pixel_bound)
+                             & (torch.minimum(source_bound, pixel_bound) <= 1e6)
+                             & (source_bound <= 1e15) & (pixel_bound <= 1e15))
+                if certified.all().item():
+                    check_range = False
+                    dtype = owner.dtype
+                    selected = tuple(value.to(dtype=dtype) for value in selected)
+                    local = local.to(dtype=dtype)
+            points = world_points(local, selected)
+            if check_range:
+                # Raw FP64 coordinates preserve the exact admission boundaries.
+                sources = frame(selected[0]).expand_as(points)
+                pitches = points.new_tensor(self.spacing[::-1])
+                source_distance = torch.linalg.vector_norm(sources / pitches, dim=-1)
+                pixel_distance = torch.linalg.vector_norm(points / pitches, dim=-1)
+                if torch.any(torch.minimum(source_distance, pixel_distance) > 1e6).item():
+                    raise ValueError(
+                        "the source or a detector pixel must lie within 1e6 voxels "
+                        "of the volume center in every view"
+                    )
+                if torch.any((source_distance > 1e15) | (pixel_distance > 1e15)).item():
+                    raise ValueError(
+                        "the source and detector surface pixels must lie within 1e15 voxels "
+                        "of the volume center"
+                    )
+            if native_surface:
+                # Native surface preprocessing composes before staging complete
+                # rays to FP32. Data transforms must preserve that operator.
+                selected = tuple(value.to(dtype=owner.dtype) for value in selected)
+                points = points.to(dtype=owner.dtype)
+            elif validate and dtype != owner.dtype:
+                # Casting the completed point misses cancellation/overflow during
+                # FP32 component arithmetic. Reuse the same sampled offsets.
+                selected = tuple(value.to(dtype=owner.dtype) for value in selected)
+                points = world_points(local.to(dtype=owner.dtype), selected)
+            if owner.beam in ("parallel", "parallel3d"):
+                origins, directions = points, frame(selected[0]).expand_as(points)
             else:
-                normal = torch.stack((axis_u[..., 1], -axis_u[..., 0]), dim=-1)
-            return points + offsets[..., 2, None] * normal
-
-        points = world_points(local, selected)
-        if (validate and owner.beam not in ("parallel", "parallel3d") and owner.dtype == torch.float32
-                and owner.detector_surface is not None):
-            # Raw FP64 coordinates preserve the exact admission boundaries.
-            sources = frame(selected[0]).expand_as(points)
-            pitches = points.new_tensor(self.spacing[::-1])
-            source_distance = torch.linalg.vector_norm(sources / pitches, dim=-1)
-            pixel_distance = torch.linalg.vector_norm(points / pitches, dim=-1)
-            if torch.any(torch.minimum(source_distance, pixel_distance) > 1e6).item():
-                raise ValueError(
-                    "the source or a detector pixel must lie within 1e6 voxels "
-                    "of the volume center in every view"
-                )
-            if torch.any((source_distance > 1e15) | (pixel_distance > 1e15)).item():
-                raise ValueError(
-                    "the source and detector surface pixels must lie within 1e15 voxels "
-                    "of the volume center"
-                )
-        if native_surface:
-            # Native surface preprocessing composes before staging complete
-            # rays to FP32. Data transforms must preserve that operator.
-            selected = tuple(value.to(dtype=owner.dtype) for value in selected)
-            points = points.to(dtype=owner.dtype)
-        elif validate and dtype != owner.dtype:
-            # Casting the completed point misses cancellation/overflow during
-            # FP32 component arithmetic. Reuse the same sampled offsets.
-            selected = tuple(value.to(dtype=owner.dtype) for value in selected)
-            points = world_points(local.to(dtype=owner.dtype), selected)
-        if owner.beam in ("parallel", "parallel3d"):
-            origins, directions = points, frame(selected[0]).expand_as(points)
-        else:
-            sources = frame(selected[0]).expand_as(points)
-            pitches = points.new_tensor(self.spacing[::-1])
-            source_distance = torch.linalg.vector_norm(sources / pitches, dim=-1)
-            pixel_distance = torch.linalg.vector_norm(points / pitches, dim=-1)
-            if validate:
-                if torch.any(torch.all(
-                    points.to(dtype=owner.dtype) == sources.to(dtype=owner.dtype), dim=-1,
-                )).item():
-                    raise ValueError("source and detector surface pixels must differ")
-            # Starting near the volume avoids subtracting its voxel planes
-            # from a distant endpoint; reversing a segment preserves [0, 1].
-            source_nearer = (source_distance <= pixel_distance).unsqueeze(-1)
-            origins = torch.where(source_nearer, sources, points)
-            difference = points - sources
-            directions = torch.where(source_nearer, difference, -difference)
-
-        pixel_ids = torch.meshgrid(*(torch.arange(pixel.start, pixel.stop, device=device)
-                                     for pixel in pixels), indexing="ij")
-        detector_ids = torch.zeros_like(pixel_ids[0])
-        for coordinate, size in zip(pixel_ids, owner.detector_shape):
-            detector_ids = detector_ids * size + coordinate
-        ray_ids = ((self._transfer(ids, device, stats) - owner.view_slice.start).reshape(
-            len(ids), *([1] * u.ndim),
-        ) * math.prod(owner.detector_shape) + detector_ids).reshape(-1)
-        return origins.reshape(-1, dimensions), directions.reshape(-1, dimensions), ray_ids
+                sources = frame(selected[0]).expand_as(points)
+                pitches = points.new_tensor(self.spacing[::-1])
+                source_distance = torch.linalg.vector_norm(sources / pitches, dim=-1)
+                pixel_distance = torch.linalg.vector_norm(points / pitches, dim=-1)
+                if validate:
+                    if torch.any(torch.all(
+                        points.to(dtype=owner.dtype) == sources.to(dtype=owner.dtype), dim=-1,
+                    )).item():
+                        raise ValueError("source and detector surface pixels must differ")
+                # Starting near the volume avoids subtracting its voxel planes
+                # from a distant endpoint; reversing a segment preserves [0, 1].
+                source_nearer = (source_distance <= pixel_distance).unsqueeze(-1)
+                origins = torch.where(source_nearer, sources, points)
+                difference = points - sources
+                directions = torch.where(source_nearer, difference, -difference)
+
+            pixel_ids = torch.meshgrid(*(torch.arange(pixel.start, pixel.stop, device=device)
+                                         for pixel in pixels), indexing="ij")
+            detector_ids = torch.zeros_like(pixel_ids[0])
+            for coordinate, size in zip(pixel_ids, owner.detector_shape):
+                detector_ids = detector_ids * size + coordinate
+            ray_ids = ((self._transfer(ids, device, stats) - owner.view_slice.start).reshape(
+                len(ids), *([1] * u.ndim),
+            ) * math.prod(owner.detector_shape) + detector_ids).reshape(-1)
+            return origins.reshape(-1, dimensions), directions.reshape(-1, dimensions), ray_ids
+
+        if admit:
+            with torch.no_grad():
+                self._check_rays(*compose(True, None)[:2])
+        return compose(validate, stats)
 
     def _intervals(self, origins, directions, tile):
         owner = self.owner
@@ -547,7 +593,7 @@ class ReferenceProjector:
         return indices.flatten()
 
     def _tensor_apply(self, tensor, is_project, reduce_cotangent=False, *, geometry=None,
-                      reduce_output=True, _plan=None):
+                      reduce_output=True, _plan=None, admit=False):
         owner = self.owner
         owner.last_execution_stats = None
         start = time.perf_counter()
@@ -588,7 +634,7 @@ class ReferenceProjector:
                     slices[axis] = slice(first, min(first + step, leading[axis]))
                     pieces.append(self._tensor_apply(tensor[tuple(slices)], is_project, reduce_cotangent,
                                                      geometry=geometry, reduce_output=reduce_output,
-                                                     _plan=plan))
+                                                     _plan=plan, admit=admit))
                     part_stats = owner.last_execution_stats
                     for key in counters:
                         stats[key] += part_stats[key]
@@ -603,13 +649,14 @@ class ReferenceProjector:
         # Empty ownership and empty leading dimensions retain a differentiable seam.
         output = output + _zero_dependency(tensor).to(dtype=owner.dtype)
         if is_project:
-            for tile, world in self._tiles(plan[0]):
+            for tile_index, (tile, world) in enumerate(self._tiles(plan[0])):
                 values = tensor[(..., *tile)]
                 if reduce_cotangent:
                     values = self._reduce(values, owner.view_group, owner.view_world_size, stats)
                 for batch, (views, pixels) in enumerate(self._batches(*plan[1:])):
                     device = self._device(batch, tensor.device, stats)
-                    origins, directions, rays = self._rays(views, pixels, device, stats=stats, geometry=geometry)
+                    origins, directions, rays = self._rays(views, pixels, device, stats=stats, geometry=geometry,
+                                                           admit=admit and tile_index == 0)
                     indices, lengths = self._intervals(origins, directions, world)
                     local = self._to(values, device, stats).reshape(*leading, math.prod(values.shape[-len(tile):]))
                     selected = torch.index_select(local, -1, indices.flatten()).reshape(*leading, *indices.shape)
@@ -632,7 +679,8 @@ class ReferenceProjector:
                 if reduce_cotangent:
                     values = self._reduce(values, owner.space_group, owner.space_world_size, stats)
                 values = self._to(values, device, stats).reshape(*leading, math.prod(values.shape[-len(shape):]))
-                origins, directions, _ = self._rays(views, pixels, device, stats=stats, geometry=geometry)
+                origins, directions, _ = self._rays(views, pixels, device, stats=stats, geometry=geometry,
+                                                    admit=admit)
                 for tile, world in self._tiles(plan[0]):
                     indices, lengths = self._intervals(origins, directions, world)
                     tile_shape = tuple(part.stop - part.start for part in tile)
@@ -657,10 +705,17 @@ class ReferenceProjector:
 
     def _apply(self, tensor, is_project):
         geometry = self._snapshot_geometry()
-        self._validate_geometry(geometry, tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu"))
-        if self.owner.world_size > 1 and torch.is_grad_enabled():
-            return _DistributedApply.apply(tensor, self, is_project, False, self.owner._participation(), *geometry)
-        return self._tensor_apply(tensor, is_project, geometry=geometry)
+        owner = self.owner
+        # A single process returns no output before every batch passes, so a
+        # parameterized surface is admitted in the numerical sampling pass and
+        # its sampler runs once per batch. Distributed calls finish all checks
+        # before any rank starts numerics.
+        admit = owner.world_size == 1 and isinstance(owner.detector_surface, ParameterizedSurface)
+        if not admit:
+            self._validate_geometry(geometry, tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu"))
+        if owner.world_size > 1 and torch.is_grad_enabled():
+            return _DistributedApply.apply(tensor, self, is_project, False, owner._participation(), *geometry)
+        return self._tensor_apply(tensor, is_project, geometry=geometry, admit=admit)
 
     def _geometry_grad(self, tensor, cotangent, is_project, reduce_cotangent, geometry):
         owner = self.owner

{diffct-2.4.0 → diffct-2.4.1}/diffct/kernels/cone_beam.py

@@ -120,6 +120,11 @@ def _cone_3d_forward_kernel(
         det_y = d_detector_positions[iview, iu, iv, 1] / voxel_spacing
         det_z = d_detector_positions[iview, iu, iv, 2] / voxel_spacing
 
+    # Preserve native endpoint composition, then clip in float64: when both
+    # endpoints are distant, float32 loses voxel widths before rebasing.
+    src_x, src_y, src_z = np.float64(src_x), np.float64(src_y), np.float64(src_z)
+    det_x, det_y, det_z = np.float64(det_x), np.float64(det_y), np.float64(det_z)
+
     # === 3D RAY DIRECTION CALCULATION ===
     # Ray direction vector from source to detector element in 3D space
     dir_x, dir_y, dir_z = det_x - src_x, det_y - src_y, det_z - src_z
@@ -133,7 +138,7 @@ def _cone_3d_forward_kernel(
 
     # === 3D RAY-VOLUME INTERSECTION CALCULATION ===
     # Measure the ray parameter from the endpoint nearer the volume centre, so a
-    # distant source does not make the float32 parameters coarse. Only the
+    # distant source does not make the parameters coarse. Only the
     # source-to-detector segment is integrated.
     if src_x * src_x + src_y * src_y + src_z * src_z <= det_x * det_x + det_y * det_y + det_z * det_z:
         org_x = src_x
@@ -170,15 +175,14 @@ def _cone_3d_forward_kernel(
     if t_min >= t_max:  # No valid 3D intersection
         d_sino[iview, iu, iv] = _ZERO; return
 
-    # The traversal restarts at the entry point, so its parameters stay
-    # within the volume diameter.
-    ent_x = org_x + t_min * dir_x
-    ent_y = org_y + t_min * dir_y
-    ent_z = org_z + t_min * dir_z
-    ray_x = dir_x
-    ray_y = dir_y
-    ray_z = dir_z
-    t_end = t_max - t_min
+    # Only setup needs float64; traverse in float32 from the volume entry.
+    ent_x = np.float32(org_x + t_min * dir_x)
+    ent_y = np.float32(org_y + t_min * dir_y)
+    ent_z = np.float32(org_z + t_min * dir_z)
+    ray_x = np.float32(dir_x)
+    ray_y = np.float32(dir_y)
+    ray_z = np.float32(dir_z)
+    t_end = np.float32(t_max - t_min)
 
     # === 3D SIDDON METHOD TRAVERSAL INITIALIZATION ===
     accum = _ZERO  # Accumulated projection value
@@ -337,6 +341,10 @@ def _cone_3d_backward_kernel(
         det_y = d_detector_positions[iview, iu, iv, 1] / voxel_spacing
         det_z = d_detector_positions[iview, iu, iv, 2] / voxel_spacing
 
+    # Match forward setup while retaining native endpoint composition.
+    src_x, src_y, src_z = np.float64(src_x), np.float64(src_y), np.float64(src_z)
+    det_x, det_y, det_z = np.float64(det_x), np.float64(det_y), np.float64(det_z)
+
     # === 3D RAY DIRECTION CALCULATION ===
     # Ray direction vector from source to detector element in 3D space
     dir_x, dir_y, dir_z = det_x - src_x, det_y - src_y, det_z - src_z
@@ -347,7 +355,7 @@ def _cone_3d_backward_kernel(
 
     # === 3D RAY-VOLUME INTERSECTION CALCULATION ===
     # Measure the ray parameter from the endpoint nearer the volume centre, so a
-    # distant source does not make the float32 parameters coarse. Only the
+    # distant source does not make the parameters coarse. Only the
     # source-to-detector segment is integrated.
     if src_x * src_x + src_y * src_y + src_z * src_z <= det_x * det_x + det_y * det_y + det_z * det_z:
         org_x = src_x
@@ -380,15 +388,14 @@ def _cone_3d_backward_kernel(
 
     if t_min >= t_max: return
 
-    # The traversal restarts at the entry point, so its parameters stay
-    # within the volume diameter.
-    ent_x = org_x + t_min * dir_x
-    ent_y = org_y + t_min * dir_y
-    ent_z = org_z + t_min * dir_z
-    ray_x = dir_x
-    ray_y = dir_y
-    ray_z = dir_z
-    t_end = t_max - t_min
+    # Only setup needs float64; traverse in float32 from the volume entry.
+    ent_x = np.float32(org_x + t_min * dir_x)
+    ent_y = np.float32(org_y + t_min * dir_y)
+    ent_z = np.float32(org_z + t_min * dir_z)
+    ray_x = np.float32(dir_x)
+    ray_y = np.float32(dir_y)
+    ray_z = np.float32(dir_z)
+    t_end = np.float32(t_max - t_min)
 
     # === 3D SIDDON METHOD TRAVERSAL INITIALIZATION ===
     t = _ZERO