plt.title(r'$\phi = \frac{\zeta_{a} g}{\omega} e^{k z} cos(\omega t + k x)$')

A very good and more detailed guide on subplots and figures can be found here.

import matplotlib.pyplot as plt
import numpy as np

# Some example data to display
x = np.linspace(0, 2 * np.pi, 400)
y = np.sin(x ** 2)

fig, ax = plt.subplots()
fig.suptitle('A single plot')
ax.plot(x, y)

fig, axs = plt.subplots(2)
fig.suptitle('Vertically stacked subplots')
axs[0].plot(x, y)
axs[1].plot(x, -y)

fig, (ax1, ax2) = plt.subplots(2)
fig.suptitle('Vertically stacked subplots')
ax1.plot(x, y)
ax2.plot(x, -y)

fig, (ax1, ax2) = plt.subplots(1, 2)
fig.suptitle('Horizontally stacked subplots')
ax1.plot(x, y)
ax2.plot(x, -y)

axes.flat is not a function, it's an atribute of the numpy.ndarray. ndarray.flat is a 1-D iterator over the array. This is a numpy.flatiter instance, which acts similarly to, but is not a subclass of Python’s built-in iterator object.

fig, axs = plt.subplots(2, 2)
fig.suptitle('Stacking subplots in two directions')
axs[0, 0].plot(x, y)
axs[0, 0].set_title('Axis [0,0]')
axs[0, 1].plot(x, y, 'tab:orange')
axs[0, 1].set_title('Axis [0,1]')
axs[1, 0].plot(x, -y, 'tab:green')
axs[1, 0].set_title('Axis [1,0]')
axs[1, 1].plot(x, -y, 'tab:red')
axs[1, 1].set_title('Axis [1,1]')


for ax in axs.flat:
    ax.set(xlabel='x-label', ylabel='y-label')

# Hide x labels and tick labels for top plots and y ticks for right plots.
# Try commenting out the next two lines to see what would happen if we did not hide the inner labels and ticks.
for ax in axs.flat:
    ax.label_outer()

To display all figures at once, only call plt.show() at the end of the last figure. As an extra practice, try moving the call or adding more, what happens, and why?

fig1, ax1 = plt.subplots()
fig1.suptitle('A single plot')
ax1.plot(x, y)

fig2, ax2 = plt.subplots()
fig2.suptitle('Another single plot')
ax2.plot(x, y)

plt.show()

plt.figure(1)                # the first figure (now current figure)
plt.subplot(211)             # the first subplot in the first figure (now current subplot)
plt.plot([1, 2, 3])
plt.subplot(212)             # the second subplot in the first figure (new current subplot)
plt.plot([4, 5, 6])

plt.figure(2)                # a second figure (new current figure)
plt.plot([4, 5, 6])          # creates a subplot(111) by default

plt.figure(1)                # figure 1 current; subplot(212) still current
plt.subplot(211)             # make subplot(211) in figure 1 current
plt.title('Easy as 1, 2, 3') # subplot 211 title

A more detailed explanation of the subplot method used above is found in the code explaining Quiver autoscaling vs manually set axes.

import numpy as np
import matplotlib.pyplot as plt

# X and Y define the arrow locations
# This setup gives us 10 arrows in width and height, as our interval is from -5 to 5 with step 1
X = np.arange(-5, 5, 1)
Y = np.arange(-5, 5, 1)

# U and V define the arrow directions, respectively in x- and y-direction
U, V = np.meshgrid(3*X, 3*Y)

plt.figure()

# Argument 121 in subplot() below denotes 1 row, 2 columns, first subplot. sublot(121) is then current.
plt.subplot(121)
plt.quiver(X, Y, U, V)
plt.title('Only autoscaling')

# Argument 122 denotes 1 row, 2 columns, second subplot. Notice that the number of rows and columns has to be equal every time, 
# where as the last number is the position where we want our subplot.
plt.subplot(122)
plt.quiver(X, Y, U, V)
# Here we specify the axes. How much extra space you need depends on the arrow size and direction,
# and must therefore be adapted each time
plt.axis([-6.5, 5.5, -6.5, 5.5])
plt.title('Manually set axes')

