ingest different df types + tests

marginaleffects/comparisons.py

@@ -40,8 +40,69 @@ def comparisons(
 ):
     """
     `comparisons()` and `avg_comparisons()` are functions for predicting the outcome variable at different regressor values and comparing those predictions by computing a difference, ratio, or some other function. These functions can return many quantities of interest, such as contrasts, differences, risk ratios, changes in log odds, lift, slopes, elasticities, etc.
-
-    # Usage:
+
+    Parameters
+    ----------
+    model : object
+        Model object fitted using the `statsmodels` formula API.
+    variables : str, list, dictionary
+        - a string, list of strings, or dictionary of variables to compute comparisons for. If `None`, comparisons are computed for all regressors in the model object. Acceptable values depend on the variable type. See the examples below.
+        - Dictionary: keys identify the subset of variables of interest, and values define the type of contrast to compute. Acceptable values depend on the variable type:
+            - Categorical variables:
+                * "reference": Each factor level is compared to the factor reference (base) level
+                * "all": All combinations of observed levels
+                * "sequential": Each factor level is compared to the previous factor level
+                * "pairwise": Each factor level is compared to all other levels
+                * "minmax": The highest and lowest levels of a factor.
+                * "revpairwise", "revreference", "revsequential": inverse of the corresponding hypotheses.
+                * Vector of length 2 with the two values to compare.
+            - Boolean variables:
+                * `None`: contrast between True and False
+            - Numeric variables:
+                * Numeric of length 1: Contrast for a gap of `x`, computed at the observed value plus and minus `x / 2`. For example, estimating a `+1` contrast compares adjusted predictions when the regressor is equal to its observed value minus 0.5 and its observed value plus 0.5.
+                * Numeric of length equal to the number of rows in `newdata`: Same as above, but the contrast can be customized for each row of `newdata`.
+                * Numeric vector of length 2: Contrast between the 2nd element and the 1st element of the `x` vector.
+                * Data frame with the same number of rows as `newdata`, with two columns of "low" and "high" values to compare.
+                * Function which accepts a numeric vector and returns a data frame with two columns of "low" and "high" values to compare. See examples below.
+                * "iqr": Contrast across the interquartile range of the regressor.
+                * "sd": Contrast across one standard deviation around the regressor mean.
+                * "2sd": Contrast across two standard deviations around the regressor mean.
+                * "minmax": Contrast between the maximum and the minimum values of the regressor.
+        - Examples:
+            + `variables = {"gear" = "pairwise", "hp" = 10}`
+            + `variables = {"gear" = "sequential", "hp" = [100, 120]}`
+    newdata : polars or pandas DataFrame, or str
+        Data frame or string specifying where statistics are evaluated in the predictor space. If `None`, unit-level contrasts are computed for each observed value in the original dataset (empirical distribution).
+    comparison : str
+        String specifying how pairs of predictions should be compared. See the Comparisons section below for definitions of each transformation.
+    transform : function
+        Function specifying a transformation applied to unit-level estimates and confidence intervals just before the function returns results. Functions must accept a full column (series) of a Polars data frame and return a corresponding series of the same length. Ex:
+            - `transform = numpy.exp`
+            - `transform = lambda x: x.exp()`
+            - `transform = lambda x: x.map_elements()`
+    equivalence : list
+        List of 2 numeric values specifying the bounds used for the two-one-sided test (TOST) of equivalence, and for the non-inferiority and non-superiority tests. See the Details section below.
+    by : bool, str
+        Logical value, list of column names in `newdata`. If `True`, estimates are aggregated for each term.
+    hypothesis : str, numpy array
+        String specifying a numeric value specifying the null hypothesis used for computing p-values.
+    conf_level : float
+        Numeric value specifying the confidence level for the confidence intervals. Default is 0.95.
+    Returns
+    -------
+    out : DataFrame
+        The functions return a data.frame with the following columns:
+        - term: the name of the variable.
+        - contrast: the comparison method used.
+        - estimate: the estimated contrast, difference, ratio, or other transformation between pairs of predictions.
+        - std_error: the standard error of the estimate.
+        - statistic: the test statistic (estimate / std.error).
+        - p_value: the p-value of the test.
+        - s_value: Shannon transform of the p value.
+        - conf_low: the lower confidence interval bound.
+        - conf_high: the upper confidence interval bound.
+    Examples
+    --------
 
         comparisons(model, variables = NULL, newdata = NULL, comparison = "difference",
                     transform = NULL, equivalence = NULL, by = FALSE, cross = FALSE,
@@ -51,60 +112,8 @@ def comparisons(
                         transform = NULL, equivalence = NULL, by = FALSE, cross = FALSE,
                         type = "response", hypothesis = 0, conf.level = 0.95, ...)
 
-    # Args:
-
-    - model (object): a model object fitted using the `statsmodels` formula API.
-    - variables (str, list, or dictionary): a string, list of strings, or dictionary of variables to compute comparisons for. If `None`, comparisons are computed for all regressors in the model object. Acceptable values depend on the variable type. See the examples below.
-    * Dictionary: keys identify the subset of variables of interest, and values define the type of contrast to compute. Acceptable values depend on the variable type:
-    - Categorical variables:
-        * "reference": Each factor level is compared to the factor reference (base) level
-        * "all": All combinations of observed levels
-        * "sequential": Each factor level is compared to the previous factor level
-        * "pairwise": Each factor level is compared to all other levels
-        * "minmax": The highest and lowest levels of a factor.
-        * "revpairwise", "revreference", "revsequential": inverse of the corresponding hypotheses.
-        * Vector of length 2 with the two values to compare.
-    - Boolean variables:
-        * `None`: contrast between True and False
-    - Numeric variables:
-        * Numeric of length 1: Contrast for a gap of `x`, computed at the observed value plus and minus `x / 2`. For example, estimating a `+1` contrast compares adjusted predictions when the regressor is equal to its observed value minus 0.5 and its observed value plus 0.5.
-        * Numeric of length equal to the number of rows in `newdata`: Same as above, but the contrast can be customized for each row of `newdata`.
-        * Numeric vector of length 2: Contrast between the 2nd element and the 1st element of the `x` vector.
-        * Data frame with the same number of rows as `newdata`, with two columns of "low" and "high" values to compare.
-        * Function which accepts a numeric vector and returns a data frame with two columns of "low" and "high" values to compare. See examples below.
-        * "iqr": Contrast across the interquartile range of the regressor.
-        * "sd": Contrast across one standard deviation around the regressor mean.
-        * "2sd": Contrast across two standard deviations around the regressor mean.
-        * "minmax": Contrast between the maximum and the minimum values of the regressor.
-    - Examples:
-        + `variables = {"gear" = "pairwise", "hp" = 10}`
-        + `variables = {"gear" = "sequential", "hp" = [100, 120]}`
-    - newdata (polars or pandas DataFrame, or str): a data frame or a string specifying where statistics are evaluated in the predictor space. If `None`, unit-level contrasts are computed for each observed value in the original dataset (empirical distribution).
-    - comparison (str): a string specifying how pairs of predictions should be compared. See the Comparisons section below for definitions of each transformation.
-    - transform (function): a function specifying a transformation applied to unit-level estimates and confidence intervals just before the function returns results. Functions must accept a full column (series) of a Polars data frame and return a corresponding series of the same length. Ex:
-        - `transform = numpy.exp`
-        - `transform = lambda x: x.exp()`
-        - `transform = lambda x: x.map_elements()`
-    - equivalence (list): a list of 2 numeric values specifying the bounds used for the two-one-sided test (TOST) of equivalence, and for the non-inferiority and non-superiority tests. See the Details section below.
-    - by (bool, str): a logical value, a list of column names in `newdata`. If `True`, estimates are aggregated for each term.
-    - hypothesis (str, numpy array): a string specifying a numeric value specifying the null hypothesis used for computing p-values.
-    - conf.level (float): a numeric value specifying the confidence level for the confidence intervals. Default is 0.95.
-
-    # Returns:
-
-    The functions return a data.frame with the following columns:
-
-    - term: the name of the variable.
-    - contrast: the comparison method used.
-    - estimate: the estimated contrast, difference, ratio, or other transformation between pairs of predictions.
-    - std_error: the standard error of the estimate.
-    - statistic: the test statistic (estimate / std.error).
-    - p_value: the p-value of the test.
-    - s_value: Shannon transform of the p value.
-    - conf_low: the lower confidence interval bound.
-    - conf_high: the upper confidence interval bound.
-
-    # Details:
+    Details
+    -------
 
     The `equivalence` argument specifies the bounds used for the two-one-sided test (TOST) of equivalence, and for the non-inferiority and non-superiority tests. The first element specifies the lower bound, and the second element specifies the upper bound. If `None`, equivalence tests are not performed.
     """

