Skip to content

Instantly share code, notes, and snippets.

@mwt

mwt/main.py

Last active Mar 6, 2021
Embed
What would you like to do?
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