Multinomial regression to support categorical variables

Author: jmafoster1

causal_testing/discovery/abstract_discovery.py

@@ -167,6 +167,15 @@ def write_dot(self, individual: CausalDAG, output_file: str):
 
         nx.drawing.nx_pydot.write_dot(individual, output_file)
 
+    def _json_stub_params(self, outcome: str) -> str:
+        if pd.api.types.is_bool_dtype(self.df[outcome]):
+            return {"estimator": "LogisticRegressionEstimator", "estimate_type": "unit_odds_ratio"}
+        if pd.api.types.is_categorical_dtype(self.df[outcome]) or pd.api.types.is_object_dtype(self.df[outcome]):
+            return {"estimator": "MultinomialRegressionEstimator", "estimate_type": "unit_odds_ratio"}
+        if pd.api.types.is_numeric_dtype(self.df[outcome]):
+            return {"estimator": "LinearRegressionEstimator", "estimate_type": "coefficient"}
+        raise ValueError(f"Invalid datatype {self.df.dtypes[outcome]}")
+
     def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
         """
         Generate and evaluate causal test cases from the supplied CausalDAG and return a list of edges for which the
@@ -188,13 +197,8 @@ def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
             {
                 "tests": [
                     relation.to_json_stub(
-                        estimator=(
-                            "LinearRegressionEstimator"
-                            if pd.api.types.is_numeric_dtype(relation.base_test_case.outcome_variable)
-                            else "LogisticRegressionEstimator"
-                        ),
-                        estimate_type="unit_odds_ratio",
                         alpha=self.alpha,
+                        **self._json_stub_params(relation.base_test_case.outcome_variable),
                     )
                     for relation in generate_metamorphic_relations(causal_dag)
                 ]
@@ -203,7 +207,9 @@ def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
 
         results = []
 
-        for test_case, result in zip(ctf.test_cases, ctf.run_tests(silent=True)):
+        # We use "silent=False" here to allow for inestimable edges, but it'd be good to have a more stringent
+        # error catching strategy to catch "genuine" problems (e.g. to do with the structure of the data)
+        for test_case, result in zip(ctf.test_cases, ctf.run_tests(silent=False)):
             if result.effect_estimate is None:
                 results.append(
                     {

causal_testing/estimation/multinomial_regression_estimator.py

@@ -0,0 +1,68 @@
+"""This module contains the LogisticRegressionEstimator class for estimating categorical outcomes."""
+
+import logging
+
+import numpy as np
+import pandas as pd
+import statsmodels.api as sm
+
+from causal_testing.estimation.abstract_regression_estimator import RegressionEstimator
+from causal_testing.estimation.effect_estimate import EffectEstimate
+
+logger = logging.getLogger(__name__)
+
+
+class MultinomialRegressionEstimator(RegressionEstimator):
+    """A Logistic Regression Estimator is a parametric estimator which restricts the variables in the data to a linear
+    combination of parameters and functions of the variables (note these functions need not be linear). It is designed
+    for estimating categorical outcomes.
+    """
+
+    regressor = sm.MNLogit
+
+    def add_modelling_assumptions(self):
+        """
+        Add modelling assumptions to the estimator. This is a list of strings which list the modelling assumptions that
+        must hold if the resulting causal inference is to be considered valid.
+        """
+        self.modelling_assumptions.append(
+            "The variables in the data must fit a shape which can be expressed as a linear"
+            "combination of parameters and functions of variables. Note that these functions"
+            "do not need to be linear."
+        )
+        self.modelling_assumptions.append("The outcome must be binary.")
+        self.modelling_assumptions.append("Independently and identically distributed errors.")
+
+    def estimate_unit_odds_ratio(self) -> EffectEstimate:
+        """Estimate the odds ratio of increasing the treatment by one. In logistic regression, this corresponds to the
+        coefficient of the treatment of interest.
+
+        :return: The odds ratio with confidence intervals.
+        """
+        model = self.fit_model(self.df)
+
+        treatment_columns = [
+            param
+            for param in model.params.index
+            if param == self.base_test_case.treatment_variable.name
+            or param.startswith(self.base_test_case.treatment_variable.name + "[")
+        ]
+
+        conf_int = model.conf_int(self.alpha)
+        levels_of_interest = [
+            (level, covariate) for level, covariate in conf_int.index if covariate in treatment_columns
+        ]
+        confidence_intervals = np.exp(conf_int.loc[levels_of_interest])
+
+        # Format the params to a MultiIndexed Series like the confidence intervals for consistent indexing
+        stacked_params = model.params.stack(dropna=False)
+        multi_indexed_params = stacked_params.swaplevel(0, 1).sort_index()
+        multi_indexed_params.index = conf_int.index
+
+        result = EffectEstimate(
+            "unit_odds_ratio",
+            pd.Series(np.exp(multi_indexed_params[levels_of_interest])),
+            pd.Series(confidence_intervals["lower"]),
+            pd.Series(confidence_intervals["upper"]),
+        )
+        return result

causal_testing/testing/metamorphic_relation.py

@@ -49,7 +49,11 @@ def to_json_stub(
         :param estimator: The name of the estimator class to use when evaluating the test
         :param alpha: The significance level to use when calculating the confidence intervals
         """
