Fix piecewise hazard prediction

.pre-commit-config.yaml

@@ -9,8 +9,3 @@ repos:
     -   id: fix-encoding-pragma
     -   id: mixed-line-ending
     -   id: trailing-whitespace
--   repo: https://github.com/ambv/black
-    rev: stable
-    hooks:
-    - id: black
-      args: ["--line-length", "130"]

lifelines/fitters/coxph_fitter.py

@@ -3126,14 +3126,21 @@ def _cumulative_hazard_with_strata(self, params, T, Xs):
 
             partial_hazard = safe_exp(anp.dot(Xs_["beta_"], params["beta_"]))
             n = T_.shape[0]
-            T_ = T_.reshape((n, 1))
+            if T.ndim == 2:
+                T_ = T_[:, 0]
+            if T_.ndim == 1:
+                T_ = T_.reshape((n, 1))
             bps = anp.append(self.breakpoints, [anp.inf])
             M = anp.minimum(anp.tile(bps, (n, 1)), T_)
             M = anp.hstack([M[:, tuple([0])], anp.diff(M, axis=1)])
             log_lambdas_ = anp.array(
                 [0] + [params[self._strata_labeler(stratum, i)][0] for i in range(2, self.n_breakpoints + 2)]
             )
-            H_ = partial_hazard * (M * anp.exp(log_lambdas_).T).sum(1)
+            lambdas_ = anp.exp(log_lambdas_)
+            if T.ndim == 1:
+                H_ = partial_hazard * (M * lambdas_.T).sum(1)
+            else:
+                H_ = anp.outer(anp.dot(M, lambdas_), partial_hazard)
 
             output.append(H_)
             start = stop
@@ -3143,12 +3150,15 @@ def _cumulative_hazard_with_strata(self, params, T, Xs):
     def _cumulative_hazard_sans_strata(self, params, T, Xs):
         partial_hazard = safe_exp(anp.dot(Xs["beta_"], params["beta_"]))
         n = T.shape[0]
+        if T.ndim == 2:
+            T = T[:, 0]
         T = T.reshape((n, 1))
         bps = anp.append(self.breakpoints, [anp.inf])
         M = anp.minimum(anp.tile(bps, (n, 1)), T)
         M = anp.hstack([M[:, tuple([0])], anp.diff(M, axis=1)])
         log_lambdas_ = anp.array([0.0] + [params[param][0] for param in self._fitted_parameter_names if param != "beta_"])
