Allow fit methods to accept pd.Series and pd.DataFrame (#62)

dte_adj/local.py

@@ -4,7 +4,7 @@
     SimpleStratifiedDistributionEstimator,
     AdjustedStratifiedDistributionEstimator,
 )
-from dte_adj.util import compute_ldte, compute_lpte
+from dte_adj.util import compute_ldte, compute_lpte, _convert_to_ndarray
 
 
 class SimpleLocalDistributionEstimator(SimpleStratifiedDistributionEstimator):
@@ -47,6 +47,7 @@ def fit(
         Returns:
             SimpleLocalDistributionEstimator: The fitted estimator.
         """
+        treatment_indicator = _convert_to_ndarray(treatment_indicator)
         super().fit(covariates, treatment_arms, outcomes, strata)
         self.treatment_indicator = treatment_indicator
 
@@ -215,6 +216,7 @@ def fit(
         Returns:
             AdjustedLocalDistributionEstimator: The fitted estimator.
         """
+        treatment_indicator = _convert_to_ndarray(treatment_indicator)
         super().fit(covariates, treatment_arms, outcomes, strata)
         self.treatment_indicator = treatment_indicator
 

dte_adj/simple.py

@@ -3,6 +3,7 @@
     SimpleStratifiedDistributionEstimator,
     AdjustedStratifiedDistributionEstimator,
 )
+from dte_adj.util import _convert_to_ndarray
 
 
 class SimpleDistributionEstimator(SimpleStratifiedDistributionEstimator):
