The following Python class computes and draws Kaplan-Meier product limit estimators for given data. An example of how to use the class follows the code.

# Code

# load useful libraries import matplotlib.pyplot as plt # # class for building Kaplan-Meier product limit estimator # class KM(object): # # constructor # def __init__(self, measured_values, censored_or_not): self.measured_values = measured_values; self.censored_or_not = censored_or_not # log measured values where event occurred self.event_occurred_times = {} self.censored_occurred_times = {} for i, mv in enumerate(self.measured_values): if self.censored_or_not[i] == 1: if not self.event_occurred_times.has_key(mv): self.event_occurred_times[mv] = 0 self.event_occurred_times[mv] += 1 else: if not self.censored_occurred_times.has_key(mv): self.censored_occurred_times[mv] = 0 self.censored_occurred_times[mv] += 1 # construct list of j values j_list = [0] j_list.extend(sorted(self.event_occurred_times.keys())) j_list.append(max(max(self.event_occurred_times.keys()), max(self.censored_occurred_times.keys())) + 1) self.j_list = j_list # count censored values in interval [j, j+1), index by j self.number_of_units_censored_in_interval_j_to_j_plus_one = {} for i in range(0, len(j_list)-1): j = j_list[i] j_plus_one = j_list[i+1] m_count = 0 for m in self.censored_occurred_times.keys(): if j <= m and m < j_plus_one: m_count = m_count + self.censored_occurred_times[m] self.number_of_units_censored_in_interval_j_to_j_plus_one[j] = m_count # calculate number of units at risk just prior to time t_j self.number_of_units_at_risk_just_prior_to_time_j = {} for i in range(0, len(j_list)-1): j = j_list[i] n_count = 0. for k in range(i, len(j_list)-1): jk = j_list[k] if self.event_occurred_times.has_key(jk): n_count = n_count + self.event_occurred_times[jk] if self.number_of_units_censored_in_interval_j_to_j_plus_one.has_key(jk): n_count = n_count + self.number_of_units_censored_in_interval_j_to_j_plus_one[jk] self.number_of_units_at_risk_just_prior_to_time_j[j] = n_count # add time zero count to self.event_occurred_times self.event_occurred_times[0] = 0 # build the estimator for each time j self.S = {} for i in range(0, len(j_list)-1): j = j_list[i] prod = 1. for k in range(0, i+1): jk = j_list[k] prod = prod * ((self.number_of_units_at_risk_just_prior_to_time_j[jk] - self.event_occurred_times[jk]) / self.number_of_units_at_risk_just_prior_to_time_j[jk]<p class="western"></p>) self.S[j] = prod # # display the estimator in the console # def display(self): print print '\ttime\tn.risk\tn.event\tsurvival' for i in range(1, len(self.j_list) - 1): j = self.j_list[i] print '\t' + str(j) + '\t' + str(int(self.number_of_units_at_risk_just_prior_to_time_j[j])) + '\t' + str(self.event_occurred_times[j]) + '\t' + str(round(self.S[j], 3)) print # # plot # def plot(self, color, xlabel, title): j_list_sans_end = self.j_list[0:-1] # plot the curve plt.figure() for i in range(0, len(j_list_sans_end) - 1): j = j_list_sans_end[i] j_plus_one = j_list_sans_end[i+1] plt.plot([j, j_plus_one], [self.S[j], self.S[j]], color=color) plt.plot([j_plus_one, j_plus_one], [self.S[j], self.S[j_plus_one]], color=color) # set the axis limits plt.xlim([0, self.j_list[-1]]) plt.ylim([-0.05, 1.05]) # add the censored cases within the curve for i in sorted(self.censored_occurred_times.keys()): last_S = 1. for j in sorted(self.S.keys()): if j >= i: plt.scatter(i, last_S, color=color) break<p class="western"></p><p class="western"></p> last_S = self.S[j] # add the censored cases beyond the curve for i in sorted(self.censored_occurred_times.keys()): max_S_key = max(self.S.keys()) if i > max_S_key: plt.scatter(i, self.S[max_S_key], color="blue") # show the plot plt.ylabel("Survival") plt.xlabel(xlabel) plt.title(title) plt.show()

# Usage Example

import matplotlib.pyplot as plt times = [3, 4, 4, 5, 7, 6, 8, 8, 12, 14, 14, 19, 20, 20] events = [1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 0] # 0 for censored cases, 1 for event occurred my_km = KM(times, events) my_km.display() my_km.plot("blue", "Time", "Kaplan-Meier Product Limit Estimator Example")