marginaleffects 0.5.1__tar.gz → 0.6.1__tar.gz

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

@@ -1,11 +1,12 @@
 Metadata-Version: 2.4
 Name: marginaleffects
-Version: 0.5.1
+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
 Requires-Dist: narwhals>=1.34.0
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: patsy>=1.0.1
@@ -14,8 +15,6 @@ Requires-Dist: pydantic>=2.10.3
 Requires-Dist: plotnine>=0.14.5
 Requires-Dist: scipy>=1.14.1
 Requires-Dist: pyarrow>=19.0.1
-Provides-Extra: autodiff
-Requires-Dist: jax>=0.4.0; extra == "autodiff"
 Provides-Extra: test
 Requires-Dist: duckdb>=1.1.2; extra == "test"
 Requires-Dist: matplotlib>=3.7.1; extra == "test"

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.5.1 → marginaleffects-0.6.1}/marginaleffects/autodiff/__init__.py

@@ -29,18 +29,14 @@ if _JAX_AVAILABLE:
 
     jax.config.update("jax_enable_x64", True)
 
-    # Import submodules to make them accessible
-    from . import linear as linear
     from . import glm as glm
+    from . import pipeline as pipeline
 
-    # Re-export types for convenience
-    from .comparisons import ComparisonType as ComparisonType
     from .glm.families import Family as Family, Link as Link
 
     __all__ = [
-        "linear",
         "glm",
-        "ComparisonType",
+        "pipeline",
         "Family",
         "Link",
     ]
@@ -50,13 +46,8 @@ else:
         def __getattr__(self, name):
             _raise_jax_error()
 
-    linear = _DummyModule()
     glm = _DummyModule()
-
-    # Create dummy enums that raise errors
-    class ComparisonType:
-        def __getattribute__(self, name):
-            _raise_jax_error()
+    pipeline = _DummyModule()
 
     class Family:
         def __getattribute__(self, name):
@@ -66,4 +57,4 @@ else:
         def __getattribute__(self, name):
             _raise_jax_error()
 
-    __all__ = []
+    __all__ = ["glm", "pipeline", "Family", "Link"]

marginaleffects-0.6.1/marginaleffects/autodiff/glm/__init__.py

@@ -0,0 +1 @@
+from . import families as families

