NNPD — a small, modular experiment framework

NNPD implements the Gaussian study of the paper Prior Dependence in Neural Ratio Estimation (Salama, Fraser, Chang, Nachman; PAI26; PDF also in docs/NNPD_paper_PAI26.pdf) on top of a small, reusable experiment framework. The framework does not import the Gaussian application or the settings files.

Everything is split into training (what is trained) and analysis (what is computed from the trained models). Start with the four files that define them:

	Settings	Implementation
Training	settings_training.py	gaussian/experiment.py
Analysis	settings_analysis.py	gaussian/analysis.py

Full documentation — the usage guides, the settings and CLI reference, and the generated Python API — is published at https://nabil-2.github.io/NNPD/.

Run it

Install uv, then from this directory:

uv sync
source .venv/bin/activate        # Windows: .venv\Scripts\activate
python run.py plan --profile smoke
python run.py run --profile smoke
python run.py verify --profile smoke

uv sync creates .venv with Python 3.11 or newer and installs NNPD and the development tools at the versions pinned in uv.lock. Instead of activating the environment, you can prefix commands with uv run.

The smoke profile trains the four prior models and exercises the complete analysis with small datasets. It checks the plumbing; it is not a scientific convergence study. The default action is plan, so python run.py cannot inadvertently start a large run. nnpd is the equivalent installed command.

For scientific settings:

python run.py plan --profile default
python run.py train --profile default
python run.py analyze --profile default

train trains the models of settings_training.py. analyze computes the metrics and plots of settings_analysis.py from the latest trained models and never trains. run does both. Rerunning reuses completed artifacts. A failed, uncommitted stage is restarted, not mistaken for a valid checkpoint.

Read the plan before running the full study. The analysis cost of every trained model is checked against the limits in settings_analysis.py before anything is computed. If any model exceeds them, run and analyze refuse to start and name the settings to reduce; nothing is reduced automatically. Truth points are Sobol points by default (1,024 at one inferred parameter, doubling per additional one). paper uses size_adaptive: the exact 25**p grid while it is small enough, Sobol above. See Truth design.

All ordinary choices live in two files

settings_training.py holds everything that defines the trained models: problem, priors, data, model, training and runtime. settings_analysis.py holds everything computed from them: inference, verification, metrics, plots and plotting. Change the analysis settings and run analyze again at any time; the models are reused. Edit these files rather than chasing defaults through implementation files.

Both files define the profiles default, paper, and smoke; --profile selects the same profile in both. default is the full study exactly as defined in the two SETTINGS dictionaries. paper makes the paper-facing choices explicit; it is not a claim of bitwise or figure-for-figure reproduction.

Sweeps: baseline plus one changed knob

"data": {"train_size": Choice([50_000, 100_000, 150_000]), ...},
"training": {"learning_rate": Choice([0.1, 0.01, 0.001]), ...},

There are five configurations, not nine:

Training rows	Learning rate
50,000	0.1
100,000	0.1
150,000	0.1
50,000	0.01
50,000	0.001

The first value is the baseline. Each alternative changes just that knob, and duplicate resolved configurations are removed. This works for nested scalar settings and for atomic list/dictionary options:

"kind": Choice(["mlp", "residual"]),
"hidden": Choice([[64, 64, 64, 64], [128, 128]]),
"infer": Choice([["mean:*"], ["mean:*", "std:1"]]),

An ordinary list is a single literal value, not a sweep. This is essential: [64, 64] means two hidden layers, and ["mean:*", "std:1"] means a single inference-parameter set. Wrap alternatives in Choice(...) to make a sweep. Choice inside an option or inside an ordinary list is deliberately rejected.

The four priors form a separately declared fixed comparison cohort repeated for each configuration. With four priors and one training replica, the five configurations above train twenty models. This is intentional and is not a Cartesian sweep of independent knobs. To make the prior itself an OFAT knob, use "priors": Choice([["uniform"], ["normal"], ["exponential"], ["grid"]]) inside cohort. replicas means independently trained models, not repeated observations in an inference case.

Choice is for training settings only. An analysis is always a single setting: change it and analyze again. Runtime settings must remain fixed within a launch; compare CPU/CUDA or jobs/ddp using separate launches. Every alternative is validated before the first job starts.

Inferred parameter sets

Indices are zero-based. mean:* expands to all observation means; std:1 adds the standard deviation of the second observation dimension. For example, with dimension=2, these alternatives have two versus three inferred parameters, and automatically construct networks with four versus five input features.

Unselected likelihood parameters stay at fixed_mean, fixed_std, or the named values in problem.fixed; they are fixed, not marginalized nuisance parameters. Bounds and prior-family parameters come from problem.parameters, with optional name-specific problem.overrides. Standard deviations must be positive. Invalid combinations, such as std:1 in a one-dimensional baseline, fail during planning. Set a compatible baseline dimension rather than relying on Cartesian combinations.

Structure and extension points

settings_training.py       What is trained: training settings and profiles
settings_analysis.py       What is computed from trained models: analysis settings and profiles
run.py                     Thin Python entry point
nnpd/
  core/
    api.py                 Experiment and Analysis; Product, Metric, Plot, CohortPlot, Context
    settings.py            Explicit Choice, OFAT expansion, settings loading
    runner.py              Planning, execution, restore, integrity verification
    storage.py             Atomic artifact store and array/checkpoint access
    runtime.py             CPU/CUDA and torchrun process management
  nre/
    distributions.py       Simulator/Prior interfaces and bounded independent priors
    data.py                Balanced joint-versus-factorized pairs
    models.py              MLP/residual implementations and factory registry
    training.py            Logit BCE training, prediction and classification
    inference.py           Ensemble log ratios and numerical HLD summaries
  results.py               Read runs, export CSV, compare arbitrary saved metrics
