marginaleffects 0.3.2__tar.gz → 0.5.0__tar.gz

{marginaleffects-0.3.2 → marginaleffects-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: marginaleffects
-Version: 0.3.2
+Version: 0.5.0
 Summary: Predictions, counterfactual comparisons, slopes, and hypothesis tests for statistical models.
 License-Expression: GPL-3.0-or-later
 Requires-Python: >=3.10
@@ -22,6 +22,7 @@ Requires-Dist: matplotlib>=3.7.1; extra == "test"
 Requires-Dist: linearmodels>=6.1; extra == "test"
 Requires-Dist: pandas>=2.2.2; extra == "test"
 Requires-Dist: pre-commit>=4.2.0; extra == "test"
+Requires-Dist: numba; extra == "test"
 Requires-Dist: pyarrow>=17.0.0; extra == "test"
 Requires-Dist: pyfixest>=0.28.0; extra == "test"
 Requires-Dist: statsmodels>=0.14.0; extra == "test"

{marginaleffects-0.3.2 → marginaleffects-0.5.0}/marginaleffects/__init__.py

@@ -7,7 +7,7 @@ _EXPORTS = {
     "avg_comparisons": ("marginaleffects.comparisons", "avg_comparisons"),
     "comparisons": ("marginaleffects.comparisons", "comparisons"),
     "datagrid": ("marginaleffects.datagrid", "datagrid"),
-    "hypotheses": ("marginaleffects.hypotheses", "hypotheses"),
+    "hypotheses": ("marginaleffects.test", "hypotheses"),
     "plot_comparisons": ("marginaleffects.plot.comparisons", "plot_comparisons"),
     "plot_predictions": ("marginaleffects.plot.predictions", "plot_predictions"),
     "plot_slopes": ("marginaleffects.plot.slopes", "plot_slopes"),
@@ -18,8 +18,8 @@ _EXPORTS = {
     "fit_statsmodels": ("marginaleffects.statsmodels.model", "fit_statsmodels"),
     "fit_sklearn": ("marginaleffects.sklearn.model", "fit_sklearn"),
     "fit_linearmodels": ("marginaleffects.linearmodels.model", "fit_linearmodels"),
-    "get_dataset": ("marginaleffects.utils", "get_dataset"),
-    "MarginaleffectsResult": ("marginaleffects.result", "MarginaleffectsResult"),
+    "get_dataset": ("marginaleffects.datasets", "get_dataset"),
+    "MarginaleffectsResult": ("marginaleffects.classes", "MarginaleffectsResult"),
     "autodiff": ("marginaleffects.settings", "autodiff"),
     "set_autodiff": ("marginaleffects.settings", "set_autodiff"),
     "get_autodiff": ("marginaleffects.settings", "get_autodiff"),

marginaleffects-0.3.2/marginaleffects/jax_dispatch.py → marginaleffects-0.5.0/marginaleffects/autodiff/dispatch.py

@@ -46,7 +46,7 @@ def try_jax_predictions(
     - by is a list/complex aggregation (only False and True supported initially)
     - vcov is None (no SEs needed anyway)
     """
-    from .settings import is_autodiff_enabled
+    from ..settings import is_autodiff_enabled
 
     # Check global setting first
     if not is_autodiff_enabled():
@@ -83,7 +83,7 @@ def try_jax_predictions(
 
         # Select appropriate autodiff module based on model type
         if config["model_type"] == "linear":
-            from .autodiff import linear as ad_module
+            from . import linear as ad_module
 
             result = ad_module.predictions.predictions(
                 beta=beta,
@@ -92,7 +92,7 @@ def try_jax_predictions(
             )
 
         elif config["model_type"] == "glm":
-            from .autodiff import glm as ad_module
+            from . import glm as ad_module
 
             result = ad_module.predictions.predictions(
                 beta=beta,
@@ -147,7 +147,7 @@ def try_jax_comparisons(
     - comparison is not a supported string type
     - cross is True (cross comparisons not supported)
     """
-    from .settings import is_autodiff_enabled
+    from ..settings import is_autodiff_enabled
 
     # Check global setting first
     if not is_autodiff_enabled():
@@ -207,9 +207,9 @@ def try_jax_comparisons(
 
         # Select appropriate autodiff module based on model type
         if config["model_type"] == "linear":
-            from .autodiff import linear as ad_module
+            from . import linear as ad_module
         elif config["model_type"] == "glm":
-            from .autodiff import glm as ad_module
+            from . import glm as ad_module
         else:
             return None
 

marginaleffects-0.5.0/marginaleffects/classes/__init__.py

@@ -0,0 +1,18 @@
+from .result import MarginaleffectsResult, MarginaleffectsDataFrame
+
+__all__ = [
+    "MarginaleffectsResult",
+    "MarginaleffectsDataFrame",
+    "ModelAbstract",
+    "ModelVault",
+]
+
+
+def __getattr__(name):
+    if name in ("ModelAbstract", "ModelVault"):
+        from .model import ModelAbstract, ModelVault
+
+        globals()["ModelAbstract"] = ModelAbstract
+        globals()["ModelVault"] = ModelVault
+        return ModelAbstract if name == "ModelAbstract" else ModelVault
+    raise AttributeError(f"module 'marginaleffects.classes' has no attribute {name!r}")

marginaleffects-0.5.0/marginaleffects/classes/model.py

@@ -0,0 +1,136 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import polars as pl
+
+from ..sanitize.validation import ModelValidation
+from .. import formula as fml
+
+
+@dataclass
+class ModelVault:
+    """Typed container for model metadata shared across all adapters."""
+
+    coef: Optional[np.ndarray] = None
+    coefnames: Optional[np.ndarray] = None
+    formula: Optional[str] = None
+    formula_engine: str = "formulaic"
+    modeldata: Optional[pl.DataFrame] = None
+    package: str = "unknown"
+    vcov: Optional[np.ndarray] = None
+    variables_type: Dict[str, str] = field(default_factory=dict)
+    variable_names: Optional[List[str]] = None
+    engine_running: Optional[Any] = None
+    # statsmodels-specific
+    design_info_patsy: Optional[Any] = None
+    pandas_categorical_orders: Dict[str, list] = field(default_factory=dict)
+    # sklearn-specific
+    model_spec: Optional[Any] = None
+    original_columns: Optional[List[str]] = None
+    # linearmodels-specific
+    multiindex: Optional[List[str]] = None
+
+
+class ModelAbstract(ModelValidation, ABC):
+    def __init__(self, model, vault: ModelVault):
+        self.model = model
+        self.vault = vault
+        self.validation()
+
+    def get_modeldata(self) -> Optional[pl.DataFrame]:
+        return self.vault.modeldata
+
+    def get_vcov(self, vcov=False) -> Optional[np.ndarray]:
+        return self.vault.vcov
+
+    def get_coef(self) -> Optional[np.ndarray]:
+        return self.vault.coef
+
+    def get_coefnames(self) -> Optional[np.ndarray]:
+        return self.vault.coefnames
+
+    def get_engine_running(self) -> Optional[Any]:
+        return self.vault.engine_running
+
+    def get_formula(self) -> Optional[str]:
+        return self.vault.formula
+
+    def get_formula_engine(self) -> str:
+        return self.vault.formula_engine
+
+    def get_package(self) -> str:
+        return self.vault.package
+
+    def get_variable_type(self, name=None) -> Dict[str, str]:
+        variables = self.vault.variables_type
+        if isinstance(name, str) and name in variables:
+            return variables[name]
+        else:
+            return variables
+
+    def set_variable_type(self, name, value):
+        self.vault.variables_type[name] = value
+
+    def find_variables(self) -> Optional[List[str]]:
+        if self.vault.variable_names is not None:
+            return self.vault.variable_names
+
+        formula = self.get_formula()
+        if isinstance(formula, str):
+            out = fml.parse_variables(self.get_formula())
+        else:
+            out = None
+
+        self.vault.variable_names = out
+
+        return out
+
+    def find_response(self) -> Optional[str]:
+        vars = self.find_variables()
+        if vars is None:
+            return None
+        else:
+            return vars[0]
+
+    def find_predictors(self) -> Optional[List[str]]:
+        vars = self.find_variables()
+        if vars is None:
+            return None
+        else:
+            return vars[1:]
+
+    def get_exog(self, newdata: pl.DataFrame):
+        """Convert newdata into the design matrix format expected by get_predict.
+
+        Subclasses may override this to handle model-specific formula engines.
+        The default implementation uses the model's formula to build design matrices.
+        """
+        from ..formula import model_matrices
+
+        if self.vault.design_info_patsy is not None:
+            f = self.vault.design_info_patsy
+        else:
+            f = self.get_formula()
+
+        if callable(f):
+            _, exog = f(newdata)
+        else:
+            _, exog = model_matrices(
+                f, newdata, formula_engine=self.get_formula_engine()
+            )
+
+        return exog
+
+    @abstractmethod
+    def get_predict(self, params: np.ndarray, newdata) -> pl.DataFrame:
+        pass
+
+    def __getattr__(self, name: str) -> Any:
+        """Forward attribute access to the underlying fitted model."""
+        try:
+            return object.__getattribute__(self, name)
+        except AttributeError:
+            # Forward to the wrapped model
+            return getattr(self.model, name)

{marginaleffects-0.3.2/marginaleffects → marginaleffects-0.5.0/marginaleffects/classes}/result.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass
 from typing import Any, Dict, Iterable, Iterator, Optional
 
@@ -218,3 +219,14 @@ class MarginaleffectsResult:
     # Helpful explicit method for conversions
     def to_polars(self) -> pl.DataFrame:
         return self._data
+
+
+# Backwards compatibility alias
+class MarginaleffectsDataFrame(MarginaleffectsResult):
+    def __init__(self, *args, **kwargs):
+        warnings.warn(
+            "MarginaleffectsDataFrame is deprecated; use MarginaleffectsResult instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(*args, **kwargs)