logit-stream 0.1.0__py3-none-any.whl

logit_stream/__init__.py

@@ -0,0 +1,77 @@
+"""logit-stream: logit models fitted out of core by streamed Newton.
+
+Binary (`logit_stream`, `probit_stream`: hdfe-stream's `feglm_stream`, with
+any number of absorbed fixed effects), multinomial (`mlogit_stream`),
+ordered (`ologit_stream`, `oprobit_stream`), conditional (`clogit_stream`),
+conditional fixed-effects (`cfelogit_stream`), mixed (`mixlogit_stream`) and
+nested (`nlogit_stream`) logit, from pyfixest-style formulas, on data that
+need not fit in memory: Parquet files, Polars LazyFrames, or DataFrames.
+
+Layout
+    api.py           the estimators
+    _frontend.py     formulas -> models (hdfe-stream's formula compiler)
+    store.py         pass 0, into hdfe-stream's block store (hdfe_stream.core.store)
+    models.py        the likelihoods: per chunk, log-likelihood, gradient,
+                     information, scores
+    newton.py        Newton's method, collinearity, the variance's middle matrix
+    estimator.py     formula + source -> result, and the inference
+    fe.py            one absorbed fixed effect (multinomial, ordered)
+    multi_fe.py      several, for the multinomial logit (a Poisson regression)
+    random_effects/  random intercepts (Laplace's approximation)
+    optimize.py      Newton and quasi-Newton steps from an analytic gradient,
+                     for the likelihoods that are not concave
+    _choice.py       what the mixed and nested logit's fits share
+    mixed.py         the mixed logit (simulated likelihood)
+    nested.py        the nested logit
+    device/          the passes on a GPU (CuPy)
+    results.py       the result classes
+    simulate.py      simulated data for examples and tests
+"""
+
+from __future__ import annotations
+
+from hdfe_stream.core.multcomp import bonferroni, holm
+from hdfe_stream.core.reporting import etable
+
+from .api import (
+    cfelogit_stream,
+    clogit_stream,
+    logit_stream,
+    mixlogit_stream,
+    mlogit_stream,
+    nlogit_stream,
+    ologit_stream,
+    oprobit_stream,
+    probit_stream,
+)
+from .results import (
+    ConditionalFELogitResult,
+    ConditionalLogitResult,
+    LogitMulti,
+    LogitResult,
+    MixedLogitResult,
+    MultinomialLogitResult,
+    NestedLogitResult,
+    OrderedLogitResult,
+    OrderedProbitResult,
+)
+
+
+def _version():
+    try:
+        from importlib.metadata import version
+        return version("logit-stream")
+    except Exception:                   # noqa: BLE001 -- not installed
+        return "0.1.0"
+
+
+__version__ = _version()
+
+__all__ = [
+    "logit_stream", "probit_stream", "mlogit_stream", "ologit_stream", "oprobit_stream",
+    "clogit_stream", "cfelogit_stream", "mixlogit_stream", "nlogit_stream",
+    "LogitResult", "MultinomialLogitResult", "OrderedLogitResult", "OrderedProbitResult",
+    "ConditionalLogitResult", "ConditionalFELogitResult",
+    "MixedLogitResult", "NestedLogitResult",
+    "LogitMulti", "etable", "bonferroni", "holm", "__version__",
+]

logit_stream/_choice.py

