hdfe-stream 0.1.0__py3-none-any.whl

hdfe_stream/__init__.py

@@ -0,0 +1,138 @@
+"""
+hdfe_stream: out-of-core OLS with several high-dimensional fixed effects, and
+Poisson, logit and probit (`fepois_stream`, `feglm_stream`) by iteratively
+reweighted least squares over the same passes (see glm.py).
+
+    from hdfe_stream import feols_stream
+    fit = feols_stream("log_earn ~ age_squared + age_cubed | worker_id + firm_id + year",
+                       "data/*.parquet", workdir="hdfe_work", vcov={"CRV1": "worker_id"})
+    fit.summary()
+
+Formulas follow pyfixest syntax and are parsed with pyfixest's own parser:
+covariates, interactions (`:` and `*`), `I()`, `log()`, `C()`, `i()`,
+several dependent variables (`y1 + y2 ~ ...`), `sw()` / `csw()` stepwise
+syntax, fixed-effect interactions (`firm_id^year`), and IV
+(`y ~ x | fe | endog ~ z`). Every covariate term is compiled to a Polars
+expression, so constructing the design never leaves the streaming engine.
+Weighted least squares via `weights=` (aweights or fweights). The
+lower-level `StreamingHDFE` class takes column names or Polars expressions
+directly.
+
+Reporting: `result.to_pyfixest()` returns a pyfixest Feols/Feiv view for
+pf.etable, pf.summary, pf.coefplot and pf.iplot; `hdfe_stream.etable(...)`
+does the conversion for you and accepts ordinary pyfixest models too.
+
+Model:  y = X b + sum_d  gamma_d[level_d(i)] + e     for d = 0 ... D-1
+
+One fixed-effect dimension (`fe[0]` internally) is *streamed*: it is never
+represented as a vector in memory. It should be the one with the most levels
+whose groups each touch only a few levels of the other dimensions (workers,
+not firms). It is chosen as: the `stream=` option if given; else the
+dimension with varying slopes; else the highest approximate cardinality
+(HyperLogLog counts gathered during pass 0's first scan), ties going to the
+dimension listed first. Every other dimension is represented by level-sized
+vectors in RAM.
+
+Varying slopes (FEIS) on the streamed dimension use fixest syntax:
+`| worker_id[t] + firm_id` gives worker effects plus worker-specific slopes on
+t (several: `worker_id[t, t2]`). Projecting out the streamed dimension
+then means removing each worker's own weighted regression on [1, slopes],
+which is still local to the worker.
+
+Fewer than two dimensions. With one, it is streamed and there is nothing
+left to solve for: step 2 is skipped and the fit is the within regression.
+With none, there is no dimension to group the rows by, so pass 0 writes them
+unpartitioned and unsorted, passes 1 and 1b and step 2 are skipped, and the
+row passes read plain batches without demeaning; the model keeps an
+"Intercept" coefficient, as in pyfixest.
+
+Memory: rows are only ever streamed. The arrays held in RAM are sized by the
+non-streamed dimensions' level counts (times a block of at most `rhs_block`
+variables at a time), plus, for the explicit solver, the sparse reduced
+matrix S, whose size depends on how levels co-occur and not on rows. Polars
+steps stream; the sort and cell group_by run one fe[0] hash bucket at a time.
+
+Pipeline
+--------
+Pass 0 (Polars, streaming)   evaluate the design as Polars expressions; drop
+                             rows with missing / non-finite values; factorize
+                             the non-streamed FE and cluster ids;
+                             hash-partition rows into fe[0] buckets; sort each
+                             bucket and assign dense fe[0] codes.
+Pass 1 (Polars, streaming)   per bucket: group_by(fe[0], other codes) -> cell
+                             table with n and sums (plus within-cell
+                             cross-products when assembly="cells").
+Pass 1b (streamed + numba)   keep cells of fe[0] groups with more than one
+                             cell as memory-mapped arrays; Jacobi diagonal.
+Step 2 (solver, per block)   solve  S Gamma = D_o' M_0 V,  S = D_o' M_0 D_o,
+                             for all variables V in blocks of `rhs_block`
+                             columns: "explicit" (sparse S built once, block
+                             PCG in memory), "stream_cg" (S applied by
+                             streaming the cells), or "within".
+Step 3 (numba + BLAS)        assemble V' M_D V (from cells, or from a row
+                             pass when there are many variables: numba
+                             residualizes each chunk in place, BLAS forms its
+                             cross-products); drop collinear covariates per
+                             model; beta.
+Step 4 (streamed, numba +    per model: stream rows for fe[0] effects,
+        BLAS)                residuals and RSS (numba), and the bread and the
+                             meat for HC1 / CRV1 (BLAS, chunk by chunk).
+
+Layout
+------
+api.py             `feols_stream`, `fepois_stream`, `feglm_stream`
+estimator.py       `StreamingHDFE`: options, layout, the driver
+glm.py             `StreamingGLM`: IRLS over the same passes
+families.py        GLM families: Poisson, logit, probit
+passes.py          pass 0/1/1b   (mixin)
+solve.py           step 2        (mixin)
+inference.py       steps 3/4     (mixin)
+kernels_base.py    numba kernels
+kernels_slopes.py  numba kernels, varying slopes
+kernels_glm.py     numba kernels, GLM cell sums and group effects
+formula.py         pyfixest formula -> Polars expressions (input side)
+reporting.py       pyfixest etable/coefplot views (output side)
+feterms.py         'worker_id[t]' / 'firm_id^year' parsing
+results.py         `HDFEResult`, `HDFEMulti`
+workspace.py       run directories, `cleanup()`
+report.py          logging / printing
+utils.py           small shared helpers
+"""
+
+from .api import feglm_stream, fepois_stream, feols_stream
+from .estimator import StreamingHDFE
+from .glm import StreamingGLM
+from .leaveout_se import ComponentSE
+from .leaveout_weakid import WeakIdDiagnostics
+from .leaveout import leave_one_out_connected, leave_out_kss
+from .reporting import etable
+from .results import HDFEMulti, HDFEResult
+from .utils import iter_group_chunks
+from .workspace import cleanup
+
+__all__ = [
+    "feols_stream",
+    "fepois_stream",
+    "feglm_stream",
+    "StreamingHDFE",
+    "StreamingGLM",
+    "HDFEResult",
+    "HDFEMulti",
+    "etable",
+    "leave_out_kss",
+    "ComponentSE",
+    "WeakIdDiagnostics",
+    "leave_one_out_connected",
+    "cleanup",
+    "iter_group_chunks",
+]
+
+def _detect_version():
+    try:
+        from importlib.metadata import version
+        return version("hdfe-stream")
+    except Exception:                 # source tree without an installed dist
+        return "0.1.0"
+
+
+__version__ = _detect_version()

