from KDEpy.BaseKDE import BaseKDE
from mpl_toolkits.mplot3d import Axes3D

n = 64
p = np.linspace(-3, 3, num=n)
obs_x_dims = np.array(np.meshgrid(p, p)).T.reshape(-1, 2)

# fig = plt.figure() is already set, adjust the size
fig.set_figwidth(7); fig.set_figheight(5);

selected_kernels = ['box', 'tri', 'exponential', 'gaussian']
for i, kernel_name in enumerate(selected_kernels, 1):

  kernel = BaseKDE._available_kernels[kernel_name]
  ax = fig.add_subplot(2, 2, i, projection='3d')
  z = kernel(obs_x_dims, norm=2).reshape((n, n))
  surf = ax.plot_surface(*np.meshgrid(p, p), z)
  ax.set_title(f"'{kernel_name}', $2$-norm")