jakelevi1996/Plotting with matplotlib.md

Plotting with `matplotlib`

Instead of writing the same plotting code using plt.plot again and again in different modules/scripts, here is a simple function which plots data and formats, saves and closes the figure; this function can be used and modified instead of writing the same lines of code again and again in different repositories. Following this simple function are some scripts which demonstrate some of the fancier features of pyplot; the resulting images are shown below the code.

Simple plotting function

import numpy as np
import matplotlib.pyplot as plt

def plot_data(x, y, name="plot"):
    plt.figure(figsize=[8, 6])
    plt.plot(x, y, "b")
    plt.grid(True)
    plt.title(name)
    plt.tight_layout()
    plt.savefig(name)
    plt.close()

x = np.linspace(-10, 10, 200)
y = np.sin(x)
plot_data(x, y, "Simple plotting function")

Sinusoids 1: variable number of lines and custom colours

See this matplotlib tutorial for a description of the different available colour maps.

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import numpy as np

# Generate data
x = np.linspace(-10, 10, 200)
num_phases = 10
phases = np.linspace(0, 2*np.pi, num_phases, endpoint=False)
y = np.sin(x.reshape(1, -1) + phases.reshape(-1, 1))
# Generate colours
colours = plt.get_cmap("hsv")(np.linspace(0, 1, num_phases, endpoint=False))
# Plot data
plt.figure(figsize=[8, 6])
handles, labels = [], []
for i in range(num_phases):
    lines = plt.plot(x, y[i], c=colours[i])

# Format, save and close
plt.title("Sinusoids 2")
handles=[Line2D(
    [], [], c=c, label="Phase = {:.2f}".format(p)
) for c, p in zip(colours, phases)]
plt.legend(handles=handles, ncol=3)
plt.xlabel("x")
plt.ylabel("y")
plt.grid(True)
plt.tight_layout()
plt.axis("equal")
plt.savefig("Sinusoids 1")
plt.close()

Sinusoids 2

This demo includes various features of matplotlib, including:

Subplots using plt.subplots, including shared axes, and non-equal sizes of subplots
Colour plots using pcolormesh, including setting chosen values to nan and chosing a colour for nan values
A colour-bar which spans multiple axes
Subplot titles
Turning on minor ticks on the axes using minorticks_on
Turning on the grid, with different line-styles for major and minor ticks
Adding horizontal and vertical lines to a subplot with axhline and axvline
Adding arrows, including controlling thickness and size of arrow-head
Setting the relative order/depth of plot components using the zorder keyword argument
Setting marker size and marker shape
Setting the background colour to a plot using set_facecolor
Turning off axis ticks and labels using tick_params
Adding custom legend entries using the Patch and Line2D classes
Setting the location of the legend
Removing the axes from a subplot altogether using axis("off") (this is especially useful for displaying images)
Setting a title over multiple subplots using fig.suptitle

NB an official matplotlib demo demonstrates an alternative snippet using plt.subplots. An alternative to plt.subplots is plt.subplot.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.patches import Patch

# Create subplots
fig, axes = plt.subplots(2, 2, sharex=True, figsize=[8, 6], gridspec_kw={
    "height_ratios": [3, 1], "width_ratios": [4, 1], "wspace": 0.3})


# UL: Generate PDF
x = np.linspace(0, 10, 100).reshape(1, -1)
y = np.linspace(-5, 5, 100).reshape(-1, 1)
z = np.sin(x)
p = np.exp(-0.5 * np.square(y - z)) / np.sqrt(2 * np.pi)

# Set some random points to nan, and set nan colour
np.random.seed(0)
nan_inds = np.random.choice(p.size, p.size//100, replace=False)
p.flat[nan_inds] = np.nan
plt.get_cmap().set_bad("r")

# Make colour plot and add colour bar and title
c = axes[0, 0].pcolormesh(x, y, p, vmin=0, vmax=1)
plt.colorbar(c, ax=axes[:, 0])
axes[0, 0].set_title("Colour plot")


# LL: Make 1D plot
axes[1, 0].plot(x.flat, z.flat, "b")
axes[1, 0].set_xlim(0, 10)
axes[1, 0].set_title("Line plot")
axes[1, 0].minorticks_on()
axes[1, 0].grid(which="major", ls="-")
axes[1, 0].grid(which="minor", ls=":", alpha=0.5)

# Add horizontal and vertical lines
axes[1, 0].axhline(0.5, c="r", ls="--")
axes[1, 0].axvline(4, c="r", ls="--")


# UR: Add arrows and markers
axes[0, 1].arrow(5, 0, 0, 1, width=1, head_length=0.3, color="g", zorder=10)
axes[0, 1].plot([5, 5], [0, 1.3], alpha=0)
axes[0, 1].plot([5], [0.25], "kd", ms=30, zorder=20)
axes[0, 1].plot([5], [0.75], "ro", ms=30, zorder=0)
axes[0, 1].set_facecolor("b")
axes[0, 1].tick_params(which="both", bottom=False, labelbottom=False,
    left=False, labelleft=False)
axes[0, 1].set_title("Arrow + markers")


# LR: create legend
axes[1, 1].legend(ncol=2, loc="center", handles=[
    Patch(color="r", label="Red"),
    Patch(color="b", label="Blue"),
    Line2D([], [], c="k", marker="d", ms=10, ls="", label="Diamond"),
    Line2D([], [], c="r", marker="o", ms=10, ls="", label="Circle"),
])
axes[1, 1].axis("off")
axes[1, 1].set_title("Legend")


# Title, save and close
fig.suptitle("Fancy plot", fontsize=20)
# fig.tight_layout(rect=[0, 0.05, 1, 0.95])
fig.savefig("Sinusoids 2")
plt.close(fig)

A note about plotting in Matlab

In Matlab, it is possible to create a figure, plot data in the figure, and save the figure to disk as an image, without displaying the figure on screen. To create a figure with user-defined position and size without displaying to the screen:

figure('Position', [1, 41, 1920, 963], 'visible', 'off');

The 'Position' argument has the form [left bottom width height]. To get the position size of the current figure (which may or may not be visible), use the following command:

get(gcf, 'Position')

After plotting data in the figure, it can be saved to disk using the following syntax:

saveas(gcf, fullfile(dir_name, plot_name), 'png');

If the directory dir_name doesn't already exist, it can be created with the following command:

if ~exist(dir_name, 'dir')
    mkdir(dir_name)
end