Skip to content

Instantly share code, notes, and snippets.

@mwt
Last active April 11, 2023 07:23
Show Gist options
  • Save mwt/25f69988105de9a801f6b7fae5e6375d to your computer and use it in GitHub Desktop.
Save mwt/25f69988105de9a801f6b7fae5e6375d to your computer and use it in GitHub Desktop.
GPU Accelerated Theil-Sen Estimator for Censored Data
## 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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment