python/dpt/examples/acquisition_design.py

1"""Rank supplied next-view candidates by expected local target-position variance.23The known attenuation field and calibrated candidate beams are supplied inputs.4The design uses independent ideal Poisson counts, a fixed local pose chart and a5strictly positive-definite current precision. It does not acquire an image or6control imaging hardware. Only compact information matrices leave the GPU.7See the examples README for recorded policy comparisons and their limitations.8"""910from __future__ import annotations1112import importlib13import math14from typing import Any, Literal, cast1516from dpt._runtime import prepare_context17from dpt.contracts import ContractError, NumericalError, finite_scalar, integer18from dpt.examples._common import example_parser, load_case, write_array19from dpt.geometry import DetectorGeometry, RigidTransform20from dpt.projection import (21    ProjectionSpec,22    prepare_projection,23    project_optical_depth,24    projection_pose_sensitivities,25)26from dpt.transmission import TransmissionSpec, prepare_transmission, transmit27from dpt.volumes import GridSpec282930def _text(value: Any, name: str) -> str:31    if not isinstance(value, str) or not value.strip():32        raise ContractError(f"{name} needs a non-empty description")33    return value343536def _positive(value: Any, name: str) -> float:37    result = finite_scalar(value, name, minimum=0.0)38    if result == 0:39        raise ContractError(f"{name} must be strictly positive")40    return result414243def _mapping(value: Any, name: str) -> dict[str, Any]:44    if not isinstance(value, dict):45        raise ContractError(f"{name} must be a JSON object with named fields")46    if any(not isinstance(key, str) for key in cast(dict[Any, Any], value)):47        raise ContractError(f"{name} must be a JSON object with named fields")48    return cast(dict[str, Any], value)495051def _grid(raw: dict[str, Any]) -> GridSpec:52    values = {key: tuple(value) for key, value in raw.items()}53    return GridSpec(**values)545556def _pose(raw: dict[str, Any]) -> RigidTransform:57    values = {key: tuple(value) for key, value in raw.items()}58    return RigidTransform(**values)596061def _geometry(raw: dict[str, Any]) -> DetectorGeometry:62    values = {key: tuple(value) for key, value in raw.items()}63    return DetectorGeometry(**values)646566# region book:example-acquisition-targets67def target_jacobians(np: Any, pose: RigidTransform, targets: Any, scales: Any) -> Any:68    """Return d(world target)/dz in mm for T_WO exp((S z)^) at z=0."""69    rotation = np.asarray(pose.rotation, dtype=np.float64).reshape(3, 3)70    result = np.empty((len(targets), 3, 6), dtype=np.float64)71    for index, (x, y, z) in enumerate(targets):72        skew = np.asarray([[0.0, -z, y], [z, 0.0, -x], [-y, x, 0.0]])73        result[index, :, :3] = rotation74        result[index, :, 3:] = -rotation @ skew75    result *= scales[None, None, :]76    if not np.isfinite(result).all():77        raise NumericalError("target derivatives exceed FP64; check target coordinates and scales")78    return result798081def target_variance(82    np: Any, precision: Any, target_jacobian: Any, maximum_condition: float83) -> tuple[float, Any, float]:84    """Mean target-position variance in mm², without forming an inverse to score."""85    if not np.isfinite(precision).all():86        raise NumericalError("pose precision is non-finite; discard this candidate score")87    values = np.linalg.eigvalsh(precision)88    if values[0] <= 0 or values[-1] / values[0] > maximum_condition:89        raise NumericalError(90            "pose precision must be positive definite and within maximum_precision_condition; "91            "check the stated uncertainty and parameter scales, without adding numerical jitter"92        )93    factor = np.linalg.cholesky(precision)94    right = target_jacobian.reshape(-1, 6).T95    whitened = np.linalg.solve(factor, right)96    if not np.isfinite(whitened).all():97        raise NumericalError(98            "target covariance solve exceeds FP64; check the precision and chart scales"99        )100    # hypot scales its sum of squares; individual products must not underflow101    # before they contribute to an otherwise representable total variance.102    normalisation = math.sqrt(len(target_jacobian))103    root_variance = math.hypot(*(float(value) / normalisation for value in whitened.flat))104    variance = root_variance * root_variance105    if not math.isfinite(variance) or variance <= 0:106        raise NumericalError(107            "positive target-position variance is outside FP64 range; check the stated "108            "uncertainty and physical units, or use a range-preserving score representation"109        )110    inverse_factor = np.linalg.solve(factor, np.eye(6))111    covariance = inverse_factor.T @ inverse_factor112    if not np.isfinite(covariance).all() or bool((np.diag(covariance) <= 0).any()):113        raise NumericalError(114            "local covariance is outside FP64 range; check the precision and chart scales "115            "before exporting this candidate"116        )117    return variance, covariance, float(values[-1] / values[0])118119120# endregion book:example-acquisition-targets121122123# region book:example-acquisition-information124def candidate_information(125    *,126    np: Any,127    torch: Any,128    ctx: Any,129    grid: GridSpec,130    geometry: DetectorGeometry,131    attenuation: Any,132    pose: Any,133    beam_host: Any,134    scales: Any,135    samples_per_ray: int,136    precision: Literal["float32", "float64"] = "float32",137    integration: Literal["midpoint", "cell_gauss"] = "midpoint",138) -> tuple[Any, dict[str, Any]]:139    """Use canonical depth derivatives and an FP64 CUDA Fisher reduction.140141    Preparation allocates one candidate's O(6P) Jacobian and weighted copy.142    Torch views alias Warp storage on the same CUDA stream. There is no143    per-pixel Python loop and no image/Jacobian download. This is an offline144    candidate evaluation, not a captured or profiled optimisation hot path.145    """146    wp, pixels = ctx.wp, geometry.pixels147    try:148        projection = prepare_projection(149            grid,150            geometry,151            ProjectionSpec(samples_per_ray, precision=precision, integration=integration),152            device=str(ctx.device),153            stream=ctx.stream,154        )155        transmission = (156            prepare_transmission(157                TransmissionSpec(beam="none"),158                max_pixels=pixels,159                device=str(ctx.device),160                stream=ctx.stream,161            )162            if precision == "float32"163            else None164        )165        with ctx.scope():166            beam = wp.array(beam_host.reshape(-1), dtype=wp.float32, device=ctx.device)167            depth = wp.empty(pixels, dtype=projection.dtype, device=ctx.device)168            log_transmission = wp.empty(169                pixels if precision == "float32" else 0, dtype=wp.float32, device=ctx.device170            )171            ones = wp.ones(pixels, dtype=projection.dtype, device=ctx.device)172            jacobian = wp.empty(6 * pixels, dtype=wp.float64, device=ctx.device)173        project_optical_depth(174            attenuation,175            pose,176            workspace=projection,177            out_L=depth,178            stream=ctx.stream,179        )180        if transmission is not None:181            transmit(182                depth,183                out_log_T=log_transmission,184                workspace=transmission,185                stream=ctx.stream,186            )187        projection_pose_sensitivities(188            attenuation,189            pose,190            ones,191            out_jacobian=jacobian,192            workspace=projection,193            stream=ctx.stream,194        )195        beam_tensor = wp.to_torch(beam)196        log_mean = beam_tensor.to(dtype=torch.float64).log_()197        if precision == "float64":198            log_mean.sub_(wp.to_torch(depth))199        else:200            log_mean.add_(wp.to_torch(log_transmission))201        illuminated = beam_tensor > 0202        # A zero supplied beam carries no information. Positive means outside203        # the normal FP64 range are rejected; no tail, floor or pixel is hidden.204        log_tiny = math.log(float(np.finfo(np.float64).tiny))205        if bool((illuminated & (log_mean < log_tiny)).any().item()):206            raise NumericalError(207                "a positive candidate count mean is below normal FP64 range; "208                "use a range-preserving design implementation before ranking this case"209            )210        root_mean = log_mean.mul_(0.5).exp_()211        depth_jacobian = wp.to_torch(jacobian).reshape(pixels, 6)212        weighted = depth_jacobian.clone()213        weighted.mul_(scales).mul_(root_mean[:, None])214        if not bool(torch.isfinite(weighted).all().item()):215            raise NumericalError("weighted pose derivatives exceed FP64; discard this candidate")216        if bool(((depth_jacobian != 0) & illuminated[:, None] & (weighted == 0)).any().item()):217            raise NumericalError(218                "a nonzero weighted derivative underflowed; discard this candidate"219            )220        magnitude = weighted.abs()221        minimum = math.sqrt(float(np.finfo(np.float64).tiny))222        maximum = math.sqrt(float(np.finfo(np.float64).max) / (6.0 * pixels))223        if bool((((magnitude > 0) & (magnitude < minimum)) | (magnitude > maximum)).any().item()):224            raise NumericalError(225                "candidate Fisher products exceed the declared FP64 accumulation range; "226                "rescale the pose chart or use a range-preserving reduction"227            )228        # Weighted depth derivatives give J_L^T diag(lambda) J_L directly.229        # Avoid a division by rounded, potentially zero count predictions.230        information = weighted.T @ weighted231        result = information.cpu().numpy().copy()232        if not np.isfinite(result).all():233            raise NumericalError("candidate information is non-finite; discard this score")234        diagnostics = {235            "pixels": pixels,236            "precision": precision,237            "integration": integration,238            "zero_beam_pixels": int((~illuminated).sum().item()),239            "omitted_positive_beam_pixels": 0,240            "jacobian_and_weighted_copy_bytes": 2 * 6 * pixels * 8,241            "fisher_reduction": "FP64 CUDA; only the 6 by 6 information matrix is downloaded",242        }243        return 0.5 * (result + result.T), diagnostics244    finally:245        # Warp owners and the host beam must survive every pending shared-stream operation.246        wp.synchronize_stream(ctx.stream)247248249# endregion book:example-acquisition-information250251252def main() -> None:253    args = example_parser(__doc__ or "Rank supplied next-view candidates").parse_args()254    case = load_case(args.case)255    cfg = _mapping(case.config["acquisition_design"], "acquisition_design")256    np: Any = importlib.import_module("numpy")257    grid = _grid(_mapping(cfg["grid"], "grid"))258    pose = _pose(_mapping(cfg["pose"], "pose"))259    attenuation_host = case.array(cfg["attenuation"], shape=grid.shape, units="mm^-1")260    if (attenuation_host < 0).any():261        raise ContractError("attenuation must be nonnegative and already converted to mm^-1")262    prior = case.array(cfg["prior_precision"], dtype="float64", shape=(6, 6), units="dimensionless")263    if not np.allclose(prior, prior.T, rtol=0.0, atol=1e-12 * float(abs(prior).max())):264        raise ContractError("prior_precision must be symmetric in the supplied scaled pose chart")265    prior = 0.5 * (prior + prior.T)266    targets = case.array(cfg["targets_object_mm"], dtype="float64", units="mm")267    if targets.ndim != 2 or targets.shape[1] != 3 or not 1 <= targets.shape[0] <= 256:268        raise ContractError(269            "targets_object_mm must contain 1 to 256 object-frame points, shape (N,3)"270        )271    raw_scales = cfg["parameter_scales"]272    if len(raw_scales) != 6:273        raise ContractError("parameter_scales needs six positive values: mm, mm, mm, rad, rad, rad")274    scales = np.asarray([_positive(value, "parameter scale") for value in raw_scales])275    target_derivatives = target_jacobians(np, pose, targets, scales)276    samples = integer(cfg["samples_per_ray"], "samples_per_ray", minimum=1)277    capacity = integer(cfg["max_candidate_pixels"], "max_candidate_pixels", minimum=1)278    maximum_cost = _positive(cfg["maximum_cost"], "maximum_cost")279    cost_unit = cfg["cost_unit"]280    if cost_unit not in ("mAs", "s", "expected_detector_photons"):281        raise ContractError(282            "cost_unit must be mAs, s or expected_detector_photons; none is absorbed dose"283        )284    current_description = _text(cfg["current_information_description"], "current information")285    maximum_condition = _positive(286        cfg.get("maximum_precision_condition", 1e12), "precision condition"287    )288    if maximum_condition < 1:289        raise ContractError("maximum_precision_condition must be at least one")290    rank_tolerance = _positive(cfg.get("rank_relative_tolerance", 1e-10), "rank tolerance")291    if rank_tolerance >= 1:292        raise ContractError("rank_relative_tolerance must be smaller than one")293    baseline, baseline_covariance, _ = target_variance(294        np, prior, target_derivatives, maximum_condition295    )296    raw_candidates = cfg["candidates"]297    if not isinstance(raw_candidates, list) or not raw_candidates:298        raise ContractError("candidates must list at least one supplied feasible acquisition")299    candidates: list[dict[str, Any]] = []300    names: set[str] = set()301    for raw in cast(list[Any], raw_candidates):302        candidate = _mapping(raw, "candidate")303        name = _text(candidate["name"], "candidate name")304        if name in names:305            raise ContractError(f"candidate name {name!r} is repeated; assign unique names")306        names.add(name)307        geometry = _geometry(_mapping(candidate["geometry"], f"{name} geometry"))308        if geometry.pixels > capacity:309            raise ContractError(310                f"{name} exceeds max_candidate_pixels; declare a sufficient memory bound"311            )312        beam = case.array(candidate["open_beam"], shape=geometry.shape, units="photons/pixel")313        if (beam < 0).any() or not (beam > 0).any():314            raise ContractError(315                f"{name} needs a nonnegative beam with at least one illuminated pixel"316            )317        cost = _positive(candidate["cost"], f"{name} cost")318        if cost_unit == "expected_detector_photons":319            expected = float(np.sum(beam, dtype=np.float64))320            if not math.isclose(cost, expected, rel_tol=1e-6, abs_tol=0.0):321                raise ContractError(f"{name} cost must equal the sum of its open-beam expectations")322        candidates.append(323            {324                "name": name,325                "geometry": geometry,326                "beam": beam,327                "cost": cost,328                "feasibility_description": _text(329                    candidate["feasibility_description"], f"{name} feasibility"330                ),331            }332        )333    eligible = [candidate for candidate in candidates if candidate["cost"] <= maximum_cost]334    if not eligible:335        raise ContractError(336            "no supplied candidate fits maximum_cost; revise the candidate set or budget"337        )338    with case.record(339        args.output,340        entrypoint=__file__,341        metadata={342            "application": "finite_candidate_acquisition_design",343            "evidence_kind": "expected local uncertainty under a fixed ideal Poisson model",344            "current_information_description": current_description,345            "pose_chart": "local right SE(3) at supplied pose; dimensionless z with supplied S",346            "execution_validation": "this run does not establish achieved registration accuracy",347            "tail_policy": (348                "reject unrepresentable positive-beam contributions; no positive tails omitted"349            ),350        },351    ) as run:352        torch: Any = importlib.import_module("torch")353        wp: Any = importlib.import_module("warp")354        torch_device = torch.device(args.device)355        if torch_device.type != "cuda":356            raise ContractError("acquisition design requires CUDA; select --device cuda:N")357        torch_stream = torch.cuda.current_stream(torch_device)358        wp.init()359        stream = wp.stream_from_torch(torch_stream)360        ctx = prepare_context(device=args.device, stream=stream)361        rows: list[dict[str, Any]] = []362        matrices: list[Any] = []363        covariances: list[Any] = []364        try:365            with (366                torch.cuda.device(torch_device),367                torch.cuda.stream(torch_stream),368                torch.no_grad(),369                ctx.scope(),370            ):371                attenuation = wp.array(372                    attenuation_host.reshape(-1), dtype=wp.float32, device=ctx.device373                )374                pose_device = wp.array(pose.packed(), dtype=wp.float64, device=ctx.device)375                scales_device = torch.as_tensor(scales, dtype=torch.float64, device=torch_device)376                # region book:example-acquisition-rank377                for candidate in eligible:378                    information, diagnostics = candidate_information(379                        np=np,380                        torch=torch,381                        ctx=ctx,382                        grid=grid,383                        geometry=candidate["geometry"],384                        attenuation=attenuation,385                        pose=pose_device,386                        beam_host=candidate["beam"],387                        scales=scales_device,388                        samples_per_ray=samples,389                        precision=cfg.get("precision", "float32"),390                        integration=cfg.get("integration", "midpoint"),391                    )392                    variance, covariance, condition = target_variance(393                        np, prior + information, target_derivatives, maximum_condition394                    )395                    eigenvalues = np.linalg.eigvalsh(information)396                    largest = float(eigenvalues[-1])397                    if eigenvalues[0] < -rank_tolerance * max(largest, np.finfo(np.float64).tiny):398                        raise NumericalError(399                            "candidate Fisher matrix is indefinite beyond rank tolerance"400                        )401                    rank = int(np.count_nonzero(eigenvalues > rank_tolerance * largest))402                    rows.append(403                        {404                            "name": candidate["name"],405                            "cost": candidate["cost"],406                            "mean_target_variance_mm2": variance,407                            "local_fisher_rank": rank,408                            "posterior_precision_condition": condition,409                            "diagnostics": diagnostics,410                        }411                    )412                    matrices.append(information)413                    covariances.append(covariance)414                # Exact score ties prefer the smaller stated cost, then the name.415                selected = min(416                    rows,417                    key=lambda row: (row["mean_target_variance_mm2"], row["cost"], row["name"]),418                )419                # endregion book:example-acquisition-rank420                wp.synchronize_stream(stream)421        finally:422            torch_stream.synchronize()423        run.set_metadata(device=str(ctx.device), warp=wp.__version__, torch=torch.__version__)424        write_array(run, "candidate_information.npy", np.stack(matrices))425        write_array(run, "candidate_local_covariance.npy", np.stack(covariances))426        write_array(run, "current_local_covariance.npy", baseline_covariance)427        run.write_json(428            "design.json",429            {430                "selected_candidate": selected["name"],431                "cost_unit": cost_unit,432                "maximum_cost": maximum_cost,433                "current_mean_target_variance_mm2": baseline,434                "score": "mean trace of local target covariance; lower is preferred",435                "rank_relative_tolerance": rank_tolerance,436                "matrix_row_order": [row["name"] for row in rows],437                "candidates": rows,438                "excluded_by_budget": [439                    candidate["name"]440                    for candidate in candidates441                    if candidate["cost"] > maximum_cost442                ],443                "interpretation": "predicted local uncertainty; no selected image acquired",444            },445        )446447448if __name__ == "__main__":449    main()450