<commit_before>import numpy as np
from numpy.linalg import LinAlgError, inv
from numpy import dot
import pandas as pd

from lifelines.plotting import plot_dataframes
from lifelines.utils import dataframe_from_events_censorship, basis

import pdb



class NelsonAalenFitter(object):
    """
    Class for fitting the Nelson-Aalen estimate for the cumulative hazard. 

    NelsonAalenFitter( alpha=0.95, nelson_aalen_smoothing=True)

    alpha: The alpha value associated with the confidence intervals. 
    nelson_aalen_smoothing: If the event times are naturally discrete (like discrete years, minutes, etc.)
      then it is advisable to turn this parameter to False. See [1], pg.84.

    """
    def __init__(self, alpha=0.95, nelson_aalen_smoothing=True):
        self.alpha = alpha
        self.nelson_aalen_smoothing = nelson_aalen_smoothing
        if nelson_aalen_smoothing:
          self._variance_f = self._variance_f_smooth
        else:
          self._variance_f = self._variance_f_discrete

    def fit(self, event_times,censorship=None, timeline=None, columns=['NA-estimate']):
        """
        Parameters:
          event_times: an (n,1) array of times that the death event occured at 
          timeline: return the best estimate at the values in timelines (postively increasing)

        Returns:
          DataFrame with index either event_times or timelines (if not None), with
          values as the NelsonAalen estimate
        """
        
        if censorship is None:
           self.censorship = np.ones_like(event_times, dtype=bool) #why boolean?
        else:
           self.censorship = censorship.copy()

        self.event_times = dataframe_from_events_censorship(event_times, self.censorship)

        if timeline is None:
           self.timeline = self.event_times.index.values.copy()
        else:
           self.timeline = timeline
        self.cumulative_hazard_, cumulative_sq_ = _additive_estimate(self.event_times, 
                                                                     self.timeline, self.censorship, 
                                                                     self._additive_f, 0, columns, self._variance_f,
                                                                     nelson_aalen_smoothing=self.nelson_aalen_smoothing )
        self.confidence_interval_ = self._bounds(cumulative_sq_)
        self.plot = plot_dataframes(self, "cumulative_hazard_")

        return

    def _additive_f(self, N,d ):
       #check it 0
       if N==d==0:
          return 0
       return 1.*d/N

    def _bounds(self, cumulative_sq_):
        inverse_norm = { 0.95:1.96, 0.99:2.57 }
        try:
          coef = inverse_norm[self.alpha]
        except:
          pass
        df = pd.DataFrame( index=self.timeline)
        df["upper_bound_%.2f"%self.alpha] = self.cumulative_hazard_.values*np.exp(coef*np.sqrt(cumulative_sq_)/self.cumulative_hazard_.values )
        df["lower_bound_%.2f"%self.alpha] = self.cumulative_hazard_.values*np.exp(-coef*np.sqrt(cumulative_sq_)/self.cumulative_hazard_.values )
        return df

    def _variance_f_smooth(self, N, d):
        if N==d==0:
            return 0
        return (1./N)**2  

    def _variance_f_discrete(self, N, d):
        if N==d==0:
            return 0
        return 1.*(N-d)*d/N**3

class KaplanMeierFitter(object):
   
  def __init__(self, alpha = 0.95):
       self.alpha = alpha

  def fit(self, event_times,censorship=None, timeline=None, columns=['KM-estimate']):
       """
       Parameters:
         event_times: an (n,1) array of times that the death event occured at 
         timeline: return the best estimate at the values in timelines (postively increasing)
         censorship: an (n,1) array of booleans -- True if the the death was observed, False if the event 
            was lost (right-censored). Defaults all True if censorship==None
       Returns:
         DataFrame with index either event_times or timelines (if not None), with
         values as the NelsonAalen estimate
       """
       #need to sort event_times
       if censorship is None:
          self.censorship = np.ones_like(event_times, dtype=bool) #why boolean?
       else:
          self.censorship = censorship.copy()

       self.event_times = dataframe_from_events_censorship(event_times, self.censorship)

       if timeline is None:
          self.timeline = self.event_times.index.values.copy()
       else:
          self.timeline = timeline
       log_surivial_function, cumulative_sq_ = _additive_estimate(self.event_times, self.timeline, 
                                                                  self.censorship, self._additive_f, 
                                                                  0, columns, self._variance_f )
       self.survival_function_ = np.exp(log_surivial_function)
       self.median_ = median_survival_times(self.survival_function_)
       self.confidence_interval_ = self._bounds(cumulative_sq_)
       self.plot = plot_dataframes(self, "survival_function_")
       return self

  def _additive_f(self,N,d):
      if N==d==0:
        return 0
      return np.log(1 - 1.*d/N)

  def _bounds(self, cumulative_sq_):
      inverse_norm = { 0.95:1.96, 0.99:2.57 }
      try:
        coef = inverse_norm[self.alpha]
      except:
        pass
      df = pd.DataFrame( index=self.timeline)
      df["upper_bound_%.2f"%self.alpha] = self.survival_function_.values**(np.exp(coef*cumulative_sq_/np.log(self.survival_function_.values)))
      df["lower_bound_%.2f"%self.alpha] = self.survival_function_.values**(np.exp(-coef*cumulative_sq_/np.log(self.survival_function_.values)))
      return df

  def _variance_f(self,N,d):
     if N==d==0:
        return 0
     return 1.*d/(N*(N-d))


