python/dpt/kernels/objectives.py

1"""Fused objective/seed evaluation with a fixed FP64 reduction tree.23The image-space derivative is caller-owned. Only one scalar partial per tile is4retained: no per-pixel loss image or global floating-point atomic is needed.5"""67# Warp annotations are executable DSL expressions; host interfaces remain strict.8# The optional GPU import is resolved only when an operator is prepared.9# pyright: reportInvalidTypeForm=false, reportUnknownParameterType=false10# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false11# pyright: reportUnknownVariableType=false, reportUntypedFunctionDecorator=false12# pyright: reportMissingImports=false, reportUntypedClassDecorator=false1314from functools import cache1516import warp as wp1718TILE = 25619OPTIONS = {"fast_math": False, "fuse_fp": True, "enable_backward": False}202122@wp.func_native("return log1p(x);")23def log_one_plus(x: wp.float64) -> wp.float64: ...242526@cache27def observation_validation_kernel(poisson: bool, weighted: bool, masked: bool = False):28    @wp.kernel(module="unique", module_options=OPTIONS)29    def validate(30        observation: wp.array(dtype=wp.float32),31        weights: wp.array(dtype=wp.float32),32        valid: wp.array(dtype=wp.uint8),33        status: wp.array(dtype=wp.int32),34    ):35        p = wp.tid()36        active = bool(True)  # noqa: UP018 - Warp runtime Boolean.37        if wp.static(masked):38            if valid[p] > wp.uint8(1):39                wp.atomic_or(status, 0, 1)40            active = valid[p] == wp.uint8(1)41        if active:42            value = observation[p]43            invalid = not wp.isfinite(value)44            if wp.static(poisson):45                invalid = invalid or value < 0.0 or wp.floor(value) != value46            if wp.static(weighted):47                invalid = invalid or not wp.isfinite(weights[p]) or weights[p] < 0.048            if invalid:49                wp.atomic_or(status, 0, 1)5051    return validate525354@cache55def validation_kernel(56    poisson: bool, weighted: bool, masked: bool = False, precision: str = "float32"57):58    value_dtype = wp.float64 if precision == "float64" else wp.float325960    @wp.kernel(module="unique", module_options=OPTIONS)61    def validate(62        prediction: wp.array(dtype=value_dtype),63        observation: wp.array(dtype=wp.float32),64        weights: wp.array(dtype=wp.float32),65        valid: wp.array(dtype=wp.uint8),66        status: wp.array(dtype=wp.int32),67    ):68        p = wp.tid()69        active = bool(True)  # noqa: UP018 - Warp runtime Boolean.70        if wp.static(masked):71            if valid[p] > wp.uint8(1):72                wp.atomic_or(status, 0, 1)73            active = valid[p] == wp.uint8(1)74        if active:75            a, b = prediction[p], observation[p]76            invalid = not wp.isfinite(a) or not wp.isfinite(b)77            if wp.static(poisson):78                invalid = invalid or a < 0.0 or b < 0.0 or wp.floor(b) != b79                invalid = invalid or (a == 0.0 and b > 0.0)80            if wp.static(weighted):81                invalid = invalid or not wp.isfinite(weights[p]) or weights[p] < 0.082            if invalid:83                wp.atomic_or(status, 0, 1)8485    return validate868788@wp.func89def poisson_half_deviance(predicted: wp.float64, observed: wp.float64) -> wp.float64:90    """Canonical cancellation-safe Poisson term shared by both compositions."""91    if observed == wp.float64(0.0):92        return predicted93    difference = predicted - observed94    loss = wp.float64(0.0)95    # Half-deviance differs from the negative log likelihood96    # only by observation constants. This form retains a small97    # residual near equality without evaluating log(0) for y=0.98    if wp.abs(difference) <= wp.float64(0.25) * observed:99        ratio = difference / observed100        if wp.abs(ratio) < wp.float64(0.001):101            # r-log1p(r) also cancels. Through degree six the102            # omitted relative term is < 3e-16 for |r|<1e-3.103            series = wp.float64(0.5) + ratio * (104                (-wp.float64(1.0) / wp.float64(3.0))105                + ratio106                * (107                    wp.float64(0.25)108                    + ratio * (wp.float64(-0.2) + ratio * (wp.float64(1.0) / wp.float64(6.0)))109                )110            )111            loss = observed * ratio * ratio * series112        else:113            loss = observed * (ratio - log_one_plus(ratio))114    else:115        loss = difference + observed * (wp.log(observed) - wp.log(predicted))116    return loss117118119# region book:objective-loss-and-seed120@cache121def evaluation_kernel(122    poisson: bool,123    weighted: bool,124    gradient: bool,125    masked: bool = False,126    precision: str = "float32",127):128    value_dtype = wp.float64 if precision == "float64" else wp.float32129130    @wp.kernel(module="unique", module_options=OPTIONS)131    def evaluate(132        prediction: wp.array(dtype=value_dtype),133        observation: wp.array(dtype=wp.float32),134        weights: wp.array(dtype=wp.float32),135        valid: wp.array(dtype=wp.uint8),136        size: int,137        normalisation: wp.float64,138        seed: wp.array(dtype=value_dtype),139        partials: wp.array(dtype=wp.float64),140        status: wp.array(dtype=wp.int32),141    ):142        block, lane = wp.tid()143        p = block * TILE + lane144        loss = wp.float64(0.0)145        if p < size:146            active = bool(True)  # noqa: UP018 - Warp runtime Boolean.147            if wp.static(masked):148                if valid[p] > wp.uint8(1):149                    wp.atomic_or(status, 0, 1)150                active = valid[p] == wp.uint8(1)151            if active:152                predicted = wp.float64(prediction[p])153                observed = wp.float64(observation[p])154                weight = normalisation155                if wp.static(weighted):156                    weight *= wp.float64(weights[p])157                difference = predicted - observed158                derivative = difference159                loss = wp.float64(0.5) * difference * difference160                if wp.static(poisson):161                    if observed == wp.float64(0.0):162                        loss = predicted163                        derivative = wp.float64(1.0)164                    else:165                        derivative = difference / predicted166                        loss = poisson_half_deviance(predicted, observed)167                if wp.static(precision == "float64" and not poisson):168                    # A small fixed weight can keep the final squared loss169                    # representable even when the unweighted square overflows.170                    loss = (wp.float64(0.5) * weight * difference) * difference171                else:172                    loss *= weight173                if not wp.isfinite(loss):174                    wp.atomic_or(status, 0, 2)175                if wp.static(gradient):176                    weighted_derivative = weight * derivative177                    if wp.static(precision == "float64" and poisson):178                        if not wp.isfinite(derivative):179                            # FP32 observations/weights bound this numerator;180                            # do not overflow a removable unweighted quotient.181                            weighted_derivative = (weight * difference) / predicted182                    result = value_dtype(weighted_derivative)183                    seed[p] = result184                    if not wp.isfinite(result):185                        wp.atomic_or(status, 0, 2)186            elif wp.static(gradient):187                seed[p] = value_dtype(0.0)188        total = wp.tile_sum(wp.tile(loss))189        wp.tile_store(partials, total, offset=block)190191    return evaluate192193194# endregion book:objective-loss-and-seed195196197@cache198def primary_evaluation_kernel(poisson: bool, counts: bool, weighted: bool, masked: bool):199    """Retain the primary signal and depth cotangent in FP64 for pose fitting."""200201    @wp.kernel(module="unique", module_options=OPTIONS)202    def evaluate(203        depth: wp.array(dtype=wp.float64),204        beam: wp.float64,205        observation: wp.array(dtype=wp.float32),206        weights: wp.array(dtype=wp.float32),207        valid: wp.array(dtype=wp.uint8),208        size: int,209        normalisation: wp.float64,210        prediction: wp.array(dtype=wp.float64),211        depth_seed: wp.array(dtype=wp.float64),212        partials: wp.array(dtype=wp.float64),213        status: wp.array(dtype=wp.int32),214    ):215        block, lane = wp.tid()216        p = block * TILE + lane217        loss = wp.float64(0.0)218        if p < size:219            optical = depth[p]220            predicted = -optical221            if wp.static(counts):222                predicted = beam * wp.exp(-optical)223                if optical > wp.float64(700.0) and beam > wp.float64(0.0):224                    # Combine exponents before a tiny transmission can underflow.225                    # Keep the direct product in the ordinary near-equality range.226                    predicted = wp.exp(wp.log(beam) - optical)227            prediction[p] = predicted228            if not wp.isfinite(optical) or optical < wp.float64(0.0) or not wp.isfinite(predicted):229                wp.atomic_or(status, 0, 2)230            active = bool(True)  # noqa: UP018 - Warp runtime Boolean.231            if wp.static(masked):232                if valid[p] > wp.uint8(1):233                    wp.atomic_or(status, 0, 1)234                active = valid[p] == wp.uint8(1)235            derivative = wp.float64(0.0)236            if active:237                observed = wp.float64(observation[p])238                difference = predicted - observed239                loss = wp.float64(0.5) * difference * difference240                derivative = -difference241                if wp.static(counts):242                    derivative *= predicted243                if wp.static(poisson):244                    if predicted == wp.float64(0.0) and observed > wp.float64(0.0):245                        wp.atomic_or(status, 0, 2)246                    loss = poisson_half_deviance(predicted, observed)247                    if (248                        observed > wp.float64(0.0)249                        and wp.abs(difference) > wp.float64(0.25) * observed250                    ):251                        # Avoid taking log of a rounded subnormal count.252                        loss = difference + observed * (wp.log(observed) - wp.log(beam) + optical)253                    # Compose (1-N/lambda)*(-lambda) before rounding/division.254                    derivative = -difference255                weight = normalisation256                if wp.static(weighted):257                    weight *= wp.float64(weights[p])258                loss *= weight259                derivative *= weight260            depth_seed[p] = derivative261            if not wp.isfinite(loss) or not wp.isfinite(derivative):262                wp.atomic_or(status, 0, 2)263        wp.tile_store(partials, wp.tile_sum(wp.tile(loss)), offset=block)264265    return evaluate266