Regression Adapters

What Is It

survey_kit.statistics.adapters wraps regression packages - four pure-Python (statsmodels, linearmodels, pyfixest, polars_ds), R (via rpy2/fixest), and Stata (via pystata) - behind one normalized return type:

AdapterStats - a StatCalculator subclass built directly from the adapter's own df_estimates/df_ses (plus df_vcov/df_tidy when the package provides them), rather than from raw microdata. Being a real StatCalculator means it already works everywhere one does with no extra step: .print(), .compare() (using df_vcov for a correct joint SE when available), survey_kit.plot, save/load, and - since mi_ses_from_function reads whatever a delegate returns generically - plugging straight into multiple imputation.

so every estimator plugs into the same downstream machinery with no special-casing for which package produced the numbers. Each adapter also has a matching mi_ses_from_<package>(...) shortcut that runs it across multiple-imputation implicates directly - the adapter's own arguments (formula, weight, vcov, ...) come through as plain keywords instead of being packed into an arguments={} dict, so your IDE shows the right parameters for the one you're actually calling.

Why Use It

The adapters handle the fiddly parts of wiring a regression package into mi_ses_from_function - term-name alignment between df_estimates/df_ses/df_vcov, a missing vcov() method, converting the data into whatever the underlying package expects - so:

• Any package plugs into MI/replicate-weight machinery identically - swap pyfixest_adapter for r_feols for stata_adapter without changing anything downstream.
• Simple call shape for the common case - mi_ses_from_pyfixest.feols(df_implicates=..., fml="y ~ x1 + x2 | firm", ...) runs the fit across every implicate and combines the results.
• Replicate-weight bootstrapping built in - pass replicates= to run the same command once per replicate weight column and get the SE from the spread across replicates, instead of the package's own vcov/cov_type - useful when you need SEs computed the same way elsewhere in a project rather than trusting a given package's own variance estimator. The data conversion to whatever the underlying package needs (pandas, an R data.frame, a Stata .dta) happens once per implicate, not once per replicate.
• An escape hatch when you need it - every language also exposes its underlying primitives (_r_interop, _stata_interop) for writing a custom delegate when the named adapters don't cover what you need. See Rolling Your Own below.

Key Features

• Same return type everywhere - an AdapterStats regardless of package, so it's a real StatCalculator with .print()/.compare()/plotting/save-load already working.
• mi_ses_from_<package> shortcuts - one call combines across implicates via Rubin's rules.
• Replicate-weight bootstrapping - replicates= on every mi_ses_from_* that has a weight argument to substitute a column into.
• No hard dependencies - none of statsmodels/linearmodels/pyfixest/polars_ds/rpy2/pystata are required by survey_kit itself; each adapter raises a clear, actionable error (with the install command) only if you actually call it without the package installed.
• Data conversion caching - each package's own "already converted, don't redo the work" passthrough (an already-pandas frame, an already-data.frame R object, Stata's reuse_data=) means repeated calls against the same implicate - the replicate-weight loop being the main case - don't redo an expensive conversion on every call.

When to Use What

Use Case	Tool	Why
Already using statsmodels/linearmodels/pyfixest/polars_ds	matching *_adapter/mi_ses_from_*	No R/Stata install needed
Fixed effects, clustered SEs	pyfixest_adapter/mi_ses_from_pyfixest	Generally the best default of the four Python adapters
Already have R code/packages you trust	r_feols/r_fixest_adapter/r_lm_adapter	Reuses your existing R model specifications
Already have Stata code/do-files	stata_adapter/mi_ses_from_stata	command= takes the Stata syntax directly
Need SEs computed the same way as elsewhere in a project	any mi_ses_from_* with replicates=	Bootstraps from replicate weights instead of the package's own vcov
A package/function with no named adapter	_r_interop/_stata_interop primitives directly	See Rolling Your Own

API

See the Regression Adapters API reference for the full parameter list of every adapter and mi_ses_from_* shortcut.

Example/Tutorial

Python (statsmodels/linearmodels/pyfixest/polars_ds)R (fixest)Stata

No R/Stata install needed - these four wrap packages already available in Python.

CodeLog

from __future__ import annotations

from survey_kit import logger
from survey_kit.statistics.adapters import (
    statsmodels_adapter,
    linearmodels_adapter,
    pyfixest_adapter,
    polars_ds_adapter,
    mi_ses_from_statsmodels,
    mi_ses_from_linearmodels,
    mi_ses_from_pyfixest,
    mi_ses_from_polars_ds,
)
from survey_kit.utilities.dataframe import summary
from survey_kit.statistics.replicates import Replicates
from sample_data import make_implicates, with_bootstrap_weights

