python/dpt/examples/reconstruction.py

1"""Fixed-geometry primary-count reconstruction; supplied inputs only.23Uses the canonical Warp projector and first-order volume VJP. Torch supplies4CUDA vector arithmetic with shared storage on the same stream. The examples5README links the recorded studies and their separate numerical acceptance.6"""78from __future__ import annotations910import importlib11import math12from collections.abc import Callable13from dataclasses import dataclass14from typing import Any1516from dpt.contracts import ContractError, NumericalError, finite_scalar, integer17from dpt.examples._common import example_parser, load_case, write_array18from dpt.geometry import DetectorGeometry, RigidTransform19from dpt.objectives import ObjectiveSpec, evaluate_objective, prepare_objective20from dpt.projection import ProjectionSpec, prepare_projection, project_optical_depth, projection_vjp21from dpt.transmission import TransmissionSpec, prepare_transmission, transmission_vjp, transmit22from dpt.volumes import GridSpec232425@dataclass(frozen=True, slots=True)26class SolverSettings:27    iterations: int28    initial_step_mm_inverse_squared: float29    backtracking_factor: float30    armijo: float31    maximum_backtracks: int32    gradient_mapping_tolerance_mm: float33    regularisation_mm: float34    mapping_step_mm_inverse_squared: float | None = None3536    @classmethod37    def read(cls, settings: dict[str, Any]) -> SolverSettings:38        result = cls(39            iterations=integer(settings["iterations"], "iterations", minimum=1),40            initial_step_mm_inverse_squared=finite_scalar(41                settings["initial_step_mm_inverse_squared"], "initial step", minimum=042            ),43            backtracking_factor=finite_scalar(settings["backtracking_factor"], "backtracking"),44            armijo=finite_scalar(settings["armijo"], "Armijo coefficient"),45            maximum_backtracks=integer(46                settings["maximum_backtracks"], "maximum_backtracks", minimum=147            ),48            gradient_mapping_tolerance_mm=finite_scalar(49                settings["gradient_mapping_tolerance_mm"], "gradient tolerance", minimum=050            ),51            regularisation_mm=finite_scalar(52                settings["regularisation_mm"], "regularisation coefficient", minimum=053            ),54            mapping_step_mm_inverse_squared=(55                finite_scalar(56                    settings["mapping_step_mm_inverse_squared"], "mapping step", minimum=057                )58                if settings.get("mapping_step_mm_inverse_squared") is not None59                else None60            ),61        )62        if result.initial_step_mm_inverse_squared <= 0:63            raise ContractError("Set initial_step_mm_inverse_squared to a positive value.")64        if (65            result.mapping_step_mm_inverse_squared is not None66            and result.mapping_step_mm_inverse_squared <= 067        ):68            raise ContractError("Set mapping_step_mm_inverse_squared to a positive value.")69        if not 0 < result.backtracking_factor < 1 or not 0 < result.armijo < 1:70            raise ContractError("Set backtracking_factor and armijo strictly between zero and one.")71        return result727374class Reconstruction:75    """Serial full-data evaluations with persistent CUDA buffers and checked VJPs."""7677    def __init__(self, case: Any, *, device: str, torch: Any, wp: Any) -> None:78        self.torch, self.wp, self.device = torch, wp, device79        settings = case.config["reconstruction"]80        self.policy = SolverSettings.read(settings["solver"])81        self.grid = GridSpec(**{key: tuple(value) for key, value in settings["grid"].items()})82        pose = RigidTransform(83            **{key: tuple(value) for key, value in settings["fixed_pose"].items()}84        )85        samples = integer(settings["samples_per_ray"], "samples_per_ray", minimum=1)86        initial = case.array(settings["initial_volume"], shape=self.grid.shape, units="mm^-1")87        if (initial < 0).any():88            raise ContractError(89                "initial_volume contains negative attenuation; supply values in mm^-1."90            )91        if not settings["views"]:92            raise ContractError("Supply at least one calibrated view in reconstruction.views.")93        self.torch_stream = torch.cuda.current_stream(device=device)94        wp.init()95        self.stream = wp.stream_from_torch(self.torch_stream)96        self.mu = torch.as_tensor(initial, device=device).flatten().clone()97        self.trial = torch.empty_like(self.mu)98        self.gradient = torch.empty_like(self.mu)99        self.delta = torch.empty_like(self.mu)100        self.edge32 = torch.empty_like(self.mu)101        self.volume64 = torch.empty(self.grid.voxels, dtype=torch.float64, device=device)102        self.edge64 = torch.empty_like(self.volume64)103        self.scalar = torch.empty((), dtype=torch.float64, device=device)104        self.pose = wp.array(pose.packed(), dtype=wp.float64, device=device)105        self.mu_wp = wp.from_torch(self.mu)106        self.trial_wp = wp.from_torch(self.trial)107        self.gradient_wp = wp.from_torch(self.gradient)108        self.views: list[dict[str, Any]] = []109        for entry in settings["views"]:110            geometry = DetectorGeometry(111                **{key: tuple(value) for key, value in entry["geometry"].items()}112            )113            observed = case.array(entry["counts"], shape=geometry.shape, units="counts")114            beam = case.array(entry["open_beam"], shape=geometry.shape, units="counts")115            if (observed < 0).any() or (observed % 1 != 0).any() or (observed > 2**24).any():116                raise ContractError(117                    "counts must be nonnegative integers no greater than 2^24; "118                    "supply unprocessed counts exactly representable in float32."119                )120            if (beam <= 0).any():121                raise ContractError("open_beam must contain positive calibrated mean counts.")122            projection = prepare_projection(123                self.grid,124                geometry,125                ProjectionSpec(samples, active_pose=False, active_volume=True),126                device=device,127                stream=self.stream,128            )129            view: dict[str, Any] = {130                "projection": projection,131                "transmission": prepare_transmission(132                    TransmissionSpec(beam="per-pixel"),133                    max_pixels=geometry.pixels,134                    device=device,135                    stream=self.stream,136                ),137                "objective": prepare_objective(138                    ObjectiveSpec(kind="poisson", domain="counts", reduction="sum"),139                    max_pixels=geometry.pixels,140                    device=device,141                    stream=self.stream,142                ),143                "observed": wp.array(observed.reshape(-1), dtype=wp.float32, device=device),144                "beam": wp.array(beam.reshape(-1), dtype=wp.float32, device=device),145                "shape": geometry.shape,146            }147            for name in ("depth", "prediction", "image_seed", "depth_seed"):148                view[name] = wp.empty(geometry.pixels, dtype=wp.float32, device=device)149            view["loss"] = wp.empty(1, dtype=wp.float64, device=device)150            self.views.append(view)151        # All uploads above are complete before either framework consumes them.152        # Subsequent Torch/Warp work shares exactly one stream.153        wp.synchronize_device(device)154        self.edges: list[tuple[Any, Any, Any, Any, Any, float]] = []155        shaped64 = self.volume64.view(self.grid.shape)156        shaped_gradient = self.gradient.view(self.grid.shape)157        cell_volume = math.prod(self.grid.spacing_mm)158        for axis, spacing in enumerate(reversed(self.grid.spacing_mm)):159            if self.grid.shape[axis] == 1:160                continue161            low = [slice(None)] * 3162            high = [slice(None)] * 3163            low[axis], high[axis] = slice(None, -1), slice(1, None)164            lower, upper = tuple(low), tuple(high)165            shape = shaped64[lower].shape166            count = shaped64[lower].numel()167            coefficient = self.policy.regularisation_mm * cell_volume / spacing**2168            if not math.isfinite(coefficient):169                raise ContractError(170                    "Grid spacing and regularisation overflow; rescale their units."171                )172            self.edges.append(173                (174                    shaped64[lower],175                    shaped64[upper],176                    shaped_gradient[lower],177                    shaped_gradient[upper],178                    (self.edge64[:count].view(shape), self.edge32[:count].view(shape)),179                    coefficient,180                )181            )182183    # region book:example-reconstruction-volume-gradient184    def evaluate(self, volume: Any, volume_wp: Any, *, gradient: bool) -> float:185        """Evaluate every supplied view before an update; add the prior once."""186        if gradient:187            self.gradient.zero_()188        loss = 0.0189        for view in self.views:190            project_optical_depth(191                volume_wp,192                self.pose,193                workspace=view["projection"],194                out_L=view["depth"],195                stream=self.stream,196            )197            transmit(198                view["depth"],199                view["beam"],200                workspace=view["transmission"],201                out_counts=view["prediction"],202                stream=self.stream,203            )204            evaluate_objective(205                view["prediction"],206                view["observed"],207                workspace=view["objective"],208                out_loss=view["loss"],209                out_seed=view["image_seed"] if gradient else None,210                stream=self.stream,211            )212            if gradient:213                transmission_vjp(214                    view["depth"],215                    view["beam"],216                    seed_counts=view["image_seed"],217                    workspace=view["transmission"],218                    out_grad_L=view["depth_seed"],219                    stream=self.stream,220                )221                projection_vjp(222                    volume_wp,223                    self.pose,224                    adj_L=view["depth_seed"],225                    workspace=view["projection"],226                    out_mu=self.gradient_wp,227                    accumulate=True,228                    stream=self.stream,229                )230            loss += float(view["loss"].numpy()[0])231        loss += self.regularise(volume, gradient=gradient)232        if not math.isfinite(loss):233            raise NumericalError("The objective is non-finite; inspect counts and model range.")234        if gradient and not bool(self.torch.isfinite(self.gradient).all().item()):235            raise NumericalError("The volume gradient overflowed; discard this evaluation.")236        return loss237238    # endregion book:example-reconstruction-volume-gradient239240    # region book:example-reconstruction-regularisation241    def regularise(self, volume: Any, *, gradient: bool) -> float:242        """Quadratic physical-space differences; omit edges beyond the grid."""243        self.volume64.copy_(volume)244        penalty = 0.0245        for low, high, grad_low, grad_high, scratch, coefficient in self.edges:246            difference, derivative = scratch247            self.torch.sub(high, low, out=difference)248            if gradient:249                # Scale in FP64 before the FP32 accumulation boundary.250                difference.mul_(coefficient)251                derivative.copy_(difference)252                grad_low.sub_(derivative)253                grad_high.add_(derivative)254                # Restore unscaled differences for the objective.255                self.torch.sub(high, low, out=difference)256            difference.square_()257            self.torch.sum(difference, dim=(0, 1, 2), out=self.scalar)258            penalty += 0.5 * coefficient * float(self.scalar.item())259        return penalty260261    # endregion book:example-reconstruction-regularisation262263    def inner_product(self, left: Any, right: Any) -> float:264        self.volume64.copy_(left)265        self.edge64.copy_(right)266        self.volume64.mul_(self.edge64)267        self.torch.sum(self.volume64, dim=0, out=self.scalar)268        return float(self.scalar.item())269270    def displacement(self, step: float) -> tuple[float, float]:271        # min(g, mu / step) evaluates the projected-gradient mapping272        # without cancellation between almost identical accepted/trial voxels.273        self.volume64.copy_(self.gradient)274        self.edge64.copy_(self.mu).div_(step)275        self.torch.minimum(self.volume64, self.edge64, out=self.volume64)276        self.torch.abs(self.volume64, out=self.edge64)277        mapping = float(self.edge64.max().item())278        self.volume64.mul_(-step).add_(self.mu).clamp_(min=0)279        self.trial.copy_(self.volume64)280        self.torch.sub(self.trial, self.mu, out=self.delta)281        slope = self.inner_product(self.gradient, self.delta)282        return mapping, slope283284    # region book:example-reconstruction-projected-armijo285    def solve(286        self, callback: Callable[[int, dict[str, Any]], None] | None = None287    ) -> dict[str, Any]:288        """Solve with optional observation of accepted updates only.289290        The callback receives the stable one-based accepted iteration and a291        detached history dictionary after ``mu`` has been updated. It may read292        or export the accepted fields but must not mutate solver state. Callback293        exceptions propagate with the last accepted volume intact.294295        An optional fixed mapping step separates the stationarity diagnostic296        from the initial Armijo trial. Omitting it preserves the original297        diagnostic at ``initial_step_mm_inverse_squared``.298        """299        policy = self.policy300        mapping_step = (301            policy.initial_step_mm_inverse_squared302            if policy.mapping_step_mm_inverse_squared is None303            else policy.mapping_step_mm_inverse_squared304        )305        history: list[dict[str, Any]] = []306        reason = "iteration_budget"307        for iteration in range(policy.iterations):308            loss = self.evaluate(self.mu, self.mu_wp, gradient=True)309            step = policy.initial_step_mm_inverse_squared310            mapping, _ = self.displacement(mapping_step)311            if not math.isfinite(mapping):312                raise NumericalError("The trial update overflowed; reduce the initial step.")313            if mapping <= policy.gradient_mapping_tolerance_mm:314                reason = "projected_gradient_tolerance"315                break316            accepted = False317            for attempt in range(policy.maximum_backtracks):318                if step == 0:319                    break320                _, slope = self.displacement(step)321                if not math.isfinite(slope) or slope >= 0:322                    step *= policy.backtracking_factor323                    continue324                # Trial value evaluation leaves the accepted gradient unchanged.325                # Domain/range errors abort with the last accepted volume intact.326                trial_loss = self.evaluate(self.trial, self.trial_wp, gradient=False)327                if trial_loss < loss and trial_loss <= loss + policy.armijo * slope:328                    self.mu.copy_(self.trial)329                    history.append(330                        {331                            "iteration": iteration + 1,332                            "loss_before": loss,333                            "loss_after": trial_loss,334                            "step_mm_inverse_squared": step,335                            "backtracks": attempt,336                            "mapping_mm_before": mapping,337                        }338                    )339                    accepted = True340                    if callback is not None:341                        callback(iteration + 1, dict(history[-1]))342                    break343                step *= policy.backtracking_factor344            if not accepted:345                reason = "line_search_failed"346                break347        # Trial buffers may describe a rejected volume; refresh accepted outputs.348        final_loss = self.evaluate(self.mu, self.mu_wp, gradient=True)349        final_mapping, _ = self.displacement(mapping_step)350        if final_mapping <= policy.gradient_mapping_tolerance_mm:351            reason = "projected_gradient_tolerance"352        return {353            "termination": reason,354            "accepted_steps": len(history),355            "final_objective": final_loss,356            "final_gradient_mapping_mm": final_mapping,357            "history": history,358        }359360    # endregion book:example-reconstruction-projected-armijo361362363def main() -> None:364    args = example_parser(365        "Reconstruct attenuation from supplied calibrated primary counts."366    ).parse_args()367    case = load_case(args.case)368    if not str(args.device).startswith("cuda:"):369        raise ContractError("Select a CUDA device, for example --device cuda:0.")370    torch: Any = importlib.import_module("torch")371    wp: Any = importlib.import_module("warp")372    with case.record(373        args.output,374        entrypoint=__file__,375        metadata={376            "application": "reconstruction",377            "execution_acceptance": "requires separate numerical and physical assessment",378        },379    ) as run:380        # Never build a Torch autograd graph around the manually supplied VJP.381        with torch.cuda.device(args.device), torch.no_grad():382            try:383                solver = Reconstruction(case, device=args.device, torch=torch, wp=wp)384                run.set_metadata(385                    device=args.device,386                    gpu_name=torch.cuda.get_device_name(args.device),387                    torch_version=torch.__version__,388                    volume_units="mm^-1",389                    objective="summed Poisson half-deviance plus quadratic regularisation",390                )391                report = solver.solve()392                wp.synchronize_stream(solver.stream)393                write_array(394                    run, "attenuation.npy", solver.mu.cpu().numpy().reshape(solver.grid.shape)395                )396                for index, view in enumerate(solver.views):397                    write_array(398                        run,399                        f"predictions/view-{index:04d}.npy",400                        view["prediction"].numpy().reshape(view["shape"]),401                    )402                run.write_json("solver.json", report)403            finally:404                # Drain pending work before exception paths release shared owners.405                torch.cuda.synchronize(args.device)406407408if __name__ == "__main__":409    main()410