Generated from the full canonical file for this source snapshot. Line numbers match the library source.
Source SHA256: f9e567ef791f1cd5d8017061136b5ce0530543971cf6af5a1acf984201c068e8
1"""Multi-view spectral pose recovery with constrained shared acquisition parameters.23This adapter composes the canonical material projector, spectral response,4spatial response, calibration and objective. The optimiser sees a fixed SE(3)5chart followed by named nuisance blocks. Detector-sized data stay on CUDA;6pinned control buffers move only pose, nuisance values, losses, gradients and7status flags. No observed image is synthesised or normalised by this adapter.8"""910from __future__ import annotations1112import math13from dataclasses import dataclass, field14from typing import Any, Literal1516from dpt._runtime import prepare_context, require_no_tape17from dpt.contracts import ContractError, NumericalError, binary32_scalar, finite_scalar, integer18from dpt.detector import (19BlurSpec,20BlurWorkspace,21CalibrationSpec,22DetectorWorkspace,23blur,24blur_transpose,25calibrate,26calibration_vjp,27prepare_blur,28prepare_detector,29)30from dpt.geometry import DetectorGeometry, RigidTransform31from dpt.material_projection import (32MaterialProjectionWorkspace,33material_projection_vjp,34prepare_material_projection,35project_material_paths,36)37from dpt.materials import Provenance38from dpt.objectives import ObjectiveSpec, ObjectiveWorkspace, evaluate_objective, prepare_objective39from dpt.projection import ProjectionSpec40from dpt.recovery import PoseChart41from dpt.registration import Evaluation, RecoveryPolicy, RecoveryResult, Vector, recover_parameters42from dpt.spectral import (43SpectralSpec,44SpectralWorkspace,45prepare_spectral,46spectral_signal,47spectral_vjp,48)49from dpt.volumes import GridSpec5051# Public Python boundaries validate runtime inputs.52# pyright: reportUnnecessaryIsInstance=false535455def _physical_parameter(value: float, name: str, *, positive: bool = False) -> float:56"""Translate shared scalar storage checks into rejectable trial-domain errors."""57try:58rounded = binary32_scalar(value, f"trial {name}")59except ContractError as error:60raise NumericalError(str(error)) from error61if positive and rounded <= 0:62raise NumericalError(f"trial {name} lies outside its representable physical domain")63return rounded646566def _scaled_product(*values: float) -> float:67"""Preserve exponent range until the final binary64 chart product is formed."""68mantissa, exponent = 1.0, 069for value in values:70if not math.isfinite(value):71raise NumericalError("non-finite physical derivative or chart scale")72part, power = math.frexp(value)73mantissa *= part74exponent += power75try:76return math.ldexp(mantissa, exponent)77except OverflowError as error:78raise NumericalError("chart derivative overflows binary64") from error798081# region book:calibration-parameter-chart82@dataclass(frozen=True, slots=True)83class CalibrationBlock:84"""Named acquisition group sharing one gain/exposure/offset across its views.8586Exactly one multiplicative scale may be fitted. Its positive chart is87reference*exp(scale_step*z); the other scale is fixed to remove the gain /88exposure gauge. An active offset is reference+offset_step*z in output units.89Optional Gaussian priors act on these dimensionless chart coordinates.9091Device parameters are rounded to binary32. The VJP differentiates the smooth92physical chart at those represented values, not the discontinuous rounding93map. Shared physical derivatives stay binary64 until the chart product is94formed, so a large positive scale can preserve a small logarithmic partial.95"""9697name: str98gain: float = 1.099exposure: float = 1.0100offset: float = 0.0101fit_scale: Literal["none", "gain", "exposure"] = "none"102fit_offset: bool = False103scale_step: float = 1.0104offset_step: float = 1.0105scale_prior_precision: float = 0.0106offset_prior_precision: float = 0.0107108def __post_init__(self) -> None:109if not self.name.strip():110raise ContractError("calibration groups need non-empty names")111if self.fit_scale not in ("none", "gain", "exposure") or type(self.fit_offset) is not bool:112raise ContractError("invalid calibration activity specification")113for name in ("gain", "exposure", "scale_step", "offset_step"):114if finite_scalar(getattr(self, name), name, minimum=0) <= 0:115raise ContractError(f"{name} must be strictly positive in the recovery chart")116finite_scalar(self.offset, "offset")117for name in ("scale_prior_precision", "offset_prior_precision"):118finite_scalar(getattr(self, name), name, minimum=0)119if self.fit_scale == "none" and self.scale_prior_precision:120raise ContractError("a scale prior requires an active scale")121if not self.fit_offset and self.offset_prior_precision:122raise ContractError("an offset prior requires an active offset")123for name in ("gain", "exposure", "offset"):124try:125_physical_parameter(getattr(self, name), name, positive=name != "offset")126except NumericalError as error:127raise ContractError(str(error)) from error128129@property130def dimension(self) -> int:131return int(self.fit_scale != "none") + int(self.fit_offset)132133def decode(self, coordinates: Vector) -> tuple[float, float, float]:134if len(coordinates) != self.dimension:135raise ContractError("nuisance coordinates differ from the declared group chart")136if not all(math.isfinite(value) for value in coordinates):137raise NumericalError("trial nuisance coordinates are non-finite")138gain, exposure, offset = self.gain, self.exposure, self.offset139cursor = 0140if self.fit_scale != "none":141try:142scale = math.exp(self.scale_step * coordinates[cursor])143except OverflowError as error:144raise NumericalError("trial logarithmic calibration scale overflowed") from error145if self.fit_scale == "gain":146gain *= scale147else:148exposure *= scale149cursor += 1150if self.fit_offset:151offset += self.offset_step * coordinates[cursor]152return (153_physical_parameter(gain, "gain", positive=True),154_physical_parameter(exposure, "exposure", positive=True),155_physical_parameter(offset, "offset"),156)157158def gradient(self, coordinates: Vector, physical_gradient: Vector) -> Vector:159if len(physical_gradient) != 3:160raise ContractError("a calibration gradient needs gain, exposure and offset partials")161gain, exposure, _ = self.decode(coordinates)162gradient: list[float] = []163if self.fit_scale != "none":164index, value = (0, gain) if self.fit_scale == "gain" else (1, exposure)165gradient.append(_scaled_product(self.scale_step, value, physical_gradient[index]))166if self.fit_offset:167gradient.append(_scaled_product(self.offset_step, physical_gradient[2]))168return tuple(gradient)169170def prior(self, coordinates: Vector) -> Evaluation:171if len(coordinates) != self.dimension:172raise ContractError("prior coordinates differ from calibration chart")173precisions = () if self.fit_scale == "none" else (self.scale_prior_precision,)174precisions += (self.offset_prior_precision,) if self.fit_offset else ()175return Evaluation(176math.fsum(0.5 * p * z * z for p, z in zip(precisions, coordinates, strict=True)),177tuple(p * z for p, z in zip(precisions, coordinates, strict=True)),178)179180181# endregion book:calibration-parameter-chart182183184@dataclass(frozen=True, slots=True)185class SpectralView:186name: str187geometry: DetectorGeometry188spectral: SpectralSpec189coefficients: Any190weights: Any191response: Any192observation: Any193calibration_group: str194observation_provenance: Provenance195objective: ObjectiveSpec = field(default_factory=ObjectiveSpec)196objective_weights: Any = None197objective_weight: float = 1.0198spatial_response: BlurSpec | None = None199200def __post_init__(self) -> None:201if self.spectral.precision != "float32" or self.objective.precision != "float32":202raise ContractError(203"spectral pose recovery requires float32 spectral and objective precision; "204"its calibration and optional spatial-response buffers are binary32"205)206if not isinstance(self.observation_provenance, Provenance):207raise ContractError("each observation requires acquisition provenance")208if not self.name.strip() or not self.calibration_group.strip():209raise ContractError("view and calibration-group names are required")210if finite_scalar(self.objective_weight, "view objective weight", minimum=0) <= 0:211raise ContractError("view objective weight must be strictly positive")212if self.objective.domain == "log_transmission":213raise ContractError(214"spectral recovery compares calibrated signal or counts, not log data"215)216if (217self.spatial_response is not None218and (self.spatial_response.height, self.spatial_response.width) != self.geometry.shape219):220raise ContractError("spatial response and detector shapes differ")221if not self.spectral.active_paths:222raise ContractError("spectral pose recovery needs active material paths")223if (224self.spectral.active_weights225or self.spectral.active_response226or self.spectral.active_coefficients227):228raise ContractError(229"this recovery adapter keeps spectrum, response and coefficients fixed"230)231if not self.objective.weighted and self.objective_weights is not None:232raise ContractError("objective weights require an explicitly weighted objective")233234235@dataclass(frozen=True, slots=True)236class SpectralPoseProblem:237grid: GridSpec238materials: int239fields: Any240fields_provenance: Provenance241views: tuple[SpectralView, ...]242calibration: tuple[CalibrationBlock, ...]243samples_per_ray: int = 256244245def __post_init__(self) -> None:246integer(self.materials, "materials", minimum=1, maximum=32)247integer(self.samples_per_ray, "samples_per_ray", minimum=1)248if not isinstance(self.fields_provenance, Provenance):249raise ContractError("fixed material fields require provenance")250if not isinstance(self.views, tuple) or not self.views:251raise ContractError(252"a spectral recovery problem needs an immutable non-empty view tuple"253)254if not isinstance(self.calibration, tuple) or not self.calibration:255raise ContractError("calibration blocks must be an immutable non-empty tuple")256if len({view.name for view in self.views}) != len(self.views):257raise ContractError("view names must be unique")258groups = {block.name: block for block in self.calibration}259if len(groups) != len(self.calibration):260raise ContractError("calibration group names must be unique")261if set(groups) != {view.calibration_group for view in self.views}:262raise ContractError("every named calibration group must be defined and used")263for group in groups:264units = {265view.spectral.output_unit for view in self.views if view.calibration_group == group266}267if len(units) != 1:268raise ContractError(269"shared calibration groups require identical detector output units"270)271for view in self.views:272if view.objective.domain == "counts" and view.spectral.output_unit != "counts":273raise ContractError("count-domain views must declare spectral output_unit='counts'")274if view.spectral.materials != self.materials:275raise ContractError("all views must use the same fixed material basis")276block = groups[view.calibration_group]277if view.objective.kind == "poisson" and (278block.gain != 1.0279or block.offset != 0.0280or block.fit_scale == "gain"281or block.fit_offset282or view.spatial_response is not None283):284raise ContractError(285"Poisson counts require unit gain, zero offset and no spatial blur"286)287288289@dataclass(slots=True)290class _PreparedView:291view: SpectralView292group: int293material: MaterialProjectionWorkspace294spectral: SpectralWorkspace295detector: DetectorWorkspace296objective: ObjectiveWorkspace297blur: BlurWorkspace | None298calibration_spec: CalibrationSpec299mean: Any300spread: Any301prediction: Any302image_seed: Any303spread_seed: Any304mean_seed: Any305loss: Any306pose_gradient: Any307nuisance_gradients: tuple[Any, Any, Any]308309310@dataclass(frozen=True, slots=True)311class CalibrationEstimate:312name: str313gain: float314exposure: float315offset: float316317318@dataclass(frozen=True, slots=True)319class SpectralRecoveryResult:320pose: RigidTransform321calibration: tuple[CalibrationEstimate, ...]322optimisation: RecoveryResult323324325class SpectralPoseEvaluator:326"""Prepared single-stream multi-view objective with manual first-order products.327328Fields, input spectra, responses, observations and weights are immutable for329this object's lifetime. Preparation checks their current values. Numerical330flags are checked through each workspace after evaluation. Every path,331image and adjoint buffer is allocated once. Concurrent/re-entrant evaluation332and ambient tape recording are rejected before any staging buffer is changed.333"""334335def __init__(336self,337problem: SpectralPoseProblem,338chart: PoseChart,339*,340device: str = "cuda:0",341stream: Any = None,342) -> None:343require_no_tape()344self.problem, self.chart = problem, chart345self.context = prepare_context(device=device, stream=stream)346ctx, wp = self.context, self.context.wp347ctx.array(348problem.fields,349"fields",350dtype=wp.float32,351shape=(problem.materials * problem.grid.voxels,),352)353self._group_slices: tuple[slice, ...] = self._parameter_slices()354self.dimension = 6 + sum(block.dimension for block in problem.calibration)355self._running = False356self.last_view_losses: tuple[float, ...] | None = None357views = len(problem.views)358with ctx.scope():359self._pose_device = wp.empty(12, dtype=wp.float64, device=ctx.device)360self._parameters_device = wp.empty(3613 * len(problem.calibration), dtype=wp.float32, device=ctx.device362)363self._control_device = wp.empty(7 * views, dtype=wp.float64, device=ctx.device)364self._nuisance_device = wp.zeros(3 * views, dtype=wp.float64, device=ctx.device)365self._parameter_views = tuple(366tuple(self._parameters_device[3 * i + j : 3 * i + j + 1] for j in range(3))367for i in range(len(problem.calibration))368)369self._pose_host = wp.empty(12, dtype=wp.float64, device="cpu", pinned=True)370self._parameters_host = wp.empty(3713 * len(problem.calibration), dtype=wp.float32, device="cpu", pinned=True372)373self._control_host = wp.empty(7 * views, dtype=wp.float64, device="cpu", pinned=True)374self._nuisance_host = wp.empty(3 * views, dtype=wp.float64, device="cpu", pinned=True)375self._pose_values = self._pose_host.numpy()376self._parameter_values = self._parameters_host.numpy()377self._control_values = self._control_host.numpy()378self._nuisance_values = self._nuisance_host.numpy()379self._prepared = tuple(380self._prepare_view(index, view) for index, view in enumerate(problem.views)381)382self._validate_fixed_inputs()383384def _parameter_slices(self) -> tuple[slice, ...]:385cursor = 6386slices: list[slice] = []387for block in self.problem.calibration:388slices.append(slice(cursor, cursor + block.dimension))389cursor += block.dimension390return tuple(slices)391392def _prepare_view(self, index: int, view: SpectralView) -> _PreparedView:393ctx, wp = self.context, self.context.wp394group = next(395i396for i, block in enumerate(self.problem.calibration)397if block.name == view.calibration_group398)399block = self.problem.calibration[group]400pixels, spec = view.geometry.pixels, view.spectral401ctx.array(402view.coefficients,403"coefficients",404dtype=wp.float32,405shape=(spec.materials * spec.energies,),406)407ctx.array(408view.weights,409"weights",410dtype=wp.float32,411shape=(spec.energies * (1 if spec.shared_weights else pixels),),412)413ctx.array(414view.response,415"response",416dtype=wp.float32,417shape=(spec.energies * (1 if spec.shared_response else pixels),),418)419ctx.array(view.observation, "observation", dtype=wp.float32, shape=(pixels,))420if view.objective.weighted:421ctx.array(422view.objective_weights, "objective_weights", dtype=wp.float32, shape=(pixels,)423)424with ctx.scope():425paths = wp.empty(spec.materials * pixels, dtype=wp.float32, device=ctx.device)426adj_paths = wp.empty_like(paths)427mean = wp.empty(pixels, dtype=wp.float32, device=ctx.device)428mean_seed = wp.empty_like(mean)429spread = wp.empty_like(mean) if view.spatial_response is not None else mean430spread_seed = wp.empty_like(mean) if view.spatial_response is not None else mean_seed431prediction, image_seed = wp.empty_like(mean), wp.empty_like(mean)432material = prepare_material_projection(433self.problem.grid,434view.geometry,435ProjectionSpec(self.problem.samples_per_ray),436materials=self.problem.materials,437fields=self.problem.fields,438out_paths=paths,439adj_paths=adj_paths,440device=ctx.device,441stream=ctx.stream,442)443spectral = prepare_spectral(spec, max_pixels=pixels, device=ctx.device, stream=ctx.stream)444detector = prepare_detector(max_pixels=pixels, device=ctx.device, stream=ctx.stream)445objective = prepare_objective(446view.objective, max_pixels=pixels, device=ctx.device, stream=ctx.stream447)448spatial = (449prepare_blur(view.spatial_response, device=ctx.device, stream=ctx.stream)450if view.spatial_response is not None451else None452)453calibration_spec = CalibrationSpec(454spec.output_unit,455active_gain=block.fit_scale == "gain",456active_exposure=block.fit_scale == "exposure",457active_offset=block.fit_offset,458)459return _PreparedView(460view,461group,462material,463spectral,464detector,465objective,466spatial,467calibration_spec,468mean,469spread,470prediction,471image_seed,472spread_seed,473mean_seed,474self._control_device[7 * index : 7 * index + 1],475self._control_device[7 * index + 1 : 7 * index + 7],476tuple(self._nuisance_device[3 * index + j : 3 * index + j + 1] for j in range(3)),477)478479def _workspaces(480self,481) -> tuple[482MaterialProjectionWorkspace483| SpectralWorkspace484| DetectorWorkspace485| ObjectiveWorkspace486| BlurWorkspace,487...,488]:489"""Use public workspace diagnostics; status encodings stay with operators."""490return tuple(491workspace492for item in self._prepared493for workspace in (494item.material,495item.spectral,496item.detector,497item.objective,498item.blur,499)500if workspace is not None501)502503def _clear_statuses(self) -> None:504for workspace in self._workspaces():505workspace.clear_status()506507def _check_statuses(self) -> None:508for workspace in self._workspaces():509workspace.check_status()510511def _validate_fixed_inputs(self) -> None:512for item in self._prepared:513item.material.validate_inputs(stream=self.context.stream)514item.spectral.validate_inputs(515item.view.coefficients,516item.view.weights,517item.view.response,518pixels=item.view.geometry.pixels,519probability_response=item.view.objective.kind == "poisson",520stream=self.context.stream,521)522item.objective.validate_observation(523item.view.observation,524weights=item.view.objective_weights,525stream=self.context.stream,526)527528# region book:spectral-recovery-composition529def __call__(self, parameters: Vector) -> Evaluation:530require_no_tape()531if self._running:532raise ContractError("a spectral evaluator cannot be used concurrently or recursively")533if len(parameters) != self.dimension:534raise ContractError("parameter count differs from pose and calibration charts")535pose = self.chart.pose(parameters[:6])536decoded = tuple(537block.decode(parameters[section])538for block, section in zip(self.problem.calibration, self._group_slices, strict=True)539)540ctx, wp = self.context, self.context.wp541self._running = True542self.last_view_losses = None543try:544with ctx.scope():545try:546self._clear_statuses()547self._pose_values[:] = pose.packed()548self._parameter_values[:] = tuple(value for block in decoded for value in block)549wp.copy(self._pose_device, self._pose_host, stream=ctx.stream)550wp.copy(self._parameters_device, self._parameters_host, stream=ctx.stream)551for item in self._prepared:552view = item.view553gain, exposure, offset = self._parameter_views[item.group]554project_material_paths(555self._pose_device,556workspace=item.material,557stream=ctx.stream,558validate=False,559)560spectral_signal(561item.material.paths,562view.coefficients,563view.weights,564view.response,565out_mean=item.mean,566workspace=item.spectral,567stream=ctx.stream,568validate=False,569)570if item.blur is not None:571blur(572item.mean,573out_signal=item.spread,574workspace=item.blur,575stream=ctx.stream,576validate=False,577)578calibrate(579item.spread,580gain,581exposure,582offset,583out_signal=item.prediction,584spec=item.calibration_spec,585workspace=item.detector,586stream=ctx.stream,587validate=False,588)589evaluate_objective(590item.prediction,591view.observation,592out_loss=item.loss,593out_seed=item.image_seed,594weights=view.objective_weights,595workspace=item.objective,596stream=ctx.stream,597validate=False,598)599calibration_vjp(600item.spread,601gain,602exposure,603offset,604seed=item.image_seed,605out_grad_mean=item.spread_seed,606out_grad_gain=item.nuisance_gradients[0]607if item.calibration_spec.active_gain608else None,609out_grad_exposure=item.nuisance_gradients[1]610if item.calibration_spec.active_exposure611else None,612out_grad_offset=item.nuisance_gradients[2]613if item.calibration_spec.active_offset614else None,615spec=item.calibration_spec,616workspace=item.detector,617stream=ctx.stream,618validate=False,619)620if item.blur is not None:621blur_transpose(622item.spread_seed,623out_grad_signal=item.mean_seed,624workspace=item.blur,625stream=ctx.stream,626validate=False,627)628spectral_vjp(629item.material.paths,630view.coefficients,631view.weights,632view.response,633seed=item.mean_seed,634out_grad_paths=item.material.adj_paths,635workspace=item.spectral,636stream=ctx.stream,637validate=False,638)639material_projection_vjp(640self._pose_device,641workspace=item.material,642out_pose=item.pose_gradient,643stream=ctx.stream,644validate=False,645)646wp.copy(self._control_host, self._control_device, stream=ctx.stream)647wp.copy(self._nuisance_host, self._nuisance_device, stream=ctx.stream)648finally:649# Rejected evaluations must finish using pinned staging before650# the optimiser can write the next trial into the same buffers.651wp.synchronize_stream(ctx.stream)652self._check_statuses()653return self._assemble(parameters)654finally:655self._running = False656657# endregion book:spectral-recovery-composition658659def _assemble(self, parameters: Vector) -> Evaluation:660losses = tuple(float(self._control_values[7 * i]) for i in range(len(self._prepared)))661pose_gradient = tuple(662math.fsum(663item.view.objective_weight * float(self._control_values[7 * i + 1 + j])664for i, item in enumerate(self._prepared)665)666for j in range(6)667)668gradient = list(self.chart.gradient(parameters[:6], pose_gradient))669prior_losses: list[float] = []670for group, (block, section) in enumerate(671zip(self.problem.calibration, self._group_slices, strict=True)672):673physical = tuple(674math.fsum(675item.view.objective_weight * float(self._nuisance_values[3 * i + j])676for i, item in enumerate(self._prepared)677if item.group == group678)679for j in range(3)680)681prior = block.prior(parameters[section])682prior_losses.append(prior.loss)683derivative = block.gradient(parameters[section], physical)684gradient.extend(a + b for a, b in zip(derivative, prior.gradient, strict=True))685loss = math.fsum(686[687*(688item.view.objective_weight * value689for item, value in zip(self._prepared, losses, strict=True)690),691*prior_losses,692]693)694if not math.isfinite(loss) or not all(math.isfinite(value) for value in gradient):695raise NumericalError("multi-view loss or chart gradient is non-finite")696self.last_view_losses = losses697return Evaluation(loss, tuple(gradient))698699def calibration_estimates(self, parameters: Vector) -> tuple[CalibrationEstimate, ...]:700if len(parameters) != self.dimension:701raise ContractError("parameter count differs from evaluator chart")702return tuple(703CalibrationEstimate(block.name, *block.decode(parameters[section]))704for block, section in zip(self.problem.calibration, self._group_slices, strict=True)705)706707@property708def predicted_images(self) -> tuple[Any, ...]:709"""Borrow the current device predictions; they may belong to a rejected trial.710711Re-evaluate accepted parameters before explicit export. No host copy is712made by this accessor, and callers must not mutate these borrowed buffers.713"""714return tuple(item.prediction for item in self._prepared)715716717def recover_spectral_pose(718evaluator: SpectralPoseEvaluator,719*,720initial: Vector | None = None,721policy: RecoveryPolicy | None = None,722) -> SpectralRecoveryResult:723"""Use the same safeguarded optimiser and immutable chart as primary recovery."""724start = (0.0,) * evaluator.dimension if initial is None else initial725optimisation = recover_parameters(evaluator, start, policy=policy)726return SpectralRecoveryResult(727evaluator.chart.pose(optimisation.parameters[:6]),728evaluator.calibration_estimates(optimisation.parameters),729optimisation,730)731