# %%
logger.info("survey_kit.statistics.adapters has four pure-Python regression")
logger.info("adapters, one per package - no R/Stata/rpy2/pystata needed. Every")
logger.info("adapter (these four, plus the R/Stata ones) returns the same")
logger.info("normalized type: an AdapterStats (a StatCalculator subclass) with")
logger.info("df_estimates/df_ses/df_vcov/df_tidy populated - so it already works")
logger.info("everywhere a StatCalculator does (.print(), .compare(), survey_kit.plot,")
logger.info("save/load) with no extra step. Each package also has a matching")
logger.info("mi_ses_from_<package>(...) - same shape as mi_ses_from_r_fixest/")
logger.info("mi_ses_from_stata - that runs it across implicates, combined via")
logger.info("Rubin's rules, taking the adapter's own arguments directly instead of")
logger.info("an arguments={} dict.")

# %%
df_implicates = make_implicates()
df = df_implicates[0]
logger.info("\n\nSample data (y = 1 + 2*x1 - 1.5*x2 + noise) - see sample_data.py:")
summary(df)

# %%
logger.info("\n\nstatsmodels - y/x column lists rather than a formula, HC3")
logger.info("(heteroskedasticity-robust) SEs by default:")
sm_result = statsmodels_adapter(df, y="y", x=["x1", "x2"])
logger.info("\n   Single df")
sm_result.print()

mi_sm = mi_ses_from_statsmodels(df_implicates=df_implicates, y="y", x=["x1", "x2"])
logger.info("\n   Multiple imputation")
mi_sm.print()

# %%
logger.info("\n\nlinearmodels - formula syntax, IV/panel-capable (plain OLS via")
logger.info("IV2SLS with no instruments, as here):")
lm_result = linearmodels_adapter(df, formula="y ~ 1 + x1 + x2")
logger.info("\n   Single df")
lm_result.print()

mi_lm = mi_ses_from_linearmodels(df_implicates=df_implicates, formula="y ~ 1 + x1 + x2")
logger.info("\n   Multiple imputation")
mi_lm.print()

# %%
logger.info("\n\npyfixest - fixest-syntax formula, fixed effects supported")
logger.info("directly (e.g. 'y ~ x1 + x2 | firm') - generally the best default")
logger.info("of the four unless you specifically need something it doesn't")
logger.info("cover (see its docstring):")
pf_result = pyfixest_adapter(df, formula="y ~ x1 + x2")
logger.info("\n   Single df")
pf_result.print()

mi_pf = mi_ses_from_pyfixest.feols(df_implicates=df_implicates, fml="y ~ x1 + x2")
logger.info("\n   Multiple imputation")
mi_pf.print()

# %%
logger.info("\n\nReplicate-weight bootstrapping instead of pyfixest's own vcov: pass")
logger.info("replicates=, and pyfixest_adapter runs once per replicate weight column")
logger.info('(point estimates only, vcov forced to "iid" since it\'s discarded')
logger.info("anyway) - the spread of estimates across replicates IS the SE, computed")
logger.info("by survey_kit's own Replicates/StatCalculator machinery, the same")
logger.info("approach mi_ses_from_stata/mi_ses_from_r_fixest's replicates= use. Every")
logger.info("mi_ses_from_<package> here (except .feglm/.femlm, whose underlying")
logger.info("estimators have no weights= argument to substitute a replicate column")
logger.info("into) supports this the same way.")

N_REPLICATES = 20
mi_pf_boot = mi_ses_from_pyfixest.feols(
    df_implicates=with_bootstrap_weights(df_implicates, n_replicates=N_REPLICATES),
    fml="y ~ x1 + x2",
    replicates=Replicates(
        weight_stub="replicate_", n_replicates=N_REPLICATES, bootstrap=True
    ),
)
mi_pf_boot.print()

# %%
logger.info("\n\npolars_ds - stays entirely in polars/narwhals, no pandas")
logger.info("conversion at all; doesn't compute a covariance matrix, so df_vcov")
logger.info("is always None here:")
pds_result = polars_ds_adapter(df, y="y", x=["x1", "x2"])
logger.info("\n   Single df")
pds_result.print()

mi_pds = mi_ses_from_polars_ds(df_implicates=df_implicates, y="y", x=["x1", "x2"])
logger.info("\n   Multiple imputation")
mi_pds.print()

# %%
logger.info("\n\nFor R or Stata instead, see basic_r.py/basic_stata.py (or")
logger.info("r_arbitrary_estimators.py/stata_arbitrary_estimators.py for the generic")
logger.info("escape hatch into any function either package provides).")

View in separate window

Requires R itself plus rpy2/rpy2-arrow (pip install survey-kit[r]) and the R fixest package. Check your setup cheaply with from survey_kit.statistics._r_interop import check_r_setup; check_r_setup(['fixest']).

CodeLog

from __future__ import annotations

from survey_kit import logger
from survey_kit.statistics.adapters import r_feols, mi_ses_from_r_fixest
from sample_data import make_implicates

# %%
logger.info("The simplest way to run an R regression from survey_kit: r_feols(),")
logger.info("a wrapper around fixest::feols() with robust SEs by default. Requires R")
logger.info("itself plus rpy2/rpy2-arrow (`pip install survey-kit[r]`) and the R")
logger.info("'fixest' package. Check your setup cheaply with:")
logger.info("    from survey_kit.statistics._r_interop import check_r_setup")
logger.info("    check_r_setup(['fixest'])")

