python/dpt/projection.py

1"""Finite-segment optical-depth projection with caller-owned device buffers.23Coordinates and quadrature accumulators are FP64; field samples remain FP32.4Depth and depth-cotangent storage is explicitly FP32 (default) or FP64.5Numerical acceptance and performance must be established on CUDA;6this interface does not equate a finite result with an accurate projection.7"""89from __future__ import annotations1011# Mathematical output names follow the optical-depth contract.12# ruff: noqa: N80313# Module-owned workspace scratch stays private to this implementation.14# pyright: reportPrivateUsage=false15import math16from dataclasses import dataclass, field17from typing import Any, Literal1819from dpt._reductions import ReductionTree, prepare_reduction20from dpt._runtime import DeviceContext, ensure_tape, load_kernels, prepare_context, require_no_tape21from dpt.contracts import ContractError, NumericalError, integer22from dpt.geometry import DetectorGeometry23from dpt.volumes import GridSpec2425_BLOCK = 128262728@dataclass(frozen=True, slots=True)29class ProjectionSpec:30    """Declared ray integration and explicitly supported active inputs.3132    Midpoint uses samples_per_ray. Cell Gauss ignores this count and splits at33    every interpolation plane, integrating each cubic segment with two nodes.3435    Grid/acquisition calibration is fixed. Pose derivatives include active36    intersection bounds. At ties, grazing rays and interpolation knots the37    returned branch derivative does not assert differentiability.38    """3940    samples_per_ray: int41    active_pose: bool = True42    active_volume: bool = False43    precision: Literal["float32", "float64"] = "float32"44    integration: Literal["midpoint", "cell_gauss"] = "midpoint"4546    def __post_init__(self) -> None:47        integer(self.samples_per_ray, "samples_per_ray", minimum=1)48        if self.integration not in ("midpoint", "cell_gauss"):49            raise ContractError("projection integration must be midpoint or cell_gauss")50        if self.precision not in ("float32", "float64"):51            raise ContractError("projection precision must be float32 or float64")52        if type(self.active_pose) is not bool or type(self.active_volume) is not bool:53            raise ContractError("active input flags must be booleans")545556@dataclass(slots=True)57class ProjectionWorkspace:58    """Persistent O(12 P/128) FP64 reduction storage, one owning CUDA stream.5960    Original volume and pose buffers must remain immutable through backward.61    One recorded forward may be outstanding per workspace. Use another62    workspace for concurrent/re-entrant evaluations. No sample history is kept.63    """6465    grid: GridSpec66    geometry: DetectorGeometry67    spec: ProjectionSpec68    context: DeviceContext69    _kernels: Any = field(repr=False)70    _geometry_kernels: Any = field(repr=False)71    _configuration: Any = field(repr=False)72    _status: Any = field(repr=False)73    _empty: Any = field(repr=False)74    _reduction: ReductionTree | None = field(repr=False)75    _empty_partials: Any = field(repr=False)76    _recorded_tape: Any = field(default=None, repr=False)7778    @property79    def dtype(self) -> Any:80        """Declared depth and depth-cotangent storage; field storage stays FP32."""81        return getattr(self.context.wp, self.spec.precision)8283    @property84    def device(self) -> Any:85        return self.context.device8687    @property88    def stream(self) -> Any:89        return self.context.stream9091    @property92    def scratch_bytes(self) -> int:93        return 4 + (self._reduction.scratch_bytes if self._reduction is not None else 0)9495    def clear_status(self) -> None:96        with self.context.scope():97            self._status.zero_()9899    def check_status(self) -> None:100        self.context.wp.synchronize_stream(self.stream)101        code = int(self._status.numpy()[0])102        if code & 1:103            raise ContractError("projection input violates finite non-negative field or rigid pose")104        if code:105            raise NumericalError("non-finite projection or gradient; discard these outputs")106107    def discard_recording(self, tape: Any) -> None:108        """Release a forward abandoned by tape.reset(), after discarding its graph."""109        if self._recorded_tape is not None and self._recorded_tape is not tape:110            raise ContractError("this workspace belongs to another tape")111        if tape.launches:112            raise ContractError("reset the abandoned tape before releasing its input lifetime")113        self._recorded_tape = None114115116def prepare_projection(117    grid: GridSpec,118    geometry: DetectorGeometry,119    spec: ProjectionSpec,120    *,121    device: str = "cuda:0",122    stream: Any = None,123) -> ProjectionWorkspace:124    """Prepare fixed metadata and device scratch; never allocate image outputs."""125    context = prepare_context(device=device, stream=stream)126    wp = context.wp127    kernels = load_kernels("dpt.kernels.projection")128    geometry_kernels = load_kernels("dpt.kernels.geometry")129    config = kernels.Configuration()130    config.shape = wp.vec3i(grid.shape[2], grid.shape[1], grid.shape[0])131    # diag(1/h) G^T: spacing follows grid axes, not object/world coordinates.132    inverse_grid = tuple(133        grid.orientation[3 * j + i] / grid.spacing_mm[i] for i in range(3) for j in range(3)134    )135    if not all(math.isfinite(value) for value in inverse_grid):136        raise ContractError("grid spacing overflows the FP64 coordinate map")137    for row in (0, geometry.shape[0] - 1):138        for column in (0, geometry.shape[1] - 1):139            endpoint = geometry.pixel_centre(row, column)140            if not all(math.isfinite(value) for value in endpoint) or not math.isfinite(141                math.dist(endpoint, geometry.source_mm)142            ):143                raise ContractError("detector extent exceeds the FP64 finite-ray range")144    config.object_to_grid = wp.mat33d(*inverse_grid)145    config.grid_origin = wp.vec3d(*grid.origin_mm)146    config.source = wp.vec3d(*geometry.source_mm)147    config.detector_origin = wp.vec3d(*geometry.origin_mm)148    config.column_step = wp.vec3d(*(x * geometry.spacing_mm[0] for x in geometry.u))149    config.row_step = wp.vec3d(*(x * geometry.spacing_mm[1] for x in geometry.v))150    config.width = geometry.shape[1]151    config.pixels = geometry.pixels152    config.samples = spec.samples_per_ray153    reduction = None154    with context.scope():155        status = wp.zeros(1, dtype=wp.int32, device=context.device)156        empty = wp.empty(0, dtype=wp.float32, device=context.device)157        empty_partials = wp.empty(0, dtype=wp.float64, device=context.device)158        if spec.active_pose:159            reduction = prepare_reduction(context, geometry.pixels, tile=_BLOCK, components=12)160    return ProjectionWorkspace(161        grid,162        geometry,163        spec,164        context,165        kernels,166        geometry_kernels,167        config,168        status,169        empty,170        reduction,171        empty_partials,172    )173174175def _inputs(mu: Any, pose: Any, workspace: ProjectionWorkspace, stream: Any) -> None:176    ctx = workspace.context177    ctx.assert_stream(stream)178    ctx.array(mu, "mu", dtype=ctx.wp.float32, shape=(workspace.grid.voxels,))179    ctx.array(pose, "pose", dtype=ctx.wp.float64, shape=(12,))180181182def _validate(mu: Any, pose: Any, workspace: ProjectionWorkspace, seed: Any = None) -> None:183    wp = workspace.context.wp184    workspace.clear_status()185    wp.launch(186        workspace._kernels.validate_nonnegative,187        dim=workspace.grid.voxels,188        inputs=[mu],189        outputs=[workspace._status],190        stream=workspace.stream,191        record_tape=False,192    )193    wp.launch(194        workspace._geometry_kernels.validate_pose,195        dim=1,196        inputs=[pose],197        outputs=[workspace._status],198        stream=workspace.stream,199        record_tape=False,200    )201    if seed is not None:202        wp.launch(203            workspace._kernels.get_validate_finite(workspace.spec.precision == "float64"),204            dim=workspace.geometry.pixels,205            inputs=[seed],206            outputs=[workspace._status],207            stream=workspace.stream,208            record_tape=False,209        )210    workspace.check_status()211212213# region book:projection-public-operator214def project_optical_depth(215    mu: Any,216    pose: Any,217    *,218    workspace: ProjectionWorkspace,219    out_L: Any,220    stream: Any = None,221    tape: Any = None,222    validate: bool = True,223) -> None:224    """Overwrite dimensionless optical depth for every finite detector ray.225226    mu is a flat non-negative FP32 field in inverse mm. pose is twelve FP64227    values (R_WO row-major, t_WO in mm). out_L is flat row-major detector228    storage in the precision declared by ProjectionSpec. A checked call229    synchronises domain diagnostics before writing and230    range diagnostics afterwards. Unchecked calls promise valid current inputs231    and require check_status() at the next acceptance checkpoint.232    """233    ensure_tape(tape)234    _inputs(mu, pose, workspace, stream)235    if workspace._recorded_tape is not None:236        raise ContractError("finish or discard the outstanding projection tape before reuse")237    ctx = workspace.context238    ctx.array(out_L, "out_L", dtype=workspace.dtype, shape=(workspace.geometry.pixels,))239    ctx.disjoint([("mu", mu), ("pose", pose)], [("out_L", out_L)])240    if tape is not None:241        _check_tape_arrays(mu, pose, out_L, workspace)242    if validate:243        _validate(mu, pose, workspace)244    ctx.wp.launch(245        workspace._kernels.get_forward(246            workspace.spec.precision == "float64", workspace.spec.integration == "cell_gauss"247        ),248        dim=workspace.geometry.pixels,249        inputs=[mu, pose, workspace._configuration],250        outputs=[out_L, workspace._status],251        stream=workspace.stream,252        block_dim=_BLOCK,253        record_tape=False,254    )255    if validate:256        workspace.check_status()257    if tape is not None:258        _record(tape, mu, pose, out_L, workspace)259260261# endregion book:projection-public-operator262263264def _check_tape_arrays(265    mu: Any, pose: Any, output: Any, workspace: ProjectionWorkspace266) -> list[Any]:267    arrays = [output]268    if workspace.spec.active_pose:269        arrays.append(pose)270    if workspace.spec.active_volume:271        arrays.append(mu)272    if len(arrays) == 1:273        raise ContractError("recording requires at least one active input")274    if any(array.grad is None for array in arrays):275        raise ContractError("active tape arrays require requires_grad=True at allocation")276    workspace.context.disjoint(277        [("mu", mu), ("pose", pose), ("output", output)],278        [("gradient", array.grad) for array in arrays],279    )280    return arrays281282283def _record(tape: Any, mu: Any, pose: Any, output: Any, workspace: ProjectionWorkspace) -> None:284    arrays = _check_tape_arrays(mu, pose, output, workspace)285    workspace._recorded_tape = tape286    if workspace.context.wp.config.verify_autograd_array_access:287        mu.mark_read()288        pose.mark_read()289        output.mark_write()290        # Register fixed inputs too, so Tape.reset() clears debug read markers.291        tape.record_launch(292            workspace._kernels.get_dependency_marker(workspace.spec.precision == "float64"),293            dim=0,294            max_blocks=0,295            inputs=[mu, pose],296            outputs=[output],297            device=workspace.device,298            block_dim=_BLOCK,299        )300301    def backward() -> None:302        # Array storage has twelve coordinates. Its cotangent is NOT the303        # six-coordinate right tangent accepted by the standalone optimiser.304        projection_vjp(305            mu,306            pose,307            adj_L=output.grad,308            workspace=workspace,309            out_pose_matrix=pose.grad if workspace.spec.active_pose else None,310            out_mu=mu.grad if workspace.spec.active_volume else None,311            accumulate=True,312            validate=False,313            _from_tape=True,314        )315        with workspace.context.scope():316            if not output.retain_grad:317                output.grad.zero_()318        workspace._recorded_tape = None319320    tape.record_func(backward, arrays)321322323def projection_vjp(324    mu: Any,325    pose: Any,326    *,327    adj_L: Any,328    workspace: ProjectionWorkspace,329    out_pose: Any = None,330    out_pose_matrix: Any = None,331    out_mu: Any = None,332    stream: Any = None,333    validate: bool = True,334    accumulate: bool = False,335    _from_tape: bool = False,336) -> None:337    """First-order discrete VJP, with explicit tangent or storage coordinates.338339    out_pose is six FP64 values for a *local right* SE(3) increment in mm/rad.340    out_pose_matrix is twelve FP64 unconstrained storage partials, for a tape341    chain rule. They are mutually exclusive. Convert local right gradients342    using geometry.chart_gradient for a fixed-anchor optimisation chart.343    out_mu is FP32 and uses non-deterministic-order scatter atomics. No grid,344    acquisition, higher-order or non-smooth boundary derivative is promised.345    """346    require_no_tape()347    _inputs(mu, pose, workspace, stream)348    if workspace._recorded_tape is not None and not _from_tape:349        raise ContractError("the outstanding tape owns the projection inputs")350    if out_pose is not None and out_pose_matrix is not None:351        raise ContractError("request either a right-tangent or a storage pose gradient")352    target = out_pose if out_pose is not None else out_pose_matrix353    if target is None and out_mu is None:354        raise ContractError("request at least one projection gradient")355    if target is not None and not workspace.spec.active_pose:356        raise ContractError("pose is fixed in this workspace")357    if out_mu is not None and not workspace.spec.active_volume:358        raise ContractError("volume is fixed in this workspace")359    ctx, wp = workspace.context, workspace.context.wp360    ctx.array(adj_L, "adj_L", dtype=workspace.dtype, shape=(workspace.geometry.pixels,))361    components = 12 if out_pose_matrix is not None else 6362    writes: list[tuple[str, Any]] = []363    if target is not None:364        ctx.array(target, "pose gradient", dtype=wp.float64, shape=(components,))365        writes.append(("pose gradient", target))366    if out_mu is not None:367        ctx.array(out_mu, "out_mu", dtype=wp.float32, shape=(workspace.grid.voxels,))368        writes.append(("out_mu", out_mu))369    ctx.disjoint([("mu", mu), ("pose", pose), ("adj_L", adj_L)], writes)370    if validate:371        _validate(mu, pose, workspace, adj_L)372    if out_mu is not None and not accumulate:373        with ctx.scope():374            out_mu.zero_()375    count = (workspace.geometry.pixels + _BLOCK - 1) // _BLOCK376    wp.launch_tiled(377        workspace._kernels.get_vjp(378            out_mu is not None,379            components == 12,380            target is not None,381            double_seed=workspace.spec.precision == "float64",382            cell_gauss=workspace.spec.integration == "cell_gauss",383        ),384        dim=count,385        inputs=[mu, pose, workspace._configuration, adj_L],386        outputs=[387            out_mu if out_mu is not None else workspace._empty,388            workspace._reduction.partials[0]389            if workspace._reduction is not None390            else workspace._empty_partials,391            workspace._status,392        ],393        stream=workspace.stream,394        device=workspace.device,395        block_dim=_BLOCK,396        record_tape=False,397    )398    if target is not None:399        assert workspace._reduction is not None400        workspace._reduction.finish(401            count,402            target,403            workspace._status,404            components=components,405            accumulate=accumulate,406        )407    if out_mu is not None:408        # Atomic accumulation may overflow after individually finite increments.409        wp.launch(410            workspace._kernels.validate_finite,411            dim=workspace.grid.voxels,412            inputs=[out_mu],413            outputs=[workspace._status],414            stream=workspace.stream,415            record_tape=False,416        )417    if validate:418        workspace.check_status()419420421# region book:projection-pose-sensitivities422def projection_pose_sensitivities(423    mu: Any,424    pose: Any,425    depth_seeds: Any,426    *,427    out_jacobian: Any,428    workspace: ProjectionWorkspace,429    stream: Any = None,430    validate: bool = True,431) -> None:432    """Export seeded per-ray local pose derivatives for bounded offline diagnostics.433434    ``out_jacobian[6*p + axis]`` is ``depth_seeds[p] * dL[p]/d(delta[axis])``435    for ``pose @ exp(delta^)`` at zero, ordered tx, ty, tz, rx, ry, rz in mm/rad.436    All arrays are caller-owned, contiguous and on the workspace's CUDA device:437    seeds follow the declared depth precision, shape ``(pixels,)``; output is438    FP64 ``(6*pixels,)``. Ones export439    the optical-depth Jacobian; canonical transmission cotangents export the440    corresponding signal sensitivities. Missed rays and zero seeds give zero.441442    The existing per-ray VJP calculation writes before reduction in O(P S)443    work, where S is samples per ray. This call allocates no device buffers and444    performs no image transfers. Checked calls synchronise for input/output445    status; unchecked calls require valid current inputs and a later446    ``workspace.check_status()``. The caller retains the O(6 P) diagnostic,447    unlike production optimisation. Ambient tapes and outstanding projection448    passes are rejected. Branch derivatives at knots/ties/grazing rays do not449    assert smoothness. Only the current local right chart is differentiated.450    """451    require_no_tape()452    _inputs(mu, pose, workspace, stream)453    if workspace._recorded_tape is not None:454        raise ContractError("the outstanding tape owns the projection inputs")455    if not workspace.spec.active_pose:456        raise ContractError("pose is fixed in this workspace")457    ctx, wp = workspace.context, workspace.context.wp458    pixels = workspace.geometry.pixels459    ctx.array(depth_seeds, "depth_seeds", dtype=workspace.dtype, shape=(pixels,))460    ctx.array(out_jacobian, "out_jacobian", dtype=wp.float64, shape=(6 * pixels,))461    ctx.disjoint(462        [("mu", mu), ("pose", pose), ("depth_seeds", depth_seeds)],463        [("out_jacobian", out_jacobian)],464    )465    if validate:466        _validate(mu, pose, workspace, depth_seeds)467    wp.launch_tiled(468        workspace._kernels.get_vjp(469            False,470            False,471            True,472            per_ray=True,473            double_seed=workspace.spec.precision == "float64",474            cell_gauss=workspace.spec.integration == "cell_gauss",475        ),476        dim=(pixels + _BLOCK - 1) // _BLOCK,477        inputs=[mu, pose, workspace._configuration, depth_seeds],478        outputs=[workspace._empty, out_jacobian, workspace._status],479        stream=workspace.stream,480        device=workspace.device,481        block_dim=_BLOCK,482        record_tape=False,483    )484    if validate:485        workspace.check_status()486487488# endregion book:projection-pose-sensitivities489