Metadata-Version: 2.4
Name: lazy-tfm
Version: 0.2.0
Summary: Lazy but Accurate z for Yinz: conditional distributions from tabular foundation models
Keywords: foundation model,tabular,in-context learning,probabilistic regression,conditional density estimation,photometric redshift,astronomy
Author: Biprateep Dey
Author-email: Biprateep Dey <biprateep@pitt.edu>
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Astronomy
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Dist: numpy>=1.26
Requires-Dist: scipy>=1.11
Requires-Dist: scikit-learn>=1.4
Requires-Dist: pandas>=2.1
Requires-Dist: matplotlib>=3.8
Requires-Dist: huggingface-hub>=0.24
Requires-Dist: h5py>=3.10
Requires-Dist: tqdm>=4.66
Requires-Dist: lazy-tfm[limix,qp,tabfm,tabicl,tabpfn] ; extra == 'all'
Requires-Dist: torch>=2.9.1 ; extra == 'limix'
Requires-Dist: einops>=0.8 ; extra == 'limix'
Requires-Dist: kditransform>=1.0 ; extra == 'limix'
Requires-Dist: nvtx>=0.2 ; extra == 'limix'
Requires-Dist: qp-prob>=1.1 ; extra == 'qp'
Requires-Dist: tabfm[pytorch]>=1.0.1 ; extra == 'tabfm'
Requires-Dist: safetensors>=0.4.3 ; extra == 'tabfm'
Requires-Dist: tabicl>=2.2.0 ; extra == 'tabicl'
Requires-Dist: tabpfn>=9.0.0 ; extra == 'tabpfn'
Requires-Python: >=3.12
Project-URL: Source Code, https://github.com/biprateep/lazy-tfm
Project-URL: Documentation, https://lazy-tfm.readthedocs.io
Project-URL: Bug Tracker, https://github.com/biprateep/lazy-tfm/issues
Project-URL: Changelog, https://github.com/biprateep/lazy-tfm/blob/main/CHANGELOG.md
Provides-Extra: all
Provides-Extra: limix
Provides-Extra: qp
Provides-Extra: tabfm
Provides-Extra: tabicl
Provides-Extra: tabpfn
Description-Content-Type: text/markdown

<h1 align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/biprateep/lazy-tfm/main/docs/_static/lazy-logo-white.svg">
    <img src="https://raw.githubusercontent.com/biprateep/lazy-tfm/main/docs/_static/lazy-logo.svg" alt="LAZY" width="440">
  </picture>
</h1>

**LAZY: Lazy but Accurate Z\* for Yinz** gives full conditional distributions
of a continuous target from pretrained tabular foundation models.

\*Where Z is whatever you want to predict from tabular features, such as a
redshift, a metallicity, a mass or a yield.

