gp-radon-model-broken

gp-radon-model-broken.py

from pymc3 import (
    NUTS,
    Deterministic,
    HalfCauchy,
    Model,
    MvNormal,
    find_MAP,
    sample,
    summary,
    traceplot,
)
import aesara.tensor as T
import numpy as np
import pymc3 as pm

np.random.seed(20090425)
n = 30
p = 10
p_dash = 2
q= 3
X = np.random.rand(n,p)
y = np.random.rand(n,p_dash)
Z = np.random.rand(n,q)
true_scale = 3
from aesara import shared

def kernel(x_1, x_2, k):
    """Kernel function."""
    z = np.dot(x_1.T, k)
    
    return np.dot(z, x_2)

def squared_distance(x,y):
    return np.array([[np.mean(x[i,:] - y[j,:]) ** 2 for i in range(len(x))] for j in range(len(y))])

with pm.Model() as model:

    sigma_p = pm.HalfCauchy("sigma_p", 5)
    non_linear = np.random.rand(n,p)
    
    # trying to add gaussian process to the model to acccount for non-linear relation between X and y
    mu = np.zeros(n)

    eta_sq   = HalfCauchy("η_sq", 5)
    rho_sq   = HalfCauchy("ρ_sq", 5)
    sigma_sq = HalfCauchy("σ_sq", 5)

    D = (squared_distance(X, X))
    non_linear = pm.MvNormal('non_linear',mu=np.zeros(n), cov=D, shape=(p_dash, n))

    beta = pm.MvNormal("beta", mu=np.zeros(p_dash), cov=sigma_p*np.eye(p_dash), shape=(p,p_dash))
    u    = pm.MvNormal("u", mu=np.zeros(p_dash), cov=sigma_p*np.eye(p_dash), shape=(q,p_dash))
    
    # Model error
    eps = pm.HalfCauchy("eps", 5)
    print(pm.math.dot(X,beta))
    print(Z.shape)
    # model
    radon_est = pm.math.dot(Z,u) + pm.math.dot(X,beta) + non_linear.T
    
    # Setup right cholesky matrix
    sd_dist = pm.HalfCauchy.dist(beta=2.5, shape=p_dash)
    colchol_packed = pm.LKJCholeskyCov('colcholpacked', n=p_dash, eta=2,
                                       sd_dist=sd_dist)
    colchol = pm.expand_packed_triangular(p_dash, colchol_packed)

    # Setup left covariance matrix
    scale = pm.Lognormal('scale', mu=np.log(true_scale), sigma=0.5)
    rowcov = T.nlinalg.diag([scale**(2*i) for i in range(n)])
    
    y_obs = pm.MatrixNormal("y_obs", mu=radon_est, colchol=colchol, rowcov=rowcov, observed=y, shape=y.shape)
    
with model:
    trace = pm.sample(2000)