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

diffct/__init__.py

@@ -20,12 +20,18 @@ 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 (
+    transform_trajectory,
+    circular_trajectory_3d_parallel,
     circular_trajectory_3d,
     random_trajectory_3d,
     spiral_trajectory_3d,
@@ -58,7 +64,7 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.2.1'
+__version__ = '2.4.0'
 
 __all__ = [
     'Projector',
@@ -67,6 +73,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',
@@ -76,6 +90,8 @@ __all__ = [
     'FanBackprojectorFunction',
     'ConeProjectorFunction',
     'ConeBackprojectorFunction',
+    'transform_trajectory',
+    'circular_trajectory_3d_parallel',
     'circular_trajectory_3d',
     'random_trajectory_3d',
     'spiral_trajectory_3d',

diffct/_solver_state.py

@@ -0,0 +1,301 @@
+"""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
+from .distribution import control_status, preflight
+
+
+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, *, cache_key=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
+    error = None if error is None else ValueError("distributed solver agreement failed: " + str(error))
+    mismatch = "distributed solver agreement failed: inconsistent checkpoint iteration or controls"
+    partition = getattr(operator, "_partition", None)
+    if partition is not None:
+        partition.preflight(error, token, cache_key=cache_key, mismatch=mismatch)
+    else:
+        preflight(error, token, process_group=operator.process_group, world_size=world,
+                  mismatch=mismatch)
+
+
+def owns_values(operator, kind):
+    """Select one copy of each physical ray or voxel for scalar reductions."""
+    partition = getattr(operator, "partition", "views")
+    rank = getattr(operator, "rank", 0)
+    return (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
+    )
+
+
+def scalar_sum(operator, value, kind):
+    """Count physical rays/voxels once, including view/space replicas."""
+    world = getattr(operator, "world_size", 1)
+    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 float(value)
+
+
+def squared_norm(operator, value, kind, elements=1048576):
+    store = TensorStore(value) if isinstance(value, torch.Tensor) else value
+    total = 0.
+    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)
+
+
+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 callbacks_enabled(operator, callback):
+    """Check callback presence once, including callbacks on only some ranks."""
+    return bool(control_status((int(callback is not None),),
+                               process_group=getattr(operator, "process_group", None),
+                               world_size=getattr(operator, "world_size", 1))[0])
+
+
+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
+    failed, stop = control_status((int(error is not None), int(stop)),
+                                  process_group=getattr(operator, "process_group", None),
+                                  world_size=getattr(operator, "world_size", 1))
+    if failed:
+        agreement(operator, error)
+    return bool(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)