Metadata-Version: 2.4
Name: priorforge
Version: 0.1.0
Summary: Reusable data-generating-process (DGP) priors for Prior-Fitted Network (PFN) training
Project-URL: Homepage, https://github.com/goku2130/priorforge
Project-URL: Repository, https://github.com/goku2130/priorforge
Project-URL: Bug Tracker, https://github.com/goku2130/priorforge/issues
Author-email: Gokul <gokul6625@gmail.com>
License: MIT
License-File: LICENSE
Keywords: bayesian,distributions,machine learning,pfn,prior-fitted networks
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: numpy>=1.24
Requires-Dist: scipy>=1.10
Requires-Dist: torch>=2.12
Provides-Extra: dev
Requires-Dist: ipykernel; extra == 'dev'
Requires-Dist: jupyter; extra == 'dev'
Requires-Dist: matplotlib; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: scikit-learn>=1.3; extra == 'dev'
Description-Content-Type: text/markdown

# PriorForge

[![CI](https://github.com/goku2130/priorforge/actions/workflows/ci.yml/badge.svg)](https://github.com/goku2130/priorforge/actions/workflows/ci.yml)

Reusable **data-generating process (DGP) priors** for Prior-Fitted Network (PFN) training.

Every PFN paper hand-crafts its own DGP from scratch — the structural generator
that turns sampled latent structure into a dataset `(X, y)`. That generator is
the hard, paper-defining part: an SCM (TabPFN), a Gaussian process (PFNs4BO), an
Error-Trend-Seasonality process (ForecastPFN), or a covariate→treatment→outcome
chain (CausalPFN). PriorForge ships these generators as composable components, so
you `import` and `train` instead of rebuilding the prior each time.

## Installation

```bash
pip install priorforge
```

For GPU support with a specific CUDA version, install PyTorch separately first:

```bash
# CUDA 13.0
pip install torch --index-url https://download.pytorch.org/whl/cu130
pip install priorforge
```

## Quick Start

A DGP emits synthetic datasets; `PriorDataLoader` turns them into the
`(context, query)` splits a PFN trains on.

```python
import torch
from priorforge import SCMPrior, PriorDataLoader, BarDistribution

dgp    = SCMPrior(n_features=5)                 # TabPFN-style structural prior
loader = PriorDataLoader(dgp, n=128, batch=32)  # context/query batches
bar    = BarDistribution(n_bins=1000)           # output head

batch = loader.sample_batch()
ctx_x, ctx_y, qry_x, qry_y = batch              # raw (X, y) context — no summary stats
# logits     = model(ctx_x, ctx_y, qry_x)
# boundaries = bar.fit_boundaries(ctx_y[0])
# loss       = bar.nll(logits, qry_y, boundaries)
```

## Generators

Each generator targets a distinct research community. All share the `DGP`
interface: `sample_dataset(n) -> (X, y)`, plus `latents()` and
`prior_predictive()` for diagnostics.

| Generator | Paper lineage | Use case |
|---|---|---|
| `SCMPrior` | TabPFN | tabular classification / regression (sparse layered DAG) |
| `MLPSCMPrior` | TabPFN-v2 / GraphPFN | tabular (dense random-MLP neural SCM) |
| `GPPrior` | PFNs4BO | Bayesian optimization, regression w/ UQ |
| `ETSPrior` | ForecastPFN | zero-shot time-series forecasting |
| `CausalDGP` | CausalPFN | treatment-effect estimation (back-door) |
| `FrontDoorDGP` | CausalFM | effect estimation via the front-door criterion |
| `IVDGP` | CausalFM | effect estimation via an instrumental variable |
| `MixtureDGP` | Mitra | mixed synthetic priors (a prior over the priors above) |

```python
from priorforge import GPPrior, ETSPrior, CausalDGP, MixtureDGP, SCMPrior, MLPSCMPrior

GPPrior(n_features=2, input_warping=True)       # ARD Matérn-3/2 + linear + warping
ETSPrior(periods=(7.0, 365.25))                 # trend × seasonality × Weibull noise
CausalDGP(n_covariates=5, confounding=1.0)      # propensity → treatment → outcome

# Mitra finding: a curated mixture generalises better than any single prior
MixtureDGP([SCMPrior(n_features=8), MLPSCMPrior(n_features=8)], weights=[0.5, 0.5])
```

For the causal generators, `latents()` carries the ground truth needed to train
and score an effect-estimation PFN: `CausalDGP` exposes `y0`, `y1`, `cate`,
`ate`, `propensity`; `FrontDoorDGP` adds the (hidden) `confounder`/`mediator`
and a closed-form `ate` with the biased `naive_diff`; `IVDGP` exposes the
structural `tau`, the biased `naive_ols`, and the `first_stage_r2` instrument
strength.

## Output heads

```python
import torch
from priorforge import BarDistribution, MartingalePosterior

bar  = BarDistribution(n_bins=1000)             # discretized predictive head
post = MartingalePosterior(bar)                 # posterior over functionals
draws  = post.quantity(logits, boundaries, fn=torch.mean)   # predictive resampling
lo, hi = post.credible_interval(draws, alpha=0.9)
```

## Training a PFN

A reference in-context regressor ships with the library, so you can go from a
generator to a trained model in one call — and see the prior actually teach a
transformer to predict held-out points.

```python
import torch
from priorforge import SCMPrior, train_pfn

dgp    = SCMPrior(n_features=5)
result = train_pfn(dgp, steps=300, batch=16, n_bins=100, generator=torch.Generator().manual_seed(0))
print(result.first_loss, result.last_loss)      # BarDistribution NLL drops over training
```

`PFNTransformer` consumes raw `(context_x, context_y, query_x)` with the
standard PFN attention mask (queries never see other queries) and predicts over
a `BarDistribution` head. Swap `SCMPrior` for any regression generator
(`GPPrior`, `MLPSCMPrior`, `MixtureDGP`, …). See `examples/train_pfn.py`.
Training auto-selects a GPU when one is available (`device=None`); pass
`device="cpu"` to force CPU.

### NanoTabPFN — a known-outcome reproduction

A faithful reimplementation of [NanoTabPFN](https://arxiv.org/abs/2511.03634)
(Pfefferle et al., 2025) ships as a classification example. Its 150×5×2
SCM-style training prior is exactly `SCMPrior(task="classification")`, so the
whole loop lives in the library:

```python
from priorforge import SCMPrior, train_nanotabpfn

dgp    = SCMPrior(n_features=5, task="classification", n_classes=2)
result = train_nanotabpfn(dgp, steps=500)   # 356,066 params — matches upstream
# held-out query accuracy climbs ~0.5 -> ~0.8
```

The model uses TabPFN-v2's two-axis attention (feature attention across columns,
datapoint attention across rows, with test rows attending to train rows only).
Two documented deviations from upstream — `AdamW` instead of
`schedulefree.AdamWScheduleFree`, plus an explicit LR warmup that ScheduleFree
would otherwise provide (without it this post-norm transformer collapses to
chance). See `examples/train_nanotabpfn.py`.

#### Zero-shot transfer to real datasets

The payoff demo. `train_nanotabpfn_multi` pretrains one model across *varied*
SCM geometries (feature counts 2–12, binary tasks), then `nanotabpfn_predict`
runs a **single forward pass per dataset** — no gradient steps, no per-dataset
fitting — on real sklearn tables:

```python
from priorforge import train_nanotabpfn_multi, nanotabpfn_predict

model = train_nanotabpfn_multi(steps=2000, feature_range=(2, 12),
                               class_range=(2, 2), n_classes=2).model
pred  = nanotabpfn_predict(model, train_x, train_y, test_x, n_classes=2)
```

Trained *only* on synthetic priors (it never sees real data), it lands ~2 pts
behind the best **fitted** logistic-regression / KNN / decision-tree /
random-forest baseline on breast-cancer, iris, wine and digits — competitive
in-context with no fitting. See `examples/eval_real_datasets.py`.

**Prior breadth matters.** `train_nanotabpfn_multi` takes a `dgp_factory`
``(k_feat, k_cls) -> DGP`` hook, so you can broaden the pretraining prior in one
line — e.g. mix the sparse-DAG `SCMPrior`, the dense `MLPSCMPrior`, and per-step
`MixtureDGP` blends. Broadening the pretraining prior **halves** the zero-shot
gap to the best fitted model — direct evidence that prior *design*, not the
training loop, is the lever.

**Faster training via prefetch.** Profiling showed CPU data generation is ~66%
of each step (the host→device copy is negligible), so the win is overlapping it
with GPU compute, not porting the generator to CUDA. Pass `num_workers=N` to
generate batches in background processes:

```python
train_nanotabpfn_multi(steps=3000, feature_range=(2, 12), num_workers=2, amp=True)
# num_workers: ~3.2x (8.4 -> 26.8 steps/s), now GPU-bound. Default 0 keeps the
#              serial, reproducible path.
# amp=True:    bf16 autocast, a further ~15% (no accuracy change).
```

(`torch.compile` was tried and dropped — for this 356k model with per-step
varying feature widths it gave no steady-state speedup over eager bf16 while
adding heavy per-width compilation cost.)

## Verification

Three levels of checks:

```python
from priorforge import verify_dgp, prior_predictive_fidelity, verify_distribution

verify_dgp(SCMPrior(n_features=5), verbose=True)     # DGP contract (shapes, task, repro)
prior_predictive_fidelity(dgp, reference_y)          # does the prior cover real data?
verify_distribution(LogNormal(), params)             # leaf-level 5-level protocol
```

## Distributions (observation/noise leaves)

The distribution library is a supporting layer — the noise/observation leaves
used *inside* a generator.

| Kind | Classes |
|---|---|
| Primitives | `LogNormal`, `NegativeBinomial`, `BetaBin` |
| Wrappers | `ZeroInflated(base)` |
| Combinators | `Mixture([d1, d2], weights)` |
| Aliases | `ZILN`, `ZINB` |

Distributions inheriting `CausalDistribution` expose `sample_with_effect` for
matched (control, treatment) pairs — a leaf used inside a causal DGP:

```python
from priorforge import ZILN
ctrl, trt = ZILN().sample_with_effect(params, N_control=500, N_treatment=500, tau=0.2)
```

## Design

- **Structural generators are the product**; the distribution catalog is a
  supporting layer. See [REDESIGN.md](REDESIGN.md).
- **Raw-context path**: PFNs consume raw `(X, y)` context, not summary
  statistics. Sufficient statistics are an opt-in leaf transform.
- **PyTorch only**: the Triton two-backend plan is deferred (REDESIGN §7).
- **Reproducible**: every generator threads a `torch.Generator`.

## License

MIT
