diffct 2.2.1__py3-none-any.whl → 2.3.0__py3-none-any.whl

diffct/__init__.py

@@ -20,11 +20,15 @@ from .surfaces import ParameterizedSurface
 from .reconstruction import (
     fbp,
     fdk,
+    fbp_into,
+    fdk_into,
     helical_fdk,
     helical_view_weights,
     helical_weighted_backproject,
 )
 
+from .solvers import SolverResult, cgls, sirt, tv_reconstruction, cgls_into, total_variation
+
 from .geometry import (
     circular_trajectory_3d,
     random_trajectory_3d,
@@ -58,7 +62,7 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.2.1'
+__version__ = '2.3.0'
 
 __all__ = [
     'Projector',
@@ -67,6 +71,14 @@ __all__ = [
     'ParameterizedSurface',
     'fbp',
     'fdk',
+    'fbp_into',
+    'fdk_into',
+    'SolverResult',
+    'cgls',
+    'sirt',
+    'tv_reconstruction',
+    'cgls_into',
+    'total_variation',
     'helical_fdk',
     'helical_view_weights',
     'helical_weighted_backproject',

diffct/_solver_state.py

@@ -0,0 +1,284 @@
+"""Bounded solver state, owned reductions and complete checkpoint publication."""
+
+import hashlib
+import json
+import math
+import numbers
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+
+from .pipeline import _fragments
+from .storage import NpyStore, TensorStore
+
+
+def count(value, name, positive=False):
+    if isinstance(value, bool) or not isinstance(value, numbers.Integral) or value < int(positive):
+        raise ValueError(f"{name} must be a {'positive' if positive else 'nonnegative'} integer")
+    return int(value)
+
+
+def finite(value, name, positive=False):
+    if isinstance(value, bool) or not isinstance(value, numbers.Real) or not math.isfinite(value) or value < 0 or (positive and value == 0):
+        raise ValueError(f"{name} must be finite and {'positive' if positive else 'nonnegative'}")
+    return float(value)
+
+
+def agreement(operator, error=None, token=None):
+    """Every rank rejects a bad request before numerical execution starts."""
+    world = getattr(operator, "world_size", 1)
+    if world == 1:
+        if error is not None:
+            raise error
+        return
+    records = [None] * world
+    dist.all_gather_object(records, (None if error is None else str(error), token),
+                           group=operator.process_group)
+    failures = [record[0] for record in records if record[0] is not None]
+    if failures or any(record[1] != records[0][1] for record in records):
+        raise ValueError("distributed solver agreement failed: " + (failures[0] if failures else "inconsistent checkpoint iteration or controls"))
+
+
+def scalar_sum(operator, value, kind):
+    """Count physical rays/voxels once, including view/space replicas."""
+    world = getattr(operator, "world_size", 1)
+    partition = getattr(operator, "partition", "views")
+    rank = getattr(operator, "rank", 0)
+    owns = (getattr(operator, "view_rank", 0) == 0 if kind == "volume"
+            else getattr(operator, "space_rank", 0) == 0) if partition == "hybrid" else (
+        rank == 0 if (kind == "volume") == (partition == "views") else True
+    )
+    value = float(value) if owns 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.tensor(value, dtype=torch.float64, device=device)
+        dist.all_reduce(accumulator, group=group)
+        value = accumulator.item()
+    return value
+
+
+def squared_norm(operator, value, kind, elements=1048576):
+    store = TensorStore(value) if isinstance(value, torch.Tensor) else value
+    total = 0.
+    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()
+        total += block.square_().sum().item()
+    return scalar_sum(operator, total, kind)
+
+
+def copy_blocks(source, target, elements):
+    for selected in _fragments(tuple(target.shape), elements):
+        target.write(selected, source.read(selected).to(device=getattr(target, "device", "cpu"), dtype=target.dtype))
+
+
+def fill_blocks(target, elements):
+    for selected in _fragments(tuple(target.shape), elements):
+        shape = tuple(part.stop - part.start for part in selected)
+        target.write(selected, torch.zeros(shape, dtype=target.dtype))
+
+
+def update_blocks(target, source, factor, elements, scale=1.):
+    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)
+        target.write(selected, block)
+
+
+def fingerprint(value, elements):
+    store = TensorStore(value.detach()) if isinstance(value, torch.Tensor) else value
+    digest = hashlib.sha256(json.dumps((tuple(store.shape), str(store.dtype))).encode())
+    for selected in _fragments(tuple(store.shape), elements):
+        block = store.read(selected).detach().to(device="cpu").contiguous()
+        # Hash storage bytes without requiring NumPy support for the original
+        # floating dtype (for example bfloat16). This view allocates no payload.
+        digest.update(block.reshape(-1).view(torch.uint8).numpy().tobytes())
+    return digest.hexdigest()
+
+
+def _effective_surface(operator, frames, elements):
+    """Hash the actual physical pixel rays, without materializing a volume."""
+    surface = getattr(operator, "detector_surface", None)
+    if surface is None:
+        return None
+    from .backends import ReferenceProjector
+    from .surfaces import ParameterizedSurface, _detector_grids
+
+    dimensions = len(operator.global_volume_shape)
+    if elements < max(3, dimensions):
+        raise ValueError("block_elements cannot hold one physical detector pixel")
+    if isinstance(surface, ParameterizedSurface):
+        if any(value.device.type != "cpu" and value.numel() > elements for value in surface.parameters):
+            raise ValueError("surface parameters cannot be staged within block_elements")
+        geometry = (*frames, *surface.parameters)
+    else:
+        # Legacy callbacks require complete grids. Their documented output can
+        # contain every view, so admit that snapshot only when it fits the
+        # bound; a smaller arbitrary callback invocation may change its meaning.
+        if operator._n_views * math.prod(operator.detector_shape) * 3 > elements:
+            raise ValueError("legacy surface fingerprint exceeds block_elements; use ParameterizedSurface")
+        pixels = tuple(slice(0, size) for size in operator.detector_shape)
+        u, v = _detector_grids(operator.detector_shape, operator.detector_spacing, pixels)
+        offsets = surface(u, v)
+        if not isinstance(offsets, torch.Tensor) or not torch.is_floating_point(offsets):
+            raise TypeError("detector_surface must return floating-point offsets")
+        if tuple(offsets.shape) not in ((*operator.detector_shape, 3),
+                                       (operator._n_views, *operator.detector_shape, 3)):
+            raise ValueError("legacy detector surface must match the complete pixel grids")
+        geometry = (*frames, offsets)
+    reference = getattr(operator, "_reference", None)
+    if reference is None:
+        reference = ReferenceProjector(operator, validate=False)
+    description = json.dumps((operator._n_views, operator.detector_shape, str(operator.dtype))).encode()
+    origins_digest, directions_digest = hashlib.sha256(description), hashlib.sha256(description)
+    for view in range(operator._n_views):
+        for pixels in _fragments(operator.detector_shape, elements // max(3, dimensions)):
+            origins, directions, _ = reference._rays(slice(view, view + 1), pixels,
+                                                      torch.device("cpu"), validate=True, geometry=geometry)
+            # Separate streams preserve row-major ray ordering across different
+            # block sizes; changing a memory ceiling does not change the hash.
+            origins_digest.update(origins.detach().contiguous().numpy().tobytes())
+            directions_digest.update(directions.detach().contiguous().numpy().tobytes())
+    return [origins_digest.hexdigest(), directions_digest.hexdigest()]
+
+
+def _signature_local(operator, source, algorithm, settings, elements, compute_dtype):
+    frames = getattr(operator, "_learnable", ()) or getattr(operator, "_trajectory", ())
+    geometry = tuple(fingerprint(value, elements) for value in frames)
+    surface = getattr(operator, "detector_surface", None)
+    surface_values = (() if surface is None else getattr(surface, "parameters", ()))
+    if isinstance(surface, torch.Tensor):
+        surface_values = (surface,)
+    return dict(
+        schema_version=1, algorithm=algorithm, settings=settings,
+        rank=getattr(operator, "rank", 0), world_size=getattr(operator, "world_size", 1),
+        partition=getattr(operator, "partition", "views"),
+        partition_grid=getattr(operator, "partition_grid", None),
+        global_volume_shape=list(operator.global_volume_shape), volume_shape=list(operator.volume_shape),
+        volume_slice=[[part.start, part.stop] for part in operator.volume_slice],
+        projection_shape=list(operator.projection_shape),
+        view_slice=[operator.view_slice.start, operator.view_slice.stop],
+        beam=operator.beam, voxel_spacing=operator.voxel_spacing,
+        detector_spacing=list(operator.detector_spacing), detector_shape=list(operator.detector_shape),
+        input_dtype=str(source.dtype), geometry=list(geometry),
+        surface=[fingerprint(value, elements) for value in surface_values],
+        effective_surface=_effective_surface(operator, frames, elements),
+        compute_dtype=str(compute_dtype), precision=getattr(operator, "precision", None),
+        backend=getattr(operator, "backend", None),
+        measurements=fingerprint(source, elements),
+    )
+
+
+def signature(operator, source, algorithm, settings, elements, *, compute_dtype=None):
+    error, metadata = None, None
+    try:
+        dtype = getattr(operator, "dtype", source.dtype) if compute_dtype is None else compute_dtype
+        metadata = _signature_local(operator, source, algorithm, settings, elements, dtype)
+    except Exception as caught:
+        error = caught
+    agreement(operator, error)
+    if getattr(operator, "world_size", 1) > 1:
+        records = [None] * operator.world_size
+        dist.all_gather_object(records, metadata, group=operator.process_group)
+        local_fields = {"rank", "volume_shape", "volume_slice", "projection_shape", "view_slice", "measurements"}
+        scans = [{key: value for key, value in record.items() if key not in local_fields} for record in records]
+        if any(scan != scans[0] for scan in scans):
+            raise ValueError("distributed checkpoint scan and compute configuration must agree")
+        replicas = {}
+        for record in records:
+            views = tuple(record["view_slice"])
+            digest = record["measurements"]
+            if views in replicas and replicas[views] != digest:
+                raise ValueError("distributed checkpoint measurement replicas must agree")
+            replicas[views] = digest
+    return metadata
+
+
+def load_checkpoint(operator, directory, expected, iterations, elements, names):
+    error, saved, states = None, None, None
+    try:
+        directory = Path(directory).resolve()
+        saved = json.loads((directory / "metadata.json").read_text(encoding="utf-8"))
+        # 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()):
+            raise ValueError("checkpoint scan, data, ownership, dtype or algorithm settings do not match")
+        completed = count(saved.get("iteration"), "checkpoint iteration")
+        history = saved.get("residual_history")
+        if completed > iterations or not isinstance(history, list) or len(history) != completed + 1:
+            raise ValueError("checkpoint iteration or residual history is invalid")
+        for value in history:
+            finite(value, "checkpoint residual")
+        finite(saved.get("measurement_norm"), "checkpoint measurement norm")
+        if expected["algorithm"] in ("cgls", "cgls_into"):
+            finite(saved.get("gamma"), "checkpoint gamma")
+        manifest = saved.get("states", {})
+        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]
+            shape = operator.projection_shape if name in ("r", "q", "row") else operator.volume_shape
+            if name == "step":
+                shape = (1,)
+            if tuple(store.shape) != tuple(shape) or list(store.shape) != specification.get("shape") or str(store.dtype) != specification.get("dtype"):
+                raise ValueError("checkpoint solver state shape or dtype does not match")
+            if fingerprint(store, elements) != specification.get("fingerprint"):
+                raise ValueError("checkpoint solver state fingerprint does not match")
+            states[name] = store
+    except (ValueError, TypeError, OSError, EOFError, KeyError) as caught:
+        error = caught
+    token = None if error is not None else (saved["iteration"], saved["residual_history"], saved.get("gamma"), saved["measurement_norm"])
+    agreement(operator, error, token)
+    return saved, states
+
+
+def save_checkpoint(operator, directory, states, metadata, iteration, history, measurement_norm, elements, **scalars):
+    error, manifest, rank = None, {}, None
+    try:
+        step = Path(directory).resolve() / f"step-{iteration:06d}"
+        step.mkdir(parents=True, exist_ok=True)
+        rank = step / f"rank-{getattr(operator, 'rank', 0):05d}"
+        rank.mkdir()
+        for name, source in states.items():
+            source = TensorStore(source.detach()) if isinstance(source, torch.Tensor) else source
+            target = NpyStore.create(rank / f"{name}.npy", source.shape, dtype=source.dtype)
+            copy_blocks(source, target, elements)
+            target.flush()
+            manifest[name] = dict(shape=list(source.shape), dtype=str(source.dtype), fingerprint=fingerprint(target, elements))
+    except (OSError, ValueError, TypeError) as caught:
+        error = caught
+    # All owned states on every rank are written and flushed before any rank
+    # atomically publishes the metadata marking a completed iteration.
+    agreement(operator, error, (iteration, history))
+    error = None
+    try:
+        document = dict(metadata, iteration=iteration, residual_history=history,
+                        measurement_norm=measurement_norm, states=manifest, **scalars)
+        temporary = rank / "metadata.tmp"
+        with temporary.open("x", encoding="utf-8") as file:
+            json.dump(document, file)
+            file.flush()
+        temporary.replace(rank / "metadata.json")
+    except Exception as caught:
+        error = caught
+    # A failed metadata write/rename must stop every rank before another
+    # projection or callback can enter a numerical collective.
+    agreement(operator, error, iteration)
+
+
+def callback_stop(operator, callback, iteration, residual, volume):
+    error, stop = None, False
+    try:
+        stop = callback is not None and callback(iteration, residual, volume) is False
+    except Exception as caught:
+        error = caught
+    agreement(operator, error)
+    if getattr(operator, "world_size", 1) > 1:
+        records = [None] * operator.world_size
+        dist.all_gather_object(records, stop, group=operator.process_group)
+        stop = any(records)
+    return stop

diffct/_solver_tv.py

@@ -0,0 +1,110 @@
+"""Physical forward-difference TV with differentiable slab boundary faces."""
+
+import math
+import numbers
+
+import torch
+import torch.distributed as dist
+
+from ._solver_state import finite, scalar_sum
+from .backends import _zero_dependency
+
+
+def _peers(operator, nonempty):
+    spaces = operator.space_world_size
+    current = operator.space_rank
+    occupied = [rank for rank in range(spaces) if rank < min(spaces, operator.global_volume_shape[0])]
+    previous = max((rank for rank in occupied if rank < current), default=None) if nonempty else None
+    following = min((rank for rank in occupied if rank > current), default=None) if nonempty else None
+    def global_rank(rank):
+        if rank is None:
+            return None
+        parent = operator.view_rank * spaces + rank
+        return operator._partition.global_rank(parent)
+    return global_rank(previous), global_rank(following)
+
+
+class _FollowingFace(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, volume, operator):
+        previous, following = _peers(operator, volume.shape[0] > 0)
+        ctx.operator, ctx.previous, ctx.following = operator, previous, following
+        ctx.shape, ctx.dtype, ctx.device = volume.shape, volume.dtype, volume.device
+        nccl = dist.get_backend(operator.space_group) == "nccl"
+        ctx.communication_device = (torch.device("cuda", torch.cuda.current_device())
+                                    if nccl else torch.device("cpu"))
+        if nccl:
+            # Initialize the subgroup communicator collectively, including
+            # empty slabs, before the first subset-only face exchange.
+            dist.barrier(group=operator.space_group)
+        face = torch.zeros(volume.shape[1:], dtype=volume.dtype, device=ctx.communication_device)
+        operations = []
+        if previous is not None:
+            first = volume[0].detach().to(device=ctx.communication_device).contiguous()
+            operations.append(dist.P2POp(dist.isend, first, previous, operator.space_group))
+        if following is not None:
+            operations.append(dist.P2POp(dist.irecv, face, following, operator.space_group))
+        if operations:
+            for work in dist.batch_isend_irecv(operations):
+                work.wait()
+        return face.to(device=volume.device)
+
+    @staticmethod
+    def backward(ctx, grad_face):
+        received = torch.zeros(ctx.shape[1:], dtype=ctx.dtype, device=ctx.communication_device)
+        operations = []
+        if ctx.following is not None:
+            outgoing = grad_face.to(device=ctx.communication_device).contiguous()
+            operations.append(dist.P2POp(dist.isend, outgoing, ctx.following, ctx.operator.space_group))
+        if ctx.previous is not None:
+            operations.append(dist.P2POp(dist.irecv, received, ctx.previous, ctx.operator.space_group))
+        if operations:
+            for work in dist.batch_isend_irecv(operations):
+                work.wait()
+        gradient = torch.zeros(ctx.shape, dtype=ctx.dtype, device=ctx.device)
+        if ctx.shape[0] > 0:
+            gradient[0] = received.to(device=ctx.device)
+        return gradient, None
+
+
+def total_variation(volume, *, operator=None, voxel_spacing=1., eps=1e-6, reduction="mean"):
+    """Isotropic TV over every voxel, with physical forward differences.
+
+    Differences vanish at the global outer boundary. Spatially partitioned
+    volumes exchange a single neighboring face in forward and backward; the
+    reported scalar counts replicated ownership once and gradients stay local.
+    """
+    if not isinstance(volume, torch.Tensor) or not torch.is_floating_point(volume) or volume.ndim < 1:
+        raise TypeError("volume must be a real floating-point tensor")
+    eps = finite(eps, "eps", positive=True)
+    if reduction not in ("mean", "sum"):
+        raise ValueError("reduction must be mean or sum")
+    spacing = (voxel_spacing,) * volume.ndim if isinstance(voxel_spacing, numbers.Real) else tuple(voxel_spacing)
+    if len(spacing) != volume.ndim:
+        raise ValueError("voxel_spacing must match the volume axes")
+    spacing = tuple(finite(pitch, "voxel_spacing", positive=True) for pitch in spacing)
+    if operator is not None and tuple(volume.shape) != tuple(operator.volume_shape):
+        raise ValueError("volume shape must match its operator ownership")
+    distributed_space = operator is not None and getattr(operator, "world_size", 1) > 1 and operator.partition in ("space", "hybrid")
+    halo = _FollowingFace.apply(volume, operator) if distributed_space else None
+    squared = torch.zeros_like(volume)
+    for axis, pitch in enumerate(spacing):
+        current, following = [slice(None)] * volume.ndim, [slice(None)] * volume.ndim
+        current[axis], following[axis] = slice(0, -1), slice(1, None)
+        difference = torch.zeros_like(volume)
+        difference[tuple(current)] = (volume[tuple(following)] - volume[tuple(current)]) / pitch
+        if axis == 0 and halo is not None and volume.shape[0] and operator.volume_slice[0].stop < operator.global_volume_shape[0]:
+            difference[-1] = (halo - volume[-1]) / pitch
+        squared = squared + difference.square()
+    local = torch.sqrt(squared + eps).sum()
+    if halo is not None:
+        # The final slab also needs backward to receive the previous slab's
+        # contribution, even though its own global outer difference is zero.
+        local = local + _zero_dependency(halo)
+    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())
+        elements = math.prod(operator.global_volume_shape)
+    return total if reduction == "sum" else total / max(elements, 1)

diffct/analytical.py

@@ -36,16 +36,29 @@ from .utils import (
     _grid_3d,
     _on_device_of,
 )
-from .kernels import (
-    _parallel_2d_fbp_backproject_kernel,
-    _fan_2d_fbp_backproject_kernel,
-    _cone_3d_fdk_backproject_kernel,
-    _parallel_2d_fbp_transpose_kernel,
-    _fan_2d_fbp_transpose_kernel,
-    _cone_3d_fdk_transpose_kernel,
+from . import kernels as _kernels
+
+
+_KERNEL_NAMES = (
+    '_parallel_2d_fbp_backproject_kernel',
+    '_fan_2d_fbp_backproject_kernel',
+    '_cone_3d_fdk_backproject_kernel',
+    '_parallel_2d_fbp_transpose_kernel',
+    '_fan_2d_fbp_transpose_kernel',
+    '_cone_3d_fdk_transpose_kernel',
 )
 
 
+def _resolve_kernel(name):
+    return globals()[name] if name in globals() else getattr(_kernels, name)
+
+
+def __getattr__(name):
+    if name not in _KERNEL_NAMES:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    return getattr(_kernels, name)
+
+
 # ============================================================================
 # Detector coordinate helper
 # ============================================================================
@@ -95,8 +108,11 @@ def angular_integration_weights(angles, redundant_full_scan=True):
     ``redundant_full_scan=True`` for automatic fan/cone scan classification.
     Setting it to ``False`` explicitly disables the full-turn ``1/2`` factor.
     Parker and angular quadrature share the same periodic-scan decision.
+    FP16 and BF16 angles are promoted to FP32; FP64 angles retain their precision.
     """
-    angles = angles.to(dtype=torch.float32).contiguous()
+    if angles.dtype not in (torch.float32, torch.float64):
+        angles = angles.to(dtype=torch.float32)
+    angles = angles.contiguous()
     n = angles.shape[0]
     if n == 0:
         return torch.zeros_like(angles)
@@ -105,6 +121,8 @@ def angular_integration_weights(angles, redundant_full_scan=True):
 
     device = angles.device
     angles_sorted, sort_idx = torch.sort(angles)
+    # Take gaps between sorted neighbours directly. Removing an origin first
+    # rounds every angle again and can collapse distinct neighbours.
     diffs = angles_sorted[1:] - angles_sorted[:-1]
 
     # Endpoint-excluded uniform circular scans are periodic quadrature
@@ -194,8 +212,11 @@ def parker_weights(angles, num_detectors, detector_spacing, sdd, detector_offset
     Only meaningful for circular (fan-shaped) trajectories — when using
     non-circular trajectories you typically want either a constant weight
     (full-coverage) or a custom redundancy scheme.
+    FP16 and BF16 angles are promoted to FP32; FP64 angles retain their precision.
     """
-    angles = angles.to(dtype=torch.float32).contiguous()
+    if angles.dtype not in (torch.float32, torch.float64):
+        angles = angles.to(dtype=torch.float32)
+    angles = angles.contiguous()
     n_ang = angles.shape[0]
     device = angles.device
 
@@ -407,7 +428,7 @@ def _parallel_weighted_backproject_raw(sinogram, ray_dir, det_origin, det_u_vec,
     pt_stream = torch.cuda.current_stream()
     numba_stream = _get_numba_external_stream_for(pt_stream)
 
-    _parallel_2d_fbp_backproject_kernel[grid, tpb, numba_stream](
+    _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),
@@ -450,7 +471,7 @@ def _analytical_geometry(src_pos, det_center, n, isocenter):
     if not (torch.isfinite(sid_n).all() and torch.isfinite(sdd_n).all()
             and (sid_n > 1e-6).all() and (sdd_n > 1e-6).all()):
         raise ValueError("Source-to-isocenter and source-to-detector distances must be positive")
-    return isocenter, float(sid_n.mean().item()), float(sdd_n.mean().item())
+    return isocenter, sid_n, sdd_n
 
 
 @_on_device_of("sinogram")
@@ -460,10 +481,10 @@ def _fan_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec,
 
     Runs the dedicated fan-beam FBP gather kernel (one thread per output
     pixel, ``(sid_n/U_n)^2`` distance weighting, linear interpolation of
-    the filtered sinogram) and multiplies by the analytical constant
-    ``sdd_mean / (2 * pi * sid_mean)`` so a unit-density disk reconstructs
-    to amplitude 1. Both means are measured in the direction of the
-    detector normal per view, so the helper reduces to the classical
+    the filtered sinogram), applying the analytical factor
+    ``sdd_n / (2 * pi * sid_n)`` to each view so a unit-density disk
+    reconstructs to amplitude 1. Both distances are measured in the
+    detector-normal direction per view, so the helper reduces to the classical
     ``sdd/(2*pi*sid)`` on a circular orbit. ``isocenter`` is a physical
     coordinate pair; by default it is inferred from the source lines along
     detector normals. Supply it explicitly for an ambiguous or non-circular
@@ -477,9 +498,10 @@ def _fan_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec,
     det_center = _as_contig_f32(det_center, device)
     det_u_vec = _as_contig_f32(det_u_vec, device)
     normal = torch.stack((-det_u_vec[:, 1], det_u_vec[:, 0]), dim=1)
-    isocenter, sid_mean, sdd_mean = _analytical_geometry(
+    isocenter, sid_n, sdd_n = _analytical_geometry(
         src_pos, det_center, normal, isocenter,
     )
+    sinogram = sinogram * (sdd_n / (2.0 * math.pi * sid_n)).view(-1, 1)
 
     n_views, n_det = sinogram.shape
     Ny, Nx = H, W
@@ -497,7 +519,7 @@ def _fan_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec,
     pt_stream = torch.cuda.current_stream()
     numba_stream = _get_numba_external_stream_for(pt_stream)
 
-    _fan_2d_fbp_backproject_kernel[grid, tpb, numba_stream](
+    _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),
@@ -505,8 +527,7 @@ def _fan_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec,
         _DTYPE(isocenter[1].item() / voxel_spacing),
     )
 
-    scale = sdd_mean / (2.0 * math.pi * sid_mean)
-    return reco * scale
+    return reco
 
 
 @_on_device_of("sinogram")
@@ -516,9 +537,9 @@ def _cone_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec, det
 
     Runs the dedicated cone-beam FDK gather kernel (one thread per output
     voxel, ``(sid_n/U_n)^2`` distance weighting, bilinear interpolation of
-    the filtered sinogram) and multiplies by ``sdd_mean/(2*pi*sid_mean)``
-    so a unit-density sphere reconstructs to amplitude 1. Both means are
-    taken in the per-view detector-normal direction, so the helper
+    the filtered sinogram), applying ``sdd_n/(2*pi*sid_n)`` to each view
+    so a unit-density sphere reconstructs to amplitude 1. Both distances are
+    measured in the per-view detector-normal direction, so the helper
     reduces to the classical ``sdd/(2*pi*sid)`` on a circular orbit.
     ``isocenter`` is a physical coordinate triple, inferred from source
     lines along detector normals by default. Supply it explicitly for an
@@ -534,9 +555,10 @@ def _cone_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec, det
     det_u_vec = _as_contig_f32(det_u_vec, device)
     det_v_vec = _as_contig_f32(det_v_vec, device)
     normal = torch.cross(det_u_vec, det_v_vec, dim=1)
-    isocenter, sid_mean, sdd_mean = _analytical_geometry(
+    isocenter, sid_n, sdd_n = _analytical_geometry(
         src_pos, det_center, normal, isocenter,
     )
+    sinogram = sinogram * (sdd_n / (2.0 * math.pi * sid_n)).view(-1, 1, 1)
 
     n_views, n_u, n_v = sinogram.shape
     Nx, Ny, Nz = W, H, D
@@ -559,7 +581,7 @@ def _cone_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec, det
     pt_stream = torch.cuda.current_stream()
     numba_stream = _get_numba_external_stream_for(pt_stream)
 
-    _cone_3d_fdk_backproject_kernel[grid, tpb, numba_stream](
+    _resolve_kernel('_cone_3d_fdk_backproject_kernel')[grid, tpb, numba_stream](
         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,
@@ -569,10 +591,8 @@ def _cone_weighted_backproject_raw(sinogram, src_pos, det_center, det_u_vec, det
         _DTYPE(isocenter[2].item() / voxel_spacing),
     )
 
-    scale = sdd_mean / (2.0 * math.pi * sid_mean)
-
     # Convert back to (D, H, W) layout for the user.
-    reco = reco_perm.permute(2, 1, 0).contiguous() * scale
+    reco = reco_perm.permute(2, 1, 0).contiguous()
     return reco
 
 
@@ -616,17 +636,18 @@ def _analytical_transpose_raw(image, beam, geometry, parameters, isocenter, sino
         source, detector, u = geometry[:3]
         normal = (torch.stack((-u[:, 1], u[:, 0]), dim=1) if beam == "fan"
                   else torch.cross(u, geometry[3], dim=1))
-        isocenter, sid_mean, sdd_mean = _analytical_geometry(source, detector, normal, isocenter)
-        scale = sdd_mean / (2.0 * math.pi * sid_mean)
+        isocenter, sid_n, sdd_n = _analytical_geometry(source, detector, normal, isocenter)
+        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)
-    image = image * scale
 
     if beam == "cone":
         D, H, W, du, dv, spacing = parameters
         image = image.permute(2, 1, 0).contiguous()
         d_image = TorchCUDABridge.tensor_to_cuda_array(image)
         grid, tpb = _grid_3d(D, H, W)
-        _cone_3d_fdk_transpose_kernel[grid, tpb, stream](
+        _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,
@@ -635,13 +656,13 @@ def _analytical_transpose_raw(image, beam, geometry, parameters, isocenter, sino
         pitch, H, W, spacing = parameters
         d_image = TorchCUDABridge.tensor_to_cuda_array(image)
         grid, tpb = _grid_2d(W, H)
-        kernel = (_parallel_2d_fbp_transpose_kernel if beam == "parallel"
-                  else _fan_2d_fbp_transpose_kernel)
+        kernel = _resolve_kernel('_parallel_2d_fbp_transpose_kernel' if beam == "parallel"
+                                 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,
         )
-    return sinogram
+    return sinogram * scale
 
 
 class _AnalyticalGather(torch.autograd.Function):