def _additive_estimate(event_times, timeline, observed, additive_f, initial, columns, variance_f,
                      nelson_aalen_smoothing=False):
    """

    nelson_aalen_smoothing: see section 3.1.3 in Survival and Event History Analysis

    """
    timeline = timeline.astype(float)
    if timeline[0] > 0:
       timeline = np.insert(timeline,0,0.)

    n = timeline.shape[0]
    _additive_estimate_ = pd.DataFrame(np.zeros((n,1)), index=timeline, columns=columns)
    _additive_var = pd.DataFrame(np.zeros((n,1)), index=timeline)

    N = event_times["removed"].sum()
    t_0 =0
    _additive_estimate_.ix[(timeline<t_0)]=initial
    _additive_var.ix[(timeline<t_0)]=0

    v = initial
    v_sq = 0
    for t, removed, observed_deaths in event_times.itertuples():
        _additive_estimate_.ix[ (t_0<=timeline)*(timeline<t) ] = v  
        _additive_var.ix[ (t_0<=timeline)*(timeline<t) ] = v_sq
        missing = removed - observed_deaths
        N -= missing
        if nelson_aalen_smoothing:
           v += np.sum([additive_f(N-i,1) for i in range(int(observed_deaths))])
           v_sq += np.sum([variance_f(N-i,1) for i in range(int(observed_deaths))])
        else:
           v += additive_f(N,observed_deaths)
           v_sq += variance_f(N,observed_deaths)
        N -= observed_deaths
        t_0 = t
    _additive_estimate_.ix[(timeline>=t)]=v
    _additive_var.ix[(timeline>=t)]=v_sq
    return _additive_estimate_, _additive_var


