Fine mapping idea

def elbo(n, p, d, xr, xrk, y, pip, mean, var, effect_var, residual_var):
  error = np.square(y - xr).sum()
  error -= np.square(xrk).sum()
  # var doesn't depend on factor?
  error += (d * (pip * (np.square(mean) + var.T)).sum(axis=0)).sum()
  error /= residual_var
  error += n * np.log(residual_var)
  error *= -.5

  # Assume prior probability 1/p for each variant
  kl_z = (pip * (np.log(pip) + np.log(p))).sum()
  kl_b = .5 * (pip * (1 + np.log(effect_var * residual_var) - np.log(var.T) + (np.square(mean) + var.T) / (effect_var * residual_var))).sum()
  return error, kl_z, kl_b

def coordinate_ascent(x, y, effect_var, residual_var, l=5, num_epochs=200):
  n, p = x.shape

  # Prior
  pi = np.ones((p, 1))

  # Variational parameters
  pip = np.zeros((l, p))
  mean = np.zeros((l, p))
  # Make sure everything is two dimensional to catch numpy broadcasting gotchas
  d = np.einsum('ij,ij->j', x, x).reshape(-1, 1)
  var = (effect_var * residual_var / (effect_var * d + 1)).reshape(-1, 1)

  # Precompute more things
  xy = x.T.dot(y)
  xr = np.dot(x, (pip * mean).sum(axis=0, keepdims=True).T)

  elbo_ = None
  trace = []
  for epoch in range(num_epochs):
    for k in range(l):
      xr -= np.dot(x, (pip * mean)[k:k + 1].T)
      mean[k:k + 1] = (var / residual_var * (xy - x.T.dot(xr))).T
      pip[k:k + 1] = (pi * np.exp(.5 * (np.log(var / (effect_var * residual_var)) + np.square(mean[k:k + 1].T) / var))).T
      pip[k] /= pip[k].sum()
      xr += np.dot(x, (pip * mean)[k:k + 1].T)
    xrk = x.dot((pip * mean).T)
    error, kl_z, kl_b = elbo(n, p, d, xr, xrk, y, pip, mean, var, effect_var, residual_var)
    update = error - kl_z - kl_b
    trace.append([update, error, kl_z, kl_b])
    elbo_ = update
  return {'pip': pip,
          'mean': pip * mean,
          'var': var,
          'elbo': elbo_,
          'trace': trace}

def plot_categorical_slab_fit(s, opt):
  plt.clf()
  q = np.logical_or(s.theta != 0, opt['pip'].sum(axis=0) > 0.1)
  fig, ax = plt.subplots(3, 1, gridspec_kw={'height_ratios': [1, 1, opt['pip'].shape[0]]})
  fig.set_size_inches(6, 8)
  ax[0].bar(np.arange(np.sum(q)), s.theta[q])
  ax[0].set_xticks(np.arange(q.sum()))
  ax[0].set_ylabel('True effect size')
  ax[1].bar(np.arange(np.sum(q)), opt['mean'].sum(axis=0)[q])
  ax[1].set_ylabel('Estimated effect size')
  im = ax[2].imshow(opt['pip'][:,q], cmap=colorcet.cm['Reds'], vmin=0, vmax=1)
  ax[2].set_yticks(np.arange(opt['pip'].shape[0]))
  ax[2].set_ylabel('Posterior inclusion probability')
  ax[2].set_xlabel('True and false positive variants')
  plt.colorbar(im, ax=ax[2], orientation='horizontal')

x = np.random.binomial(2, 0.5, size=(500, 1000))
y = np.random.normal(size=(500, 1))
opt = coordinate_ascent(x, y, effect_var=1, residual_var=y.var(), l=5, num_epochs=10)
pd.DataFrame(opt['trace'], columns=['elbo', 'error', 'kl_z', 'kl_b'])

plot_categorical_slab_fit(s, opt)

with nwas.simulation.simulation(p=1000, pve=0.5, annotation_params=[(10, 1)], seed=0) as s:
  x, y = s.sample_gaussian(n=500)
  x = x.astype('float32')
  y = y.reshape(-1, 1).astype('float32')

opt = coordinate_ascent(x, y, effect_var=1, residual_var=s.residual_var, l=5, num_epochs=10)
pd.DataFrame(opt['trace'], columns=['elbo', 'error', 'kl_z', 'kl_b'])

plot_categorical_slab_fit(s, opt)

opt = coordinate_ascent(x, y, effect_var=1, residual_var=s.residual_var, l=1, num_epochs=10)
pd.DataFrame(opt['trace'], columns=['elbo', 'error', 'kl_z', 'kl_b'])

plot_categorical_slab_fit(s, opt)

opt = coordinate_ascent(x, y, effect_var=1, residual_var=s.residual_var, l=2, num_epochs=10)
pd.DataFrame(opt['trace'], columns=['elbo', 'error', 'kl_z', 'kl_b'])

plot_categorical_slab_fit(s, opt)

srun --partition=broadwl --mem=8G plink --memory 8000 --vcf /project/compbio/geuvadis/genotypes/GEUVADIS.chr1.PH1PH2_465.IMPFRQFILT_BIALLELIC_PH.annotv2.genotypes.vcf.gz --make-bed --out 1

geuvadis_chr1 = pyplink.PyPlink(os.path.join(os.getenv('SCRATCH'), 'spikeslab', '1'))
x = np.zeros((geuvadis_chr1.get_nb_samples(), s.p), dtype=np.float32)
for i, (_, geno) in enumerate(geuvadis_chr1):
  if i >= x.shape[1]:
    break
  x[:,i] = geno.astype(np.float32)
x = np.ma.masked_equal(x, -1)

def plot_ld(x):
  w = x - x.mean(axis=0)
  w /= 1e-8 + w.std(axis=0)
  corr = w.T.dot(w) / w.shape[0]
  plt.clf()
  plt.gcf().set_size_inches(8, 8)
  plt.imshow(np.triu(corr), cmap='RdBu_r', vmin=-1, vmax=1)
  cb = plt.colorbar()
  cb.ax.set_xlabel('Correlation')

plot_ld(x)

s.estimate_mafs(x)
annotation = np.zeros(s.p)
annotation[600:610] = np.ones(10)
s.load_annotations(annotation)
s.sample_effects(pve=0.5, annotation_params=[(0, 1), (2, 1)], permute=True)
x = (x - x.mean(axis=0)).filled(0)
y = s.compute_liabilities(x).reshape(-1, 1)

opt = coordinate_ascent(x, y, effect_var=1, residual_var=s.residual_var, l=5, num_epochs=10)
pd.DataFrame(opt['trace'], columns=['elbo', 'llik', 'kl_z', 'kl_b'])

plot_categorical_slab_fit(s, opt)