# %%
df_implicates = make_implicates()
logger.info(
    f"\n\nSample data: {len(df_implicates)} implicates, {df_implicates[0].height}"
)
logger.info("rows each (y = 1 + 2*x1 - 1.5*x2 + noise) - see sample_data.py.")

# %%
logger.info("\n\nOn one dataset, standalone - no MI at all. r_feols() returns an")
logger.info("AdapterStats (a StatCalculator subclass), so it already works with")
logger.info(".print(), .compare(), survey_kit.plot, save/load - no extra step:")
r_result = r_feols(df_implicates[0], formula="y ~ x1 + x2")
r_result.print()


# %%
logger.info("\n\nAcross multiple imputed datasets, combined via Rubin's rules -")
logger.info("mi_ses_from_r_fixest.feols(...) runs r_feols once per implicate and")
logger.info("combines the results, taking r_feols's own arguments directly:")
mi_reg = mi_ses_from_r_fixest.feols(
    df_implicates=df_implicates,
    formula="y ~ x1 + x2",
    round_output=False,
)
mi_reg.print(round_output=False)


# %%
logger.info("\n\nThat's it for the common case. mi_ses_from_r_fixest also has")
logger.info(".feglm/.fepois/.femlm for fixest's other estimators, all with the same")
logger.info("shape. r_lm_adapter (base R lm()/glm()) and r_fixest_adapter (any")
logger.info("fixest estimator via a func= string, including ones .feglm/.fepois/")
logger.info(".femlm don't cover) are the lower-level pieces these are built from -")
logger.info("see r_arbitrary_estimators.py for rolling your own with those, plus a")
logger.info("generic escape hatch into any R package/function at all.")

View in separate window

Requires Stata 17+ plus pip install survey-kit[stata]. Check your setup cheaply with from survey_kit.statistics._stata_interop import check_stata_setup; check_stata_setup(stata_path=r"C:\Program Files\Stata18").

CodeLog

from __future__ import annotations

import os

import numpy as np
import polars as pl
from dotenv import load_dotenv

from survey_kit import logger
from survey_kit.statistics.adapters import stata_adapter

# %%
logger.info("The simplest way to run a Stata regression from survey_kit:")
logger.info("stata_adapter() - pass any e-class command as a plain string.")

#   Machine-specific - set these in a local ".env" file (see .gitignore,
#   which excludes it from git) in the repo root rather than editing this
#   file or exporting them yourself:
#       _survey_kit_stata_path_=C:\Program Files\Stata17
#       _survey_kit_stata_edition_=se
#   Or, just as easily, set them directly in code instead of via env vars:
#       from survey_kit import config
#       config.stata_path = r"C:\Program Files\Stata17"
#       config.stata_edition = "se"
load_dotenv()

# %%
rng = np.random.default_rng(0)
n = 300
x1 = rng.normal(size=n)
x2 = rng.normal(size=n)
y = 1 + 2 * x1 - 1.5 * x2 + rng.normal(size=n) * 0.4
df = pl.DataFrame({"x1": x1, "x2": x2, "y": y})

# %%
logger.info("\n\nOn one dataset, standalone - no MI at all. stata_adapter() returns")
logger.info("an AdapterStats (a StatCalculator subclass), so it already works with")
logger.info(".print(), .compare(), survey_kit.plot, save/load - no extra step:")
stata_result = stata_adapter(df, command="regress y x1 x2")
stata_result.print()


# %%
logger.info(
    "\n\nThat's it for the basics. For multiple imputation (mi_ses_from_stata),"
)
logger.info("weighted/survey designs, other commands (svy:/xtreg/areg/logit/a")
logger.info("community-installed ado/...), replicate-weight bootstrapping, or pulling")
logger.info("back custom r()/e() results via stata_results_adapter instead of the")
logger.info("usual e(b)/e(V)/r(table), see stata_arbitrary_estimators.py.")

Rolling Your Own

Each language's tutorial above only covers the named adapters (statsmodels_adapter, r_feols, stata_adapter, ...). For a package or function with no named wrapper, every language exposes the primitives those adapters are themselves built from, so writing a custom delegate is a few lines rather than a new adapter:

RStata

_r_interop's get_library/dataframe_to_r/formula/extract_fit/coef_table etc. - any R function becomes a Python-callable attribute via rpy2's dot-calling, with no R call string to build. Covers reaching any R package/function at all, plus how to convert an implicate to R once and reuse it across several model calls (or once per replicate weight) instead of re-converting on every call.

CodeLog

from __future__ import annotations

import numpy as np
import polars as pl

from survey_kit import logger
from survey_kit.statistics.multiple_imputation import mi_ses_from_function
from survey_kit.statistics.adapters import r_feols, mi_ses_from_r_fixest
from survey_kit.statistics.adapter_stats import AdapterStats
from survey_kit.statistics import _r_interop as _r
from survey_kit.statistics.replicates import Replicates
from sample_data import make_implicates, with_bootstrap_weights


