python/dpt/transport/estimators.py

1"""Original-history uncertainty and independent products for stochastic inversion.23Detector misses count as zero histories. Photon branches are not independent4samples, and plug-in squared sample means are not unbiased squared expectations.5The product estimators below explicitly require disjoint random-stream identities.6"""78from __future__ import annotations910# Workspace internals are shared only within this package.11# pyright: reportPrivateUsage=false12import math13from dataclasses import dataclass, field14from typing import Any1516from dpt._runtime import require_no_tape17from dpt.statistics import mean_standard_error1819from .forward import TransportWorkspace20from .model import TransportError21from .rng import HistoryBatch, require_independent2223_NO_VARIANCE = object()242526@dataclass(slots=True)27class EstimatorWorkspace:28    """Reusable detector moments; no per-call allocations or hidden downloads."""2930    transport: TransportWorkspace31    _sum: Any = field(repr=False)32    _hits: Any = field(repr=False)33    _scale: Any = field(repr=False)3435    @property36    def scratch_bytes(self) -> int:37        return 20 * self.transport.spec.detector.pixels383940def prepare_estimators(workspace: TransportWorkspace) -> EstimatorWorkspace:41    """Allocate binary64 sum/scale images and int32 hit counts on the owning stream."""42    require_no_tape()43    context = workspace.context44    with context.scope():45        sums = context.wp.empty(46            workspace.spec.detector.pixels, dtype=context.wp.float64, device=context.device47        )48        hits = context.wp.empty(49            workspace.spec.detector.pixels, dtype=context.wp.int32, device=context.device50        )51        scale = context.wp.empty_like(sums)52    return EstimatorWorkspace(workspace, sums, hits, scale)535455def _accumulate_history_mean(56    pixel: Any,57    score: Any,58    history_status: Any,59    *,60    batch: HistoryBatch,61    workspace: EstimatorWorkspace,62    out_mean: Any,63    out_variance_of_mean: Any = _NO_VARIANCE,64    stream: Any,65    validate: bool,66) -> None:67    """Validate shared ownership, clear scratch and launch scaled original-history means."""68    transport = workspace.transport69    context = transport.context70    context.assert_stream(stream)71    require_no_tape()72    wp = context.wp73    pixels = transport.spec.detector.pixels74    for name, value, dtype in (75        ("pixel", pixel, wp.int32),76        ("score", score, wp.float64),77        ("history_status", history_status, wp.int32),78    ):79        context.array(value, name, dtype=dtype, shape=(batch.count,))80    outputs = [("out_mean", out_mean)]81    scratch = [("moment_sum", workspace._sum), ("moment_scale", workspace._scale)]82    if out_variance_of_mean is not _NO_VARIANCE:83        outputs.append(("out_variance_of_mean", out_variance_of_mean))84        scratch.append(("moment_hits", workspace._hits))85    for name, value in outputs:86        context.array(value, name, dtype=wp.float64, shape=(pixels,))87    context.disjoint(88        [89            *transport._reads(),90            ("pixel", pixel),91            ("score", score),92            ("history_status", history_status),93        ],94        outputs + scratch,95    )96    if validate:97        transport.check_status()98    with context.scope():99        workspace._sum.zero_()100        workspace._scale.zero_()101        if out_variance_of_mean is not _NO_VARIANCE:102            workspace._hits.zero_()103            out_variance_of_mean.zero_()104    transport._launch(105        transport._kernels.find_tally_scale,106        batch.count,107        [pixel, score, history_status, pixels, workspace._scale, transport._status],108    )109    transport._launch(110        transport._kernels.tally_sum,111        batch.count,112        [pixel, score, history_status, pixels, workspace._scale, workspace._sum, transport._status],113    )114    transport._launch(115        transport._kernels.finish_mean,116        pixels,117        [workspace._sum, workspace._scale, batch.count, out_mean, transport._status],118    )119120121# region book:transport-original-history-moments122def history_moments(123    pixel: Any,124    score: Any,125    history_status: Any,126    *,127    batch: HistoryBatch,128    workspace: EstimatorWorkspace,129    out_mean: Any,130    out_variance_of_mean: Any,131    stream: Any = None,132    validate: bool = True,133) -> None:134    """Mean and unbiased variance of that mean, for IID original-history scores.135136    Source histories must be independently identically distributed; a repeated137    deterministic ray also qualifies. This sampling-law precondition cannot be138    verified from the realised score arrays.139    Deterministically stratified sources require variance across independent140    *replicated complete batches*, not this within-history formula. The two141    output images contain marginal variances; they do not claim independent142    detector pixels or supply a covariance-free loss-error bar.143144    FP64 scale, mean and centred-deviation passes preserve the result range. Their145    floating summation order is not bitwise reproducible. This analogue engine146    produces at most one contribution per original history; descendant splitting147    must first combine contributions before using these moment operations.148    """149    if batch.count < 2:150        raise TransportError("variance of a mean requires at least two original histories")151    _accumulate_history_mean(152        pixel,153        score,154        history_status,155        batch=batch,156        workspace=workspace,157        out_mean=out_mean,158        out_variance_of_mean=out_variance_of_mean,159        stream=stream,160        validate=validate,161    )162    transport = workspace.transport163    pixels = transport.spec.detector.pixels164    transport._launch(165        transport._kernels.centred_moments,166        batch.count,167        [168            pixel,169            score,170            out_mean,171            workspace._scale,172            out_variance_of_mean,173            workspace._hits,174            transport._status,175        ],176    )177    transport._launch(178        transport._kernels.finish_variance,179        pixels,180        [181            out_mean,182            workspace._scale,183            workspace._hits,184            batch.count,185            out_variance_of_mean,186            transport._status,187        ],188    )189    if validate:190        transport.check_status()191192193# endregion book:transport-original-history-moments194195196def _product(197    mean_a: Any,198    other_b: Any,199    observed: Any,200    weights: Any,201    *,202    batches: tuple[HistoryBatch, HistoryBatch],203    workspace: TransportWorkspace,204    out_components: Any,205    loss: bool,206    stream: Any,207    validate: bool,208) -> None:209    require_independent(*batches)210    require_no_tape()211    context = workspace.context212    context.assert_stream(stream)213    arrays = [214        ("mean_a", mean_a),215        ("other_b", other_b),216        ("observed", observed),217        ("weights", weights),218    ]219    shape = (workspace.spec.detector.pixels,)220    for name, value in [*arrays, ("out_components", out_components)]:221        context.array(value, name, dtype=context.wp.float64, shape=shape)222    context.disjoint(workspace._reads() + arrays, [("out_components", out_components)])223    if validate:224        workspace.check_status()225    workspace._launch(226        workspace._kernels.independent_product,227        shape[0],228        [mean_a, other_b, observed, weights, int(loss), out_components, workspace._status],229    )230    if validate:231        workspace.check_status()232233234# region book:transport-independent-loss-gradient235def independent_squared_gradient(236    mean_a: Any,237    derivative_mean_b: Any,238    observed: Any,239    weights: Any,240    *,241    batches: tuple[HistoryBatch, HistoryBatch],242    workspace: TransportWorkspace,243    out_components: Any,244    stream: Any = None,245    validate: bool = True,246) -> None:247    """Write unbiased per-pixel contributions to the loss-of-expected-score gradient.248249    For fixed finite observation y and nonnegative fixed pixel weight w, the250    target is 1/2 sum_p w_p (E[S_p]-y_p)^2. Independent unbiased estimates A of251    E[S] and B of its derivative give E[w(A-y)B] = w(E[S]-y)dE[S]. Using the same252    histories in both factors generally adds their covariance and is rejected.253    Sum components on device using the shared reduction operator. Estimate the254    scalar gradient's uncertainty across independent product replicates, because255    pixel covariances generally do not vanish.256    """257    _product(258        mean_a,259        derivative_mean_b,260        observed,261        weights,262        batches=batches,263        workspace=workspace,264        out_components=out_components,265        loss=False,266        stream=stream,267        validate=validate,268    )269270271# endregion book:transport-independent-loss-gradient272273274def independent_squared_loss(275    mean_a: Any,276    mean_b: Any,277    observed: Any,278    weights: Any,279    *,280    batches: tuple[HistoryBatch, HistoryBatch],281    workspace: TransportWorkspace,282    out_components: Any,283    stream: Any = None,284    validate: bool = True,285) -> None:286    """Write unbiased loss components from two independent score estimates.287288    Individual estimates may be negative even though the expected squared-error289    loss is nonnegative. Clamping a noisy estimate to zero would introduce bias.290    A candidate-minus-incumbent difference may use common random numbers between291    the two parameter values, provided the two product factors remain independent.292    Independent replicate differences provide the acceptance-policy uncertainty.293    """294    _product(295        mean_a,296        mean_b,297        observed,298        weights,299        batches=batches,300        workspace=workspace,301        out_components=out_components,302        loss=True,303        stream=stream,304        validate=validate,305    )306307308@dataclass(frozen=True, slots=True)309class ReplicateEstimate:310    """A scalar checkpoint, after device reduction and explicit host transfer."""311312    value: float313    batches: tuple[HistoryBatch, ...]314315    def __post_init__(self) -> None:316        if not math.isfinite(self.value) or not self.batches:317            raise TransportError(318                "replicate estimates need a finite value and their batch identities"319            )320        require_independent(*self.batches)321322323@dataclass(frozen=True, slots=True)324class MeanUncertainty:325    mean: float326    standard_error: float327    replicates: int328329330def summarise_replicates(estimates: tuple[ReplicateEstimate, ...]) -> MeanUncertainty:331    """Estimate uncertainty across genuinely independent complete scalar replicates.332333    The caller may pass paired candidate-minus-incumbent objective changes as334    values. Within each pair, common random numbers are permitted; record each335    *distinct* batch once. Across replicate pairs no streams may overlap.336    """337    if len(estimates) < 2:338        raise TransportError("uncertainty requires at least two independent replicates")339    for index, estimate in enumerate(estimates):340        for other in estimates[index + 1 :]:341            require_independent(*estimate.batches, *other.batches)342    count = len(estimates)343    mean, standard_error = mean_standard_error(tuple(estimate.value for estimate in estimates))344    return MeanUncertainty(mean, standard_error, count)345346347def history_mean(348    pixel: Any,349    score: Any,350    history_status: Any,351    *,352    batch: HistoryBatch,353    workspace: EstimatorWorkspace,354    out_mean: Any,355    stream: Any = None,356    validate: bool = True,357) -> None:358    """Accumulate a mean without inventing a within-batch uncertainty estimate.359360    This supports single-history batches and independent replicated batch designs.361    It avoids both the extra squared-score atomic and avoidable score-square362    overflow when only the mean is requested by a nonlinear product estimator.363    """364    _accumulate_history_mean(365        pixel,366        score,367        history_status,368        batch=batch,369        workspace=workspace,370        out_mean=out_mean,371        stream=stream,372        validate=validate,373    )374    if validate:375        workspace.transport.check_status()376