simple_example_prior_posterior.py

import numpy as np
from matplotlib import pyplot as plt

plt.rcParams["text.usetex"] = True


def f_(x, y):
    return x**3 + y**2


def manifold_density_reweight_desired_(x, y):
    return ((3 * x**2) ** 2 + (y) ** 2) ** (-1 / 2)


x_max = 2
x_min = -2
y_max = np.sqrt(1 - (-2) ** 3)
y_min = -np.sqrt(1 - (-2) ** 3)

general_max = np.maximum(y_max, x_max)
general_min = np.minimum(y_min, x_min)

n = 500
x = np.linspace(general_min, general_max, n)
y = np.linspace(general_min, general_max, n)
X, Y = np.meshgrid(x, y)
Z = manifold_density_reweight_desired_(X, Y)
Z = (Z - Z.min()) / (Z.max() - Z.min())
Z = Z * (2 - Z)

mask = f_(X, Y)

for Zm, name in [(Z, "prior"), (np.ma.masked_where((1.3 < mask) | (mask < 0.7), Z), "posterior_trivial")]:
    plt.pcolormesh(X, Y, Zm, shading="gouraud", color="blue", cmap="Blues")
    plt.xlim(general_min, general_max)
    plt.ylim(general_min, general_max)
    plt.xticks(fontsize=18)
    plt.yticks(fontsize=18)
    plt.xlabel(r"$\theta_1$", loc="center", fontsize=18)
    plt.ylabel(r"$\theta_2$", loc="center", fontsize=18)
    plt.gca().set_aspect("equal", adjustable="box")
    plt.tight_layout()
    plt.savefig(f"{name}.png")
    plt.show()