# %%
logger.info("Every regression adapter in survey_kit.statistics.adapters is a plain")
logger.info("function returning an AdapterStats (a StatCalculator subclass built")
logger.info("directly from df_estimates/df_ses, optionally df_vcov/df_tidy too - see")
logger.info("survey_kit.statistics.adapter_stats). mi_ses_from_function() calls it")
logger.info("once per implicate and combines the results via Rubin's rules - there's")
logger.info("no special-casing for which package produced the estimates, and no")
logger.info("other return shape is accepted (build an AdapterStats even for a")
logger.info("one-off custom delegate, as Part 2 below does).")
logger.info("")
logger.info("This tutorial has two parts:")
logger.info("  1. A quick reminder of using a *named* R adapter (r_feols).")
logger.info("  2. How to reach ANY R function survey_kit hasn't wrapped - fixest and")
logger.info("     base lm()/glm() are just the ones with named wrappers; the plumbing")
logger.info("     underneath (survey_kit.statistics._r_interop) works for anything.")
logger.info("")
logger.info("Requires: R itself, plus rpy2/rpy2-arrow on the Python side")
logger.info("(`pip install survey-kit[r]`). Check your setup cheaply with:")
logger.info("    from survey_kit.statistics._r_interop import check_r_setup")
logger.info("    check_r_setup()")


# %%
df_implicates = make_implicates()


# %%
logger.info("\n\nPart 1: a named adapter (r_feols), via mi_ses_from_r_fixest - nothing")
logger.info("new here, just a reminder of the shape everything in this tutorial")
logger.info("produces.")

mi_feols = mi_ses_from_r_fixest.feols(
    df_implicates=df_implicates,
    formula="y ~ x1 + x2",
    round_output=False,
)
mi_feols.print(round_output=False)


# %%
logger.info("\n\nReplicate-weight bootstrapping instead of fixest's own vcov: pass")
logger.info("replicates=, and r_feols runs once per replicate weight column (point")
logger.info("estimates only, vcov forced to \"iid\" since it's discarded anyway) rather")
logger.info("than reading fixest's own SE - the spread of estimates across replicates")
logger.info("IS the SE, computed by survey_kit's own Replicates/StatCalculator")
logger.info("machinery. Each replicate weight column is passed via r_feols's own")
logger.info("`weight=` argument (no \"{weight}\" string placeholder needed in `formula`")
logger.info("the way Stata's raw `command` string needs one) - and each implicate is")
logger.info("converted to an R data.frame once, not once per replicate:")
logger.info("dataframe_to_r()'s passthrough-if-already-converted behavior (see the")
logger.info("caching section further below) makes that free to do with no")
logger.info("special-casing in r_feols itself.")

N_REPLICATES = 20

mi_boot = mi_ses_from_r_fixest.feols(
    df_implicates=with_bootstrap_weights(df_implicates, n_replicates=N_REPLICATES),
    formula="y ~ x1 + x2",
    replicates=Replicates(weight_stub="replicate_", n_replicates=N_REPLICATES, bootstrap=True),
    round_output=False,
)
mi_boot.print(round_output=False)


# %%
logger.info("\n\nPart 2: calling an R function survey_kit has no named adapter for.")
logger.info("Example: MASS::rlm() - robust (M-estimation) linear regression. MASS")
logger.info("ships with R itself, so this needs no extra R package install.")
logger.info("")
logger.info("rpy2's get_library(name) (a thin wrapper over importr()) returns an")
logger.info("object that already exposes every R function in that package as a")
logger.info("Python-callable attribute - dot-calling, the same way rpy2 always works.")
logger.info("So the pattern is just three calls into _r_interop:")
logger.info("  1. get_library(name)       - import (and cache) an R package; call")
logger.info("                                 its functions directly as attributes")
logger.info("  2. dataframe_to_r / formula - convert a polars df / formula string")
logger.info("                                 into the R objects the function needs")
logger.info("  3. extract_fit(fit)          - pull back coef()/vcov() from the")
logger.info("                                 already-fitted R object (works for ANY")
logger.info("                                 R model with those two generics - true")
logger.info("                                 of almost every R estimator)")
logger.info("  4. coef_table/ses_from_vcov/vcov_table - convert the raw rpy2")
logger.info("                                 objects into survey_kit's normalized")
logger.info("                                 tables")
logger.info("")
logger.info("No R call string to build, no RRaw escaping, no scratch global-env")
logger.info("variable - `fit` below is already a plain Python reference to the")
logger.info("fitted R object. Wrap that in a plain function that builds and returns")
logger.info("an AdapterStats and it's already a valid mi_ses_from_function delegate")
logger.info("- no need to add it to adapters.py unless you want to reuse it")
logger.info("elsewhere.")


def rlm_adapter(
    df,
    formula: str,
    join_on_name: str = "Variable",
    value_name: str = "estimate",
) -> AdapterStats:
    mass = _r.get_library("MASS")

    fit = mass.rlm(_r.formula(formula), data=_r.dataframe_to_r(df))
    coef, vcov, _tidy = _r.extract_fit(fit)

    df_estimates = _r.coef_table(coef, join_on_name, value_name)
    df_ses = _r.ses_from_vcov(vcov, join_on_name, value_name)
    df_vcov = _r.vcov_table(vcov, join_on_name, value_name)

    return AdapterStats(
        df_estimates, df_ses, variable_ids=join_on_name, df_vcov=df_vcov, display=False
    )


