python/dpt/transport/kernels.py

1"""Canonical complete-history Warp kernels, including the likelihood-ratio replay.23No path-by-event storage, queue allocation, roulette or hidden truncation. One4thread owns one original history and writes one sparse detector contribution.5Derivative replay runs the *same* flight function with a selected active material.6This initial scheduling candidate is unprofiled: divergence, register pressure,7FP64 cost and tally contention require actual-device acceptance before any claim.8"""910# pyright: reportInvalidTypeForm=false, reportUnknownParameterType=false11# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false12# pyright: reportUnknownVariableType=false, reportUntypedFunctionDecorator=false13# pyright: reportMissingImports=false, reportUntypedClassDecorator=false14# Integer constructors declare mutable Warp loop variables.15# ruff: noqa: UP018, RUF04616import warp as wp1718from dpt.kernels.random import random41920from ._constants import LOG_SOURCE_AMPLITUDE, RANDOM_NAMESPACE_BIT, SOURCE_AMPLITUDE2122wp.set_module_options({"fast_math": False, "fuse_fp": True, "enable_backward": False})2324MAX_DISTANCE = wp.constant(wp.float64(1.7976931348623157e308))25TWO_PI = wp.constant(wp.float64(6.2831853071795864769))26ELECTRON_REST_ENERGY_KEV = wp.constant(wp.float64(510.99895069))27MISS_PIXEL = wp.constant(-1)28INVALID_PIXEL = wp.constant(-2)293031@wp.func_native("""32#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 70033    // Every non-exited lane reaches this ballot, including misses and invalid34    // histories. Membership is established before the participating branch.35    const unsigned members = __ballot_sync(0xffffffffu, participating != 0);36    if (!participating) return wp::vec_t<2, double>(0.0, 0.0);37    const unsigned peers = __match_any_sync(members, target);38    const int leader = __ffs(peers) - 1;39    const int lane = threadIdx.x & 31;40    const int count = __popc(peers);41    double total = value;42    for (int delta = 1; delta < count; delta <<= 1) {43        // __fns selects the delta-th participating lane after this lane, even44        // for interleaved detector addresses or inactive/missing histories.45        const unsigned source = __fns(peers, lane, delta + 1);46        const double next = __shfl_sync(peers, total, source < 32 ? source : lane);47        if (source < 32)48            total = maximum ? fmax(total, next) : total + next;49    }50    return wp::vec_t<2, double>(total, lane == leader ? double(count) : 0.0);51#else52    return wp::vec_t<2, double>(participating ? value : 0.0, participating ? 1.0 : 0.0);53#endif54""")55def grouped_tally(value: wp.float64, target: int, maximum: int, participating: int) -> wp.vec2d:56    """Combine equal detector addresses selected by a converged warp ballot.5758    Call unconditionally from every live kernel lane, after computing its input59    and participation flag. Every lane named by a peer mask then executes the60    same shuffles. Misses and the final partial warp contribute nothing. Only the61    peer leader updates global storage; unrelated detector addresses stay separate.62    Inputs are nonnegative magnitudes or bounded normalised scores/deviations.63    Summation order remains unspecified, as it was for per-history atomics.64    CUDA targets below SM 70 retain individual atomics; no unsupported match65    intrinsic is emitted there. The measured optimisation targets SM 121.66    """67    ...686970@wp.func_native("""71if (a == 0.0 || b == 0.0 || c == 0.0) return 0.0;72int ea, eb, ec;73double ma = frexp(a, &ea);74double mb = frexp(b, &eb);75double mc = frexp(c, &ec);76return ldexp((ma * mb) * mc, ea + eb + ec);77""")78def product3(a: wp.float64, b: wp.float64, c: wp.float64) -> wp.float64:79    """Form a finite-factor product without intermediate exponent overflow/underflow."""80    ...818283@wp.func_native("""84if (a == 0.0 || b == 0.0 || c == 0.0 || d == 0.0) return 0.0;85int ea, eb, ec, ed;86double ma = frexp(a, &ea);87double mb = frexp(b, &eb);88double mc = frexp(c, &ec);89double md = frexp(d, &ed);90return ldexp(((ma * mb) * mc) * md, ea + eb + ec + ed);91""")92def product4(a: wp.float64, b: wp.float64, c: wp.float64, d: wp.float64) -> wp.float64:93    """A derivative product has its own range, independent of a rounded primal score."""94    ...959697@wp.func_native("""98if (a == 0.0 || b == 0.0 || c == 0.0 || d == 0.0) return 0.0;99int ea, eb, ec, ed;100const double ma = frexp(a, &ea);101const double mb = frexp(b, &eb);102const double mc = frexp(c, &ec);103const double md = frexp(d, &ed);104const double mantissa = ((ma * mb) * mc) * md;105// Four finite binary64 factors cannot rescue attenuation above this bound.106// Check before converting the optical depth to an integer exponent.107if (tau > 8192.0) return copysign(0.0, mantissa);108const int attenuation_exponent = static_cast<int>(floor(tau * 1.4426950408889634074));109// Split ln(2) and one FMA retain the residual when tau is near 800 or larger.110const double residual = fma(double(attenuation_exponent), 0.69314718055994530942, -tau)111    + double(attenuation_exponent) * 2.3190468138462995584e-17;112return ldexp(mantissa * exp(residual), ea + eb + ec + ed - attenuation_exponent);113""")114def attenuated_product4(115    tau: wp.float64, a: wp.float64, b: wp.float64, c: wp.float64, d: wp.float64116) -> wp.float64:117    """Multiply original factors by exp(-tau) before the final binary64 rounding.118119    Finite factors and nonnegative finite optical depth are checked by callers.120    An underflowed attenuation or forward score never determines derivative range.121    """122    ...123124125@wp.struct126class Parameters:127    origin: wp.vec3d128    spacing: wp.vec3d129    shape: wp.vec3i130    detector_lower: wp.vec2d131    detector_spacing: wp.vec2d132    detector_shape: wp.vec2i133    detector_z: wp.float64134    source_energy: wp.float64135    energy_bins: int136    max_events: int137    max_crossings: int138    max_angle_trials: int139    compton: int140    energy_score: int141    continuous_absorption: int142143144@wp.struct145class Trace:146    pixel: int147    events: int148    status: int149    energy: wp.float64150    score: wp.float64151    density_score: wp.float64152    absorption_depth: wp.float64153154155@wp.struct156class GridEntry:157    position: wp.vec3d158    cell: wp.vec3i159    alive: int160    status: int161162163@wp.struct164class FaceCrossing:165    position: wp.vec3d166    cell: wp.vec3i167    inside: int168169170@wp.func171def detector_pixel(position: wp.vec3d, direction: wp.vec3d, p: Parameters) -> int:172    pixel = int(MISS_PIXEL)173    if direction[2] > wp.float64(0.0):174        distance = (p.detector_z - position[2]) / direction[2]175        if not wp.isfinite(distance):176            return INVALID_PIXEL177        if distance >= wp.float64(0.0):178            x = position[0] + distance * direction[0] - p.detector_lower[0]179            y = position[1] + distance * direction[1] - p.detector_lower[1]180            if not wp.isfinite(x) or not wp.isfinite(y):181                return INVALID_PIXEL182            # Test support before converting potentially huge floating coordinates.183            if (184                x >= wp.float64(0.0)185                and y >= wp.float64(0.0)186                and x < wp.float64(p.detector_shape[0]) * p.detector_spacing[0]187                and y < wp.float64(p.detector_shape[1]) * p.detector_spacing[1]188            ):189                # Physical support has already been tested; quotient rounding190                # at its upper edge must not alias a neighbouring row.191                column = wp.min(int(wp.floor(x / p.detector_spacing[0])), p.detector_shape[0] - 1)192                row = wp.min(int(wp.floor(y / p.detector_spacing[1])), p.detector_shape[1] - 1)193                pixel = row * p.detector_shape[0] + column194    return pixel195196197@wp.func198def cell_index(position: wp.vec3d, direction: wp.vec3d, p: Parameters) -> wp.vec3i:199    cell = wp.vec3i(0)200    for axis in range(3):201        coordinate = (position[axis] - p.origin[axis]) / p.spacing[axis]202        index = int(wp.floor(coordinate))203        # A point on a face belongs to the downstream cell. No epsilon displaces204        # a physical path or silently changes a thin voxel's optical thickness.205        if direction[axis] < wp.float64(0.0) and coordinate == wp.float64(index):206            index -= 1207        cell[axis] = wp.clamp(index, 0, p.shape[axis] - 1)208    return cell209210211@wp.func212def coefficients(213    material: int,214    energy: wp.float64,215    energies: wp.array(dtype=wp.float64),216    absorption: wp.array(dtype=wp.float64),217    scattering: wp.array(dtype=wp.float64),218    p: Parameters,219) -> wp.vec2d:220    offset = material * p.energy_bins221    if p.energy_bins == 1:222        return wp.vec2d(absorption[offset], scattering[offset])223    lower = int(0)224    upper = p.energy_bins - 1225    while upper - lower > 1:226        middle = (lower + upper) // 2227        if energies[middle] <= energy:228            lower = middle229        else:230            upper = middle231    fraction = (energy - energies[lower]) / (energies[upper] - energies[lower])232    return wp.vec2d(233        (wp.float64(1.0) - fraction) * absorption[offset + lower]234        + fraction * absorption[offset + upper],235        (wp.float64(1.0) - fraction) * scattering[offset + lower]236        + fraction * scattering[offset + upper],237    )238239240@wp.func241def scatter_direction(direction: wp.vec3d, cosine: wp.float64, azimuth: wp.float64) -> wp.vec3d:242    # Choose the auxiliary axis away from collinearity; polar singularities do243    # not justify dividing by sin(theta) of the incoming direction.244    auxiliary = wp.vec3d(wp.float64(0.0), wp.float64(0.0), wp.float64(1.0))245    if wp.abs(direction[2]) > wp.float64(0.9):246        auxiliary = wp.vec3d(wp.float64(1.0), wp.float64(0.0), wp.float64(0.0))247    tangent = wp.normalize(wp.cross(auxiliary, direction))248    bitangent = wp.cross(direction, tangent)249    sine = wp.sqrt(wp.max(wp.float64(0.0), wp.float64(1.0) - cosine * cosine))250    return wp.normalize(251        cosine * direction + sine * (wp.cos(azimuth) * tangent + wp.sin(azimuth) * bitangent)252    )253254255# region book:transport-compton-law256@wp.func257def compton_scatter(258    energy: wp.float64,259    seed: wp.uint64,260    history: wp.uint64,261    event: wp.uint32,262    max_trials: int,263) -> wp.vec4d:264    """Cosine, azimuth, outgoing energy and status from the conditional free-electron law."""265    result = wp.vec4d(wp.float64(0.0), wp.float64(0.0), energy, wp.float64(6.0))266    for trial in range(max_trials):267        angular = random4(seed, history, event, wp.uint32(RANDOM_NAMESPACE_BIT) + wp.uint32(trial))268        cosine = wp.float64(2.0) * angular[0] - wp.float64(1.0)269        ratio = wp.float64(1.0) / (270            wp.float64(1.0) + energy / ELECTRON_REST_ENERGY_KEV * (wp.float64(1.0) - cosine)271        )272        # Klein-Nishina density divided by its envelope 2. This conditional273        # scattering law is independent of the active material-density scale.274        acceptance = wp.float64(0.5) * (275            ratio * ratio * ratio + ratio - ratio * ratio * (wp.float64(1.0) - cosine * cosine)276        )277        if angular[1] < acceptance:278            result = wp.vec4d(279                cosine,280                TWO_PI * angular[2],281                energy * ratio,282                wp.float64(0.0),283            )284            break285    return result286287288# endregion book:transport-compton-law289290291# region book:transport-grid-traversal292@wp.func293def enter_grid(position: wp.vec3d, direction: wp.vec3d, p: Parameters) -> GridEntry:294    """Intersect a forward ray and assign its downstream cell without an epsilon."""295    result = GridEntry()296    result.status = 0297    result.cell = wp.vec3i(0)298    entry = wp.float64(0.0)299    exit_distance = wp.float64(MAX_DISTANCE)300    intersects = int(1)301    for axis in range(3):302        lower = p.origin[axis]303        upper = lower + wp.float64(p.shape[axis]) * p.spacing[axis]304        if direction[axis] == wp.float64(0.0):305            if position[axis] < lower or position[axis] >= upper:306                intersects = 0307        else:308            first = (lower - position[axis]) / direction[axis]309            second = (upper - position[axis]) / direction[axis]310            if not wp.isfinite(first) or not wp.isfinite(second):311                result.status = 4312                intersects = 0313                break314            entry = wp.max(entry, wp.min(first, second))315            exit_distance = wp.min(exit_distance, wp.max(first, second))316    if exit_distance <= entry:317        intersects = 0318    result.alive = intersects319    if result.alive != 0:320        position += entry * direction321        # Only round the intersection coordinate to a face, not the travelled322        # distance. The entry came from these same faces and cannot be outside.323        for axis in range(3):324            position[axis] = wp.clamp(325                position[axis],326                p.origin[axis],327                p.origin[axis] + wp.float64(p.shape[axis]) * p.spacing[axis],328            )329    if result.alive != 0:330        result.cell = cell_index(position, direction, p)331    result.position = position332    return result333334335@wp.func336def distances_to_faces(337    position: wp.vec3d, direction: wp.vec3d, cell: wp.vec3i, p: Parameters338) -> wp.vec3d:339    """Distances to the next cell face on each axis, retaining exact ties."""340    distances = wp.vec3d(MAX_DISTANCE)341    for axis in range(3):342        if direction[axis] != wp.float64(0.0):343            face_index = cell[axis]344            if direction[axis] > wp.float64(0.0):345                face_index += 1346            face = p.origin[axis] + wp.float64(face_index) * p.spacing[axis]347            distances[axis] = (face - position[axis]) / direction[axis]348    return distances349350351@wp.func352def cross_faces(353    position: wp.vec3d,354    direction: wp.vec3d,355    cell: wp.vec3i,356    face_distances: wp.vec3d,357    distance: wp.float64,358    p: Parameters,359) -> FaceCrossing:360    """Cross every tied face and snap only its coordinate to the known cell face."""361    result = FaceCrossing()362    result.position = position363    result.cell = cell364    result.inside = 1365    for axis in range(3):366        if face_distances[axis] == distance:367            if direction[axis] > wp.float64(0.0):368                result.cell[axis] += 1369                result.position[axis] = (370                    p.origin[axis] + wp.float64(result.cell[axis]) * p.spacing[axis]371                )372            else:373                result.position[axis] = (374                    p.origin[axis] + wp.float64(result.cell[axis]) * p.spacing[axis]375                )376                result.cell[axis] -= 1377        if result.cell[axis] < 0 or result.cell[axis] >= p.shape[axis]:378            result.inside = 0379    return result380381382# endregion book:transport-grid-traversal383384385# region book:transport-residual-flight386@wp.func387def walk(388    initial_position: wp.vec3d,389    initial_direction: wp.vec3d,390    source_weight: wp.float64,391    source_amplitude: wp.float64,392    seed: wp.uint64,393    history: wp.uint64,394    active_material: int,395    material_ids: wp.array(dtype=wp.int32),396    energies: wp.array(dtype=wp.float64),397    absorption: wp.array(dtype=wp.float64),398    scattering: wp.array(dtype=wp.float64),399    density: wp.array(dtype=wp.float64),400    p: Parameters,401) -> Trace:402    result = Trace()403    result.pixel = MISS_PIXEL404    result.events = 0405    result.status = 0406    result.energy = p.source_energy407    result.score = wp.float64(0.0)408    result.density_score = wp.float64(0.0)409    result.absorption_depth = wp.float64(0.0)410    position = initial_position411    direction = initial_direction412    entry = enter_grid(position, direction, p)413    position = entry.position414    cell = entry.cell415    alive = entry.alive416    result.status = entry.status417    crossings = int(0)418    while alive != 0:419        if result.events >= p.max_events:420            result.status = 2421            break422        if result.energy < energies[0] or result.energy > energies[p.energy_bins - 1]:423            result.status = 5424            break425        draws = random4(seed, history, wp.uint32(result.events), wp.uint32(0))426        residual = -wp.log(draws[0])427        collision = int(0)428        material = int(0)429        interaction = wp.vec2d(wp.float64(0.0))430        while alive != 0 and collision == 0:431            flat = (cell[2] * p.shape[1] + cell[1]) * p.shape[0] + cell[0]432            material = material_ids[flat]433            interaction = coefficients(material, result.energy, energies, absorption, scattering, p)434            extinction = density[material] * (interaction[0] + interaction[1])435            absorption_rate = wp.float64(0.0)436            if p.continuous_absorption != 0:437                extinction = density[material] * interaction[1]438                absorption_rate = density[material] * interaction[0]439            face_distance = distances_to_faces(position, direction, cell, p)440            distance = wp.min(face_distance[0], wp.min(face_distance[1], face_distance[2]))441            if distance < wp.float64(0.0) or not wp.isfinite(distance):442                result.status = 4443                alive = 0444                break445            optical_distance = extinction * distance446            if (447                not wp.isfinite(extinction)448                or not wp.isfinite(optical_distance)449                or not wp.isfinite(absorption_rate)450            ):451                result.status = 4452                alive = 0453                break454            absorption_distance = wp.float64(0.0)455            if p.continuous_absorption != 0:456                segment_distance = distance457                if extinction > wp.float64(0.0) and residual < optical_distance:458                    segment_distance = residual / extinction459                # Use the realised segment, not the entire distance to the face:460                # scattering can interrupt this segment before the face is reached.461                absorption_distance = absorption_rate * segment_distance462                result.absorption_depth += absorption_distance463                if not wp.isfinite(result.absorption_depth):464                    result.status = 4465                    alive = 0466                    break467            if extinction > wp.float64(0.0) and residual < optical_distance:468                distance = residual / extinction469                position += distance * direction470                if material == active_material:471                    result.density_score += wp.float64(1.0) - residual472                    if p.continuous_absorption != 0:473                        result.density_score -= absorption_distance474                collision = 1475            else:476                if material == active_material:477                    result.density_score -= optical_distance478                    if p.continuous_absorption != 0:479                        result.density_score -= absorption_distance480                residual -= optical_distance481                position += distance * direction482                crossing = cross_faces(position, direction, cell, face_distance, distance, p)483                position = crossing.position484                cell = crossing.cell485                alive = crossing.inside486                crossings += 1487                if alive != 0 and crossings >= p.max_crossings:488                    result.status = 3489                    alive = 0490        if collision != 0:491            result.events += 1492            absorption_probability = wp.float64(0.0)493            if p.continuous_absorption == 0:494                absorption_probability = interaction[0] / (interaction[0] + interaction[1])495            if p.continuous_absorption == 0 and draws[1] < absorption_probability:496                result.status = 1497                alive = 0498            else:499                cosine = wp.float64(2.0) * draws[2] - wp.float64(1.0)500                azimuth = TWO_PI * draws[3]501                if p.compton != 0:502                    angular = compton_scatter(503                        result.energy,504                        seed,505                        history,506                        wp.uint32(result.events - 1),507                        p.max_angle_trials,508                    )509                    cosine = angular[0]510                    azimuth = angular[1]511                    result.energy = angular[2]512                    if angular[3] != wp.float64(0.0):513                        result.status = 6514                        alive = 0515                if alive != 0:516                    direction = scatter_direction(direction, cosine, azimuth)517    if result.status == 0:518        result.pixel = detector_pixel(position, direction, p)519        if result.pixel == INVALID_PIXEL:520            result.pixel = MISS_PIXEL521            result.status = 4522        if result.pixel >= 0:523            score_energy = wp.float64(1.0)524            if p.energy_score != 0:525                score_energy = result.energy526            result.score = product3(source_weight, source_amplitude, score_energy)527            if p.continuous_absorption != 0:528                result.score = attenuated_product4(529                    result.absorption_depth,530                    source_weight,531                    source_amplitude,532                    score_energy,533                    wp.float64(1.0),534                )535    if not wp.isfinite(result.score) or not wp.isfinite(result.density_score):536        result.status = 4537    return result538539540# endregion book:transport-residual-flight541542543@wp.kernel544def validate_model(545    material_ids: wp.array(dtype=wp.int32),546    absorption: wp.array(dtype=wp.float64),547    scattering: wp.array(dtype=wp.float64),548    materials: int,549    status: wp.array(dtype=wp.int32),550):551    index = wp.tid()552    if index < material_ids.shape[0]:553        if material_ids[index] < 0 or material_ids[index] >= materials:554            wp.atomic_or(status, 0, 1)555    if index < absorption.shape[0]:556        a = absorption[index]557        s = scattering[index]558        if (559            not wp.isfinite(a)560            or not wp.isfinite(s)561            or a < wp.float64(0.0)562            or s < wp.float64(0.0)563            or not wp.isfinite(a + s)564        ):565            wp.atomic_or(status, 0, 1)566567568@wp.kernel569def validate_sources(570    positions: wp.array(dtype=wp.vec3d),571    directions: wp.array(dtype=wp.vec3d),572    weights: wp.array(dtype=wp.float64),573    density: wp.array(dtype=wp.float64),574    detector_z: wp.float64,575    status: wp.array(dtype=wp.int32),576):577    index = wp.tid()578    if index < positions.shape[0]:579        position = positions[index]580        direction = directions[index]581        for axis in range(3):582            if not wp.isfinite(position[axis]) or not wp.isfinite(direction[axis]):583                wp.atomic_or(status, 0, 1)584        if (585            wp.abs(wp.dot(direction, direction) - wp.float64(1.0)) > wp.float64(1.0e-12)586            or not wp.isfinite(weights[index])587            or weights[index] < wp.float64(0.0)588            or position[2] >= detector_z589        ):590            wp.atomic_or(status, 0, 1)591    if index < density.shape[0]:592        if not wp.isfinite(density[index]) or density[index] <= wp.float64(0.0):593            wp.atomic_or(status, 0, 1)594595596@wp.kernel597def trace_histories(598    positions: wp.array(dtype=wp.vec3d),599    directions: wp.array(dtype=wp.vec3d),600    weights: wp.array(dtype=wp.float64),601    density: wp.array(dtype=wp.float64),602    material_ids: wp.array(dtype=wp.int32),603    energies: wp.array(dtype=wp.float64),604    absorption: wp.array(dtype=wp.float64),605    scattering: wp.array(dtype=wp.float64),606    parameters: Parameters,607    seed: wp.uint64,608    first_history: wp.uint64,609    source_amplitude: wp.float64,610    out_pixel: wp.array(dtype=wp.int32),611    out_score: wp.array(dtype=wp.float64),612    out_energy: wp.array(dtype=wp.float64),613    out_events: wp.array(dtype=wp.int32),614    out_status: wp.array(dtype=wp.int32),615    status: wp.array(dtype=wp.int32),616):617    index = wp.tid()618    result = walk(619        positions[index],620        directions[index],621        weights[index],622        source_amplitude,623        seed,624        first_history + wp.uint64(index),625        SOURCE_AMPLITUDE,626        material_ids,627        energies,628        absorption,629        scattering,630        density,631        parameters,632    )633    out_pixel[index] = result.pixel634    out_score[index] = result.score635    if not wp.isfinite(out_score[index]):636        wp.atomic_or(status, 0, 16)637    out_energy[index] = result.energy638    out_events[index] = result.events639    out_status[index] = result.status640    if result.status >= 2:641        wp.atomic_or(status, 0, 1 << result.status)642643644# region book:transport-density-score645@wp.kernel646def derivative_histories(647    positions: wp.array(dtype=wp.vec3d),648    directions: wp.array(dtype=wp.vec3d),649    weights: wp.array(dtype=wp.float64),650    density: wp.array(dtype=wp.float64),651    material_ids: wp.array(dtype=wp.int32),652    energies: wp.array(dtype=wp.float64),653    absorption: wp.array(dtype=wp.float64),654    scattering: wp.array(dtype=wp.float64),655    parameters: Parameters,656    seed: wp.uint64,657    first_history: wp.uint64,658    active_material: int,659    source_amplitude: wp.float64,660    out_pixel: wp.array(dtype=wp.int32),661    out_derivative: wp.array(dtype=wp.float64),662    out_status: wp.array(dtype=wp.int32),663    status: wp.array(dtype=wp.int32),664):665    index = wp.tid()666    amplitude = source_amplitude667    base_weight = weights[index]668    if active_material >= 0:669        # A density derivative is formed directly below. An unused primal670        # overflow/underflow must not determine its representability.671        base_weight = wp.float64(0.0)672    if active_material == SOURCE_AMPLITUDE:673        amplitude = wp.float64(1.0)674    result = walk(675        positions[index],676        directions[index],677        base_weight,678        amplitude,679        seed,680        first_history + wp.uint64(index),681        active_material,682        material_ids,683        energies,684        absorption,685        scattering,686        density,687        parameters,688    )689    out_pixel[index] = result.pixel690    derivative = wp.float64(0.0)691    if active_material >= 0 and result.pixel >= 0:692        score_energy = wp.float64(1.0)693        if parameters.energy_score != 0:694            score_energy = result.energy695        derivative = product4(weights[index], source_amplitude, score_energy, result.density_score)696        if parameters.continuous_absorption != 0:697            derivative = attenuated_product4(698                result.absorption_depth,699                weights[index],700                source_amplitude,701                score_energy,702                result.density_score,703            )704    if active_material == SOURCE_AMPLITUDE:705        # d(a * base_score)/da, including at a=0; never divide by amplitude.706        derivative = result.score707    if active_material == LOG_SOURCE_AMPLITUDE:708        derivative = result.score709    out_derivative[index] = derivative710    out_status[index] = result.status711    if result.status >= 2:712        wp.atomic_or(status, 0, 1 << result.status)713    if not wp.isfinite(derivative):714        wp.atomic_or(status, 0, 16)715716717# endregion book:transport-density-score718719720@wp.kernel721def centred_moments(722    pixel: wp.array(dtype=wp.int32),723    score: wp.array(dtype=wp.float64),724    mean: wp.array(dtype=wp.float64),725    scale: wp.array(dtype=wp.float64),726    out_variance: wp.array(dtype=wp.float64),727    out_hits: wp.array(dtype=wp.int32),728    status: wp.array(dtype=wp.int32),729):730    history = wp.tid()731    target = pixel[history]732    participating = int(0)733    contribution = wp.float64(0.0)734    if target >= 0 and target < mean.shape[0]:735        residual = wp.float64(0.0)736        if scale[target] > wp.float64(0.0):737            value = score[history]738            average = mean[target]739            if (value >= wp.float64(0.0)) == (average >= wp.float64(0.0)):740                residual = (value - average) / scale[target]741            else:742                residual = value / scale[target] - average / scale[target]743        contribution = residual * residual744        if not wp.isfinite(contribution):745            wp.atomic_or(status, 0, 16)746        else:747            participating = 1748    aggregate = grouped_tally(contribution, target, 0, participating)749    if aggregate[1] > wp.float64(0.0):750        wp.atomic_add(out_variance, target, aggregate[0])751        wp.atomic_add(out_hits, target, int(aggregate[1]))752753754@wp.kernel755def finish_variance(756    mean: wp.array(dtype=wp.float64),757    scale: wp.array(dtype=wp.float64),758    hits: wp.array(dtype=wp.int32),759    count: int,760    out_variance: wp.array(dtype=wp.float64),761    status: wp.array(dtype=wp.int32),762):763    pixel = wp.tid()764    n = wp.float64(count)765    missing = count - hits[pixel]766    centred = out_variance[pixel]767    if missing > 0 and scale[pixel] > wp.float64(0.0):768        residual = mean[pixel] / scale[pixel]769        centred += wp.float64(missing) * residual * residual770    # Accumulate normalised squares before restoring their exponent. Squaring771    # each tiny history first can turn a representable aggregate variance to zero.772    sigma = scale[pixel] * wp.sqrt(centred / n / (n - wp.float64(1.0)))773    variance = sigma * sigma774    out_variance[pixel] = variance775    if missing < 0 or not wp.isfinite(variance):776        wp.atomic_or(status, 0, 16)777778779@wp.kernel780def independent_product(781    mean_a: wp.array(dtype=wp.float64),782    other_b: wp.array(dtype=wp.float64),783    observed: wp.array(dtype=wp.float64),784    weights: wp.array(dtype=wp.float64),785    loss_mode: int,786    out_components: wp.array(dtype=wp.float64),787    status: wp.array(dtype=wp.int32),788):789    pixel = wp.tid()790    a = mean_a[pixel]791    b = other_b[pixel]792    y = observed[pixel]793    w = weights[pixel]794    if (795        not wp.isfinite(a)796        or not wp.isfinite(b)797        or not wp.isfinite(y)798        or not wp.isfinite(w)799        or w < wp.float64(0.0)800    ):801        wp.atomic_or(status, 0, 1)802    value = w * (a - y) * b803    if loss_mode != 0:804        value = wp.float64(0.5) * w * (a - y) * (b - y)805    out_components[pixel] = value806    if not wp.isfinite(value):807        wp.atomic_or(status, 0, 16)808809810@wp.kernel811def sample_parallel_source(812    lower: wp.vec3d,813    extent: wp.vec2d,814    direction: wp.vec3d,815    weight: wp.float64,816    seed: wp.uint64,817    first_history: wp.uint64,818    out_position: wp.array(dtype=wp.vec3d),819    out_direction: wp.array(dtype=wp.vec3d),820    out_weight: wp.array(dtype=wp.float64),821):822    index = wp.tid()823    # The high event bit separates source draws from every supported collision.824    draw = random4(825        seed, first_history + wp.uint64(index), wp.uint32(RANDOM_NAMESPACE_BIT), wp.uint32(0)826    )827    out_position[index] = lower + wp.vec3d(828        extent[0] * draw[0], extent[1] * draw[1], wp.float64(0.0)829    )830    out_direction[index] = direction831    out_weight[index] = weight832833834@wp.kernel835def update_density_chart(836    parameters: wp.array(dtype=wp.float64),837    material_parameter: wp.array(dtype=wp.int32),838    base_density: wp.array(dtype=wp.float64),839    out_density: wp.array(dtype=wp.float64),840    status: wp.array(dtype=wp.int32),841):842    material = wp.tid()843    active = material_parameter[material]844    value = base_density[material]845    if active >= 0:846        value = wp.exp(parameters[active])847    out_density[material] = value848    if not wp.isfinite(value) or value <= wp.float64(0.0):849        wp.atomic_or(status, 0, 16)850851852@wp.kernel853def find_tally_scale(854    pixel: wp.array(dtype=wp.int32),855    score: wp.array(dtype=wp.float64),856    history_status: wp.array(dtype=wp.int32),857    pixels: int,858    out_scale: wp.array(dtype=wp.float64),859    status: wp.array(dtype=wp.int32),860):861    history = wp.tid()862    target = pixel[history]863    value = score[history]864    participating = int(0)865    terminal = history_status[history]866    if terminal >= 2 and terminal <= 6:867        wp.atomic_or(status, 0, 1 << terminal)868    elif (869        target < MISS_PIXEL870        or target >= pixels871        or not wp.isfinite(value)872        or terminal < 0873        or terminal > 6874        or (target == MISS_PIXEL and value != wp.float64(0.0))875        or (terminal == 1 and target != MISS_PIXEL)876    ):877        wp.atomic_or(status, 0, 1)878    elif target >= 0:879        participating = 1880    aggregate = grouped_tally(wp.abs(value), target, 1, participating)881    if aggregate[1] > wp.float64(0.0):882        wp.atomic_max(out_scale, target, aggregate[0])883884885@wp.kernel886def tally_sum(887    pixel: wp.array(dtype=wp.int32),888    score: wp.array(dtype=wp.float64),889    history_status: wp.array(dtype=wp.int32),890    pixels: int,891    scale: wp.array(dtype=wp.float64),892    out_sum: wp.array(dtype=wp.float64),893    status: wp.array(dtype=wp.int32),894):895    history = wp.tid()896    target = pixel[history]897    participating = int(0)898    contribution = wp.float64(0.0)899    if (900        history_status[history] <= 1901        and history_status[history] >= 0902        and target >= 0903        and target < pixels904        and scale[target] > wp.float64(0.0)905    ):906        contribution = score[history] / scale[target]907        if not wp.isfinite(contribution):908            wp.atomic_or(status, 0, 16)909        else:910            participating = 1911    aggregate = grouped_tally(contribution, target, 0, participating)912    if aggregate[1] > wp.float64(0.0):913        wp.atomic_add(out_sum, target, aggregate[0])914915916@wp.kernel917def finish_mean(918    sums: wp.array(dtype=wp.float64),919    scale: wp.array(dtype=wp.float64),920    count: int,921    out_mean: wp.array(dtype=wp.float64),922    status: wp.array(dtype=wp.int32),923):924    pixel = wp.tid()925    # Scale after dividing the bounded sum: neither a huge raw sum nor a926    # subnormal score divided prematurely by the history count is required.927    value = scale[pixel] * (sums[pixel] / wp.float64(count))928    out_mean[pixel] = value929    if not wp.isfinite(value):930        wp.atomic_or(status, 0, 16)931932933@wp.kernel934def validate_measurement(935    observation: wp.array(dtype=wp.float64),936    weights: wp.array(dtype=wp.float64),937    status: wp.array(dtype=wp.int32),938):939    pixel = wp.tid()940    if (941        not wp.isfinite(observation[pixel])942        or not wp.isfinite(weights[pixel])943        or weights[pixel] < wp.float64(0.0)944    ):945        wp.atomic_or(status, 0, 1)946