hdfe_stream/_types.py

@@ -0,0 +1,18 @@
+"""Type aliases shared by the public signatures."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Mapping, Sequence, Union
+
+import polars as pl
+
+PathLike = Union[str, Path]
+# a Parquet path or glob, or a Polars LazyFrame
+Source = Union[PathLike, pl.LazyFrame]
+# 'iid', 'hetero', 'HC1', 'CRV1:var', or {'CRV1': 'var'}
+Vcov = Union[str, Mapping[str, str]]
+# column name, (name, expression), or a mapping / sequence of them
+Variables = Union[
+    str, tuple, Sequence[Union[str, tuple]], Mapping[str, pl.Expr]]
+Options = Any

hdfe_stream/am_interval.py

@@ -0,0 +1,119 @@
+"""The Andrews-Mikusheva interval for a weakly identified quadratic form.
+
+Pure numerics, no data: given the leading eigenvalue lambda_1 and the joint
+estimate (b1-hat, theta1-hat) with its 2x2 covariance, return the q = 1
+confidence interval of Kline, Saggio and Solvsten (2020), Section 6.1:
+
+    C = [ min, max ] of  lambda_1 b^2 + t  over the ellipse
+        { (b1-hat - b, theta1-hat - t) Sigma^-1 (.)' <= z_{alpha,kappa}^2 }
+
+Projecting a 2-d ellipse through a quadratic map would ordinarily need the
+chi-squared(2) critical value to control size; Andrews and Mikusheva (2016)
+show the curvature of the map lets it come down, toward the chi-squared(1) value
+when the map is nearly linear. KSS's curvature for q = 1 is
+
+    kappa = 2 |lambda_1| V[b1] / ( V[theta1]^(1/2) (1 - rho^2)^(1/2) )
+
+and z_{alpha,kappa} is the (1 - alpha) quantile of
+
+    rho(kappa) = sqrt(chi_1^2 + (chi_1' + 1/kappa)^2) - 1/kappa
+
+for independent chi variates (Appendix C.6.1). Saggio's reference tabulates that
+quantile by simulation; for q = 1 it has an exact one-dimensional integral form,
+used here. With a, b independent half-normals and c = 1/kappa,
+
+    rho <= z  <=>  a^2 + (b + c)^2 <= (z + c)^2
+    P(rho <= z) = int_0^z 2 phi(b) [2 Phi(sqrt((z - b)(z + b + 2c))) - 1] db
+
+written in factored form so nothing cancels as kappa -> 0. The limits are the
+chi-squared(1) quantile at kappa = 0 and chi-squared(2) as kappa -> infinity.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def am_cdf(z: float, kappa: float) -> float:
+    """P(rho(kappa) <= z) for q = 1."""
+    from scipy import integrate, stats
+
+    if z <= 0:
+        return 0.0
+    if kappa <= 0:
+        return float(2.0 * stats.norm.cdf(z) - 1.0)
+    c = 1.0 / kappa
+
+    def integrand(b):
+        inner = np.sqrt(max((z - b) * (z + b + 2.0 * c), 0.0))
+        return 2.0 * stats.norm.pdf(b) * (2.0 * stats.norm.cdf(inner) - 1.0)
+
+    value, _err = integrate.quad(integrand, 0.0, z, epsabs=1e-13, epsrel=1e-12,
+                                 limit=200)
+    return float(value)
+
+
+def am_critical_value(kappa: float, alpha: float = 0.05) -> float:
+    """z_{alpha, kappa} for q = 1: between sqrt(chi2_1) and sqrt(chi2_2)."""
+    from scipy import optimize, stats
+
+    lo = float(np.sqrt(stats.chi2.ppf(1.0 - alpha, 1)))
+    if kappa <= 0:
+        return lo
+    hi = float(np.sqrt(stats.chi2.ppf(1.0 - alpha, 2)))
+    target = 1.0 - alpha
+    # the quantile rises with kappa from lo to hi; bracket just outside both
+    return float(optimize.brentq(lambda z: am_cdf(z, kappa) - target,
+                                 lo * (1 - 1e-9), hi * (1 + 1e-9),
+                                 xtol=1e-13, rtol=1e-13))
+
+
+def curvature(lam1: float, sigma: np.ndarray) -> float:
+    """KSS's q = 1 curvature from lambda_1 and the 2x2 covariance."""
+    v_b, v_t, cov = float(sigma[0, 0]), float(sigma[1, 1]), float(sigma[0, 1])
+    rho2 = cov * cov / (v_b * v_t)
+    return 2.0 * abs(lam1) * v_b / np.sqrt(v_t * (1.0 - rho2))
+
+
+def am_interval(lam1: float, b1: float, theta1: float, sigma: np.ndarray,
+                alpha: float = 0.05, grid: int = 4096) -> tuple[float, float, dict[str, float]]:
+    """(lower, upper, info) for the q = 1 weak-identification interval.
+
+    `sigma` is the 2x2 covariance of (b1-hat, theta1-hat), which must be
+    positive definite. The extremes of lambda_1 b^2 + t over the ellipse lie on
+    its boundary -- the map is linear in t with coefficient one, so any interior
+    point can be improved by moving t -- so the problem is one-dimensional in
+    the boundary angle. The objective there is a trigonometric polynomial of
+    degree two, with at most four stationary points, so a dense grid followed by
+    a bounded Brent refinement finds both global extremes. KSS's closed form
+    (Appendix C.6.2) solves the same first-order condition as a quartic; the
+    tests check the two agree.
+    """
+    from scipy import optimize
+
+    sigma = np.asarray(sigma, dtype=np.float64)
+    kappa = curvature(lam1, sigma)
+    z = am_critical_value(kappa, alpha)
+    L = np.linalg.cholesky(sigma)
+
+    def g(t):
+        b = b1 + z * L[0, 0] * np.cos(t)
+        th = theta1 + z * (L[1, 0] * np.cos(t) + L[1, 1] * np.sin(t))
+        return lam1 * b * b + th
+
+    ts = np.linspace(0.0, 2.0 * np.pi, grid, endpoint=False)
+    values = g(ts)
+    step = ts[1] - ts[0]
+
+    def refine(index, sign):
+        t0 = ts[index]
+        res = optimize.minimize_scalar(lambda t: sign * g(t),
+                                       bounds=(t0 - step, t0 + step),
+                                       method="bounded",
+                                       options={"xatol": 1e-14})
+        return float(g(res.x)) if sign > 0 else float(g(res.x))
+
+    lower = min(refine(int(np.argmin(values)), 1.0), float(values.min()))
+    upper = max(refine(int(np.argmax(values)), -1.0), float(values.max()))
+    rho = float(sigma[0, 1] / np.sqrt(sigma[0, 0] * sigma[1, 1]))
+    return lower, upper, {"kappa": float(kappa), "z": float(z), "rho": rho}

hdfe_stream/api.py

@@ -0,0 +1,184 @@
+"""`feols_stream`, `fepois_stream`, `feglm_stream`: fit a pyfixest-style
+formula out of core."""
+
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import polars as pl
+
+from ._types import PathLike, Source, Vcov
+
+from .estimator import StreamingHDFE
+from .feterms import _parse_fe_term
+from .glm import StreamingGLM
+from .report import _log
+from .formula import _plan_formula, _pyfixest_formula_api
+from .results import HDFEMulti, HDFEResult, _vcov_key
+
+
+class _FormulaEstimator:
+    def __init__(self, fml: str, workdir: PathLike | None, options: dict[str, Any]) -> None:
+        self.fml, self.workdir, self.options = fml, workdir, options
+
+    def fit(self, data: Source, vcov: Vcov | None = None, cluster: Sequence[str] | str = (),
+            fe_dof: str = "exact") -> HDFEResult | HDFEMulti:
+        return feols_stream(self.fml, data, self.workdir, vcov=vcov, cluster=cluster,
+                            fe_dof=fe_dof, **self.options)
+
+
+def feols_stream(fml: str, data: Source, workdir: PathLike | None = None,
+                 vcov: Vcov | None = None, cluster: Sequence[str] | str = (),
+                 fe_dof: str = "exact", **options: Any) -> HDFEResult | HDFEMulti:
+    """
+    Out-of-core OLS with high-dimensional fixed effects from a pyfixest-style
+    formula, e.g. "y ~ x1 + i(year, treat, ref=2010) | worker_id + firm_id^year".
+    Any number of fixed effects works, including one ("y ~ x | worker_id",
+    a within regression) and none ("y ~ x", OLS with an intercept).
+
+    data    : Parquet path/glob or Polars LazyFrame.
+    workdir : directory under which run directories are created (default:
+              $HDFE_STREAM_WORKDIR if set, else the system temporary
+              directory); see StreamingHDFE for the
+              outputs=, save_resid= and keep_intermediates= options.
+    vcov    : as in pyfixest; default {'CRV1': <first FE>} like pyfixest, or
+              iid without fixed effects. {'CRV3': var} (one-way, OLS) needs
+              every fixed effect nested within the clusters, or none.
+    cluster : extra cluster specs to compute CRV1 for (one-way or 'a+b').
+    **options : passed to StreamingHDFE (stream, weights, weights_type,
+              solver, precond, keep, verbose, logger, log_level, ...).
+
+    Varying slopes on the streamed dimension: "y ~ x | worker_id[t] + firm_id".
+    The worker FE file then has fe_worker_id (intercept) and fe_worker_id[t] (slope)
+    columns; the residual file's fe_worker_id column is the worker's total
+    contribution for that row.
+
+    IV formulas ("y ~ x | fe | endog ~ z") are estimated by 2SLS; each result
+    carries its first-stage regressions (`first_stage`) and the Wald F of the
+    excluded instruments computed with the same vcov (`f_stat_1st_stage`).
+
+    All models with the same fixed effects share one set of passes and one
+    solve, and are estimated on the rows where *all* of their variables are
+    non-missing (pyfixest drops missing values model by model).
+    """
+    def estimators(fe, g):
+        yield StreamingHDFE(y=list(g["y"].items()), x=list(g["x"].items()), fe=list(fe),
+                            workdir=workdir, models=g["models"], **options)
+    return _fit_formula(fml, data, estimators, vcov, cluster, fe_dof, options)
+
+
+def _fit_formula(fml, data, estimators, vcov, cluster, fe_dof, options):
+    """Plan the formula, fit each estimator that `estimators(fe, group)`
+    makes for each set of fixed effects, and return the results in the
+    formula's order."""
+    lf = data if isinstance(data, pl.LazyFrame) else pl.scan_parquet(data)
+    _log(options.get("verbose", True), f"planning {fml!r} (formula expansion, level discovery)",
+         options.get("logger"), options.get("log_level"))
+    groups = _plan_formula(fml, lf)
+    results = []
+    try:
+        for fe, g in groups.items():
+            # pyfixest's default: cluster by the first fixed effect as written,
+            # or iid when there is none
+            default = {"CRV1": _parse_fe_term(fe[0])[0]} if fe else "iid"
+            for est in estimators(fe, g):
+                res = est.fit(lf, vcov=vcov if vcov is not None else default, cluster=cluster,
+                              fe_dof=fe_dof)
+                results += list(res) if isinstance(res, HDFEMulti) else [res]
+    except BaseException:
+        # a later fit failed: don't leave the earlier fits' files behind
+        if not options.get("keep_intermediates"):
+            for r in results:
+                r.cleanup()
+        raise
+    Formula, _ = _pyfixest_formula_api()
+    order = {s.formula: i for i, s in enumerate(Formula.parse(fml))}
+    results.sort(key=lambda r: order.get(r.fml, len(order)))
+    return results[0] if len(results) == 1 else HDFEMulti(results)
+
+
+def fepois_stream(fml: str, data: Source, workdir: PathLike | None = None,
+                  vcov: Vcov | None = None, cluster: Sequence[str] | str = (),
+                  fe_dof: str = "exact", offset: str | pl.Expr | None = None,
+                  iwls_tol: float = 1e-8, iwls_maxiter: int = 25,
+                  separation_check: bool = True, **options: Any) -> HDFEResult | HDFEMulti:
+    """
+    Out-of-core Poisson regression (log link) with high-dimensional fixed
+    effects, from a pyfixest-style formula, e.g.
+    "trade ~ log_dist | exporter^year + importer^year". Estimated by
+    iteratively reweighted least squares; each step is one read of the rows
+    and one solve of the fixed effects (see `StreamingGLM`).
+
+    data, workdir, vcov, cluster, fe_dof, **options : as in `feols_stream`,
+              except that the vcov may not be CRV3, and IV formulas and
+              varying slopes are not available. weights= and weights_type=
+              work as in pyfixest's fepois.
+    offset  : column name added to the linear predictor with a coefficient of
+              one (e.g. log exposure).
+    iwls_tol, iwls_maxiter : convergence tolerance on the relative change in
+              deviance, and the most IRLS steps (pyfixest's defaults).
+    separation_check : drop fixed-effect levels whose outcome is zero on
+              every row (their effect would be minus infinity), as pyfixest's
+              "fe" check, which it runs by default. Covariates that separate
+              the outcome (its "ir" check) are not detected.
+
+    The dependent variable must be nonnegative. Each dependent variable is
+    its own set of IRLS fits (the separation check depends on it); models of
+    the same outcome share pass 0.
+    """
+    return _fit_glm(fml, data, "poisson", workdir, vcov, cluster, fe_dof,
+                    dict(options, offset=offset, iwls_tol=iwls_tol, iwls_maxiter=iwls_maxiter,
+                         separation_check=separation_check))
+
+
+def feglm_stream(fml: str, data: Source, family: str, workdir: PathLike | None = None,
+                 vcov: Vcov | None = None, cluster: Sequence[str] | str = (),
+                 fe_dof: str = "exact", iwls_tol: float = 1e-8, iwls_maxiter: int = 25,
+                 separation_check: bool = True, **options: Any) -> HDFEResult | HDFEMulti:
+    """
+    Out-of-core logit or probit regression with high-dimensional fixed
+    effects, from a pyfixest-style formula: `family` is "logit" or "probit".
+    Estimated by iteratively reweighted least squares, with step-halving, as
+    pyfixest's feglm (see `StreamingGLM`).
+
+    Arguments as in `fepois_stream`, without the offset. The dependent
+    variable must be 0 or 1. The separation check drops fixed-effect levels
+    whose outcome is the same on every row (all 0 or all 1), repeating until
+    none is left; pyfixest's "fe" check drops only levels whose outcome is
+    all 0, once.
+
+    Weights (weights=, weights_type=) multiply each row's log-likelihood, as
+    in fepois; pyfixest's feglm takes none. Frequency weights give exactly
+    the fit of the data with each row repeated that many times.
+
+    No correction is made for the incidental parameter problem: with fixed
+    effects estimated from few observations each, the coefficients are
+    biased, as in pyfixest and fixest.
+    """
+    if str(family).lower() not in ("logit", "probit"):
+        raise ValueError(f"family must be 'logit' or 'probit', got {family!r}"
+                         + ("; use fepois_stream" if str(family).lower() == "poisson" else "")
+                         + ("; for a linear model use feols_stream"
+                            if str(family).lower() == "gaussian" else ""))
+    return _fit_glm(fml, data, str(family).lower(), workdir, vcov, cluster, fe_dof,
+                    dict(options, iwls_tol=iwls_tol, iwls_maxiter=iwls_maxiter,
+                         separation_check=separation_check))
+
+
+def _fit_glm(fml, data, family, workdir, vcov, cluster, fe_dof, options):
+    key = _vcov_key(vcov)
+    if key is not None and key.startswith("CRV3:"):
+        raise ValueError("CRV3 standard errors are not available for GLMs; use CRV1")
+
+    def estimators(fe, g):
+        by_y = {}
+        for model in g["models"]:
+            if model.get("iv"):
+                raise ValueError(f"{model['fml']}: IV (2SLS) is not available for GLMs")
+            by_y.setdefault(model["y"], []).append(model)
+        for yname, models in by_y.items():
+            xs = list(dict.fromkeys(x for model in models for x in model["x"]))
+            yield StreamingGLM(y=[(yname, g["y"][yname])], x=[(x, g["x"][x]) for x in xs],
+                               fe=list(fe), family=family, workdir=workdir, models=models,
+                               **options)
+    return _fit_formula(fml, data, estimators, vcov, cluster, fe_dof, options)