marginaleffects 0.1.5__tar.gz → 0.2.1__tar.gz

{marginaleffects-0.1.5 → marginaleffects-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: marginaleffects
-Version: 0.1.5
+Version: 0.2.1
 Summary: Predictions, counterfactual comparisons, slopes, and hypothesis tests for statistical models.
 License-Expression: GPL-3.0-or-later
 Requires-Python: >=3.10
@@ -14,6 +14,8 @@ 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.1.5 → marginaleffects-0.2.1}/marginaleffects/__init__.py

@@ -10,6 +10,16 @@ from .plot_slopes import plot_slopes
 from .predictions import avg_predictions, predictions
 from .slopes import avg_slopes, slopes
 from .utils import get_dataset
+from .result import MarginaleffectsResult
+
+# Conditionally import autodiff module if JAX is available
+try:
+    from . import autodiff
+
+    _AUTODIFF_AVAILABLE = True
+except ImportError:
+    _AUTODIFF_AVAILABLE = False
+    autodiff = None
 
 __all__ = [
     "avg_comparisons",
@@ -27,4 +37,8 @@ __all__ = [
     "fit_sklearn",
     "fit_linearmodels",
     "get_dataset",
+    "MarginaleffectsResult",
 ]
+
+if _AUTODIFF_AVAILABLE:
+    __all__.append("autodiff")

marginaleffects-0.2.1/marginaleffects/autodiff/__init__.py

@@ -0,0 +1,63 @@
+"""
+JAX-based automatic differentiation module for marginal effects.
+
+This module provides high-performance computation of predictions, comparisons,
+and standard errors using JAX's automatic differentiation capabilities.
+
+Note: This module requires JAX to be installed. Install with:
+    pip install marginaleffects[autodiff]
+    or
+    pip install jax
+"""
+
+import importlib.util
+
+_JAX_AVAILABLE = importlib.util.find_spec("jax") is not None
+
+
+def _raise_jax_error():
+    raise ImportError(
+        "The autodiff module requires JAX to be installed. "
+        "Install with: pip install marginaleffects[autodiff] or pip install jax"
+    )
+
+
+if _JAX_AVAILABLE:
+    # Import submodules to make them accessible
+    from . import linear as linear
+    from . import glm as glm
+
+    # 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",
+        "Family",
+        "Link",
+    ]
+else:
+    # Create dummy module objects that raise helpful errors
+    class _DummyModule:
+        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()
+
+    class Family:
+        def __getattribute__(self, name):
+            _raise_jax_error()
+
+    class Link:
+        def __getattribute__(self, name):
+            _raise_jax_error()
+
+    __all__ = []

marginaleffects-0.2.1/marginaleffects/autodiff/comparisons.py

@@ -0,0 +1,58 @@
+"""Comparison types and functions using enum-based approach for JAX compatibility."""
+
+import jax.numpy as jnp
+from jax import lax
+from enum import IntEnum
+
+
+class ComparisonType(IntEnum):
+    """Comparison types for marginal effects."""
+
+    DIFFERENCE = 0
+    RATIO = 1
+    LNRATIO = 2
+    LNOR = 3
+    LIFT = 4
+    DIFFERENCEAVG = 5
+
+
+def _compute_comparison_vector(
+    comparison_type: int, pred_hi: jnp.ndarray, pred_lo: jnp.ndarray
+) -> jnp.ndarray:
+    """Apply comparison function element-wise (returns N-length array)."""
+    return lax.switch(
+        comparison_type,
+        [
+            lambda hi, lo: hi - lo,  # difference
+            lambda hi, lo: hi / lo,  # ratio
+            lambda hi, lo: jnp.log(hi / lo),  # lnratio
+            lambda hi, lo: jnp.log((hi / (1 - hi)) / (lo / (1 - lo))),  # lnor
+            lambda hi, lo: (hi - lo) / lo,  # lift
+        ],
+        pred_hi,
+        pred_lo,
+    )
+
+
+def _compute_comparison_scalar(
+    comparison_type: int, pred_hi: jnp.ndarray, pred_lo: jnp.ndarray
+) -> jnp.ndarray:
+    """Apply comparison function with aggregation (returns scalar)."""
+    return lax.switch(
+        comparison_type,
+        [
+            lambda hi, lo: jnp.mean(hi) - jnp.mean(lo),  # differenceavg
+        ],
+        pred_hi,
+        pred_lo,
+    )
+
+
+def _compute_comparison(
+    comparison_type: int, pred_hi: jnp.ndarray, pred_lo: jnp.ndarray
+) -> jnp.ndarray:
+    """Apply comparison function element-wise (returns N-length array).
+
+    Note: DIFFERENCEAVG should use _compute_comparison_scalar directly.
+    """
+    return _compute_comparison_vector(comparison_type, pred_hi, pred_lo)

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