plt.show()

Further demos on contour plots and contour labels.

import matplotlib.pyplot as plt
import numpy as np

def f(x, y):
    return np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x)
 
x = np.linspace(0, 5, 60)
y = np.linspace(0, 5, 50)

X, Y = np.meshgrid(x, y)
Z = f(X, Y)

fig, axs = plt.subplots(1,3)
fig.suptitle('Three versions of the same contour plot')
axs[0].contour(X, Y, Z)
axs[1].contourf(X, Y, Z)
axs[2].contour(X, Y, Z, colors='black')

plt.show()

3D Plots as a tool of learning may be especially useful when learning about linear wave theory in the course TMR4247 - Marine Technology - Hydrodynamics.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()

# Setting the axes to 3D
ax = Axes3D(fig)        
X, Y = np.meshgrid(np.arange(0, 2*np.pi, 0.2), np.arange(0, 2*np.pi, 0.2))
Z = np.sin(X)

# Making the 3D plot as a surface plot, where X, Y and Z is data values as 2D arrays
# The arguments here are in the same format as used in Contour plots
# Also specifing a colormap to make our surface easier to interpet
ax.plot_surface(X, Y, Z, cmap='Blues')
ax.set_xlabel('X-label')
ax.set_ylabel('Y-label')
ax.set_zlabel('Z-label')
ax.set_title('3D Plot')

plt.show()

Keep in mind that this is not the best approach when making good, smooth animations, but rather a fast and easy setup.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


fig = plt.figure()
ax = Axes3D(fig)
X, Y = np.meshgrid(np.arange(0, 2*np.pi, 0.2), np.arange(0, 2*np.pi, 0.2))

# Setting a time interval
dt = 0.1                

# We have set our range from 0 to 10000 to make sure it runs as long as we want, and we'll then close it ourself
for i in range(0, 10000):
    # Our function for Z has to be dependent on some non-constant value
    # We have chosen to increase 'dt' each loop
    dt = dt + 0.1
    Z = np.sin(X + dt)

    # Our plotting has to be done inside the loop, as we want to redraw it mulitple times
    ax.plot_surface(X, Y, Z, cmap='Blues')

    # Since we are using the 'ax.clear()' at the end to clear all previous drawings, we need to specfiy all axes inside the loop
    # If we choose not to clear, we can move this outside the loop
    ax.set_xlabel('X-label')
    ax.set_ylabel('Y-label')
    ax.set_zlabel('Z-label')
    ax.set_title('3D Animation')

    # Here we pause the system for some time, in order for us to be able to watch as the plot changes
    plt.pause(1/60)

    # Lastly, we clear all previous drawings, as we usually only want to see the plot at the current time value
    ax.clear()
    
plt.show()

The animation below is slightly altered from the code above in order to make it a GIF.

"""
Matplotlib Animation Example

author: Jake Vanderplas
email: vanderplas@astro.washington.edu
website: http://jakevdp.github.com
license: BSD
Please feel free to use and modify this, but keep the above information. Thanks!
"""

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation

# First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
ax = plt.axes(xlim=(0, 2), ylim=(-2, 2))
line, = ax.plot([], [], lw=2)

# Initialization function: plot the background of each frame
def init():
    line.set_data([], [])
    return line,

# Animation function. This is called sequentially
def animate(i):
    x = np.linspace(0, 2, 1000)
    y = np.sin(2 * np.pi * (x - 0.01 * i))
    line.set_data(x, y)
    return line,

# Call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=200, interval=20, blit=True)

# The line below will save the animation as a GIF if uncommented.
# anim.save('basic_animation.gif', fps=30, writer='imagemagick')

plt.show()

More examples by the same author is found at this page, e.g. a double pendulum.

The exercises provided will have a varying degree of difficulty, thus including more advanced and detailed methods than described here.