# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Paper Diagrams
#
# Code for making Figures 1, 2, 3, 4 and 5 in the paper.

# ### Set up

# +
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import bsr.basis_functions as bf
# %matplotlib inline
# %autoreload 2

# Set plot size, font and fontsize to match LaTeX template
# --------------------------------------------------------
# NB A4 paper is 8.27 × 11.69 inches (=210 × 297 mm)
# Font: \T1/ntxtlf/m/n/9
# Caption font: \T1/ntxtlf/m/n/8
# Footnote font: \T1/ntxtlf/m/n/8
# Abstract font: \T1/ntxtlf/m/n/10
fontsize = 8  # matplotlib default is 10
textwidth = 6.97522 * 0.99  # make 1% smaller to ensure everything fits
textheight = 9.43869 * 0.99  # make 1% smaller to ensure everything fits
colwidth = 3.32153
matplotlib.rc('text', usetex=True)
matplotlib.rc('font', family='serif', serif='Times New Roman', size=fontsize)
matplotlib.rc('axes', titlesize='medium')
# To print params, use "matplotlib.rcParams" or for a specific key "matplotlib.rcParams['font']"
# -

# ### Free-form decomposition diagrams (Figures 1 and 2)

# Settings
# --------
sigma = 0.15
amps_gau = [0.1, 0.5, 0.1, 0.35, 0.5, 0.4, 0.25]
figsize = (colwidth, 1.5)
ymax = 1.25
# Plotting
# --------
npoints = len(amps_gau)
pixels = np.linspace(0, 1, npoints)
pixel_width = pixels[1] - pixels[0]
x = np.linspace(0 - (1 / npoints), 1 + (1 / npoints), npoints * 20)
y = np.zeros(x.shape)
# make gaussian figure
fig_gau, ax_gau = plt.subplots(figsize=figsize)
fig_bar, ax_bar = plt.subplots(figsize=figsize)
for i, pixel in enumerate(pixels):
    comp = amps_gau[i] * np.exp(-((x - pixel) ** 2) / (2 * (sigma ** 2)))
    ax_gau.plot(x, comp, linestyle='dashed', color='black', linewidth=0.6)
    y += comp
    ax_gau.axvline(x=pixel, ymin=0, ymax=amps_gau[i] / ymax, color='black')
ax_gau.plot(x, y, label='$y=f(x)$')
ax_gau.legend()
# make bar figure, calculating amplitudes from y
amps_bar = np.interp(pixels, x, y)
# make adjustments so they are less smooth
amps_bar[2] *= 0.4
amps_bar[5] *=0.2
bar_x = [x.min(), pixels[0] - 0.5 * pixel_width]
bar_y = [0, 0]
for i, pixel in enumerate(pixels):
    bar_x += [pixel - 0.5 * pixel_width, pixel + 0.5 * pixel_width]
    bar_y += [amps_bar[i]] * 2
    ax_bar.axvline(x=pixel, ymin=0, ymax=amps_bar[i] / ymax, color='black')
ax_bar.bar(pixels, amps_bar, pixel_width, color='white',
           linewidth=0.6, edgecolor='black', linestyle='dashed')
bar_x += [pixels[-1] + 0.5 * pixel_width, x.max()]
bar_y += [0, 0]
ax_bar.plot(bar_x, bar_y, label='$y=f(x)$')
ax_bar.legend()
# add labels
for i, ax in enumerate([ax_gau, ax_bar]):
    ax.set_ylim([0, ymax])
    ax.set_xlim([x.min(), x.max()])
    ax.set_xticks(pixels)
    ax.set_yticks([])
    ax.set_xlabel('$x$')
    ax.set_ylabel('$y$')
    ind_j = npoints // 2
    labels = [''] * pixels.shape[0]
    labels[0] = '$x_1$'
    labels[-1] = '$x_M$'
    labels[ind_j] = '$x_j$'
    ax.set_xticklabels(labels)
    if i == 0:
        amps_plot = amps_gau
    else:
        amps_plot = amps_bar
    ax.text(pixels[0] - 0.02, amps_plot[0] + 0.04, '$a_1$')
    ax.text(pixels[-1] - 0.02, amps_plot[-1] + 0.04, '$a_M$')
    ax.text(pixels[ind_j] - 0.02, amps_plot[ind_j] + 0.04, '$a_j$')
adjust = {'left': 0.05, 'top': 0.99, 'right': 0.99, 'bottom': 0.24}
fig_bar.subplots_adjust(**adjust)
fig_gau.subplots_adjust(**adjust)
fig_bar.savefig('plots/freeform_bar.pdf')
fig_gau.savefig('plots/freeform_gau.pdf')

# ### 3d Lp norm plots (Figure 3)

# +
from mpl_toolkits.mplot3d import Axes3D  # needed to use projection='3d' in add_subplot
# Settings
# --------
figsize = (colwidth, 0.75)
npoints = 100  # paper uses 200 - reduce this to go faster
p_list = [0, 0.5, 1, 2, 4]
color = 'lightblue'