@@ -0,0 +1,64 @@
+"""What the fits on choice data with `optimize._GradientDriver` (the mixed
+and the nested logit) share: the one formula, the conditional logit as the
+starting point, and the diagnostics of the fit."""
+
+from __future__ import annotations
+
+import time
+
+import numpy as np
+
+
+def plan_one(fml, lf, engine, log, model):
+    """The formula's one model, without fixed effects (`model` names the
+    estimator in the errors)."""
+    from ._frontend import _plan
+    log(f"planning {fml!r}")
+    specs = _plan(fml, lf, intercept=False, collect=lambda q, what: engine.collect(q, what))
+    if len(specs) != 1:
+        raise ValueError(f"multiple estimation is not available for the {model}")
+    spec = specs[0]
+    if spec.fe:
+        raise ValueError(f"{spec.fml}: fixed effects are not available for the {model}")
+    return spec
+
+
+def clogit_start(p0, xnames, spec, logger, log):
+    """The conditional logit on the store: the starting point (it refuses
+    collinear covariates, which these fits do not drop)."""
+    from .models import _Conditional
+    from .newton import _Driver
+    log("starting values: the conditional logit")
+    K = len(xnames)
+    cdrv = _Driver(p0.store, _Conditional(xnames, K), tol=1e-8, maxiter=50, l2=0.0,
+                   chunk_mb=64.0, log=lambda m: None, logger=logger)
+    cfit = cdrv.run(np.zeros(K), spec.fml)
+    if cfit.dropped:
+        raise ValueError(f"{spec.fml}: collinear covariates {cfit.dropped}; drop them (they "
+                         "are not dropped for this model)")
+    return cfit
+
+
+def fit_diagnostics(fit, cfit, p0, *, storage, dtype, tol, t0, device=None, drv=None,
+                    **more):
+    """The diagnostics of a `_GradientFit` started from `cfit`."""
+    out = {
+        "optimizer": {"iterations": fit.iterations, "bhhh_steps": fit.bhhh_steps,
+                      "newton_steps": fit.newton_steps, "quasi_steps": fit.quasi_steps,
+                      "converged": fit.converged, "decrement": fit.decrement,
+                      "loglik": fit.path, "passes": fit.passes, "seconds": fit.seconds,
+                      "tol": tol},
+        "start": {"loglik_clogit": float(cfit.ll), "passes": cfit.passes},
+        "gradient": [float(v) for v in fit.g],
+        "pass0_seconds": p0.seconds, "n_rows": p0.n_obs, "n_raw": p0.n_raw,
+        "store": {"storage": storage, "dtype": dtype, "blocks": len(p0.store.blocks),
+                  "bytes": int(p0.store.bytes)},
+        "disk_peak_gb": round((p0.store.bytes if storage == "disk" else 0) / 1e9, 4),
+        "seconds_total": round(time.time() - t0, 2),
+        **more,
+    }
+    if device is not None:
+        ds = drv._dstore
+        out["device"] = {"name": device.name, "resident": bool(ds and ds.resident),
+                         "bytes": int(ds.bytes) if ds else 0}
+    return out

logit_stream/_frontend.py

@@ -0,0 +1,78 @@
+"""Formula front end: a pyfixest-style formula -> the models to fit.
+
+The right-hand side is compiled by hdfe-stream's formula compiler
+(`hdfe_stream.core.formula`), which runs pyfixest's own parser and model
+matrix on a tiny synthetic frame and turns each column back into a Polars
+expression. Coefficient names and the coding of categorical variables are
+therefore exactly pyfixest's. The outcome is handled here: for the
+multinomial and ordered models it is a categorical variable, not a number,
+so it is kept as a raw column and coded once its levels are known.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import polars as pl
+from hdfe_stream.core.feterms import parse_fe
+from hdfe_stream.core.formula import (
+    categorical_vars,
+    compile_design,
+    numeric_expr,
+    pyfixest_formula_api,
+)
+
+
+@dataclass
+class _Spec:
+    """One model: its formula, outcome, covariates and fixed effects."""
+
+    fml: str
+    depvar: str
+    y: pl.Expr                      # the outcome, uncast
+    x: list = field(default_factory=list)       # [(coefficient name, Float64 expr)]
+    fe: tuple = ()
+
+
+def _outcome_expr(dep, schema):
+    """A column is taken as it is (strings, categoricals, integers);
+    anything else is a numeric expression such as ``I(educ > 12)``."""
+    if dep in schema:
+        return pl.col(dep)
+    return numeric_expr(dep)
+
+
+def _plan(fml: str, lf: pl.LazyFrame, intercept: bool, collect) -> list[_Spec]:
+    """Expand `fml` into its models (pyfixest's multiple-estimation syntax:
+    several outcomes, ``sw()``, ``csw()``, ...).
+
+    `intercept`: whether a model without fixed effects gets the model
+    matrix's constant (multinomial logit does; ordered and conditional logit
+    do not, their cutpoints and choice sets absorbing it). Without it the
+    categorical coding is still reduced rank, as with it.
+    `collect(lf, what)` runs a query (so that the GPU engine applies).
+    """
+    Formula, _ = pyfixest_formula_api()
+    schema = lf.collect_schema()
+    specs = Formula.parse(fml)
+    cats = set()
+    for s in specs:
+        if s.first_stage is not None:
+            raise ValueError(f"{s.formula}: IV formulas are not available for logit models")
+        cats |= categorical_vars(s.second_stage.split("~", 1)[1], schema)
+    levels = {}
+    for v in sorted(cats):
+        u = collect(lf.select(pl.col(v).drop_nulls().unique().sort()), f"levels of {v}")
+        levels[v] = u[v].to_list()
+    out = []
+    for s in specs:
+        dep, rhs = (t.strip() for t in s.second_stage.split("~", 1))
+        fe = tuple(parse_fe(s.fixed_effects))
+        for d in fe:
+            if "[" in d:
+                raise ValueError(f"{s.formula}: varying slopes are not available for logit models")
+        _, cols = compile_design(dep, rhs, schema, levels, intercept=intercept and not fe)
+        if not intercept:
+            cols = [(nm, ex) for nm, ex in cols if nm != "Intercept"]
+        out.append(_Spec(fml=s.formula, depvar=dep, y=_outcome_expr(dep, schema), x=cols, fe=fe))
+    return out