@@ -0,0 +1,5 @@
+# GLM model implementations with switchable link functions
+
+from . import predictions as predictions
+from . import comparisons as comparisons
+from . import families as families

marginaleffects-0.2.1/marginaleffects/autodiff/glm/comparisons.py

@@ -0,0 +1,193 @@
+import jax.numpy as jnp
+import numpy as np
+from jax import jit, jacrev
+from .families import linkinv, resolve_link
+from ..comparisons import (
+    _compute_comparison,
+    _compute_comparison_scalar,
+    ComparisonType,
+)
+from ..utils import group_reducer, standard_errors
+
+
+def _comparison_core(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    comparison_type: int,
+    family_type: int,
+    link_type: int,
+) -> jnp.ndarray:
+    """Core comparison function - single source of truth for comparison vector computation."""
+    pred_hi = linkinv(link_type, X_hi @ beta)
+    pred_lo = linkinv(link_type, X_lo @ beta)
+    return _compute_comparison(comparison_type, pred_hi, pred_lo)
+
+
+@jit
+def _comparison_byT(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    comparison_type: int,
+    family_type: int,
+    link_type: int = None,
+) -> jnp.ndarray:
+    comp = _comparison_core(beta, X_hi, X_lo, comparison_type, family_type, link_type)
+    return jnp.mean(comp)
+
+
+def _comparison_byG(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    groups: jnp.ndarray,
+    num_groups: int,
+    comparison_type: int,
+    family_type: int,
+    link_type: int = None,
+) -> jnp.ndarray:
+    comp = _comparison_core(beta, X_hi, X_lo, comparison_type, family_type, link_type)
+    return group_reducer(comp, groups, num_groups)
+
+
+@jit
+def _comparisons_core(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    comparison_type: int,
+    family_type: int,
+    link_type: int = None,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    comp = _comparison_core(beta, X_hi, X_lo, comparison_type, family_type, link_type)
+    jac = jacrev(
+        lambda b: _comparison_core(
+            b, X_hi, X_lo, comparison_type, family_type, link_type
+        )
+    )(beta)
+    return comp, jac
+
+
+def comparisons(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    vcov: jnp.ndarray,
+    comparison_type: int,
+    family_type: int,
+    link_type: int = None,
+) -> dict[str, np.ndarray]:
+    link_type = resolve_link(family_type, link_type)
+
+    # Handle DIFFERENCEAVG separately (returns scalar)
+    if comparison_type == ComparisonType.DIFFERENCEAVG:
+
+        @jit
+        def _scalar_core(b, X_h, X_l, lt):
+            pred_hi = linkinv(lt, X_h @ b)
+            pred_lo = linkinv(lt, X_l @ b)
+            return _compute_comparison_scalar(0, pred_hi, pred_lo)
+
+        comp = _scalar_core(beta, X_hi, X_lo, link_type)
+        jac = jacrev(lambda b: _scalar_core(b, X_hi, X_lo, link_type))(beta)
+        se = standard_errors(jac.reshape(1, -1), vcov)
+        return {
+            "estimate": np.array(comp),
+            "jacobian": np.array(jac),
+            "std_error": se[0],
+        }
+
+    # Handle element-wise comparisons
+    comp, jac = _comparisons_core(
+        beta, X_hi, X_lo, comparison_type, family_type, link_type
+    )
+    se = standard_errors(jac, vcov)
+    return {
+        "estimate": np.array(comp),
+        "jacobian": np.array(jac),
+        "std_error": se,
+    }
+
+
+@jit
+def _comparisons_byT_core(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    comparison_type: int,
+    family_type: int,
+    link_type: int = None,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    comp = _comparison_byT(beta, X_hi, X_lo, comparison_type, family_type, link_type)
+    jac = jacrev(
+        lambda b: _comparison_byT(
+            b, X_hi, X_lo, comparison_type, family_type, link_type
+        )
+    )(beta)
+    return comp, jac
+
+
+def comparisons_byT(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    vcov: jnp.ndarray,
+    comparison_type: int,
+    family_type: int,
+    link_type: int = None,
+) -> dict[str, np.ndarray]:
+    link_type = resolve_link(family_type, link_type)
+    comp, jac = _comparisons_byT_core(
+        beta, X_hi, X_lo, comparison_type, family_type, link_type
+    )
+    se = standard_errors(jac.reshape(1, -1), vcov)
+    return {
+        "estimate": np.array(comp),
+        "jacobian": np.array(jac),
+        "std_error": se[0],
+    }
+
+
+def _comparisons_byG_core(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    groups: jnp.ndarray,
+    num_groups: int,
+    comparison_type: int,
+    family_type: int,
+    link_type: int = None,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    comp = _comparison_byG(
+        beta, X_hi, X_lo, groups, num_groups, comparison_type, family_type, link_type
+    )
+    jac = jacrev(
+        lambda b: _comparison_byG(
+            b, X_hi, X_lo, groups, num_groups, comparison_type, family_type, link_type
+        )
+    )(beta)
+    return comp, jac
+
+
+def comparisons_byG(
+    beta: jnp.ndarray,
+    X_hi: jnp.ndarray,
+    X_lo: jnp.ndarray,
+    vcov: jnp.ndarray,
+    groups: jnp.ndarray,
+    num_groups: int,
+    comparison_type: int,
+    family_type: int,
+    link_type: int = None,
+) -> dict[str, np.ndarray]:
+    link_type = resolve_link(family_type, link_type)
+    comp, jac = _comparisons_byG_core(
+        beta, X_hi, X_lo, groups, num_groups, comparison_type, family_type, link_type
+    )
+    se = standard_errors(jac, vcov)
+    return {
+        "estimate": np.array(comp),
+        "jacobian": np.array(jac),
+        "std_error": se,
+    }

marginaleffects-0.2.1/marginaleffects/autodiff/glm/families.py

@@ -0,0 +1,108 @@
+"""GLM families and link functions using enum-based approach for JAX compatibility."""
+
+import jax.numpy as jnp
+from jax import lax
+from jax.scipy.stats import norm
+from enum import IntEnum
+
+
+class Family(IntEnum):
+    """GLM family types."""
+
+    GAUSSIAN = 0
+    BINOMIAL = 1
+    POISSON = 2
+    GAMMA = 3
+    INVERSE_GAUSSIAN = 4
+
+
+class Link(IntEnum):
+    """Link function types."""
+
+    IDENTITY = 0
+    LOG = 1
+    LOGIT = 2
+    PROBIT = 3
+    INVERSE = 4
+    SQRT = 5
+    CLOGLOG = 6
+
+
+# Default links for each family
+DEFAULT_LINKS = {
+    Family.GAUSSIAN: Link.IDENTITY,
+    Family.BINOMIAL: Link.LOGIT,
+    Family.POISSON: Link.LOG,
+    Family.GAMMA: Link.INVERSE,
+    Family.INVERSE_GAUSSIAN: Link.INVERSE,
+}
+
+
+def linkinv(link_type: int, eta: jnp.ndarray) -> jnp.ndarray:
+    """Inverse link function: transform linear predictor to mean."""
+    return lax.switch(
+        link_type,
+        [
+            lambda x: x,  # identity
+            lambda x: jnp.exp(x),  # log
+            lambda x: 1 / (1 + jnp.exp(-x)),  # logit
+            lambda x: norm.cdf(x),  # probit
+            lambda x: 1.0 / x,  # inverse
+            lambda x: x**2,  # sqrt
+            lambda x: 1 - jnp.exp(-jnp.exp(x)),  # cloglog
+        ],
+        eta,
+    )
+
+
+def linkfun(link_type: int, mu: jnp.ndarray) -> jnp.ndarray:
+    """Link function: transform mean to linear predictor."""
+    return lax.switch(
+        link_type,
+        [
+            lambda x: x,  # identity
+            lambda x: jnp.log(x),  # log
+            lambda x: jnp.log(x / (1 - x)),  # logit
+            lambda x: norm.ppf(x),  # probit
+            lambda x: 1.0 / x,  # inverse
+            lambda x: jnp.sqrt(x),  # sqrt
+            lambda x: jnp.log(-jnp.log(1 - x)),  # cloglog
+        ],
+        mu,
+    )
+
+
+# Valid link functions for each family
+VALID_LINKS = {
+    Family.GAUSSIAN: [Link.IDENTITY, Link.LOG, Link.INVERSE],
+    Family.BINOMIAL: [Link.LOGIT, Link.PROBIT, Link.CLOGLOG, Link.LOG],
+    Family.POISSON: [Link.LOG, Link.IDENTITY, Link.SQRT],
+    Family.GAMMA: [Link.INVERSE, Link.IDENTITY, Link.LOG],
+    Family.INVERSE_GAUSSIAN: [Link.INVERSE, Link.IDENTITY, Link.LOG],
+}
+
+
+def validate_family_link(family_type: int, link_type: int) -> bool:
+    """Check if link function is valid for the given family."""
+    if family_type not in VALID_LINKS:
+        return False
+    return link_type in VALID_LINKS[family_type]
+
+
+def resolve_link(family_type: int, link_type: int = None) -> int:
+    """Resolve link type, using default if None, and validate the combination."""
+    if link_type is None:
+        return DEFAULT_LINKS.get(family_type, Link.IDENTITY)
+    if not validate_family_link(family_type, link_type):
+        default = DEFAULT_LINKS.get(family_type, Link.IDENTITY)
+        raise ValueError(
+            f"Invalid link {link_type} for family {family_type}. "
+            f"Valid links: {VALID_LINKS.get(family_type, [])}. "
+            f"Using default: {default}"
+        )
+    return link_type
+
+
+# Convenience instances for direct use
+family = Family
+link = Link

marginaleffects-0.2.1/marginaleffects/autodiff/glm/predictions.py

@@ -0,0 +1,131 @@
+import jax.numpy as jnp
+import numpy as np
+from jax import jacfwd, jacrev, jit
+from .families import linkinv, resolve_link
+from ..utils import group_reducer, standard_errors
+
+
+def _predict_core(
+    beta: jnp.ndarray,
+    X: jnp.ndarray,
+    family_type: int,
+    link_type: int,
+) -> jnp.ndarray:
+    """Core prediction function - single source of truth for prediction computation."""
+    linear_pred = X @ beta
+    return linkinv(link_type, linear_pred)
+
+
+@jit
+def _predict(
+    beta: jnp.ndarray, X: jnp.ndarray, family_type: int, link_type: int = None
+) -> jnp.ndarray:
+    return _predict_core(beta, X, family_type, link_type)
+
+
+@jit
+def _predict_byT(
+    beta: jnp.ndarray, X: jnp.ndarray, family_type: int, link_type: int = None
+) -> jnp.ndarray:
+    pred = _predict_core(beta, X, family_type, link_type)
+    return jnp.mean(pred)
+
+
+def _predict_byG(
+    beta: jnp.ndarray,
+    X: jnp.ndarray,
+    groups: jnp.ndarray,
+    num_groups: int,
+    family_type: int,
+    link_type: int = None,
+) -> jnp.ndarray:
+    pred = _predict_core(beta, X, family_type, link_type)
+    return group_reducer(pred, groups, num_groups)
+
+
+@jit
+def _predictions_core(
+    beta: jnp.ndarray, X: jnp.ndarray, family_type: int, link_type: int = None
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    pred = _predict_core(beta, X, family_type, link_type)
+    jac = jacfwd(_predict, argnums=0)(beta, X, family_type, link_type)
+    return pred, jac
+
+
+def predictions(
+    beta: jnp.ndarray,
+    X: jnp.ndarray,
+    vcov: jnp.ndarray,
+    family_type: int,
+    link_type: int = None,
+) -> dict[str, np.ndarray]:
+    link_type = resolve_link(family_type, link_type)
+    pred, jac = _predictions_core(beta, X, family_type, link_type)
+    se = standard_errors(jac, vcov)
+    return {
+        "estimate": np.array(pred),
+        "jacobian": np.array(jac),
+        "std_error": se,
+    }
+
+
+@jit
+def _predictions_byT_core(
+    beta: jnp.ndarray, X: jnp.ndarray, family_type: int, link_type: int = None
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    pred = _predict_byT(beta, X, family_type, link_type)
+    jac = jacrev(_predict_byT, argnums=0)(beta, X, family_type, link_type)
+    return pred, jac
+
+
+def predictions_byT(
+    beta: jnp.ndarray,
+    X: jnp.ndarray,
+    vcov: jnp.ndarray,
+    family_type: int,
+    link_type: int = None,
+) -> dict[str, np.ndarray]:
+    link_type = resolve_link(family_type, link_type)
+    pred, jac = _predictions_byT_core(beta, X, family_type, link_type)
+    se = standard_errors(jac.reshape(1, -1), vcov)
+    return {
+        "estimate": np.array(pred),
+        "jacobian": np.array(jac),
+        "std_error": se[0],
+    }
+
+
+def _predictions_byG_core(
+    beta: jnp.ndarray,
+    X: jnp.ndarray,
+    groups: jnp.ndarray,
+    num_groups: int,
+    family_type: int,
+    link_type: int = None,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    pred = _predict_byG(beta, X, groups, num_groups, family_type, link_type)
+    jac = jacrev(
+        lambda b: _predict_byG(b, X, groups, num_groups, family_type, link_type)
+    )(beta)
+    return pred, jac
+
+
+def predictions_byG(
+    beta: jnp.ndarray,
+    X: jnp.ndarray,
+    vcov: jnp.ndarray,
+    groups: jnp.ndarray,
+    num_groups: int,
+    family_type: int,
+    link_type: int = None,
+) -> dict[str, np.ndarray]:
+    link_type = resolve_link(family_type, link_type)
+    pred, jac = _predictions_byG_core(
+        beta, X, groups, num_groups, family_type, link_type
+    )
+    se = standard_errors(jac, vcov)
+    return {
+        "estimate": np.array(pred),
+        "jacobian": np.array(jac),
+        "std_error": se,
+    }

marginaleffects-0.2.1/marginaleffects/autodiff/linear/__init__.py

@@ -0,0 +1,4 @@
+# Linear model implementations
+
+from . import predictions as predictions
+from . import comparisons as comparisons