fitters: Make _hazard() a normal function
_hazard() was a dynamic lambda function - make it a normal
function so it can be pickled.
@@ -301,9 +301,6 @@ class ParametericUnivariateFitter(UnivariateFitter):
def __init__(self, *args, **kwargs):
super(ParametericUnivariateFitter, self).__init__(*args, **kwargs)
self._estimate_name = "cumulative_hazard_"
- if not hasattr(self, "_hazard"):
- # pylint: disable=no-value-for-parameter,unexpected-keyword-arg
- self._hazard = egrad(self._cumulative_hazard, argnum=1)
if not hasattr(self, "_bounds"):
self._bounds = [(0.0, None)] * len(self._fitted_parameter_names)
self._bounds = list(self._buffer_bounds(self._bounds))
@@ -376,6 +373,10 @@ class ParametericUnivariateFitter(UnivariateFitter):
def _cumulative_hazard(self, params, times):
return -anp.log(self._survival_function(params, times))
+ def _hazard(self, *args, **kwargs):
+ # pylint: disable=no-value-for-parameter,unexpected-keyword-arg
+ return egrad(self._cumulative_hazard, argnum=1)(*args, **kwargs)
+
def _survival_function(self, params, times):
return anp.exp(-self._cumulative_hazard(params, times))
@@ -1175,7 +1176,6 @@ class ParametricRegressionFitter(BaseFitter):
def __init__(self, alpha=0.05, penalizer=0.0):
super(ParametricRegressionFitter, self).__init__(alpha=alpha)
- self._hazard = egrad(self._cumulative_hazard, argnum=1) # pylint: disable=unexpected-keyword-arg
self.penalizer = penalizer
def _check_values(self, df, T, E, weights, entries):
@@ -1202,6 +1202,9 @@ class ParametricRegressionFitter(BaseFitter):
if self.entry_col:
utils.check_entry_times(T, entries)
+ def _hazard(self, *args, **kwargs):
+ return egrad(self._cumulative_hazard, argnum=1)(*args, **kwargs) # pylint: disable=unexpected-keyword-arg
+
def _log_hazard(self, params, T, Xs):
# can be overwritten to improve convergence, see example in WeibullAFTFitter
hz = self._hazard(params, T, Xs)
@@ -2135,13 +2138,16 @@ class ParametericAFTRegressionFitter(ParametricRegressionFitter):
def __init__(self, alpha=0.05, penalizer=0.0, l1_ratio=0.0, fit_intercept=True, model_ancillary=False):
super(ParametericAFTRegressionFitter, self).__init__(alpha=alpha)
- self._hazard = egrad(self._cumulative_hazard, argnum=1) # pylint: disable=unexpected-keyword-arg
+ # self._hazard = egrad(self._cumulative_hazard, argnum=1) # pylint: disable=unexpected-keyword-arg
self._fitted_parameter_names = [self._primary_parameter_name, self._ancillary_parameter_name]
self.penalizer = penalizer
self.l1_ratio = l1_ratio
self.fit_intercept = fit_intercept
self.model_ancillary = model_ancillary
+ def _hazard(self, *args, **kwargs):
+ return egrad(self._cumulative_hazard, argnum=1)(*args, **kwargs) # pylint: disable=unexpected-keyword-arg
+
@utils.CensoringType.right_censoring
def fit(
self,