marginaleffects/hypotheses_joint.py

@@ -1,7 +1,6 @@
 import numpy as np
 import scipy.stats as stats
 import polars as pl
-import pandas as pd
 
 from .sanity import sanitize_hypothesis_null
 from .classes import MarginaleffectsDataFrame
@@ -10,9 +9,6 @@
 def joint_hypotheses(obj, joint_index=None, joint_test="f", hypothesis=0):
     assert joint_test in ["f", "chisq"], "`joint_test` must be `f` or `chisq`"
 
-    if isinstance(obj, pd.DataFrame):
-        obj = pl.DataFrame(obj)
-
     # theta_hat: P x 1 vector of estimated parameters
     theta_hat = obj.get_coef()
 

marginaleffects/sanity.py

@@ -6,6 +6,7 @@
 
 from .datagrid import datagrid
 from .estimands import estimands
+from .utils import ingest, ArrowStreamExportable
 
 
 def sanitize_vcov(vcov, model):
@@ -67,14 +68,14 @@ def sanitize_newdata(model, newdata, wts, by=[]):
 
     else:
         try:
-            import pandas as pd
-
-            if isinstance(newdata, pd.DataFrame):
-                out = pl.from_pandas(newdata)
+            if isinstance(newdata, ArrowStreamExportable):
+                out = ingest(newdata)
             else:
