python/dpt/transport/derivatives.py

1"""First-order expected-score derivatives with a reviewed parameter boundary.23For analogue sampling and log density eta_m, the path-density derivative is4N_m - integral_path rho_m (mu_a,m(E) + mu_s,m(E)) ds. Conditional interaction-type5and angular/energy distributions do not depend on rho_m, so their *combined*6collision density contributes one per collision. The integral includes every7survival segment, including the final escape. Ignoring that final term biases8transmission derivatives even when no collision occurs.910Continuous absorption instead samples only scattering events. For a fixed11scattering path its likelihood derivative is N_s,m - integral rho_m mu_s,m(E) ds;12the explicit absorption weight contributes -integral rho_m mu_a,m(E) ds.13The resulting expected-score derivative is X times their sum. Conditional14Compton angles and the energy sequence are independent of density at a fixed15path, so energy-dependent coefficients enter each segment's integral without an16extra angular derivative. This assumes fixed coefficient zeros, positive active17densities, fixed source/geometry and sufficient integrability to differentiate18the expectation. No derivative of a realised sampled collision location is used.1920This is a likelihood-ratio estimator, not automatic differentiation through a21realised random trace. The same counter addresses permit deterministic replay;22pathwise support/boundary derivatives are not implied by that replay facility.23"""2425from __future__ import annotations2627# Workspace internals are shared only within this package.28# pyright: reportPrivateUsage=false29from dataclasses import dataclass30from typing import Any, Literal3132from ._constants import LOG_SOURCE_AMPLITUDE, SOURCE_AMPLITUDE33from .forward import TransportWorkspace, _validate_call34from .model import TransportError35from .rng import HistoryBatch363738@dataclass(frozen=True, slots=True)39class DerivativeCapability:40    parameter: str41    supported: bool42    density_term: str43    pathwise_term: str44    support_assumptions: str454647CAPABILITIES = (48    DerivativeCapability(49        "log-material-density",50        True,51        "analogue: collisions minus total depth; "52        "continuous absorption: scatterings minus scattering depth",53        "analogue: zero; continuous absorption: minus absorption depth times weighted score",54        "fixed geometry; positive density; fixed partial-coefficient ratios and conditional laws",55    ),56    DerivativeCapability(57        "source-amplitude",58        True,59        "zero",60        "base detector score, including at zero amplitude",61        "fixed source sampling and importance weights; amplitude multiplies measurement only",62    ),63    DerivativeCapability(64        "log-source-amplitude",65        True,66        "zero",67        "amplitude times base detector score",68        "strictly positive amplitude; fixed source distribution and importance weights",69    ),70    DerivativeCapability(71        "moving-geometry",72        False,73        "not implemented",74        "not implemented",75        "requires a boundary-aware estimator for material visibility and detector edges",76    ),77    DerivativeCapability(78        "energy-or-angular-law",79        False,80        "not implemented",81        "not implemented",82        "requires derivatives of the conditional law, coefficient interpolation and energy score",83    ),84    DerivativeCapability(85        "source-distribution",86        False,87        "not implemented",88        "not implemented",89        "sampling-density and source-support terms must be derived before activation",90    ),91)929394@dataclass(frozen=True, slots=True)95class TransportParameter:96    kind: Literal["log-material-density", "source-amplitude", "log-source-amplitude"]97    material: int | None = None9899    def __post_init__(self) -> None:100        if self.kind == "log-material-density":101            if type(self.material) is not int or self.material < 0:102                raise TransportError("a log-density derivative requires a nonnegative material ID")103        elif self.kind in ("source-amplitude", "log-source-amplitude"):104            if self.material is not None:105                raise TransportError("source amplitude is global and has no material ID")106        else:107            raise TransportError(108                f"unsupported transport derivative: {self.kind!r}; see CAPABILITIES"109            )110111112# region book:transport-derivative-contract113def derivative_histories(114    positions: Any,115    directions: Any,116    weights: Any,117    density: Any,118    *,119    parameter: TransportParameter,120    batch: HistoryBatch,121    workspace: TransportWorkspace,122    out_pixel: Any,123    out_derivative: Any,124    out_status: Any,125    source_amplitude: float = 1.0,126    stream: Any = None,127    validate: bool = True,128) -> None:129    """Replay complete histories and write their sparse expected-score derivatives.130131    One selected parameter costs one complete-history launch and O(H) caller132    outputs; no history-by-event tape or material-by-history scratch is retained.133    Several parameters can use sequential calls and reuse these buffers. This134    trades recomputation for an explicit bounded memory footprint; its useful135    parameter-count range requires profiling. Binary64 signed scores are kept136    before reduction so uncertainty remains defined at the original-history level.137138    Forward execution is not a prerequisite. Supplying an earlier batch identity139    and unchanged inputs reproduces its paths; independent score and derivative140    estimates for nonlinear losses must instead use disjoint batches.141    The prepared estimator selects both the flight law and derivative measure.142    Each output combines original factors and absorption log weight independently143    so forward-score underflow cannot erase a representable derivative.144    """145    material = SOURCE_AMPLITUDE if parameter.material is None else parameter.material146    if parameter.kind == "log-source-amplitude":147        if source_amplitude <= 0:148            raise TransportError("log-amplitude derivatives require strictly positive amplitude")149        material = LOG_SOURCE_AMPLITUDE150    if material >= workspace.spec.grid.materials:151        raise TransportError("active material is outside this workspace")152    wp = workspace.context.wp153    _validate_call(154        workspace,155        batch,156        positions,157        directions,158        weights,159        density,160        [161            ("out_pixel", out_pixel, wp.int32),162            ("out_derivative", out_derivative, wp.float64),163            ("out_status", out_status, wp.int32),164        ],165        source_amplitude,166        stream,167        validate,168    )169    workspace._launch(170        workspace._kernels.derivative_histories,171        batch.count,172        [173            positions,174            directions,175            weights,176            density,177            *workspace._model_inputs(),178            wp.uint64(batch.seed),179            wp.uint64(batch.first_history),180            material,181            source_amplitude,182            out_pixel,183            out_derivative,184            out_status,185            workspace._status,186        ],187    )188    if validate:189        workspace.check_status()190191192# endregion book:transport-derivative-contract193