Skip to content

Commit

Permalink
Refactor Exponential, Weibull and Gamma to be special cases of Genera…
Browse files Browse the repository at this point in the history 
…lizedGamma
  • Loading branch information
Erik Bernhardsson committed Apr 10, 2018
1 parent 4279860 commit 2eaf8fc
Show file tree
Hide file tree
Showing 4 changed files with 58 additions and 90 deletions.
3 changes: 2 additions & 1 deletion convoys/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import random
import seaborn
from matplotlib import pyplot
from convoys.multi import Exponential, Weibull, Gamma, KaplanMeier, Nonparametric
from convoys.multi import *


def get_timescale(t):
Expand Down Expand Up @@ -59,6 +59,7 @@ def get_groups(data, group_min_size, max_groups):
'exponential': Exponential,
'weibull': Weibull,
'gamma': Gamma,
'generalized-gamma': GeneralizedGamma,
}


Expand Down
4 changes: 4 additions & 0 deletions convoys/multi.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,10 @@ class Gamma(RegressionToMulti):
base_model_cls = regression.Gamma


class GeneralizedGamma(RegressionToMulti):
base_model_cls = regression.GeneralizedGamma


class KaplanMeier(SingleToMulti):
base_model_cls = single.KaplanMeier

Expand Down
124 changes: 35 additions & 89 deletions convoys/regression.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,57 +38,36 @@ class RegressionModel:
pass


class Exponential(RegressionModel):
def fit(self, X, B, T):
n, k = X.shape
class GeneralizedGamma(RegressionModel):
# https://en.wikipedia.org/wiki/Generalized_gamma_distribution
# Note however that lambda is a^-1 in WP's notation
# Note also that k = d/p so d = k*p
def fit(self, X, B, T, k=None, p=None):
n_features = X.shape[1]
X_batch, B_batch, T_batch = tf_utils.get_batch_placeholders((X, B, T))

X_batch = tf.nn.dropout(X_batch, keep_prob=0.5)
a = LinearCombination(X_batch, k)
b = LinearCombination(X_batch, k)
a = LinearCombination(X_batch, n_features)
b = LinearCombination(X_batch, n_features)
lambd = tf.exp(a.y)
c = tf.sigmoid(b.y)

log_pdf = tf.log(lambd) - T_batch*lambd
cdf = 1 - tf.exp(-(T_batch * lambd))

LL_observed = tf.log(c) + log_pdf
LL_censored = tf.log((1-c) + c * (1 - cdf))

LL = tf.reduce_sum(B_batch * LL_observed + (1 - B_batch) * LL_censored, 0)

with tf.Session() as sess:
feed_dict = {X_batch: X, B_batch: B, T_batch: T}
tf_utils.optimize(sess, LL, feed_dict)
self.params = {
'a': a.params(sess, LL, feed_dict),
'b': b.params(sess, LL, feed_dict),
}

def cdf(self, x, t, ci=None, n=1000):
t = numpy.array(t)
a = LinearCombination.sample(self.params['a'], x, ci, n)
b = LinearCombination.sample(self.params['b'], x, ci, n)
return tf_utils.predict(expit(b) * (1 - numpy.exp(numpy.multiply.outer(-t, numpy.exp(a)))), ci)


class Weibull(RegressionModel):
def fit(self, X, B, T):
n, k = X.shape
X_batch, B_batch, T_batch = tf_utils.get_batch_placeholders((X, B, T))
if k is None:
k = tf.Variable(1.0, name='k', trainable=False)
should_update_k = True
else:
k = tf.constant(k, tf.float32)
should_update_k = False

log_k_var = tf.Variable(tf.zeros([]), name='log_k')
X_batch = tf.nn.dropout(X_batch, keep_prob=0.5)
a = LinearCombination(X_batch, k)
b = LinearCombination(X_batch, k)
k = tf.exp(log_k_var)
lambd = tf.exp(a.y)
c = tf.sigmoid(b.y)
if p is None:
log_p_var = tf.Variable(tf.zeros([]), name='log_p')
p = tf.exp(log_p_var)
else:
p = tf.constant(p, tf.float32)

# PDF of Weibull: k * lambda * (x * lambda)^(k-1) * exp(-(t * lambda)^k)
log_pdf = tf.log(k) + tf.log(lambd) + (k-1)*(tf.log(T_batch) + tf.log(lambd)) - (T_batch*lambd)**k
# CDF of Weibull: 1 - exp(-(t * lambda)^k)
cdf = 1 - tf.exp(-(T_batch * lambd)**k)
# PDF: p*lambda^(k*p) / gamma(k) * t^(k*p-1) * exp(-(x*lambda)^p)
log_pdf = tf.log(p) + (k*p) * tf.log(lambd) - tf.lgamma(k) + (k*p-1) * tf.log(T_batch) - (T_batch*lambd)**p
cdf = tf.igamma(k, (T_batch*lambd)**p)

