Skip to content

Commit

Permalink
Rename example files.
Browse files Browse the repository at this point in the history
  • Loading branch information
@alicjapolanska
alicjapolanska committed Nov 15, 2023
1 parent 23b2e10 commit 28662c7
Show file tree
Hide file tree
Showing 11 changed files with 1,097 additions and 1,097 deletions.
0 examples/gaussian_diagcov.py → examples/gaussian_diagcov_legacy.py
View file
File renamed without changes.
278 changes: 107 additions & 171 deletions examples/gaussian_nondiagcov.py
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,10 @@
import numpy as np
import emcee
import sys
import time
import matplotlib.pyplot as plt
from functools import partial
from matplotlib import cm
import harmonic as hm

sys.path.append("examples")
import ex_utils
import jax
import jax.numpy as jnp
import emcee


def ln_analytic_evidence(ndim, cov):
Expand All @@ -30,6 +26,7 @@ def ln_analytic_evidence(ndim, cov):
return ln_norm_lik


# @partial(jax.jit, static_argnums=(1,))
def ln_posterior(x, inv_cov):
"""Compute log_e of posterior.
Expand All @@ -45,7 +42,7 @@ def ln_posterior(x, inv_cov):
"""

return -np.dot(x, np.dot(inv_cov, x)) / 2.0
return -jnp.dot(x, jnp.dot(inv_cov, x)) / 2.0


def init_cov(ndim):
Expand Down Expand Up @@ -73,84 +70,119 @@ def init_cov(ndim):


def run_example(
flow_type,
ndim=2,
nchains=100,
samples_per_chain=1000,
nburn=500,
plot_corner=False,
plot_surface=False,
):
"""Run Gaussian example with non-diagonal covariance matrix.
Args:
flow_type: Which flow model to use, "RealNVP" or "RQSpline".
ndim: Dimension.
nchains: Number of chains.
samples_per_chain: Number of samples per chain.
nburn: Number of burn in samples for each chain.
plot_corner: Plot marginalised distributions if true.
plot_surface: Plot surface and samples if true.
"""

savefigs = True

# Initialise covariance matrix.
cov = init_cov(ndim)
inv_cov = np.linalg.inv(cov)
inv_cov = jnp.linalg.inv(cov)
training_proportion = 0.5
if flow_type == "RealNVP":
epochs_num = 5
elif flow_type == "RQSpline":
epochs_num = 3

save_name_start = "examples/plots/" + flow_type

temperature = 0.8
standardize = True
verbose = True

# Spline params
n_layers = 5
n_bins = 5
hidden_size = [32, 32]
spline_range = (-10.0, 10.0)

# Start timer.
clock = time.process_time()

# Run multiple realisations.
n_realisations = 2
n_realisations = 1
evidence_inv_summary = np.zeros((n_realisations, 3))
for i_realisation in range(n_realisations):
if n_realisations > 0:
hm.logs.info_log(
"Realisation = {}/{}".format(i_realisation + 1, n_realisations)
)

# Set up and run sampler.
hm.logs.info_log("Run sampling...")

pos = np.random.rand(ndim * nchains).reshape((nchains, ndim))

sampler = emcee.EnsembleSampler(nchains, ndim, ln_posterior, args=[inv_cov])
rstate = np.random.get_state() # Set random state to repeatable
# across calls.
(pos, prob, state) = sampler.run_mcmc(pos, samples_per_chain, rstate0=rstate)
samples = np.ascontiguousarray(sampler.chain[:, nburn:, :])
lnprob = np.ascontiguousarray(sampler.lnprobability[:, nburn:])
# Define the number of dimensions and the mean of the Gaussian
num_samples = nchains * samples_per_chain
# Initialize a PRNG key (you can use any valid key)
key = jax.random.PRNGKey(0)
mean = jnp.zeros(ndim)

# Generate random samples from the 2D Gaussian distribution
samples = jax.random.multivariate_normal(key, mean, cov, shape=(num_samples,))
lnprob = jax.vmap(ln_posterior, in_axes=(0, None))(samples, jnp.array(inv_cov))
samples = jnp.reshape(samples, (nchains, -1, ndim))
lnprob = jnp.reshape(lnprob, (nchains, -1))

MCMC = False
if MCMC:
nburn = 500
# Set up and run sampler.
hm.logs.info_log("Run sampling...")

pos = np.random.rand(ndim * nchains).reshape((nchains, ndim))

