marginaleffects 0.6.0__tar.gz → 0.6.1__tar.gz

{marginaleffects-0.6.0 → marginaleffects-0.6.1}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: marginaleffects
-Version: 0.6.0
+Version: 0.6.1
 Summary: Predictions, counterfactual comparisons, slopes, and hypothesis tests for statistical models.
 License-Expression: GPL-3.0-or-later
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 Requires-Dist: formulaic>=1.0.2
 Requires-Dist: jax>=0.4.0

marginaleffects-0.6.1/benchmarks/benchmark_autodiff.py

@@ -0,0 +1,228 @@
+"""
+Benchmark comparing autodiff (JAX) vs finite differences.
+
+Run with:
+    uv run --all-extras python benchmarks/benchmark_autodiff.py
+"""
+
+import time
+
+import numpy as np
+import polars as pl
+import statsmodels.api as sm
+from marginaleffects import (
+    autodiff,
+    avg_comparisons,
+    avg_predictions,
+    comparisons,
+    predictions,
+)
+from marginaleffects.uncertainty import get_jacobian, get_se
+
+
+def generate_data(
+    n_obs: int = 5000, n_predictors: int = 50, seed: int = 42
+) -> pl.DataFrame:
+    """Generate simulated data."""
+    rng = np.random.default_rng(seed)
+    X = rng.normal(size=(n_obs, n_predictors))
+    beta = rng.normal(size=n_predictors + 1)
+    y = beta[0] + X @ beta[1:] + rng.normal(size=n_obs) * 0.5
+    data = {"y": y}
+    for i in range(n_predictors):
+        data[f"x{i}"] = X[:, i]
+    return pl.DataFrame(data)
+
+
+def _time_call(func, n_runs: int):
+    out = []
+    for _ in range(n_runs):
+        start = time.perf_counter()
+        func()
+        out.append(time.perf_counter() - start)
+    return out
+
+
+def benchmark_public_api(func, model, n_runs: int = 10, **kwargs):
+    """
+    Benchmark public API calls with and without autodiff.
+
+    This includes shared overhead such as model sanitization, formula/design
+    matrix construction, vcov retrieval, confidence intervals, and result
+    formatting.
+    """
+    results = {"autodiff": [], "finite_diff": []}
+
+    autodiff(True)
+    func(model, **kwargs)
+    autodiff(False)
+    func(model, **kwargs)
+
+    autodiff(True)
+    results["autodiff"] = _time_call(lambda: func(model, **kwargs), n_runs=n_runs)
+
+    autodiff(False)
+    results["finite_diff"] = _time_call(lambda: func(model, **kwargs), n_runs=n_runs)
+
+    autodiff(None)
+    return results
+
+
+def _engine_cases(X: np.ndarray, beta: np.ndarray):
+    from marginaleffects.autodiff import pipeline
+
+    V = np.eye(beta.size)
+    n_obs = X.shape[0]
+    segments = np.zeros(n_obs, dtype=np.int32)
+
+    X_hi = X.copy()
+    X_lo = X.copy()
+    X_hi[:, 1] += 0.5
+
+    def jax_predictions():
+        result = pipeline.compute(beta=beta, model_type="linear", X=X)
+        return get_se(result["jacobian"], V)
+
+    def fd_predictions():
+        J = get_jacobian(lambda b: X @ b, beta)
+        return get_se(J, V)
+
+    def jax_avg_predictions():
+        result = pipeline.compute(
+            beta=beta,
+            model_type="linear",
+            X=X,
+            agg_segments=segments,
+            agg_num_segments=1,
+        )
+        return get_se(result["jacobian"], V)
+
+    def fd_avg_predictions():
+        J = get_jacobian(lambda b: np.array([(X @ b).mean()]), beta)
+        return get_se(J, V)
+
+    def jax_comparisons():
+        result = pipeline.compute(
+            beta=beta,
+            model_type="linear",
+            X_hi=X_hi,
+            X_lo=X_lo,
+            ops=[{"op": "difference", "n": n_obs, "w": None}],
+        )
+        return get_se(result["jacobian"], V)
+
+    def fd_comparisons():
+        J = get_jacobian(lambda b: (X_hi @ b) - (X_lo @ b), beta)
+        return get_se(J, V)
+
+    def jax_avg_comparisons():
+        result = pipeline.compute(
+            beta=beta,
+            model_type="linear",
+            X_hi=X_hi,
+            X_lo=X_lo,
+            ops=[{"op": "differenceavg", "n": n_obs, "w": None}],
+        )
+        return get_se(result["jacobian"], V)
+
+    def fd_avg_comparisons():
+        J = get_jacobian(
+            lambda b: np.array([np.mean((X_hi @ b) - (X_lo @ b))]),
+            beta,
+        )
+        return get_se(J, V)
+
+    return {
+        "OLS predictions()": (jax_predictions, fd_predictions),
+        "OLS avg_predictions()": (jax_avg_predictions, fd_avg_predictions),
+        "OLS comparisons()": (jax_comparisons, fd_comparisons),
+        "OLS avg_comparisons()": (jax_avg_comparisons, fd_avg_comparisons),
+    }
+
+
+def benchmark_engine(X: np.ndarray, beta: np.ndarray, n_runs: int = 20):
+    """
+    Benchmark only the Jacobian engine on cached design matrices.
+
+    This isolates the finite-difference loop from JAX autodiff, excluding public
+    API setup and result formatting.
+    """
+    out = {}
+    for name, (jax_func, fd_func) in _engine_cases(X, beta).items():
+        # Warm-up. The JAX call also compiles the traced function.
+        jax_func()
+        fd_func()
+        out[name] = {
+            "autodiff": _time_call(jax_func, n_runs=n_runs),
+            "finite_diff": _time_call(fd_func, n_runs=n_runs),
+        }
+    return out
+
+
+def print_results(name: str, results: dict):
+    """Print benchmark results."""
+    ad_mean = np.mean(results["autodiff"]) * 1000
+    fd_mean = np.mean(results["finite_diff"]) * 1000
+    speedup = fd_mean / ad_mean
+    print(f"{name:30s}  {ad_mean:6.1f} ms  {fd_mean:6.1f} ms  {speedup:.2f}x")
+
+
+if __name__ == "__main__":
+    try:
+        from marginaleffects.autodiff import _JAX_AVAILABLE
+
+        if not _JAX_AVAILABLE:
+            print("JAX not available")
+            exit(1)
+    except ImportError:
+        print("autodiff module not found")
+        exit(1)
+
+    data = generate_data()
+    formula = "y ~ " + " + ".join([f"x{i}" for i in range(50)])
+    X = data.select([f"x{i}" for i in range(50)]).to_numpy()
+    X = np.column_stack([np.ones(X.shape[0]), X])
+    beta = np.random.default_rng(42).normal(size=X.shape[1])
+
+    print("Autodiff Public API Benchmark (5000 obs, 50 predictors)")
+    print("Includes model wrapping, formula/design matrices, vcov, and formatting")
+    print("=" * 78)
+    print(f"{'Function':30s}  {'JAX':>8s}  {'Finite':>8s}  Speedup")
+    print("-" * 78)
+
+    ols = sm.OLS.from_formula(formula, data.to_pandas()).fit()
+    print_results("OLS predictions()", benchmark_public_api(predictions, ols))
+    print_results("OLS avg_predictions()", benchmark_public_api(avg_predictions, ols))
+    print_results(
+        "OLS comparisons()",
+        benchmark_public_api(comparisons, ols, variables="x0"),
+    )
+    print_results(
+        "OLS avg_comparisons()",
+        benchmark_public_api(avg_comparisons, ols, variables="x0"),
+    )
+
+    glm = sm.GLM.from_formula(
+        formula, data.to_pandas(), family=sm.families.Gaussian()
+    ).fit()
+    print_results("GLM predictions()", benchmark_public_api(predictions, glm))
+    print_results("GLM avg_predictions()", benchmark_public_api(avg_predictions, glm))
+    print_results(
+        "GLM comparisons()",
+        benchmark_public_api(comparisons, glm, variables="x0"),
+    )
+    print_results(
+        "GLM avg_comparisons()",
+        benchmark_public_api(avg_comparisons, glm, variables="x0"),
+    )
+    print("=" * 78)
+
+    print()
+    print("Autodiff Jacobian Engine Benchmark (cached 5000x51 design matrix)")
+    print("Excludes shared public API setup; row-level quantities have many outputs")
+    print("=" * 78)
+    print(f"{'Function':30s}  {'JAX':>8s}  {'Finite':>8s}  Speedup")
+    print("-" * 78)
+    for name, results in benchmark_engine(X, beta).items():
+        print_results(name, results)
+    print("=" * 78)