marginaleffects-0.6.1/marginaleffects/autodiff/lower.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import warnings
+
+import numpy as np
+
+from ..settings import is_autodiff_enabled, is_autodiff_forced
+from ..uncertainty import get_se
+from .ops import COMPARISON_OPS
+
+
+@dataclass
+class Lowered:
+    ok: bool
+    kwargs: dict | None = None
+    reason: str = ""
+
+
+@dataclass
+class AutodiffResult:
+    std_error: np.ndarray
+    jacobian: np.ndarray
+
+
+def _fail(reason: str) -> Lowered:
+    return Lowered(False, kwargs=None, reason=reason)
+
+
+def _model_args(model):
+    args = model.get_autodiff_args()
+    if args is None:
+        return None, _fail("")
+    if isinstance(args, str):
+        return None, _fail(args)
+    return args, None
+
+
+def _coefs_or_failure(model):
+    coefs = np.asarray(model.get_coef(), dtype=float).reshape(-1)
+    if np.isnan(coefs).any():
+        return None, _fail("models with NA coefficients")
+    return coefs, None
+
+
+def _design_or_failure(X, coefs, n_pred):
+    X = np.asarray(X, dtype=float)
+    if X.ndim != 2 or X.shape[1] != coefs.size or X.shape[0] != n_pred:
+        return None, _fail("this model/data configuration")
+    return X, None
+
+
+def _hypothesis_or_failure(plan):
+    if plan.hyp is None:
+        return None, None
+    if plan.hyp.kind != "matrix":
+        return None, _fail("this form of the `hypothesis` argument")
+    return np.asarray(plan.hyp.H, dtype=float), None
+
+
+def _has_nan(x) -> bool:
+    return x is not None and np.isnan(np.asarray(x, dtype=float)).any()
+
+
+def lower_predictions(plan, model) -> Lowered:
+    args, failure = _model_args(model)
+    if failure is not None:
+        return failure
+    coefs, failure = _coefs_or_failure(model)
+    if failure is not None:
+        return failure
+    X, failure = _design_or_failure(plan.exog, coefs, plan.n_pred)
+    if failure is not None:
+        return failure
+    if plan.align is not None:
+        return _fail("models with grouped/multi-equation outcomes")
+    if plan.has_na:
+        return _fail("missing values in predictions")
+    H, failure = _hypothesis_or_failure(plan)
+    if failure is not None:
+        return failure
+
+    agg_segments = None
+    agg_num_segments = None
+    agg_weights = None
+    if plan.agg is not None:
+        agg_segments = np.empty(plan.n_pred, dtype=np.int32)
+        agg_weights = np.ones(plan.n_pred, dtype=float)
+        for i, group in enumerate(plan.agg):
+            agg_segments[group.idx] = i
+            if group.w is not None:
+                if _has_nan(group.w):
+                    return _fail("missing values in weights")
+                agg_weights[group.idx] = np.asarray(group.w, dtype=float)
+        agg_num_segments = len(plan.agg)
+
+    kwargs = {
+        **args,
+        "X": X,
+        "agg_segments": agg_segments,
+        "agg_num_segments": agg_num_segments,
+        "agg_weights": agg_weights if agg_segments is not None else None,
+        "H": H,
+    }
+    return Lowered(True, kwargs=kwargs)
+
+
+def lower_comparisons(plan, model) -> Lowered:
+    args, failure = _model_args(model)
+    if failure is not None:
+        return failure
+    coefs, failure = _coefs_or_failure(model)
+    if failure is not None:
+        return failure
+    X_hi, failure = _design_or_failure(plan.exog_hi, coefs, plan.n_pred)
+    if failure is not None:
+        return failure
+    X_lo, failure = _design_or_failure(plan.exog_lo, coefs, plan.n_pred)
+    if failure is not None:
+        return failure
+    if plan.align is not None:
+        return _fail("models with grouped/multi-equation outcomes")
+    if plan.has_na:
+        return _fail("missing values in predictions")
+
+    if any(group.fun_key is None for group in plan.groups):
+        return _fail("custom comparison functions")
+    if plan.need_y:
+        return _fail("elasticities")
+    for group in plan.groups:
+        if group.fun_key not in COMPARISON_OPS:
+            return _fail(f"comparison='{group.fun_key}'")
+
+    H, failure = _hypothesis_or_failure(plan)
+    if failure is not None:
+        return failure
+
+    if not plan.groups:
+        order = np.asarray([], dtype=int)
+    else:
+        order = np.concatenate([group.idx for group in plan.groups]).astype(int)
+    if order.size != plan.n_pred:
+        return _fail("this model/data configuration")
+
+    ops = []
+    for group in plan.groups:
+        spec = COMPARISON_OPS[group.fun_key]
+        w = None
+        if spec.weighted:
+            if _has_nan(group.w):
+                return _fail("missing values in weights")
+            w = None if group.w is None else np.asarray(group.w, dtype=float)
+        ops.append({"op": spec.pipeline_op, "n": len(group.idx), "w": w})
+
+    kwargs = {
+        **args,
+        "X_hi": X_hi[order],
+        "X_lo": X_lo[order],
+        "ops": ops,
+        "H": H,
+    }
+    return Lowered(True, kwargs=kwargs)
+
+
+def _warn_unsupported(reason):
+    if reason:
+        warnings.warn(
+            "Automatic differentiation does not support "
+            f"{reason}. Reverting to finite differences.",
+            UserWarning,
+            stacklevel=3,
+        )
+
+
+def autodiff_try(plan, model, V, estimate, kind):
+    if plan is None or V is None or not is_autodiff_enabled():
+        return None
+
+    warn_on_fallback = is_autodiff_forced()
+    lowered = (
+        lower_predictions(plan, model)
+        if kind == "predictions"
+        else lower_comparisons(plan, model)
+    )
+    if not lowered.ok:
+        if warn_on_fallback:
+            _warn_unsupported(lowered.reason)
+        return None
+
+    try:
+        from . import pipeline
+
+        result = pipeline.compute(beta=model.get_coef(), **lowered.kwargs)
+    except Exception as exc:
+        if warn_on_fallback:
+            warnings.warn(
+                "Automatic differentiation failed "
+                f"({exc}). Reverting to finite differences.",
+                UserWarning,
+                stacklevel=3,
+            )
+        return None
+
+    estimate = np.asarray(estimate, dtype=float).reshape(-1)
+    if not np.allclose(result["estimate"], estimate, rtol=1e-8, atol=1e-8):
+        if warn_on_fallback:
+            warnings.warn(
+                "Automatic differentiation estimates did not match the standard "
+                "pipeline. Reverting to finite differences.",
+                UserWarning,
+                stacklevel=3,
+            )
+        return None
+
+    # Coef/vcov positional alignment is guaranteed by each adapter vault.
+    J = np.asarray(result["jacobian"], dtype=float)
+    se = get_se(J, V)
+    se[se == 0] = np.nan
+    return AutodiffResult(std_error=se, jacobian=J)

marginaleffects-0.6.1/marginaleffects/autodiff/ops.py

@@ -0,0 +1,60 @@
+"""Shared comparison operation registry for autodiff lowering and execution."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+
+ArrayFn = Callable[[list[Any]], Any]
+WMeanFn = Callable[[Any, Any | None], Any]
+EstimateFn = Callable[[Any, Any, Any | None, ArrayFn, WMeanFn], Any]
+
+
+@dataclass(frozen=True)
+class PipelineOp:
+    estimate: EstimateFn
+    scalar: bool
+
+    def output_size(self, n: int) -> int:
+        return 1 if self.scalar else n
+
+
+@dataclass(frozen=True)
+class ComparisonOp:
+    pipeline_op: str
+    weighted: bool
+
+
+def _difference(hi, lo, _w, _array, _wmean):
+    return hi - lo
+
+
+def _ratio(hi, lo, _w, _array, _wmean):
+    return hi / lo
+
+
+def _differenceavg(hi, lo, w, array, wmean):
+    return array([wmean(hi, w) - wmean(lo, w)])
+
+
+def _ratioavg(hi, lo, w, array, wmean):
+    return array([wmean(hi, w) / wmean(lo, w)])
+
+
+PIPELINE_OPS = {
+    "difference": PipelineOp(estimate=_difference, scalar=False),
+    "ratio": PipelineOp(estimate=_ratio, scalar=False),
+    "differenceavg": PipelineOp(estimate=_differenceavg, scalar=True),
+    "ratioavg": PipelineOp(estimate=_ratioavg, scalar=True),
+}
+
+COMPARISON_OPS = {
+    "difference": ComparisonOp(pipeline_op="difference", weighted=False),
+    "ratio": ComparisonOp(pipeline_op="ratio", weighted=False),
+    "differenceavg": ComparisonOp(pipeline_op="differenceavg", weighted=False),
+    "ratioavg": ComparisonOp(pipeline_op="ratioavg", weighted=False),
+    "differenceavgwts": ComparisonOp(pipeline_op="differenceavg", weighted=True),
+    "ratioavgwts": ComparisonOp(pipeline_op="ratioavg", weighted=True),
+}