convergence issues with new fitter class

[dhubb314]

Hi there. I am trying to implement a new class based on the BetaGeometric Survival model popularized by Fader & Hardie. I have a start on that class here, and I think the model is doing close to the right thing and converged, but I get an error still even when the model has found the values of alpha & beta that I expect. Here's the class and an example:

-- beta-geometric class --

import autograd.numpy as np
from lifelines.fitters import KnownModelParametricUnivariateFitter
from lifelines.utils.safe_exp import safe_exp
from lifelines import utils
from autograd import numpy as anp

class BetaGeometricFitter(KnownModelParametricUnivariateFitter):

    r"""

    This class implements a Beta Geometric model for univariate data. The model has parameterized
    form:

    .. math::  S(T>t) = 1 - \sum_\forall t>=1{P(T=t)},   t >= 1, \alpha > 0, \beta > 0

    where
    .. math:: P(T=1) = \frac{\alpha}{\alpha + \beta}
    .. math:: P(T=t) = P(T=t-1)*\frac{\beta + t - 2}{\alpha + \beta + t - 1}, t > 1 
      
    which implies the cumulative hazard rate is

    .. math:: H(t) = - \log(S(t)),

    and the hazard rate is:

    .. math::  h(t) = \frac{\P(t=T)}{S(t=T)}

    After calling the `.fit` method, you have access to properties like: ``cumulative_hazard_``, ``survival_function_``, ``alpha_`` and ``beta_``.
    A summary of the fit is available with the method ``print_summary()``.

    Parameters
    -----------
    alpha: float, optional (default=0.05)
        the level in the confidence intervals.

    Important
    ----------
    Examples
    --------

    >>> from lifelines import BetaGeometricFitter
    >>> df = pd.DataFrame({'T':[1,2,3,4,5,6,7,7],'observed':[1,1,1,1,1,1,1,0],'num':[131,126,90,60,42,34,26,491]})
    >>> fit(df['T'], df['observed'],weights=df['num'])
    >>> bgf = BetaGeometricFitter()
    >>> bgf.fit(df['T'], df['observed'],weights=df['num'])
    >>> bgf.plot()
    >>> print(bgf.alpha_)

    Attributes
    ----------
    cumulative_hazard_ : DataFrame
        The estimated cumulative hazard (with custom timeline if provided)
    hazard_ : DataFrame
        The estimated hazard (with custom timeline if provided)
    survival_function_ : DataFrame
        The estimated survival function (with custom timeline if provided)
    cumumlative_density_ : DataFrame
        The estimated cumulative density function (with custom timeline if provided)
    variance_matrix_ : numpy array
        The variance matrix of the coefficients
    median_: float
        The median time to event
    alpha_: float
        The fitted parameter in the model
    beta_: float
        The fitted parameter in the model
    durations: array
        The durations provided
    event_observed: array
        The event_observed variable provided
    timeline: array
        The time line to use for plotting and indexing
    entry: array or None
        The entry array provided, or None
    """

    _fitted_parameter_names = ["alpha_", "beta_"]
    _bounds = [(1e-2, None), (1e-2, None)]
    _scipy_fit_method = "SLSQP"
    _scipy_fit_options = {"ftol": 1e-3, "maxiter": 500}

    def _create_initial_point(self, Ts, E, entry, weights):
        return np.array([.01, .01])

    def _cumulative_hazard(self,params, times):
        return -np.log(self._p_and_s_t(params,times)[1])

    def _survival_function(self,params, times):
        return self._p_and_s_t(params,times)[1]

    def _log_hazard(self,params,times):
        (p,s)=self._p_and_s_t(params,times)
        return np.log(np.clip(p/s,1e-50,np.inf))

    def _p_and_s_t(self,params,times):
        alpha_, beta_ = params
        pout = np.zeros(len(times))
        sout = np.zeros(len(times)) 
        #print(alpha_)
        #print(beta_)
        if len(times) > 0:
            s,p = self._predict_survival_from_alpha_beta(alpha_,beta_,int(times.max()))
            #print(s)
            #print(p)
            #print(times)
            for i,t in enumerate(times):
                pout[i] = p[t]
                sout[i] = s[t]
        else:
            s = []
            p = []
        return (pout,sout)
    
    def extract_value(self,box):
        while isinstance(box, np.numpy_boxes.ArrayBox):
            box = box._value
        return box
    
    def _predict_survival_from_alpha_beta(self,alpha, beta, max_horizon):
        survival_curve = np.zeros((max_horizon + 1), dtype=np.float32)
        p_curve = np.zeros((max_horizon + 1), dtype=np.float32)
        alpha = self.extract_value(alpha) #._value
        beta = self.extract_value(beta) #._value
        p = alpha / (alpha + beta)
        s = 1 - p
        p_curve[0] = 0
        p_curve[1] = p
        survival_curve[0] = 1
        survival_curve[1] = s
        for i in range(2, max_horizon + 1):
            p = p * (beta + i - 2) / (alpha + beta + i - 1)
            s = s - p
            p_curve[i] = p
            survival_curve[i] = s
        return survival_curve, p_curve

    def _negative_log_likelihood_(self,params,times,observed,num):
        (p,s) = self._p_and_s_t(params,times)
        return -np.sum(num*(np.log(p*observed + s * (1-observed))))