# Do the plot
# -----------
# Plot via spherical coordinates - based on
# https://stackoverflow.com/questions/7819498/plotting-ellipsoid-with-matplotlib
# Set of all spherical angles:
u = np.linspace(0, 2 * np.pi, npoints)
v = np.linspace(0, np.pi, npoints)
# set up array of xyz coords
x = np.outer(np.cos(u), np.sin(v))
y = np.outer(np.sin(u), np.sin(v))
z = np.outer(np.ones_like(u), np.cos(v))
xyz = np.stack([x, y, z])
fig = plt.figure(figsize=figsize)
for i, p in enumerate(p_list):
    ax = fig.add_subplot(1, len(p_list), i + 1, projection='3d', alpha=1)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title('$p={}$'.format(p).replace('0.5', '\\frac{1}{2}'))
    if p == 0:
        lw = 0.7
        ax.plot3D((0, 0), (0, 0), (-1, 1), color=color, lw=lw)
        ax.plot3D((0, 0), (-1, 1), (0, 0), color=color, lw=lw)
        ax.plot3D((-1, 1), (0, 0), (0, 0), color=color, lw=lw)
    else:
        pnorm = np.sum(np.abs(xyz) ** p, axis=0) ** (1 / p)
        xyz_plot = xyz / pnorm
        ax.plot_surface(xyz_plot[0], xyz_plot[1], xyz_plot[2],  rstride=1, cstride=1,
                    linewidth=0, antialiased=False, color=color, rasterized=True)
fig.subplots_adjust(left=0, right=1, bottom=-0.3, wspace=0)
fig.savefig('plots/pnorm.pdf', dpi=400)
# -

# ### Sparsity promotion plot (Figure 4)

# +
# Settings
figsize=(colwidth, 0.5*colwidth)
fig, axes = plt.subplots(ncols=2, figsize=figsize)
arrow_kwargs = {'head_width': 0.05, 'head_length': 0.05, 'fc': 'black', 'ec': None}
marker_kwargs = {'marker': 'X', 'color': 'red'}
fill_kwargs = {'color': 'lightblue', 'zorder': 0}
e_centre = (0.2, 0.7)
e_kwargs={'fc': None, 'ec': 'black', 'fill': False, 'angle': 60}
e_widths = [0.1, 0.2, 0.3, 0.4]
e_ratio = 3

for i, ax in enumerate(axes):
    ax.set_aspect('equal')
    ax.set_xlim([-1.1, 1.1])
    ax.set_ylim([-1.1, 1.1])
    # ax.xaxis.set_visible(False)
    ax.axis('off')
    ax.text(0, -1.15, '$a_1$', ha='center', ma='center')
    ax.text(-1.15, 0, '$a_2$', ha='center', ma='center')
    ax.set_title('${{||\mathbf{{a}}||}}_{}$'.format(2-i))
    if i == 0:
        cross_y = 0.475
        rad = np.sqrt(cross_y ** 2 + e_centre[1] ** 2)
        ax.add_artist(plt.Circle((0.0, 0.0), 0.5, **fill_kwargs))
        ax.scatter(e_centre[0], cross_y, **marker_kwargs)
    elif i == 1:
        cross_y = 0.57
        ax.scatter(0, cross_y, **marker_kwargs)
        ax.add_artist(plt.Polygon([[0, cross_y], [cross_y, 0], [0, -cross_y], [-cross_y, 0]],
                                  **fill_kwargs))
    al = 0.9
    for arrow_end in [(0, al), (0, -al), (al, 0), (-al, 0)]:
        ax.arrow(0, 0, *arrow_end, **arrow_kwargs)
    for width in e_widths:
        ax.add_artist(Ellipse(e_centre, width, width * e_ratio, **e_kwargs))
plt.subplots_adjust(left=0.05, right=1, bottom=0.05, top=0.85)
fig.savefig('plots/pnorm_sparsity.pdf')
# -

# ### 1d generalised Gaussian and tanh basis function diagrams (Figures 5a and 5b)

figsize = (colwidth * 0.6, 2)
xmax = 3
x = np.linspace(-xmax, xmax, 100)
adjust = {'top': 0.68, 'bottom': 0.17, 'left': 0.23, 'right': 0.85}
anchor = (0.5, 1.68)
def ylims_given_ymin_ymax(ymin, ymax):
    delta = ymax - ymin
    return [ymin - 0.1 * delta, ymax + 0.1 * delta]
# gg
# --
fig_gg, ax = plt.subplots(figsize=figsize)
for beta in [0.5, 1, 2, 4, 8]:
    ax.plot(x, bf.gg_1d(x, 1, 0, 1, beta), label=r'$\beta={}$'.format(beta).replace('0.5', r'\frac{1}{2}'))
ax.set_xlim([-xmax, xmax])
ax.set_ylim(ylims_given_ymin_ymax(0, 1))
fig_gg.subplots_adjust(**adjust)
ax.legend(loc='upper center', ncol=2, bbox_to_anchor=anchor)
ax.set_xlabel('$x$')
ax.set_ylabel(r'$y = \mathrm{e}^{- {(-|x-\mu|/\sigma)}^\beta }$')
fig_gg.savefig('plots/gg_demo.pdf')
# tanh
# ----
fig_ta, ax = plt.subplots(figsize=figsize)
for w1 in [-8, -2, -0.5, 0.5, 2, 8]:
    ax.plot(x, bf.ta_1d(x, 1, 0, w1), label=r'$w={}$'.format(w1).replace('0.5', r'\frac{1}{2}'))
ax.set_xlim([-xmax, xmax])
ax.set_ylim(ylims_given_ymin_ymax(-1, 1))
fig_ta.subplots_adjust(**adjust)
ax.legend(loc='upper center', ncol=2, bbox_to_anchor=anchor)
ax.set_xlabel('$x$')
ax.set_ylabel(r'$y = \mathrm{tanh}(w x + b)$')
fig_ta.savefig('plots/ta_demo.pdf')

StarcoderdataJupyterScriptsDedupFiltered