-        return partial_hazard * (M * anp.exp(log_lambdas_).T).sum(1)
+        lambdas_ = anp.exp(log_lambdas_)
+        return anp.outer(anp.dot(M, lambdas_), partial_hazard)
 
     def predict_cumulative_hazard(self, df, times=None, conditional_after=None) -> pd.DataFrame:
         """

lifelines/fitters/piecewise_exponential_regression_fitter.py

@@ -73,12 +73,14 @@ def _add_penalty(self, params, neg_ll):
 
     def _cumulative_hazard(self, params, T, Xs):
         n = T.shape[0]
+        if T.ndim == 2:
+            T = T[:, 0]
         T = T.reshape((n, 1))
         bps = np.append(self.breakpoints, [np.inf])
         M = np.minimum(np.tile(bps, (n, 1)), T)
         M = np.hstack([M[:, tuple([0])], np.diff(M, axis=1)])
         lambdas_ = np.array([safe_exp(-np.dot(Xs[param], params[param])) for param in self._fitted_parameter_names])
-        return (M * lambdas_.T).sum(1)
+        return np.dot(M, lambdas_)
 
     def _prep_inputs_for_prediction_and_return_parameters(self, X):
         X = X.copy()

lifelines/tests/test_estimation.py

@@ -2873,7 +2873,7 @@ def test_efron_newtons_method(self, data_nus, cph):
         assert np.abs(newton(X, T, E, W, entries)[0] - -0.0335) < 0.0001
 
 
-class TestCoxPHFitterPeices:
+class TestCoxPHFitterPieces:
     @pytest.fixture
     def cph(self):
         return CoxPHFitter(baseline_estimation_method="piecewise", breakpoints=[25])
@@ -2904,6 +2904,80 @@ def test_trivial_model_doesnt_fail(self, cph, rossi):
         cph.fit(rossi[["week", "arrest"]], "week", "arrest")
         cph.baseline_hazard_
 
+    def test_predict_hazard_sans_strata(self, cph, rossi):
+        cph = cph.fit(rossi, "week", "arrest", formula='fin')
+        samples = rossi[["week", "arrest", "fin"]].head(5)
+        hazards = cph.predict_hazard(samples)
+        npt.assert_equal(hazards.isnull().sum().sum(), 0)
+        assert (hazards >= 0).all().all()
+        print(hazards)
+        hazard_last = cph.predict_hazard(samples.iloc[4:5])
+        print(hazard_last)
+
+        cumulative_hazard_last = cph.predict_cumulative_hazard(samples.iloc[4:5])
+        print(cumulative_hazard_last)
+
+        npt.assert_array_equal(hazard_last[4].values, hazards[4].values)
+
+    def test_predict_hazard_with_strata(self, cph, rossi):
+        cph = cph.fit(rossi, "week", "arrest", formula="fin", strata=["mar"])
+        samples = rossi[["week", "arrest", "fin", "mar"]].head(5)
+        hazards = cph.predict_hazard(samples)
+        npt.assert_equal(hazards.isnull().sum().sum(), 0)
+        assert (hazards >= 0).all().all()
+
+    def test_predict_cumulative_hazard_sans_strata(self, cph, rossi):
+        cph = cph.fit(rossi, "week", "arrest", formula='fin')
+        samples = rossi[["week", "arrest", "fin"]].head(5)
+        cumulative_hazards = cph.predict_cumulative_hazard(samples)
+        npt.assert_equal(cumulative_hazards.isnull().sum().sum(), 0)
+        assert (cumulative_hazards >= 0).all().all()
+
+    def test_predict_cumulative_hazard_with_strata(self, cph, rossi):
+        cph = cph.fit(rossi, "week", "arrest", formula='fin', strata=["mar"])
+        samples = rossi[["week", "arrest", "fin", "mar"]].head(5)
+        cumulative_hazards = cph.predict_cumulative_hazard(samples)
+        npt.assert_equal(cumulative_hazards.isnull().sum().sum(), 0)
+        assert (cumulative_hazards >= 0).all().all()
+
+
+class TestCoxPHFitterSpline:
+    @pytest.fixture
+    def cph(self):
+        return CoxPHFitter(baseline_estimation_method="spline", n_baseline_knots=2)
+
+    def test_predict_hazard_sans_strata(self, cph, rossi):
+        cph = cph.fit(rossi, "week", "arrest", formula='fin')
+        samples = rossi[["week", "arrest", "fin"]].head(5)
+        hazards = cph.predict_hazard(samples)
+        npt.assert_equal(hazards.isnull().sum().sum(), 0)
+        assert (hazards >= 0).all().all()
+        print(hazards)
+        hazard_last = cph.predict_hazard(samples.iloc[4:5])
+        print(hazard_last)
+
+    def test_predict_hazard_with_strata(self, cph, rossi):
+        cph = cph.fit(rossi, "week", "arrest", formula="fin", strata=["mar"])
+        samples = rossi[["week", "arrest", "fin", "mar"]].head(5)
+        hazards = cph.predict_hazard(samples)
+        npt.assert_equal(hazards.isnull().sum().sum(), 0)
+        assert (hazards >= 0).all().all()
+
+    def test_predict_cumulative_hazard_sans_strata(self, cph, rossi):
+        cph = cph.fit(rossi, "week", "arrest", formula='fin')
+        samples = rossi[["week", "arrest", "fin"]].head(5)
+        cumulative_hazards = cph.predict_cumulative_hazard(samples)
+        npt.assert_equal(cumulative_hazards.isnull().sum().sum(), 0)
+        assert (cumulative_hazards >= 0).all().all()
+
+    def test_predict_cumulative_hazard_with_strata(self, cph, rossi):
+        cph = cph.fit(rossi, "week", "arrest", formula='fin', strata=["mar"])
+        samples = rossi[["week", "arrest", "fin", "mar"]].head(5)
+        cumulative_hazards = cph.predict_cumulative_hazard(samples)
+        npt.assert_equal(cumulative_hazards.isnull().sum().sum(), 0)
+        assert (cumulative_hazards >= 0).all().all()
+
+
 
 class TestCoxPHFitter:
     @pytest.fixture