python/dpt/kernels/spectral.py

1"""Spectral primary signal and explicit first-order products.23Paths are material-major (M,P); coefficients are (M,K); weights/response are4shared (K) or energy-major (K,P). FP64 registers hold depths, products and sums.5There is no pixel-by-energy retained tensor. Global parameter gradients use6bounded split reductions; only diagnostic flags use atomics.7"""89# Warp annotations are executable DSL expressions; host interfaces remain strict.10# The optional GPU import is resolved only when an operator is prepared.11# pyright: reportInvalidTypeForm=false, reportUnknownParameterType=false12# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false13# pyright: reportUnknownVariableType=false, reportUntypedFunctionDecorator=false14# pyright: reportMissingImports=false, reportUntypedClassDecorator=false1516from functools import cache1718import warp as wp1920STRICT = {"fast_math": False, "fuse_fp": True, "enable_backward": False}21wp.set_module_options(STRICT)22TILE = 256232425@wp.func26def checked_store(value: wp.float64, status: wp.array(dtype=wp.int32)) -> wp.float32:27    rounded = wp.float32(value)28    if not wp.isfinite(rounded):29        wp.atomic_or(status, 0, 2)30    return rounded313233@cache34def get_value_check(wide: bool, nonnegative: bool):35    dtype = wp.float64 if wide else wp.float323637    @wp.kernel(module="unique", module_options=STRICT)38    def check(values: wp.array(dtype=dtype), status: wp.array(dtype=wp.int32)):39        index = wp.tid()40        invalid = not wp.isfinite(values[index])41        if wp.static(nonnegative):42            invalid = invalid or values[index] < dtype(0.0)43        if invalid:44            wp.atomic_or(status, 0, 1)4546    return check474849check_nonnegative = get_value_check(False, True)50check_finite = get_value_check(False, False)515253@cache54def get_signal_store(wide: bool):55    dtype = wp.float64 if wide else wp.float325657    @wp.func58    def store(value: wp.float64, status: wp.array(dtype=wp.int32)) -> dtype:59        result = dtype(value)60        if not wp.isfinite(result):61            wp.atomic_or(status, 0, 2)62        return result6364    return store656667@cache68def get_optical_depth(wide: bool):69    dtype = wp.float64 if wide else wp.float327071    @wp.func72    def depth_at(73        paths: wp.array(dtype=dtype),74        coefficients: wp.array(dtype=wp.float32),75        pixel: int,76        energy: int,77        pixels: int,78        materials: int,79        energies: int,80    ) -> wp.float64:81        depth = wp.float64(0.0)82        for material in range(materials):83            depth = depth + wp.float64(paths[material * pixels + pixel]) * wp.float64(84                coefficients[material * energies + energy]85            )86        return depth8788    return depth_at899091@cache92def get_attenuation(wide: bool):93    @wp.func94    def attenuate(depth: wp.float64, status: wp.array(dtype=wp.int32)) -> wp.float64:95        factor = wp.exp(-depth)96        if wp.static(wide):97            # A zero exponential may hide a representable weighted tail. This98            # explicit range boundary avoids silently disagreeing with the VJP.99            if factor == wp.float64(0.0):100                wp.atomic_or(status, 0, 4)101        return factor102103    return attenuate104105106@cache107def get_product(wide: bool):108    @wp.func109    def product(a: wp.float64, b: wp.float64, status: wp.array(dtype=wp.int32)) -> wp.float64:110        result = a * b111        if wp.static(wide):112            # A later large factor could restore a representable value. Reject113            # loss of an intermediate instead of silently returning a zero VJP.114            if a != wp.float64(0.0) and b != wp.float64(0.0) and result == wp.float64(0.0):115                wp.atomic_or(status, 0, 4)116        return result117118    return product119120121# region book:spectral-primary-sum122@cache123def get_forward(124    materials: int, energies: int, shared_weights: bool, shared_response: bool, wide: bool = False125):126    dtype = wp.float64 if wide else wp.float32127    optical_depth = get_optical_depth(wide)128    attenuate = get_attenuation(wide)129    product = get_product(wide)130    store_signal = get_signal_store(wide)131132    @wp.kernel(module="unique", module_options=STRICT)133    def forward(134        paths: wp.array(dtype=dtype),135        coefficients: wp.array(dtype=wp.float32),136        weights: wp.array(dtype=wp.float32),137        response: wp.array(dtype=wp.float32),138        pixels: int,139        mean: wp.array(dtype=dtype),140        status: wp.array(dtype=wp.int32),141    ):142        pixel = wp.tid()143        total = wp.float64(0.0)144        compensation = wp.float64(0.0)145        for energy in range(energies):146            wi = energy * pixels + pixel147            ri = wi148            if wp.static(shared_weights):149                wi = energy150            if wp.static(shared_response):151                ri = energy152            depth = optical_depth(paths, coefficients, pixel, energy, pixels, materials, energies)153            contribution = product(154                wp.float64(weights[wi]) * wp.float64(response[ri]), attenuate(depth, status), status155            )156            # All terms are non-negative, but compensation retains small bins157            # when a broad response places many decades in the same sum.158            corrected = contribution - compensation159            updated = total + corrected160            compensation = (updated - total) - corrected161            total = updated162        mean[pixel] = store_signal(total, status)163164    return forward165166167# endregion book:spectral-primary-sum168169170# region book:spectral-recomputed-adjoint171@cache172def get_pixel_vjp(173    materials: int,174    energies: int,175    shared_weights: bool,176    shared_response: bool,177    write_paths: bool,178    write_weights: bool,179    write_response: bool,180    wide: bool = False,181):182    dtype = wp.float64 if wide else wp.float32183    optical_depth = get_optical_depth(wide)184    attenuate = get_attenuation(wide)185    product = get_product(wide)186    store_signal = get_signal_store(wide)187    gradient_vector = wp.types.vector(length=materials, dtype=wp.float64)188189    @wp.kernel(module="unique", module_options=STRICT)190    def vjp(191        paths: wp.array(dtype=dtype),192        coefficients: wp.array(dtype=wp.float32),193        weights: wp.array(dtype=wp.float32),194        response: wp.array(dtype=wp.float32),195        seed: wp.array(dtype=dtype),196        pixels: int,197        grad_paths: wp.array(dtype=dtype),198        grad_weights: wp.array(dtype=wp.float32),199        grad_response: wp.array(dtype=wp.float32),200        status: wp.array(dtype=wp.int32),201    ):202        pixel = wp.tid()203        path_gradient = gradient_vector()204        path_compensation = gradient_vector()205        for energy in range(energies):206            wi = energy * pixels + pixel207            ri = wi208            if wp.static(shared_weights):209                wi = energy210            if wp.static(shared_response):211                ri = energy212            depth = optical_depth(paths, coefficients, pixel, energy, pixels, materials, energies)213            weighted_seed = product(wp.float64(seed[pixel]), attenuate(depth, status), status)214            if wp.static(write_weights):215                grad_weights[wi] = checked_store(216                    product(weighted_seed, wp.float64(response[ri]), status), status217                )218            if wp.static(write_response):219                grad_response[ri] = checked_store(220                    product(weighted_seed, wp.float64(weights[wi]), status), status221                )222            if wp.static(write_paths):223                common = product(224                    product(weighted_seed, wp.float64(weights[wi]), status),225                    wp.float64(response[ri]),226                    status,227                )228                for material in range(materials):229                    term = product(230                        -common, wp.float64(coefficients[material * energies + energy]), status231                    )232                    corrected = term - path_compensation[material]233                    updated = path_gradient[material] + corrected234                    path_compensation[material] = (updated - path_gradient[material]) - corrected235                    path_gradient[material] = updated236        if wp.static(write_paths):237            for material in range(materials):238                grad_paths[material * pixels + pixel] = store_signal(239                    path_gradient[material], status240                )241242    return vjp243244245# endregion book:spectral-recomputed-adjoint246247248@cache249def get_shared_partials(250    materials: int,251    energies: int,252    shared_weights: bool,253    shared_response: bool,254    kind: int,255    wide: bool = False,256):257    """kind=0 weights or 1 response; coefficients have their own shared kernel."""258    dtype = wp.float64 if wide else wp.float32259    optical_depth = get_optical_depth(wide)260    attenuate = get_attenuation(wide)261    product = get_product(wide)262    if kind not in (0, 1):263        raise ValueError("shared partial kind must select weights or response")264265    @wp.kernel(module="unique", module_options=STRICT)266    def partials(267        paths: wp.array(dtype=dtype),268        coefficients: wp.array(dtype=wp.float32),269        weights: wp.array(dtype=wp.float32),270        response: wp.array(dtype=wp.float32),271        seed: wp.array(dtype=dtype),272        pixels: int,273        groups: int,274        output: wp.array(dtype=wp.float64),275        status: wp.array(dtype=wp.int32),276    ):277        group, parameter, lane = wp.tid()278        energy = parameter279        total = wp.float64(0.0)280        compensation = wp.float64(0.0)281        pixel = group * TILE + lane282        while pixel < pixels:283            wi = energy * pixels + pixel284            ri = wi285            if wp.static(shared_weights):286                wi = energy287            if wp.static(shared_response):288                ri = energy289            depth = optical_depth(paths, coefficients, pixel, energy, pixels, materials, energies)290            term = product(wp.float64(seed[pixel]), attenuate(depth, status), status)291            if wp.static(kind == 0):292                term = product(term, wp.float64(response[ri]), status)293            else:294                term = product(term, wp.float64(weights[wi]), status)295            corrected = term - compensation296            updated = total + corrected297            compensation = (updated - total) - corrected298            total = updated299            # The final stride may exceed int32 even when every input index fits.300            # Stop before adding it; pixels - pixel is non-negative and bounded.301            if pixels - pixel <= groups * TILE:302                break303            pixel = pixel + groups * TILE304        values = wp.tile(total)305        result = wp.tile_sum(values)306        wp.tile_store(output, result, offset=parameter * groups + group)307308    return partials309310311@cache312def get_coefficient_partials(313    materials: int, energies: int, shared_weights: bool, shared_response: bool, wide: bool = False314):315    """Reuse each energy's depth across its material cotangents, without a P*K tensor."""316    dtype = wp.float64 if wide else wp.float32317    optical_depth = get_optical_depth(wide)318    attenuate = get_attenuation(wide)319    product = get_product(wide)320    accumulator = wp.types.vector(length=materials, dtype=wp.float64)321322    @wp.kernel(module="unique", module_options=STRICT)323    def coefficient_partials(324        paths: wp.array(dtype=dtype),325        coefficients: wp.array(dtype=wp.float32),326        weights: wp.array(dtype=wp.float32),327        response: wp.array(dtype=wp.float32),328        seed: wp.array(dtype=dtype),329        pixels: int,330        groups: int,331        output: wp.array(dtype=wp.float64),332        status: wp.array(dtype=wp.int32),333    ):334        group, energy, lane = wp.tid()335        totals = accumulator()336        compensations = accumulator()337        pixel = group * TILE + lane338        while pixel < pixels:339            wi = energy * pixels + pixel340            ri = wi341            if wp.static(shared_weights):342                wi = energy343            if wp.static(shared_response):344                ri = energy345            depth = optical_depth(paths, coefficients, pixel, energy, pixels, materials, energies)346            # Keep the same product order as the scalar-parameter reduction.347            common = product(-wp.float64(seed[pixel]), attenuate(depth, status), status)348            common = product(349                product(common, wp.float64(weights[wi]), status), wp.float64(response[ri]), status350            )351            for material in range(materials):352                term = product(common, wp.float64(paths[material * pixels + pixel]), status)353                corrected = term - compensations[material]354                updated = totals[material] + corrected355                compensations[material] = (updated - totals[material]) - corrected356                totals[material] = updated357            # The final stride may exceed int32 even when every input index fits.358            # Stop before adding it; pixels - pixel is non-negative and bounded.359            if pixels - pixel <= groups * TILE:360                break361            pixel = pixel + groups * TILE362        # Each parameter retains the original lane tree and split-sum order.363        for material in range(materials):364            values = wp.tile(totals[material])365            result = wp.tile_sum(values)366            wp.tile_store(output, result, offset=(material * energies + energy) * groups + group)367368    return coefficient_partials369370371@cache372def get_finish_shared(output64: bool):373    """One reduction implementation with explicitly chosen scalar destination precision."""374    output_dtype = wp.float64 if output64 else wp.float32375376    @wp.kernel(module="unique", module_options=STRICT)377    def finish_shared(378        partials: wp.array(dtype=wp.float64),379        groups: int,380        output: wp.array(dtype=output_dtype),381        status: wp.array(dtype=wp.int32),382    ):383        parameter = wp.tid()384        total = wp.float64(0.0)385        compensation = wp.float64(0.0)386        for group in range(groups):387            corrected = partials[parameter * groups + group] - compensation388            updated = total + corrected389            compensation = (updated - total) - corrected390            total = updated391        if wp.static(output64):392            output[parameter] = total393            if not wp.isfinite(total):394                wp.atomic_or(status, 0, 2)395        else:396            output[parameter] = checked_store(total, status)397398    return finish_shared399400401finish_shared = get_finish_shared(False)402403404@wp.kernel405def check_probability(values: wp.array(dtype=wp.float32), status: wp.array(dtype=wp.int32)):406    p = wp.tid()407    if not wp.isfinite(values[p]) or values[p] < wp.float32(0.0) or values[p] > wp.float32(1.0):408        wp.atomic_or(status, 0, 1)409410411@cache412def get_bin_counts(413    materials: int, energies: int, shared_weights: bool, shared_efficiency: bool, wide: bool = False414):415    dtype = wp.float64 if wide else wp.float32416    optical_depth = get_optical_depth(wide)417    attenuate = get_attenuation(wide)418    product = get_product(wide)419    store_signal = get_signal_store(wide)420421    @wp.kernel(module="unique", module_options=STRICT)422    def counts(423        paths: wp.array(dtype=dtype),424        coefficients: wp.array(dtype=wp.float32),425        weights: wp.array(dtype=wp.float32),426        efficiency: wp.array(dtype=wp.float32),427        pixels: int,428        output: wp.array(dtype=dtype),429        status: wp.array(dtype=wp.int32),430    ):431        energy, p = wp.tid()432        wi = energy * pixels + p433        ri = wi434        if wp.static(shared_weights):435            wi = energy436        if wp.static(shared_efficiency):437            ri = energy438        depth = optical_depth(paths, coefficients, p, energy, pixels, materials, energies)439        rate = product(440            wp.float64(weights[wi]) * wp.float64(efficiency[ri]), attenuate(depth, status), status441        )442        output[energy * pixels + p] = store_signal(rate, status)443444    return counts445