gaussian/
  experiment.py            What is trained: problem, prior, data, model, training
  problem.py               Gaussian likelihood and selected parameter layout
  sampling.py              Training-data builder
  evaluation.py            Shared truth ensembles, candidates, inference products
  diagnostics.py           Pairwise ratios, reweighting, normalization datasets
  metrics.py               Independent metric functions
  plots.py                 Matplotlib figures comparing the priors of a configuration
  analysis.py              What is computed: named products, metrics and plots
  extensions.py            Working custom metric/model metric/plot example
notebooks/                 Three thin executable notebooks; no function definitions
tests/                     Unit, science, integration, notebook and distributed tests

An Experiment, named by "experiment" in the training settings, supplies the problem, prior, data sampler, model builder and trainer. An Analysis, named by "analysis" in the analysis settings, registers products, metrics and plots. A Product builds a reusable dataset once. Each metric or plot declares which products it needs, and receives those artifacts plus the full context: model, simulator, prior, both settings, runtime, run directory and artifact store. There is no fixed number of hooks and no hard-coded metric switch in the runner.

Analysis can be added after training. gaussian/extensions.py adds a median-bias metric sharing the existing inference data, a direct model-parameter-count metric, and a custom observation plot. To compute them for already trained models, change settings_analysis.py and analyze again; nothing is retrained:

"analysis": "gaussian.extensions:ExtendedAnalysis",
"metrics": [..., "median_absolute_bias", "parameter_count"],
"plots": [..., "observation_histogram"],

notebooks/03_extensions.ipynb analyzes trained models again and shows these metrics on a saved run. See Extending the framework for contracts and examples.

Notebooks and Python are the same workflow

Execute notebooks/01_run.ipynb, then 02_inspect.ipynb, then 03_extensions.ipynb. The first trains and analyzes a study; the second reloads saved models/data, calls metrics and plots, and exports CSV; the third analyzes the trained models again and tries new metrics on them. Each starts with PROFILE = "smoke"; set it to another profile in all three to use that one. The second and third notebooks never train.

All of this is also ordinary Python, reading the same two settings files:

from nnpd import execute, plan

jobs = plan("smoke")
run_directories = execute("run", profile="smoke")  # or "train" / "analyze"

CUDA and multiple GPUs

On Linux, the PyTorch that uv sync installs is a CUDA build; for other hardware see the CUDA guide. runtime.device="auto" chooses CUDA when available, otherwise CPU. Explicit "cuda" fails clearly when CUDA is unavailable; it never silently falls back. "cuda:1" selects a single device outside a distributed launch.

There are two simple parallel modes:

jobs assigns different complete model/prior/configuration jobs to workers, including their analysis. For these small networks this is usually the useful first option; each job fits on one device.

# settings_training.py: runtime.parallel="jobs", device="auto" or "cuda"
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 run.py run --profile smoke

ddp trains each model collectively using DistributedDataParallel, then runs that model's analysis on rank zero. Batch size is global, not multiplied by the GPU count. The trainer handles unequal and empty final shards without duplicating or dropping real training examples.

# settings_training.py: runtime.parallel="ddp"
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 run.py run --profile smoke

Use one worker per visible GPU and homogeneous GPU types within a DDP group. CPU distributed runs use Gloo; CUDA DDP uses NCCL. Multi-node launch requires correct torchrun rendezvous settings and a shared output filesystem; it is not covered by the tests. DDP analysis is deliberately not sharded within a model. float16 AMP requires CUDA; bfloat16 is also supported on capable CPUs/GPUs. A stochastic model or different numerical backend need not reproduce identical weights across world sizes.

The distributed tests run two-process Gloo DDP and independent jobs on CPU. Single- and multi-GPU tests are included; they run only where CUDA hardware is available. See Testing and validation.

Outputs, correctness and scope

Run folders contain the resolved training settings, cohort member, training status, the analysis settings and state, source-snapshot references, metrics, plots and relative artifact references. Each configuration folder also holds the figures that compare its priors side by side. All figures are PDFs. Shared datasets and model weights live in a content-addressed store; multiple metrics do not duplicate them. Numeric arrays use .npy with optional memory mapping; metadata uses JSON; models use tensor state dictionaries loaded with weights_only=True. Each profile's output folder holds one experiment: running again with changed settings replaces the earlier results, copies of both settings files are saved there, and the cache keeps what the last three distinct training settings used. See Storage and reproducibility.

The example distinguishes the normalized ensemble likelihood ratio from a posterior. Bias, projected HLD coverage, widths, candidate-resolution checks, analytic-likelihood references, pairwise log-ratio errors, sliced Wasserstein reweighting and evidence-normalization checks are separate metrics. Raw inference scores can be kept for a few diagnostic cases or for every case.

Several choices need explicit interpretation: the exponential sign (the default and paper presets differ), the normal-prior scale, the truth design and the discrete integration measure. Read the Scientific notes before comparing results with the paper. Additional priors, publication-specific panels and scheduler-specific launch policy belong in the application's experiment or analysis, not the framework.

Tests

python -m pytest -q
python -m pytest -q --cov=nnpd --cov=gaussian --cov-report=term-missing
python -m pytest -q -m distributed
python -m pytest -q -m cuda       # executes on suitably provisioned CUDA hardware

The tests include mathematical reference cases, exact OFAT expansion, bounded prior sampling, arbitrary parameter layouts, both networks, shared caching, corruption detection, interrupted-stage recovery, safe reload/relocation, independent non-Gaussian framework use, notebook execution and real distributed processes. They are evidence for the tested implementation—not a guarantee of scientific convergence for arbitrary high-dimensional configurations. See Testing and validation for coverage details and the full-size reference run.