marginal likelihood sampling

dummy_margl_sampling.py

+import os
+import pickle
+import random
+
+import numpy as np
+
+import matplotlib
+from matplotlib import pyplot as plt
+
+import random_fields as rf
+
+
+def main():
+    with open('./data/field.pickle', 'rb') as f:
+        xs, ys, field, nu, sigma, l = pickle.load(f)
+
+    print("Pre-evaluating distance matrix")
+    dist_mtx = rf.dist_matrix(xs, ys)
+    flat_gt_field = np.ravel(field)
+
+    noise_sigma = 1e-2
+    noise = np.random.normal(scale=noise_sigma, size=flat_gt_field.shape)
+
+    observed_field = flat_gt_field + noise
+    initial_theta = np.array([nu, sigma, l]) + 1e-2
+
+    proposal_cov = 1e-2 * np.eye(3)
+    lower = np.array([0.5, 0.5, 0.5])
+    upper = np.array([3, 3, 3])
+
+    print("Starting MCMC")
+    chain_res = rf.metropolis_sample_matern_theta(
+        initial_theta,
+        proposal_cov,
+        observed_field,
+        dist_mtx,
+        noise_sigma,
+        lower=lower, upper=upper
+    )
+
+    with open('./data/matern_chain.pickle', 'wb') as f:
+        pickle.dump(chain_res, f)
+
+    chain, p_accs, nlls = chain_res
+    figs_outpath = './figs/margl/'
+    os.makedirs(figs_outpath, exist_ok=True)
+    rf.show_theta_chain(chain, p_accs, nlls, figs_outpath)
+
+
+if __name__ == '__main__':
+    random.seed(42)
+    np.random.seed(42)
+
+    main()

random_fields.py

@@ -6,6 +6,9 @@ import scipy as sp
 
 from scipy.special import kv, gamma
 
+import matplotlib
+from matplotlib import pyplot as plt
+
 
 def eps():
     return 1e-2
@@ -64,15 +67,15 @@ def matern_parameter_nlpost(observed, field, field_cov, noise_sigma):
 
     chol_cov = np.linalg.cholesky(field_cov)
     cc_diag = np.diag(chol_cov)
-    log_det_prior = 2 * np.sum(np.log(cc_diag))
+    log_det_cov = 2 * np.sum(np.log(cc_diag))
 
-    return nll + field_nl_prior + log_det_prior
+    return nll + field_nl_prior + log_det_cov
 
 
 def gibbs_within_metropolis(initial_theta, observed_field, noise_sigma,
                             proposal_cov, lower, upper,
                             dist_mtx, initial_field=None,
-                            adaptation_start=1,
+                            adaptation_start=None,
                             n_samples=10000, n_save=1000):
     chain = []
     nlls = []
@@ -137,7 +140,7 @@ def gibbs_within_metropolis(initial_theta, observed_field, noise_sigma,
                 "Gibbs-Metropolis [{}]: rejected with p={}".format(step, acceptance))
 
         if step > (n_samples - n_save):
-            chain.append((state, field))
+            chain.append((np.exp(state), field))
 
         n_mean_intext = p_mean_intext + (state - p_mean_intext) / (step + 1)
 
