python/dpt/recovery.py

1"""Composition of the canonical projector, transmission, objective and optimiser.23The fixed volume, observation and reusable image buffers remain on CUDA. Each4evaluation explicitly uploads a twelve-value rigid transform and downloads one5loss and six local pose derivatives, plus explicit numerical-status checkpoints.6This is a deterministic primary-model7adapter; spectral/nuisance adapters use the same scaled optimiser contract.8"""910from __future__ import annotations1112import math13from dataclasses import dataclass, field14from typing import Any, Literal1516from dpt._runtime import require_no_tape17from dpt.contracts import ContractError, NumericalError, finite_scalar18from dpt.geometry import DetectorGeometry, RigidTransform, chart_gradient, compose_pose19from dpt.objectives import (20    ObjectiveSpec,21    evaluate_objective,22    evaluate_primary_objective,23    prepare_objective,24)25from dpt.projection import ProjectionSpec, prepare_projection, project_optical_depth, projection_vjp26from dpt.registration import Evaluation, RecoveryPolicy, RecoveryResult, Vector, recover_parameters27from dpt.transmission import TransmissionSpec, prepare_transmission, transmission_vjp, transmit28from dpt.volumes import GridSpec293031@dataclass(frozen=True, slots=True)32class PoseChart:33    """Dimensionless optimiser coordinates mapped into one fixed SE(3) chart."""3435    anchor: RigidTransform = field(default_factory=RigidTransform)36    scales: tuple[float, float, float, float, float, float] = (1.0, 1.0, 1.0, 0.01, 0.01, 0.01)37    rotation_radius_radians: float = math.pi3839    def __post_init__(self) -> None:40        if not isinstance(self.anchor, RigidTransform):  # pyright: ignore[reportUnnecessaryIsInstance]41            raise ContractError("pose anchor must be a validated RigidTransform")42        if len(self.scales) != 6:43            raise ContractError("pose scales must contain three mm and three radian values")44        for value in self.scales:45            if finite_scalar(value, "pose scale", minimum=0.0) == 0:46                raise ContractError("pose scales must be strictly positive")47        object.__setattr__(self, "scales", tuple(float(value) for value in self.scales))48        radius = finite_scalar(self.rotation_radius_radians, "rotation radius", minimum=0.0)49        if not 0 < radius <= math.pi:50            raise ContractError("fixed pose chart rotation radius must lie in (0, pi]")5152    def increment(self, parameters: Vector) -> Vector:53        if len(parameters) != 6:54            raise ContractError("a rigid-pose evaluator needs six chart parameters")55        values = tuple(a * b for a, b in zip(parameters, self.scales, strict=True))56        if (57            not all(map(math.isfinite, values))58            or math.hypot(*values[3:]) >= self.rotation_radius_radians59        ):60            raise NumericalError("trial pose lies outside the declared fixed rotation chart")61        return values6263    def pose(self, parameters: Vector) -> RigidTransform:64        increment = self.increment(parameters)65        try:66            return compose_pose(self.anchor, increment)67        except (ValueError, OverflowError) as error:68            raise NumericalError("trial transform exceeds the finite rigid-pose range") from error6970    def gradient(self, parameters: Vector, local_gradient: Vector) -> Vector:71        gradient = chart_gradient(self.increment(parameters), local_gradient)72        return tuple(a * b for a, b in zip(gradient, self.scales, strict=True))737475@dataclass(frozen=True, slots=True)76class PrimaryPoseProblem:77    grid: GridSpec78    geometry: DetectorGeometry79    attenuation: Any80    observation: Any81    objective: ObjectiveSpec = field(default_factory=lambda: ObjectiveSpec(domain="counts"))82    open_beam: float = 1.083    weights: Any = None84    samples_per_ray: int = 25685    precision: Literal["float32", "float64"] = "float64"86    integration: Literal["midpoint", "cell_gauss"] = "midpoint"878889@dataclass(frozen=True, slots=True)90class PoseRecoveryResult:91    pose: RigidTransform92    optimisation: RecoveryResult939495class PrimaryPoseEvaluator:96    """Prepared single-stream composition; never share concurrently.9798    Attenuation, observations, weights and chart are immutable for a solve.99    The first evaluation checks their device values; later calls rely on that100    immutability and check numerical output flags at explicit completion points.101    FP64 depth, signal and depth cotangents are retained by default so a pose102    line search can resolve the loss changes indicated by its gradient. Explicit103    precision="float32" reproduces the earlier composed storage contract.104    Scratch contents after return may correspond to a rejected trial; the result105    returned by recover_pose identifies the accepted pose and objective. Call106    evaluate explicitly at that pose if a final predicted image is needed.107    """108109    def __init__(110        self,111        problem: PrimaryPoseProblem,112        chart: PoseChart,113        *,114        device: str = "cuda:0",115        stream: Any = None,116    ) -> None:117        self.problem, self.chart = problem, chart118        if problem.objective.domain not in ("counts", "log_transmission"):119            raise ContractError("primary pose recovery requires counts or log-transmission data")120        self.projection = prepare_projection(121            problem.grid,122            problem.geometry,123            ProjectionSpec(124                problem.samples_per_ray,125                precision=problem.precision,126                integration=problem.integration,127            ),128            device=device,129            stream=stream,130        )131        self.context = self.projection.context132        ctx, wp, count = self.context, self.context.wp, problem.geometry.pixels133        ctx.array(134            problem.attenuation, "attenuation", dtype=wp.float32, shape=(problem.grid.voxels,)135        )136        ctx.array(problem.observation, "observation", dtype=wp.float32, shape=(count,))137        if problem.objective.weighted:138            ctx.array(problem.weights, "weights", dtype=wp.float32, shape=(count,))139        elif problem.weights is not None:140            raise ContractError("weights require an explicitly weighted objective")141        self.transmission = (142            prepare_transmission(143                TransmissionSpec(beam="scalar" if problem.objective.domain == "counts" else "none"),144                max_pixels=count,145                device=device,146                stream=ctx.stream,147            )148            if problem.precision == "float32"149            else None150        )151        self.objective = prepare_objective(152            problem.objective, max_pixels=count, device=device, stream=ctx.stream153        )154        with ctx.scope():155            self.pose_device = wp.empty(12, dtype=wp.float64, device=ctx.device)156            self.optical_depth = wp.empty(count, dtype=self.projection.dtype, device=ctx.device)157            self.prediction = wp.empty(count, dtype=self.projection.dtype, device=ctx.device)158            self.image_seed = wp.empty(159                count if problem.precision == "float32" else 0, dtype=wp.float32, device=ctx.device160            )161            self.depth_seed = wp.empty(count, dtype=self.projection.dtype, device=ctx.device)162            self.pose_seed = wp.empty(6, dtype=wp.float64, device=ctx.device)163            self.loss = wp.empty(1, dtype=wp.float64, device=ctx.device)164        # Pinned host staging has a persistent owner and a fixed-size NumPy view.165        # These are explicit control-plane transfers, not per-pixel Python work.166        self._pose_host = wp.empty(12, dtype=wp.float64, device="cpu", pinned=True)167        self._loss_host = wp.empty(1, dtype=wp.float64, device="cpu", pinned=True)168        self._gradient_host = wp.empty(6, dtype=wp.float64, device="cpu", pinned=True)169        self._pose_view = self._pose_host.numpy()170        self._loss_view = self._loss_host.numpy()171        self._gradient_view = self._gradient_host.numpy()172        self._inputs_checked = False173174    # region book:recovery-device-composition175    def __call__(self, parameters: Vector) -> Evaluation:176        require_no_tape()177        try:178            return self._evaluate(parameters)179        finally:180            # Even a failed launch can follow an enqueued asynchronous H2D copy.181            # Release its read of pinned staging before a retry may overwrite it.182            self.context.wp.synchronize_stream(self.context.stream)183184    def _evaluate(self, parameters: Vector) -> Evaluation:185        pose = self.chart.pose(parameters)186        ctx, wp, problem = self.context, self.context.wp, self.problem187        validate = not self._inputs_checked188        self.projection.clear_status()189        if self.transmission is not None:190            self.transmission.clear_status()191        self.objective.clear_status()192        with ctx.scope():193            self._pose_view[:] = pose.packed()194            wp.copy(self.pose_device, self._pose_host, stream=ctx.stream)195            project_optical_depth(196                problem.attenuation,197                self.pose_device,198                out_L=self.optical_depth,199                workspace=self.projection,200                stream=ctx.stream,201                validate=validate,202            )203            if problem.precision == "float64":204                evaluate_primary_objective(205                    self.optical_depth,206                    problem.observation,207                    open_beam=problem.open_beam,208                    weights=problem.weights,209                    out_prediction=self.prediction,210                    out_loss=self.loss,211                    out_depth_seed=self.depth_seed,212                    workspace=self.objective,213                    stream=ctx.stream,214                    validate=validate,215                )216            else:217                assert self.transmission is not None218                counts = problem.objective.domain == "counts"219                transmit(220                    self.optical_depth,221                    problem.open_beam if counts else None,222                    out_counts=self.prediction if counts else None,223                    out_log_T=None if counts else self.prediction,224                    workspace=self.transmission,225                    stream=ctx.stream,226                    validate=validate,227                )228                evaluate_objective(229                    self.prediction,230                    problem.observation,231                    weights=problem.weights,232                    out_loss=self.loss,233                    out_seed=self.image_seed,234                    workspace=self.objective,235                    stream=ctx.stream,236                    validate=validate,237                )238                transmission_vjp(239                    self.optical_depth,240                    problem.open_beam if counts else None,241                    seed_counts=self.image_seed if counts else None,242                    seed_log_T=None if counts else self.image_seed,243                    out_grad_L=self.depth_seed,244                    workspace=self.transmission,245                    stream=ctx.stream,246                    validate=validate,247                )248            projection_vjp(249                problem.attenuation,250                self.pose_device,251                adj_L=self.depth_seed,252                out_pose=self.pose_seed,253                workspace=self.projection,254                stream=ctx.stream,255                validate=validate,256            )257            self.projection.check_status()258            if self.transmission is not None:259                self.transmission.check_status()260            self.objective.check_status()261            self._inputs_checked = True262            wp.copy(self._loss_host, self.loss, stream=ctx.stream)263            wp.copy(self._gradient_host, self.pose_seed, stream=ctx.stream)264            wp.synchronize_stream(ctx.stream)265        gradient = tuple(float(value) for value in self._gradient_view)266        return Evaluation(float(self._loss_view[0]), self.chart.gradient(parameters, gradient))267268    # endregion book:recovery-device-composition269270271def recover_pose(272    evaluator: PrimaryPoseEvaluator,273    *,274    initial: Vector = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0),275    policy: RecoveryPolicy | None = None,276) -> PoseRecoveryResult:277    """Run the canonical safeguarded driver without rebasing its curvature chart."""278    result = recover_parameters(evaluator, initial, policy=policy)279    return PoseRecoveryResult(evaluator.chart.pose(result.parameters), result)280