# %%
logger.info("\n\nUse it exactly like any other delegate - standalone on one dataset:")
rlm_result = rlm_adapter(df_implicates[0], formula="y ~ x1 + x2")
rlm_result.print()

# %%
logger.info("\n\n...or across implicates, combined via Rubin's rules:")
mi_rlm = mi_ses_from_function(
    delegate=rlm_adapter,
    df_implicates=df_implicates,
    join_on=["Variable"],
    arguments={"formula": "y ~ x1 + x2"},
    round_output=False,
)
mi_rlm.print(round_output=False)


# %%
logger.info("\n\nIf you're going to fit SEVERAL models against the same implicates -")
logger.info("comparing specifications is a common workflow - converting each")
logger.info("implicate to R once and reusing that list is cheaper than letting every")
logger.info("mi_ses_from_function call re-convert the same data from scratch.")
logger.info("dataframe_to_r() materializes a real R data.frame (not a free")
logger.info("Arrow-backed view), and it passes an R object straight through")
logger.info("unchanged if you hand it one - so this needs no adapter code changes:")
logger.info("rlm_adapter works identically whether df is a polars frame or an")
logger.info("already-converted one.")

df_implicates_r = [_r.dataframe_to_r(dfi) for dfi in df_implicates]

mi_rlm_x1_only = mi_ses_from_function(
    delegate=rlm_adapter,
    df_implicates=df_implicates_r,
    join_on=["Variable"],
    arguments={"formula": "y ~ x1"},
    round_output=False,
)
mi_rlm_x1_only.print(round_output=False)

mi_rlm_both = mi_ses_from_function(
    delegate=rlm_adapter,
    df_implicates=df_implicates_r,
    join_on=["Variable"],
    arguments={"formula": "y ~ x1 + x2"},
    round_output=False,
)
mi_rlm_both.print(round_output=False)

logger.info("Same df_implicates_r list, two different formulas, no re-conversion")
logger.info("in between - dataframe_to_r() only ran once per implicate, back when")
logger.info("df_implicates_r was built above.")


# %%
logger.info("\n\nHow much does that actually save? Time it directly: run several")
logger.info("formulas against the same 5 implicates once with raw polars implicates")
logger.info("(re-converted to R on every single call) and once with the")
logger.info("df_implicates_r list already built above (converted once, total).")

import time

formulas_to_compare = ["y ~ x1", "y ~ x2", "y ~ x1 + x2", "y ~ x1 - 1", "y ~ x2 - 1"]

start = time.perf_counter()
for formulai in formulas_to_compare:
    mi_ses_from_function(
        delegate=rlm_adapter,
        df_implicates=df_implicates,
        join_on=["Variable"],
        arguments={"formula": formulai},
        round_output=False,
    )
elapsed_raw = time.perf_counter() - start

start = time.perf_counter()
for formulai in formulas_to_compare:
    mi_ses_from_function(
        delegate=rlm_adapter,
        df_implicates=df_implicates_r,
        join_on=["Variable"],
        arguments={"formula": formulai},
        round_output=False,
    )
elapsed_cached = time.perf_counter() - start

logger.info(
    f"{len(formulas_to_compare)} formulas x 5 implicates, raw polars "
    f"(re-converted every call): {elapsed_raw:.3f}s"
)
logger.info(
    f"{len(formulas_to_compare)} formulas x 5 implicates, pre-converted R "
    f"(converted once, reused): {elapsed_cached:.3f}s"
)
logger.info(f"Speedup: {elapsed_raw / elapsed_cached:.2f}x")
logger.info("")
logger.info("This tutorial's data is tiny (300 rows x 5 implicates), so rlm()'s own")
logger.info("fit time dominates and the gap here is modest - the win scales with")
logger.info("real data size (Arrow->R materialization cost grows with row/column")
logger.info("count) and with how many separate model calls reuse the same")
logger.info("implicates, which is exactly the shape of a real specification search.")


# %%
logger.info("\n\nA fancier example: quantreg::rq() for quantile regression. Needs an")
logger.info("extra R package (quantreg) - check first rather than finding out from a")
logger.info("cryptic mid-fit error:")
logger.info("    from survey_kit.statistics._r_interop import check_r_setup")
logger.info("    check_r_setup(['quantreg'])")
logger.info("")
logger.info("quantreg::rq() has no vcov() method at all (extract_fit tolerates that")
logger.info("gracefully - see its docstring) - the SE lives in summary() instead, so")
logger.info("this delegate builds both df_estimates and df_ses from a matrix table via")
logger.info("matrix_table() instead of coef_table()/ses_from_vcov(), using")
logger.info("extract_fit's tidy_fn= to call summary() directly (R's usual S3 dispatch")
logger.info("still applies to a dot-called generic like stats.summary(), so this")
logger.info("reaches summary.rq() exactly as calling summary(fit) would in R). Note")
logger.info("se='nid' - rq()'s summary() defaults to a rank-based CI table (columns:")
logger.info("coefficients/lower bd/upper bd) rather than the classical Value/Std.")
logger.info("Error/t value layout every other adapter in this tutorial produces;")
logger.info("se='nid' asks for that classical layout instead.")


