Merge pull request #1576 from Pearcekieser/master

Add throw_on_fail setting to check_assumptions

Author: CamDavidsonPilon

lifelines/exceptions.py

@@ -5,6 +5,10 @@ class StatError(Exception):
     pass
 
 
+class ProportionalHazardAssumptionError(Exception):
+    pass
+
+
 class ConvergenceError(ValueError):
     # inherits from ValueError for backwards compatibility reasons
     def __init__(self, msg, original_exception=""):

lifelines/fitters/mixins.py

@@ -4,6 +4,7 @@
 from autograd import numpy as anp
 import numpy as np
 from pandas import DataFrame, Series
+from lifelines.exceptions import ProportionalHazardAssumptionError
 from lifelines.statistics import proportional_hazard_test, TimeTransformers
 from lifelines.utils import format_p_value
 from lifelines.utils.lowess import lowess
@@ -28,6 +29,7 @@ def check_assumptions(
         p_value_threshold: float = 0.01,
         plot_n_bootstraps: int = 15,
         columns: Optional[List[str]] = None,
+        raise_on_fail: bool = False,
     ) -> None:
         """
         Use this function to test the proportional hazards assumption. See usage example at
@@ -51,6 +53,8 @@ def check_assumptions(
             the function significantly.
         columns: list, optional
             specify a subset of columns to test.
+        raise_on_fail: bool, optional
+            throw a ``ProportionalHazardAssumptionError`` if the test fails. Default: False.
 
         Returns
         --------
@@ -107,7 +111,7 @@ def check_assumptions(
 
         for variable in self.params_.index.intersection(columns or self.params_.index):
             minumum_observed_p_value = test_results.summary.loc[variable, "p"].min()
-            
+
             # plot is done (regardless of test result) whenever `show_plots = True`
             if show_plots:
                 axes.append([])
@@ -224,9 +228,8 @@ def check_assumptions(
                         ),
                         end="\n\n",
                     )
-#################
+        #################
 
-            
         if advice and counter > 0:
             print(
                 dedent(
@@ -243,6 +246,8 @@ def check_assumptions(
 
         if counter == 0:
             print("Proportional hazard assumption looks okay.")
+        elif raise_on_fail:
+            raise ProportionalHazardAssumptionError()
         return axes
 
     @property

lifelines/tests/test_estimation.py

@@ -34,7 +34,14 @@
     qth_survival_time,
 )
 
-from lifelines.exceptions import StatisticalWarning, ApproximationWarning, StatError, ConvergenceWarning, ConvergenceError
+from lifelines.exceptions import (
+    ProportionalHazardAssumptionError,
+    StatisticalWarning,
+    ApproximationWarning,
+    StatError,
+    ConvergenceWarning,
+    ConvergenceError,
+)
 from lifelines.fitters import BaseFitter, ParametricUnivariateFitter, ParametricRegressionFitter, RegressionFitter
 from lifelines.fitters.coxph_fitter import SemiParametricPHFitter
 
@@ -3119,9 +3126,14 @@ def test_formulas_can_be_used_with_prediction(self, rossi, cph):
 
     def test_formulas_handles_categories_at_inference(self, cph):
         # Create a dummy dataset with some one continuous and one categorical features
-        df = pd.DataFrame({
-            'time': [1, 2, 3, 1, 2, 3], 'event': [0, 1, 1, 1, 0, 0],
-            'cov_cont':[0.1, 0.2, 0.3, 0.1, 0.2, 0.3], 'cov_categ': ['A', 'A', 'B', 'B', 'C', 'C']})
+        df = pd.DataFrame(
+            {
+                "time": [1, 2, 3, 1, 2, 3],
+                "event": [0, 1, 1, 1, 0, 0],
+                "cov_cont": [0.1, 0.2, 0.3, 0.1, 0.2, 0.3],
+                "cov_categ": ["A", "A", "B", "B", "C", "C"],
+            }
+        )
         cph.fit(df, "time", "event", formula="cov_cont + C(cov_categ)")
         cph.predict_survival_function(df.iloc[:4])
 
@@ -3402,6 +3414,11 @@ def test_check_assumptions(self, cph, rossi):
         cph.fit(rossi, "week", "arrest")
         cph.check_assumptions(rossi)
 
+    def test_check_assumptions_thows_if_raise_on_fail_enalbed(self, cph, rossi):
+        cph.fit(rossi, "week", "arrest")
+        with pytest.raises(ProportionalHazardAssumptionError):
+            cph.check_assumptions(rossi, p_value_threshold=0.05, raise_on_fail=True)
+
     def test_check_assumptions_for_subset_of_columns(self, cph, rossi):
         cph.fit(rossi, "week", "arrest")
         cph.check_assumptions(rossi, columns=["age"])