<jupyter_start><jupyter_text>Variational inference for mean, dispersion, and correlationSame generative model, but with log-likelihood prior and variational inference.For now, using ADVI and log-normal likelihoods.<jupyter_code>import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt
import scipy as sp
from statsmodels.stats.moment_helpers import cov2corr
import pymc3 as pm
import pandas as pd
import theano
import theano.tensor as tt<jupyter_output><empty_output><jupyter_text>Scratch/Testing<jupyter_code>mean = [1, 1.5]
cov =[[3, 1], [1, 2]]
normal_rvs = stats.multivariate_normal.rvs(mean=mean, cov=cov, size=100000)
lognormal_rvs = np.exp(normal_rvs)


log_mean = np.log(lognormal_rvs.mean(axis=0))
log_disp = np.log(lognormal_rvs.var(axis=0)) - log_mean
corr = cov2corr(np.cov(lognormal_rvs, rowvar=False))
a, b = moment_to_lognorm_params(log_mean, log_disp, corr)
a.eval()
b.eval()<jupyter_output><empty_output><jupyter_text>Define custom random variables in PyMC3 and scipy<jupyter_code>def generalized_binom_coef(x,y):
    
#     if type(x) == int and type(y) == np.ndarray:
#         x = np.full_like(y, x)
#     if type(y) == int and type(y) == np.ndarray:
#         y = np.full_like(x, y)

    return sp.special.gamma(x+1) / (sp.special.gamma(y+1) * sp.special.gamma(x-y+1))

class gen_binom(stats.rv_discrete):
    
    def _pmf(self, k, n, p):
        
        return generalized_binom_coef(n, k) * p ** (k) * (1-p)**(n-k)
import theano.tensor as tt

class Generalized_binomial(pm.Discrete):
    
    def __init__(self, n, p, *args, **kwargs):
        super(Generalized_binomial, self).__init__(*args, **kwargs)
        self.n = n
        self.p = p
        self.mean = n*p
    
    def logp(self, value):
        
        tt.gammaln(self.n+1) - \
        tt.gammaln(value+1) - \
        tt.gammaln(self.n-value+1) + \
        (self.n-value)*tt.log(1-self.p)  + \
        value*tt.log(self.p)<jupyter_output><empty_output><jupyter_text>ADVI<jupyter_code>true_log_mean = 1
true_log_disp = 2
N=10000
capture_eff = 0.1
true_mean = np.exp(true_log_mean)
true_disp = np.exp(true_log_disp)
true_sigma_sq = np.log(true_disp/true_mean + 1)
true_mu = true_log_mean - true_sigma_sq/2
z_true = stats.lognorm.rvs(
    s=np.sqrt(true_sigma_sq), 
    scale=np.exp(true_mu),
    size=N)
x_obs = gen_binom().rvs(n=z_true, p=capture_eff)
plt.hist(x_obs, bins=50);
import theano.tensor as tt

class Generalized_binomial(pm.Discrete):
    
    def __init__(self, n, p, *args, **kwargs):
        super(Generalized_binomial, self).__init__(*args, **kwargs)
        self.n = n
        self.p = p
        self.mean = n*p
    
    def logp(self, value):
        
        return tt.gammaln(self.n+1) - \
        tt.gammaln(value+1) - \
        tt.gammaln(self.n-value+1) + \
        (self.n-value)*tt.log(1-self.p)  + \
        value*tt.log(self.p)
with pm.Model() as model:
    
    log_mean = pm.Normal('log_mean', mu=true_log_mean, sigma=0.5, testval=true_log_mean)
    log_disp = pm.Normal('log_disp', mu=true_log_disp, sigma=0.5, testval=true_log_disp)
    
    mean = pm.math.exp(log_mean)
    disp = pm.math.exp(log_disp)
    
    sigma_sq = pm.math.log(disp/mean + 1)
    mu = log_mean - sigma_sq/2
    
    z = pm.Normal('z', mu=mu, sigma=pm.math.sqrt(sigma_sq), shape=N)

    x = Generalized_binomial('x', n=pm.math.exp(z), p=capture_eff, observed=x_obs, shape=N)
model.check_test_point()
with model:
    mean_field = pm.fit(method='advi')
plt.hist(mean_field.sample(4000)['log_disp'], bins=100);
plt.hist(mean_field.sample(3000)['log_mean'], bins=50);<jupyter_output><empty_output><jupyter_text>Auxillary functions to handle multivariate<jupyter_code>def expand_upper_triangular(n, packed):
    
    out = tt.ones((n, n), dtype=theano.config.floatX)
    idxs1 = np.triu_indices(n, k=1)
    idxs2 = np.tril_indices(n, k=-1)
    out = tt.set_subtensor(out[tuple(idxs1)], packed)
    out = tt.set_subtensor(out[tuple(idxs2)], packed)
    return out