def quantreg_adapter(
    df,
    formula: str,
    tau: float = 0.5,
    join_on_name: str = "Variable",
    value_name: str = "estimate",
) -> AdapterStats:
    qr = _r.get_library("quantreg")

    fit = qr.rq(_r.formula(formula), data=_r.dataframe_to_r(df), tau=tau)
    #   rq() has no vcov() method at all - extract_fit() tolerates that and
    #   returns vcov=None rather than raising. Pull the coefficient table
    #   from summary(fit, se="nid") instead via tidy_fn - se="nid" asks for
    #   the classical Value/Std. Error/t value layout rather than rq()'s
    #   default rank-based CI table.
    base = _r.get_library("base")
    _coef, _vcov, tidy_matrix = _r.extract_fit(
        fit, tidy_fn=lambda f: base.summary(f, se="nid").rx2("coefficients")
    )
    df_tidy = _r.matrix_table(tidy_matrix, join_on_name)

    df_estimates = df_tidy.select(join_on_name, pl.col("Value").alias(value_name))
    df_ses = df_tidy.select(join_on_name, pl.col("Std. Error").alias(value_name))

    #   No vcov here (rq() doesn't have one - see above).
    return AdapterStats(
        df_estimates, df_ses, variable_ids=join_on_name, df_tidy=df_tidy, display=False
    )


# %%
logger.info("\n\nOnly run this cell if quantreg is installed (install.packages('quantreg')")
logger.info("in R) - it's not part of base R the way MASS is.")
setup = _r.check_r_setup(["quantreg"])
if setup["r_packages"].get("quantreg"):
    #   Still using df_implicates_r (built above) rather than df_implicates -
    #   quantreg_adapter's dataframe_to_r(df) call gets the passthrough for
    #   free, same as rlm_adapter did.
    mi_qr = mi_ses_from_function(
        delegate=quantreg_adapter,
        df_implicates=df_implicates_r,
        join_on=["Variable"],
        arguments={"formula": "y ~ x1 + x2", "tau": 0.5},
        round_output=False,
    )
    mi_qr.print(round_output=False)
else:
    logger.info("quantreg isn't installed - skipping (see the warning above for how).")


# %%
logger.info("\n\nSummary - to wire up any R estimator survey_kit doesn't already wrap:")
logger.info("  1. Find its R function and confirm it has coef()/vcov() methods (most")
logger.info("     do) - or a summary()/tidy-style table if not (see quantreg above).")
logger.info("  2. Write a small Python function: get_library(name) to import the R")
logger.info("     package, call its function directly as a Python attribute (formulas")
logger.info("     via formula(), data via dataframe_to_r() - everything else is a")
logger.info("     plain keyword argument, converted by rpy2 automatically), then")
logger.info("     extract_fit(fit) plus coef_table/ses_from_vcov/vcov_table/")
logger.info("     matrix_table to normalize the result, and build an AdapterStats")
logger.info("     from the pieces (that's the one shape mi_ses_from_function accepts).")
logger.info("  3. That function is already a valid mi_ses_from_function delegate, and")
logger.info("     already usable standalone on a single dataset with no MI at all.")
logger.info("")
logger.info("A few R arguments still need special handling rather than a plain")
logger.info("Python value: a one-sided formula like weights=~column (build it with")
logger.info("formula('~column')), or an argument that's itself an R function call")
logger.info("(e.g. fixest's ssc(fixef.K=\"full\") - call get_library('fixest').ssc(**")
logger.info("{'fixef.K': 'full'}) directly rather than writing it as a string). See")
logger.info("r_fixest_adapter/_fixest_fit in adapters.py for the older string-based")
logger.info("(_r_interop.call/RRaw/fit_r_model) approach those still use - it")
logger.info("remains available for cases where building a raw R call string is")
logger.info("genuinely more convenient than assembling the equivalent rpy2 objects.")
logger.info("")
logger.info("One more thing worth knowing if you call a delegate like rlm_adapter")
logger.info("repeatedly on the SAME underlying data - the multiple-formulas example")
logger.info("earlier in Part 2, or once per replicate weight via")
logger.info("StatCalculator.from_function, where the df object passed to the")
logger.info("delegate is identical across replicates and only which weight column is")
logger.info("referenced changes: dataframe_to_r() is real work (R materializes an")
logger.info("actual data.frame, not a free Arrow-backed view), and it passes an R")
logger.info("object straight through unchanged if you hand it one - convert once,")
logger.info("hold the result, and pass that instead of re-converting every call. No")
logger.info("cache to manage or clear - it's a plain Python variable, freed the")
logger.info("normal way (falls out of scope / reassigned) once you're done with it.")
logger.info("Stata's escape hatch (adapters.stata_adapter/stata_results_adapter) has")
logger.info("an analogous reuse_data=True option, but that one DOES need an explicit")
logger.info("clear_stata_cache() call afterward - Stata holds one single shared")
logger.info("in-memory dataset (a stateful resource), unlike R where dataframe_to_r")
logger.info("is just a pure conversion you can hold a reference to.")