{marginaleffects-0.6.0 → marginaleffects-0.6.1}/marginaleffects/comparisons.py

@@ -20,6 +20,7 @@ from .plan import (
     ComparisonPlan,
     comparison_plan_apply,
     comparison_plan_predict,
+    plan_values_allclose,
 )
 from .utils import (
     get_pad,
@@ -546,12 +547,9 @@ def _comparisons_build(
 
     hi, lo, y = plan_predictions
     replay = comparison_plan_apply(plan, hi, lo, y if plan.need_y else None)
-    if not np.allclose(
+    if not plan_values_allclose(
         replay,
         tmp["estimate"].to_numpy(),
-        rtol=1e-12,
-        atol=1e-12,
-        equal_nan=True,
     ):
         raise RuntimeError(
             "marginaleffects internal error: comparison plan baseline check failed"

{marginaleffects-0.6.0 → marginaleffects-0.6.1}/marginaleffects/plan.py

@@ -98,6 +98,19 @@ def _estimand_to_numpy(x) -> np.ndarray:
     return _as_float_array(x)
 
 
+def plan_values_allclose(replay, target) -> bool:
+    """Compare plan replay values to baseline values at backend precision."""
+    replay = np.asarray(replay)
+    target = np.asarray(target)
+    tol = 1e-12
+    if any(
+        np.issubdtype(x.dtype, np.floating) and x.dtype.itemsize <= 4
+        for x in (replay, target)
+    ):
+        tol = 1e-6
+    return np.allclose(replay, target, rtol=tol, atol=tol, equal_nan=True)
+
+
 def prediction_plan_predict(plan: PredictionPlan, model, coefs) -> np.ndarray:
     out = model.get_predict(params=np.asarray(coefs), newdata=plan.exog)
     return _as_float_array(out["estimate"].to_numpy())

{marginaleffects-0.6.0 → marginaleffects-0.6.1}/marginaleffects/predictions.py

@@ -6,7 +6,12 @@ from .by import get_by_plan
 from .hypothesis_compile import hypothesis_compile
 from .uncertainty import get_jacobian, get_se
 from .classes import MarginaleffectsResult
-from .plan import PredictionPlan, prediction_plan_apply, prediction_plan_predict
+from .plan import (
+    PredictionPlan,
+    plan_values_allclose,
+    prediction_plan_apply,
+    prediction_plan_predict,
+)
 from .utils import prepare_base_inputs, finalize_result, call_avg
 from .sanitize import handle_deprecated_hypotheses_argument
 from .docstrings import doc
@@ -146,12 +151,9 @@ def _predictions_build(model, exog, newdata, by, wts, hypothesis):
     replay = prediction_plan_apply(
         plan, prediction_plan_predict(plan, model, model.get_coef())
     )
-    if not np.allclose(
+    if not plan_values_allclose(
         replay,
         out["estimate"].to_numpy(),
-        rtol=1e-12,
-        atol=1e-12,
-        equal_nan=True,
     ):
         raise RuntimeError(
             "marginaleffects internal error: prediction plan baseline check failed"

{marginaleffects-0.6.0 → marginaleffects-0.6.1}/marginaleffects.egg-info/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: marginaleffects
-Version: 0.6.0
+Version: 0.6.1
 Summary: Predictions, counterfactual comparisons, slopes, and hypothesis tests for statistical models.
 License-Expression: GPL-3.0-or-later
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 Requires-Dist: formulaic>=1.0.2
 Requires-Dist: jax>=0.4.0

{marginaleffects-0.6.0 → marginaleffects-0.6.1}/pyproject.toml

@@ -1,10 +1,10 @@
 [project]
 name = "marginaleffects"
-version = "0.6.0"
+version = "0.6.1"
 license = "GPL-3.0-or-later"
 description = "Predictions, counterfactual comparisons, slopes, and hypothesis tests for statistical models."
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 dependencies = [
     "formulaic >=1.0.2",
     "jax >=0.4.0",

{marginaleffects-0.6.0 → marginaleffects-0.6.1}/tests/test_sklearn.py

@@ -146,3 +146,46 @@ def test_categorical_formula_equivalence():
     branch_categorical = pred_categorical["branch"].to_list()
     branch_direct = pred_direct["branch"].to_list()
     assert branch_categorical == branch_direct
+
+
+def test_callable_selector_xgboost_float32_predictions_by():
+    """Callable sklearn selectors can replay float32 predictions with by groups."""
+    import polars.selectors as cs
+    from sklearn.compose import make_column_transformer
+    from sklearn.pipeline import make_pipeline
+    from sklearn.preprocessing import FunctionTransformer, OneHotEncoder
+    from xgboost import XGBRegressor
+
+    rng = np.random.default_rng(1)
+    data = pl.DataFrame(
+        {
+            "price": rng.normal(size=80),
+            "unit_type": np.where(np.arange(80) % 2 == 0, "a", "b"),
+            "neighborhood": np.where(np.arange(80) % 3 == 0, "n1", "n2"),
+            "bedrooms": rng.integers(1, 4, size=80),
+            "x": rng.normal(size=80),
+        }
+    )
+    catvar = data.select(~cs.numeric()).columns
+    data = data.with_columns([pl.col(c).cast(pl.Categorical) for c in catvar])
+
+    preprocessor = make_column_transformer(
+        (OneHotEncoder(), catvar),
+        remainder=FunctionTransformer(lambda x: x.to_numpy()),
+    )
+    pipeline = make_pipeline(
+        preprocessor,
+        XGBRegressor(n_estimators=10, max_depth=2, random_state=1),
+    )
+
+    def selector(data):
+        y = data.select(cs.by_name("price", require_all=False))
+        X = data.select(~cs.by_name("price", require_all=False))
+        return y, X
+
+    mod = me.fit_sklearn(selector, data=data[:60], engine=pipeline)
+    out = me.avg_predictions(mod, newdata=data[60:], by="unit_type")
+
+    assert isinstance(out, MarginaleffectsResult)
+    assert out.height == 2
+    assert set(out["unit_type"].to_list()) == {"a", "b"}

marginaleffects-0.6.0/benchmarks/benchmark_autodiff.py

@@ -1,111 +0,0 @@
-"""
-Benchmark comparing autodiff (JAX) vs finite differences.
-
-Run with:
-    uv run --all-extras python benchmarks/benchmark_autodiff.py
-"""
-
-import time
-import numpy as np
-import polars as pl
-import statsmodels.api as sm
-from marginaleffects import (
-    predictions,
-    avg_predictions,
-    comparisons,
-    avg_comparisons,
-    autodiff,
-)
-
-
-def generate_data(
-    n_obs: int = 5000, n_predictors: int = 50, seed: int = 42
-) -> pl.DataFrame:
-    """Generate simulated data."""
-    np.random.seed(seed)
-    X = np.random.randn(n_obs, n_predictors)
-    beta = np.random.randn(n_predictors + 1)
-    y = beta[0] + X @ beta[1:] + np.random.randn(n_obs) * 0.5
-    data = {"y": y}
-    for i in range(n_predictors):
-        data[f"x{i}"] = X[:, i]
-    return pl.DataFrame(data)
-
-
-def benchmark(func, model, n_runs: int = 10, **kwargs):
-    """Benchmark a function with and without autodiff."""
-    results = {"autodiff": [], "finite_diff": []}
-
-    # Warm-up
-    autodiff(True)
-    func(model, **kwargs)
-    autodiff(False)
-    func(model, **kwargs)
-
-    # Benchmark with autodiff
-    autodiff(True)
-    for _ in range(n_runs):
-        start = time.perf_counter()
-        func(model, **kwargs)
-        results["autodiff"].append(time.perf_counter() - start)
-
-    # Benchmark without autodiff
-    autodiff(False)
-    for _ in range(n_runs):
-        start = time.perf_counter()
-        func(model, **kwargs)
-        results["finite_diff"].append(time.perf_counter() - start)
-
-    autodiff(None)
-    return results
-
-
-def print_results(name: str, results: dict):
-    """Print benchmark results."""
-    ad_mean = np.mean(results["autodiff"]) * 1000
-    fd_mean = np.mean(results["finite_diff"]) * 1000
-    speedup = fd_mean / ad_mean
-    print(f"{name:30s}  {ad_mean:6.1f} ms  {fd_mean:6.1f} ms  {speedup:.2f}x")
-
-
-if __name__ == "__main__":
-    # Check JAX
-    try:
-        from marginaleffects.autodiff import _JAX_AVAILABLE
-
-        if not _JAX_AVAILABLE:
-            print("JAX not available")
-            exit(1)
-    except ImportError:
-        print("autodiff module not found")
-        exit(1)
-
-    print("Autodiff Benchmark (5000 obs, 50 predictors)")
-    print("=" * 65)
-    print(f"{'Function':30s}  {'JAX':>8s}  {'Finite':>8s}  Speedup")
-    print("-" * 65)
-
-    data = generate_data()
-    formula = "y ~ " + " + ".join([f"x{i}" for i in range(50)])
-
-    # OLS
-    ols = sm.OLS.from_formula(formula, data.to_pandas()).fit()
-    print_results("OLS predictions()", benchmark(predictions, ols))
-    print_results("OLS avg_predictions()", benchmark(avg_predictions, ols))
-    print_results("OLS comparisons()", benchmark(comparisons, ols, variables="x0"))
-    print_results(
-        "OLS avg_comparisons()", benchmark(avg_comparisons, ols, variables="x0")
-    )
-
-    # GLM
-    glm = sm.GLM.from_formula(
-        formula, data.to_pandas(), family=sm.families.Gaussian()
-    ).fit()
-    print_results("GLM predictions()", benchmark(predictions, glm))
-    print_results("GLM avg_predictions()", benchmark(avg_predictions, glm))
-    print_results("GLM comparisons()", benchmark(comparisons, glm, variables="x0"))
-    print_results(
-        "GLM avg_comparisons()", benchmark(avg_comparisons, glm, variables="x0")
-    )
-
-    print("=" * 65)