GPU Accelerated Theil-Sen Estimator for Censored Data

main.py

## Define Akritas et al T-S estimator
def censoredts(data):
  """
  This function takes a three-column cupy array where the first column is the
  independent variable, the second is the dependent variable, and the third is 
  an indicator which tells us if the data has been censored.
  """
  # Make an index for pairwise treatment of data
  from itertools import combinations
  indit = zip(*combinations(range(n),2))
  ind0 = list(next(indit))
  ind1 = list(next(indit))
  x0,y0,d0 = data[:,ind0]
  x1,y1,d1 = data[:,ind1]
  import cupy as cp
  # Run with Numpy if nparray given
  xp = cp.get_array_module(data)
  # Generate a list of betas to try
  bs = xp.sort(((y0-y1)/(x0-x1))[xp.logical_or(d0,d1)]).reshape(-1,1)
  # Residuals
  r0 = y0 - bs*x0
  r1 = y1 - bs*x1
  # Compute the objective function for each beta
  Tnb = xp.sum((2*((x0 < x1) - 1/2 )) * ( d0*(r0 < r1) - d1*(r1 < r0) ),axis=1)
  return float((bs[Tnb>0][-1]+bs[Tnb<0][0])/2)

#%% Example
n = 100
beta = 2
import cupy as cp
# Generate x and error term
x = cp.random.exponential(scale=10, size=n)
u = cp.random.normal(loc=0, scale=3, size=n)
# Set the y value acording to
yt = (beta*x+u)
# Compute the 80th percentile
yc = cp.percentile(yt,80)
# Generate indicator variable for truth
d = (yt < yc)
# Censor values above 80% cutoff
y = (yt+(yt > yc)*(yc-yt))
# Combine into a single array
data = cp.vstack((x,y,d))
# Get our result
print(censoredts(data))