View in separate window

stata_adapter's generic command= string already covers any e-class Stata command by itself (svy:/xtreg/areg/logit/a community-installed ado/...) - there's no separate named-wrapper-per-estimator layer to route around the way fixest has on the R side. _stata_interop's lower-level primitives (run_stata_model, run_stata_results) are what stata_adapter/stata_results_adapter are themselves built from, for reaching something below the command-string level (a custom pre/post-processing step, or pulling back an arbitrary r()/e() result).

CodeLog

from __future__ import annotations

import os
import numpy as np
import polars as pl
from dotenv import load_dotenv


from survey_kit import logger
from survey_kit.statistics.adapters import stata_adapter, mi_ses_from_stata
from survey_kit.statistics import _stata_interop as _st
from survey_kit.statistics.replicates import Replicates
from survey_kit.utilities.random import set_seed, RandomNumberGenerator
from sample_data import make_implicates, with_bootstrap_weights

#   Machine-specific - set these in a local ".env" file (see .gitignore,
#   which excludes it from git) in the repo root rather than editing this
#   file or exporting them yourself:
#       _survey_kit_stata_path_=C:\Program Files\Stata17
#       _survey_kit_stata_edition_=se
#   Or, just as easily, set them directly in code instead of via env vars:
#       from survey_kit import config
#       config.stata_path = r"C:\Program Files\Stata17"
#       config.stata_edition = "se"
load_dotenv()


# %%
logger.info("Before running any cell here, point survey_kit at your Stata install and")
logger.info("check what's importable:")
logger.info("    from survey_kit.statistics._stata_interop import check_stata_setup")
logger.info('    check_stata_setup(stata_path=r"C:\\Program Files\\Stata18")')
logger.info("")
logger.info("Requires `pip install survey-kit[stata]` (polars_readstat, for writing")
logger.info(".dta files) plus Stata 17+ for pystata itself.")


# %%
df_implicates = make_implicates()


# %%
logger.info(
    "\n\nPart 1: mi_ses_from_stata - mi_ses_from_function(delegate=stata_adapter,"
)
logger.info("...), with stata_adapter's own arguments (command, edition, stata_path,")
logger.info("...) taken directly instead of packed into an arguments={} dict. Returns")
logger.info("the same shape every other adapter in survey_kit.statistics.adapters")
logger.info("does: an AdapterStats (df_estimates/df_ses/df_vcov/df_tidy), combined")
logger.info("across implicates via Rubin's rules.")
logger.info("")
logger.info("Data goes into Stata as a .dta file written by polars_readstat, not")
logger.info("through pystata's own DataFrame transfer - this part IS verified (the")
logger.info(".dta round-trips correctly through polars_readstat's own reader), even")
logger.info("though the Stata-side execution below isn't.")

mi_reg = mi_ses_from_stata(
    df_implicates=df_implicates,
    command="regress y x1 x2",
    round_output=False,
)
mi_reg.print(round_output=False)


# %%
logger.info("\n\nWeighted regression, and a survey design set up per-implicate:")
logger.info("`command` can be a list of Stata commands run in order (run after")
logger.info("`use` but before the estimation command) instead of a single string -")
logger.info("only the LAST command's e(b)/e(V)/r(table) are read back.")

mi_svy = mi_ses_from_stata(
    df_implicates=df_implicates,
    #   svyset's pweight must be a variable, not a literal - replace
    #   with your actual design (real psu/weight/strata variables).
    command=[
        "gen _svy_weight = 1",
        "svyset _n [pw=_svy_weight]",
        "svy: regress y x1 x2",
    ],
    round_output=False,
)
mi_svy.print(round_output=False)


# %%
logger.info("\n\nReplicate-weight bootstrapping instead of Stata's own e(V): pass")
logger.info("replicates=, and `command` runs once per replicate weight column (via")
logger.info("stata_results_adapter under the hood, reading back e(b) only) rather")
logger.info("than Stata's own bootstrap/brr/jackknife prefix - the spread of")
logger.info("estimates across replicates IS the SE, computed in Python by")
logger.info("survey_kit's own Replicates/StatCalculator machinery. Stata's own")
logger.info("replicate-estimation commands are usually faster/more idiomatic if")
logger.info("you're already set up for them - this is for matching SEs computed the")
logger.info('same way elsewhere in a project instead. `command` needs a "{weight}"')
logger.info("placeholder here, unlike the e(V)-based calls above.")

N_REPLICATES = 20

mi_boot = mi_ses_from_stata(
    df_implicates=with_bootstrap_weights(df_implicates, n_replicates=N_REPLICATES),
    command="regress y x1 x2 [pw={weight}]",
    replicates=Replicates(
        weight_stub="replicate_", n_replicates=N_REPLICATES, bootstrap=True
    ),
    round_output=False,
)
mi_boot.print(round_output=False)