class AalenAdditiveFitter(object):

  def __init__(self,fit_intercept=True, alpha=0.95, penalizer = 0.0):
    self.fit_intercept = fit_intercept
    self.alpha = alpha
    self.penalizer = penalizer

  def fit(self, event_times, X, timeline = None, censorship=None, columns=None):
    """currently X is a static (n,d) array<commit_msg>refactoring _addtive_f for NA, which helps _additive_estimate<commit_after>import numpy as np
from numpy.linalg import LinAlgError, inv
from numpy import dot
import pandas as pd

from lifelines.plotting import plot_dataframes
from lifelines.utils import dataframe_from_events_censorship, basis

import pdb



class NelsonAalenFitter(object):
    """
    Class for fitting the Nelson-Aalen estimate for the cumulative hazard. 

    NelsonAalenFitter( alpha=0.95, nelson_aalen_smoothing=True)

    alpha: The alpha value associated with the confidence intervals. 
    nelson_aalen_smoothing: If the event times are naturally discrete (like discrete years, minutes, etc.)
      then it is advisable to turn this parameter to False. See [1], pg.84.

    """
    def __init__(self, alpha=0.95, nelson_aalen_smoothing=True):
        self.alpha = alpha
        self.nelson_aalen_smoothing = nelson_aalen_smoothing
        if nelson_aalen_smoothing:
          self._variance_f = self._variance_f_smooth
          self._additive_f = self._additive_f_smooth
        else:
          self._variance_f = self._variance_f_discrete
          self._additive_f = self._additive_f_discrete

    def fit(self, event_times,censorship=None, timeline=None, columns=['NA-estimate']):
        """
        Parameters:
          event_times: an (n,1) array of times that the death event occured at 
          timeline: return the best estimate at the values in timelines (postively increasing)

        Returns:
          DataFrame with index either event_times or timelines (if not None), with
          values as the NelsonAalen estimate
        """
        
        if censorship is None:
           self.censorship = np.ones_like(event_times, dtype=bool) #why boolean?
        else:
           self.censorship = censorship.copy()

        self.event_times = dataframe_from_events_censorship(event_times, self.censorship)

        if timeline is None:
           self.timeline = self.event_times.index.values.copy()
        else:
           self.timeline = timeline
        self.cumulative_hazard_, cumulative_sq_ = _additive_estimate(self.event_times, 
                                                                     self.timeline, self.censorship, 
                                                                     self._additive_f, 0, columns, self._variance_fgi )
        self.confidence_interval_ = self._bounds(cumulative_sq_)
        self.plot = plot_dataframes(self, "cumulative_hazard_")

        return

    def _bounds(self, cumulative_sq_):
        inverse_norm = { 0.95:1.96, 0.99:2.57 }
        try:
          coef = inverse_norm[self.alpha]
        except:
          pass
        df = pd.DataFrame( index=self.timeline)
        df["upper_bound_%.2f"%self.alpha] = self.cumulative_hazard_.values*np.exp(coef*np.sqrt(cumulative_sq_)/self.cumulative_hazard_.values )
        df["lower_bound_%.2f"%self.alpha] = self.cumulative_hazard_.values*np.exp(-coef*np.sqrt(cumulative_sq_)/self.cumulative_hazard_.values )
        return df

    def _variance_f_smooth(self, N, d):
        if N==d==0:
            return 0
        return np.sum([1./(N-i)**2 for i in range(d)])

    def _variance_f_discrete(self, N, d):
        if N==d==0:
            return 0
        return 1.*(N-d)*d/N**3

    def _additive_f_smooth(self, N, d):
        if N==d==0:
          return 0
        return np.sum([1./(N-i) for i in range(d)])

    def _additive_f_discrete(self, N,d ):
       #check it 0
       if N==d==0:
          return 0
       return 1.*d/N


class KaplanMeierFitter(object):
   
  def __init__(self, alpha = 0.95):
       self.alpha = alpha

  def fit(self, event_times,censorship=None, timeline=None, columns=['KM-estimate']):
       """
       Parameters:
         event_times: an (n,1) array of times that the death event occured at 
         timeline: return the best estimate at the values in timelines (postively increasing)
         censorship: an (n,1) array of booleans -- True if the the death was observed, False if the event 
            was lost (right-censored). Defaults all True if censorship==None
       Returns:
         DataFrame with index either event_times or timelines (if not None), with
         values as the NelsonAalen estimate
       """
       #need to sort event_times
       if censorship is None:
          self.censorship = np.ones_like(event_times, dtype=bool) #why boolean?
       else:
          self.censorship = censorship.copy()

       self.event_times = dataframe_from_events_censorship(event_times, self.censorship)

       if timeline is None:
          self.timeline = self.event_times.index.values.copy()
       else:
          self.timeline = timeline
       log_surivial_function, cumulative_sq_ = _additive_estimate(self.event_times, self.timeline, 
                                                                  self.censorship, self._additive_f, 
                                                                  0, columns, self._variance_f )
       self.survival_function_ = np.exp(log_surivial_function)
       self.median_ = median_survival_times(self.survival_function_)
       self.confidence_interval_ = self._bounds(cumulative_sq_)
       self.plot = plot_dataframes(self, "survival_function_")
       return self

  def _additive_f(self,N,d):
      if N==d==0:
        return 0
      return np.log(1 - 1.*d/N)

  def _bounds(self, cumulative_sq_):
      inverse_norm = { 0.95:1.96, 0.99:2.57 }
      try:
        coef = inverse_norm[self.alpha]
      except:
        pass
      df = pd.DataFrame( index=self.timeline)
      df["upper_bound_%.2f"%self.alpha] = self.survival_function_.values**(np.exp(coef*cumulative_sq_/np.log(self.survival_function_.values)))
      df["lower_bound_%.2f"%self.alpha] = self.survival_function_.values**(np.exp(-coef*cumulative_sq_/np.log(self.survival_function_.values)))
      return df

  def _variance_f(self,N,d):
     if N==d==0:
        return 0
     return 1.*d/(N*(N-d))


def _additive_estimate(event_times, timeline, observed, additive_f, initial, columns, variance_f):
    """

    nelson_aalen_smoothing: see section 3.1.3 in Survival and Event History Analysis

    """
    timeline = timeline.astype(float)
    if timeline[0] > 0:
       timeline = np.insert(timeline,0,0.)

    n = timeline.shape[0]
    _additive_estimate_ = pd.DataFrame(np.zeros((n,1)), index=timeline, columns=columns)
    _additive_var = pd.DataFrame(np.zeros((n,1)), index=timeline)

    N = event_times["removed"].sum()
    t_0 =0
    _additive_estimate_.ix[(timeline<t_0)]=initial
    _additive_var.ix[(timeline<t_0)]=0

    v = initial
    v_sq = 0
    for t, removed, observed_deaths in event_times.itertuples():
        _additive_estimate_.ix[ (t_0<=timeline)*(timeline<t) ] = v  
        _additive_var.ix[ (t_0<=timeline)*(timeline<t) ] = v_sq
        missing = removed - observed_deaths
        N -= missing
        v += additive_f(N,observed_deaths)
        v_sq += variance_f(N,observed_deaths)
        N -= observed_deaths
        t_0 = t
    _additive_estimate_.ix[(timeline>=t)]=v
    _additive_var.ix[(timeline>=t)]=v_sq
    return _additive_estimate_, _additive_var


class AalenAdditiveFitter(object):

  def __init__(self,fit_intercept=True, alpha=0.95, penalizer = 0.0):
    self.fit_intercept = fit_intercept
    self.alpha = alpha
    self.penalizer = penalizer

  def fit(self, event_times, X, timeline = None, censorship=None, columns=None):
    """currently X is a static (n,d) array
{'fraction_non_alphanumeric': 0.09875557330139083, 'fraction_numerical': 0.01204498569242031, 'mean_word_length': 2.5177902621722845, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 25, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3271883', 'n_tokens_mistral': 4992, 'n_tokens_neox': 4655, 'n_words': 1040}