sampler = emcee.EnsembleSampler(nchains, ndim, ln_posterior, args=[inv_cov])
rstate = np.random.get_state() # Set random state to repeatable
# across calls.
(pos, prob, state) = sampler.run_mcmc(
pos, samples_per_chain, rstate0=rstate
)
samples = np.ascontiguousarray(sampler.chain[:, nburn:, :])
lnprob = np.ascontiguousarray(sampler.lnprobability[:, nburn:])

# Calculate evidence using harmonic....

# Set up chains.
chains = hm.Chains(ndim)
chains.add_chains_3d(samples, lnprob)
chains_train, chains_test = hm.utils.split_data(
chains, training_proportion=0.05
)

# Fit model.
hm.logs.info_log("Fit model...")

r_scale = np.sqrt(ndim - 1)
hm.logs.debug_log("r scale = {}".format(r_scale))
domains = [r_scale * np.array([1e-1, 1e0])]
hm.logs.debug_log("Domain = {}".format(domains))
model = hm.model_legacy.HyperSphere(ndim, domains)
fit_success, objective = model.fit(
chains_train.samples, chains_train.ln_posterior
chains, training_proportion=training_proportion
)

hm.logs.info_log("Fit success = {}".format(fit_success))
hm.logs.info_log("Objective = {}".format(objective))
# =======================================================================
# Fit model
# =======================================================================
hm.logs.info_log("Fit model for {} epochs...".format(epochs_num))
if flow_type == "RealNVP":
model = hm.model.RealNVPModel(
ndim, standardize=standardize, temperature=temperature
)
if flow_type == "RQSpline":
model = hm.model.RQSplineModel(
ndim,
n_layers=n_layers,
n_bins=n_bins,
hidden_size=hidden_size,
spline_range=spline_range,
standardize=standardize,
temperature=temperature,
)
model.fit(chains_train.samples, epochs=epochs_num, verbose=verbose)