# %%
logger.info("\n\nPart 2: calling this like any other delegate, standalone on one")
logger.info("dataset with no MI at all - stata_adapter() returns an AdapterStats")
logger.info("(a StatCalculator subclass), so it already works with .print(),")
logger.info(".compare(), survey_kit.plot, save/load - no extra step:")
stata_single = stata_adapter(df_implicates[0], command="regress y x1 x2")
stata_single.print()
logger.info("df_tidy is Stata's own r(table), transposed to one row per term:")
logger.info(stata_single.replicate_stats.df_tidy)


# %%
logger.info("\n\nPart 3: reaching a Stata command stata_adapter's generic `command=`")
logger.info("string already covers by itself - `stata_adapter` IS the 'arbitrary")
logger.info("estimator' entry point here (unlike the R side, there's no separate")
logger.info("named-wrapper-per-function layer to route around): any e-class command")
logger.info("works by just changing the `command` string, including svy: prefixes,")
logger.info("xtreg/xtlogit/areg, or a community-installed command (`ssc install`)")
logger.info("survey_kit has never heard of. Two examples:")

logger.info("\n  Fixed effects via areg:")
mi_areg = mi_ses_from_stata(
    df_implicates=[
        d.with_columns((pl.arange(0, pl.len()) % 10).alias("firm"))
        for d in df_implicates
    ],
    command="areg y x1 x2, absorb(firm)",
    round_output=False,
)
mi_areg.print(round_output=False)

logger.info("\n  Logit (any e-class command works the same way):")


def _with_ybin(d: pl.DataFrame, seed: int) -> pl.DataFrame:
    if seed > 0:
        set_seed(seed)
    rng = RandomNumberGenerator()
    p = 1 / (1 + np.exp(-d["x1"].to_numpy()))
    ybin = rng.binomial(1, p).astype(np.int8)
    return d.with_columns(pl.Series("ybin", ybin))


mi_logit = mi_ses_from_stata(
    df_implicates=[_with_ybin(d, seed) for seed, d in enumerate(df_implicates)],
    command="logit ybin x1 x2",
    round_output=False,
)
mi_logit.print(round_output=False)


# %%
logger.info("\n\nPart 4: if you need something below stata_adapter's command-string")
logger.info("level - a custom pre/post-processing step around the .dta write, or a")
logger.info("different way of pulling results back than e(b)/e(V)/r(table) - the")
logger.info("primitives in _stata_interop are the same ones stata_adapter itself is")
logger.info("built from, and you can call them directly:")
logger.info("  - dataframe_to_dta(df, path)        - write via polars_readstat")
logger.info("  - require_pystata(edition, stata_path) - get a live Stata session")
logger.info("  - run_stata_model(df, command, ...) - the full write+use+run+extract")
logger.info("                                         stata_adapter wraps")
logger.info("  - run_stata_results(df, command, results, ...) - like")
logger.info("                                         run_stata_model, but for ANY")
logger.info("                                         named r()/e() result rather")
logger.info("                                         than the fixed e(b)/e(V)/")
logger.info("                                         r(table) triplet")
logger.info("")
logger.info("A custom example - grabbing a scalar result that isn't part of")
logger.info("e(b)/e(V)/r(table) at all (say, e(N) or e(r2)):")

raw = _st.run_stata_results(
    df_implicates[0],
    command="regress y x1 x2",
    results=["e(N)", "e(r2)"],
)
logger.info(raw)

logger.info("")
logger.info("adapters.stata_results_adapter wraps run_stata_results the same way")
logger.info("stata_adapter wraps run_stata_model - it works for ANY command")
logger.info("(r-class or e-class) and any named result, not just an e-class fit's")
logger.info("coefficient table. It's shaped as a StatCalculator.from_function")
logger.info("delegate (point estimates only, no vcov - the SE comes from")
logger.info("resampling across replicate weights, not Stata's own e(V)) - it's what")
logger.info("mi_ses_from_stata uses under the hood for the replicates= case above.")


# %%
logger.info("\n\nSummary / troubleshooting checklist for getting this working:")
logger.info("  1. check_stata_setup(stata_path=...) - confirms pystata and")
logger.info("     polars_readstat both import; doesn't guarantee config.init()")
logger.info("     succeeds too.")
logger.info("  2. If a command errors immediately with just a bare 'r(####);' and no")
logger.info("     explanation, pass quietly=False to stata_adapter/run_stata_model/")
logger.info("     run_stata_results to see Stata's own error text in the console.")
logger.info("  3. If e(b)/e(V) come back empty, the command you ran probably isn't")
logger.info("     e-class (didn't leave results in e()) - `ereturn list` right after")
logger.info("     running it interactively in Stata will confirm.")
logger.info("  4. If r(table) shapes/names look different than expected, `matrix list")
logger.info("     r(table)` interactively will show you its real shape.")

view in separate window if available.