-        if estimator not in ["LinearRegressionEstimator", "LogisticRegressionEstimator"]:
+        if estimator not in [
+            "LinearRegressionEstimator",
+            "LogisticRegressionEstimator",
+            "MultinomialRegressionEstimator",
+        ]:
             raise ValueError(
                 f"Unsupported estimator {estimator}. "
                 "We only support autogeneration using LinearRegressionEstimator or LogisticRegressionEstimator."

pyproject.toml

@@ -73,6 +73,7 @@ LogisticRegressionEstimator = "causal_testing.estimation.logistic_regression_est
 CubicSplineEstimator = "causal_testing.estimation.cubic_spline_estimator:CubicSplineEstimator"
 InstrumentalVariableEstimator = "causal_testing.estimation.instrumental_variable_estimator:InstrumentalVariableEstimator"
 IPCWEstimator = "causal_testing.estimation.ipcw_estimator:IPCWEstimator"
+MultinomialRegressionEstimator = "causal_testing.estimation.multinomial_regression_estimator:MultinomialRegressionEstimator"
 
 [project.entry-points."causal_effects"]
 NoEffect = "causal_testing.testing.causal_effect:NoEffect"

tests/estimation_tests/test_multinomial_regression_estimator.py

@@ -0,0 +1,33 @@
+import unittest
+import pandas as pd
+from causal_testing.estimation.multinomial_regression_estimator import MultinomialRegressionEstimator
+from causal_testing.testing.base_test_case import BaseTestCase
+from causal_testing.specification.variable import Input, Output
+
+
+class TestMultinomialRegressionEstimator(unittest.TestCase):
+    """Test the multinomial regression estimator against the scarf example from
+    https://investigate.ai/regression/multinomial-regression/.
+    (For binary categories, this should behave the same as for logistic regression)
+    """
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv")
+
+    def test_odds_ratio(self):
+        df = self.scarf_df.copy()
+        multinomial_regression_estimator = MultinomialRegressionEstimator(
+            BaseTestCase(Input("length_in", float), Output("completed", bool)), 65, 55, set(), df
+        )
+        effect_estimate = multinomial_regression_estimator.estimate_unit_odds_ratio()
+        self.assertEqual(round(effect_estimate.value.iloc[0], 4), 0.8948)
+
+    def test_odds_ratio_data(self):
+        df = self.scarf_df.copy()
+        multinomial_regression_estimator = MultinomialRegressionEstimator(
+            BaseTestCase(Input("length_in", float), Output("completed", bool)), 65, 55, set()
+        )
+        multinomial_regression_estimator.df = df
+        effect_estimate = multinomial_regression_estimator.estimate_unit_odds_ratio()
+        self.assertEqual(round(effect_estimate.value.iloc[0], 4), 0.8948)