-- example (with estimated alpha & beta & log likelihood to compare) --

df_low = pd.DataFrame({'T':[1,2,3,4,5,6,7,7],'observed':[1,1,1,1,1,1,1,0],'num':[369,163,86,56,37,27,21,241],'type':'low',
                      'alpha_':.704,'beta_':1.182,'ll':-1680.3})
df_high = pd.DataFrame({'T':[1,2,3,4,5,6,7,7],'observed':[1,1,1,1,1,1,1,0],'num':[131,126,90,60,42,34,26,491],'type':'high',
                       'alpha_':.668,'beta_':3.806,'ll':-1611.2})

for df,color in zip([df_low,df_high],['orange','blue']):
    bgf = BetaGeometricFitter()
    bgf.fit(df['T'], df['observed'],weights=df['num'] )
    bgf.plot_survival_function(color=color,linestyle='dashed',label='BG %s' % df['type'].values[0])
    print("alpha %s, beta %s" % (bgf.alpha_, bgf.beta_))
    kmf = KaplanMeierFitter()
    
    # Fit the data into the model
    kmf.fit(durations=df['T'], event_observed=df['observed'], weights=df['num'])
    # Plot the Kaplan-Meier curve
    kmf.plot(color=color,label='KM %s' % df['type'].values[0])

[dhubb314]

Update here. I still seem to be unable to control the convergence by setting ftol, but it turns out that the error was not due to convergence issues (that was only a warning). The error was that in the very last iteration, some non-integer times get passed to the function, but it is a discrete time model. So I added a hack to convert them to integers (see below). I don't understand where those times came from since the data we input only contains integers. Thanks again for the help here!

import autograd.numpy as np
from lifelines.fitters import KnownModelParametricUnivariateFitter
from lifelines.utils.safe_exp import safe_exp
from lifelines import utils
from autograd import numpy as anp