[![PyPI](https://img.shields.io/pypi/v/lazy-tfm)](https://pypi.org/project/lazy-tfm/)
[![Python](https://img.shields.io/pypi/pyversions/lazy-tfm)](https://pypi.org/project/lazy-tfm/)
[![Unit test and code coverage](https://github.com/biprateep/lazy-tfm/actions/workflows/testing-and-coverage.yml/badge.svg)](https://github.com/biprateep/lazy-tfm/actions/workflows/testing-and-coverage.yml)
[![Documentation](https://readthedocs.org/projects/lazy-tfm/badge/?version=latest)](https://lazy-tfm.readthedocs.io)
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://github.com/biprateep/lazy-tfm/blob/main/LICENSE)

LAZY is a unified, scikit-learn-compatible wrapper around several pretrained
tabular foundation models (TabPFN, TabICL, TabFM and LimiX-2). The models are
pretrained and never fine-tuned: we hand them labeled rows as *context*, and
they answer queries in one forward pass, with no training loop, no
hyperparameter search and no per-dataset retraining. This is the sense in
which the package is lazy.

The package grew out of photometric redshifts, so its worked examples and
benchmark tools (the LSST DESC DC1 catalog, the HSC spectroscopic selection
and the Data Challenge metrics) come from photo-z, while the estimators work
on any tabular regression problem.

**Documentation:** [lazy-tfm.readthedocs.io](https://lazy-tfm.readthedocs.io) ·
**Tutorials:** [gallery](https://lazy-tfm.readthedocs.io/en/latest/tutorials/index.html), or [open the first in Colab](https://colab.research.google.com/github/biprateep/lazy-tfm/blob/tutorials/basic_usage.ipynb) ·
**Changelog:** [CHANGELOG.md](https://github.com/biprateep/lazy-tfm/blob/main/CHANGELOG.md) · **Contributing:** [CONTRIBUTING.md](https://github.com/biprateep/lazy-tfm/blob/main/CONTRIBUTING.md)

## Install

```console
pip install lazy-tfm              # grid, metrics, plots, datasets
pip install 'lazy-tfm[tabfm]'     # + the TabFM backend
pip install 'lazy-tfm[tabicl]'    # + the TabICLv2 backend
pip install 'lazy-tfm[tabpfn]'    # + the TabPFN backend (v2 to v3.5)
pip install 'lazy-tfm[limix]'     # + the LimiX-2 backend (its code installs separately)
pip install 'lazy-tfm[qp]'        # + qp interoperability, for RAIL
pip install 'lazy-tfm[all]'       # + everything above
```

The pretrained weights are not bundled with the package. They are fetched
from the Hugging Face Hub at the first `fit` and cached thereafter. TabFM's
classification checkpoint is ~6.6 GB and carries a **non-commercial** license
from Google, and TabICLv2's is ~100 MB (BSD-3-Clause). The TabPFN checkpoints
run from ~41 MB to ~880 MB and are **non-commercial** from Prior Labs, except
`v2`, which is Apache-2.0 with an attribution clause. LimiX-2's is ~1.6 GB and
**non-commercial** from Stable AI, with the attribution "Built with StableAI
LimiX". Every license is quoted in `lazy.CHECKPOINTS`, and
`lazy.download_checkpoint("tabfm")` warms the cache ahead of time.

## Use

```python
import numpy as np

import lazy
from lazy import LazyModel, Grid
from lazy.datasets import fetch_dc1

train, test = fetch_dc1(split=True)        # or fetch_dc1() for both, concatenated
# 391k test rows in all, sorted by redshift, so take a random 20k.
test = test.take(np.random.default_rng(0).choice(len(test), 20_000, replace=False))
X_train, X_test = train.features("mag-color"), test.features("mag-color")
z_test = test.redshift

model = LazyModel()                          # TabPFN-3.5, the default; wants a GPU
model.fit(X_train, train.redshift)

pdfs = model.predict_proba(X_test, Grid.linear(0, 2, 200))
z = model.predict(X_test, method="mode")  # or mean, peak_mean, median
lo, med, hi = model.predict_quantiles(X_test, [0.16, 0.5, 0.84]).T
print(model.evaluate(X_test, z_test, y_grid=lazy.datasets.DC1_GRID, scale="1+y"))  # DC1's photo-z convention
```

Without a GPU, TabPFN is slow and refuses contexts above 5,000 rows. In that
case `LazyModel("tabicl")` is small, BSD-licensed and quick on a laptop, and
`LazyModel("tabpfn", version="v3.5-fast")` is TabPFN's smaller, faster
checkpoint, usable on a CPU for contexts up to 5,000 rows. The
[supported models](https://lazy-tfm.readthedocs.io/en/latest/models/index.html)
page compares them all.

`LazyModel(name, ...)` picks the backend by name, so switching between models
is a change of string. The concrete classes (`TabPFNBarDistribution`,
`LimiXBarDistribution`, `TabICLQuantile`, `TabFMHistogram`) are importable and
identical. The API is scikit-learn's, except that `predict_proba` returns a
density on a grid of the target rather than class probabilities, which is the
natural output of a model that knows its own uncertainty. Since
`get_params`/`set_params`/`clone` work, the models drop into scikit-learn
pipelines and search objects unmodified.

**`y_grid` belongs to the prediction rather than to the fit.** Nothing about
fitting depends on the output binning, since these models place their internal
bins by the distribution of the context targets and the grid only enters at
the final exact rebinning step. Therefore, one fitted model answers on as many
grids as you like without refitting:

```python
import numpy as np

model.predict_proba(X_test)                                  # the native grid
model.predict_proba(X_test, Grid.linear(0, 3, 300))
model.predict_proba(X_test, np.linspace(0.005, 2.995, 300))  # or bin centers
```

With no grid, a model answers on its *native* grid, the one it thinks in (for
TabPFN and LimiX-2, their 5,000 buckets), so nothing is lost to rebinning.
`predict_distribution(X)` returns that native answer itself, with exact
`pdf`/`cdf`/`ppf`/`mean`/`interval`/`rvs` in the vocabulary of scipy.stats and
LSST DESC's qp, and `.to_qp()` for RAIL.

| Backend  | Method                                                           | Versions | Weights  |
| -------- | ---------------------------------------------------------------- | -------- | -------- |
| `tabpfn` | Bucket masses of the bar distribution | `v2` … `v3.5` | 41 MB – 880 MB |
| `limix`  | Bucket masses of LimiX-2's 5,000-bucket head | `v2` | ~1.6 GB |
| `tabicl` | Quantiles of an in-context regression head | `v2` | ~100 MB |
| `tabfm`  | Hierarchy of in-context classifiers over equal-mass bins of the target | `v1.0` | ~6.6 GB  |

**One set of parameters, with one set of defaults, on every model.** The
parameters `n_estimators`, `transforms`, `feature_shuffle`, `bag_size`,
`kv_cache`, `random_state`, `chunk_size`, `softmax_temperature`,
`mixed_precision`, `outlier_threshold` and `y_grid` mean the same thing on
every backend. Of these, `transforms` takes a shared vocabulary (`power`,
`quantile`, `robust`, etc.) plus recipes such as `"limix"`, with `"auto"`
being each model's own recipe, pinned in `lazy`; `bag_size` draws per-member
subsets of the context; and `kv_cache` processes the context once and is on by
default. We translate each parameter to the model's own machinery where it has
it and build it around the model where it does not, we pin every other
upstream setting that changes an answer, and registering a new backend without
these parameters fails. LimiX-2 degrades above ~20,000 context rows and warns
when a larger context arrives without `bag_size`.

Memory is bounded by `chunk_size` on every backend, and chunking is exact: a
query row's answer never depends on which other query rows share its chunk
(bit for bit on TabPFN and TabICL, to float rounding on LimiX-2).

### Finishing the TabFM and LimiX-2 installs

PyPI cannot install everything these two backends need, so after their
extras, one command completes them:

```bash
pip install 'lazy-tfm[tabfm,limix]'
lazy setup
```

TabFM's PyPI release has no KV-cache API, so every chunk of query rows
re-encodes the whole training context, which takes about **13.7 ms per
member-row against 0.53 ms**, roughly **26× the compute** for identical
answers. `lazy setup` installs the repository build, which has the API; on
the release, `fit` warns with a `TabFMPerformanceWarning` rather than silently
taking 26× longer, and `kv_cache=False` accepts the slow path and silences
it. LimiX's code is not on PyPI, so `lazy` downloads it at the commit it was
validated on, into its cache, the first time a LimiX-2 model is fitted;
`lazy setup` does it ahead of time. A checkout of this repository
(`uv sync --extra all`) gets the TabFM build automatically. In your own uv
project, add the build in the same command instead, so that `uv sync` keeps
it:

```bash
uv add 'lazy-tfm[tabfm]' 'tabfm[pytorch] @ git+https://github.com/google-research/tabfm@fbb665569425fd2f490c6576b3af967876fe11ff'
```

Upstream LimiX-2 has no key/value cache, and its answers depend on how the
queries are chunked. The backend runs its network with a ported cache,
context-only preprocessing and a dedicated random generator, which removes
both problems. LimiX's network needs `triton`, so the backend is Linux-only
for now. Built with StableAI LimiX.

All the models write onto any `Grid` you ask for, with any number of bins, any
spacing and any range. (TabFM's ten-class ceiling constrains its internal
hierarchy, never your output grid.)

### Progress bars

Every backend draws a `tqdm` bar while it predicts, counting whatever that
model's loop actually iterates over. For 20,000 DC1 test galaxies on a
1,000-galaxy context (one GB10 GPU), the bars read:

```
TabPFN v3.5: 100%|██████████| 20000/20000 [00:15<00:00, 1279.46row/s, buckets=5000, context=1000]
TabICL v2: 100%|██████████| 20000/20000 [00:10<00:00, 1968.82row/s, context=1000, quantiles=999]
TabFM v1.0 (20,000 rows): 100%|██████████| 33/33 [08:47<00:00, 16.00s/stage, context=34, dither=3/3, level=fine 10/10]
```

TabPFN, LimiX and TabICL answer in chunks of query rows, so their bars count
rows. TabFM is a hierarchy of in-context classifications, each over every
query row, so its bar counts *stages* (`n_dither × (1 + n_coarse_bins)` of
them) and shows which dither and level is running on how many context rows.

With `progress="auto"` (the default), the bar is shown on a terminal or in a
notebook and stays silent when output goes to a file, so batch logs stay
clean. Pass `progress=True` to force it on or `progress=False` to turn it off.
The bar goes to stderr, while `verbose=True` log messages go to stdout.

## The harder benchmark: a biased training set

DC1 hands every training galaxy a redshift. Real spectroscopic samples do not
look like that, since they are bright, incomplete and cut in redshift by which
features fall in the observed window, and that mismatch, rather than the
estimator, is usually what limits a survey's redshifts. `fetch_dc1_biased`
merges both DC1 files and cuts the realistic case out of them:

```python
from lazy.datasets import fetch_dc1_biased

split = fetch_dc1_biased()          # 35,011 biased / 10,000 calibration / 19,383 test
split.biased                        # what a HSC-like campaign would have got
split.calibration                   # representative, yours to repair the model with
split.test                          # representative, held out from all of it
print(split.summary())              # each subset against the parent catalog
```

We first merge and reshuffle both files, and then cut them in two. The first
part is run through the selection, and what it keeps is the training set,
while the second is divided at random into the calibration sample and the test
set. Where the cut falls is solved for rather than chosen, because the
selection keeps a near-fixed fraction of whatever it sees, so `n_train` is the
knob. The calibration sample comes *out* of the hold-out rather than on top of
it, so no galaxy is both given to a model and scored on.

`n_train` cannot reach the DC1 training file's 43,486. The selection keeps
8.66 percent of a representative sample, so all 434,476 galaxies yield at most
37,626, and that with nothing left to test on. Therefore, 35,000 is the
largest round number that still leaves a usable hold-out.

`control=True` adds `split.unbiased`, exactly as many galaxies drawn at random
from the same pool, which is the control that separates the selection from the
sample size. It is drawn last, so asking for it changes nothing else in the
split.

The selection is a port of RAIL's HSC `GridSelection`, in which each galaxy is
kept with the HSC spectroscopic success rate of its (i, g−z) pixel, below a
color-dependent redshift ceiling. It is reproduced galaxy-for-galaxy against
the runs this work reports. It also acts on any photometry, so a catalog of
your own can be biased the same way:

```python
from lazy.selection import grid_selection, selection_summary

keep, diagnostics = grid_selection(catalog.raw, catalog.redshift)
print(selection_summary(keep, catalog.raw["I"], catalog.redshift))
#  i in [16.0, 20.0)  ...  fraction 0.661
#  i in [24.0, 25.3)  ...  fraction 0.0005
```

`diagnostics` carries each galaxy's pixel success rate and redshift ceiling,
which is everything the selection knew, for a method that tries to estimate it
back.

## Which model, exactly

A backbone is a family rather than a single model, so every backend takes a
`version`, and each version is a separately pinned checkpoint:

```python
lazy.list_versions("tabpfn")
# ['v2', 'v2.5', 'v2.6', 'v3', 'v3.5', 'v3.5-fast']

model = LazyModel("tabpfn", version="v2.5")
model.name_          # 'tabpfn:v2.5' -- also the label evaluate() puts on its row
```

Sweeping `version` compares model versions on equal terms, and a results
table says which one produced each row rather than only which family. What
actually answered is recorded on the fitted model, ready to be written out
beside the numbers:

```python
model.fit(X_train, train.redshift).provenance_
# {'backend': 'tabpfn', 'version': 'v2.5',
#  'repo_id': 'Prior-Labs/tabpfn_2_5',
#  'filename': 'tabpfn-v2.5-regressor-v2.5_default.ckpt',
#  'revision': '6c45f3a6d0d07c6c5f62572e04a0c2929de91b8b',
#  'package': 'tabpfn 9.0.0', 'lazy': 'lazy-tfm 0.1.0', 'device': 'cuda',
#  'n_estimators': 8, 'transforms': ['auto', ...], 'kv_cache': True, ...}
```

The provenance holds the weights, the code that read them and the ensemble
that answered, and no local paths, so it means the same thing on another
machine and survives a trip through JSON. The rest of the settings are in
`get_params()`. `CHECKPOINTS` is keyed `"backend:version"` and is the
authority on what each name loads, and `download_checkpoint("tabpfn", "v2.5")`
warms exactly that one.

Versions are not interchangeable. TabPFN `v2` is pretrained for at most 10,000
context rows and `v2.5` for 50,000, against a million for `v3`, and `v2` is
also the only one under a commercial-use license. `lazy.get_checkpoint(name,
version)` carries both notes.

## Caches and offline use

Nothing platform-specific is baked in. The wheel is pure Python
(`py3-none-any`), the lockfile resolves for macOS/Linux/Windows on Python
3.12–3.14, and the device is chosen at run time (`device="auto"`).

| Cache | Default | Override |
| ----- | ------- | -------- |
| Pretrained weights | `~/.cache/huggingface/hub` | `HF_HOME` |
| Benchmark catalogs | `~/.cache/lazy-tfm` | `LAZY_DATA_HOME`, else `XDG_CACHE_HOME` |

For a compute node with no network, warm both caches on a login node first,
then run with `download_checkpoint`/`is_cached` and
`fetch_dc1(..., download_if_missing=False)` so that a missing file fails
immediately instead of hanging. The ~14 MB HSC selection grid is cached beside
the catalogs, and `fetch_hsc_grid` takes the same flag. The checkpoint a model
loads is the one `download_checkpoint` fetches, since the pinned revision is
handed to the backend rather than left to its default.

Memory is bounded by default on every backend via `chunk_size` (8,192 query
rows), and a row's answer does not depend on which rows share its chunk, so
the default costs nothing.

## Package contents

| Module           | Contents                                                                     |
| ---------------- | ---------------------------------------------------------------------------- |
| `lazy.models`    | `LazyModel`, the concrete backends, and the name registry                    |
| `lazy.grid`      | `Grid`: binning, normalization, mass-conserving rebinning            |
| `lazy.metrics`   | LSST DESC PZ Data Challenge point and PDF metrics, and `summarize`           |
| `lazy.selection` | The HSC spectroscopic selection function, for biasing a catalog of your own |
| `lazy.plotting`  | Publication figure style, and the standard diagnostic figures                |
| `lazy.datasets`  | The DC1 catalog (download, checksum, cache), `Catalog.features`, `fetch_dc1_biased` |

## Develop

```console
git clone https://github.com/biprateep/lazy-tfm
cd lazy-tfm
uv sync                  # locked environment, including dev tooling
uv run pre-commit install
uv run pytest
```

The test suite needs neither a GPU nor a checkpoint. See
[CONTRIBUTING.md](https://github.com/biprateep/lazy-tfm/blob/main/CONTRIBUTING.md)
for the checks CI runs, building the docs and making a release.

## Repository layout

This repository is the library alone, that is, the public, installable
package, its tests and its documentation. The production runs and figures
behind the paper live in a separate repository that uses this package through
its public API.

## Citing

A paper is in preparation. Until it appears, cite the repository and the
checkpoint revisions recorded in `lazy.CHECKPOINTS`, since a foundation
model's weights are part of the method. The `provenance_` of a fitted model is
that citation for the run you actually did.

Built from the
[LINCC Frameworks Python Project Template](https://github.com/lincc-frameworks/python-project-template).

## License

MIT, for this code. The pretrained checkpoints carry their own licenses.
TabFM's and LimiX-2's are non-commercial, and LimiX-2's requires the
attribution "Built with StableAI LimiX". The HSC selection grid is
redistributed by DESC's
[rail_astro_tools](https://github.com/LSSTDESC/rail_astro_tools) (MIT) and
derives from HSC PDR2 (Aihara et al. 2019); `lazy.selection` downloads it from
there, at a pinned commit, rather than bundling it.