# Use chains and model to compute inverse evidence.
hm.logs.info_log("Compute evidence...")
Expand Down Expand Up @@ -249,135 +281,37 @@ def run_example(
# Create corner/triangle plot.
# ======================================================================
if plot_corner and i_realisation == 0:
ex_utils.plot_corner(samples.reshape((-1, ndim)))
if savefigs:
plt.savefig(
"examples/plots/gaussian_nondiagcov_corner.png", bbox_inches="tight"
)

hm.utils.plot_getdist(samples.reshape((-1, ndim)))
if savefigs:
plt.savefig(
"examples/plots/gaussian_nondiagcov_getdist.png",
bbox_inches="tight",
)

plt.show()
save_name = save_name_start + "_gaussian_nondiagcov_getdist.png"
plt.savefig(save_name, bbox_inches="tight")

# ======================================================================
# In 2D case, plot surface/image and samples.
# ======================================================================
if plot_surface and ndim == 2 and i_realisation == 0:
# ==================================================================
# Define plot parameters.
# ==================================================================
nx = 50
xmin = -3.0
xmax = 3.0

# ==================================================================
# 2D surface plot of posterior.
# ==================================================================
ln_posterior_func = partial(ln_posterior, inv_cov=inv_cov)
ln_posterior_grid, x_grid, y_grid = ex_utils.eval_func_on_grid(
ln_posterior_func,
xmin=xmin,
xmax=xmax,
ymin=xmin,
ymax=xmax,
nx=nx,
ny=nx,
)
i_chain = 0
ax = ex_utils.plot_surface(
np.exp(ln_posterior_grid),
x_grid,
y_grid,
samples[i_chain, :, :].reshape((-1, ndim)),
np.exp(lnprob[i_chain, :].reshape((-1, 1))),
contour_z_offset=-0.5,
alpha=0.3,
)
num_samp = chains_train.samples.shape[0]
samps_compressed = model.sample(num_samp)

ax.set_zlabel(r"$\mathcal{L}$")
hm.utils.plot_getdist_compare(chains_train.samples, samps_compressed)
plt.ticklabel_format(style="sci", axis="y", scilimits=(0, 0))

# Save.
if savefigs:
plt.savefig(
"examples/plots/gaussian_nondiagcov_posterior_surface.png",
bbox_inches="tight",
save_name = (
save_name_start
+ "_gaussian_nondiagcov_corner_all_{}D.png".format(ndim)
)
plt.savefig(save_name, bbox_inches="tight", dpi=300)

plt.show(block=False)

# ==================================================================
# Image of posterior samples overlayed with contour plot.
# ==================================================================
# Plot posterior image.
ax = ex_utils.plot_image(
np.exp(ln_posterior_grid),
x_grid,
y_grid,
samples[i_chain].reshape((-1, ndim)),
colorbar_label="$\mathcal{L}$",
plot_contour=True,
markersize=1.0,
)
# Save.
hm.utils.plot_getdist(samps_compressed)
if savefigs:
plt.savefig(
"examples/plots/gaussian_nondiagcov_posterior_image.png",
bbox_inches="tight",
save_name = (
save_name_start
+ "_gaussian_nondiagcov_flow_getdist_{}D.png".format(ndim)
)

plt.show(block=False)

# ==================================================================
# Learnt model of the posterior
# ==================================================================
# Evaluate ln_posterior and model over grid.
x = np.linspace(xmin, xmax, nx)
y = np.linspace(xmin, xmax, nx)
x, y = np.meshgrid(x, y)
ln_model_grid = np.zeros((nx, nx))
for i in range(nx):
for j in range(nx):
ln_model_grid[i, j] = model.predict(np.array([x[i, j], y[i, j]]))

i_chain = 0
ax = ex_utils.plot_surface(
np.exp(ln_model_grid), x_grid, y_grid, contour_z_offset=-0.075
)
ax.set_zlabel(r"$\mathcal{L}$")

# Save.
if savefigs:
plt.savefig(
"examples/plots/gaussian_nondiagcov_surface.png",
save_name,
bbox_inches="tight",
dpi=300,
)

plt.show(block=False)

# ==================================================================
# Projection of posteior onto x1,x2 plane with contours.
# ==================================================================
# Plot posterior image.
ax = ex_utils.plot_image(
np.exp(ln_model_grid),
x_grid,
y_grid,
colorbar_label="$\mathcal{L}$",
plot_contour=True,
)
# Save.
if savefigs:
plt.savefig(
"examples/plots/gaussian_nondiagcov_image.png", bbox_inches="tight"
)

plt.show(block=False)
# ==================================================================
plt.show()

evidence_inv_summary[i_realisation, 0] = ev.evidence_inv
evidence_inv_summary[i_realisation, 1] = ev.evidence_inv_var
Expand All @@ -387,16 +321,20 @@ def run_example(
hm.logs.info_log("Execution_time = {}s".format(clock))

if n_realisations > 1:
np.savetxt(
"examples/data/gaussian_nondiagcov_evidence_inv" + "_realisations.dat",
evidence_inv_summary,
save_name = (
save_name_start
+ "_gaussian_nondiagcov_evidence_inv"
+ "_realisations_{}D.dat".format(ndim)
)
np.savetxt(save_name, evidence_inv_summary)
evidence_inv_analytic_summary = np.zeros(1)
evidence_inv_analytic_summary[0] = np.exp(-ln_evidence_analytic)
np.savetxt(
"examples/data/gaussian_nondiagcov_evidence_inv" + "_analytic.dat",
evidence_inv_analytic_summary,
save_name = (
save_name_start
+ "_gaussian_nondiagcov_evidence_inv"
+ "_analytic_{}D.dat".format(ndim)
)
np.savetxt(save_name, evidence_inv_analytic_summary)

created_plots = True
if created_plots:
Expand All @@ -408,11 +346,12 @@ def run_example(
hm.logs.setup_logging()

# Define parameters.
ndim = 2
ndim = 5
nchains = 100
samples_per_chain = 5000
nburn = 500
np.random.seed(10)
flow_str = "RealNVP"
# flow_str = "RQSpline"
np.random.seed(10) # used for initializing covariance matrix

hm.logs.info_log("Non-diagonal Covariance Gaussian example")

Expand All @@ -421,11 +360,8 @@ def run_example(
hm.logs.debug_log("Dimensionality = {}".format(ndim))
hm.logs.debug_log("Number of chains = {}".format(nchains))
hm.logs.debug_log("Samples per chain = {}".format(samples_per_chain))
hm.logs.debug_log("Burn in = {}".format(nburn))

hm.logs.debug_log("-------------------------")

# Run example.
run_example(
ndim, nchains, samples_per_chain, nburn, plot_corner=True, plot_surface=True
)
run_example(flow_str, ndim, nchains, samples_per_chain, plot_corner=False)
Loading

0 comments on commit 28662c7

Please sign in to comment.