<s> import numpy as np <EOL> import pandas as pd <EOL> from scipy . stats . distributions import chi2 , norm <EOL> from statsmodels . graphics import utils <EOL> def _calc_survfunc_right ( time , status ) : <EOL> """<STR_LIT>""" <EOL> time = np . asarray ( time ) <EOL> status = np . asarray ( status ) <EOL> time , rtime = np . unique ( time , return_inverse = True ) <EOL> d = np . bincount ( rtime , weights = status ) <EOL> n = np . bincount ( rtime ) <EOL> n = np . cumsum ( n [ : : - <NUM_LIT:1> ] ) [ : : - <NUM_LIT:1> ] <EOL> ii = np . flatnonzero ( d > <NUM_LIT:0> ) <EOL> d = d [ ii ] <EOL> n = n [ ii ] <EOL> time = time [ ii ] <EOL> sp = <NUM_LIT:1> - d / n . astype ( np . float64 ) <EOL> sp = np . log ( sp ) <EOL> sp = np . cumsum ( sp ) <EOL> sp = np . exp ( sp ) <EOL> se = d / ( n * ( n - d ) ) . astype ( np . float64 ) <EOL> se = np . cumsum ( se ) <EOL> se = np . sqrt ( se ) <EOL> se *= sp <EOL> return sp , se , time , n , d <EOL> class SurvfuncRight ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , time , status , title = None ) : <EOL> self . time = np . asarray ( time ) <EOL> self . status = np . asarray ( status ) <EOL> m = len ( status ) <EOL> x = _calc_survfunc_right ( time , status ) <EOL> self . surv_prob = x [ <NUM_LIT:0> ] <EOL> self . surv_prob_se = x [ <NUM_LIT:1> ] <EOL> self . surv_times = x [ <NUM_LIT:2> ] <EOL> self . n_risk = x [ <NUM_LIT:3> ] <EOL> self . n_events = x [ <NUM_LIT:4> ] <EOL> self . title = "<STR_LIT>" if not title else title <EOL> def plot ( self , ax = None ) : <EOL> """<STR_LIT>""" <EOL> return plot_survfunc ( self , ax ) <EOL> def quantile ( self , p ) : <EOL> """<STR_LIT>""" <EOL> ii = np . flatnonzero ( self . surv_prob < <NUM_LIT:1> - p ) <EOL> if len ( ii ) == <NUM_LIT:0> : <EOL> return np . nan <EOL> return self . surv_times [ ii [ <NUM_LIT:0> ] ] <EOL> def quantile_ci ( self , p , alpha = <NUM_LIT> , method = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> tr = norm . ppf ( <NUM_LIT:1> - alpha / <NUM_LIT:2> ) <EOL> method = method . lower ( ) <EOL> if method == "<STR_LIT>" : <EOL> g = lambda x : np . log ( - np . log ( x ) ) <EOL> gprime = lambda x : - <NUM_LIT:1> / ( x * np . log ( x ) ) <EOL> elif method == "<STR_LIT>" : <EOL> g = lambda x : x <EOL> gprime = lambda x : <NUM_LIT:1> <EOL> elif method == "<STR_LIT>" : <EOL> g = lambda x : np . log ( x ) <EOL> gprime = lambda x : <NUM_LIT:1> / x <EOL> elif method == "<STR_LIT>" : <EOL> g = lambda x : np . log ( x / ( <NUM_LIT:1> - x ) ) <EOL> gprime = lambda x : <NUM_LIT:1> / ( x * ( <NUM_LIT:1> - x ) ) <EOL> elif method == "<STR_LIT>" : <EOL> g = lambda x : np . arcsin ( np . sqrt ( x ) ) <EOL> gprime = lambda x : <NUM_LIT:1> / ( <NUM_LIT:2> * np . sqrt ( x ) * np . sqrt ( <NUM_LIT:1> - x ) ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> r = g ( self . surv_prob ) - g ( <NUM_LIT:1> - p ) <EOL> r /= ( gprime ( self . surv_prob ) * self . surv_prob_se ) <EOL> ii = np . flatnonzero ( np . abs ( r ) <= tr ) <EOL> if len ( ii ) == <NUM_LIT:0> : <EOL> return np . nan , np . nan <EOL> lb = self . surv_times [ ii [ <NUM_LIT:0> ] ] <EOL> if ii [ - <NUM_LIT:1> ] == len ( self . surv_times ) - <NUM_LIT:1> : <EOL> ub = np . inf <EOL> else : <EOL> ub = self . surv_times [ ii [ - <NUM_LIT:1> ] + <NUM_LIT:1> ] <EOL> return lb , ub <EOL> def summary ( self ) : <EOL> """<STR_LIT>""" <EOL> df = pd . DataFrame ( index = self . surv_times ) <EOL> df . index . name = "<STR_LIT>" <EOL> df [ "<STR_LIT>" ] = self . surv_prob <EOL> df [ "<STR_LIT>" ] = self . surv_prob_se <EOL> df [ "<STR_LIT>" ] = self . n_risk <EOL> df [ "<STR_LIT>" ] = self . n_events <EOL> return df <EOL> def simultaneous_cb ( self , alpha = <NUM_LIT> , method = "<STR_LIT>" , transform = "<STR_LIT>" ) : <EOL> """<STR_LIT>""" <EOL> method = method . lower ( ) <EOL> if method != "<STR_LIT>" : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if alpha != <NUM_LIT> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> transform = transform . lower ( ) <EOL> s2 = self . surv_prob_se ** <NUM_LIT:2> / self . surv_prob ** <NUM_LIT:2> <EOL> nn = self . n_risk <EOL> if transform == "<STR_LIT>" : <EOL> denom = np . sqrt ( nn ) * np . log ( self . surv_prob ) <EOL> theta = <NUM_LIT> * ( <NUM_LIT:1> + nn * s2 ) / denom <EOL> theta = np . exp ( theta ) <EOL> lcb = self . surv_prob ** ( <NUM_LIT:1> / theta ) <EOL> ucb = self . surv_prob ** theta <EOL> elif transform == "<STR_LIT>" : <EOL> k = <NUM_LIT> <EOL> k *= ( <NUM_LIT:1> + nn * s2 ) / ( <NUM_LIT:2> * np . sqrt ( nn ) ) <EOL> k *= np . sqrt ( self . surv_prob / ( <NUM_LIT:1> - self . surv_prob ) ) <EOL> f = np . arcsin ( np . sqrt ( self . surv_prob ) ) <EOL> v = np . clip ( f - k , <NUM_LIT:0> , np . inf ) <EOL> lcb = np . sin ( v ) ** <NUM_LIT:2> <EOL> v = np . clip ( f + k , - np . inf , np . pi / <NUM_LIT:2> ) <EOL> ucb = np . sin ( v ) ** <NUM_LIT:2> <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> return lcb , ucb <EOL> def survdiff ( time , status , group , weight_type = None , strata = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> time = np . asarray ( time ) <EOL> status = np . asarray ( status ) <EOL> group = np . asarray ( group ) <EOL> gr = np . unique ( group ) <EOL> if len ( gr ) != <NUM_LIT:2> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if strata is None : <EOL> obs , var = _survdiff ( time , status , group , weight_type , gr , <EOL> ** kwargs ) <EOL> else : <EOL> strata = np . asarray ( strata ) <EOL> stu = np . unique ( strata ) <EOL> obs , var = <NUM_LIT:0.> , <NUM_LIT:0.> <EOL> for st in stu : <EOL> ii = ( strata == st ) <EOL> obs1 , var1 = _survdiff ( time [ ii ] , status [ ii ] , group [ ii ] , <EOL> weight_type , gr , ** kwargs ) <EOL> obs += obs1 <EOL> var += var1 <EOL> zstat = obs / np . sqrt ( var ) <EOL> chisq = zstat ** <NUM_LIT:2> <EOL> pvalue = <NUM_LIT:1> - chi2 . cdf ( chisq , <NUM_LIT:1> ) <EOL> return chisq , pvalue <EOL> def _survdiff ( time , status , group , weight_type , gr , ** kwargs ) : <EOL> ii = ( group == gr [ <NUM_LIT:0> ] ) <EOL> time1 = time [ ii ] <EOL> status1 = status [ ii ] <EOL> ii = ( group == gr [ <NUM_LIT:1> ] ) <EOL> time2 = time [ ii ] <EOL> status2 = status [ ii ] <EOL> utimes = np . unique ( time ) <EOL> status1 = status1 . astype ( np . bool ) <EOL> status2 = status2 . astype ( np . bool ) <EOL> ix1 = np . searchsorted ( utimes , time1 [ status1 ] ) <EOL> ix2 = np . searchsorted ( utimes , time2 [ status2 ] ) <EOL> obs1 = np . bincount ( ix1 , minlength = len ( utimes ) ) <EOL> obs2 = np . bincount ( ix2 , minlength = len ( utimes ) ) <EOL> obs = obs1 + obs2 <EOL> nvec = [ ] <EOL> for time0 in time1 , time2 : <EOL> ix = np . searchsorted ( utimes , time0 ) <EOL> n = np . bincount ( ix , minlength = len ( utimes ) ) <EOL> n = np . cumsum ( n ) <EOL> n = np . roll ( n , <NUM_LIT:1> ) <EOL> n [ <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> n = len ( time0 ) - n <EOL> nvec . append ( n ) <EOL> n1 , n2 = tuple ( nvec ) <EOL> n = n1 + n2 <EOL> r = n1 / n . astype ( np . float64 ) <EOL> var = obs * r * ( <NUM_LIT:1> - r ) * ( n - obs ) / ( n - <NUM_LIT:1> ) <EOL> exp1 = obs * r <EOL> weights = None <EOL> if weight_type is not None : <EOL> weight_type = weight_type . lower ( ) <EOL> if weight_type == "<STR_LIT>" : <EOL> weights = n <EOL> elif weight_type == "<STR_LIT>" : <EOL> weights = np . sqrt ( n ) <EOL> elif weight_type == "<STR_LIT>" : <EOL> if "<STR_LIT>" not in kwargs : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> fh_p = kwargs [ "<STR_LIT>" ] <EOL> sp = <NUM_LIT:1> - obs / n . astype ( np . float64 ) <EOL> sp = np . log ( sp ) <EOL> sp = np . cumsum ( sp ) <EOL> sp = np . exp ( sp ) <EOL> weights = sp ** fh_p <EOL> weights = np . roll ( weights , <NUM_LIT:1> ) <EOL> weights [ <NUM_LIT:0> ] = <NUM_LIT:1> <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> ix = np . flatnonzero ( n > <NUM_LIT:1> ) <EOL> if weights is None : <EOL> obs = np . sum ( obs1 [ ix ] - exp1 [ ix ] ) <EOL> var = np . sum ( var [ ix ] ) <EOL> else : <EOL> obs = np . dot ( weights [ ix ] , obs1 [ ix ] - exp1 [ ix ] ) <EOL> var = np . dot ( weights [ ix ] ** <NUM_LIT:2> , var [ ix ] ) <EOL> return obs , var <EOL> def plot_survfunc ( survfuncs , ax = None ) : <EOL> """<STR_LIT>""" <EOL> fig , ax = utils . create_mpl_ax ( ax ) <EOL> try : <EOL> assert ( type ( survfuncs [ <NUM_LIT:0> ] ) is SurvfuncRight ) <EOL> except : <EOL> survfuncs = [ survfuncs ] <EOL> for gx , sf in enumerate ( survfuncs ) : <EOL> surv_times = np . concatenate ( ( [ <NUM_LIT:0> ] , sf . surv_times ) ) <EOL> surv_prob = np . concatenate ( ( [ <NUM_LIT:1> ] , sf . surv_prob ) ) <EOL> mxt = max ( sf . time ) <EOL> if mxt > surv_times [ - <NUM_LIT:1> ] : <EOL> surv_times = np . concatenate ( ( surv_times , [ mxt ] ) ) <EOL> surv_prob = np . concatenate ( ( surv_prob , [ surv_prob [ - <NUM_LIT:1> ] ] ) ) <EOL> label = getattr ( sf , "<STR_LIT:title>" , "<STR_LIT>" % ( gx + <NUM_LIT:1> ) ) <EOL> li , = ax . step ( surv_times , surv_prob , '<STR_LIT:->' , label = label , lw = <NUM_LIT:2> , where = '<STR_LIT>' ) <EOL> ii = np . flatnonzero ( np . logical_not ( sf . status ) ) <EOL> ti = sf . time [ ii ] <EOL> jj = np . searchsorted ( surv_times , ti ) - <NUM_LIT:1> <EOL> sp = surv_prob [ jj ] <EOL> ax . plot ( ti , sp , '<STR_LIT:+>' , ms = <NUM_LIT:12> , color = li . get_color ( ) , <EOL> label = label + "<STR_LIT>" ) <EOL> ax . set_ylim ( <NUM_LIT:0> , <NUM_LIT> ) <EOL> return fig </s>
