Plot using Python
In this post we record some Python scripts for generating 1D or 2D plots. The focus is on the commands that generate plots of different types and control adjustable parameters (e.g. font sizes, ticks, etc).
First, we import some useful libraries and also set fonts to display latex nicely.
#!/usr/bin/env python
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cm
import numpy as np
import math
plt.rc('text', usetex=True)
plt.rc('font', family='serif')
1. An example of 1D plot
fig = plt.figure(figsize=(13,5))
ax = fig.add_subplot(1, 2, 1)
# prepare data
pi = math.pi
x = np.linspace(0, 2 * pi)
y = np.sin(x)
# plot
ax.plot(x, y)
# set various parameters below in order to make the plot nice
ax.set_aspect('equal')
ax.set_xlabel(r'$x$',fontsize=20)
ax.set_ylabel(r'$y$',fontsize=20, rotation=0)
ax.tick_params(axis='both', labelsize=20)
ax.set_xticks([0, 2, 4, 6])
ax.set_yticks([-1,-0.5, 0, 0.5, 1])
ax.set_xlim([0, 2*pi])
ax.set_ylim([-1,1])
ax.set_title(r'$\sin(x)$',fontsize=25)
# save the plot to file
fig_name = f"1d.png"
fig.savefig(fig_name, bbox_inches='tight', dpi=200)
Here is the output:
2. Examples of 2D plots
Before plotting, let us first define some functions used to generate data.
# a potential function
def V(X):
return (X[0]**2 - 1)**2 + 1.0 / 0.5 * (X[0]**2 + X[1] - 1)**2
# gradient of potential function
def gradV(X):
return np.array(( 4.0 * X[0] * (X[0]**2 - 1.0 + 1.0 / 0.5* (X[0]**2 + X[1] - 1)), 2.0 / 0.5 * (X[0]**2 + X[1] - 1)) )
# function to generate trajectory data using Euler-Maruyama scheme
def sample():
r = np.random.RandomState(100)
X = [-1, 0]
dim = 2
traj = [X]
beta = 1
delta_t = 0.001
save = 10
N = 100000
for i in range(N):
b = r.normal(size=(dim,))
X = X - gradV(X) * delta_t + np.sqrt(2 * delta_t/beta) * b
if i % save==0:
traj.append(X)
return np.array(traj)
We also set some common parameters for different plots below.
nx = 100
ny = 100
x_domain = [-2.5, 2.5]
y_domain = [-2.5, 2.5]
contour_levels = [0.0, 1.0, 1.5, 2.0, 3.0, 4.0]
v_min_max = [0,4]
xticks = [-2.0, -1.0, 0, 1.0, 2.0]
yticks = [-2.0, -1.0, 0, 1.0, 2.0]
dx = (x_domain[1] - x_domain[0]) / nx
dy = (y_domain[1] - y_domain[0]) / ny
gridx = np.linspace(x_domain[0], x_domain[1], nx)
gridy = np.linspace(y_domain[0], y_domain[1], ny)
x_plot = np.outer(gridx, np.ones(ny))
y_plot = np.outer(gridy, np.ones(nx)).T
x2d = np.concatenate((x_plot.reshape(nx * ny, 1), y_plot.reshape(nx * ny, 1)), axis=1)
# generate 2D trajectory
trajectory = sample()
# get the potential on the 2D grid
V_on_grid = np.array([V(x) for x in x2d]).reshape(nx, ny)
# compute values of potential along the trajectory
v_traj = np.array([V(x) for x in trajectory])
2.1. 2D Profile with contours
In the first example, we plot the profile of the function $V: \mathbb{R}^2\rightarrow \mathbb{R}$ together with its contour lines.
fig = plt.figure(figsize=(7,5))
ax = fig.add_subplot(1, 1, 1)
# plot profile by pcolormesh
im = ax.pcolormesh(x_plot, y_plot, V_on_grid, cmap='coolwarm',shading='auto', vmin=v_min_max[0], vmax=v_min_max[1])
# show contour lines
contours = ax.contour(x_plot, y_plot, V_on_grid, contour_levels)
ax.clabel(contours, inline=True, fontsize=13,colors='black')
ax.set_aspect('equal')
ax.set_xlabel(r'$x_1$',fontsize=20)
ax.set_ylabel(r'$x_2$',fontsize=20, rotation=0)
ax.tick_params(axis='both', labelsize=20)
ax.set_xticks(xticks)
ax.set_yticks(yticks)
ax.set_xlim([x_domain[0], x_domain[1]])
ax.set_ylim([y_domain[0], y_domain[1]])
ax.set_title('V',fontsize=25)
# show colorbar
cbar = fig.colorbar(im, ax=ax, shrink=1.0)
cbar.ax.tick_params(labelsize=15)
fig_name = f"V2d.png"
fig.savefig(fig_name, bbox_inches='tight', dpi=200)
Here is the output:
2.2. 2D histogram plot
In the second example, we plot histogram of the 2d trajectory data.
fig = plt.figure(figsize=(7,5))
ax = fig.add_subplot(1, 1, 1)
# compute histogram statistics
h, xedges, yedges = np.histogram2d(trajectory[:,0], trajectory[:,1], bins=[nx, ny], range=[[x_domain[0],x_domain[1]],[y_domain[0],y_domain[1]]], density=True)
# get the meshgrid
X, Y = np.meshgrid(xedges, yedges)
# plot the histogram, specify the colormap and log scale
im = ax.pcolormesh(X, Y, h.T, \
cmap='coolwarm', shading='auto', \
norm=colors.LogNorm(0.1, 1) )
# show colorbar
cbar = fig.colorbar(im, ax=ax, shrink=1.0)
cbar.ax.tick_params(labelsize=15)
ax.set_aspect('equal')
ax.set_xlim([x_domain[0], x_domain[1]])
ax.set_ylim([y_domain[0], y_domain[1]])
ax.set_xlabel(r'$x_1$',fontsize=20)
ax.set_ylabel(r'$x_2$',fontsize=20, rotation=0)
ax.tick_params(axis='both', labelsize=20)
ax.set_title('trajectory data',fontsize=25)
ax.set_xticks(xticks)
ax.set_yticks(yticks)
fig_name = f"traj_hist.png"
fig.savefig(fig_name, bbox_inches='tight')
plt.close()
Here is the output:
2.3. 2D scatter plot
In the third example, we plot the function $V$ on scattered trajectory data.
fig=plt.figure(figsize=(7,5))
ax = plt.gca()
# draw trajectory data according to the values of V.
ret = ax.scatter(trajectory[:,0], trajectory[:,1], c=v_traj)
cbar = fig.colorbar(ret, ax=ax, shrink=1.0)
cbar.ax.tick_params(labelsize=15)
ax.set_xlabel(r'$x_1$',fontsize=20)
ax.set_ylabel(r'$x_2$',fontsize=20, rotation=0)
ax.tick_params(axis='both', labelsize=20)
ax.set_title('trajectory data',fontsize=25)
ax.set_xticks(xticks)
ax.set_yticks(yticks)
fig_name = f"v_scatter.png"
fig.savefig(fig_name, bbox_inches='tight', dpi=200)
plt.close()
Here is the output:
