python/dpt/transport/forward.py

1"""Prepared CUDA transport, sparse history scoring and mandatory completion checks.23Device arrays are caller-owned. Model buffers remain immutable for the workspace4lifetime; dynamic density/source arrays remain unchanged until the owning stream5finishes. Preparation uploads only the small declared energy grid. Repeated calls6allocate no device arrays. Explicitly unchecked calls require ``check_status``7before accepting outputs; incomplete histories never become absorbed histories.8"""910from __future__ import annotations1112# Workspace internals are shared only within this package.13# pyright: reportPrivateUsage=false14import math15from dataclasses import dataclass, field16from typing import Any1718from dpt._runtime import DeviceContext, load_kernels, prepare_context, require_no_tape19from dpt.contracts import NumericalError2021from .model import IncompleteHistoryError, TransportError, TransportSpec, _physical22from .rng import HistoryBatch232425@dataclass(slots=True)26class TransportWorkspace:27    """One stream, immutable material tables, small diagnostics and no history tape."""2829    spec: TransportSpec30    max_histories: int31    context: DeviceContext32    material_ids: Any33    absorption: Any34    scattering: Any35    _energies: Any = field(repr=False)36    _parameters: Any = field(repr=False)37    _status: Any = field(repr=False)38    _kernels: Any = field(repr=False)3940    @property41    def scratch_bytes(self) -> int:42        return 4 + 8 * len(self.spec.energy_nodes)4344    def clear_status(self) -> None:45        """Reset at a declared independent run boundary, after discarding failures."""46        with self.context.scope():47            self._status.zero_()4849    def check_status(self) -> None:50        """Synchronise this stream and transfer only the four-byte status word."""51        self.context.wp.synchronize_stream(self.context.stream)52        code = int(self._status.numpy()[0])53        if code & (4 | 8 | 64):54            raise IncompleteHistoryError(55                "event, voxel-crossing or angular-rejection budget exhausted; discard the estimate"56            )57        if code & 32:58            raise TransportError("a history left the supplied coefficient energy support")59        if code & ~(4 | 8 | 16 | 32 | 64):60            raise TransportError("invalid device transport inputs")61        if code & 16:62            raise NumericalError("nonfinite transport arithmetic or an invalid moment")6364    def _launch(self, kernel: Any, dim: int, inputs: list[Any]) -> None:65        self.context.wp.launch(66            kernel,67            dim=dim,68            inputs=inputs,69            device=self.context.device,70            stream=self.context.stream,71            block_dim=self.spec.block_dim,72            record_tape=False,73        )7475    def _model_inputs(self) -> list[Any]:76        return [77            self.material_ids,78            self._energies,79            self.absorption,80            self.scattering,81            self._parameters,82        ]8384    def _reads(self) -> list[tuple[str, Any]]:85        return [86            ("material_ids", self.material_ids),87            ("absorption", self.absorption),88            ("scattering", self.scattering),89            ("energy_nodes", self._energies),90            ("workspace_status", self._status),91        ]929394def prepare_transport(95    spec: TransportSpec,96    *,97    material_ids: Any,98    absorption: Any,99    scattering: Any,100    max_histories: int,101    device: str = "cuda:0",102    stream: Any = None,103    validate: bool = True,104) -> TransportWorkspace:105    """Bind CUDA buffers and validate supplied immutable coefficients explicitly.106107    Material IDs are int32, flattened x fastest. Coefficients are binary64,108    material-major with energy fastest. Their partial sum is total extinction;109    no inconsistent independent total coefficient is accepted. The tables are110    already linear coefficients in inverse mm. For Compton mode they describe111    free-electron scattering in the user's declared model, not bound-electron112    corrections inferred by this package.113    """114    if type(max_histories) is not int or not 0 < max_histories < 2**31:115        raise TransportError("max_histories must be a positive signed 32-bit integer")116    context = prepare_context(device=device, stream=stream)117    require_no_tape()118    wp = context.wp119    context.array(material_ids, "material_ids", dtype=wp.int32, shape=(math.prod(spec.grid.shape),))120    table_size = spec.grid.materials * len(spec.energy_nodes)121    if table_size >= 2**31:122        raise TransportError("material-energy table exceeds signed 32-bit indexing")123    for name, value in (("absorption", absorption), ("scattering", scattering)):124        context.array(value, name, dtype=wp.float64, shape=(table_size,))125    kernels = load_kernels("dpt.transport.kernels")126    with context.scope():127        energies = wp.array(list(spec.energy_nodes), dtype=wp.float64, device=context.device)128        status = wp.zeros(1, dtype=wp.int32, device=context.device)129    parameters = kernels.Parameters()130    parameters.origin = wp.vec3d(*spec.grid.origin_mm)131    parameters.spacing = wp.vec3d(*spec.grid.spacing_mm)132    parameters.shape = wp.vec3i(*spec.grid.shape_xyz)133    parameters.detector_lower = wp.vec2d(*spec.detector.lower_xy_mm)134    parameters.detector_spacing = wp.vec2d(*spec.detector.spacing_xy_mm)135    parameters.detector_shape = wp.vec2i(*spec.detector.shape)136    parameters.detector_z = spec.detector.z_mm137    parameters.source_energy = spec.energy_kev138    parameters.energy_bins = len(spec.energy_nodes)139    parameters.max_events = spec.max_events140    parameters.max_crossings = spec.max_crossings141    parameters.max_angle_trials = spec.max_angle_trials142    parameters.compton = int(spec.scattering_law == "free-electron-compton")143    parameters.energy_score = int(spec.scoring == "energy-kev")144    parameters.continuous_absorption = int(spec.estimator == "continuous-absorption")145    workspace = TransportWorkspace(146        spec,147        max_histories,148        context,149        material_ids,150        absorption,151        scattering,152        energies,153        parameters,154        status,155        kernels,156    )157    if validate:158        workspace._launch(159            kernels.validate_model,160            max(int(material_ids.size), table_size),161            [material_ids, absorption, scattering, spec.grid.materials, status],162        )163        workspace.check_status()164    return workspace165166167def _validate_call(168    workspace: TransportWorkspace,169    batch: HistoryBatch,170    positions: Any,171    directions: Any,172    weights: Any,173    density: Any,174    writes: list[tuple[str, Any, Any]],175    source_amplitude: float,176    stream: Any,177    validate: bool,178) -> None:179    require_no_tape()180    context = workspace.context181    context.assert_stream(stream)182    wp = context.wp183    if batch.count > workspace.max_histories:184        raise TransportError("history batch exceeds workspace capacity")185    _physical(source_amplitude, "source amplitude", nonnegative=True)186    source_arrays = [("positions", positions), ("directions", directions), ("weights", weights)]187    for name, value in source_arrays:188        context.array(189            value,190            name,191            dtype=wp.float64 if name == "weights" else wp.vec3d,192            shape=(batch.count,),193        )194    context.array(density, "density", dtype=wp.float64, shape=(workspace.spec.grid.materials,))195    for name, value, dtype in writes:196        context.array(value, name, dtype=dtype, shape=(batch.count,))197    context.disjoint(198        workspace._reads() + source_arrays + [("density", density)],199        [(name, value) for name, value, _ in writes],200    )201    # Unchecked calls retain prior failure status. A successful later launch202    # cannot launder an incomplete earlier estimate into an accepted one.203    if validate:204        workspace.check_status()205        workspace._launch(206            workspace._kernels.validate_sources,207            max(batch.count, workspace.spec.grid.materials),208            [209                positions,210                directions,211                weights,212                density,213                workspace.spec.detector.z_mm,214                workspace._status,215            ],216        )217        workspace.check_status()218219220# region book:transport-history-execution221def trace_histories(222    positions: Any,223    directions: Any,224    weights: Any,225    density: Any,226    *,227    batch: HistoryBatch,228    workspace: TransportWorkspace,229    out_pixel: Any,230    out_score: Any,231    out_energy: Any,232    out_events: Any,233    out_status: Any,234    source_amplitude: float = 1.0,235    stream: Any = None,236    validate: bool = True,237) -> None:238    """Write one sparse detector score per independent original history.239240    Binary64 positions are in mm; directions must be unit vectors to 1e-12 in241    squared norm. Source positions lie below the detector plane. Nonnegative242    base weights encode the caller's fixed source importance sampling, while243    ``source_amplitude`` scales the expected measurement. ``density`` contains244    strictly positive dimensionless scales multiplying both partial coefficients.245246    The score is per launched history; averaging includes misses and absorption247    as zeros. It is not normalised by detected photons. `out_pixel=-1` means no248    detector hit. Score, final energy (keV), collision count and status are always249    written, including on failure; invalid outputs must be discarded as a batch.250    Continuous absorption samples scattering flights and weights each detector251    hit by absorption along its entire path; its event count counts scatterings.252    The analogue default retains sampled absorption and total-extinction flights.253    """254    wp = workspace.context.wp255    _validate_call(256        workspace,257        batch,258        positions,259        directions,260        weights,261        density,262        [263            ("out_pixel", out_pixel, wp.int32),264            ("out_score", out_score, wp.float64),265            ("out_energy", out_energy, wp.float64),266            ("out_events", out_events, wp.int32),267            ("out_status", out_status, wp.int32),268        ],269        source_amplitude,270        stream,271        validate,272    )273    workspace._launch(274        workspace._kernels.trace_histories,275        batch.count,276        [277            positions,278            directions,279            weights,280            density,281            *workspace._model_inputs(),282            wp.uint64(batch.seed),283            wp.uint64(batch.first_history),284            source_amplitude,285            out_pixel,286            out_score,287            out_energy,288            out_events,289            out_status,290            workspace._status,291        ],292    )293    if validate:294        workspace.check_status()295296297# endregion book:transport-history-execution298