python/dpt/kernels/detector.py

1"""Detector calibration, zero-extended spatial response and separate observations.23No estimator differentiates a realised count draw. Poisson draws use counter4addresses shared with transport, a bounded rejection loop and an explicit error5on budget exhaustion. There is no Gaussian replacement for a Poisson tail.6"""78# Warp annotations are executable DSL expressions; host interfaces remain strict.9# The optional GPU import is resolved only when an operator is prepared.10# pyright: reportInvalidTypeForm=false, reportUnknownParameterType=false11# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false12# pyright: reportUnknownVariableType=false, reportUntypedFunctionDecorator=false13# pyright: reportMissingImports=false, reportUntypedClassDecorator=false1415from functools import cache1617import warp as wp1819from dpt.kernels.random import random420from dpt.kernels.spectral import checked_store2122STRICT = {"fast_math": False, "fuse_fp": True, "enable_backward": False}23wp.set_module_options(STRICT)24TILE = 256252627@wp.kernel28def check_positive(values: wp.array(dtype=wp.float32), status: wp.array(dtype=wp.int32)):29    p = wp.tid()30    if not wp.isfinite(values[p]) or values[p] <= wp.float32(0.0):31        wp.atomic_or(status, 0, 1)323334@wp.kernel35def check_rates(values: wp.array(dtype=wp.float32), status: wp.array(dtype=wp.int32)):36    p = wp.tid()37    if not wp.isfinite(values[p]) or values[p] < wp.float32(0.0) or values[p] > wp.float32(1.0e9):38        wp.atomic_or(status, 0, 1)394041# region book:detector-calibration42@cache43def get_calibration(shared_gain: bool, shared_offset: bool):44    @wp.kernel(module="unique", module_options=STRICT)45    def calibration(46        mean: wp.array(dtype=wp.float32),47        gain: wp.array(dtype=wp.float32),48        exposure: wp.array(dtype=wp.float32),49        offset: wp.array(dtype=wp.float32),50        output: wp.array(dtype=wp.float32),51        status: wp.array(dtype=wp.int32),52    ):53        p = wp.tid()54        gi = p55        oi = p56        if wp.static(shared_gain):57            gi = 058        if wp.static(shared_offset):59            oi = 060        signal = wp.float64(gain[gi]) * wp.float64(exposure[0]) * wp.float64(mean[p]) + wp.float64(61            offset[oi]62        )63        output[p] = checked_store(signal, status)6465    return calibration666768# endregion book:detector-calibration697071@cache72def get_calibration_pixel_vjp(shared_gain: bool, write_mean: bool):73    @wp.kernel(module="unique", module_options=STRICT)74    def vjp(75        gain: wp.array(dtype=wp.float32),76        exposure: wp.array(dtype=wp.float32),77        seed: wp.array(dtype=wp.float32),78        output: wp.array(dtype=wp.float32),79        status: wp.array(dtype=wp.int32),80    ):81        p = wp.tid()82        gi = p83        if wp.static(shared_gain):84            gi = 085        if wp.static(write_mean):86            output[p] = checked_store(87                wp.float64(seed[p]) * wp.float64(gain[gi]) * wp.float64(exposure[0]),88                status,89            )9091    return vjp929394@cache95def get_calibration_partials(shared_gain: bool, kind: int):96    @wp.kernel(module="unique", module_options=STRICT)97    def partials(98        mean: wp.array(dtype=wp.float32),99        gain: wp.array(dtype=wp.float32),100        exposure: wp.array(dtype=wp.float32),101        seed: wp.array(dtype=wp.float32),102        pixels: int,103        groups: int,104        output: wp.array(dtype=wp.float64),105    ):106        group, lane = wp.tid()107        total = wp.float64(0.0)108        compensation = wp.float64(0.0)109        p = group * TILE + lane110        while p < pixels:111            gi = p112            if wp.static(shared_gain):113                gi = 0114            term = wp.float64(seed[p])115            if wp.static(kind == 0):116                term = term * wp.float64(exposure[0]) * wp.float64(mean[p])117            elif wp.static(kind == 1):118                term = term * wp.float64(gain[gi]) * wp.float64(mean[p])119            corrected = term - compensation120            updated = total + corrected121            compensation = (updated - total) - corrected122            total = updated123            # Do not overflow the final int32 stride at maximum image capacity.124            if pixels - p <= groups * TILE:125                break126            p = p + groups * TILE127        values = wp.tile(total)128        total_tile = wp.tile_sum(values)129        wp.tile_store(output, total_tile, offset=group)130131    return partials132133134# region book:detector-spatial-transpose135@cache136def get_blur(height: int, width: int, kernel_height: int, kernel_width: int, transpose: bool):137    @wp.kernel(module="unique", module_options=STRICT)138    def blur(139        source: wp.array(dtype=wp.float32),140        weights: wp.array(dtype=wp.float64),141        output: wp.array(dtype=wp.float32),142        status: wp.array(dtype=wp.int32),143    ):144        p = wp.tid()145        row = p // width146        column = p % width147        total = wp.float64(0.0)148        compensation = wp.float64(0.0)149        for kr in range(kernel_height):150            for kc in range(kernel_width):151                dr = kr - kernel_height // 2152                dc = kc - kernel_width // 2153                if wp.static(transpose):154                    dr = -dr155                    dc = -dc156                sr = row + dr157                sc = column + dc158                # Zero extension loses signal crossing the finite detector edge.159                # Do not renormalise a boundary row: that changes both B and Bᵀ.160                if sr >= 0 and sr < height and sc >= 0 and sc < width:161                    term = wp.float64(weights[kr * kernel_width + kc]) * wp.float64(162                        source[sr * width + sc]163                    )164                    corrected = term - compensation165                    updated = total + corrected166                    compensation = (updated - total) - corrected167                    total = updated168        output[p] = checked_store(total, status)169170    return blur171172173# endregion book:detector-spatial-transpose174175176@wp.func177def uniform_pair(seed: wp.uint64, identity: wp.uint64, event: wp.uint32, domain: wp.uint32):178    """Two open uniforms using 53 counter bits each, without endpoint clamping."""179    draw = random4(seed, identity, event, domain)180    scale32 = wp.float64(4294967296.0)181    scale21 = wp.float64(2097152.0)182    denominator = wp.float64(9007199254740994.0)183    first = (184        wp.floor(draw[0] * scale32) * scale21 + wp.floor(draw[1] * scale21) + wp.float64(1.0)185    ) / denominator186    second = (187        wp.floor(draw[2] * scale32) * scale21 + wp.floor(draw[3] * scale21) + wp.float64(1.0)188    ) / denominator189    return wp.vec2d(first, second)190191192@wp.func_native("return lgamma(value);")193def log_gamma(value: wp.float64) -> wp.float64:194    """Native double log-gamma for the Poisson rejection acceptance inequality."""195    ...196197198@wp.func199def poisson_log_probability(count: wp.float64, rate: wp.float64) -> wp.float64:200    """Stable log mass for a non-negative integer and a positive Poisson mean.201202    Loader (2002), equations (4), (6), (7): express the mass through Stirling203    error and deviance before evaluating it. Subtracting count*log(rate) and204    lgamma(count+1) directly discards useful digits near a large mean.205    https://www.r-project.org/doc/reports/CLoader-dbinom-2002.pdf206    """207    if count == wp.float64(0.0):208        return -rate209    if count < wp.float64(16.0):210        return -rate + count * wp.log(rate) - log_gamma(count + wp.float64(1.0))211    inverse = wp.float64(1.0) / count212    square = inverse * inverse213    # Cast both operands: casting a quotient lets Warp divide in FP32 first.214    # Six Bernoulli terms; at count >= 16 the next term bounds the absolute215    # truncation error by 7/(1092*16**13) < 1.5e-18 before FP64 rounding.216    correction = inverse * (217        (wp.float64(1.0) / wp.float64(12.0))218        - square219        * (220            (wp.float64(1.0) / wp.float64(360.0))221            - square222            * (223                (wp.float64(1.0) / wp.float64(1260.0))224                - square225                * (226                    (wp.float64(1.0) / wp.float64(1680.0))227                    - square228                    * (229                        (wp.float64(1.0) / wp.float64(1188.0))230                        - square * (wp.float64(691.0) / wp.float64(360360.0))231                    )232                )233            )234        )235    )236    difference = count - rate237    deviance = wp.float64(0.0)238    if wp.abs(difference) < wp.float64(0.1) * (count + rate):239        ratio = difference / (count + rate)240        ratio_squared = ratio * ratio241        deviance = difference * ratio242        term = wp.float64(2.0) * count * ratio243        # |ratio| < .1 gives geometric convergence; 32 terms make the244        # omitted relative tail smaller than binary64 precision throughout.245        for order in range(1, 33):246            term = term * ratio_squared247            updated = deviance + term / wp.float64(2 * order + 1)248            if updated == deviance:249                break250            deviance = updated251    else:252        deviance = count * wp.log(count / rate) - difference253    return (254        -wp.float64(0.5) * wp.log(wp.float64(6.283185307179586476925286766559) * count)255        - correction256        - deviance257    )258259260# region book:detector-poisson-observation261@wp.func262def poisson_draw(263    rate: wp.float64,264    seed: wp.uint64,265    identity: wp.uint64,266    domain: wp.uint32,267    budget: int,268    status: wp.array(dtype=wp.int32),269) -> wp.uint64:270    if rate == wp.float64(0.0):271        return wp.uint64(0)272    if not wp.isfinite(rate) or rate < wp.float64(0.0) or rate > wp.float64(1.0e9):273        wp.atomic_or(status, 0, 1)274        return wp.uint64(0)275    if rate < wp.float64(10.0):276        u = uniform_pair(seed, identity, wp.uint32(0), domain)[0]277        probability = wp.exp(-rate)278        cumulative = probability279        count = int(0)  # noqa: UP018, RUF046 - mutable Warp loop variable280        # One inverse-CDF proposal. Its recurrence bound is independent of281        # draw_budget: a small rate does not imply a bounded Poisson count.282        # At rate < 10 the tail beyond 128 is far below the 2^-33 closest283        # approach of our open-interval 32-bit uniform to one.284        while u > cumulative and count < 128:285            count = count + 1286            probability = probability * rate / wp.float64(count)287            cumulative = cumulative + probability288        if u <= cumulative:289            return wp.uint64(count)290    else:291        # Hörmann's transformed rejection with squeeze (PTRS), 1993,292        # doi:10.1016/0167-6687(93)90997-4. All acceptance arithmetic is FP64.293        root = wp.sqrt(rate)294        b = wp.float64(0.931) + wp.float64(2.53) * root295        a = wp.float64(-0.059) + wp.float64(0.02483) * b296        inverse_alpha = wp.float64(1.1239) + wp.float64(1.1328) / (b - wp.float64(3.4))297        squeeze = wp.float64(0.9277) - wp.float64(3.6224) / (b - wp.float64(2.0))298        for trial in range(budget):299            uniforms = uniform_pair(seed, identity, wp.uint32(trial), domain)300            u = uniforms[0] - wp.float64(0.5)301            v = uniforms[1]302            distance = wp.float64(0.5) - wp.abs(u)303            candidate = wp.floor((wp.float64(2.0) * a / distance + b) * u + rate + wp.float64(0.43))304            if candidate >= wp.float64(0.0):305                if distance >= wp.float64(0.07) and v <= squeeze:306                    return wp.uint64(candidate)307                if not (distance < wp.float64(0.013) and v > distance):308                    lhs = wp.log(v * inverse_alpha / (a / (distance * distance) + b))309                    rhs = poisson_log_probability(candidate, rate)310                    if lhs <= rhs:311                        return wp.uint64(candidate)312    # A bounded failure is an invalid realisation, never an observed zero.313    wp.atomic_or(status, 0, 4)314    return wp.uint64(0)315316317@wp.kernel318def poisson_counts(319    mean: wp.array(dtype=wp.float32),320    seed: wp.uint64,321    observation: wp.uint32,322    pixel_offset: wp.uint32,323    energy_offset: wp.uint32,324    budget: int,325    output: wp.array(dtype=wp.uint64),326    status: wp.array(dtype=wp.int32),327):328    p = wp.tid()329    identity = (wp.uint64(observation) << wp.uint64(32)) | wp.uint64(pixel_offset + wp.uint32(p))330    output[p] = poisson_draw(331        wp.float64(mean[p]), seed, identity, wp.uint32(1) + energy_offset, budget, status332    )333334335@cache336def get_compound_poisson(energies: int, shared_scores: bool):337    @wp.kernel(module="unique", module_options=STRICT)338    def compound(339        rates: wp.array(dtype=wp.float32),340        scores: wp.array(dtype=wp.float32),341        pixels: int,342        seed: wp.uint64,343        observation: wp.uint32,344        pixel_offset: wp.uint32,345        energy_offset: wp.uint32,346        budget: int,347        output: wp.array(dtype=wp.float32),348        status: wp.array(dtype=wp.int32),349    ):350        p = wp.tid()351        identity = (wp.uint64(observation) << wp.uint64(32)) | wp.uint64(352            pixel_offset + wp.uint32(p)353        )354        total = wp.float64(0.0)355        compensation = wp.float64(0.0)356        for energy in range(energies):357            si = energy * pixels + p358            if wp.static(shared_scores):359                si = energy360            count = poisson_draw(361                wp.float64(rates[energy * pixels + p]),362                seed,363                identity,364                wp.uint32(1) + energy_offset + wp.uint32(energy),365                budget,366                status,367            )368            corrected = wp.float64(count) * wp.float64(scores[si]) - compensation369            updated = total + corrected370            compensation = (updated - total) - corrected371            total = updated372        output[p] = checked_store(total, status)373374    return compound375376377# endregion book:detector-poisson-observation378379380@wp.kernel381def gaussian_read_noise(382    signal: wp.array(dtype=wp.float32),383    sigma: wp.float64,384    seed: wp.uint64,385    observation: wp.uint32,386    pixel_offset: wp.uint32,387    output: wp.array(dtype=wp.float32),388    status: wp.array(dtype=wp.int32),389):390    p = wp.tid()391    identity = (wp.uint64(observation) << wp.uint64(32)) | wp.uint64(pixel_offset + wp.uint32(p))392    uniforms = uniform_pair(seed, identity, wp.uint32(0), wp.uint32(0x7FFFFFFF))393    normal = wp.sqrt(-wp.float64(2.0) * wp.log(uniforms[0])) * wp.cos(394        wp.float64(6.283185307179586476925286766559) * uniforms[1]395    )396    output[p] = checked_store(wp.float64(signal[p]) + sigma * normal, status)397