plot_cartoon_objective.py

import numpy as np

import matplotlib.pyplot as plt
from hips.plotting.layout import create_figure, create_axis_at_location
from hips.plotting.colormaps import gradient_cmap

import seaborn as sns
color_names = ["windows blue",
               "amber",
               "crimson",
               "faded green",
               "dusty purple",
               "greyish"]
colors = sns.xkcd_palette(color_names)
sns.set(style="ticks", palette=sns.xkcd_palette(color_names))

# def make_joint_objective_figure():
if __name__ == "__main__":

    # Lay down a grid of X,Y points
    N = 101
    x = np.linspace(0,2,num=N)
    y = np.linspace(0,2,num=N)
    X,Y = np.meshgrid(x, y)
    xy = np.column_stack((X.ravel(), Y.ravel()))

    # Suppose the true objective looks like a warped Gaussian density
    mu = [1., 1.]
    Sigma = np.array([[1., 0.5],
                      [0.5, 1.]])
    L_true = np.exp(-0.5 * ((xy - mu) * (np.linalg.solve(Sigma, (xy - mu).T)).T).sum(1))
    L_true *= np.exp(-0.5 * xy[:, 0] ** 3) * np.exp(-0.5 * xy[:, 1] ** 3)
    L_true = L_true.reshape((N,N))

    # Set the SVAE objective as a rough Gaussian approximation
    mu_svae = np.array([0.6, 0.6])
    L_svae = np.exp(-0.5 * ((xy - mu_svae) * (np.linalg.solve(.8*Sigma, (xy - mu_svae).T)).T).sum(1))
    L_svae = L_svae.reshape((N,N))

    # Plot the objective function over \eta_theta (x axis) and \eta_x (y axis)
    fig = create_figure(figsize=(2.6,2.6))
    ax = create_axis_at_location(fig, .7,.7, 1.8, 1.8)
    ax.contour(X, Y, L_svae, 8,
               cmap=gradient_cmap((np.ones(3), colors[1])))

    ax.contour(X, Y, L_true, 8,
               cmap=gradient_cmap((np.ones(3), colors[0])),
               )

    # Now find the optimal y for fixed x
    xi_star = 50
    x_star = X[0, xi_star]
    y_true = Y[np.argmax(L_true[:, xi_star]), xi_star]
    y_svae = Y[np.argmax(L_svae[:, xi_star]), xi_star]

    # Lines
    # plt.plot([x,x], [0,y_true], '-', color=colors[0], lw=2)
    # plt.plot([0,x], [y_true,y_true], '-', color=colors[0], lw=2)
    plt.plot([0, x_star], [y_true, y_true], ':', color='k', lw=1)

    # plt.plot([x,x], [0,y_svae], '-', color=colors[1], lw=2)
    # plt.plot([0,x], [y_svae,y_svae], '-', color=colors[1], lw=2)
    plt.plot([0, x_star], [y_svae, y_svae], ':', color='k', lw=1)

    plt.plot([x_star, x_star], [0, max(y_true, y_svae)], ':', color='k', lw=1)

    # Dots
    plt.plot(x_star, y_true, 'o', markersize=8, color=colors[0])
    plt.plot(x_star, y_svae, 'o', markersize=8, color=colors[1])

    # Labels
    plt.xlabel("$\\widetilde{\\eta}_{\\theta}$", labelpad=10)
    plt.ylabel("$\\widetilde{\\eta}_{x}$", labelpad=1)

    # Legend
    proxy = [plt.plot([-1,-1], [-1,-1], '-', color=colors[0], lw=2, label="$\\mathcal{L}$"),
             plt.plot([-1,-1], [-1,-1], '-', color=colors[1], lw=2, label="$\\widehat{\\mathcal{L}}$")]
    legend = plt.legend(frameon=True)

    plt.xlim(0,2)
    plt.ylim(0,2)
    plt.xticks([0,1,2])
    plt.yticks([0,1,2])
    sns.despine(ax=ax)

    plt.savefig("cartoon_2d.png")
    plt.savefig("cartoon_2d.pdf")

    # Plot the partially optimized objective function
    fig = create_figure(figsize=(2.6,2.))
    ax = create_axis_at_location(fig, .7,.7, 1.8, 1.2)
    Lp_true = np.max(L_true, axis=0)
    Lp_svae = np.max(L_svae, axis=0)
    plt.plot(x, Lp_true, color=colors[0])
    plt.plot(x, Lp_svae, color=colors[1])

    # Labels
    plt.xlabel("$\\widetilde{\\eta}_{\\theta}$", labelpad=10)
    plt.ylabel("$\\mathcal{L}$", labelpad=1)

    plt.ylim(0,1.01)
    plt.yticks([0,0.5,1])
    # plt.legend(frameon=True)

    sns.despine(ax=ax)
    plt.savefig("cartoon_1d.png")
    plt.savefig("cartoon_1d.pdf")

    plt.show()