<s> from statsmodels . compat . python import range <EOL> import numpy as np <EOL> from math import sqrt <EOL> import matplotlib . pyplot as plt <EOL> class KAPLAN_MEIER ( object ) : <EOL> def __init__ ( self , data , timesIn , groupIn , censoringIn ) : <EOL> raise RuntimeError ( '<STR_LIT>' ) <EOL> self . data = data <EOL> self . timesIn = timesIn <EOL> self . groupIn = groupIn <EOL> self . censoringIn = censoringIn <EOL> def fit ( self ) : <EOL> groups = list ( set ( self . data [ : , self . groupIn ] ) ) <EOL> groupList = [ ] <EOL> for i in range ( len ( groups ) ) : <EOL> groupList . append ( [ ] ) <EOL> for i in range ( len ( groups ) ) : <EOL> for j in range ( len ( self . data ) ) : <EOL> if self . data [ j , self . groupIn ] == groups [ i ] : <EOL> groupList [ i ] . append ( self . data [ j ] ) <EOL> timeList = [ ] <EOL> for i in range ( len ( groupList ) ) : <EOL> times = [ ] <EOL> for j in range ( len ( groupList [ i ] ) ) : <EOL> times . append ( groupList [ i ] [ j ] [ self . timesIn ] ) <EOL> times = list ( sorted ( set ( times ) ) ) <EOL> timeList . append ( times ) <EOL> timeCounts = [ ] <EOL> points = [ ] <EOL> censoredPoints = [ ] <EOL> for i in range ( len ( groupList ) ) : <EOL> survival = <NUM_LIT:1> <EOL> varSum = <NUM_LIT:0> <EOL> riskCounter = len ( groupList [ i ] ) <EOL> counts = [ ] <EOL> x = [ ] <EOL> y = [ ] <EOL> for j in range ( len ( timeList [ i ] ) ) : <EOL> if j != <NUM_LIT:0> : <EOL> if j == <NUM_LIT:1> : <EOL> groupInd = <NUM_LIT:1> <EOL> x . append ( <NUM_LIT:0> ) <EOL> y . append ( <NUM_LIT:1> ) <EOL> x . append ( timeList [ i ] [ j - <NUM_LIT:1> ] ) <EOL> y . append ( <NUM_LIT:1> ) <EOL> x . append ( timeList [ i ] [ j - <NUM_LIT:1> ] ) <EOL> y . append ( survival ) <EOL> y . append ( survival ) <EOL> else : <EOL> groupInd = <NUM_LIT:0> <EOL> y . append ( survival ) <EOL> y . append ( survival ) <EOL> x . append ( timeList [ i ] [ j - <NUM_LIT:1> ] ) <EOL> x . append ( timeList [ i ] [ j - <NUM_LIT:1> ] ) <EOL> censoredPoints . append ( [ timeList [ i ] [ j - <NUM_LIT:1> ] , <EOL> censoringNum , survival , groupInd ] ) <EOL> counts . append ( [ timeList [ i ] [ j - <NUM_LIT:1> ] , riskCounter , <EOL> eventCounter , survival , <EOL> sqrt ( ( ( survival ) ** <NUM_LIT:2> ) * varSum ) ] ) <EOL> riskCounter += - <NUM_LIT:1> * ( riskChange ) <EOL> riskChange = <NUM_LIT:0> <EOL> eventCounter = <NUM_LIT:0> <EOL> censoringCounter = <NUM_LIT:0> <EOL> censoringNum = <NUM_LIT:0> <EOL> for k in range ( len ( groupList [ i ] ) ) : <EOL> if ( groupList [ i ] [ k ] [ self . timesIn ] ) == ( timeList [ i ] [ j ] ) : <EOL> riskChange += <NUM_LIT:1> <EOL> if groupList [ i ] [ k ] [ self . censoringIn ] == <NUM_LIT:1> : <EOL> eventCounter += <NUM_LIT:1> <EOL> else : <EOL> censoringNum += <NUM_LIT:1> <EOL> if eventCounter != censoringCounter : <EOL> censoringCounter = eventCounter <EOL> survival *= ( ( float ( riskCounter ) - <EOL> eventCounter ) / ( riskCounter ) ) <EOL> try : <EOL> varSum += ( eventCounter ) / ( ( riskCounter ) <EOL> * ( float ( riskCounter ) - <EOL> eventCounter ) ) <EOL> except ZeroDivisionError : <EOL> varSum = <NUM_LIT:0> <EOL> counts . append ( [ timeList [ i ] [ len ( timeList [ i ] ) - <NUM_LIT:1> ] , <EOL> riskCounter , eventCounter , survival , <EOL> sqrt ( ( ( survival ) ** <NUM_LIT:2> ) * varSum ) ] ) <EOL> x . append ( timeList [ i ] [ len ( timeList [ i ] ) - <NUM_LIT:1> ] ) <EOL> x . append ( timeList [ i ] [ len ( timeList [ i ] ) - <NUM_LIT:1> ] ) <EOL> y . append ( survival ) <EOL> censoredPoints . append ( [ timeList [ i ] [ len ( timeList [ i ] ) - <NUM_LIT:1> ] , <EOL> censoringNum , survival , <NUM_LIT:1> ] ) <EOL> timeCounts . append ( np . array ( counts ) ) <EOL> points . append ( [ x , y ] ) <EOL> self . results = timeCounts <EOL> self . points = points <EOL> self . censoredPoints = censoredPoints <EOL> def plot ( self ) : <EOL> x = [ ] <EOL> for i in range ( len ( self . points ) ) : <EOL> plt . plot ( np . array ( self . points [ i ] [ <NUM_LIT:0> ] ) , np . array ( self . points [ i ] [ <NUM_LIT:1> ] ) ) <EOL> x += self . points [ i ] [ <NUM_LIT:0> ] <EOL> for j in range ( len ( self . censoredPoints ) ) : <EOL> if ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] != <NUM_LIT:0> ) : <EOL> if ( self . censoredPoints [ j ] [ <NUM_LIT:3> ] == <NUM_LIT:1> ) and ( j == <NUM_LIT:0> ) : <EOL> dx = ( ( <NUM_LIT:1.> / ( ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] ) + <NUM_LIT:1.> ) ) <EOL> * ( float ( self . censoredPoints [ j ] [ <NUM_LIT:0> ] ) ) ) <EOL> for k in range ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] ) : <EOL> plt . vlines ( ( <NUM_LIT:1> + ( ( k + <NUM_LIT:1> ) * dx ) ) , <EOL> self . censoredPoints [ j ] [ <NUM_LIT:2> ] - <NUM_LIT> , <EOL> self . censoredPoints [ j ] [ <NUM_LIT:2> ] + <NUM_LIT> ) <EOL> elif ( ( self . censoredPoints [ j ] [ <NUM_LIT:3> ] == <NUM_LIT:1> ) and <EOL> ( self . censoredPoints [ j - <NUM_LIT:1> ] [ <NUM_LIT:3> ] == <NUM_LIT:1> ) ) : <EOL> dx = ( ( <NUM_LIT:1.> / ( ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] ) + <NUM_LIT:1.> ) ) <EOL> * ( float ( self . censoredPoints [ j ] [ <NUM_LIT:0> ] ) ) ) <EOL> for k in range ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] ) : <EOL> plt . vlines ( ( <NUM_LIT:1> + ( ( k + <NUM_LIT:1> ) * dx ) ) , <EOL> self . censoredPoints [ j ] [ <NUM_LIT:2> ] - <NUM_LIT> , <EOL> self . censoredPoints [ j ] [ <NUM_LIT:2> ] + <NUM_LIT> ) <EOL> elif j == ( len ( self . censoredPoints ) - <NUM_LIT:1> ) : <EOL> dx = ( ( <NUM_LIT:1.> / ( ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] ) + <NUM_LIT:1.> ) ) <EOL> * ( float ( self . censoredPoints [ j ] [ <NUM_LIT:0> ] ) ) ) <EOL> for k in range ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] ) : <EOL> plt . vlines ( ( self . censoredPoints [ j - <NUM_LIT:1> ] [ <NUM_LIT:0> ] + ( ( k + <NUM_LIT:1> ) * dx ) ) , <EOL> self . censoredPoints [ j ] [ <NUM_LIT:2> ] - <NUM_LIT> , <EOL> self . censoredPoints [ j ] [ <NUM_LIT:2> ] + <NUM_LIT> ) <EOL> else : <EOL> dx = ( ( <NUM_LIT:1.> / ( ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] ) + <NUM_LIT:1.> ) ) <EOL> * ( float ( self . censoredPoints [ j + <NUM_LIT:1> ] [ <NUM_LIT:0> ] ) <EOL> - self . censoredPoints [ j ] [ <NUM_LIT:0> ] ) ) <EOL> for k in range ( self . censoredPoints [ j ] [ <NUM_LIT:1> ] ) : <EOL> plt . vlines ( ( self . censoredPoints [ j ] [ <NUM_LIT:0> ] + ( ( k + <NUM_LIT:1> ) * dx ) ) , <EOL> self . censoredPoints [ j ] [ <NUM_LIT:2> ] - <NUM_LIT> , <EOL> self . censoredPoints [ j ] [ <NUM_LIT:2> ] + <NUM_LIT> ) <EOL> plt . xlim ( ( <NUM_LIT:0> , np . max ( x ) ) ) <EOL> plt . ylim ( ( <NUM_LIT:0> , <NUM_LIT> ) ) <EOL> plt . xlabel ( '<STR_LIT:time>' ) <EOL> plt . ylabel ( '<STR_LIT>' ) <EOL> plt . show ( ) <EOL> def show_results ( self ) : <EOL> resultsString = '<STR_LIT>' <EOL> for i in range ( len ( self . results ) ) : <EOL> resultsString += ( '<STR_LIT>' . format ( i ) + <EOL> '<STR_LIT>' ) <EOL> for j in self . results [ i ] : <EOL> resultsString += ( <EOL> '<STR_LIT>' . format ( <EOL> int ( j [ <NUM_LIT:0> ] ) , int ( j [ <NUM_LIT:1> ] ) , int ( j [ <NUM_LIT:2> ] ) , j [ <NUM_LIT:3> ] , j [ <NUM_LIT:4> ] ) ) <EOL> print ( resultsString ) </s>
