def define_model(x, n_total, K):
"""
Model definition in Edward2
Parameters
----------
x -- numpy array [NxD]
covariate matrix
n_total -- numpy array [N]
number of cells per sample
K -- int
Number of cell types
"""
N, D = x.shape
dtype = tf.float32
# normal prior on bias
alpha = ed.Normal(loc=tf.zeros([K]),
scale=tf.ones([K]) * 5,
name="alpha")
# Noncentered parametrization for raw slopes (before spike-and-slab)
mu_b = ed.Normal(loc=tf.zeros(1, dtype=dtype),
scale=tf.ones(1, dtype=dtype),
name="mu_b")
sigma_b = ed.HalfCauchy(tf.zeros(1, dtype=dtype),
tf.ones(1, dtype=dtype),
name="sigma_b")
b_offset = ed.Normal(loc=tf.zeros([D, K], dtype=dtype),
scale=tf.ones([D, K], dtype=dtype),
name="b_offset")
b_raw = mu_b + sigma_b * b_offset
# Spike-and-slab priors
sigma_ind_raw = ed.Normal(loc=tf.zeros(shape=[D, K], dtype=dtype),
scale=tf.ones(shape=[D, K], dtype=dtype),
name='sigma_ind_raw')
ind_t = sigma_ind_raw * 50
ind = tf.exp(ind_t) / (1 + tf.exp(ind_t))
# Calculate betas
beta = ind * b_raw
# Concentration vector from intercepts, slopes
concentration_ = tf.exp(alpha + tf.matmul(x, beta))
# Cell count prediction via DirMult
predictions = ed.DirichletMultinomial(n_total,
concentration=concentration_,
name="predictions")
return predictions
def define_model(x, n_total, K):
N, D = x.shape
dtype = tf.float32
alpha = ed.Normal(loc=tf.zeros([K], dtype=dtype),
scale=tf.ones([K], dtype=dtype),
name="alpha")
beta = ed.Normal(loc=tf.zeros([D, K], dtype=dtype),
scale=tf.ones([D, K], dtype=dtype),
name="beta")
concentration_ = tf.exp(alpha + tf.matmul(x, beta))
# Likelihood
predictions = ed.DirichletMultinomial(n_total,
concentration=concentration_,
name="predictions")
return predictions