experiments/worked-reconstruction/run.py

1"""A fixed CT-derived material-reconstruction example with independent acceptance.23The three commands separate acquisition, fitting and post-freeze evaluation.4The matched coarse generating/inverse basis is intentional and declared; this5is an instructional correctness study, not patient-material validation.6"""78from __future__ import annotations910import argparse11import hashlib12import json13import sys14import time15from dataclasses import asdict, replace16from pathlib import Path17from typing import Any1819import numpy as np2021from dpt.experiments import RunRecorder, experiment_sources, private_output, repository_root22from dpt.geometry import DetectorGeometry, RigidTransform23from dpt.material_reconstruction import (24    MaterialReconstruction,25    MaterialReconstructionSettings,26    MaterialReconstructionView,27)28from dpt.validation.projection import integrate_sampled_field29from dpt.volumes import GridSpec3031STUDY = Path(__file__).resolve().parents[1] / "reconstruction-study"32sys.path.insert(0, str(STUDY))33from _resolution import check_coverage  # noqa: E40234from _study import (  # noqa: E40235    coarse_cells,36    complete_record,37    stationarity_gate,38    stationarity_result,39)40from probe_acquisition import detector, fixed_metric, helpers, pose_at  # noqa: E4024142LIB = helpers()434445def digest(path: Path) -> str:46    return hashlib.sha256(path.read_bytes()).hexdigest()474849def grid_from(record: dict[str, Any]) -> GridSpec:50    return GridSpec(**{key: tuple(value) for key, value in record.items()})515253def checked_array(folder: Path, name: str, record: dict[str, Any]) -> np.ndarray:54    path = folder / name55    if digest(path) != record["output_sha256"][name]:56        raise ValueError(f"recorded input changed: {name}")57    return np.load(path, allow_pickle=False)585960def checked_json(folder: Path, name: str, record: dict[str, Any]) -> dict[str, Any]:61    path = folder / name62    if digest(path) != record["output_sha256"][name]:63        raise ValueError(f"recorded metadata changed: {name}")64    return json.loads(path.read_text())656667def require_physics_identity(folder: Path, public: dict[str, Any]) -> None:68    if (69        digest(folder / "physics.npz") != public["physics_sha256"]70        or digest(folder / "metadata.json") != public["physics_metadata_sha256"]71    ):72        raise ValueError("physical arrays or provenance changed after acquisition")737475def sources(config: Path, extra: dict[str, Path]) -> dict[str, Path]:76    result = experiment_sources(__file__, config, extra=extra)77    for name in ("_study.py", "_resolution.py", "probe_acquisition.py"):78        result[f"helpers/{name}"] = STUDY / name79    result["helpers/spectral-run.py"] = Path(LIB.__file__)80    return result818283def trajectories(config: dict[str, Any]) -> dict[str, list[dict[str, Any]]]:84    rows = {}85    for role, count in (86        ("fitting", config["fitting_views_per_ring"]),87        ("withheld", config["withheld_views_per_ring"]),88    ):89        offset = config["withheld_yaw_offset_degrees"] if role == "withheld" else 0.090        rows[role] = [91            {92                "tilt_degrees": tilt,93                "yaw_degrees": offset + j * 360.0 / count,94                "pose": asdict(pose_at(tilt, offset + j * 360.0 / count)),95            }96            for tilt in config["tilts_degrees"]97            for j in range(count)98        ]99    return rows100101102def pose_from(row: dict[str, Any]) -> RigidTransform:103    return RigidTransform(**{key: tuple(value) for key, value in row["pose"].items()})104105106def poisson(107    means: np.ndarray, config: dict[str, Any], phase: int, replicate: int, role: int108) -> tuple[np.ndarray, list[list[int]]]:109    values = np.empty_like(means, dtype=np.float32)110    keys = []111    for view in range(means.shape[0]):112        for channel in range(means.shape[1]):113            key = [config["noise_root_seed"], phase, replicate, role, view, channel]114            sample = np.random.Generator(np.random.PCG64(np.random.SeedSequence(key))).poisson(115                means[view, channel].astype(np.float64)116            )117            if np.any(sample > 2**24):118                raise ValueError("Poisson count exceeds exact FP32 integer representation")119            values[view, channel] = sample120            keys.append(key)121    return values, keys122123124def mean_error(actual: np.ndarray, reference: np.ndarray) -> dict[str, float]:125    reference = reference.astype(np.float64)126    error = (actual.astype(np.float64) - reference) / np.sqrt(reference)127    if not np.isfinite(error).all():128        raise ValueError("nonfinite Poisson-scaled mean error")129    return {130        "rms_poisson_sd": float(np.sqrt(np.mean(error**2))),131        "p95_poisson_sd": float(np.quantile(np.abs(error), 0.95)),132        "maximum_poisson_sd": float(np.max(np.abs(error))),133    }134135136def independent_means(137    fields: np.ndarray,138    grid: GridSpec,139    geometry: DetectorGeometry,140    pose: RigidTransform,141    physics: dict[str, Any],142    pixels: list[list[int]],143) -> np.ndarray:144    flat = fields.astype(np.float64).reshape(2, -1)145    paths = np.array(146        [147            [148                integrate_sampled_field(flat[m], grid, geometry, pose, row, col, validate=False)149                for row, col in pixels150            ]151            for m in range(2)152        ]153    )154    spectrum = physics["weights"].astype(np.float64) * physics["response"].astype(np.float64)155    return spectrum @ np.exp(-physics["mu_mm_inv"].astype(np.float64).T @ paths)156157158def require_sampling(error: dict[str, float], config: dict[str, Any]) -> None:159    if (160        error["p95_poisson_sd"] > config["quadrature_p95_poisson_sd"]161        or error["maximum_poisson_sd"] > config["quadrature_maximum_poisson_sd"]162    ):163        raise ValueError(f"prescribed quadrature failed: {error}")164165166def independent_mapping(fields: np.ndarray, gradient: np.ndarray, step: float) -> float:167    """CPU projection onto the two-material nonnegative unit simplex."""168    original = fields.astype(np.float64).reshape(2, -1)169    values = original - step * gradient.astype(np.float64).reshape(2, -1)170    projected = np.maximum(values, 0.0)171    cap = projected.sum(axis=0) > 1.0172    projected[0, cap] = np.clip(0.5 * (1.0 + values[0, cap] - values[1, cap]), 0.0, 1.0)173    projected[1, cap] = 1.0 - projected[0, cap]174    return float(np.max(np.abs((projected - original) / step)))175176177def acquisition_anatomy(178    folder: Path, config: dict[str, Any]179) -> tuple[dict[str, Any], np.ndarray, GridSpec]:180    """Admit either the original field or its hash-pinned portable coarse derivative."""181    anatomy = json.loads((folder / "anatomy.json").read_text())182    native = LIB.checked_array(folder, "forward-fractions.npy", anatomy["outputs"])183    if anatomy.get("format") == "worked-material-basis-v1":184        if (185            digest(folder / "anatomy.json") != config.get("coarse_anatomy_metadata_sha256")186            or digest(folder / "forward-fractions.npy") != config.get("coarse_fractions_sha256")187            or anatomy["original_forward_fractions_sha256"] != config["forward_fractions_sha256"]188        ):189            raise ValueError("portable anatomy differs from its admitted derivative identity")190        fields, grid = native, grid_from(anatomy["forward_grid"])191        if tuple(grid.shape) != tuple(config["shape_zyx"]):192            raise ValueError("portable material basis has the wrong shape")193    else:194        if digest(folder / "forward-fractions.npy") != config["forward_fractions_sha256"]:195            raise ValueError("native anatomy differs from the prospective input identity")196        fields, grid = coarse_cells(197            native, grid_from(anatomy["forward_grid"]), tuple(config["shape_zyx"])198        )199    if {name: row["sha256"] for name, row in anatomy["source_files"].items()} != config[200        "source_sha256"201    ] or anatomy["native_crop_xyz"] != config["native_crop_xyz"]:202        raise ValueError("anatomical source or crop differs from the prospective identity")203    return anatomy, fields, grid204205206def observation_pins(config: dict[str, Any]) -> dict[str, str]:207    """Admit the exact preserved observations required by the precision repair."""208    if config.get("generation_precision", "float32") != "float32":209        raise ValueError("observation generation must retain the original float32 precision")210    precision = config.get("calculation_precision", "float32")211    if precision not in ("float32", "float64"):212        raise ValueError("calculation precision must be float32 or float64")213    pins = config.get("final_observation_sha256", {})214    names = {215        f"{folder}/rep{replicate}-{role}-counts.npy"216        for replicate in (0, 1)217        for folder, role in (("fitting", "fitting"), ("evaluation", "withheld"))218    }219    if (precision == "float64" or "final_observation_sha256" in config) and (220        not isinstance(pins, dict)221        or set(pins) != names222        or any(223            not isinstance(value, str)224            or len(value) != 64225            or any(character not in "0123456789abcdef" for character in value)226            for value in pins.values()227        )228    ):229        raise ValueError("the precision repair must pin all four original final count arrays")230    return pins231232233def verify_observation_identity(234    folder: Path, config: dict[str, Any], development: bool235) -> dict[str, Any]:236    """Check generation output before either fit; no evaluation values enter fitting."""237    pins = observation_pins(config)238    actual = {name: digest(folder / name) for name in pins}239    if pins and not development and actual != pins:240        raise ValueError("generated counts differ from the preserved final observations")241    return {242        "generation_precision": "float32",243        "count_files_sha256": actual,244        "matches_preserved_final_observations": bool(pins) and not development,245        "development": development,246    }247248249def acquire(args: argparse.Namespace) -> None:250    config = json.loads(args.config.read_text())251    observation_pins(config)252    anatomy, fields, grid = acquisition_anatomy(args.anatomy, config)253    metadata, physics, spec = LIB.load_physics(args.physics)254    spec = replace(spec, precision="float32")255    if metadata.get("model_id") != config["physics_model_id"]:256        raise ValueError("physical inputs do not match the prospective protocol")257    if digest(args.physics / "physics.npz") != config["physics_npz_sha256"]:258        raise ValueError("physics differs from the prospective input identity")259    geometry, rows = detector(config), trajectories(config)260    coverage = {role: check_coverage(grid, geometry, records) for role, records in rows.items()}261    phase = config["development_noise_phase"] if args.development else config["final_noise_phase"]262    public = {263        "protocol": config,264        "protocol_sha256": digest(args.config),265        "grid": asdict(grid),266        "geometry": asdict(geometry),267        "trajectories": rows,268        "metric": fixed_metric(physics),269        "noise_phase": phase,270        "development": args.development,271        "physics_sha256": digest(args.physics / "physics.npz"),272        "physics_metadata_sha256": digest(args.physics / "metadata.json"),273        "coverage": coverage,274        "anatomy_source_files": anatomy["source_files"],275        "assignment": config["representation"],276    }277    extra = {278        "inputs/anatomy.json": args.anatomy / "anatomy.json",279        "inputs/forward-fractions.npy": args.anatomy / "forward-fractions.npy",280        "rights/anatomy.json": args.anatomy / "source-license-manifest.json",281        "inputs/physics.npz": args.physics / "physics.npz",282        "inputs/physics-metadata.json": args.physics / "metadata.json",283        **LIB.physics_preparation_sources(args.physics, metadata),284    }285    if (286        digest(args.anatomy / "source-license-manifest.json")287        != anatomy["implementation"]["source_license_manifest_sha256"]288    ):289        raise ValueError("anatomy rights record changed")290    with RunRecorder(291        private_output(args.output, repository_root(__file__)),292        configuration=public,293        sources=sources(args.config, extra),294    ) as run:295        LIB.write_array(run, "evaluation/reference-fractions.npy", fields)296        run.write_json("fitting.json", public)297        forwards = [298            LIB.Forward(grid, geometry, fields, physics, spec, samples, args.device)299            for samples in (config["fitting_samples"], config["generating_samples"])300        ]301        means_by_role = {}302        checks = []303        pixels = config["independent_pixels_rc"]304        for role, records in rows.items():305            means = []306            for index, row in enumerate(records):307                pose = pose_from(row)308                fitting, generating = [forward.project(pose) for forward in forwards]309                error = mean_error(fitting, generating)310                require_sampling(error, config)311                cpu = independent_means(fields, grid, geometry, pose, physics, pixels)312                selected = np.array(313                    [[generating[c, r, k] for r, k in pixels] for c in range(generating.shape[0])]314                )315                cpu_error = mean_error(selected, cpu)316                require_sampling(cpu_error, config)317                checks.append(318                    {"role": role, "view": index, "quadrature": error, "independent_cpu": cpu_error}319                )320                means.append(generating)321                print(f"qualified {role} view {index + 1}/{len(records)}", flush=True)322            means_by_role[role] = np.stack(means)323            LIB.write_array(run, f"evaluation/{role}-means.npy", means_by_role[role])324        run.write_json(325            "qualification.json", {"passed": True, "checks": checks, "before_any_count_draw": True}326        )327        for rep in config["replicates"]:328            for role_id, (role, means) in enumerate(means_by_role.items()):329                counts, keys = poisson(means, config, phase, rep, role_id)330                prefix = "fitting" if role == "fitting" else "evaluation"331                LIB.write_array(run, f"{prefix}/rep{rep}-{role}-counts.npy", counts)332                run.write_json(f"{prefix}/rep{rep}-{role}-noise.json", {"keys": keys})333        run.write_json(334            "observation-identity.json",335            verify_observation_identity(args.output, config, args.development),336        )337338339class SolveBudgetError(Exception):340    """An accepted state exhausted the prospectively declared soft solve budget."""341342343def make_solver(344    public: dict[str, Any],345    counts: np.ndarray,346    physics: dict[str, Any],347    spec: Any,348    settings: MaterialReconstructionSettings,349    device: str,350    initial: np.ndarray | None = None,351) -> MaterialReconstruction:352    grid, config = grid_from(public["grid"]), public["protocol"]353    if initial is None:354        initial = np.empty((2, *grid.shape), dtype=np.float32)355        initial[0].fill(config["initial_water"])356        initial[1].fill(config["initial_bone"])357    geometry = DetectorGeometry(**{key: tuple(value) for key, value in public["geometry"].items()})358    views = tuple(359        MaterialReconstructionView(360            geometry, pose_from(row), counts[i], physics["weights"], physics["response"]361        )362        for i, row in enumerate(public["trajectories"]["fitting"])363    )364    h = public["metric"]["matrix"]365    return MaterialReconstruction(366        grid=grid,367        views=views,368        coefficients=physics["mu_mm_inv"],369        spectral_spec=spec,370        initial_fractions=initial,371        settings=settings,372        samples_per_ray=config["fitting_samples"],373        device=device,374        material_metric=(h[0][0], h[0][1], h[1][1]),375    )376377378def derivative_check(solver: MaterialReconstruction, config: dict[str, Any]) -> dict[str, Any]:379    """Feasible central differences check the complete fitting objective/VJP."""380    original = solver.fractions_numpy().copy()381    interior = np.empty_like(original)382    interior[0].fill(0.65)383    interior[1].fill(0.2)384    solver.set_fractions(interior)385    solver.evaluate(gradient=True)386    gradient = solver.gradient.numpy().astype(np.float64)387    checks = []388    for seed in (7051, 7052, 7053):389        direction = (390            np.random.default_rng(seed).uniform(-0.1, 0.1, interior.shape).astype(np.float32)391        )392        analytic = float(gradient @ direction.ravel().astype(np.float64))393        differences = []394        for step in (0.02, 0.01, 0.005):395            high = np.ascontiguousarray(interior + step * direction, dtype=np.float32)396            low = np.ascontiguousarray(interior - step * direction, dtype=np.float32)397            solver.set_fractions(high)398            upper = solver.evaluate(gradient=False)399            solver.set_fractions(low)400            lower = solver.evaluate(gradient=False)401            differences.append({"step": step, "finite_difference": (upper - lower) / (2 * step)})402        error = min(abs(x["finite_difference"] - analytic) for x in differences)403        limit = max(404            config["directional_derivative_absolute_tolerance"],405            abs(analytic) * config["directional_derivative_relative_tolerance"],406        )407        checks.append(408            {409                "seed": seed,410                "analytic": analytic,411                "differences": differences,412                "minimum_error": error,413                "limit": limit,414                "passed": error <= limit,415            }416        )417    solver.set_fractions(original)418    if not all(row["passed"] for row in checks):419        raise ValueError(f"full-objective directional derivative check failed: {checks}")420    return {"passed": True, "checks": checks}421422423def independent_derivative_check(424    public: dict[str, Any], counts: np.ndarray, physics: dict[str, Any], spec: Any, device: str425) -> dict[str, Any]:426    """Compare the CUDA VJP with exact CPU path integrals on fixed actual rays.427428    The artificial interior field and directions are derivative test inputs,429    never example anatomy or optimisation initialisation. Counts come only430    from the fitting bundle. No withheld/reference data enter this check.431    """432    config, grid = public["protocol"], grid_from(public["grid"])433    geometry = DetectorGeometry(**{key: tuple(value) for key, value in public["geometry"].items()})434    fields = np.empty((2, *grid.shape), dtype=np.float32)435    fields[0].fill(0.65)436    fields[1].fill(0.2)437    selection = [438        (view, row, col) for view in (0, 17, 34) for row, col in ((16, 16), (24, 24), (32, 32))439    ]440    views = []441    for view, row, col in selection:442        pixel_origin = tuple(443            geometry.origin_mm[i]444            + col * geometry.spacing_mm[0] * geometry.u[i]445            + row * geometry.spacing_mm[1] * geometry.v[i]446            for i in range(3)447        )448        pixel_geometry = DetectorGeometry(449            geometry.source_mm, pixel_origin, geometry.u, geometry.v, geometry.spacing_mm, (1, 1)450        )451        observed = np.ascontiguousarray(counts[view, :, row : row + 1, col : col + 1])452        views.append(453            MaterialReconstructionView(454                pixel_geometry,455                pose_from(public["trajectories"]["fitting"][view]),456                observed,457                physics["weights"],458                physics["response"],459            )460        )461    solver = MaterialReconstruction(462        grid=grid,463        views=tuple(views),464        coefficients=physics["mu_mm_inv"],465        spectral_spec=spec,466        initial_fractions=fields,467        settings=MaterialReconstructionSettings(iterations=1),468        samples_per_ray=config["fitting_samples"],469        device=device,470    )471    solver.evaluate(gradient=True)472    gradient = solver.gradient.numpy().astype(np.float64)473    coefficients = physics["mu_mm_inv"].astype(np.float64)474    spectral_weights = physics["weights"].astype(np.float64) * physics["response"].astype(475        np.float64476    )477    checks = []478    for seed in (8041, 8042, 8043):479        direction = np.random.default_rng(seed).uniform(-0.1, 0.1, fields.shape).astype(np.float32)480        cpu = 0.0481        for view, row, col in selection:482            pose = pose_from(public["trajectories"]["fitting"][view])483            paths = np.array(484                [485                    integrate_sampled_field(486                        fields[m].ravel(), grid, geometry, pose, row, col, validate=False487                    )488                    for m in (0, 1)489                ]490            )491            direction_paths = np.array(492                [493                    integrate_sampled_field(494                        direction[m].ravel(), grid, geometry, pose, row, col, validate=False495                    )496                    for m in (0, 1)497                ]498            )499            contributions = spectral_weights * np.exp(-(paths @ coefficients))500            mean = contributions.sum(axis=1)501            derivative = -(contributions @ (direction_paths @ coefficients))502            cpu += float(np.sum((1.0 - counts[view, :, row, col] / mean) * derivative))503        cuda = float(gradient @ direction.ravel().astype(np.float64))504        limit = max(505            config["directional_derivative_absolute_tolerance"],506            abs(cpu) * config["directional_derivative_relative_tolerance"],507        )508        checks.append(509            {510                "seed": seed,511                "cpu_exact_integral_derivative": cpu,512                "cuda_sampled_derivative": cuda,513                "absolute_error": abs(cuda - cpu),514                "limit": limit,515                "passed": abs(cuda - cpu) <= limit,516            }517        )518    if not all(row["passed"] for row in checks):519        raise ValueError(f"independent CPU derivative comparison failed: {checks}")520    return {521        "passed": True,522        "selected_fitting_rays": selection,523        "checks": checks,524        "scope": "exact CPU trilinear integration plus analytic spectral Poisson derivative",525    }526527528def prepare_metric(solver: MaterialReconstruction, run: RunRecorder, prefix: str) -> None:529    """Prepare an owned metric from this accepted field, outside a solve."""530    from dpt.material_preconditioning import prepare_fisher_scaling531532    scaling = prepare_fisher_scaling(solver)533    solver.set_voxel_metric_scale(scaling.scale)534    run.write_json(f"{prefix}preconditioning.json", scaling.diagnostics)535    LIB.write_array(536        run,537        f"{prefix}voxel-metric-scale.npy",538        solver.voxel_metric_scale.numpy().reshape(solver.grid.shape),539    )540541542def solve_stages(543    solver: MaterialReconstruction,544    run: RunRecorder,545    config: dict[str, Any],546    iterations: int,547    seconds: float,548    refresh_steps: list[int],549) -> dict[str, Any]:550    """Keep one accepted trajectory/gate across prescribed metric refreshes.551552    Each canonical solve starts fresh BB/inertial histories. Refreshes only553    occur after a completed update block; no field or tolerance is changed.554    All stages share one soft time budget and global checkpoint numbering.555    """556    history: list[dict[str, Any]] = []557    stages: list[dict[str, Any]] = []558    started = time.perf_counter()559    boundaries = [step for step in refresh_steps if step < iterations] + [iterations]560    result: dict[str, Any] = {}561    for stage, boundary in enumerate(boundaries):562        offset = len(history)563        if stage:564            if time.perf_counter() - started >= seconds:565                result["termination"] = "wall_time_budget"566                break567            prepare_metric(solver, run, f"stages/after-{offset:04d}/")568            if time.perf_counter() - started >= seconds:569                result["termination"] = "wall_time_budget"570                break571        solver.settings = replace(solver.settings, iterations=boundary - offset)572573        def callback(574            iteration: int, row: dict[str, Any], *, offset: int = offset, stage: int = stage575        ) -> None:576            global_iteration = offset + iteration577            record = {578                **row,579                "iteration": global_iteration,580                "stage": stage,581                "stage_iteration": iteration,582                "wall_seconds": time.perf_counter() - started,583            }584            history.append(record)585            run.write_json(f"history/accepted-{global_iteration:04d}.json", record)586            if global_iteration in config["checkpoint_steps"]:587                LIB.write_array(588                    run,589                    f"checkpoints/fields-{global_iteration:04d}.npy",590                    solver.fractions_numpy(),591                )592            if global_iteration % 25 == 0:593                print(594                    f"accepted {global_iteration}: objective={row['objective']:.8g}, "595                    f"mapping_before={row['gradient_mapping_before']:.8g}",596                    flush=True,597                )598            if time.perf_counter() - started >= seconds:599                raise SolveBudgetError600601        try:602            result = solver.solve(callback)603        except SolveBudgetError:604            result = {605                "termination": "wall_time_budget",606                "accepted_steps": len(history) - offset,607                "history": history[offset:],608            }609        run.write_json(f"stages/stage-{stage:02d}.json", result)610        stages.append(611            {612                "stage": stage,613                "accepted_updates_before": offset,614                "accepted_steps": result["accepted_steps"],615                "termination": result["termination"],616                "metric_refreshed": stage > 0,617                "secant_and_inertial_histories": "fresh",618            }619        )620        if result["termination"] != "iteration_budget":621            break622    # Failed trials remain in the corresponding stage record. The global623    # trajectory contains accepted updates only, with no duplicate boundary.624    result.update(625        accepted_steps=len(history),626        history=history,627        stages=stages,628        solve_seconds=time.perf_counter() - started,629    )630    return result631632633def fit(args: argparse.Namespace) -> None:634    record = complete_record(args.acquisition)635    public = checked_json(args.acquisition, "fitting.json", record)636    config = public["protocol"]637    pins = observation_pins(config)638    if (639        digest(args.config) != public["protocol_sha256"]640        or args.replicate not in config["replicates"]641    ):642        raise ValueError("protocol or replicate differs from frozen acquisition")643    require_physics_identity(args.physics, public)644    qualification = checked_json(args.acquisition, "qualification.json", record)645    if not qualification["passed"]:646        raise ValueError("acquisition qualification did not pass")647    name = f"fitting/rep{args.replicate}-fitting-counts.npy"648    counts = checked_array(args.acquisition, name, record)649    if pins and not public["development"]:650        identity = checked_json(args.acquisition, "observation-identity.json", record)651        if (652            identity["count_files_sha256"] != pins653            or identity["matches_preserved_final_observations"] is not True654            or identity["development"] is not False655            or identity["generation_precision"] != "float32"656            or record["output_sha256"][name] != pins[name]657        ):658            raise ValueError("fitting requires the exact preserved final observations")659    _, physics, spec = LIB.load_physics(args.physics)660    spec = replace(spec, precision=config.get("calculation_precision", "float32"))661    development = public["development"]662    iterations = (663        config["development_accepted_updates"]664        if development665        else config["maximum_accepted_updates"]666    )667    seconds = (668        config["development_soft_solve_seconds"] if development else config["soft_solve_seconds"]669    )670    if args.development_seconds is not None:671        if not development or not 0 < args.development_seconds <= config["soft_solve_seconds"]:672            raise ValueError("a development time override must stay within the final budget")673        seconds = args.development_seconds674    if args.development_iterations is not None:675        if (676            not development677            or not 0 < args.development_iterations <= config["maximum_accepted_updates"]678        ):679            raise ValueError("a development update override must stay within the final budget")680        iterations = args.development_iterations681    step_selection = config.get("step_selection", "geometric")682    acceleration = config.get("acceleration", "none")683    preconditioner = config.get("preconditioner", "none")684    if args.development_step_selection is not None:685        if not development:686            raise ValueError("a step-selection override is allowed only for development")687        step_selection = args.development_step_selection688    if args.development_acceleration is not None:689        if not development:690            raise ValueError("an acceleration override is allowed only for development")691        acceleration = args.development_acceleration692    if args.development_preconditioner is not None:693        if not development:694            raise ValueError("a preconditioner override is allowed only for development")695        preconditioner = args.development_preconditioner696    if preconditioner not in ("none", "fisher"):697        raise ValueError("preconditioner must be none or fisher")698    refresh_steps = config.get("preconditioner_refresh_after", [])699    if (700        not isinstance(refresh_steps, list)701        or any(type(step) is not int or step <= 0 for step in refresh_steps)702        or refresh_steps != sorted(set(refresh_steps))703        or (refresh_steps and preconditioner != "fisher")704    ):705        raise ValueError(706            "metric refreshes require ordered unique positive steps and Fisher scaling"707        )708    settings = MaterialReconstructionSettings(709        iterations=iterations,710        initial_step=config["initial_step"],711        mapping_step=config["mapping_step"],712        regularisation_mm_inverse=config["regularisation_mm_inverse"],713        step_selection=step_selection,714        acceleration=acceleration,715    )716    extra = {717        "inputs/acquisition-run.json": args.acquisition / "run.json",718        "inputs/fitting.json": args.acquisition / "fitting.json",719        "inputs/counts.npy": args.acquisition / name,720        "inputs/qualification.json": args.acquisition / "qualification.json",721        "inputs/physics.npz": args.physics / "physics.npz",722        "inputs/physics-metadata.json": args.physics / "metadata.json",723    }724    if pins:725        extra["inputs/observation-identity.json"] = args.acquisition / "observation-identity.json"726    resumed = None727    if args.resume_development_fit is not None:728        if not development:729            raise ValueError("checkpoint continuation is allowed only for development")730        previous = complete_record(args.resume_development_fit)731        if (732            not previous["configuration"]["development"]733            or previous["configuration"]["replicate"] != args.replicate734            or previous["source_sha256"]["inputs/acquisition-run.json"]735            != digest(args.acquisition / "run.json")736        ):737            raise ValueError("continuation must use the same development acquisition and replicate")738        resumed = {739            "fields": checked_array(args.resume_development_fit, "fields.npy", previous),740            "gate": checked_json(args.resume_development_fit, "stationarity-gate.json", previous),741            "result": checked_json(args.resume_development_fit, "solver-result.json", previous),742        }743        for name in ("run.json", "fields.npy", "stationarity-gate.json", "solver-result.json"):744            extra[f"continuation/{name}"] = args.resume_development_fit / name745    with RunRecorder(746        private_output(args.output, repository_root(__file__)),747        configuration={748            "protocol": config,749            "development": development,750            "replicate": args.replicate,751            "maximum_updates": iterations,752            "soft_seconds": seconds,753            "step_selection": step_selection,754            "acceleration": acceleration,755            "preconditioner": preconditioner,756            "preconditioner_refresh_after": refresh_steps,757            "fitting_reference_access": False,758            "continuation": str(args.resume_development_fit) if resumed is not None else None,759            "secant_history_restart": resumed is not None,760        },761        sources=sources(args.config, extra),762    ) as run:763        run.write_json(764            "independent-derivatives.json",765            independent_derivative_check(public, counts, physics, spec, args.device),766        )767        solver = make_solver(public, counts, physics, spec, settings, args.device)768        run.set_metadata(device=str(solver.device), device_name=solver.device.name)769        run.write_json("derivatives.json", derivative_check(solver, config))770        initial_objective = solver.evaluate(gradient=True)771        initial_mapping, _ = solver.projected_trial(config["mapping_step"], use_metric=False)772        gate = stationarity_gate(773            initial_mapping,774            config["mapping_step"],775            config["relative_mapping_tolerance"],776            config["maximum_mapped_fraction_displacement"],777        )778        continuation = None779        if resumed is not None:780            # Round-off in a repeated GPU reduction cannot alter the admitted781            # original-start threshold. Keep the actual recorded gate exactly.782            if not np.isclose(783                gate["initial_mapping"], resumed["gate"]["initial_mapping"], rtol=1e-6784            ):785                raise ValueError("continued solver does not reproduce the original start mapping")786            gate = resumed["gate"]787            solver.set_fractions(resumed["fields"])788            continuation = {789                "previous_accepted_steps": resumed["result"]["accepted_steps"],790                "start_objective": solver.evaluate(gradient=True),791                "secant_history": "fresh restart at the recorded accepted field",792                "original_gate_retained": True,793            }794        solver.settings = replace(795            settings,796            gradient_mapping_tolerance=gate["absolute_threshold"],797            relative_gradient_mapping_tolerance=0.0,798        )799        run.write_json("stationarity-gate.json", gate)800        if preconditioner == "fisher":801            prepare_metric(solver, run, "")802        LIB.write_array(run, "checkpoints/fields-0000.npy", solver.fractions_numpy())803        started = time.perf_counter()804        result = solve_stages(solver, run, config, iterations, seconds, refresh_steps)805        final_objective = solver.evaluate(gradient=True)806        mapping, _ = solver.projected_trial(config["mapping_step"], use_metric=False)807        result.update(808            initial_objective=initial_objective,809            initial_gradient_mapping=initial_mapping,810            final_objective=final_objective,811            stationarity=stationarity_result(mapping, gate),812            solve_seconds=time.perf_counter() - started,813            reference_or_withheld_data_access=False,814            continuation=continuation,815        )816        LIB.write_array(run, "fields.npy", solver.fractions_numpy())817        LIB.write_array(run, "gradient.npy", solver.gradient.numpy())818        LIB.write_array(run, "fitting-predictions.npy", np.stack(solver.predictions_numpy()))819        run.write_json("solver-result.json", result)820        print(821            json.dumps(822                {823                    key: result[key]824                    for key in ("termination", "accepted_steps", "stationarity", "solve_seconds")825                }826            ),827            flush=True,828        )829830831def evaluate(args: argparse.Namespace) -> dict[str, Any]:832    acquisition = complete_record(args.acquisition)833    fitting = complete_record(args.fit)834    public = checked_json(args.acquisition, "fitting.json", acquisition)835    config = public["protocol"]836    if digest(args.config) != public["protocol_sha256"]:837        raise ValueError("evaluation protocol changed")838    # Freeze and verify actual accepted output before opening the assigned reference.839    fields = checked_array(args.fit, "fields.npy", fitting)840    result = checked_json(args.fit, "solver-result.json", fitting)841    if fitting["source_sha256"]["inputs/acquisition-run.json"] != digest(842        args.acquisition / "run.json"843    ):844        raise ValueError("fit is not bound to this acquisition")845    frozen = {846        name: digest(args.fit / name)847        for name in ("run.json", "fields.npy", "gradient.npy", "solver-result.json")848    }849    gradient = checked_array(args.fit, "gradient.npy", fitting)850    cpu_mapping = independent_mapping(fields, gradient, config["mapping_step"])851    if not np.isclose(cpu_mapping, result["stationarity"]["final_mapping"], rtol=1e-10, atol=1e-7):852        raise ValueError("independent CPU projected mapping disagrees with the final CUDA state")853    reference = checked_array(args.acquisition, "evaluation/reference-fractions.npy", acquisition)854    grid = grid_from(public["grid"])855    geometry = DetectorGeometry(**{key: tuple(value) for key, value in public["geometry"].items()})856    require_physics_identity(args.physics, public)857    _, physics, spec = LIB.load_physics(args.physics)858    extra = {859        "inputs/acquisition-run.json": args.acquisition / "run.json",860        "inputs/fit-run.json": args.fit / "run.json",861        "inputs/fields.npy": args.fit / "fields.npy",862        "inputs/gradient.npy": args.fit / "gradient.npy",863        "inputs/solver-result.json": args.fit / "solver-result.json",864        "inputs/reference.npy": args.acquisition / "evaluation/reference-fractions.npy",865        "inputs/physics.npz": args.physics / "physics.npz",866        "inputs/physics-metadata.json": args.physics / "metadata.json",867    }868    with RunRecorder(869        private_output(args.output, repository_root(__file__)),870        configuration={"protocol": config, "accepted_output_freeze": frozen},871        sources=sources(args.config, extra),872    ) as run:873        forward = LIB.Forward(874            grid, geometry, fields, physics, spec, config["generating_samples"], args.device875        )876        errors = fields.astype(np.float64) - reference.astype(np.float64)877        rmse = np.sqrt(np.mean(errors**2, axis=(1, 2, 3)))878        metrics = {879            "material_rmse_water_bone": rmse.tolist(),880            "stationarity": result["stationarity"],881            "termination": result["termination"],882            "accepted_updates": result["accepted_steps"],883            "replicate": fitting["configuration"]["replicate"],884            "development": public["development"],885            "independent_cpu_mapping": cpu_mapping,886        }887        for role, rows in public["trajectories"].items():888            predictions = np.stack([forward.project(pose_from(row)) for row in rows])889            means = checked_array(args.acquisition, f"evaluation/{role}-means.npy", acquisition)890            metrics[role] = mean_error(predictions, means)891            LIB.write_array(run, f"{role}-predictions.npy", predictions)892            checks = []893            for index, row in enumerate(rows):894                pixels = config["independent_pixels_rc"]895                cpu = independent_means(fields, grid, geometry, pose_from(row), physics, pixels)896                selected = np.array(897                    [898                        [predictions[index, c, r, k] for r, k in pixels]899                        for c in range(predictions.shape[1])900                    ]901                )902                error = mean_error(selected, cpu)903                require_sampling(error, config)904                checks.append({"view": index, **error})905            run.write_json(906                f"{role}-independent-final-field.json", {"passed": True, "checks": checks}907            )908        domain_passed = bool(909            np.isfinite(fields).all()910            and np.all(fields >= 0)911            and np.all(fields.astype(np.float64).sum(axis=0) <= 1 + 2**-24)912        )913        metrics["domain_passed"] = domain_passed914        metrics["numerical_passed"] = bool(915            domain_passed916            and result["stationarity"]["passed"]917            and np.all(rmse <= config["maximum_material_rmse"])918            and metrics["withheld"]["rms_poisson_sd"]919            <= config["maximum_withheld_generating_mean_error_rms_poisson_sd"]920        )921        metrics["accepted"] = metrics["numerical_passed"] and not public["development"]922        metrics["scope"] = (923            "Matched coarse CT-derived assigned composition; ideal primary-only physics, "924            "not patient material validation"925        )926        run.write_json("evaluation.json", metrics)927        LIB.write_array(run, "signed-material-error.npy", errors)928        print(json.dumps(metrics), flush=True)929    return metrics930931932# region book:worked-reconstruction-workflow933def run_worked_example(args: argparse.Namespace) -> dict[str, Any]:934    """Generate once, freeze both fits, then assess every prescribed replicate."""935    config = json.loads(args.config.read_text())936    if config["replicates"] != [0, 1]:937        raise ValueError("the complete worked example requires both prescribed replicates [0, 1]")938    with RunRecorder(939        private_output(args.output, repository_root(__file__)),940        configuration={"protocol": config, "development": args.development},941        sources=sources(args.config, {}),942    ) as run:943        acquisition = args.output / "acquisition"944        acquisition_args = argparse.Namespace(**vars(args))945        acquisition_args.output = acquisition946        acquire(acquisition_args)947        fits = {}948        for replicate in config["replicates"]:949            fit_args = argparse.Namespace(**vars(args))950            fit_args.acquisition = acquisition951            fit_args.replicate = replicate952            fit_args.output = args.output / f"fit-rep{replicate}"953            fit(fit_args)954            fits[replicate] = fit_args.output955        # No reference or withheld assessment is opened until both fits exist.956        outcomes = []957        for replicate in config["replicates"]:958            evaluation_args = argparse.Namespace(**vars(args))959            evaluation_args.acquisition = acquisition960            evaluation_args.fit = fits[replicate]961            evaluation_args.output = args.output / f"evaluation-rep{replicate}"962            outcomes.append(evaluate(evaluation_args))963        child_records = {964            str(path.relative_to(args.output)): digest(path)965            for path in sorted(args.output.glob("*/run.json"))966        }967        summary = {968            "development": args.development,969            "replicates": outcomes,970            "numerical_passed": all(row["numerical_passed"] for row in outcomes),971            "accepted": all(row["accepted"] for row in outcomes),972            "all_fits_frozen_before_reference_evaluation": True,973            "child_records_sha256": child_records,974        }975        run.write_json("summary.json", summary)976    return summary977978979# endregion book:worked-reconstruction-workflow980981982def main() -> None:983    parser = argparse.ArgumentParser(description=__doc__)984    parser.add_argument(985        "command", choices=("all", "acquire", "fit", "evaluate"), nargs="?", default="all"986    )987    parser.add_argument("--config", type=Path, default=Path(__file__).with_name("protocol-v3.json"))988    bundle = (989        Path(__file__).resolve().parents[2]990        / "public/generated/worked-examples/inputs/reconstruction"991    )992    parser.add_argument("--anatomy", type=Path, default=bundle / "anatomy")993    parser.add_argument("--physics", type=Path, default=bundle / "physics")994    parser.add_argument("--acquisition", type=Path)995    parser.add_argument("--fit", type=Path)996    parser.add_argument("--replicate", type=int, default=0)997    parser.add_argument("--development", action="store_true")998    parser.add_argument("--development-seconds", type=float)999    parser.add_argument("--development-iterations", type=int)1000    parser.add_argument("--development-step-selection", choices=("geometric", "bb"))1001    parser.add_argument("--development-acceleration", choices=("none", "inertial"))1002    parser.add_argument("--development-preconditioner", choices=("none", "fisher"))1003    parser.add_argument("--resume-development-fit", type=Path)1004    parser.add_argument("--output", type=Path, required=True)1005    parser.add_argument("--device", default="cuda:0")1006    args = parser.parse_args()1007    required = {1008        "all": (),1009        "acquire": ("anatomy",),1010        "fit": ("acquisition",),1011        "evaluate": ("acquisition", "fit"),1012    }1013    for name in required[args.command]:1014        if getattr(args, name) is None:1015            parser.error(f"{args.command} requires --{name}")1016    result = {"all": run_worked_example, "acquire": acquire, "fit": fit, "evaluate": evaluate}[1017        args.command1018    ](args)1019    if args.command in ("all", "evaluate") and not result["accepted"] and not result["development"]:1020        parser.exit(1021            1,1022            "The retained final outcomes did not pass scientific acceptance; "1023            "inspect the recorded report.\n",1024        )102510261027if __name__ == "__main__":1028    main()1029