python/dpt/material_projection.py

1"""Fixed-density material-fraction paths using the canonical sampled-field projector.23Fields are material-major, flat FP32 arrays with M*V values; paths have M*P4values in millimetres, using the supplied ProjectionSpec precision for both5paths and their cotangents. Field-gradient storage remains FP32.6Each basis field is a volume fraction. Fractions are7non-negative and sum to at most one at every sample; the remainder is vacuum.8The FP64 sum admits 2^-24 excess for independently rounded binary32 fractions.9The explicit nonnegative domain instead accepts dimensionless equivalent-basis10fields without a sum cap. It uses the same integration and derivatives; basis11scales belong in the supplied attenuation coefficients. There is no hidden12normalisation or inferred CT-to-material conversion.13"""1415from __future__ import annotations1617from dataclasses import dataclass, field18from typing import Any, Literal1920from dpt._runtime import DeviceContext, load_kernels, require_no_tape21from dpt.contracts import ContractError, integer22from dpt.geometry import DetectorGeometry23from dpt.projection import ProjectionSpec, ProjectionWorkspace, prepare_projection24from dpt.projection import project_optical_depth as _project25from dpt.projection import projection_vjp as _vjp26from dpt.volumes import GridSpec272829@dataclass(slots=True)30class MaterialProjectionWorkspace:31    """Prepared views retain their caller-owned parent buffers for the whole solve.3233    M serial launches reuse a single projection workspace, avoiding duplicated34    ray geometry and quadrature code. This is an execution candidate, not a35    measured throughput claim. No per-evaluation views or device arrays are made.36    """3738    projection: ProjectionWorkspace39    materials: int40    fields: Any41    paths: Any42    adj_paths: Any43    grad_fields: Any44    field_views: tuple[Any, ...]45    path_views: tuple[Any, ...]46    seed_views: tuple[Any, ...]47    gradient_views: tuple[Any, ...]48    kernels: Any = field(repr=False)49    status: Any = field(repr=False)50    field_domain: Literal["fractions", "nonnegative"] = "fractions"5152    @property53    def context(self) -> DeviceContext:54        return self.projection.context5556    @property57    def scratch_bytes(self) -> int:58        return self.projection.scratch_bytes + 45960    def clear_status(self) -> None:61        self.projection.clear_status()62        with self.context.scope():63            self.status.zero_()6465    def check_status(self) -> None:66        self.projection.check_status()67        if int(self.status.numpy()[0]):68            name = "fractions" if self.field_domain == "fractions" else "nonnegative fields"69            raise ContractError(f"invalid material {name} or path cotangents")7071    def validate_inputs(self, *, seeds: bool = False, stream: Any = None) -> None:72        """Check current bound fractions and optional cotangents without projecting.7374        Fraction-domain roundoff permits a total up to 1 + 2^-24. Nonnegative75        fields have no sum cap. Values are never changed. This checks no pose76        and does not certify later field mutations.77        """78        require_no_tape()79        self.context.assert_stream(stream)80        if seeds and self.adj_paths is None:81            raise ContractError("material validation requires a bound path cotangent")82        _validate(self, seeds=seeds)838485def prepare_material_projection(86    grid: GridSpec,87    geometry: DetectorGeometry,88    spec: ProjectionSpec,89    *,90    materials: int,91    fields: Any,92    out_paths: Any,93    adj_paths: Any = None,94    out_fields: Any = None,95    device: str = "cuda:0",96    stream: Any = None,97    field_domain: Literal["fractions", "nonnegative"] = "fractions",98) -> MaterialProjectionWorkspace:99    """Bind fields; nonnegative equivalent-basis fields have no fraction sum cap.100101    Both domains use dimensionless fields, paths in mm and supplied basis102    attenuation in mm^-1. The caller records any mass-density normalisation.103    ``spec.precision`` selects FP32 or FP64 paths and path cotangents; fields104    and their accumulated volume gradients remain FP32.105    """106    if field_domain not in ("fractions", "nonnegative"):107        raise ContractError("field_domain must be fractions or nonnegative")108    integer(materials, "materials", minimum=1, maximum=32)109    integer(materials * grid.voxels, "material field length", minimum=1)110    integer(materials * geometry.pixels, "material path length", minimum=1)111    projection = prepare_projection(grid, geometry, spec, device=device, stream=stream)112    ctx, wp = projection.context, projection.context.wp113    ctx.array(fields, "fields", dtype=wp.float32, shape=(materials * grid.voxels,))114    path_dtype = wp.float64 if spec.precision == "float64" else wp.float32115    ctx.array(out_paths, "out_paths", dtype=path_dtype, shape=(materials * geometry.pixels,))116    reads = [("fields", fields)]117    writes = [("out_paths", out_paths)]118    if adj_paths is not None:119        ctx.array(adj_paths, "adj_paths", dtype=path_dtype, shape=(materials * geometry.pixels,))120        reads.append(("adj_paths", adj_paths))121    if out_fields is not None:122        if not spec.active_volume:123            raise ContractError("material fields are fixed in this projection specification")124        if adj_paths is None:125            raise ContractError("material field gradients require a bound path cotangent")126        ctx.array(out_fields, "out_fields", dtype=wp.float32, shape=(materials * grid.voxels,))127        writes.append(("out_fields", out_fields))128    ctx.disjoint(reads, writes)129    with ctx.scope():130        status = wp.zeros(1, dtype=wp.int32, device=ctx.device)131    field_views = tuple(fields[m * grid.voxels : (m + 1) * grid.voxels] for m in range(materials))132    path_views = tuple(133        out_paths[m * geometry.pixels : (m + 1) * geometry.pixels] for m in range(materials)134    )135    seeds = (136        tuple(adj_paths[m * geometry.pixels : (m + 1) * geometry.pixels] for m in range(materials))137        if adj_paths is not None138        else ()139    )140    gradients = (141        tuple(out_fields[m * grid.voxels : (m + 1) * grid.voxels] for m in range(materials))142        if out_fields is not None143        else ()144    )145    return MaterialProjectionWorkspace(146        projection,147        materials,148        fields,149        out_paths,150        adj_paths,151        out_fields,152        field_views,153        path_views,154        seeds,155        gradients,156        load_kernels("dpt.kernels.material_projection"),157        status,158        field_domain,159    )160161162def _validate(workspace: MaterialProjectionWorkspace, *, seeds: bool) -> None:163    ctx, projection = workspace.context, workspace.projection164    workspace.clear_status()165    ctx.wp.launch(166        workspace.kernels.validate_fractions167        if workspace.field_domain == "fractions"168        else workspace.kernels.validate_nonnegative,169        dim=projection.grid.voxels,170        inputs=[workspace.fields, workspace.materials, projection.grid.voxels],171        outputs=[workspace.status],172        stream=ctx.stream,173        record_tape=False,174    )175    if seeds:176        ctx.wp.launch(177            workspace.kernels.validate_seeds(projection.spec.precision),178            dim=workspace.materials * projection.geometry.pixels,179            inputs=[workspace.adj_paths],180            outputs=[workspace.status],181            stream=ctx.stream,182            record_tape=False,183        )184    workspace.check_status()185186187# region book:material-path-composition188def project_material_paths(189    pose: Any,190    *,191    workspace: MaterialProjectionWorkspace,192    stream: Any = None,193    validate: bool = True,194) -> None:195    """Overwrite material-major path lengths in mm, using one canonical projector.196197    All material data stay on CUDA. The host loop schedules M operations and198    never visits a voxel, pixel or quadrature sample. Explicit VJPs compose199    with spectral_signal; ambient tape recording is rejected.200    """201    require_no_tape()202    ctx, projection = workspace.context, workspace.projection203    ctx.assert_stream(stream)204    ctx.array(pose, "pose", dtype=ctx.wp.float64, shape=(12,))205    ctx.disjoint([("pose", pose)], [("paths", workspace.paths)])206    if validate:207        workspace.validate_inputs(stream=ctx.stream)208    for material in range(workspace.materials):209        _project(210            workspace.field_views[material],211            pose,212            workspace=projection,213            out_L=workspace.path_views[material],214            stream=ctx.stream,215            validate=validate and material == 0,216        )217    if validate:218        workspace.check_status()219220221# endregion book:material-path-composition222223224def material_projection_vjp(225    pose: Any,226    *,227    workspace: MaterialProjectionWorkspace,228    out_pose: Any = None,229    stream: Any = None,230    validate: bool = True,231    accumulate: bool = False,232) -> None:233    """Accumulate material cotangents into one local-right six-coordinate pose VJP.234235    The optional bound field gradient is an ambient derivative. Feasible236    fraction perturbations must still respect non-negativity and the simplex;237    an optimiser must provide its own declared constrained parameterisation.238    Field scatter order follows the canonical projector's atomic contract.239    """240    require_no_tape()241    ctx, projection = workspace.context, workspace.projection242    ctx.assert_stream(stream)243    ctx.array(pose, "pose", dtype=ctx.wp.float64, shape=(12,))244    if not workspace.seed_views:245        raise ContractError("prepare a path cotangent before requesting material VJPs")246    writes: list[tuple[str, Any]] = []247    if out_pose is not None:248        if not projection.spec.active_pose:249            raise ContractError("pose is fixed in this projection specification")250        ctx.array(out_pose, "out_pose", dtype=ctx.wp.float64, shape=(6,))251        writes.append(("out_pose", out_pose))252    if workspace.grad_fields is not None:253        writes.append(("out_fields", workspace.grad_fields))254    if not writes:255        raise ContractError("request a pose or material field gradient")256    ctx.disjoint(257        [258            ("fields", workspace.fields),259            ("paths", workspace.paths),260            ("adj_paths", workspace.adj_paths),261            ("pose", pose),262        ],263        writes,264    )265    if validate:266        workspace.validate_inputs(seeds=True, stream=ctx.stream)267    for material in range(workspace.materials):268        # Pose contributions share one destination; distinct field slices must269        # each honour the caller's overwrite/accumulate request independently.270        field_gradient = workspace.gradient_views[material] if workspace.gradient_views else None271        if out_pose is not None and field_gradient is not None and material > 0 and not accumulate:272            with ctx.scope():273                field_gradient.zero_()274        _vjp(275            workspace.field_views[material],276            pose,277            adj_L=workspace.seed_views[material],278            workspace=projection,279            out_pose=out_pose,280            out_mu=field_gradient,281            stream=ctx.stream,282            validate=validate and material == 0,283            accumulate=accumulate or (out_pose is not None and material > 0),284        )285    if validate:286        workspace.check_status()287