Getting started

Installation

From PyPI:

pip install genriesz

From a local checkout (editable install):

python -m pip install -U pip
pip install -e .

Optional extras:

# scikit-learn integrations (tree-based feature maps)
pip install "genriesz[sklearn]"

# PyTorch integration (neural network feature maps)
pip install "genriesz[torch]"

Notes

  • scikit-learn is optional and only required for genriesz.sklearn_basis.

  • PyTorch is optional and only required for genriesz.torch_basis.

Quickstart: ATE

The following example assumes your regressor matrix is X = [D, Z] where D is the (0/1) treatment and Z are covariates.

import numpy as np
from genriesz import (
    grr_ate,
    PolynomialBasis,
    TreatmentInteractionBasis,
    UKLGenerator,
)

# Synthetic data: X = [D, Z]
n, d_z = 2000, 5
rng = np.random.default_rng(0)
Z = rng.normal(size=(n, d_z))
D = (rng.normal(size=n) > 0).astype(float)
Y = 2.0 * D + Z[:, 0] + rng.normal(size=n)

X = np.column_stack([D, Z])

# Basis on Z, then interact with D (ATE-friendly)
psi = PolynomialBasis(degree=2, include_bias=True)
phi = TreatmentInteractionBasis(base_basis=psi)

# Note: you do not need to call phi.fit(X) manually.
# grr_ate and grr_functional will copy and fit the basis inside each training fold.

# UKL generator induces the link function automatically
gen = UKLGenerator(C=1.0, branch_fn=lambda x: int(x[0] == 1.0)).as_generator()

res = grr_ate(
    X=X,
    Y=Y,
    basis=phi,
    generator=gen,
    cross_fit=True,
    folds=5,
    estimators=("ra", "rw", "arw", "tmle"),
    riesz_penalty="l2",
    riesz_lam=1e-3,
)

print(res.summary_text())

Next steps

  • See User guide for details on bases, generators, estimators, and cross fitting.

  • See diagnostics for guidance on interpreting alpha statistics, covariate balance, and Love plots.

  • See Examples for runnable scripts and notebooks.

  • See API reference for the full API reference.