LL_observed = tf.log(c) + log_pdf
LL_censored = tf.log((1-c) + c * (1 - cdf))
Expand All @@ -97,66 +76,33 @@ def fit(self, X, B, T):

with tf.Session() as sess:
feed_dict = {X_batch: X, B_batch: B, T_batch: T}
tf_utils.optimize(sess, LL, feed_dict)
tf_utils.optimize(sess, LL, feed_dict,
update_callback=(tf_utils.get_tweaker(sess, LL, k, feed_dict)
if should_update_k else None))
self.params = {
'a': a.params(sess, LL, feed_dict),
'b': b.params(sess, LL, feed_dict),
'k': sess.run(k),
'p': sess.run(p),
}

def cdf(self, x, t, ci=None, n=1000):
t = numpy.array(t)
a = LinearCombination.sample(self.params['a'], x, ci, n)
b = LinearCombination.sample(self.params['b'], x, ci, n)
return tf_utils.predict(expit(b) * (1 - numpy.exp(-numpy.multiply.outer(t, numpy.exp(a))**self.params['k'])), ci)
return tf_utils.predict(expit(b) * gammainc(self.params['k'], numpy.multiply.outer(t, numpy.exp(a))**self.params['p']), ci)


class Gamma(RegressionModel):
class Exponential(GeneralizedGamma):
def fit(self, X, B, T):
n, k = X.shape
X_batch, B_batch, T_batch = tf_utils.get_batch_placeholders((X, B, T))

X_batch = tf.nn.dropout(X_batch, keep_prob=0.5)
a = LinearCombination(X_batch, k)
b = LinearCombination(X_batch, k)
k = tf.Variable(2.0, name='k', trainable=False)
lambd = tf.exp(a.y)
c = tf.sigmoid(b.y)

# PDF of gamma: 1.0 / gamma(k) * lambda ^ k * t^(k-1) * exp(-t * lambda)
log_pdf = -tf.lgamma(k) + k*tf.log(lambd) + (k-1)*tf.log(T_batch) - lambd*T_batch
# CDF of gamma: gammainc(k, lambda * t)
cdf = tf.igamma(k, lambd * T_batch)

LL_observed = tf.log(c) + log_pdf
LL_censored = tf.log((1-c) + c * (1 - cdf))

LL = tf.reduce_sum(B_batch * LL_observed + (1 - B_batch) * LL_censored, 0)
feed_dict = {X_batch: X, B_batch: B, T_batch: T}
super(Exponential, self).fit(X, B, T, k=1, p=1)

new_k = tf.placeholder(tf.float32, shape=[])
assign_k = tf.assign(k, new_k)

def update_k(sess):
# tf.igamma doesn't compute the gradient wrt a properly
# So let's just try small perturbations
k_value = sess.run(k)
res = {}
for k_mult in [0.97, 1.0, 1.03]:
sess.run(assign_k, feed_dict={new_k: k_value * k_mult})
res[k_value * k_mult] = sess.run(LL, feed_dict=feed_dict)
sess.run(assign_k, feed_dict={new_k: max(res.keys(), key=res.get)})
class Weibull(GeneralizedGamma):
def fit(self, X, B, T):
super(Weibull, self).fit(X, B, T, k=1)

with tf.Session() as sess:
tf_utils.optimize(sess, LL, feed_dict, update_callback=update_k)
self.params = {
'a': a.params(sess, LL, feed_dict),
'b': b.params(sess, LL, feed_dict),
'k': sess.run(k),
}

def cdf(self, x, t, ci=None, n=1000):
t = numpy.array(t)
a = LinearCombination.sample(self.params['a'], x, ci, n)
b = LinearCombination.sample(self.params['b'], x, ci, n)
return tf_utils.predict(expit(b) * gammainc(self.params['k'], numpy.multiply.outer(t, numpy.exp(a))), ci)
class Gamma(GeneralizedGamma):
def fit(self, X, B, T):
super(Gamma, self).fit(X, B, T, p=1)
17 changes: 17 additions & 0 deletions convoys/tf_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,23 @@ def optimize(sess, target, placeholders, batch_size=1024, update_callback=None):
update_callback(sess)


def get_tweaker(sess, target, z, feed_dict):
new_z = tf.placeholder(tf.float32, shape=z.shape)
assign_z = tf.assign(z, new_z)

def tweak_z(sess):
# tf.igamma doesn't compute the gradient wrt a properly
# So let's just try small perturbations
z_value = sess.run(z)
res = {}
for z_mult in [0.97, 1.0, 1.03]:
sess.run(assign_z, feed_dict={new_z: z_value * z_mult})
res[z_value * z_mult] = sess.run(target, feed_dict=feed_dict)
sess.run(assign_z, feed_dict={new_z: max(res.keys(), key=res.get)})

return tweak_z


def predict(func_values, ci):
if ci is None:
return func_values
Expand Down

0 comments on commit 2eaf8fc

Please sign in to comment.