def cov_from_corr(corr, normal_var):
    """
        Convert a covariance matrix of the exponentiated data
        to a covariance matrix of the original data.
    """

    base_col = tt.reshape(normal_var, (-1, 1))
    base_row = tt.reshape(normal_var, (1, -1))
    base_mtx = tt.sqrt(tt.dot(base_col, base_row))
    return corr*base_mtx
def get_ln_corr(normal_var, normal_cov):
    """
        Convert the correlation in normal space to the log normal space.
        
        corr is the correlation in logged (normal) space.
    """
    
    cv = tt.exp(normal_var)-1
    base_col = tt.reshape(cv, (-1, 1))
    base_row = tt.reshape(cv, (1, -1))
    base_mtx = tt.sqrt(tt.dot(base_col, base_row))
    ln_corr = (tt.exp(normal_cov)-1) / base_mtx
    
    return ln_corr
def moment_to_params(log_mean, log_disp, corr):
    """
        Convert a vector of log means, log indices of dispersion, and a correlation
        matrix to the parameters of the lognormal distribution, i.e. the
        parameters of the underlying normal distribution.
        
        corr is the correlation in the logged (normal) space.
    """
    
    mean = tt.exp(log_mean)
    disp = tt.exp(log_disp)
    var = tt.exp(log_disp + log_mean)
    normal_var = tt.log(1+var/mean**2)
    
    mu = log_mean - tt.log(disp/mean + 1)/2
    sigma = cov_from_corr(corr, normal_var)
    
    ln_corr = get_ln_corr(normal_var, sigma)
    return mu, sigma, ln_corr<jupyter_output><empty_output><jupyter_text>Multivariate case<jupyter_code>N = 1000
d = 3
p=0.1
random_mtx = np.random.normal(size=(d,d))
cov_mtx = random_mtx.T@random_mtx
var_mv = np.diag(cov_mtx)
log_mean_mv = np.random.normal(size=d)
log_disp_mv = np.log(var_mv)-log_mean_mv
corr_mv = cov2corr(cov_mtx)
corr_mv_vals = corr_mv[np.triu_indices(d, k=1)]
print('log-mean', log_mean_mv)
print('log-disp', log_disp_mv)
print('corr\n', corr_mv)
print('corr values', corr_mv_vals)

mu_mv, sigma_mv, ln_corr = moment_to_params(
    log_mean_mv,
    log_disp_mv,
    corr_mv)
mu_mv, sigma_mv, ln_corr = mu_mv.eval(), sigma_mv.eval(), ln_corr.eval()
print('ln_corr\n', ln_corr)

sigma_mv
with pm.Model() as model:
    
    log_mean = pm.Normal(
        'log_mean', 
        mu=log_mean_mv, 
        sigma=0.2, 
        shape=d,
        testval=log_mean_mv)
    
    tt.printing.Print('temp')(log_disp)
    
    log_disp = pm.Normal(
        'log_disp', 
        mu=log_disp_mv,
        sigma=0.2,
        shape=d,
        testval=log_disp_mv)
    
    tt.printing.Print('logdisp')(log_disp)
    corr = pm.LKJCorr(
        'corr',
        eta=0.1,
        n=d,
        testval=corr_mv_vals)
    
    tt.printing.Print('corr')(corr)
    
    corr_mat = expand_upper_triangular(n=d, packed=corr)
    
    mu, sigma, adjusted_corr = moment_to_lognorm_params(log_mean, log_disp, corr_mat)
    
    #ln_corr = pm.Deterministic('ln_corr', adjusted_corr)
    
    z = pm.MvNormal('z', mu=mu, cov=sigma, shape=(N,d), testval=mu)

    x = Generalized_binomial('x', n=pm.math.exp(z), p=p, observed=x_obs, shape=(N,d))
    
    
model.check_test_point()
with model:
    model = pm.sample(100)
    #mean_field = pm.fit(method='advi')<jupyter_output>Only 100 samples in chain.
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
/anaconda3/envs/sc_ht/lib/python3.7/site-packages/theano/tensor/basic.py:6611: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.
  result[diagonal_slice] = x
/anaconda3/envs/sc_ht/lib/python3.7/site-packages/theano/tensor/basic.py:6611: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.
  result[diagonal_slice] = x
/anaconda3/envs/sc_ht/lib/python3.7/site-packages/theano/tensor/subtensor.py:2339: FutureWarning: Using a non-tuple sequence for multidim[...]
{'fraction_non_alphanumeric': 0.13262045231071778, 'fraction_numerical': 0.015363815142576204, 'mean_word_length': 3.498065229408513, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '650089', 'license': 'MIT', 'n_tokens_mistral': 3269, 'n_tokens_neox': 3107, 'n_words': 620, 'path': 'analysis/simulation/legacy/simple/variational_inference.ipynb', 'repo_name': 'yelabucsf/scrna-parameter-estimation', 'size': '8136'}