API reference

This page lists the main public objects in genriesz.

High-level estimation

genriesz.grr_functional(*, X, Y, m, basis, generator=None, g=None, grad_g=None, inv_grad_g=None, grad2_g=None, riesz_method='grr', riesz_penalty='l2', riesz_lam=0.001, riesz_p_norm=None, M=1, local_poly_degree=1, standardize_for_matching=True, outcome_models='auto', outcome_basis=None, outcome_link=None, outcome_penalty='l2', outcome_lam=0.001, outcome_p_norm=None, cross_fit=True, folds=5, random_state=0, estimators=('ra', 'rw', 'arw', 'tmle'), alpha=0.05, null=0.0, max_iter=500, tol=1e-08, verbose=False)[source]

Estimate a linear functional using generalized Riesz regression.

Parameters:

X (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Regressors and outcome.
Y (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Regressors and outcome.
m (Union[LinearFunctional, Callable]) – Either a LinearFunctional instance (recommended) or a plain callable m(x_row, gamma) -> float.
basis (Union[Basis, Callable]) – Basis used for Riesz regression (and for the outcome regression when outcome_models='shared').
generator (BregmanGenerator | None) – Bregman generator used for GRR. If None, you can pass a generator function via g (and optionally grad_g, inv_grad_g, grad2_g).
riesz_method (Literal['grr', 'nn_matching', 'local_poly_nn_lsif']) –
- “grr” : solve the GRR optimization problem
- ”nn_matching” : NN-matching inverse propensity weights (ATE-only convenience)
- ”local_poly_nn_lsif” : local polynomial NN-LSIF weights (ATE-only convenience)
Matching-based Riesz methods currently require cross_fit=False and do not support TMLE (because they do not provide a function-valued representer that can be evaluated counterfactually).
outcome_link (str | None) – If None, inferred as ‘logit’ for outcomes bounded in [0, 1], else ‘identity’. TMLE likelihood is inferred from this link.
g (Callable | None)
grad_g (Callable | None)
inv_grad_g (Callable | None)
grad2_g (Callable | None)
riesz_penalty (str | None)
riesz_lam (float)
riesz_p_norm (float | None)
M (int)
local_poly_degree (int)
standardize_for_matching (bool)
outcome_models (Literal['shared', 'separate', 'both', 'none', 'auto'])
outcome_basis (Basis | Callable | None)
outcome_penalty (str | None)
outcome_lam (float)
outcome_p_norm (float | None)
cross_fit (bool)
folds (int)
random_state (int | None)
estimators (Sequence[str])
alpha (float)
null (float)
max_iter (int)
tol (float)
verbose (bool)

Return type:

FunctionalEstimate

genriesz.grr_ate(*, X, Y, treatment_index=0, basis, generator=None, **kwargs)[source]

Estimate ATE with the GRR API.

Return type:

FunctionalEstimate

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
Y (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
treatment_index (int)
basis (Basis | Callable)
generator (BregmanGenerator | None)

genriesz.grr_att(*, X, Y, treatment_index=0, basis, generator=None, **kwargs)[source]

Estimate ATT with the GRR API.

Return type:

FunctionalEstimate

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
Y (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
treatment_index (int)
basis (Basis | Callable)
generator (BregmanGenerator | None)

genriesz.grr_did(*, X, Y0, Y1, treatment_index=0, basis, generator=None, **kwargs)[source]

Panel DID implemented as ATT on ΔY = Y1-Y0.

Return type:

FunctionalEstimate

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
Y0 (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
Y1 (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
treatment_index (int)
basis (Basis | Callable)
generator (BregmanGenerator | None)

genriesz.grr_ame(*, X, Y, coordinate=0, basis, generator=None, **kwargs)[source]

Estimate an average marginal effect (average derivative) wrt x_coordinate.

Return type:

FunctionalEstimate

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
Y (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
coordinate (int)
basis (Basis | Callable)
generator (BregmanGenerator | None)

Density ratio and covariate shift

genriesz.fit_density_ratio(X_num, X_den, *, basis=None, n_centers=200, sigma=1.0, standardize=True, generator=None, g=None, grad_g=None, inv_grad_g=None, grad2_g=None, penalty='l2', lam=0.01, p_norm=None, cv=False, folds=5, sigma_grid=None, lam_grid=None, random_state=0, max_iter=500, tol=1e-08, verbose=False)[source]

Estimate a density ratio under covariate shift.

Parameters:

X_num (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Samples from p (numerator) and q (denominator), respectively.
X_den (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Samples from p (numerator) and q (denominator), respectively.
basis (Union[BaseBasis, CallableBasis, Callable, None]) – Feature map phi(X) used for the linear predictor. If None, a Gaussian-kernel RKHS basis is used.
n_centers (int) – Parameters of the default Gaussian RKHS basis.
sigma (float | None) – Parameters of the default Gaussian RKHS basis.
standardize (bool) – Parameters of the default Gaussian RKHS basis.
generator (BregmanGenerator | str | None) – Either a BregmanGenerator instance or one of the built-in names: 'sq', 'ukl', 'bkl', 'bp', 'pu'. If None and g is also None, defaults to 'sq'.
g (Optional[Callable]) – Custom generator specification (same conventions as grr_functional()).
grad_g (Optional[Callable]) – Custom generator specification (same conventions as grr_functional()).
inv_grad_g (Optional[Callable]) – Custom generator specification (same conventions as grr_functional()).
grad2_g (Optional[Callable]) – Custom generator specification (same conventions as grr_functional()).
penalty (str | None) – Regularization on beta.
lam (float) – Regularization on beta.
p_norm (float | None) – Regularization on beta.
cv (bool) – If cv=True, choose (sigma, lam) by cross validation. This is currently implemented only for the default Gaussian RKHS basis.
folds (int) – If cv=True, choose (sigma, lam) by cross validation. This is currently implemented only for the default Gaussian RKHS basis.
sigma_grid (Optional[Iterable[float]]) – If cv=True, choose (sigma, lam) by cross validation. This is currently implemented only for the default Gaussian RKHS basis.
lam_grid (Optional[Iterable[float]]) – If cv=True, choose (sigma, lam) by cross validation. This is currently implemented only for the default Gaussian RKHS basis.
max_iter (int) – L-BFGS-B controls for the general (non-squared) case.
tol (float) – L-BFGS-B controls for the general (non-squared) case.
verbose (bool) – L-BFGS-B controls for the general (non-squared) case.
random_state (int | None)

Returns:

A fitted model with predict_ratio().

Return type:

DensityRatioResult

class genriesz.DensityRatioResult(basis, generator, beta, penalty, lam, p_norm, centers=None, sigma=None, standardize=None, class_prior_ratio=None)[source]

Result of fit_density_ratio().

Parameters:

basis (BaseBasis | CallableBasis)
generator (BregmanGenerator)
beta (ndarray[tuple[Any, ...], dtype[float64]])
penalty (str | None)
lam (float)
p_norm (float)
centers (ndarray[tuple[Any, ...], dtype[float64]] | None)
sigma (float | None)
standardize (bool | None)
class_prior_ratio (float | None)

basis: Fitted basis object used for the linear predictor.

generator: Bregman generator defining the loss and the link.

beta: Coefficients of the linear predictor v(x) = phi(x)^T beta.

penalty, lam, p_norm: Regularization specification for beta.

centers, sigma, standardize: Convenience metadata when the default Gaussian RKHS basis is used. These are None when a custom basis is provided.

predict_v(X)[source]

Predict the linear score v(x) = phi(x)^T beta.

Return type:: ndarray[tuple[Any, ...], dtype[float64]]
Parameters:: X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

predict_ratio(X, *, clip_nonnegative=True)[source]

Predict the density ratio r_hat(x).

Parameters:

X (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Points at which to evaluate the ratio.
clip_nonnegative (bool) – If True, clip predictions at 0. This is often used for squared-loss density ratio estimators.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Functionals

class genriesz.LinearFunctional(name)[source]

Base class for linear functionals used by GRR.

Parameters:: name (str)

m_from_function(X, *, predict, derivative=None)[source]

Apply the functional to a generic function.

This is mostly used for TMLE updates where we need m(X, alpha_hat).

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
predict (Callable[[ndarray[tuple[Any, ...], dtype[float64]]], ndarray[tuple[Any, ...], dtype[float64]]])
derivative (Callable[[ndarray[tuple[Any, ...], dtype[float64]], int], ndarray[tuple[Any, ...], dtype[float64]]] | None)

class genriesz.CallableFunctional(m, *, name='custom')[source]

Wrap a Python callable m(x, gamma) as a LinearFunctional.

This adapter is meant for quick experiments and README-style usage. The callable must implement a linear map in gamma.

Parameters:

m (Callable[[ndarray[tuple[Any, ...], dtype[float64]], Callable[[ndarray[tuple[Any, ...], dtype[float64]]], float]], float]) –
A callable of the form m(x_row, gamma) -> float where:
- x_row is a 1D regressor vector,
- gamma is a callable that maps a 1D regressor vector to a scalar.
The callable is assumed to be linear in the function argument gamma. If it is not, the GRR objective will generally be invalid.
name (str) – Name used in results.

Notes

To solve GRR in a finite-dimensional model alpha(x)=phi(x)^T beta, the solver needs the matrix M such that m(X_i, phi(\cdot)^T beta)=M_i^T beta. For a callable functional we construct M by repeatedly applying m to each basis coordinate function.

This is flexible but can be slower than implementing a dedicated LinearFunctional.

class genriesz.ATEFunctional(treatment_index=0)[source]

Average treatment effect: E[ gamma(1,Z) - gamma(0,Z) ].

Parameters:: treatment_index (int)

class genriesz.ATTFunctional(*, treatment_index=0, pi)[source]

Average treatment effect on the treated.

theta = E[ Y(1) - Y(0) | D=1 ]: = E[ D * (gamma(1,Z) - gamma(0,Z)) ] / E[D].

We treat this as a plug-in linear functional given a fixed value of pi = E[D] (estimated from the sample in the wrapper).

Parameters:

treatment_index (int)
pi (float)

class genriesz.DIDFunctional(*, treatment_index=0, pi)[source]

Difference-in-differences as ATT on delta outcomes.

In the notebooks we treat DID as an ATT estimand on the panel difference

ΔY = Y_post - Y_pre.

The functional form is identical to ATT (with the same pi), but we keep a separate name for clarity.

Parameters:

treatment_index (int)
pi (float)

class genriesz.AMEFunctional(coordinate=0)[source]

Average marginal effect (average derivative) of gamma wrt x_k.

Parameters:: coordinate (int)

m_from_function(X, *, predict, derivative=None)[source]

Apply the functional to a generic function.

This is mostly used for TMLE updates where we need m(X, alpha_hat).

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
predict (Callable[[ndarray[tuple[Any, ...], dtype[float64]]], ndarray[tuple[Any, ...], dtype[float64]]])
derivative (Callable[[ndarray[tuple[Any, ...], dtype[float64]], int], ndarray[tuple[Any, ...], dtype[float64]]] | None)

Low-level solvers

class genriesz.GRRGLM(*, basis, generator, functional, penalty='l2', lam=0.001, p_norm=None)[source]

Finite-dimensional generalized Riesz regression (GLM form).

Parameters:

basis (Basis)
generator (BregmanGenerator)
functional (LinearFunctional)
penalty (str | None)
lam (float)
p_norm (float | None)

derivative_alpha(X, coordinate)[source]

Derivative of alpha(x) wrt x_coordinate.

Uses the identity grad_g(alpha(x)) = v(x) and the inverse function theorem:

g’’(alpha) * d alpha/dx = d v/dx.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
coordinate (int)

Bases

The abstract base class all bases inherit from:

class genriesz.BaseBasis[source]: Convenience base class implementing copy and a no-op fit.

Built-in bases:

class genriesz.PolynomialBasis(degree=2, *, include_bias=True)[source]

Full polynomial features up to a given total degree.

This class intentionally implements polynomial features without relying on scikit-learn so that the core package can depend only on NumPy and SciPy. The enumeration of monomials is deterministic and ordered by total degree, then lexicographic within each total degree.

Derivatives are implemented analytically via the monomial exponent table.

Parameters:

degree (int)
include_bias (bool)

derivative(X, coordinate)[source]

Derivative of the feature map wrt X[:, coordinate].

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
coordinate (int)

class genriesz.TreatmentInteractionBasis(*, base_basis, treatment_index=0)[source]

Interaction basis for binary-treatment functionals.

Given a base basis on covariates Z (excluding the treatment), this basis maps X = [D, Z] to

phi(X) = [ D * psi(Z) , (1 - D) * psi(Z) ].

This is a convenient default for ATE/ATT/DID-style functionals.

Parameters:

base_basis (BaseBasis)
treatment_index (int)

class genriesz.RBFRandomFourierBasis(*, n_features=500, sigma=1.0, include_bias=True, standardize=True, random_state=None)[source]

RBF random Fourier features (Rahimi-Recht) with optional standardization.

Parameters:

n_features (int)
sigma (float)
include_bias (bool)
standardize (bool)
random_state (int | None)

derivative(X, coordinate)[source]

Derivative of the feature map wrt X[:, coordinate].

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
coordinate (int)

d/dx_k [ sqrt(2/D) cos(W[:,j]^T (x-mu)/std + b_j) ]: = -sqrt(2/D) * W[k,j]/std[k] * sin(W[:,j]^T (x-mu)/std + b_j)

class genriesz.GaussianRKHSBasis(*, n_centers=300, sigma=1.0, include_bias=True, standardize=True, random_state=None, centers=None)[source]

Gaussian-kernel RKHS basis using kernel evaluations at fitted center points.

The feature map is

phi_j(x) = K(x, c_j),

where {c_j} are fitted centers (one basis function per center).

This is a convenient default for RKHS regression in moderate dimensions.

Parameters:

n_centers (int)
sigma (float)
include_bias (bool)
standardize (bool)
random_state (int | None)
centers (ArrayLike | None)

class genriesz.RBFNystromBasis(*, n_centers=300, sigma=1.0, include_bias=True, standardize=True, random_state=None, centers=None, jitter=1e-08)[source]

Nyström feature map for the RBF kernel.

This class computes a Nyström approximation to the implicit RBF-RKHS feature map. Compared to GaussianRKHSBasis, Nyström features apply an additional whitening transform based on the kernel matrix among centers.

The feature map is

Phi(x) = K(x, C) (K(C, C) + jitter I)^{-1/2}.

Notes

This is an O(m^3) preprocessing step in the number of centers m.
For large m, consider RBFRandomFourierBasis.

Parameters:

n_centers (int)
sigma (float)
include_bias (bool)
standardize (bool)
random_state (int | None)
centers (ArrayLike | None)
jitter (float)

class genriesz.KNNCatchmentBasis(*, n_neighbors=1, include_bias=False, standardize=True, random_state=None)[source]

kNN nearest-neighbor indicator basis.

After fitting on a set of centers, evaluating on query points returns a (dense) indicator matrix whose entry (i, j) is 1 if center j is among the k nearest neighbors of query i, and 0 otherwise.

Notes

This is mainly intended for small-to-medium center sets used in notebooks.
For large-scale matching, prefer the dedicated NN/LSIF utilities.

Parameters:

n_neighbors (int)
include_bias (bool)
standardize (bool)
random_state (int | None)

class genriesz.CallableBasis(func, *, derivative=None)[source]

Wrap a Python callable as a basis.

This is useful when you want to define a custom feature map without creating a full class.

Parameters:

func (Callable[[Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]], Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]]) – A callable returning a feature matrix. It may accept either a 2D array X or a 1D array x.
derivative (Optional[Callable[[Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]], int], Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]]]) – Optional callable implementing the derivative feature map required by AME-type functionals.

Notes

We infer n_features during fit (or the first call).

Optional bases (scikit-learn)

Requires pip install "genriesz[sklearn]".

class genriesz.sklearn_basis.RandomForestLeafBasis(model, include_bias=False, normalize=True)[source]

Leaf encodings from a fitted RandomForest.

Parameters:

model (object) – A scikit-learn estimator with a fit method and an apply method (e.g., sklearn.ensemble.RandomForestRegressor).
include_bias (bool) – If True, prepend a constant-1 column.
normalize (bool) – If True (default), divide each leaf encoding by sqrt(n_estimators) so that the row L2-norm stays O(1) regardless of forest size. Without this, the norm grows as sqrt(T) and makes the effective regularisation scale T-dependent.

n_output_: Alias for the number of output features (including bias if enabled). This attribute is provided for compatibility with older example code.

Optional bases (PyTorch)

Requires pip install "genriesz[torch]".

Generators

class genriesz.BregmanGenerator(*, g=None, grad=None, inv_grad=None, grad2=None, name='Custom', C=0.0, branch_fn=None, finite_diff_eps=1e-06, newton_max_iter=60, newton_tol=1e-10)[source]

A Bregman generator with an (optional) automatic link.

Parameters:

g (Optional[Callable]) – Generator function g(x, alpha) or g(alpha).
grad (Optional[Callable]) – First derivative wrt alpha, ∂g(x, alpha)/∂alpha. If omitted, it is approximated by finite differences.
inv_grad (Optional[Callable]) – Inverse derivative (link) alpha = (∂g)^{-1}(x, v). If omitted, it is computed by Newton iterations using grad and grad2.
grad2 (Optional[Callable]) – Second derivative wrt alpha (elementwise). If omitted, it is approximated from g via a second-order finite difference.
name (str) – Display name.
C (float) – Optional domain parameter used by some generator families. The generic implementation uses it only as a soft domain guard.
branch_fn (Optional[Callable[[ndarray[tuple[Any, ...], dtype[float64]]], int]]) – Optional branch selector returning 1 (positive) or 0 (negative). Built-in UKL/BP generators use this to choose the sign branch.
finite_diff_eps (float)
newton_max_iter (int)
newton_tol (float)

Notes

The generic (user-specified) generator supports regressor-dependent generators via g(x, alpha). For performance and numerical stability, providing analytic grad and especially inv_grad is strongly recommended.

as_generator()[source]

For API compatibility with earlier drafts.

Return type:: BregmanGenerator

g(X, alpha)[source]

Evaluate g(x, alpha) row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad(X, alpha)[source]

Evaluate the derivative ∂g/∂alpha row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad2(X, alpha)[source]

Evaluate the second derivative ∂²g/∂alpha² row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

inv_grad(X, v)[source]

Inverse derivative map alpha = (∂g)^{-1}(x, v).

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
v (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

conjugate(X, v)[source]

Return (g*(v), alpha) evaluated row-wise.

Return type:

tuple[ndarray[tuple[Any, ...], dtype[float64]], ndarray[tuple[Any, ...], dtype[float64]]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
v (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

class genriesz.SquaredGenerator(C=0.0)[source]

Squared generator (SQ-Riesz).

g(alpha) = (alpha - C)^2.

This generator has no strict domain constraints and induces a linear link

alpha = C + 0.5 * v.

Parameters:: C (float)

inv_grad(X, v)[source]

Inverse derivative map alpha = (∂g)^{-1}(x, v).

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
v (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

g(X, alpha)[source]

Evaluate g(x, alpha) row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad(X, alpha)[source]

Evaluate the derivative ∂g/∂alpha row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad2(X, alpha)[source]

Evaluate the second derivative ∂²g/∂alpha² row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

class genriesz.UKLGenerator(C=1.0, *, branch_fn=None)[source]

Unnormalized KL generator (UKL-Riesz).

The inverse gradient is branch-wise:

positive branch (sign +1): alpha = C + exp(v)
negative branch (sign -1): alpha = -C - exp(-v)

If branch_fn is provided, it determines which branch is used for each observation.

Parameters:

C (float)
branch_fn (BranchFn | None)

inv_grad(X, v)[source]

Branch-wise inverse gradient alpha = (g’)^{-1}(v).

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
v (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

g(X, alpha)[source]

Evaluate g(x, alpha) row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad(X, alpha)[source]

Evaluate the derivative ∂g/∂alpha row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad2(X, alpha)[source]

Evaluate the second derivative ∂²g/∂alpha² row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

class genriesz.BKLGenerator(C=1.0, *, branch_fn=None)[source]

Binary KL generator (BKL-Riesz).

The generator is

g(alpha) = (|alpha| - C) log(|alpha| - C) - (|alpha| + C) log(|alpha| + C),

with domain |alpha| > C and C > 0.

Its derivative is

g’(alpha) = sign(alpha) * log( (|alpha|-C) / (|alpha|+C) ).

The inverse gradient is branch-wise. Let s be the desired sign branch (+1 or -1) and let u = s * v. Since the log-ratio is always negative, the theoretical domain is u < 0. In finite samples (and especially under cross fitting), u may violate this. Instead of raising an exception, we clip u to a small negative value to keep the link well-defined.

If branch_fn is provided, it selects the sign branch.

Parameters:

C (float)
branch_fn (BranchFn | None)

inv_grad(X, v)[source]

Inverse derivative map alpha = (∂g)^{-1}(x, v).

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
v (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

g(X, alpha)[source]

Evaluate g(x, alpha) row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad(X, alpha)[source]

Evaluate the derivative ∂g/∂alpha row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad2(X, alpha)[source]

Evaluate the second derivative ∂²g/∂alpha² row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

class genriesz.BPGenerator(C=1.0, *, omega=0.5, power=None, branch_fn=None)[source]

Box-Power generator (BP-Riesz).

A smooth family interpolating between UKL-like (small power) and squared-like (power near 1) behavior.

We use the parametrization, for t = |alpha| - C,

g(alpha) = ( t^{1+omega} - (1+omega) t ) / omega,

with domain |alpha| > C and omega > 0.

The derivative is

g’(alpha) = sign(alpha) * (1+omega)/omega * ( t^omega - 1 ),

so the inverse gradient (branch-wise) can be written as

k = (1+omega)/omega u = 1 + sign * v / k (must be > 0) alpha = sign * ( C + u^{1/omega} ).

As with UKL, branch_fn can be supplied to select the sign.

Parameters:

C (float)
omega (float)
power (float | None)
branch_fn (BranchFn | None)

inv_grad(X, v)[source]

Branch-wise inverse gradient map for BP.

The theoretical domain restriction is t = 1 + sign*v/k > 0. In finite samples (and especially under cross fitting) the linear predictor can violate this constraint. Instead of raising an exception, we clip t to a small positive value.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
v (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

g(X, alpha)[source]

Evaluate g(x, alpha) row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad(X, alpha)[source]

Evaluate the derivative ∂g/∂alpha row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad2(X, alpha)[source]

Evaluate the second derivative ∂²g/∂alpha² row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

class genriesz.PUGenerator(C=1.0, *, branch_fn=None)[source]

PU generator (PU-Riesz).

This generator is based on the binary-entropy potential

g(alpha) = C * [ |alpha| log|alpha| + (1-|alpha|) log(1-|alpha|) ],

with domain |alpha| in (0, 1) and C > 0.

The derivative is

g’(alpha) = sign(alpha) * C * log( |alpha| / (1-|alpha|) ).

The inverse gradient is a (scaled) logistic map.

Notes

This generator is primarily useful when you want the representer to be bounded (in absolute value) by 1.

Parameters:

C (float)
branch_fn (BranchFn | None)

inv_grad(X, v)[source]

Inverse derivative map alpha = (∂g)^{-1}(x, v).

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
v (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

g(X, alpha)[source]

Evaluate g(x, alpha) row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad(X, alpha)[source]

Evaluate the derivative ∂g/∂alpha row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

grad2(X, alpha)[source]

Evaluate the second derivative ∂²g/∂alpha² row-wise.

Return type:

ndarray[tuple[Any, ...], dtype[float64]]

Parameters:

X (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])
alpha (_Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str])

Matching helpers

genriesz.nn_matching_inverse_propensity_weights(X, D, M, *, standardize=True, metric='euclidean', algorithm='auto', n_jobs=None)[source]

Compute NN-matching inverse-propensity weights for ATE.

This reproduces the weights used by bias-corrected kNN matching:

w_hat(1, X_i) = 1 + K1M(i) w_hat(0, X_i) = 1 + K0M(i)

where K1M(i) (resp. K0M(i)) is the number of times treated (resp. control) unit i is used as one of the M nearest neighbors for a unit in the opposite group, divided by M.

Parameters:

X (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Covariates of shape (n, d). Do not include the treatment indicator.
D (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Treatment indicator of shape (n,), with values 0/1.
M (int) – Number of nearest neighbors.
standardize (bool) – If True, standardize X column-wise using the full sample (recommended).
metric (str) – Parameters for the kNN search. Without scikit-learn, only metric="euclidean" is supported and algorithm is ignored. n_jobs is passed to SciPy’s cKDTree as workers.
algorithm (str) – Parameters for the kNN search. Without scikit-learn, only metric="euclidean" is supported and algorithm is ignored. n_jobs is passed to SciPy’s cKDTree as workers.
n_jobs (int | None) – Parameters for the kNN search. Without scikit-learn, only metric="euclidean" is supported and algorithm is ignored. n_jobs is passed to SciPy’s cKDTree as workers.

Returns:

Contains full-sample weights w (aligned with rows of X), plus group-specific pieces.

Return type:

NNMatchingWeights

Notes

This matches the logic in Lin et al.’s R code:

Index1 = knnx.index(X1, X0, M) K1M = tabulate(Index1) / M

and similarly for controls.

genriesz.local_polynomial_nn_lsif_density_ratio(numerator, denominator, eval_points, M, *, degree=0, kernel='ball', exclude_self=False, ridge=1e-08, metric='euclidean', algorithm='auto', n_jobs=None, verbose=False)[source]

Local-polynomial NN–LSIF density-ratio estimation.

Implements equations (8) in Section 6.3:

beta_hat(x) = H_hat(x)^{-1} h_hat(x) r_hat(x) = e0^T beta_hat(x)

with the M-NN ball kernel.

Parameters:

numerator (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Sample from f(1), shape (N1, d).
denominator (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Sample from f(0), shape (N0, d).
eval_points (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Points x at which to estimate r0(x) = f(1)(x)/f(0)(x), shape (n_eval, d).
M (int) – Neighborhood size.
degree (int) – Polynomial degree p.
kernel (Literal['ball']) – Currently only “ball” (M-NN ball kernel) is implemented.
exclude_self (bool) – If True, then when an eval point equals a denominator sample point, we attempt to exclude that identical point when defining the M-th neighbor radius. This is mostly relevant for in-sample evaluation.
ridge (float) – A small ridge term added to H_hat(x) for numerical stability.
metric (str) – Parameters for the kNN search. Without scikit-learn, only metric="euclidean" is supported and algorithm is ignored. n_jobs is passed to SciPy’s cKDTree as workers.
algorithm (str) – Parameters for the kNN search. Without scikit-learn, only metric="euclidean" is supported and algorithm is ignored. n_jobs is passed to SciPy’s cKDTree as workers.
n_jobs (int | None) – Parameters for the kNN search. Without scikit-learn, only metric="euclidean" is supported and algorithm is ignored. n_jobs is passed to SciPy’s cKDTree as workers.
verbose (bool) – If True, prints progress every ~500 eval points.

Returns:

Array of shape (n_eval,).

Return type:

r_hat

Notes

This implementation is intentionally explicit (a loop over eval points) to keep the logic faithful to the definition. For large-scale simulations you may want to vectorize and/or parallelize externally.

genriesz.local_polynomial_nn_lsif_inverse_propensity_weights(X, D, M, *, degree=1, standardize=True, kernel='ball', exclude_self=True, ridge=1e-08, clip_min=1e-08, metric='euclidean', algorithm='auto', n_jobs=None, verbose=False)[source]

Local-polynomial NN–LSIF inverse-propensity weight estimation for ATE.

This uses the density ratio representation of the ATE Riesz representer:

numerator sample: f(X) (all units)
denominator sample: f(X | D=d)

For each d in {0,1}, local-polynomial NN–LSIF estimates f(x)/f(x | D=d).

Parameters:

X (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Covariates of shape (n, d). Do not include the treatment indicator.
D (Union[_Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Treatment indicator of shape (n,), with values 0/1.
M (int) – Neighborhood size.
degree (int) – Polynomial degree p (p=0: local constant, p=1: local linear, …).
standardize (bool) – If True, standardize X using the full sample.
clip_min (float | None) – Optional lower bound for weights. If not None, we set w = max(w, clip_min). This is a pragmatic safeguard because unconstrained LSIF may yield negative estimates.
kernel (Literal['ball'])
exclude_self (bool)
ridge (float)
metric (str)
algorithm (str)
n_jobs (int | None)
verbose (bool)

Returns:

w aligns with the original sample order.

Return type:

LocalPolynomialLSIFWeights

Results

class genriesz.FunctionalEstimate(n, alpha, null, estimand, estimates, diagnostics)[source]

Container for estimates for the same estimand.

Parameters:

n (int)
alpha (float)
null (float)
estimand (str)
estimates (dict[str, SingleEstimate])
diagnostics (dict[str, object])

love_plot_data(*, as_pandas=None)[source]

Return covariate-balance diagnostics used by love_plot().

This method is available only when the underlying estimand is of treatment-effect type (ATE / ATT / DID). The returned object contains standardized mean differences (SMDs) before and after weighting.

Parameters:: as_pandas (bool | None) – If True, return a pandas.DataFrame. If False, return a list of row dicts. If None, return a DataFrame when pandas is available.

love_plot(*, threshold=0.1, max_covariates=30, sort_by='weighted', absolute=True, ax=None)[source]

Create a Love plot (covariate-balance plot).

The plot compares standardized mean differences (SMDs) before vs. after weighting induced by the estimated Riesz representer.

Notes

Requires matplotlib (not a hard dependency).
Only available for ATE / ATT / DID.

Parameters:

threshold (float)
max_covariates (int | None)
sort_by (str)
absolute (bool)

class genriesz.SingleEstimate(name, estimate, se, ci_low, ci_high, p_value)[source]

A single point estimate with asymptotic (normal) inference.

Parameters:

name (str)
estimate (float)
se (float)
ci_low (float)
ci_high (float)
p_value (float)