experiments/worked-applications/run.py

1"""Complete a CT-derived registration and then select, observe and fit a new view.23All projection, derivatives and optimisation use the canonical dpt CUDA code.4The known case is deliberately a teaching example, not a held-out clinical study.5Future noisy observations are generated only after the view decision is saved.6"""78from __future__ import annotations910import argparse11import importlib.util12import itertools13import json14import math15import sys16import time17from dataclasses import asdict18from pathlib import Path1920import numpy as np2122from dpt.experiments import RunRecorder, experiment_sources, private_output, repository_root23from dpt.geometry import RigidTransform, compose_pose24from dpt.projection import ProjectionSpec, prepare_projection, project_optical_depth25from dpt.validation.projection import integrate_sampled_field26from dpt.validation.recovery import detector_coordinate, pose_error2728# Reuse the supplied-input experiment orchestration; there is no second solver.29APPLICATION = Path(__file__).resolve().parents[1] / "application-study"30sys.path.insert(0, str(APPLICATION))31from _study import (  # noqa: E40232    checked_array,33    digest,34    geometry,35    grid_spec,36    keyed_rng,37    read_json,38    registration_case,39    rigid,40    save_array,41)42from batch import rank_design  # noqa: E402434445def _load_fit():46    # This entrypoint is itself called run.py: load its sibling driver by path.47    spec = importlib.util.spec_from_file_location("application_fit_driver", APPLICATION / "run.py")48    assert spec is not None and spec.loader is not None49    module = importlib.util.module_from_spec(spec)50    spec.loader.exec_module(module)51    return module.fit_case525354def validate_protocol(protocol):55    """Reject changes that would invalidate this experiment's shared contracts."""56    acquisition = protocol["acquisition"]57    angles = acquisition["candidate_angles_degrees"]58    if protocol["task_coordinates_mm"] != [-40.0, 40.0]:59        raise ValueError("the reused design driver requires the eight ±40-mm task corners")60    if acquisition["base_view_degrees"] != 0.0:61        raise ValueError("this worked sequence starts from the declared zero-degree base view")62    if (63        not angles64        or any(not math.isfinite(a) or a != round(a) for a in angles)65        or len(set(angles)) != len(angles)66        or acquisition["fixed_view_degrees"] not in angles67    ):68        raise ValueError(69            "candidate angles must be distinct integral degrees including the baseline"70        )71    if protocol["randomness"]["replicates"] != 8:72        raise ValueError("the paired summary is declared for eight independent replicates")73    if (74        protocol["numerics"]["policy"]["gradient_tolerance"]75        != protocol["acceptance"]["gradient_tolerance"]76    ):77        raise ValueError("the solver and acceptance gradient thresholds must agree")78    if set(protocol["held_out_views_degrees"]) & {79        0.0,80        *angles,81        *protocol["registration"]["views_degrees"],82    }:83        raise ValueError("reserved views must be disjoint from every fitting candidate")848586def _project(wp, field_device, grid, protocol, angle, pose, samples, device, *, generating=False):87    geom = geometry(protocol, angle)88    numerics = protocol["numerics"]89    precision = numerics["generation_precision" if generating else "precision"]90    integration = numerics["generation_integration" if generating else "integration"]91    workspace = prepare_projection(92        grid,93        geom,94        ProjectionSpec(samples, precision=precision, integration=integration),95        device=device,96    )97    output = wp.empty(geom.pixels, dtype=workspace.dtype, device=device)98    packed = wp.array(pose.packed(), dtype=wp.float64, device=device)99    project_optical_depth(field_device, packed, workspace=workspace, out_L=output)100    return output.numpy().reshape(geom.shape).astype(np.float64)101102103def _sampling(wp, field, field_device, grid, protocol, target, angles, device):104    """Qualify every prospective angle before producing any noisy data."""105    shape = protocol["geometry"]["detector_shape_hw"]106    rows = tuple(min(shape[0] - 1, round(shape[0] * x)) for x in (0.125, 0.375, 0.625, 0.875))107    cols = tuple(min(shape[1] - 1, round(shape[1] * x)) for x in (0, 0.25, 0.4375, 0.5, 0.75, 1))108    pixels = list(itertools.product(rows, cols))109    values = field.ravel().tolist()110    means = {}111    diagnostics = []112    cfg = protocol["numerics"]113    beam = cfg["sampling_gate_open_beam_per_pixel"]114    for angle in angles:115        geom = geometry(protocol, angle)116        corners = [117            grid.grid_to_object(tuple(index))118            for index in itertools.product(119                (-0.5, grid.shape[2] - 0.5),120                (-0.5, grid.shape[1] - 0.5),121                (-0.5, grid.shape[0] - 0.5),122            )123        ]124        coverage = [detector_coordinate(geom, target.point(corner)) for corner in corners]125        if any(126            not (-0.5 <= c <= geom.shape[1] - 0.5 and -0.5 <= r <= geom.shape[0] - 0.5)127            for c, r in coverage128        ):129            raise ValueError("the detector does not cover the declared transformed field")130        generated = _project(131            wp,132            field_device,133            grid,134            protocol,135            angle,136            target,137            cfg["generation_samples_per_ray"],138            device,139            generating=True,140        )141        fitted = _project(142            wp,143            field_device,144            grid,145            protocol,146            angle,147            target,148            cfg["fitting_samples_per_ray"],149            device,150        )151        exact = np.array(152            [153                integrate_sampled_field(154                    values, grid, geometry(protocol, angle), target, r, c, validate=False155                )156                for r, c in pixels157            ]158        )159        exact_mean = beam * np.exp(-exact)160        independent_error = np.abs(161            beam * np.exp(-np.array([generated[r, c] for r, c in pixels])) - exact_mean162        ) / np.sqrt(exact_mean)163        fitting_independent_error = np.abs(164            beam * np.exp(-np.array([fitted[r, c] for r, c in pixels])) - exact_mean165        ) / np.sqrt(exact_mean)166        reference_mean = beam * np.exp(-generated)167        full_error = np.abs(beam * np.exp(-fitted) - reference_mean) / np.sqrt(reference_mean)168        row = {169            "angle_degrees": angle,170            "projected_support_corners_column_row": coverage,171            "pixels": pixels,172            "independent_exact_depths": exact.tolist(),173            "generation_independent_p95_sd": float(np.quantile(independent_error, 0.95)),174            "generation_independent_max_sd": float(independent_error.max()),175            "generation_integration": cfg["generation_integration"],176            "fitting_integration": cfg["integration"],177            "fitting_independent_p95_sd": float(np.quantile(fitting_independent_error, 0.95)),178            "fitting_independent_max_sd": float(fitting_independent_error.max()),179            "fitting_full_image_p95_sd": float(np.quantile(full_error, 0.95)),180            "fitting_full_image_max_sd": float(full_error.max()),181        }182        row["passed"] = (183            max(184                row["generation_independent_p95_sd"],185                row["fitting_independent_p95_sd"],186                row["fitting_full_image_p95_sd"],187            )188            <= cfg["sampling_error_poisson_sd_p95_max"]189            and max(190                row["generation_independent_max_sd"],191                row["fitting_independent_max_sd"],192                row["fitting_full_image_max_sd"],193            )194            <= cfg["sampling_error_poisson_sd_max"]195        )196        diagnostics.append(row)197        means[angle] = np.exp(-generated)198        print(f"sampling angle={angle:g} passed={row['passed']}", flush=True)199    return means, diagnostics200201202def _observation(record, public, means, protocol, angle, beam, role, replicate, seed):203    code = round(angle) + 180204    identifier = f"r{role}-n{replicate:02d}-a{code:03d}"205    if identifier in public["views"]:206        return identifier207    counts64 = keyed_rng(seed, role, replicate, code).poisson(beam * means[angle])208    if counts64.max() > 2**24:209        raise ValueError("generated counts exceed exact binary32 integers")210    counts = counts64.astype(np.float32)211    mask = np.ones(counts.shape, dtype=np.float32)212    observation = save_array(213        record.output,214        f"observations/{identifier}.npy",215        counts,216        units="counts",217        role="simulated fitting observation",218        recorder=record,219    )220    mask_record = save_array(221        record.output,222        f"masks/{identifier}.npy",223        mask,224        units="dimensionless",225        role="fixed binary fitting mask",226        recorder=record,227    )228    public["views"][identifier] = {229        "angle_degrees": angle,230        "geometry": asdict(geometry(protocol, angle)),231        "open_beam_counts": beam,232        "observation": observation,233        "mask": mask_record,234    }235    return identifier236237238def _fit(known, output, public, protocol, name, identifiers, anchor, device):239    case_path = output / "cases" / f"{name}.json"240    registration_case(known, output, public, protocol, identifiers, anchor, case_path)241    fit_case = _load_fit()242    result = fit_case(243        case_path,244        output / "fits" / name,245        device=device,246        deadline=time.perf_counter() + protocol["numerics"]["solve_soft_seconds"],247    )248    print(249        f"fit {name}: {result['result']['reason']} "250        f"gradient={max(map(abs, result['result']['evaluation']['gradient'])):.6g}",251        flush=True,252    )253    return result254255256def _assess(wp, field_device, grid, protocol, means, target, outcome, points, device):257    recovered = rigid(outcome["pose_object_to_world"])258    errors = asdict(pose_error(recovered, target, tuple(map(tuple, points))))259    initial = asdict(260        pose_error(rigid(outcome["chart"]["anchor"]), target, tuple(map(tuple, points)))261    )262    reserved = []263    beam = protocol["acquisition"]["base_and_candidate_open_beam_per_pixel"]264    for angle in protocol["held_out_views_degrees"]:265        depth = _project(266            wp,267            field_device,268            grid,269            protocol,270            angle,271            recovered,272            protocol["numerics"]["fitting_samples_per_ray"],273            device,274        )275        expected = beam * means[angle]276        discrepancy = (beam * np.exp(-depth) - expected) / np.sqrt(expected)277        reserved.append(278            {279                "angle_degrees": angle,280                "expected_error_poisson_sd_rms": float(np.sqrt(np.mean(discrepancy**2))),281            }282        )283    gate = protocol["acceptance"]284    # pose_error exposes radians; the protocol intentionally uses reader-facing degrees.285    geometric = (286        errors["landmark_rms_mm"] <= gate["rms_probe_error_mm_max"]287        and errors["landmark_max_mm"] <= gate["maximum_probe_error_mm_max"]288        and math.degrees(errors["rotation_radians"]) <= gate["rotation_error_degrees_max"]289    )290    return {291        "stationary": outcome["stationary"],292        "termination": outcome["result"]["reason"],293        "errors": errors,294        "initial_errors": initial,295        "held_out": reserved,296        "geometric_pass": geometric,297        "held_out_pass": all(298            x["expected_error_poisson_sd_rms"] <= gate["held_out_expected_error_poisson_sd_rms_max"]299            for x in reserved300        ),301    }302303304def execute(known_root, output, config, device, development=False):305    import warp as wp306307    wp.init()308    known_root = known_root.resolve()309    protocol = read_json(config)310    validate_protocol(protocol)311    known = read_json(known_root / "known.json")312    if (313        known["case"] != protocol["case"]314        or known["attenuation"]["sha256"] != protocol["known_attenuation_sha256"]315    ):316        raise ValueError("this worked example requires its declared CT-derived field")317    field = checked_array(known_root, known["attenuation"])318    grid = grid_spec(known["grid"])319    target = compose_pose(RigidTransform(), tuple(protocol["generating_local_se3_mm_rad"]))320    points = list(itertools.product(protocol["task_coordinates_mm"], repeat=3))321    angles = sorted(322        set(323            [324                *protocol["registration"]["views_degrees"],325                protocol["acquisition"]["base_view_degrees"],326                *protocol["acquisition"]["candidate_angles_degrees"],327                *protocol["held_out_views_degrees"],328            ]329        )330    )331    replicates = 1 if development else protocol["randomness"]["replicates"]332    seed = protocol["randomness"]["development_root_seed" if development else "root_seed"]333    sources = experiment_sources(334        __file__,335        config,336        extra={337            "known.json": known_root / "known.json",338            "known-attenuation.npy": known_root / known["attenuation"]["file"],339            **{f"application-study/{p.name}": p for p in APPLICATION.glob("*.py")},340        },341    )342    output = private_output(output, repository_root(__file__))343    with RunRecorder(344        output,345        configuration={346            "protocol": protocol,347            "development": development,348            "replicates": replicates,349            "root_seed": seed,350            "scope": "CT-derived simulated teaching example; not clinical validation",351        },352        sources=sources,353    ) as record:354        record.set_metadata(device=str(wp.get_device(device)), warp=wp.__version__)355        with wp.ScopedDevice(device):356            field_device = wp.array(field.ravel(), dtype=wp.float32, device=device)357            means, sampling = _sampling(358                wp, field, field_device, grid, protocol, target, angles, device359            )360            record.write_json("sampling.json", sampling)361            if not all(x["passed"] for x in sampling):362                raise ValueError("prospective quadrature checks failed; no observations generated")363            record.write_json(364                "evaluation-reference.json",365                {366                    "pose_object_to_world": asdict(target),367                    "task_points_object_mm": points,368                    "role": "generation/evaluation only; not supplied to fit or selector",369                },370            )371            public = {372                "known_sha256": digest(known_root / "known.json"),373                "views": {},374                "effective_fitting_samples": protocol["numerics"]["fitting_samples_per_ray"],375            }376            (output / "cases").mkdir()377            (output / "fits").mkdir()378            reg_ids = [379                _observation(380                    record,381                    public,382                    means,383                    protocol,384                    a,385                    protocol["registration"]["open_beam_per_view"],386                    1,387                    0,388                    seed,389                )390                for a in protocol["registration"]["views_degrees"]391            ]392            # region book:worked-registration393            registration = _fit(394                known_root,395                output,396                public,397                protocol,398                "registration",399                reg_ids,400                RigidTransform(),401                device,402            )403            # endregion book:worked-registration404            results = []405            # region book:worked-application-sequence406            for replicate in range(replicates):407                beam = protocol["acquisition"]["base_and_candidate_open_beam_per_pixel"]408                base_id = _observation(409                    record, public, means, protocol, 0.0, beam, 2, replicate, seed410                )411                base = _fit(412                    known_root,413                    output,414                    public,415                    protocol,416                    f"base-{replicate:02d}",417                    [base_id],418                    RigidTransform(),419                    device,420                )421                accepted = rigid(base["pose_object_to_world"])422                decision_path = output / f"decision-{replicate:02d}"423                selection = rank_design(424                    known_root,425                    output,426                    public,427                    protocol,428                    {429                        "angles": protocol["acquisition"]["candidate_angles_degrees"],430                    },431                    accepted,432                    decision_path,433                    device,434                )435                # Selection is on disk before either future count image exists.436                record.write_json(f"decision-{replicate:02d}-binding.json", selection)437                pair = {"replicate": replicate, "base": base, "selection": selection}438                for policy, angle in (439                    ("selected", selection["selected_angle"]),440                    ("near_parallel", protocol["acquisition"]["fixed_view_degrees"]),441                ):442                    future_id = _observation(443                        record, public, means, protocol, angle, beam, 3, replicate, seed444                    )445                    pair[policy] = _fit(446                        known_root,447                        output,448                        public,449                        protocol,450                        f"{policy}-{replicate:02d}",451                        [base_id, future_id],452                        accepted,453                        device,454                    )455                results.append(pair)456            # endregion book:worked-application-sequence457            record.write_json("fitting-observations.json", public)458            evaluations = {459                "registration": _assess(460                    wp, field_device, grid, protocol, means, target, registration, points, device461                ),462                "pairs": [],463            }464            for pair in results:465                evaluations["pairs"].append(466                    {467                        "replicate": pair["replicate"],468                        "selection": pair["selection"],469                        **{470                            policy: _assess(471                                wp,472                                field_device,473                                grid,474                                protocol,475                                means,476                                target,477                                pair[policy],478                                points,479                                device,480                            )481                            for policy in ("base", "selected", "near_parallel")482                        },483                    }484                )485            selected = np.array(486                [p["selected"]["errors"]["landmark_rms_mm"] ** 2 for p in evaluations["pairs"]]487            )488            fixed = np.array(489                [p["near_parallel"]["errors"]["landmark_rms_mm"] ** 2 for p in evaluations["pairs"]]490            )491            differences = selected - fixed492            ratio = float(selected.mean() / fixed.mean())493            upper = (494                None495                if development496                else float(497                    differences.mean()498                    + protocol["acceptance"]["paired_t_critical_7df"]499                    * differences.std(ddof=1)500                    / np.sqrt(len(differences))501                )502            )503            all_outcomes = [504                evaluations["registration"],505                *[506                    p[k]507                    for p in evaluations["pairs"]508                    for k in ("base", "selected", "near_parallel")509                ],510            ]511            completed = all(512                x["stationary"] and x["geometric_pass"] and x["held_out_pass"] for x in all_outcomes513            )514            benefit = (515                ratio516                <= protocol["acceptance"]["selected_vs_fixed_mean_squared_task_error_ratio_max"]517                and upper is not None518                and upper < protocol["acceptance"]["paired_error_difference_t95_upper_max"]519            )520            record.write_json("evaluation.json", evaluations)521            summary = {522                "development": development,523                "replicates": replicates,524                "all_solves_stationary_and_accurate": completed,525                "selected_mean_squared_task_error_mm2": float(selected.mean()),526                "near_parallel_mean_squared_task_error_mm2": float(fixed.mean()),527                "selected_to_near_parallel_ratio": ratio,528                "paired_mean_difference_mm2": float(differences.mean()),529                "paired_descriptive_t95_upper_mm2": upper,530                "benefit_pass": bool(benefit),531                "passed": bool(not development and completed and benefit),532                "uncertainty_scope": (533                    "Eight independent prescribed Poisson pairs on one teaching anatomy; "534                    "descriptive t interval, not a population or clinical claim."535                ),536            }537            record.write_json("summary.json", summary)538            children = {539                str(p.relative_to(output)): digest(p)540                for p in sorted(output.rglob("run.json"))541                if p != output / "run.json" and "sources" not in p.parts542            }543            record.write_json("child-records.json", children)544            print(json.dumps(summary, indent=2), flush=True)545    return summary546547548def main():549    parser = argparse.ArgumentParser(description=__doc__)550    parser.add_argument("--known", type=Path, required=True)551    parser.add_argument("--config", type=Path, default=Path(__file__).with_name("config.json"))552    parser.add_argument("--output", type=Path, required=True)553    parser.add_argument("--device", default="cuda:0")554    parser.add_argument(555        "--development",556        action="store_true",557        help="Separate prescribed development seed and one pair; never a final passing result.",558    )559    args = parser.parse_args()560    result = execute(args.known, args.output, args.config.resolve(), args.device, args.development)561    if not args.development and not result["passed"]:562        raise SystemExit("worked-example acceptance failed; retained all outcomes")563564565if __name__ == "__main__":566    main()567