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