I must be doing something dumb then:
import gzip
import numpy as np
import pymc3 as pm
import pandas as pd
import theano.tensor as tt
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
class MvNormalReg(pm.Continuous):
"""Multivariate normal for a regularized decomposed covariance.
The covariance matrix has the form
Sigma = σ1(A + σ2 I)
where A = Q @ np.diag(diag) @ Q.T for orthogonal Q.
"""
def __init__(self, mu, Q, diag, sigma1, sigma2, *args, **kwargs):
self._mu = tt.as_tensor_variable(mu)
self._Q = tt.as_tensor_variable(Q)
self._diag = tt.as_tensor_variable(diag)
self._sigma1 = tt.as_tensor_variable(sigma1)
self._sigma2 = tt.as_tensor_variable(sigma2)
self.mean = self.median = self._mu
super(MvNormalReg, self).__init__(*args, **kwargs)
def logp(self, value):
if value.ndim != 1:
raise ValueError('Only one-dimensional values are supported.')
y = value - self._mu
Q = self._Q
n = value.shape[-1]
d = self._diag + self._sigma2
logdet = tt.sum(tt.log(d))
logdet += n * tt.log(self._sigma1)
trans = tt.dot(Q, tt.dot(Q.T, y) / d) / self._sigma1
quaddist = (y * trans).sum()
norm = n * pm.floatX(np.log(2 * np.pi))
return norm - 0.5 * (quaddist + logdet)
with pm.Model():
"""# my data; ignore this
df = pd.read_csv(gzip.open('phi2.gz'), delim_whitespace=True)
df.columns = ('id1', 'id2', 'phi2', 'd7')
df = pd.pivot_table(df, values='phi2', index='id1', columns='id2')
df = df.fillna(0)
A = df.values
"""
# randomly generated A
d, k = 1000, 1500
W = np.random.randn(d, k)
S = np.dot(W, W.T) + np.diag(np.random.rand(1, d))
A = np.dot(np.diag(1. / np.sqrt(np.diag(S))),
np.dot(S, np.diag(1. / np.sqrt(np.diag(S)))))
n = A.shape[-1]
X = np.random.normal(size=(n, 2)) # fixed-effect covariates
true_β = [0.1, 0.5]
true_σ_A = .5
true_σ_E = .5
eigs, Q = np.linalg.eigh(A)
true_μ = np.dot(X, true_β)
true_Σ = true_σ_A * A + true_σ_E * np.identity(A.shape[-1])
y = multivariate_normal.rvs(true_μ, true_Σ)
print(multivariate_normal.logpdf(y, true_μ, true_Σ))
print(MvNormalReg.dist(true_μ, Q, eigs, true_σ_A, true_σ_E / true_σ_A).logp(y).eval({}))
β = pm.Cauchy(name='β', alpha=0., beta=10., shape=2)
σ_A = pm.HalfCauchy(name='σ_A', beta=10.)
σ_E = pm.HalfCauchy(name='σ_E', beta=10.)
MvNormalReg('y', tt.dot(X, β), Q, eigs, sigma1=σ_A, sigma2=σ_E / σ_A,
observed=y)
trace = pm.sample()
pm.traceplot(trace)
plt.savefig('traceplot.png')
Obviously this returns something different every time because the data are different, but here is the last thing:
-1377.36542449
1379.4501627243912