@@ -156,3 +159,119 @@ def gibbs_within_metropolis(initial_theta, observed_field, noise_sigma,
         nlls.append(state_nlp)
 
     return chain, p_accs, nlls
+
+
+def metropolis_sample_matern_theta(initial_theta, proposal_cov, observed_field,
+                                   dist_mtx, noise_sigma,
+                                   n_samples=10, burn_in=0,
+                                   adaptation_start=None,
+                                   lower=None, upper=None):
+    chain = []
+    nlls = []
+    p_accs = []
+
+    state = np.log(initial_theta)
+    nu, sigma, l = initial_theta
+    state_nll = matern_marginal_likelihood(
+        observed_field, nu, sigma, l, dist_mtx, noise_sigma)
+
+    p_mean_intext = np.copy(state)
+    sd = 5.76 / state.size
+    eps = np.finfo(np.float32).eps
+    eps_eye = eps * np.eye(state.size)
+    proposal_cov += eps_eye
+    if adaptation_start is None:
+        adaptation_start = n_samples
+
+    for step in range(n_samples):
+        pr_chol = np.linalg.cholesky(proposal_cov)
+        dstate = np.matmul(pr_chol, np.random.normal(size=state.shape))
+        proposal_state = state + dstate
+
+        proposal_theta = np.exp(proposal_state)
+        up_lower = True if lower is None else np.all(proposal_theta >= lower)
+        lo_upper = True if upper is None else np.all(proposal_theta <= upper)
+        if not (lo_upper and up_lower):
+            print(
+                "MargL-Metropolis [{}]: rejected with out of bounds".format(step))
+            continue
+
+        nu, sigma, l = proposal_theta
+        proposal_nll = matern_marginal_likelihood(
+            observed_field, nu, sigma, l, dist_mtx, noise_sigma)
+
+        acceptance = np.exp(state_nll - proposal_nll)
+        accepted = np.random.random() < acceptance
+        if accepted:
+            state = proposal_state
+            state_nll = proposal_nll
+            print(
+                "MargL-Metropolis [{}]: acceptance with p={}".format(step, acceptance))
+        else:
+            print(
+                "MargL-Metropolis [{}]: rejected with p={}".format(step, acceptance))
+
+        n_mean_intext = p_mean_intext + (state - p_mean_intext) / (step + 1)
+
+        if step >= adaptation_start:
+            base = (step - 1) / step * proposal_cov
+            pmmt = np.outer(p_mean_intext, p_mean_intext) * step
+            nmmt = np.outer(n_mean_intext, n_mean_intext) * (step + 1)
+            sst = np.outer(state, state)
+            upd_cov = pmmt - nmmt + sst + eps_eye
+            proposal_cov = base + sd / step * upd_cov
+
+        if step >= burn_in:
+            chain.append(np.exp(state))
+            nlls.append(state_nll)
+            p_accs.append((acceptance, accepted))
+
+    return chain, p_accs, nlls
+
+
+def matern_marginal_likelihood(observed_field, nu, sigma, l, dist_mtx, noise_sigma):
+    cov_mtx = matern_cov_dm(dist_mtx, nu, sigma, l)
+    i_mtx = np.eye(cov_mtx.shape[0])
+    cov = cov_mtx + (noise_sigma**2) * i_mtx
+    inv_cov = np.linalg.inv(cov)
+
+    chol_cov = np.linalg.cholesky(cov)
+    cc_diag = np.diag(chol_cov)
+    log_det_cov = 2 * np.sum(np.log(cc_diag))
+
+    return log_det_cov + np.matmul(np.matmul(observed_field, inv_cov), observed_field)
+
+
+def plot_and_save_hist(xs, title, outpath):
+    plt.figure()
+    plt.title(title)
+    plt.hist(xs, bins='auto')
+    plt.savefig(outpath)
+    plt.close()
+
+
+def plot_and_save_plot(ys, title, outpath):
+    plt.figure()
+    plt.title(title)
+    plt.plot(ys)
+    plt.savefig(outpath)
+    plt.close()
+
+
+def show_theta_chain(chain, p_accs, nlls, outpath):
+    arr_chain = np.array(chain)
+    plot_and_save_hist(arr_chain[:, 0], "nu",
+                       os.path.join(outpath, "nu_hist.png"))
+    plot_and_save_hist(arr_chain[:, 1], "sigma",
+                       os.path.join(outpath, "sigma_hist.png"))
+    plot_and_save_hist(arr_chain[:, 2], "l",
+                       os.path.join(outpath, "l_hist.png"))
+
+    p_acceptances, acceptances = zip(*p_accs)
+
+    plot_and_save_plot(p_acceptances, "p_acceptances",
+                       os.path.join(outpath, "p_acceptances.png"))
+    plot_and_save_plot(acceptances, "accepted",
+                       os.path.join(outpath, "accepted.png"))
+    plot_and_save_plot(nlls, "nlls",
+                       os.path.join(outpath, "nlls.png"))