Skip to content

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

Open
okiner-3 wants to merge 2 commits intomainfrom
feat/input/arrylike_to_ndarray
Open

Allow fit methods to accept pd.Series and pd.DataFrame (#62)#92
okiner-3 wants to merge 2 commits intomainfrom
feat/input/arrylike_to_ndarray

Conversation

@okiner-3
Copy link
Collaborator

close #62

Summary

  • Allow fit methods of all estimators to accept pd.Series and pd.DataFrame in addition to np.ndarray
  • Inputs are automatically converted to np.ndarray via to_numpy() inside the fit method
  • Any object that implements to_numpy() (e.g. polars) is also supported

Changes

  • Add _convert_to_ndarray helper in dte_adj/util.py
  • Apply conversion in fit methods of all 6 estimator classes:
    • SimpleDistributionEstimator
    • AdjustedDistributionEstimator
    • SimpleStratifiedDistributionEstimator
    • AdjustedStratifiedDistributionEstimator
    • SimpleLocalDistributionEstimator
    • AdjustedLocalDistributionEstimator
  • Add tests/test_pandas_input.py with 8 test cases

@okiner-3 okiner-3 self-assigned this Feb 28, 2026
@okiner-3 okiner-3 requested a review from TomeHirata February 28, 2026 02:54
)


def _convert_to_ndarray(data: object) -> np.ndarray:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we update the type hint to Dataframe, ndarray and Series only?

@@ -0,0 +1,161 @@
import unittest
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we use test_utils.py?

Copy link
Collaborator

@TomeHirata TomeHirata Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's fine to test _convert_to_ndarray only. Alternativel, we can add a test case for each estimator class.


def fit(
self,
covariates: np.ndarray,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we update the type hint for fit?

"""
super().__init__()

def fit(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto

"""
super().__init__()

def fit(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Allow to pass pd.Series and pd.DataFrame

2 participants