python/dpt/detector.py

1"""Deterministic detector operators and explicitly separate observation models.23Expected response, spatial spreading, electronic calibration and random draws4are distinct operations with caller-owned buffers. An energy-integrating mean5is not passed to a Poisson count sampler by the library. Imports are CPU-safe;6all numerical image work uses the single CUDA implementation.7"""89from __future__ import annotations1011# Public Python boundaries validate runtime inputs; workspace scratch stays module-owned.12# pyright: reportPrivateUsage=false, reportUnnecessaryIsInstance=false13import math14from dataclasses import dataclass, field15from typing import Any, Literal1617from dpt._runtime import DeviceContext, load_kernels, prepare_context, require_no_tape18from dpt.contracts import ContractError, NumericalError, finite_scalar, finite_tuple, integer19from dpt.materials import Provenance202122@dataclass(slots=True)23class DetectorWorkspace:24    """Reusable diagnostics and scalar-reduction scratch on one CUDA stream."""2526    max_pixels: int27    context: DeviceContext28    _kernels: Any = field(repr=False)29    _checks: Any = field(repr=False)30    _status: Any = field(repr=False)31    _partials: Any = field(repr=False)3233    @property34    def reduction_groups(self) -> int:35        """Read-only launch bound derived from the allocated scratch capacity."""36        return int(self._partials.size)3738    @property39    def scratch_bytes(self) -> int:40        return 4 + 8 * int(self._partials.size)4142    def clear_status(self) -> None:43        with self.context.scope():44            self._status.zero_()4546    def check_status(self) -> None:47        """Synchronise and reject a numerical failure; no failed draw is an observation."""48        self.context.wp.synchronize_stream(self.context.stream)49        code = int(self._status.numpy()[0])50        if code & 1:51            raise ContractError("detector device inputs violate their finite/domain/rate contract")52        if code & 2:53            raise NumericalError("detector output or gradient overflows its declared precision")54        if code & 4:55            raise NumericalError("Poisson draw budget exhausted; discard the entire realisation")5657    def _values(self, arrays: list[tuple[Any, str]]) -> None:58        self.clear_status()59        for array, domain in arrays:60            if not array.size:61                continue62            kernel = {63                "finite": self._checks.check_finite,64                "nonnegative": self._checks.check_nonnegative,65                "positive": self._kernels.check_positive,66                "rate": self._kernels.check_rates,67            }[domain]68            self.context.wp.launch(69                kernel,70                dim=array.size,71                inputs=[array],72                outputs=[self._status],73                stream=self.context.stream,74                record_tape=False,75            )76        self.check_status()7778    def _pixels(self, values: Any, name: str) -> int:79        count = self.context.array(values, name, dtype=self.context.wp.float32)80        if count > self.max_pixels:81            raise ContractError(f"{name} exceeds workspace capacity")82        return count838485def prepare_detector(86    *,87    max_pixels: int,88    device: str = "cuda:0",89    stream: Any = None,90    reduction_groups: int = 64,91) -> DetectorWorkspace:92    """Prepare O(groups) device scratch; do not allocate any detector image."""93    integer(max_pixels, "max_pixels")94    integer(reduction_groups, "reduction_groups", minimum=1, maximum=256)95    context = prepare_context(device=device, stream=stream)96    kernels = load_kernels("dpt.kernels.detector")97    checks = load_kernels("dpt.kernels.spectral")98    with context.scope():99        status = context.wp.zeros(1, dtype=context.wp.int32, device=context.device)100        partials = context.wp.empty(101            reduction_groups, dtype=context.wp.float64, device=context.device102        )103    return DetectorWorkspace(max_pixels, context, kernels, checks, status, partials)104105106# region book:calibration-identifiability107@dataclass(frozen=True, slots=True)108class CalibrationSpec:109    """y = gain * exposure * mean + offset, in the declared output unit.110111    Gain is strictly positive; exposure is non-negative; offset is any finite112    electronic baseline. A fixed flat-field may vary by pixel. Optimised gain113    and offset are shared scalars: unconstrained correction images are excluded.114    Exposure and gain cannot both be active because their scale is unidentifiable.115    """116117    output_unit: str118    shared_gain: bool = True119    shared_offset: bool = True120    active_mean: bool = True121    active_gain: bool = False122    active_exposure: bool = False123    active_offset: bool = False124125    def __post_init__(self) -> None:126        if not self.output_unit.strip():127            raise ContractError("calibration output unit is required")128        for name in (129            "shared_gain",130            "shared_offset",131            "active_mean",132            "active_gain",133            "active_exposure",134            "active_offset",135        ):136            if type(getattr(self, name)) is not bool:137                raise ContractError(f"{name} must be a boolean")138        if self.active_gain and self.active_exposure:139            raise ContractError("fix gain or exposure before fitting the other scale")140        if self.active_gain and not self.shared_gain:141            raise ContractError("a fitted gain must be a shared acquisition parameter")142        if self.active_offset and not self.shared_offset:143            raise ContractError("a fitted offset must be a shared acquisition parameter")144145146# endregion book:calibration-identifiability147148149def _calibration_inputs(150    mean: Any,151    gain: Any,152    exposure: Any,153    offset: Any,154    spec: CalibrationSpec,155    workspace: DetectorWorkspace,156) -> tuple[int, list[tuple[str, Any]]]:157    ctx = workspace.context158    pixels = workspace._pixels(mean, "mean")159    for name, array, size in (160        ("gain", gain, 1 if spec.shared_gain else pixels),161        ("exposure", exposure, 1),162        ("offset", offset, 1 if spec.shared_offset else pixels),163    ):164        ctx.array(array, name, dtype=ctx.wp.float32, shape=(size,))165    return pixels, [("mean", mean), ("gain", gain), ("exposure", exposure), ("offset", offset)]166167168def calibrate(169    mean: Any,170    gain: Any,171    exposure: Any,172    offset: Any,173    *,174    out_signal: Any,175    spec: CalibrationSpec,176    workspace: DetectorWorkspace,177    stream: Any = None,178    validate: bool = True,179) -> None:180    """Apply deterministic calibration with one fused pixel launch and FP64 intermediates."""181    require_no_tape()182    ctx = workspace.context183    ctx.assert_stream(stream)184    pixels, reads = _calibration_inputs(mean, gain, exposure, offset, spec, workspace)185    ctx.array(out_signal, "out_signal", dtype=ctx.wp.float32, shape=(pixels,))186    ctx.disjoint(reads, [("out_signal", out_signal)])187    with ctx.scope():188        if validate:189            workspace._values(190                [191                    (mean, "nonnegative"),192                    (gain, "positive"),193                    (exposure, "nonnegative"),194                    (offset, "finite"),195                ]196            )197        if pixels:198            ctx.wp.launch(199                workspace._kernels.get_calibration(spec.shared_gain, spec.shared_offset),200                dim=pixels,201                inputs=[mean, gain, exposure, offset],202                outputs=[out_signal, workspace._status],203                stream=ctx.stream,204                record_tape=False,205            )206        if validate:207            workspace.check_status()208209210def calibration_vjp(211    mean: Any,212    gain: Any,213    exposure: Any,214    offset: Any,215    *,216    seed: Any,217    out_grad_mean: Any = None,218    out_grad_gain: Any = None,219    out_grad_exposure: Any = None,220    out_grad_offset: Any = None,221    spec: CalibrationSpec,222    workspace: DetectorWorkspace,223    stream: Any = None,224    validate: bool = True,225) -> None:226    """Overwrite first-order products; shared nuisance reductions use FP64 trees.227228    Per-image cotangents use binary32. Each requested shared scalar destination229    may be binary32 or binary64. Use binary64 before a logarithmic parameter230    chart: its scale factor can recover a derivative too small for binary32.231    """232    require_no_tape()233    ctx = workspace.context234    ctx.assert_stream(stream)235    pixels, reads = _calibration_inputs(mean, gain, exposure, offset, spec, workspace)236    ctx.array(seed, "seed", dtype=ctx.wp.float32, shape=(pixels,))237    writes: list[tuple[str, Any]] = []238    for name, output, active, size in (239        ("mean", out_grad_mean, spec.active_mean, pixels),240        ("gain", out_grad_gain, spec.active_gain, 1),241        ("exposure", out_grad_exposure, spec.active_exposure, 1),242        ("offset", out_grad_offset, spec.active_offset, 1),243    ):244        if output is not None:245            if not active:246                raise ContractError(f"calibration {name} was declared fixed")247            dtype = ctx.wp.float32248            if name != "mean":249                dtype = getattr(output, "dtype", None)250                if dtype not in (ctx.wp.float32, ctx.wp.float64):251                    raise ContractError("shared calibration gradients must be binary32 or binary64")252            ctx.array(output, f"out_grad_{name}", dtype=dtype, shape=(size,))253            writes.append((f"out_grad_{name}", output))254    if not writes:255        raise ContractError("request at least one active calibration gradient")256    ctx.disjoint([*reads, ("seed", seed)], writes)257    with ctx.scope():258        if validate:259            workspace._values(260                [261                    (mean, "nonnegative"),262                    (gain, "positive"),263                    (exposure, "nonnegative"),264                    (offset, "finite"),265                    (seed, "finite"),266                ]267            )268        if pixels and out_grad_mean is not None:269            ctx.wp.launch(270                workspace._kernels.get_calibration_pixel_vjp(spec.shared_gain, True),271                dim=pixels,272                inputs=[gain, exposure, seed],273                outputs=[out_grad_mean, workspace._status],274                stream=ctx.stream,275                record_tape=False,276            )277        groups = min(workspace.reduction_groups, max(1, (pixels + 255) // 256))278        for kind, output in enumerate((out_grad_gain, out_grad_exposure, out_grad_offset)):279            if output is None:280                continue281            if not pixels:282                output.zero_()283                continue284            ctx.wp.launch_tiled(285                workspace._kernels.get_calibration_partials(spec.shared_gain, kind),286                dim=groups,287                block_dim=256,288                inputs=[mean, gain, exposure, seed, pixels, groups],289                outputs=[workspace._partials],290                stream=ctx.stream,291                record_tape=False,292            )293            ctx.wp.launch(294                workspace._checks.get_finish_shared(output.dtype == ctx.wp.float64),295                dim=1,296                inputs=[workspace._partials, groups],297                outputs=[output, workspace._status],298                stream=ctx.stream,299                record_tape=False,300            )301        if validate:302            workspace.check_status()303304305@dataclass(frozen=True, slots=True)306class BlurSpec:307    """Finite detector response: Bx[r,c] = sum_ij h[i,j] x[r+i-a,c+j-b].308309    Samples outside the detector are zero. The non-negative odd-sized kernel310    sums to at most one, representing a passive spread with optional loss.311    The edge is never renormalised; transpose therefore reverses the offsets.312    """313314    height: int315    width: int316    kernel_height: int317    kernel_width: int318    weights: tuple[float, ...]319    provenance: Provenance320    boundary: Literal["zero"] = "zero"321322    def __post_init__(self) -> None:323        for name in ("height", "width", "kernel_height", "kernel_width"):324            integer(getattr(self, name), name, minimum=1)325        if self.height * self.width > 2**31 - 1:326            raise ContractError("detector dimensions exceed the flat index range")327        if self.kernel_height % 2 != 1 or self.kernel_width % 2 != 1:328            raise ContractError("spatial kernel dimensions must be odd, with a unique centre")329        finite_tuple(self.weights, "blur weights")330        if len(self.weights) != self.kernel_height * self.kernel_width:331            raise ContractError("blur weights must match the declared kernel footprint")332        if math.fsum(self.weights) > 1.0:333            raise ContractError("passive blur weights must sum to at most one")334        if self.boundary != "zero":335            raise ContractError("only explicit zero-extension boundaries are supported")336        if not isinstance(self.provenance, Provenance):337            raise ContractError("spatial-response provenance is required")338339340@dataclass(slots=True)341class BlurWorkspace:342    spec: BlurSpec343    detector: DetectorWorkspace344    _weights: Any = field(repr=False)345346    @property347    def scratch_bytes(self) -> int:348        return self.detector.scratch_bytes + 8 * int(self._weights.size)349350    def clear_status(self) -> None:351        self.detector.clear_status()352353    def check_status(self) -> None:354        self.detector.check_status()355356357def prepare_blur(spec: BlurSpec, *, device: str = "cuda:0", stream: Any = None) -> BlurWorkspace:358    """Upload the fixed stencil in binary64, preserving its declared host coefficients.359360    Stencils are small immutable calibration data. Narrowing their coefficients361    to binary32 can erase a positive tail whose product with a bright input is362    representable, and can change passive total mass. Images still use binary32.363    """364    detector = prepare_detector(max_pixels=spec.height * spec.width, device=device, stream=stream)365    ctx = detector.context366    with ctx.scope():367        weights = ctx.wp.array(spec.weights, dtype=ctx.wp.float64, device=ctx.device)368    return BlurWorkspace(spec, detector, weights)369370371def _apply_blur(372    source: Any,373    output: Any,374    workspace: BlurWorkspace,375    transpose: bool,376    stream: Any,377    validate: bool,378) -> None:379    require_no_tape()380    spec, detector = workspace.spec, workspace.detector381    ctx = detector.context382    ctx.assert_stream(stream)383    size = spec.height * spec.width384    ctx.array(source, "source", dtype=ctx.wp.float32, shape=(size,))385    ctx.array(output, "output", dtype=ctx.wp.float32, shape=(size,))386    ctx.disjoint([("source", source), ("weights", workspace._weights)], [("output", output)])387    with ctx.scope():388        if validate:389            detector._values([(source, "finite")])390        ctx.wp.launch(391            detector._kernels.get_blur(392                spec.height, spec.width, spec.kernel_height, spec.kernel_width, transpose393            ),394            dim=size,395            inputs=[source, workspace._weights],396            outputs=[output, detector._status],397            stream=ctx.stream,398            record_tape=False,399        )400        if validate:401            detector.check_status()402403404def blur(405    source: Any,406    *,407    out_signal: Any,408    workspace: BlurWorkspace,409    stream: Any = None,410    validate: bool = True,411) -> None:412    """Apply the declared finite spatial response; inputs may be signed real signals."""413    _apply_blur(source, out_signal, workspace, False, stream, validate)414415416def blur_transpose(417    seed: Any,418    *,419    out_grad_signal: Any,420    workspace: BlurWorkspace,421    stream: Any = None,422    validate: bool = True,423) -> None:424    """Apply the exact algebraic transpose of the same zero-extended discrete matrix."""425    _apply_blur(seed, out_grad_signal, workspace, True, stream, validate)426427428# region book:observation-identity429@dataclass(frozen=True, slots=True)430class ObservationIdentity:431    """A reproducible random experiment independent of launch tiling.432433    Counter identity = (observation_id << 32) | global_pixel. Poisson bin k has434    domain 1 + energy_offset + k; electronic read noise has domain 0x7fffffff.435    Transport owns domain 0 and domains >= 0x80000000. Reusing an identity is436    common random numbers, not an independent replication. Keep bin identities437    stable when batching, and assign a fresh observation_id for a fresh image.438    ``draw_budget`` bounds rejection proposals, never the sampled count. Below439    rate 10 an inverse-CDF draw uses one proposal with a fixed internal work440    bound; higher rates use at most ``draw_budget`` transformed proposals.441    """442443    seed: int444    observation_id: int445    pixel_offset: int = 0446    energy_offset: int = 0447    draw_budget: int = 256448449    def __post_init__(self) -> None:450        integer(self.seed, "seed", maximum=2**64 - 1)451        integer(self.observation_id, "observation_id", maximum=2**32 - 1)452        integer(self.pixel_offset, "pixel_offset", maximum=2**32 - 1)453        integer(self.energy_offset, "energy_offset", maximum=2**31 - 3)454        integer(self.draw_budget, "draw_budget", minimum=1)455456    def validate_extent(self, pixels: int, energies: int = 1) -> None:457        integer(pixels, "pixels")458        integer(energies, "energies", minimum=1)459        if self.pixel_offset + pixels > 2**32:460            raise ContractError("pixel identities would wrap")461        if self.energy_offset + energies > 2**31 - 2:462            raise ContractError("Poisson domains would overlap the read-noise/transport namespace")463464465# endregion book:observation-identity466467468def sample_poisson_counts(469    expected_counts: Any,470    *,471    out_counts: Any,472    identity: ObservationIdentity,473    workspace: DetectorWorkspace,474    stream: Any = None,475    validate: bool = True,476) -> None:477    """Draw uint64 independent counts from count-domain means in [0, 10⁹].478479    This is an observation generator, not a differentiable forward operator.480    Its input must already be a photon-counting mean. A zero mean returns zero;481    rejection-budget exhaustion invalidates the realisation. No noise or482    non-negativity clipping is applied to calibrated electronic signals here.483    The declared upper rate is an implementation boundary, not an executed484    acceptance result; large-rate mass accuracy and sampling remain CUDA gates.485    """486    require_no_tape()487    ctx = workspace.context488    ctx.assert_stream(stream)489    pixels = workspace._pixels(expected_counts, "expected_counts")490    identity.validate_extent(pixels)491    ctx.array(out_counts, "out_counts", dtype=ctx.wp.uint64, shape=(pixels,))492    ctx.disjoint([("expected_counts", expected_counts)], [("out_counts", out_counts)])493    with ctx.scope():494        if validate:495            workspace._values([(expected_counts, "rate")])496        if pixels:497            ctx.wp.launch(498                workspace._kernels.poisson_counts,499                dim=pixels,500                inputs=[501                    expected_counts,502                    ctx.wp.uint64(identity.seed),503                    ctx.wp.uint32(identity.observation_id),504                    ctx.wp.uint32(identity.pixel_offset),505                    ctx.wp.uint32(identity.energy_offset),506                    identity.draw_budget,507                ],508                outputs=[out_counts, workspace._status],509                stream=ctx.stream,510                record_tape=False,511            )512        if validate:513            workspace.check_status()514515516def sample_compound_poisson(517    detected_bin_means: Any,518    photon_scores: Any,519    *,520    energies: int,521    shared_scores: bool,522    out_signal: Any,523    identity: ObservationIdentity,524    workspace: DetectorWorkspace,525    stream: Any = None,526    validate: bool = True,527) -> None:528    """Draw sum_k score[k,p] N[k,p], N[k,p] ~ Poisson(detected_bin_means[k,p]).529530    The independent rate array is energy-major (K,P). Scores are deterministic531    output per detected photon; an arbitrary mean response does not specify this532    law. The variance before spatial spread is sum_k rate[k,p] score[k,p]².533    Apply spatial spread to this realisation, then add independent electronic534    noise if that order describes the actual detector. Energy-response variance535    within a bin requires a different, explicitly supplied observation law.536    """537    require_no_tape()538    integer(energies, "energies", minimum=1, maximum=65536)539    if type(shared_scores) is not bool:540        raise ContractError("shared_scores must be a boolean")541    ctx = workspace.context542    ctx.assert_stream(stream)543    size = ctx.array(detected_bin_means, "detected_bin_means", dtype=ctx.wp.float32)544    if size % energies:545        raise ContractError("detected bin means must have energy-major shape (K,P)")546    pixels = size // energies547    if pixels > workspace.max_pixels or size > 2**31 - 1:548        raise ContractError("detected bin means exceed the declared pixel/index capacity")549    identity.validate_extent(pixels, energies)550    ctx.array(551        photon_scores,552        "photon_scores",553        dtype=ctx.wp.float32,554        shape=(energies * (1 if shared_scores else pixels),),555    )556    ctx.array(out_signal, "out_signal", dtype=ctx.wp.float32, shape=(pixels,))557    ctx.disjoint(558        [("detected_bin_means", detected_bin_means), ("photon_scores", photon_scores)],559        [("out_signal", out_signal)],560    )561    with ctx.scope():562        if validate:563            workspace._values([(detected_bin_means, "rate"), (photon_scores, "nonnegative")])564        if pixels:565            ctx.wp.launch(566                workspace._kernels.get_compound_poisson(energies, shared_scores),567                dim=pixels,568                inputs=[569                    detected_bin_means,570                    photon_scores,571                    pixels,572                    ctx.wp.uint64(identity.seed),573                    ctx.wp.uint32(identity.observation_id),574                    ctx.wp.uint32(identity.pixel_offset),575                    ctx.wp.uint32(identity.energy_offset),576                    identity.draw_budget,577                ],578                outputs=[out_signal, workspace._status],579                stream=ctx.stream,580                record_tape=False,581            )582        if validate:583            workspace.check_status()584585586def add_gaussian_read_noise(587    signal: Any,588    *,589    standard_deviation: float,590    out_signal: Any,591    identity: ObservationIdentity,592    workspace: DetectorWorkspace,593    stream: Any = None,594    validate: bool = True,595) -> None:596    """Add independent electronic Gaussian noise in the signal's units, without clipping."""597    require_no_tape()598    sigma = finite_scalar(standard_deviation, "standard_deviation", minimum=0)599    ctx = workspace.context600    ctx.assert_stream(stream)601    pixels = workspace._pixels(signal, "signal")602    identity.validate_extent(pixels)603    ctx.array(out_signal, "out_signal", dtype=ctx.wp.float32, shape=(pixels,))604    ctx.disjoint([("signal", signal)], [("out_signal", out_signal)])605    with ctx.scope():606        if validate:607            workspace._values([(signal, "finite")])608        if pixels:609            ctx.wp.launch(610                workspace._kernels.gaussian_read_noise,611                dim=pixels,612                inputs=[613                    signal,614                    sigma,615                    ctx.wp.uint64(identity.seed),616                    ctx.wp.uint32(identity.observation_id),617                    ctx.wp.uint32(identity.pixel_offset),618                ],619                outputs=[out_signal, workspace._status],620                stream=ctx.stream,621                record_tape=False,622            )623        if validate:624            workspace.check_status()625