python/dpt/transport/recovery_experiments.py

1"""Private, preregistered recovery comparisons using the production CUDA oracle.23Independent layer integrals supply observations and assess answers, never runtime4proposals. The sampled reference for scattering uses a separate seed pool; it is5not fed back into model construction or candidate acceptance.6"""78# This private experiment driver inspects prepared workspaces for execution evidence.9# pyright: reportPrivateUsage=false10from __future__ import annotations1112import argparse13import hashlib14import json15import math16from dataclasses import asdict17from pathlib import Path18from time import perf_counter19from typing import Any2021from dpt.contracts import ContractError, integer22from dpt.experiments import RunRecorder23from dpt.statistics import mean_standard_error24from dpt.stochastic_recovery import StochasticPolicy, recover_expected_signal2526from .derivatives import TransportParameter27from .experiments import AbsorptionExperiment, _device_metadata, _recovery28from .forward import prepare_transport29from .inverse import prepare_transport_inverse30from .model import MaterialGrid, PlanarDetector, TransportSpec31from .rng import HistoryBatch32from .source import ParallelBeam3334# region book:transport-scattering-worked-inputs35_ABSORPTION = 0.336_SCATTERING = 0.137_OBSERVATION = 0.0838_INITIAL_AMPLITUDE = 0.08394041def _scattering_problem() -> tuple[TransportSpec, ParallelBeam]:42    return (43        TransportSpec(44            MaterialGrid((-0.5, -0.5, 0.0), (1.0, 1.0, 1.0), (1, 1, 1), 1),45            PlanarDetector((-2.0, -2.0), (4.0, 4.0), (1, 1), 2.0),46            80.0,47            "declared isotropic scattering cube for recovery validation",48            estimator="continuous-absorption",49        ),50        ParallelBeam((0.0, 0.0, -1.0), (0.0, 0.0)),51    )525354# endregion book:transport-scattering-worked-inputs555657def scattering_policy(sampling_multiplier: int = 1) -> StochasticPolicy:58    """Scale original-history work without changing the objective or its gates."""59    factor = integer(sampling_multiplier, "sampling_multiplier", minimum=1, maximum=32767)60    return StochasticPolicy(61        proposal="quadratic",62        replicates=8,63        initial_batch=4096 * factor,64        maximum_batch=65536 * factor,65        final_validation_batch=65536 * factor,66        unique_history_budget=16_000_000 * factor,67        numerical_gradient_allowance=1e-12,68    )697071def scattering_configuration(sampling_multiplier: int = 1) -> dict[str, Any]:72    policy = scattering_policy(sampling_multiplier)73    spec, source = _scattering_problem()74    return {75        "sampling_multiplier": sampling_multiplier,76        "policy": asdict(policy),77        "workspace_max_histories": policy.maximum_batch,78        "reference_batch": policy.maximum_batch,79        "reference_replicates": 256,80        "physical_spec": asdict(spec),81        "source": asdict(source),82        "material_ids": [0],83        "absorption_mm_inverse": [_ABSORPTION],84        "scattering_mm_inverse": [_SCATTERING],85        "active_parameters": ["log-source-amplitude"],86        "base_density": [1.0],87        "pixel_weights": [1.0],88        "observation": _OBSERVATION,89        "initial_amplitude": _INITIAL_AMPLITUDE,90    }919293def validate_scattering_configuration(configuration: dict[str, Any], multiplier: int) -> None:94    """New references declare the complete model; admit legacy default references."""95    if "physical_spec" not in configuration:96        if multiplier != 1:97            raise ContractError("scaled reference lacks its physical/sampling configuration")98        return99    expected = {"seed": 419003, **scattering_configuration(multiplier)}100    if any(101        json.dumps(configuration.get(key), sort_keys=True) != json.dumps(value, sort_keys=True)102        for key, value in expected.items()103    ):104        raise ContractError("reference physical/sampling configuration differs from recovery")105106107def validate_scattering_reference(reference: dict[str, Any], sampling_multiplier: int) -> None:108    """Require the separately generated reference to use the declared work scale."""109    policy = scattering_policy(sampling_multiplier)110    if (111        reference.get("batch") != policy.maximum_batch112        or len(reference.get("replicates", [])) != 256113        or reference.get("histories_traced") != 256 * policy.maximum_batch114        or reference.get("seed") != 419003115    ):116        raise ContractError("reference does not match the declared independent sampling design")117    mean, standard_error = mean_standard_error(tuple(reference["replicates"]))118    uncertainty = reference["absolute_uncertainty"]119    if (120        not math.isfinite(uncertainty)121        or uncertainty < 0122        or not math.isclose(mean, reference["mean"], rel_tol=1e-12, abs_tol=1e-15)123        or not math.isclose(standard_error, reference["standard_error"], rel_tol=1e-12)124        or not math.isclose(uncertainty, 7 * standard_error, rel_tol=1e-12)125        or mean <= uncertainty126    ):127        raise ContractError("reference statistics or positive mean interval are invalid")128129130def _scattering_run(131    seed: int,132    recorder: RunRecorder,133    *,134    reference: dict[str, Any] | None,135    sampling_multiplier: int = 1,136) -> None:137    policy = scattering_policy(sampling_multiplier)138    if reference is not None:139        validate_scattering_reference(reference, sampling_multiplier)140    from dpt._runtime import prepare_context141142    context = prepare_context()143    wp = context.wp144    # Fixed mathematical cube, not a surrogate tissue or measured acquisition.145    spec, source = _scattering_problem()146    workspace = prepare_transport(147        spec,148        material_ids=wp.array([0], dtype=wp.int32, device=context.device),149        absorption=wp.array([_ABSORPTION], dtype=wp.float64, device=context.device),150        scattering=wp.array([_SCATTERING], dtype=wp.float64, device=context.device),151        max_histories=policy.maximum_batch,152        device=str(context.device),153        stream=context.stream,154    )155    # Observation is declared directly. The independent reference locates its156    # optimum afterwards; no analytic/reference signal enters the oracle.157    observation = _OBSERVATION158    oracle = prepare_transport_inverse(159        workspace,160        source=source,161        parameters=(TransportParameter("log-source-amplitude"),),162        observation=wp.array([observation], dtype=wp.float64, device=context.device),163        pixel_weights=wp.ones(1, dtype=wp.float64, device=context.device),164        base_density=(1.0,),165        local_model=True,166    )167    if reference is None:168        means: list[float] = []169        for replicate in range(256):170            oracle._chart((0.0,))171            oracle._mean(172                HistoryBatch(seed, replicate * policy.maximum_batch, policy.maximum_batch),173                1.0,174                oracle._arrays["mean_a"],175            )176            workspace.check_status()177            means.append(float(oracle._arrays["mean_a"].numpy()[0]))178        mean, standard_error = mean_standard_error(tuple(means))179        recorder.set_metadata(**_device_metadata(workspace))180        recorder.write_json(181            "reference.json",182            {183                "mean": mean,184                "standard_error": standard_error,185                "absolute_uncertainty": 7 * standard_error,186                "uncertainty_method": (187                    "seven replicate standard errors; heuristic, not a finite-sample bound"188                ),189                "seed": seed,190                "batch": policy.maximum_batch,191                "sampling_multiplier": sampling_multiplier,192                "allocated_device_bytes": oracle.allocated_bytes,193                "replicates": means,194                "histories_traced": oracle.histories_traced,195                "purpose": "independent held-out forward mean; never optimiser input",196            },197        )198        return199    start = perf_counter()200    try:201        result = recover_expected_signal(202            oracle, (math.log(_INITIAL_AMPLITUDE),), seed=seed, policy=policy203        )204    except Exception as error:205        snapshot = getattr(error, "recovery_diagnostics", None)206        recorder.write_json(207            "failure-diagnostics.json",208            {209                "error": repr(error),210                "controller": asdict(snapshot) if snapshot is not None else None,211                "histories_submitted_including_replay": oracle.histories_traced,212                "parameter_upload_bytes": oracle.parameter_upload_bytes,213                "scalar_download_bytes": oracle.scalar_download_bytes,214            },215        )216        raise217    elapsed = perf_counter() - start218    amplitude = math.exp(result.parameters[0])219    mu = reference["mean"]220    error = reference["absolute_uncertainty"]221    # Monotone in this positive-signal neighbourhood; evaluate both interval ends222    # and its possible quadratic vertex to enclose the reference gradient.223    means = [mu - error, mu + error]224    vertex = observation / (2 * amplitude)225    if means[0] < vertex < means[1]:226        means.append(vertex)227    gradients = [(amplitude * value - observation) * amplitude * value for value in means]228    recorder.set_metadata(**_device_metadata(workspace))229    recorder.write_json(230        "recovery.json",231        {232            "seed": seed,233            "sampling_multiplier": sampling_multiplier,234            "policy": asdict(policy),235            "result": asdict(result),236            "amplitude": amplitude,237            "independent_optimum_interval": [238                observation / (mu + error),239                observation / (mu - error),240            ],241            "reference_gradient_interval": [min(gradients), max(gradients)],242            "reference_stationarity_pass": max(abs(value) for value in gradients)243            <= policy.gradient_tolerance,244            "histories_traced_including_replay": oracle.histories_traced,245            "allocated_device_bytes": oracle.allocated_bytes,246            "scalar_download_bytes": oracle.scalar_download_bytes,247            "parameter_upload_bytes": oracle.parameter_upload_bytes,248            "elapsed_seconds_contended": elapsed,249            "deterministic_sampling": oracle.deterministic_sampling,250        },251    )252253254def main(script: str) -> None:255    parser = argparse.ArgumentParser(description=__doc__)256    parser.add_argument("--output", type=Path, required=True)257    parser.add_argument(258        "--mode",259        choices=("original", "ablation", "family", "reference", "stochastic"),260        required=True,261    )262    parser.add_argument("--reference", type=Path)263    parser.add_argument("--seed-base", type=int, default=9031001)264    parser.add_argument(265        "--sampling-multiplier",266        type=int,267        default=1,268        help="Scale all scattering batches, reference work and unique-history budget together.",269    )270    parser.add_argument(271        "--repetitions",272        type=int,273        default=32,274        help="Number of consecutive stochastic recovery runs (1-32); final evaluation uses 32.",275    )276    args = parser.parse_args()277    if not 1 <= args.repetitions <= 32:278        parser.error("repetitions must lie between 1 and 32")279    if args.repetitions != 32 and args.mode != "stochastic":280        parser.error("repetitions only applies to stochastic recovery")281    if args.sampling_multiplier != 1 and args.mode not in ("reference", "stochastic"):282        parser.error("sampling-multiplier only applies to scattering reference/recovery")283    try:284        sampling = scattering_configuration(args.sampling_multiplier)285    except ContractError as error:286        parser.error(str(error))287    if not 0 <= args.seed_base <= 2**64 - 1 - (args.repetitions - 1) * 104729:288        parser.error("the complete seed sequence must fit unsigned 64 bits")289    root = Path(script).resolve().parents[2]290    output = args.output.resolve()291    if output.is_relative_to(root):292        raise ValueError("raw execution records must remain outside the authoring tree")293    sources = {str(path.relative_to(root)): path for path in (root / "python/dpt").rglob("*.py")}294    sources.update({str(Path(script).resolve().relative_to(root)): Path(script).resolve()})295    sources.update({name: root / name for name in ("pyproject.toml", "uv.lock")})296    cases: list[tuple[str, AbsorptionExperiment]] = []297    if args.mode == "original":298        cases = [299            (300                "original",301                AbsorptionExperiment(302                    seed=23971,303                    histories=256,304                    recovery_maximum_batch=256,305                    recovery_history_budget=1_000_000,306                    estimator="continuous-absorption",307                    proposal="quadratic",308                    numerical_gradient_allowance=1e-12,309                ),310            )311        ]312    if args.mode == "ablation":313        for estimator in ("analogue", "continuous-absorption"):314            for proposal in ("linear", "quadratic"):315                cases.append(316                    (317                        f"{estimator}-{proposal}",318                        AbsorptionExperiment(319                            seed=23971,320                            histories=256,321                            estimator=estimator,322                            proposal=proposal,323                            numerical_gradient_allowance=1e-12,324                            final_validation_batch=262144,325                        ),326                    )327                )328    if args.mode == "family":329        for depth in (0.1, 1.0, 3.0, 6.0):330            for target in (0.5, 1.25, 2.0):331                for initial in (0.6, 1.0, 1.8):332                    cases.append(333                        (334                            f"depth-{depth}-target-{target}-initial-{initial}",335                            AbsorptionExperiment(336                                layer_optical_depths=(depth,),337                                target_density=target,338                                initial_density=initial,339                                histories=16,340                                replicates=2,341                                recovery_maximum_batch=16,342                                recovery_history_budget=1_000_000,343                                estimator="continuous-absorption",344                                proposal="quadratic",345                                numerical_gradient_allowance=1e-12,346                            ),347                        )348                    )349    for name, config in cases:350        with RunRecorder(output / name, configuration=asdict(config), sources=sources) as recorder:351            _recovery(config, recorder)352        record = json.loads((output / name / "recovery.json").read_text())353        print(354            name,355            record["result"]["reason"],356            record["absolute_density_error"],357            record["independent_true_gradient"],358            flush=True,359        )360    if args.mode == "reference":361        with RunRecorder(362            output / "scattering-reference",363            configuration={"seed": 419003, **sampling},364            sources=sources,365        ) as recorder:366            _scattering_run(367                419003, recorder, reference=None, sampling_multiplier=args.sampling_multiplier368            )369    if args.mode == "stochastic":370        if args.reference is None:371            raise ValueError("stochastic evaluation requires a frozen independent reference")372        reference: dict[str, Any] = json.loads(args.reference.read_text())373        validate_scattering_reference(reference, args.sampling_multiplier)374        reference_manifest_path = args.reference.parent / "run.json"375        reference_manifest = json.loads(reference_manifest_path.read_text())376        validate_scattering_configuration(377            reference_manifest["configuration"], args.sampling_multiplier378        )379        reference_sha = hashlib.sha256(args.reference.read_bytes()).hexdigest()380        if (381            reference_manifest["status"] != "complete"382            or not reference_manifest["sources_unchanged"]383            or not reference_manifest["recorded_files_unchanged"]384            or reference_manifest["output_sha256"].get(args.reference.name) != reference_sha385        ):386            raise ContractError("reference must belong to a completed unchanged run record")387        sources["evaluation-reference.json"] = args.reference388        sources["evaluation-reference-run.json"] = reference_manifest_path389        for i in range(args.repetitions):390            seed = args.seed_base + 104729 * i391            evaluation_config: dict[str, Any] = {392                "seed": seed,393                "evaluation_index": i,394                "repetitions": args.repetitions,395                "reference": reference,396                **sampling,397            }398            with RunRecorder(399                output / f"replicate-{i:02d}", configuration=evaluation_config, sources=sources400            ) as recorder:401                _scattering_run(402                    seed,403                    recorder,404                    reference=reference,405                    sampling_multiplier=args.sampling_multiplier,406                )407            print(f"completed independent evaluation {i + 1}/{args.repetitions}", flush=True)408