diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 030a8452..94854f20 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -12,7 +12,7 @@ Changelog **New features** -* Rename ``_treatment_variants_indices`` to ``_treatment_variants_mask``in ``metalearner``, ``xlearner``, ``tlearner``, and ``test_learner``. +* Rename ``_treatment_variants_indices`` to ``_treatment_variants_mask``in ``metalearner``, ``xlearner``, ``tlearner``, ``drlearner``, and ``test_learner``. 0.11.0 (2024-09-xx) diff --git a/docs/api/metalearners.cross_fit_estimator.rst b/docs/api/metalearners.cross_fit_estimator.rst new file mode 100644 index 00000000..ea548a4c --- /dev/null +++ b/docs/api/metalearners.cross_fit_estimator.rst @@ -0,0 +1,7 @@ +metalearners.cross\_fit\_estimator module +========================================= + +.. automodule:: metalearners.cross_fit_estimator + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.data_generation.rst b/docs/api/metalearners.data_generation.rst new file mode 100644 index 00000000..fdd67f37 --- /dev/null +++ b/docs/api/metalearners.data_generation.rst @@ -0,0 +1,7 @@ +metalearners.data\_generation module +==================================== + +.. automodule:: metalearners.data_generation + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.drlearner.rst b/docs/api/metalearners.drlearner.rst new file mode 100644 index 00000000..4745ae9c --- /dev/null +++ b/docs/api/metalearners.drlearner.rst @@ -0,0 +1,7 @@ +metalearners.drlearner module +============================= + +.. automodule:: metalearners.drlearner + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.explainer.rst b/docs/api/metalearners.explainer.rst new file mode 100644 index 00000000..4a550204 --- /dev/null +++ b/docs/api/metalearners.explainer.rst @@ -0,0 +1,7 @@ +metalearners.explainer module +============================= + +.. automodule:: metalearners.explainer + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.grid_search.rst b/docs/api/metalearners.grid_search.rst new file mode 100644 index 00000000..1f1d8113 --- /dev/null +++ b/docs/api/metalearners.grid_search.rst @@ -0,0 +1,7 @@ +metalearners.grid\_search module +================================ + +.. automodule:: metalearners.grid_search + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.metalearner.rst b/docs/api/metalearners.metalearner.rst new file mode 100644 index 00000000..d0d7845b --- /dev/null +++ b/docs/api/metalearners.metalearner.rst @@ -0,0 +1,7 @@ +metalearners.metalearner module +=============================== + +.. automodule:: metalearners.metalearner + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.outcome_functions.rst b/docs/api/metalearners.outcome_functions.rst new file mode 100644 index 00000000..f2462d58 --- /dev/null +++ b/docs/api/metalearners.outcome_functions.rst @@ -0,0 +1,7 @@ +metalearners.outcome\_functions module +====================================== + +.. automodule:: metalearners.outcome_functions + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.rlearner.rst b/docs/api/metalearners.rlearner.rst new file mode 100644 index 00000000..1707a551 --- /dev/null +++ b/docs/api/metalearners.rlearner.rst @@ -0,0 +1,7 @@ +metalearners.rlearner module +============================ + +.. automodule:: metalearners.rlearner + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.rst b/docs/api/metalearners.rst new file mode 100644 index 00000000..df0906ef --- /dev/null +++ b/docs/api/metalearners.rst @@ -0,0 +1,29 @@ +metalearners package +==================== + +Submodules +---------- + +.. toctree:: + :maxdepth: 4 + + metalearners.cross_fit_estimator + metalearners.data_generation + metalearners.drlearner + metalearners.explainer + metalearners.grid_search + metalearners.metalearner + metalearners.outcome_functions + metalearners.rlearner + metalearners.slearner + metalearners.tlearner + metalearners.utils + metalearners.xlearner + +Module contents +--------------- + +.. automodule:: metalearners + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.slearner.rst b/docs/api/metalearners.slearner.rst new file mode 100644 index 00000000..abfa3b70 --- /dev/null +++ b/docs/api/metalearners.slearner.rst @@ -0,0 +1,7 @@ +metalearners.slearner module +============================ + +.. automodule:: metalearners.slearner + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.tlearner.rst b/docs/api/metalearners.tlearner.rst new file mode 100644 index 00000000..086debc4 --- /dev/null +++ b/docs/api/metalearners.tlearner.rst @@ -0,0 +1,7 @@ +metalearners.tlearner module +============================ + +.. automodule:: metalearners.tlearner + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.utils.rst b/docs/api/metalearners.utils.rst new file mode 100644 index 00000000..03682711 --- /dev/null +++ b/docs/api/metalearners.utils.rst @@ -0,0 +1,7 @@ +metalearners.utils module +========================= + +.. automodule:: metalearners.utils + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/metalearners.xlearner.rst b/docs/api/metalearners.xlearner.rst new file mode 100644 index 00000000..1e7ded8b --- /dev/null +++ b/docs/api/metalearners.xlearner.rst @@ -0,0 +1,7 @@ +metalearners.xlearner module +============================ + +.. automodule:: metalearners.xlearner + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/modules.rst b/docs/api/modules.rst new file mode 100644 index 00000000..c24a10e8 --- /dev/null +++ b/docs/api/modules.rst @@ -0,0 +1,7 @@ +metalearners +============ + +.. toctree:: + :maxdepth: 4 + + metalearners diff --git a/docs/examples/model.onnx b/docs/examples/model.onnx new file mode 100644 index 00000000..99cecadc Binary files /dev/null and b/docs/examples/model.onnx differ diff --git a/metalearners/drlearner.py b/metalearners/drlearner.py index 6c03ab36..f2c67e15 100644 --- a/metalearners/drlearner.py +++ b/metalearners/drlearner.py @@ -150,12 +150,12 @@ def fit_all_nuisance( self._validate_treatment(w) self._validate_outcome(y, w) - self._treatment_variants_indices = [] + self._treatment_variants_mask = [] qualified_fit_params = self._qualified_fit_params(fit_params) for treatment_variant in range(self.n_variants): - self._treatment_variants_indices.append(w == treatment_variant) + self._treatment_variants_mask.append(w == treatment_variant) self._cv_split_indices: SplitIndices | None @@ -168,10 +168,8 @@ def fit_all_nuisance( for treatment_variant in range(self.n_variants): nuisance_jobs.append( self._nuisance_joblib_specifications( - X=index_matrix( - X, self._treatment_variants_indices[treatment_variant] - ), - y=y[self._treatment_variants_indices[treatment_variant]], + X=index_matrix(X, self._treatment_variants_mask[treatment_variant]), + y=y[self._treatment_variants_mask[treatment_variant]], model_kind=VARIANT_OUTCOME_MODEL, model_ord=treatment_variant, n_jobs_cross_fitting=n_jobs_cross_fitting,