@@ -58,6 +59,10 @@ def fit(
         Returns:
             SimpleDistributionEstimator: The fitted estimator.
         """
+        covariates = _convert_to_ndarray(covariates)
+        treatment_arms = _convert_to_ndarray(treatment_arms)
+        outcomes = _convert_to_ndarray(outcomes)
+
         if covariates.shape[0] != treatment_arms.shape[0]:
             raise ValueError("The shape of covariates and treatment_arm should be same")
 
@@ -118,6 +123,10 @@ def fit(
         Returns:
             AdjustedDistributionEstimator: The fitted estimator.
         """
+        covariates = _convert_to_ndarray(covariates)
+        treatment_arms = _convert_to_ndarray(treatment_arms)
+        outcomes = _convert_to_ndarray(outcomes)
+
         if covariates.shape[0] != treatment_arms.shape[0]:
             raise ValueError("The shape of covariates and treatment_arm should be same")
 

dte_adj/stratified.py

@@ -2,6 +2,7 @@
 from typing import Tuple, Any
 from copy import deepcopy
 from dte_adj.base import DistributionEstimatorBase
+from dte_adj.util import _convert_to_ndarray
 
 
 class SimpleStratifiedDistributionEstimator(DistributionEstimatorBase):
@@ -25,6 +26,11 @@ def fit(
         Returns:
             DistributionEstimatorBase: The fitted estimator.
         """
+        covariates = _convert_to_ndarray(covariates)
+        treatment_arms = _convert_to_ndarray(treatment_arms)
+        outcomes = _convert_to_ndarray(outcomes)
+        strata = _convert_to_ndarray(strata)
+
         if covariates.shape[0] != treatment_arms.shape[0]:
             raise ValueError("The shape of covariates and treatment_arm should be same")
 
@@ -184,6 +190,11 @@ def fit(
         Returns:
             DistributionEstimatorBase: The fitted estimator.
         """
+        covariates = _convert_to_ndarray(covariates)
+        treatment_arms = _convert_to_ndarray(treatment_arms)
+        outcomes = _convert_to_ndarray(outcomes)
+        strata = _convert_to_ndarray(strata)
+
         if covariates.shape[0] != treatment_arms.shape[0]:
             raise ValueError("The shape of covariates and treatment_arm should be same")
 

dte_adj/util.py

@@ -9,6 +9,13 @@
     )
 
 
+def _convert_to_ndarray(data: object) -> np.ndarray:
+    """Convert pd.Series or pd.DataFrame to np.ndarray if needed."""
+    if hasattr(data, "to_numpy"):
+        return data.to_numpy()
+    return data
+
+
 def compute_confidence_intervals(
     vec_y: np.ndarray,
     vec_d: np.ndarray,

pyproject.toml

@@ -30,7 +30,8 @@ dev = [
   "ruff>=0.12.2,<0.16.0",
   "sphinx>=7.3.7,<8.2.0",
   "scikit-learn>=1.5,<1.9",
-  "pre-commit>=4.0.1,<4.6.0"
+  "pre-commit>=4.0.1,<4.6.0",
+  "pandas>=2.0"
 ]
 
 [tool.setuptools.packages.find]
@@ -47,7 +48,8 @@ dev-dependencies = [
   "ruff>=0.12.2,<0.16.0",
   "sphinx>=7.3.7,<8.2.0",
   "scikit-learn>=1.5,<1.9",
-  "pre-commit>=4.0.1,<4.6.0"
+  "pre-commit>=4.0.1,<4.6.0",
+  "pandas>=2.0"
 ]
 
 [tool.ruff.lint]

tests/test_pandas_input.py

@@ -0,0 +1,161 @@
+import unittest
+import numpy as np
+import pandas as pd
+from unittest.mock import MagicMock
+from sklearn.linear_model import LogisticRegression
+from dte_adj import (
+    SimpleDistributionEstimator,
+    AdjustedDistributionEstimator,
+    SimpleStratifiedDistributionEstimator,
+    AdjustedStratifiedDistributionEstimator,
+    SimpleLocalDistributionEstimator,
+    AdjustedLocalDistributionEstimator,
+)
+
+
+class TestPandasInputSimple(unittest.TestCase):
+    """Test that Simple/Adjusted DistributionEstimator accept pandas inputs."""
+
+    def setUp(self):
+        np.random.seed(42)
+        n = 20
+        self.covariates_df = pd.DataFrame(np.zeros((n, 5)), columns=[f"x{i}" for i in range(5)])
+        self.treatment_arms_series = pd.Series(np.hstack([np.zeros(10), np.ones(10)]))
+        self.outcomes_series = pd.Series(np.arange(n, dtype=float))
+
+    def test_simple_estimator_with_dataframe_and_series(self):
+        estimator = SimpleDistributionEstimator()
+        result = estimator.fit(
+            self.covariates_df, self.treatment_arms_series, self.outcomes_series
+        )
+
+        self.assertIsInstance(result.covariates, np.ndarray)
+        self.assertIsInstance(result.treatment_arms, np.ndarray)
+        self.assertIsInstance(result.outcomes, np.ndarray)
+
+    def test_simple_estimator_predict_after_pandas_fit(self):
+        estimator = SimpleDistributionEstimator()
+        estimator.fit(self.covariates_df, self.treatment_arms_series, self.outcomes_series)
+
+        output = estimator.predict(0, np.array([3, 6]))
+        expected = np.array([0.4, 0.7])
+        np.testing.assert_array_almost_equal(output, expected, decimal=2)
+
+    def test_adjusted_estimator_with_dataframe_and_series(self):
+        base_model = MagicMock()
+        base_model.predict_proba.side_effect = lambda x, y: x
+        estimator = AdjustedDistributionEstimator(base_model, folds=2)
+        result = estimator.fit(
+            self.covariates_df, self.treatment_arms_series, self.outcomes_series
+        )
+
+        self.assertIsInstance(result.covariates, np.ndarray)
+        self.assertIsInstance(result.treatment_arms, np.ndarray)
+        self.assertIsInstance(result.outcomes, np.ndarray)
+
+
+class TestPandasInputStratified(unittest.TestCase):
+    """Test that Stratified estimators accept pandas inputs."""
+
+    def setUp(self):
+        np.random.seed(42)
+        n = 100
+        self.covariates_df = pd.DataFrame(
+            np.random.randn(n, 5), columns=[f"x{i}" for i in range(5)]
+        )
+        self.treatment_arms_series = pd.Series(np.random.choice([0, 1], size=n))
+        self.outcomes_series = pd.Series(np.random.randn(n))
+        self.strata_series = pd.Series(np.random.choice([0, 1, 2], size=n))
+
+    def test_simple_stratified_with_pandas(self):
+        estimator = SimpleStratifiedDistributionEstimator()
+        result = estimator.fit(
+            self.covariates_df,
+            self.treatment_arms_series,
+            self.outcomes_series,
+            self.strata_series,
+        )
+
+        self.assertIsInstance(result.covariates, np.ndarray)
+        self.assertIsInstance(result.treatment_arms, np.ndarray)
+        self.assertIsInstance(result.outcomes, np.ndarray)
+        self.assertIsInstance(result.strata, np.ndarray)
+
+    def test_adjusted_stratified_with_pandas(self):
+        base_model = LogisticRegression(random_state=42)
+        estimator = AdjustedStratifiedDistributionEstimator(base_model, folds=2)
+        result = estimator.fit(
+            self.covariates_df,
+            self.treatment_arms_series,
+            self.outcomes_series,
+            self.strata_series,
+        )
+
+        self.assertIsInstance(result.covariates, np.ndarray)
+        self.assertIsInstance(result.treatment_arms, np.ndarray)
+        self.assertIsInstance(result.outcomes, np.ndarray)
+        self.assertIsInstance(result.strata, np.ndarray)
+
+
+class TestPandasInputLocal(unittest.TestCase):
+    """Test that Local estimators accept pandas inputs."""
+
+    def setUp(self):
+        np.random.seed(42)
+        n = 100
+        self.covariates_df = pd.DataFrame(
+            np.random.randn(n, 3), columns=[f"x{i}" for i in range(3)]
+        )
+        self.treatment_arms_series = pd.Series(np.random.choice([0, 1], size=n))
+        self.treatment_indicator_series = pd.Series(np.random.choice([0, 1], size=n))
+        self.outcomes_series = pd.Series(np.random.randn(n))
+        self.strata_series = pd.Series(np.random.choice([0, 1], size=n))
+
+    def test_simple_local_with_pandas(self):
+        estimator = SimpleLocalDistributionEstimator()
+        result = estimator.fit(
+            self.covariates_df,
+            self.treatment_arms_series,
+            self.treatment_indicator_series,
+            self.outcomes_series,
+            self.strata_series,
+        )
+
+        self.assertIsInstance(result.covariates, np.ndarray)
+        self.assertIsInstance(result.treatment_arms, np.ndarray)
+        self.assertIsInstance(result.treatment_indicator, np.ndarray)
+        self.assertIsInstance(result.outcomes, np.ndarray)
+        self.assertIsInstance(result.strata, np.ndarray)
+
+    def test_adjusted_local_with_pandas(self):
+        base_model = LogisticRegression(random_state=42)
+        estimator = AdjustedLocalDistributionEstimator(base_model=base_model)
+        result = estimator.fit(
+            self.covariates_df,
+            self.treatment_arms_series,
+            self.treatment_indicator_series,
+            self.outcomes_series,
+            self.strata_series,
+        )
+
+        self.assertIsInstance(result.covariates, np.ndarray)
+        self.assertIsInstance(result.treatment_arms, np.ndarray)
+        self.assertIsInstance(result.treatment_indicator, np.ndarray)
+        self.assertIsInstance(result.outcomes, np.ndarray)
+        self.assertIsInstance(result.strata, np.ndarray)
+
+
+class TestNumpyInputStillWorks(unittest.TestCase):
+    """Verify that np.ndarray inputs continue to work as before."""
+
+    def test_simple_estimator_with_numpy(self):
+        estimator = SimpleDistributionEstimator()
+        covariates = np.zeros((20, 5))
+        treatment_arms = np.hstack([np.zeros(10), np.ones(10)])
+        outcomes = np.arange(20, dtype=float)
+
+        result = estimator.fit(covariates, treatment_arms, outcomes)
+
+        self.assertIsInstance(result.covariates, np.ndarray)
+        self.assertIsInstance(result.treatment_arms, np.ndarray)
+        self.assertIsInstance(result.outcomes, np.ndarray)