class BetaGeometricFitter(KnownModelParametricUnivariateFitter):

    r"""

    This class implements a Beta Geometric model for univariate data. The model has parameterized
    form:

    .. math::  S(T>t) = 1 - \sum_\forall t>=1{P(T=t)},   t >= 1, \alpha > 0, \beta > 0

    where
    .. math:: P(T=1) = \frac{\alpha}{\alpha + \beta}
    .. math:: P(T=t) = P(T=t-1)*\frac{\beta + t - 2}{\alpha + \beta + t - 1}, t > 1 
      
    which implies the cumulative hazard rate is

    .. math:: H(t) = - \log(S(t)),

    and the hazard rate is:

    .. math::  h(t) = \frac{\P(t=T)}{S(t=T)}

    After calling the `.fit` method, you have access to properties like: ``cumulative_hazard_``, ``survival_function_``, ``alpha_`` and ``beta_``.
    A summary of the fit is available with the method ``print_summary()``.

    Parameters
    -----------
    alpha: float, optional (default=0.05)
        the level in the confidence intervals.

    Important
    ----------
    Examples
    --------

    >>> from lifelines import BetaGeometricFitter
    >>> df = pd.DataFrame({'T':[1,2,3,4,5,6,7,7],'observed':[1,1,1,1,1,1,1,0],'num':[131,126,90,60,42,34,26,491]})
    >>> fit(df['T'], df['observed'],weights=df['num'])
    >>> bgf = BetaGeometricFitter()
    >>> bgf.fit(df['T'], df['observed'],weights=df['num'])
    >>> bgf.plot()
    >>> print(bgf.alpha_)

    Attributes
    ----------
    cumulative_hazard_ : DataFrame
        The estimated cumulative hazard (with custom timeline if provided)
    hazard_ : DataFrame
        The estimated hazard (with custom timeline if provided)
    survival_function_ : DataFrame
        The estimated survival function (with custom timeline if provided)
    cumumlative_density_ : DataFrame
        The estimated cumulative density function (with custom timeline if provided)
    variance_matrix_ : numpy array
        The variance matrix of the coefficients
    median_: float
        The median time to event
    alpha_: float
        The fitted parameter in the model
    beta_: float
        The fitted parameter in the model
    durations: array
        The durations provided
    event_observed: array
        The event_observed variable provided
    timeline: array
        The time line to use for plotting and indexing
    entry: array or None
        The entry array provided, or None
    """
    alpha_: float
    beta_: float
    _fitted_parameter_names = ["alpha_", "beta_"]
    _bounds = [(1e-2, None), (1e-2, None)]
    _scipy_fit_method = "SLSQP"
    _scipy_fit_options = {"tol": 1e-2, "maxiter": 500}

    def _create_initial_point(self, Ts, E, entry, weights):
        return np.array([1, 1])

    def _cumulative_hazard(self,params, times):
        return -np.log(self._p_and_s_t(params,times)[1])

    def _survival_function(self,params, times):
        return self._p_and_s_t(params,times)[1]

    def _log_hazard(self,params,times):
        (p,s)=self._p_and_s_t(params,times)
        return np.log(np.clip(p/s,1e-50,np.inf))

    def _p_and_s_t(self,params,times):
        alpha_, beta_ = params
        pout = np.zeros(len(times))
        sout = np.zeros(len(times)) 
        if len(times) > 0:
            s,p = self._predict_survival_from_alpha_beta(alpha_,beta_,int(np.max(self.extract_value(times))))  
            for i,t in enumerate(times):
                # @todo ensure t is an integer
                t = int(np.floor(t))
                pout[i] = p[t]
                sout[i] = s[t]
        else:
            s = []
            p = []
        return (pout,sout)
    
    def extract_value(self,box):
        while isinstance(box, np.numpy_boxes.ArrayBox):
            box = box._value
        return box
    
    def _predict_survival_from_alpha_beta(self,alpha, beta, max_horizon):
        # Make use of recursion formula from Fader & Hardie: 
        # 'How to Project Customer Retention' 2007.  
        # More details here: https://arxiv.org/pdf/1905.03818
        #
        survival_curve = np.zeros((max_horizon + 1), dtype=np.float32)
        p_curve = np.zeros((max_horizon + 1), dtype=np.float32)
        alpha = self.extract_value(alpha) #._value
        beta = self.extract_value(beta) #._value
        p = alpha / (alpha + beta)
        s = 1 - p
        p_curve[0] = 0
        p_curve[1] = p
        survival_curve[0] = 1
        survival_curve[1] = s
        for i in range(2, max_horizon + 1):
            p = p * (beta + i - 2) / (alpha + beta + i - 1)
            s = s - p
            p_curve[i] = p
            survival_curve[i] = s
        return survival_curve, p_curve

    def _negative_log_likelihood_(self,params,times,observed,num):
        (p,s) = self._p_and_s_t(params,times)
        return -np.sum(num*(np.log(p*observed + s * (1-observed))))