python/dpt/kernels/projection.py

1"""Streaming finite-ray quadrature and an explicit first-order discrete VJP.23One lane traces one ray. Values and pose derivatives use FP64 registers with4FP32 field storage. Reverse recomputes each sample; no ray-by-sample tensor or5image Jacobian is retained. Pose partials use a fixed tree; active-volume6cotangents use FP32 scatter atomics and are not bitwise deterministic.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 wp1920from dpt.kernels.geometry import unpack_rotation, unpack_translation21from dpt.kernels.projection_checks import (22    get_dependency_marker as get_dependency_marker,23)24from dpt.kernels.projection_checks import (25    get_validate_finite as get_validate_finite,26)27from dpt.kernels.projection_checks import (28    validate_finite as validate_finite,29)30from dpt.kernels.projection_checks import (31    validate_nonnegative as validate_nonnegative,32)3334OPTIONS = {"fast_math": False, "fuse_fp": True, "enable_backward": False}35wp.set_module_options(OPTIONS)36BLOCK = 12837Vec12d = wp.types.vector(length=12, dtype=wp.float64)383940@wp.struct41class Configuration:42    shape: wp.vec3i43    object_to_grid: wp.mat33d44    grid_origin: wp.vec3d45    source: wp.vec3d46    detector_origin: wp.vec3d47    column_step: wp.vec3d48    row_step: wp.vec3d49    width: int50    pixels: int51    samples: int525354@wp.struct55class Ray:56    origin: wp.vec3d57    direction: wp.vec3d58    object_source: wp.vec3d59    object_direction: wp.vec3d60    world_source_relative: wp.vec3d61    world_direction: wp.vec3d62    length: wp.float64636465@wp.func66def ray_for_pixel(config: Configuration, pose: wp.array(dtype=wp.float64), pixel: int) -> Ray:67    ray = Ray()68    row, column = pixel // config.width, pixel % config.width69    end = (70        config.detector_origin71        + wp.float64(column) * config.column_step72        + wp.float64(row) * config.row_step73    )74    inverse = wp.transpose(unpack_rotation(pose))75    ray.world_source_relative = config.source - unpack_translation(pose)76    ray.world_direction = end - config.source77    ray.object_source = inverse * ray.world_source_relative78    ray.object_direction = inverse * ray.world_direction79    ray.origin = config.object_to_grid * (ray.object_source - config.grid_origin)80    ray.direction = config.object_to_grid * ray.object_direction81    ray.length = wp.length(ray.world_direction)82    return ray838485@wp.func86def finite_point(point: wp.vec3d) -> bool:87    return wp.isfinite(point[0]) and wp.isfinite(point[1]) and wp.isfinite(point[2])888990@wp.func91def valid_ray(ray: Ray) -> bool:92    return (93        finite_point(ray.origin)94        and finite_point(ray.direction)95        and finite_point(ray.object_source)96        and finite_point(ray.object_direction)97        and finite_point(ray.world_source_relative)98        and finite_point(ray.world_direction)99        and wp.isfinite(ray.length)100        and ray.length > wp.float64(0.0)101    )102103104@wp.func105def inside_support(point: wp.vec3d, shape: wp.vec3i) -> bool:106    return (107        finite_point(point)108        and point[0] >= wp.float64(-0.5)109        and point[1] >= wp.float64(-0.5)110        and point[2] >= wp.float64(-0.5)111        and point[0] <= wp.float64(shape[0]) - wp.float64(0.5)112        and point[1] <= wp.float64(shape[1]) - wp.float64(0.5)113        and point[2] <= wp.float64(shape[2]) - wp.float64(0.5)114    )115116117@wp.struct118class Interval:119    lower: wp.float64120    upper: wp.float64121    lower_origin_gradient: wp.vec3d122    lower_direction_gradient: wp.vec3d123    upper_origin_gradient: wp.vec3d124    upper_direction_gradient: wp.vec3d125126127# region book:projection-active-intersection128@wp.func129def finite_interval(origin: wp.vec3d, direction: wp.vec3d, shape: wp.vec3i) -> Interval:130    """Clip a finite segment and retain the active face derivatives.131132    Exact parallelism is handled separately; arbitrary epsilon thresholds would133    change the physical interval. At tied faces a derivative is not unique:134    strict comparisons choose the first active axis, a documented branch value.135    """136    interval = Interval()137    interval.lower = wp.float64(0.0)138    interval.upper = wp.float64(1.0)139    for axis in range(3):140        low = wp.float64(-0.5)141        high = wp.float64(shape[axis]) - wp.float64(0.5)142        velocity = direction[axis]143        if velocity == wp.float64(0.0):144            if origin[axis] < low or origin[axis] > high:145                interval.upper = wp.float64(-1.0)146        else:147            first = (low - origin[axis]) / velocity148            last = (high - origin[axis]) / velocity149            if first > last:150                temporary = first151                first = last152                last = temporary153            if first > interval.lower:154                interval.lower = first155                interval.lower_origin_gradient = wp.vec3d(0.0)156                interval.lower_direction_gradient = wp.vec3d(0.0)157                interval.lower_origin_gradient[axis] = -wp.float64(1.0) / velocity158                interval.lower_direction_gradient[axis] = -first / velocity159            if last < interval.upper:160                interval.upper = last161                interval.upper_origin_gradient = wp.vec3d(0.0)162                interval.upper_direction_gradient = wp.vec3d(0.0)163                interval.upper_origin_gradient[axis] = -wp.float64(1.0) / velocity164                interval.upper_direction_gradient[axis] = -last / velocity165    return interval166167168# endregion book:projection-active-intersection169170171@wp.struct172class Sample:173    value: wp.float64174    gradient: wp.vec3d175176177# region book:projection-clamped-trilinear178@wp.func179def sample_field(field: wp.array(dtype=wp.float32), point: wp.vec3d, shape: wp.vec3i) -> Sample:180    """Value and grid-coordinate slope of the half-cell-extended sampled field."""181    result = Sample()182    if not inside_support(point, shape):183        return result184    lower = wp.vec3i()185    fraction = wp.vec3d()186    slope = wp.vec3d()187    for axis in range(3):188        coordinate = wp.clamp(point[axis], wp.float64(0.0), wp.float64(shape[axis] - 1))189        lower[axis] = wp.min(int(wp.floor(coordinate)), wp.max(shape[axis] - 2, 0))190        fraction[axis] = coordinate - wp.float64(lower[axis])191        if point[axis] >= wp.float64(0.0) and point[axis] < wp.float64(shape[axis] - 1):192            slope[axis] = wp.float64(1.0)193    # Separable interpolation reuses the same eight loads for value and slope.194    # Form differences in FP64: FP32 subtraction would lose small field slopes.195    # The compiler removes slope work from the forward-only caller.196    x0, y0, z0 = lower[0], lower[1], lower[2]197    x1 = wp.min(x0 + 1, shape[0] - 1)198    y1 = wp.min(y0 + 1, shape[1] - 1)199    z1 = wp.min(z0 + 1, shape[2] - 1)200    row00 = (z0 * shape[1] + y0) * shape[0]201    row10 = (z0 * shape[1] + y1) * shape[0]202    row01 = (z1 * shape[1] + y0) * shape[0]203    row11 = (z1 * shape[1] + y1) * shape[0]204    v000 = wp.float64(field[row00 + x0])205    v100 = wp.float64(field[row00 + x1])206    v010 = wp.float64(field[row10 + x0])207    v110 = wp.float64(field[row10 + x1])208    v001 = wp.float64(field[row01 + x0])209    v101 = wp.float64(field[row01 + x1])210    v011 = wp.float64(field[row11 + x0])211    v111 = wp.float64(field[row11 + x1])212    dx00, dx10 = v100 - v000, v110 - v010213    dx01, dx11 = v101 - v001, v111 - v011214    fx, fy, fz = fraction[0], fraction[1], fraction[2]215    # Weighted lerp preserves a small endpoint next to a huge neighbour.216    # a+t*(b-a) would erase b at t=1 when the difference rounds to -a.217    ax, ay, az = wp.float64(1.0) - fx, wp.float64(1.0) - fy, wp.float64(1.0) - fz218    x00, x10 = ax * v000 + fx * v100, ax * v010 + fx * v110219    x01, x11 = ax * v001 + fx * v101, ax * v011 + fx * v111220    xy0, xy1 = ay * x00 + fy * x10, ay * x01 + fy * x11221    result.value = az * xy0 + fz * xy1222    # Differentiate corner values before interpolation. Subtracting two223    # already interpolated values can erase a small slope beside a huge224    # orthogonal background, even when its cotangent is representable.225    dx0, dx1 = ay * dx00 + fy * dx10, ay * dx01 + fy * dx11226    dy0 = ax * (v010 - v000) + fx * (v110 - v100)227    dy1 = ax * (v011 - v001) + fx * (v111 - v101)228    dz0 = ax * (v001 - v000) + fx * (v101 - v100)229    dz1 = ax * (v011 - v010) + fx * (v111 - v110)230    result.gradient = wp.vec3d(231        slope[0] * (az * dx0 + fz * dx1),232        slope[1] * (az * dy0 + fz * dy1),233        slope[2] * (ay * dz0 + fy * dz1),234    )235    return result236237238# endregion book:projection-clamped-trilinear239240241@wp.func242def scatter_field(243    gradient: wp.array(dtype=wp.float32),244    point: wp.vec3d,245    shape: wp.vec3i,246    seed: wp.float64,247    status: wp.array(dtype=wp.int32),248):249    if not finite_point(point):250        wp.atomic_or(status, 0, 2)251    if not inside_support(point, shape):252        return253    lower = wp.vec3i()254    fraction = wp.vec3d()255    for axis in range(3):256        coordinate = wp.clamp(point[axis], wp.float64(0.0), wp.float64(shape[axis] - 1))257        lower[axis] = wp.min(int(wp.floor(coordinate)), wp.max(shape[axis] - 2, 0))258        fraction[axis] = coordinate - wp.float64(lower[axis])259    for corner in range(8):260        bit = wp.vec3i(corner & 1, (corner >> 1) & 1, (corner >> 2) & 1)261        index = wp.vec3i()262        contribution = seed263        for axis in range(3):264            index[axis] = wp.min(lower[axis] + bit[axis], shape[axis] - 1)265            if bit[axis] == 0:266                contribution *= wp.float64(1.0) - fraction[axis]267            else:268                contribution *= fraction[axis]269        rounded = wp.float32(contribution)270        if not wp.isfinite(rounded):271            wp.atomic_or(status, 0, 2)272        wp.atomic_add(gradient, (index[2] * shape[1] + index[1]) * shape[0] + index[0], rounded)273274275@wp.struct276class CellTraversal:277    plane: wp.vec3i278    step: wp.vec3i279    next: wp.vec3d280281282@wp.func283def cell_crossing(284    ray: Ray, axis: int, plane: int, shape: wp.vec3i, upper: wp.float64285) -> wp.float64:286    if ray.direction[axis] == wp.float64(0.0) or plane < 0 or plane >= shape[axis]:287        return upper288    return wp.min((wp.float64(plane) - ray.origin[axis]) / ray.direction[axis], upper)289290291@wp.func292def begin_cells(ray: Ray, interval: Interval, shape: wp.vec3i) -> CellTraversal:293    """Next strictly forward interpolation plane, including half-cell edge slabs."""294    traversal = CellTraversal()295    entry = ray.origin + interval.lower * ray.direction296    for axis in range(3):297        velocity = ray.direction[axis]298        if velocity > wp.float64(0.0):299            traversal.step[axis] = 1300            traversal.plane[axis] = wp.max(int(wp.ceil(entry[axis])), 0)301        elif velocity < wp.float64(0.0):302            traversal.step[axis] = -1303            traversal.plane[axis] = wp.min(int(wp.floor(entry[axis])), shape[axis] - 1)304        crossing = cell_crossing(ray, axis, traversal.plane[axis], shape, interval.upper)305        if crossing <= interval.lower:306            traversal.plane[axis] += traversal.step[axis]307            crossing = cell_crossing(ray, axis, traversal.plane[axis], shape, interval.upper)308        traversal.next[axis] = crossing309    return traversal310311312@wp.func313def advance_cells(314    traversal: CellTraversal, ray: Ray, shape: wp.vec3i, end: wp.float64, upper: wp.float64315) -> CellTraversal:316    for axis in range(3):317        # Advance every tied plane; no epsilon changes which interval is integrated.318        if traversal.next[axis] <= end:319            traversal.plane[axis] += traversal.step[axis]320            traversal.next[axis] = cell_crossing(ray, axis, traversal.plane[axis], shape, upper)321    return traversal322323324@wp.func325def support_value(326    field: wp.array(dtype=wp.float32), point: wp.vec3d, shape: wp.vec3i327) -> wp.float64:328    # A clipped endpoint can round a few ulps outside its active support face.329    bounded = point330    for axis in range(3):331        bounded[axis] = wp.clamp(332            point[axis], wp.float64(-0.5), wp.float64(shape[axis]) - wp.float64(0.5)333        )334    return sample_field(field, bounded, shape).value335336337# region book:projection-streamed-quadrature338@cache339def get_forward(double_output: bool = False, cell_gauss: bool = False):340    dtype = wp.float64 if double_output else wp.float32341342    @wp.kernel(module="unique", module_options=OPTIONS)343    def forward(344        field: wp.array(dtype=wp.float32),345        pose: wp.array(dtype=wp.float64),346        config: Configuration,347        output: wp.array(dtype=dtype),348        status: wp.array(dtype=wp.int32),349    ):350        pixel = wp.tid()351        ray = ray_for_pixel(config, pose, pixel)352        if not valid_ray(ray):353            wp.atomic_or(status, 0, 2)354            output[pixel] = dtype(0.0)355            return356        interval = finite_interval(ray.origin, ray.direction, config.shape)357        integral = wp.float64(0.0)358        if interval.upper > interval.lower:359            if wp.static(cell_gauss):360                traversal = begin_cells(ray, interval, config.shape)361                start = interval.lower362                segments = wp.int64(0)363                limit = (364                    wp.int64(config.shape[0])365                    + wp.int64(config.shape[1])366                    + wp.int64(config.shape[2])367                    + wp.int64(1)368                )369                while start < interval.upper and segments < limit:370                    end = wp.min(traversal.next[0], wp.min(traversal.next[1], traversal.next[2]))371                    if end <= start:372                        wp.atomic_or(status, 0, 2)373                        break374                    radius = (end - start) * wp.float64(0.5)375                    centre = start + radius376                    offset = radius / wp.sqrt(wp.float64(3.0))377                    for node in range(2):378                        position = centre + wp.float64(2 * node - 1) * offset379                        integral += (380                            radius381                            * sample_field(382                                field, ray.origin + position * ray.direction, config.shape383                            ).value384                        )385                    start = end386                    traversal = advance_cells(traversal, ray, config.shape, end, interval.upper)387                    segments += wp.int64(1)388                if start < interval.upper:389                    wp.atomic_or(status, 0, 2)390                integral *= ray.length391            else:392                step = (interval.upper - interval.lower) / wp.float64(config.samples)393                for sample in range(config.samples):394                    position = interval.lower + (wp.float64(sample) + wp.float64(0.5)) * step395                    point = ray.origin + position * ray.direction396                    if not finite_point(point):397                        wp.atomic_or(status, 0, 2)398                    integral += sample_field(field, point, config.shape).value399                integral *= ray.length * step400        output[pixel] = dtype(integral)401        if not wp.isfinite(output[pixel]):402            wp.atomic_or(status, 0, 2)403404    return forward405406407# endregion book:projection-streamed-quadrature408409410# region book:projection-discrete-vjp411@cache412def get_vjp(413    active_volume: bool,414    matrix_gradient: bool,415    active_pose: bool,416    *,417    per_ray: bool = False,418    double_seed: bool = False,419    cell_gauss: bool = False,420):421    if per_ray and (active_volume or matrix_gradient or not active_pose):422        raise ValueError("per-ray diagnostics require only the six local pose coordinates")423    components = 12 if matrix_gradient else 6424    seed_dtype = wp.float64 if double_seed else wp.float32425426    @wp.kernel(module="unique", module_options=OPTIONS)427    def vjp(428        field: wp.array(dtype=wp.float32),429        pose: wp.array(dtype=wp.float64),430        config: Configuration,431        seeds: wp.array(dtype=seed_dtype),432        volume_gradient: wp.array(dtype=wp.float32),433        partials: wp.array(dtype=wp.float64),434        status: wp.array(dtype=wp.int32),435    ):436        block, lane = wp.tid()437        pixel = block * BLOCK + lane438        gradient = Vec12d()439        if pixel < config.pixels:440            ray = ray_for_pixel(config, pose, pixel)441            if not valid_ray(ray):442                wp.atomic_or(status, 0, 2)443            else:444                interval = finite_interval(ray.origin, ray.direction, config.shape)445                span = interval.upper - interval.lower446                if span > wp.float64(0.0):447                    origin_gradient = wp.vec3d()448                    direction_gradient = wp.vec3d()449                    factor = wp.float64(seeds[pixel]) * ray.length450                    if wp.static(cell_gauss):451                        traversal = begin_cells(ray, interval, config.shape)452                        start = interval.lower453                        segments = wp.int64(0)454                        limit = (455                            wp.int64(config.shape[0])456                            + wp.int64(config.shape[1])457                            + wp.int64(config.shape[2])458                            + wp.int64(1)459                        )460                        while start < interval.upper and segments < limit:461                            end = wp.min(462                                traversal.next[0], wp.min(traversal.next[1], traversal.next[2])463                            )464                            if end <= start:465                                wp.atomic_or(status, 0, 2)466                                break467                            radius = (end - start) * wp.float64(0.5)468                            centre = start + radius469                            offset = radius / wp.sqrt(wp.float64(3.0))470                            for node in range(2):471                                position = centre + wp.float64(2 * node - 1) * offset472                                point = ray.origin + position * ray.direction473                                if wp.static(active_pose):474                                    value = sample_field(field, point, config.shape)475                                    origin_gradient += radius * value.gradient476                                    direction_gradient += radius * position * value.gradient477                                if wp.static(active_volume):478                                    scatter_field(479                                        volume_gradient,480                                        point,481                                        config.shape,482                                        factor * radius,483                                        status,484                                    )485                            start = end486                            traversal = advance_cells(487                                traversal, ray, config.shape, end, interval.upper488                            )489                            segments += wp.int64(1)490                        if start < interval.upper:491                            wp.atomic_or(status, 0, 2)492                        if wp.static(active_pose):493                            # Field continuity cancels all internal moving-cell terms.494                            # The zero-extended support still has external endpoint terms.495                            lower_value = support_value(496                                field, ray.origin + interval.lower * ray.direction, config.shape497                            )498                            upper_value = support_value(499                                field, ray.origin + interval.upper * ray.direction, config.shape500                            )501                            origin_gradient += (502                                upper_value * interval.upper_origin_gradient503                                - lower_value * interval.lower_origin_gradient504                            )505                            direction_gradient += (506                                upper_value * interval.upper_direction_gradient507                                - lower_value * interval.lower_direction_gradient508                            )509                    else:510                        mean = wp.float64(0.0)511                        direct_origin = wp.vec3d()512                        direct_direction = wp.vec3d()513                        shift_lower = wp.float64(0.0)514                        shift_upper = wp.float64(0.0)515                        for sample in range(config.samples):516                            alpha = (wp.float64(sample) + wp.float64(0.5)) / wp.float64(517                                config.samples518                            )519                            position = interval.lower + alpha * span520                            point = ray.origin + position * ray.direction521                            if not finite_point(point):522                                wp.atomic_or(status, 0, 2)523                            if wp.static(active_pose):524                                value = sample_field(field, point, config.shape)525                                mean += value.value526                                direct_origin += value.gradient527                                direct_direction += position * value.gradient528                                along = wp.dot(value.gradient, ray.direction)529                                shift_lower += (wp.float64(1.0) - alpha) * along530                                shift_upper += alpha * along531                            if wp.static(active_volume):532                                scatter_field(533                                    volume_gradient,534                                    point,535                                    config.shape,536                                    wp.float64(seeds[pixel])537                                    * ray.length538                                    * span539                                    / wp.float64(config.samples),540                                    status,541                                )542                        if wp.static(active_pose):543                            lower_seed = -mean + span * shift_lower544                            upper_seed = mean + span * shift_upper545                            origin_gradient = (546                                span * direct_origin547                                + lower_seed * interval.lower_origin_gradient548                                + upper_seed * interval.upper_origin_gradient549                            )550                            direction_gradient = (551                                span * direct_direction552                                + lower_seed * interval.lower_direction_gradient553                                + upper_seed * interval.upper_direction_gradient554                            )555                        factor /= wp.float64(config.samples)556                    if wp.static(active_pose):557                        source_adjoint = factor * (558                            wp.transpose(config.object_to_grid) * origin_gradient559                        )560                        direction_adjoint = factor * (561                            wp.transpose(config.object_to_grid) * direction_gradient562                        )563                        if wp.static(matrix_gradient):564                            for row in range(3):565                                for column in range(3):566                                    gradient[3 * row + column] = (567                                        ray.world_source_relative[row] * source_adjoint[column]568                                        + ray.world_direction[row] * direction_adjoint[column]569                                    )570                            translation_adjoint = -(unpack_rotation(pose) * source_adjoint)571                            for axis in range(3):572                                gradient[9 + axis] = translation_adjoint[axis]573                        else:574                            rotation_adjoint = wp.cross(575                                source_adjoint, ray.object_source576                            ) + wp.cross(direction_adjoint, ray.object_direction)577                            for axis in range(3):578                                gradient[axis] = -source_adjoint[axis]579                                gradient[axis + 3] = rotation_adjoint[axis]580        # The diagnostic writes the same ray derivative before any reduction.581        # Its caller owns O(6 P) storage; the production specialisation retains582        # only reduction partials and performs exactly its existing reduction.583        if wp.static(per_ray):584            if pixel < config.pixels:585                for component in range(components):586                    component_value = gradient[component]587                    partials[wp.int64(pixel) * wp.int64(components) + wp.int64(component)] = (588                        component_value589                    )590                    if not wp.isfinite(component_value):591                        wp.atomic_or(status, 0, 2)592        # All lanes, including padded rays, participate in the fixed reduction.593        elif wp.static(active_pose):594            for component in range(components):595                values = wp.tile(gradient[component])596                total = wp.tile_sum(values)597                wp.tile_store(partials, total, offset=block * components + component)598599    return vjp600601602# endregion book:projection-discrete-vjp603