Implemeting a custom AFT regression model for bathtub-like hazard shapes

[Argantonio65]

Dear lifelines users,

I'm working on a survival regression problem where I have what is probably a bathtub-like hazard rate function. Unfortunately, none of the implemented AFT regression models support this form which is quite common in reliability problems. See below a comparison of Kaplan-Meier, a Cox-Proportional Haz and an AFT with a Generalized Gamma for my dataset:

I have been trying to implement a custom AFT model that can cover bathtub-like hazard functions. In particular, I have been trying to implement the log-normal power model (proposed in Reed, W. J. (2011). A flexible parametric survival model which allows a bathtub-shaped hazard rate function. Journal of Applied Statistics, 38(8), 1665-1680.). Following the lifelines documentation on user-defined AFT models, it seems it is sufficient to implement the cumulative hazard function in the general class.

see my implementation here:

from lifelines.fitters import ParametricRegressionFitter
from autograd import numpy as np
from lifelines.datasets import load_rossi
from scipy.stats import norm

class lognormal_power_Reed(ParametricRegressionFitter):
    '''Implements cumulative hazard for Reed, W. J. (2011). A flexible parametric survival model which allows a bathtub-shaped hazard rate function. Journal of Applied Statistics, 38(8), 1665-1680.'''
    # this class property is necessary, and should always be a non-empty list of strings. 
    _fitted_parameter_names = ['mu_','sigma_','beta_', 'theta_']
    
    def _cumulative_hazard(self, params, T, Xs):
        # params is a dictionary that maps unknown parameters to a numpy vector. 
        # Xs is a dictionary that maps unknown parameters to a numpy 2d array 

        ## lognormal-power 3-params
        mu = params['mu_']
        sigma = params['sigma_']
        beta = params['beta_']

        ## Covariates regression (AFT)
        theta = np.dot(Xs['theta_'], params['theta_']) #

        k = (np.log(np.clip(T, 1e-100, np.inf)) - mu)/sigma # k = (log(t) - mu)/sigma

        def R_mills(x):
            '''Mill's ratio for a standard Normal distribution'''
            return (1 - norm.cdf(x))/norm.pdf(x)
        
        # Survival function as in Reeds 2011
        surv = (norm.pdf(k) * (R_mills(k) - R_mills(beta*sigma + k)))**theta   # Equation(4) and link with external covariates through theta at Equation (11)

        return -np.log(surv)  # transform to cumulative hazard
    
rossi = load_rossi()

# the below variables maps {dataframe columns, formulas} to parameters
regressors = {
    'mu_': '1',
    'sigma_':'1',
    'beta_':'1',
    'theta_': "age + fin + mar + paro + prio + race + wexp + 1"
}

reed_aft= lognormal_power_Reed().fit(rossi, 'week', 'arrest', regressors=regressors)
reed_aft.print_summary()

Unfortunately, I get the following error:

Output exceeds the [size limit](command:workbench.action.openSettings?[). Open the full output data [in a text editor](command:workbench.action.openLargeOutput?45811057-4e4e-47ea-bf28-ee26c2ee678a)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-224-f0066ce22238> in <module>
---> 43 eaf = lognormal_power_Reed().fit(rossi, 'week', 'arrest', regressors=regressors)
     44 eaf.print_summary()
lib\site-packages\lifelines\utils\__init__.py in f(model, *args, **kwargs)
     54         def f(model, *args, **kwargs):
     55             cls.set_censoring_type(model, cls.RIGHT)
---> 56             return function(model, *args, **kwargs)
     58         return f

lib\site-packages\lifelines\fitters\__init__.py in fit(self, df, duration_col, event_col, regressors, show_progress, timeline, weights_col, robust, initial_point, entry_col)
   1707         self.durations = utils.pass_for_numeric_dtypes_or_raise_array(df.pop(duration_col)).astype(float)
   1708 
-> 1709         self._fit(
   1710             self._log_likelihood_right_censoring,
   1711             df,
\lib\site-packages\lifelines\fitters\__init__.py in _fit(self, log_likelihood_function, df, Ts, regressors, event_col, show_progress, timeline, weights_col, robust, initial_point, entry_col)
   1795 
   1796         self._pre_fit_model(Ts, E, Xs)
-> 1840         putmask(output, (1-cond0)+np.isnan(x), self.badvalue)
   1841         if np.any(cond):
   1842             goodargs = argsreduce(cond, *((x,)+args+(scale,)))
TypeError: ufunc 'isnan' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

Does anyone with experience implementing lifelines AFT models have some advice on how to follow?

Is there more documentation on how to implement custom AFT models? It would be great to be able to contribute with new distributions.

Also, did anyone experience fitting regression models on bathtub like hazard functions in lifelines?

Many thanks
Regards

[CamDavidsonPilon]

Try using the norm from autograd, too: from autograd.scipy.stats import norm

[CamDavidsonPilon]

Also, -np.log(surv), which is -np.log((norm.pdf(k) * (R_mills(k) - R_mills(beta*sigma + k)))**theta) can be simplified to help with numerical stability. The norm namespace has logpdf, too, for example.