-                out = newdata
-        except ImportError:
-            out = newdata
+                raise RuntimeError(
+                    "Unable to ingest newdata data provided. If it is a DataFrame, make sure it implements the ArrowStreamExportable interface."
+                )
+        except Exception as e:
+            raise e
 
     reserved_names = {
         "rowid",

marginaleffects/utils.py

@@ -1,7 +1,17 @@
 import itertools
-
+import narwhals as nw
 import numpy as np
 import polars as pl
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class ArrowStreamExportable(Protocol):
+    def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: ...
+
+
+def ingest(df: ArrowStreamExportable):
+    return nw.from_arrow(df, native_namespace=pl).to_native()
 
 
 def sort_columns(df, by=None, newdata=None):

pyproject.toml

@@ -5,21 +5,23 @@ description = "Predictions, counterfactual comparisons, slopes, and hypothesis t
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
+    "narwhals >=1.10.0",
     "numpy >=2.0.0",
     "patsy >=0.5.6",
     "polars >=1.7.0",
-    "pyarrow >=17.0.0",
-    "scipy >=1.14.1",
     "plotnine >=0.13.6",
+    "scipy >=1.14.1",
 ]
 
 [project.optional-dependencies]
 test = [
-    "pandas >=2.2.2",
+    "duckdb >=1.1.2",
     "matplotlib >=3.7.1",
-    "typing-extensions >=4.7.0",
-    "statsmodels >=0.14.0",
+    "pandas >=2.2.2",
+    "pyarrow >=17.0.0",
     "pyfixest >=0.24.2",
+    "statsmodels >=0.14.0",
+    "typing-extensions >=4.7.0",
 ]
 
 [tool.uv]

tests/test_utils.py

@@ -0,0 +1,70 @@
+import pytest
+import pandas as pd
+import polars as pl
+import duckdb
+import pyarrow as pa
+import narwhals as nw
+from marginaleffects.utils import ingest
+from typing import Callable
+
+
+def get_sample_data():
+    return pd.DataFrame(
+        {
+            "id": [1, 2, 3],
+            "name": ["Alice", "Bob", "Charlie"],
+            "age": [25, 30, 35],
+            "score": [85.5, 90.0, 95.5],
+        }
+    )
+
+
+sample_data = get_sample_data()
+
+
+@pytest.fixture
+def sample_pandas_df():
+    return get_sample_data()
+
+
+@pytest.fixture
+def sample_polars_df():
+    pd_df = get_sample_data()
+    return pl.from_pandas(pd_df)
+
+
+@pytest.fixture
+def sample_duckdb_df():
+    # Using DuckDB to create a DataFrame
+    con = duckdb.connect()
+    return con.execute("SELECT * FROM sample_data").df()
+
+
+def test_ingest_pandas(sample_pandas_df):
+    result = ingest(sample_pandas_df)
+    assert isinstance(result, pl.DataFrame), "Result should be a Polars DataFrame"
+    # Verify contents
+    expected = pl.from_pandas(sample_pandas_df)
+    assert result.equals(
+        expected
+    ), "Ingested DataFrame does not match expected Polars DataFrame"
+
+
+def test_ingest_polars(sample_polars_df):
+    result = ingest(sample_polars_df)
+    assert isinstance(result, pl.DataFrame), "Result should be a Polars DataFrame"
+    # Verify contents
+    expected = sample_polars_df
+    assert result.equals(
+        expected
+    ), "Ingested DataFrame does not match expected Polars DataFrame"
+
+
+def test_ingest_duckdb(sample_duckdb_df):
+    result = ingest(sample_duckdb_df)
+    assert isinstance(result, pl.DataFrame), "Result should be a Polars DataFrame"
+    # Verify contents
+    expected = pl.from_pandas(sample_duckdb_df)
+    assert result.equals(
+        expected
+    ), "Ingested DataFrame does not match expected Polars DataFrame"