python/dpt/spectral.py

1"""GPU-resident polychromatic primary transmission from material paths.23Importing this module is CPU-safe. Explicit host data preparation precedes CUDA4uploads; paths and signals have an explicit retained precision. Reverse products5recompute energy depths and preserve original inputs, with no P-by-K tape.6The API intentionally requires explicit VJP composition and rejects ambient tape7recording. It does not differentiate material interpolation or energy nodes.8"""910from __future__ import annotations1112# Public Python boundaries validate runtime inputs; workspace scratch stays module-owned.13# pyright: reportPrivateUsage=false, reportUnnecessaryIsInstance=false14import math15from dataclasses import dataclass, field16from typing import Any, Literal1718from dpt._runtime import DeviceContext, load_kernels, prepare_context, require_no_tape19from dpt.contracts import ContractError, NumericalError, finite_tuple, integer20from dpt.materials import Provenance212223# region book:spectrum-units24@dataclass(frozen=True, slots=True)25class Spectrum:26    """A supplied open-beam spectrum at the detector, before response losses.2728    Density values are photons/keV and require positive quadrature weights in29    keV. Bin values are already integrated expected photon populations and must30    not receive a second energy weight. No normalisation or inverse-square31    correction is guessed. Discrete lines may be included as integrated bins.32    """3334    energies_kev: tuple[float, ...]35    values: tuple[float, ...]36    representation: Literal["density", "bin-fluence"]37    provenance: Provenance38    quadrature_weights_kev: tuple[float, ...] | None = None3940    def __post_init__(self) -> None:41        finite_tuple(self.energies_kev, "energies_kev", positive=True)42        finite_tuple(self.values, "spectrum")43        if len(self.values) != len(self.energies_kev):44            raise ContractError("spectrum values and energies must have equal length")45        if any(a >= b for a, b in zip(self.energies_kev, self.energies_kev[1:], strict=False)):46            raise ContractError("spectrum energy nodes must be strictly increasing")47        if not isinstance(self.provenance, Provenance):48            raise ContractError("spectrum provenance is required")49        if self.representation == "density":50            if self.quadrature_weights_kev is None:51                raise ContractError("a density needs explicit energy quadrature weights")52            finite_tuple(self.quadrature_weights_kev, "quadrature weights", positive=True)53            if len(self.quadrature_weights_kev) != len(self.values):54                raise ContractError("one quadrature weight is required for each density value")55        elif self.representation == "bin-fluence":56            if self.quadrature_weights_kev is not None:57                raise ContractError("bin-integrated fluence must not be weighted twice")58        else:59            raise ContractError("spectrum representation must be density or bin-fluence")6061    @property62    def bin_fluence(self) -> tuple[float, ...]:63        """Materialise energy-integrated host values at the explicit ingestion boundary."""64        if self.quadrature_weights_kev is None:65            return self.values66        result = tuple(67            value * weight68            for value, weight in zip(self.values, self.quadrature_weights_kev, strict=True)69        )70        if any(not math.isfinite(value) for value in result):71            raise ContractError("energy integration overflowed")72        if any(73            source > 0 and target == 0 for source, target in zip(self.values, result, strict=True)74        ):75            raise ContractError("energy integration underflowed a positive bin population")76        return result777879# endregion book:spectrum-units808182@dataclass(frozen=True, slots=True)83class SpectralSpec:84    """Physical array layout and parameter capabilities, fixed at preparation.8586    The coefficients are total primary attenuation in mm⁻¹, with material paths87    in mm. ``input_description`` records the basis mixing/density assumption and88    detector acceptance used to turn source output into per-pixel fluence.89    Active weights are bin populations, not unconstrained normalised fractions.90    ``precision`` selects paths, means, image seeds and path gradients. Supplied91    coefficients, weights, responses and their gradients remain binary32.92    Retained binary64 rejects underflow of a nonzero exponential or intermediate93    product: later multiplication could otherwise conceal a representable tail.94    """9596    materials: int97    energies: int98    coefficients_provenance: tuple[Provenance, ...]99    spectrum_provenance: Provenance100    response_provenance: Provenance101    output_unit: str102    input_description: str103    shared_weights: bool = True104    shared_response: bool = True105    active_paths: bool = True106    active_weights: bool = False107    active_response: bool = False108    active_coefficients: bool = False109    reduction_groups: int = 64110    precision: Literal["float32", "float64"] = "float32"111112    def __post_init__(self) -> None:113        # Reverse keeps two M-element FP64 vectors per pixel. This bound is an114        # explicit implementation envelope, not an empirical performance claim.115        if self.precision not in ("float32", "float64"):116            raise ContractError("spectral precision must be float32 or float64")117        integer(self.materials, "materials", minimum=1, maximum=32)118        integer(self.energies, "energies", minimum=1, maximum=65536)119        integer(self.reduction_groups, "reduction_groups", minimum=1, maximum=256)120        if (121            not isinstance(self.coefficients_provenance, tuple)122            or len(self.coefficients_provenance) != self.materials123            or any(not isinstance(p, Provenance) for p in self.coefficients_provenance)124        ):125            raise ContractError("one immutable provenance record is required per material")126        if not isinstance(self.spectrum_provenance, Provenance):127            raise ContractError("spectrum provenance is required")128        if not isinstance(self.response_provenance, Provenance):129            raise ContractError("response provenance is required")130        if not self.output_unit.strip() or not self.input_description.strip():131            raise ContractError("spectral output units and input assumptions must be recorded")132        for name in (133            "shared_weights",134            "shared_response",135            "active_paths",136            "active_weights",137            "active_response",138            "active_coefficients",139        ):140            if type(getattr(self, name)) is not bool:141                raise ContractError(f"{name} must be a boolean")142143144@dataclass(slots=True)145class SpectralWorkspace:146    """One stream, persistent diagnostics and bounded shared-parameter reduction."""147148    spec: SpectralSpec149    max_pixels: int150    context: DeviceContext151    _kernels: Any = field(repr=False)152    _status: Any = field(repr=False)153    _empty: Any = field(repr=False)154    _empty_signal: Any = field(repr=False)155    _partials: Any = field(repr=False)156157    @property158    def dtype(self) -> Any:159        return (160            self.context.wp.float64 if self.spec.precision == "float64" else self.context.wp.float32161        )162163    @property164    def scratch_bytes(self) -> int:165        return 4 + 8 * int(self._partials.size)166167    def clear_status(self) -> None:168        with self.context.scope():169            self._status.zero_()170171    def check_status(self) -> None:172        """Explicit completion point; output/gradient overflow invalidates the pass."""173        self.context.wp.synchronize_stream(self.context.stream)174        code = int(self._status.numpy()[0])175        if code & 1:176            raise ContractError("spectral inputs contain nonfinite or negative values")177        if code & 4:178            raise NumericalError(179                "spectral intermediate underflow: unsupported exponent/product range"180            )181        if code & 2:182            raise NumericalError(183                "spectral output or gradient cannot be represented in its destination precision"184            )185186    def _check_values(self, arrays: list[tuple[Any, bool]]) -> None:187        self.clear_status()188        for values, nonnegative in arrays:189            if values.size:190                self.context.wp.launch(191                    self._kernels.get_value_check(192                        values.dtype == self.context.wp.float64, nonnegative193                    ),194                    dim=values.size,195                    inputs=[values],196                    outputs=[self._status],197                    stream=self.context.stream,198                    record_tape=False,199                )200        self.check_status()201202    def validate_inputs(203        self,204        coefficients: Any,205        weights: Any,206        response: Any,207        *,208        pixels: int,209        paths: Any = None,210        probability_response: bool = False,211        stream: Any = None,212    ) -> None:213        """Validate current supplied arrays without evaluating a spectral image.214215        A count observation requires probabilities; a general first-moment216        response only requires nonnegative values and its declared output unit.217        Omit paths when preparing fixed inputs before the material projection.218        """219        require_no_tape()220        ctx, spec = self.context, self.spec221        ctx.assert_stream(stream)222        integer(pixels, "pixels", maximum=self.max_pixels)223        if type(probability_response) is not bool:224            raise ContractError("probability_response must be a boolean")225        arrays = [226            (values, True)227            for _, values in _fixed_inputs(coefficients, weights, response, pixels, self)228        ]229        if paths is not None:230            ctx.array(paths, "paths", dtype=self.dtype, shape=(spec.materials * pixels,))231            arrays.append((paths, True))232        with ctx.scope():233            self._check_values(arrays)234            if probability_response and response.size:235                ctx.wp.launch(236                    self._kernels.check_probability,237                    dim=response.size,238                    inputs=[response],239                    outputs=[self._status],240                    stream=ctx.stream,241                    record_tape=False,242                )243                self.check_status()244245246def prepare_spectral(247    spec: SpectralSpec,248    *,249    max_pixels: int,250    device: str = "cuda:0",251    stream: Any = None,252) -> SpectralWorkspace:253    """Allocate scratch only; callers own source coefficients, weights and destinations."""254    integer(max_pixels, "max_pixels", maximum=(2**31 - 1) // max(spec.materials, spec.energies))255    context = prepare_context(device=device, stream=stream)256    kernels = load_kernels("dpt.kernels.spectral")257    parameters = 0258    if (spec.active_weights and spec.shared_weights) or (259        spec.active_response and spec.shared_response260    ):261        parameters = spec.energies262    if spec.active_coefficients:263        parameters = spec.materials * spec.energies264    with context.scope():265        status = context.wp.zeros(1, dtype=context.wp.int32, device=context.device)266        empty = context.wp.empty(0, dtype=context.wp.float32, device=context.device)267        empty_signal = context.wp.empty(268            0,269            dtype=context.wp.float64 if spec.precision == "float64" else context.wp.float32,270            device=context.device,271        )272        partials = context.wp.empty(273            parameters * spec.reduction_groups,274            dtype=context.wp.float64,275            device=context.device,276        )277    return SpectralWorkspace(278        spec, max_pixels, context, kernels, status, empty, empty_signal, partials279    )280281282def _inputs(283    paths: Any,284    coefficients: Any,285    weights: Any,286    response: Any,287    workspace: SpectralWorkspace,288) -> tuple[int, list[tuple[str, Any]]]:289    ctx, spec = workspace.context, workspace.spec290    size = ctx.array(paths, "paths", dtype=workspace.dtype)291    if size % spec.materials:292        raise ContractError("material-major paths must contain M*P values")293    pixels = size // spec.materials294    if pixels > workspace.max_pixels:295        raise ContractError("material paths exceed workspace pixel capacity")296    return pixels, [297        ("paths", paths),298        *_fixed_inputs(coefficients, weights, response, pixels, workspace),299    ]300301302def _fixed_inputs(303    coefficients: Any,304    weights: Any,305    response: Any,306    pixels: int,307    workspace: SpectralWorkspace,308) -> list[tuple[str, Any]]:309    ctx, spec = workspace.context, workspace.spec310    ctx.array(311        coefficients, "coefficients", dtype=ctx.wp.float32, shape=(spec.materials * spec.energies,)312    )313    ctx.array(314        weights,315        "weights",316        dtype=ctx.wp.float32,317        shape=(spec.energies * (1 if spec.shared_weights else pixels),),318    )319    ctx.array(320        response,321        "response",322        dtype=ctx.wp.float32,323        shape=(spec.energies * (1 if spec.shared_response else pixels),),324    )325    return [326        ("coefficients", coefficients),327        ("weights", weights),328        ("response", response),329    ]330331332# region book:spectral-call-contract333def spectral_signal(334    paths: Any,335    coefficients: Any,336    weights: Any,337    response: Any,338    *,339    out_mean: Any,340    workspace: SpectralWorkspace,341    stream: Any = None,342    validate: bool = True,343) -> None:344    """Compute sum_k weights[k,p]*response[k,p]*exp(-sum_m mu[m,k]*A[m,p]).345346    Arrays are flat contiguous CUDA buffers with the prepared precision. ``weights`` are347    integrated expected photons per bin, never a spectral density. Checked calls348    scan current inputs before writing and synchronise for final range status.349    ``validate=False`` promises valid current inputs and defers range status to350    ``workspace.check_status()``. No input may alias a destination.351    """352    require_no_tape()353    ctx, spec = workspace.context, workspace.spec354    ctx.assert_stream(stream)355    pixels, reads = _inputs(paths, coefficients, weights, response, workspace)356    ctx.array(out_mean, "out_mean", dtype=workspace.dtype, shape=(pixels,))357    ctx.disjoint(reads, [("out_mean", out_mean)])358    with ctx.scope():359        if validate:360            workspace.validate_inputs(361                coefficients, weights, response, pixels=pixels, paths=paths, stream=ctx.stream362            )363        if pixels:364            ctx.wp.launch(365                workspace._kernels.get_forward(366                    spec.materials,367                    spec.energies,368                    spec.shared_weights,369                    spec.shared_response,370                    spec.precision == "float64",371                ),372                dim=pixels,373                # Smaller blocks pack the register-heavy spectral work more evenly.374                # The reduction kernels retain their separate 256-lane layout.375                block_dim=128,376                inputs=[paths, coefficients, weights, response, pixels],377                outputs=[out_mean, workspace._status],378                stream=ctx.stream,379                record_tape=False,380            )381        if validate:382            workspace.check_status()383384385# endregion book:spectral-call-contract386387388def spectral_vjp(389    paths: Any,390    coefficients: Any,391    weights: Any,392    response: Any,393    *,394    seed: Any,395    out_grad_paths: Any = None,396    out_grad_weights: Any = None,397    out_grad_response: Any = None,398    out_grad_coefficients: Any = None,399    workspace: SpectralWorkspace,400    stream: Any = None,401    validate: bool = True,402) -> None:403    """Overwrite selected first-order products, recomputing from unchanged inputs.404405    Zero source/response values retain their one-sided algebraic derivatives;406    no derivative divides by a weight or by the rounded output. Shared parameter407    sums have a fixed split/tree order, and never use floating-point atomics.408    Caller-supplied output seeds remain unchanged. Higher derivatives and ambient409    tape recording raise rather than returning an incomplete product.410    """411    require_no_tape()412    ctx, spec = workspace.context, workspace.spec413    ctx.assert_stream(stream)414    pixels, reads = _inputs(paths, coefficients, weights, response, workspace)415    ctx.array(seed, "seed", dtype=workspace.dtype, shape=(pixels,))416    destinations = (417        ("paths", out_grad_paths, spec.active_paths, spec.materials * pixels),418        (419            "weights",420            out_grad_weights,421            spec.active_weights,422            spec.energies * (1 if spec.shared_weights else pixels),423        ),424        (425            "response",426            out_grad_response,427            spec.active_response,428            spec.energies * (1 if spec.shared_response else pixels),429        ),430        (431            "coefficients",432            out_grad_coefficients,433            spec.active_coefficients,434            spec.materials * spec.energies,435        ),436    )437    writes: list[tuple[str, Any]] = []438    for name, output, active, length in destinations:439        if output is not None:440            if not active:441                raise ContractError(f"{name} was declared fixed")442            ctx.array(443                output,444                f"out_grad_{name}",445                dtype=workspace.dtype if name == "paths" else ctx.wp.float32,446                shape=(length,),447            )448            writes.append((f"out_grad_{name}", output))449    if not writes:450        raise ContractError("request at least one active spectral gradient")451    ctx.disjoint([*reads, ("seed", seed)], writes)452    with ctx.scope():453        if validate:454            workspace._check_values([*((value, True) for _, value in reads), (seed, False)])455        pixel_weights = out_grad_weights is not None and not spec.shared_weights456        pixel_response = out_grad_response is not None and not spec.shared_response457        if pixels and (out_grad_paths is not None or pixel_weights or pixel_response):458            ctx.wp.launch(459                workspace._kernels.get_pixel_vjp(460                    spec.materials,461                    spec.energies,462                    spec.shared_weights,463                    spec.shared_response,464                    out_grad_paths is not None,465                    pixel_weights,466                    pixel_response,467                    spec.precision == "float64",468                ),469                dim=pixels,470                block_dim=128,471                inputs=[paths, coefficients, weights, response, seed, pixels],472                outputs=[473                    out_grad_paths if out_grad_paths is not None else workspace._empty_signal,474                    out_grad_weights if pixel_weights else workspace._empty,475                    out_grad_response if pixel_response else workspace._empty,476                    workspace._status,477                ],478                stream=ctx.stream,479                record_tape=False,480            )481        groups = min(spec.reduction_groups, max(1, (pixels + 255) // 256))482        for kind, output, shared, parameters in (483            (0, out_grad_weights, spec.shared_weights, spec.energies),484            (1, out_grad_response, spec.shared_response, spec.energies),485            (2, out_grad_coefficients, True, spec.materials * spec.energies),486        ):487            if output is None or not shared:488                continue489            if pixels:490                partial_kernel = (491                    workspace._kernels.get_coefficient_partials(492                        spec.materials,493                        spec.energies,494                        spec.shared_weights,495                        spec.shared_response,496                        spec.precision == "float64",497                    )498                    if kind == 2499                    else workspace._kernels.get_shared_partials(500                        spec.materials,501                        spec.energies,502                        spec.shared_weights,503                        spec.shared_response,504                        kind,505                        spec.precision == "float64",506                    )507                )508                ctx.wp.launch_tiled(509                    partial_kernel,510                    dim=(groups, spec.energies if kind == 2 else parameters),511                    block_dim=256,512                    inputs=[paths, coefficients, weights, response, seed, pixels, groups],513                    outputs=[workspace._partials, workspace._status],514                    stream=ctx.stream,515                    record_tape=False,516                )517                ctx.wp.launch(518                    workspace._kernels.finish_shared,519                    dim=parameters,520                    inputs=[workspace._partials, groups],521                    outputs=[output, workspace._status],522                    stream=ctx.stream,523                    record_tape=False,524                )525            else:526                output.zero_()527        if validate:528            workspace.check_status()529530531def spectral_bin_counts(532    paths: Any,533    coefficients: Any,534    weights: Any,535    detection_probability: Any,536    *,537    out_bin_counts: Any,538    workspace: SpectralWorkspace,539    stream: Any = None,540    validate: bool = True,541) -> None:542    """Materialise detected count means in explicit caller-owned (K,P) storage.543544    This observation-preparation boundary uses probabilities in [0,1], rather545    than an integrator's first-moment response. Feed these independent Poisson546    rates and deterministic per-photon scores to ``sample_compound_poisson``.547    It deliberately allocates no P*K tensor unless the caller requests that548    observation representation. This helper exposes no derivative or tape rule;549    differentiate the expected image with ``spectral_vjp`` instead.550    """551    require_no_tape()552    ctx, spec = workspace.context, workspace.spec553    ctx.assert_stream(stream)554    pixels, reads = _inputs(paths, coefficients, weights, detection_probability, workspace)555    ctx.array(556        out_bin_counts, "out_bin_counts", dtype=workspace.dtype, shape=(spec.energies * pixels,)557    )558    ctx.disjoint(reads, [("out_bin_counts", out_bin_counts)])559    with ctx.scope():560        if validate:561            workspace.validate_inputs(562                coefficients,563                weights,564                detection_probability,565                pixels=pixels,566                paths=paths,567                probability_response=True,568                stream=ctx.stream,569            )570        if pixels:571            ctx.wp.launch(572                workspace._kernels.get_bin_counts(573                    spec.materials,574                    spec.energies,575                    spec.shared_weights,576                    spec.shared_response,577                    spec.precision == "float64",578                ),579                dim=(spec.energies, pixels),580                inputs=[paths, coefficients, weights, detection_probability, pixels],581                outputs=[out_bin_counts, workspace._status],582                stream=ctx.stream,583                record_tape=False,584            )585        if validate:586            workspace.check_status()587