python/dpt/examples/registration.py

1"""Fit one rigid pose to supplied calibrated primary-count radiographs.23All image formation and derivatives use the canonical CUDA operators. NumPy4handles supplied-file validation and final reports, never an image optimiser.5See the examples README for the provenance-bearing JSON/NPY input contract6and the separately recorded registration study.7"""89from __future__ import annotations1011import importlib12import math13import re14from dataclasses import asdict, dataclass15from typing import Any, cast1617from dpt.contracts import ContractError, finite_scalar, integer18from dpt.examples._common import example_parser, load_case, write_array19from dpt.geometry import RigidTransform, Vector320from dpt.objectives import ObjectiveSpec21from dpt.recovery import PoseChart, PrimaryPoseEvaluator, PrimaryPoseProblem22from dpt.registration import Evaluation, RecoveryPolicy, Vector, recover_parameters23from dpt.validation.recovery import pose_error24from dpt.volumes import GridSpec252627def _mapping(value: Any, name: str) -> dict[str, Any]:28    if not isinstance(value, dict):29        raise ContractError(f"{name} must be a JSON object; see the examples README")30    return cast(dict[str, Any], value)313233def _required(values: dict[str, Any], name: str) -> Any:34    if name not in values:35        raise ContractError(f"registration configuration is missing {name!r}")36    return values[name]373839def _immutable_geometry(values: dict[str, Any]) -> dict[str, Any]:40    """Decode JSON arrays at the immutable host-geometry boundary."""41    return {42        key: tuple(cast(list[Any], value)) if isinstance(value, list) else value43        for key, value in values.items()44    }454647@dataclass(frozen=True, slots=True)48class PreparedView:49    identifier: str50    evaluator: PrimaryPoseEvaluator51    included_pixels: int525354# region book:example-registration-shared-objective55class SharedPoseObjective:56    """Sum views in one immutable chart; each evaluator owns its image scratch."""5758    def __init__(self, views: tuple[PreparedView, ...], chart: PoseChart) -> None:59        if not views or any(view.evaluator.chart is not chart for view in views):60            raise ContractError("all views must share the same PoseChart object")61        self.views, self.chart = views, chart6263    def __call__(self, parameters: Vector) -> Evaluation:64        values = tuple(view.evaluator(parameters) for view in self.views)65        return Evaluation(66            math.fsum(value.loss for value in values),67            tuple(math.fsum(value.gradient[k] for value in values) for k in range(6)),68        )697071# endregion book:example-registration-shared-objective727374def prepare_views(75    case: Any, configuration: dict[str, Any], chart: PoseChart, device: str76) -> tuple[PreparedView, ...]:77    """Validate supplied arrays and upload fixed inputs once on a shared stream."""78    from dpt.geometry import DetectorGeometry7980    wp: Any = importlib.import_module("warp")81    np: Any = importlib.import_module("numpy")82    grid = GridSpec(83        **_immutable_geometry(_mapping(_required(configuration, "grid"), "registration.grid"))84    )85    attenuation_host = case.array(86        _required(configuration, "attenuation"), shape=grid.shape, units="mm^-1"87    )88    if not bool(np.all(np.isfinite(attenuation_host) & (attenuation_host >= 0))):89        raise ContractError("attenuation must contain finite nonnegative values in mm^-1")90    samples = integer(_required(configuration, "samples_per_ray"), "samples_per_ray", minimum=1)91    raw_views = _required(configuration, "views")92    if not isinstance(raw_views, list) or not raw_views:93        raise ContractError("registration.views needs at least one calibrated count view")94    views: list[dict[str, Any]] = [95        _mapping(value, "registration.views entry") for value in cast(list[Any], raw_views)96    ]97    identifiers: set[str] = set()98    prepared: list[PreparedView] = []99    with wp.ScopedDevice(device):100        stream = wp.Stream(device=device)101        with wp.ScopedStream(stream):102            try:103                attenuation = wp.array(104                    attenuation_host.reshape(-1), dtype=wp.float32, device=device105                )106                for view in views:107                    identifier = _required(view, "id")108                    if not isinstance(identifier, str) or not re.fullmatch(109                        r"[A-Za-z0-9_-]+", identifier110                    ):111                        raise ContractError(112                            "view id must use letters, digits, underscores or hyphens"113                        )114                    if identifier in identifiers:115                        raise ContractError(116                            f"duplicate view id {identifier!r}; give each view its own id"117                        )118                    identifiers.add(identifier)119                    geometry = DetectorGeometry(120                        **_immutable_geometry(121                            _mapping(_required(view, "geometry"), "view.geometry")122                        )123                    )124                    observed = case.array(125                        _required(view, "observation"), shape=geometry.shape, units="counts"126                    )127                    mask = case.array(128                        _required(view, "mask"), shape=geometry.shape, units="dimensionless"129                    )130                    valid_counts = np.isfinite(observed) & (observed >= 0) & (observed <= 2**24)131                    if not bool(np.all(valid_counts & (observed == np.floor(observed)))):132                        raise ContractError(133                            f"view {identifier}: counts must be integers in [0, 2**24], stored as "134                            "float32; processed display images cannot use this Poisson example"135                        )136                    if not bool(np.all((mask == 0) | (mask == 1))) or not bool(np.any(mask == 1)):137                        raise ContractError(138                            f"view {identifier}: mask must be binary with included pixels"139                        )140                    if bool(np.any((mask == 0) & (observed != 0))):141                        raise ContractError(142                            f"view {identifier}: set every mask=0 observation to zero in the "143                            "supplied preprocessed array; preserve the raw source and record "144                            "this preprocessing in its provenance"145                        )146                    beam = finite_scalar(_required(view, "open_beam_counts"), "open_beam_counts")147                    if beam <= 0:148                        raise ContractError(f"view {identifier}: open_beam_counts must be positive")149                    # region book:example-registration-prepare-view150                    problem = PrimaryPoseProblem(151                        grid=grid,152                        geometry=geometry,153                        attenuation=attenuation,154                        observation=wp.array(observed.reshape(-1), dtype=wp.float32, device=device),155                        objective=ObjectiveSpec(156                            kind="poisson", domain="counts", reduction="sum", weighted=True157                        ),158                        open_beam=beam,159                        weights=wp.array(mask.reshape(-1), dtype=wp.float32, device=device),160                        samples_per_ray=samples,161                        precision=configuration.get("precision", "float64"),162                        integration=configuration.get("integration", "midpoint"),163                    )164                    evaluator = PrimaryPoseEvaluator(problem, chart, device=device, stream=stream)165                    # endregion book:example-registration-prepare-view166                    prepared.append(167                        PreparedView(identifier, evaluator, int(np.count_nonzero(mask)))168                    )169            finally:170                # Drain uploads and setup work before their owners can be released.171                wp.synchronize_stream(stream)172    return tuple(prepared)173174175def evaluate_reference(case: Any, configuration: dict[str, Any], recovered: RigidTransform) -> Any:176    """Read the optional reference only after the accepted pose has been fixed."""177    if "evaluation" not in configuration:178        return {"status": "no independent geometric reference supplied"}179    evaluation = _mapping(configuration["evaluation"], "registration.evaluation")180    for key in ("source", "uncertainty"):181        if not isinstance(evaluation.get(key), str) or not evaluation[key].strip():182            raise ContractError(f"evaluation.{key} must describe the reference and its limits")183    reference = RigidTransform(184        **_immutable_geometry(_mapping(_required(evaluation, "reference_pose"), "reference_pose"))185    )186    landmarks = case.array(_required(evaluation, "landmarks"), dtype="float64", units="mm")187    np: Any = importlib.import_module("numpy")188    if landmarks.ndim != 2 or landmarks.shape[1] != 3 or landmarks.shape[0] < 1:189        raise ContractError("evaluation landmarks must have shape (N, 3), with N >= 1")190    if not bool(np.all(np.isfinite(landmarks))):191        raise ContractError(192            "evaluation landmarks must contain finite object-frame coordinates in mm"193        )194    points = [cast(Vector3, tuple(float(value) for value in point)) for point in landmarks]195    return {196        "status": "evaluated against supplied reference",197        "source": evaluation["source"],198        "uncertainty": evaluation["uncertainty"],199        "reference_pose": asdict(reference),200        "errors": asdict(pose_error(recovered, reference, points)),201    }202203204def main() -> None:205    parser = example_parser("Register a known attenuation volume to supplied primary-count views.")206    args = parser.parse_args()207    case = load_case(args.case)208    configuration = _mapping(case.config.get("registration"), "registration")209    chart_values = dict(_mapping(_required(configuration, "chart"), "registration.chart"))210    anchor = RigidTransform(211        **_immutable_geometry(_mapping(_required(chart_values, "anchor"), "chart.anchor"))212    )213    chart_values["anchor"] = anchor214    _required(chart_values, "scales")215    chart = PoseChart(**chart_values)216    policy = RecoveryPolicy(**_mapping(_required(configuration, "policy"), "registration.policy"))217    with case.record(218        args.output,219        entrypoint=__file__,220        metadata={221            "application": "registration",222            "device_requested": args.device,223            "measurement_model": "independent primary counts with scalar open beam per view",224            "objective": "sum of masked Poisson half-deviances",225            "evaluation_reference_used_for_fit": False,226        },227    ) as run:228        views = prepare_views(case, configuration, chart, args.device)229        objective = SharedPoseObjective(views, chart)230        run.set_metadata(device=str(views[0].evaluator.context.device))231        # region book:example-registration-solve232        initial: Vector = (0.0,) * 6233        result = recover_parameters(objective, initial, policy=policy)234        recovered = chart.pose(result.parameters)235        run.write_json(236            "optimisation.json",237            {238                "pose_object_to_world": asdict(recovered),239                "chart": asdict(chart),240                "policy": asdict(policy),241                "precision": views[0].evaluator.problem.precision,242                "integration": views[0].evaluator.problem.integration,243                "result": asdict(result),244                "stationary": result.stationary,245            },246        )247        # endregion book:example-registration-solve248        final_losses: dict[str, float] = {}249        # region book:example-registration-export250        for view in views:251            evaluator = view.evaluator252            # A rejected line-search trial may be the last occupant of scratch.253            final_value = evaluator(result.parameters)254            final_losses[view.identifier] = final_value.loss255            prediction = evaluator.prediction.numpy().reshape(evaluator.problem.geometry.shape)256            write_array(run, f"prediction-{view.identifier}.npy", prediction)257        run.write_json(258            "fit-report.json",259            {260                "final_half_deviance_by_view": final_losses,261                "included_pixels_by_view": {v.identifier: v.included_pixels for v in views},262                "solver_evaluations": result.evaluations,263                # A rejected trial can fail before some view evaluators are called.264                "solver_view_evaluation_upper_bound": result.evaluations * len(views),265                "additional_final_export_evaluations": len(views),266                "evaluation": evaluate_reference(case, configuration, recovered),267            },268        )269        # endregion book:example-registration-export270271272if __name__ == "__main__":273    main()274