Here we demonstrate various options for plotting in Jupyter notebooks, with a special focus on monitoring time-dependent processes in simulations.
In [1]:
import mmf_setup;mmf_setup.nbinit(theme='mmf)')
Test Problem¶
Here we will mock up a simulation. The visualization consists of a main panel showing the evolution of a wavefunction (possibly a couple of components) which needs to be large to display details. In addition, we would like to have various
Modular Plotting¶
Often one has several components one would like to (optionally) plot. These often need to be laid out in a stack of plots. Here we discuss how to do this.
Matplotlib¶
In [51]:
%pylab inline --no-import-all
ts = np.linspace(0, 10, 100)
a = np.exp(-(ts-2)**2/2.0/1.0)
b = np.sin(ts)
c = ts**2/50 - 1.0
Here is a manual construction of the desired plot:
In [255]:
plt.subplot(gs[0,0])
plt.plot(ts, ts)
# Now we construct our desired layout in the right half
gs = GridSpecFromSubplotSpec(
3, 1,
hspace=0.0,
subplot_spec=gs[0, 1])
# Plot a
ax = plt.subplot(gs[0,0])
ax.plot(ts, a)
ax.set_ylabel('a')
ax.xaxis.set_ticklabels([])
ax.tick_params(axis='x', direction='in')
ax.yaxis.tick_right()
ax.yaxis.set_label_position("right")
# Plot b
ax = plt.subplot(gs[1,0])
ax.plot(ts, b)
ax.set_ylabel('b')
ax.xaxis.set_ticklabels([])
ax.tick_params(axis='x', direction='in')
ax.yaxis.tick_right()
ax.yaxis.set_label_position("right")
# Plot c
ax = plt.subplot(gs[2,0])
ax.plot(ts, c)
ax.set_ylabel('c')
ax.yaxis.tick_right()
ax.yaxis.set_label_position("right")
ax.set_xlabel('t')
# Adjust layout if desired
plt.tight_layout()
Much of the repeated axis manipulation can be done by a function. Here is an idea using a generator to allow the plots to be added:
In [256]:
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec
def vgrid(N, subplot_spec=None, right=False, **kw):
if subplot_spec is None:
subplot_spec = plt.subplot(111)
gs = GridSpecFromSubplotSpec(
N, 1, subplot_spec=subplot_spec, **kw)
for n in range(N):
ax = plt.subplot(gs[n, 0])
if right:
ax.yaxis.tick_right()
ax.yaxis.set_label_position("right")
if n < N - 1:
ax.xaxis.set_ticklabels([])
ax.tick_params(axis='x', direction='in')
yield ax
In [257]:
# Here is the overall grid layout
gs = GridSpec(1, 2)
# In the left half we have some plot:
plt.subplot(gs[0,0])
plt.plot(ts, ts)
vg = vgrid(N=3, hspace=0.0, right=True, subplot_spec=gs[0, 1])
# Now we construct our desired layout in the right half
ax = vg.next()
ax.plot(ts, a)
ax.set_ylabel('a')
# Plot b
ax = vg.next()
ax.plot(ts, b)
ax.set_ylabel('b')
# Plot c
ax = vg.next()
ax.plot(ts, c)
ax.set_ylabel('c')
ax.set_xlabel('t')
# Adjust layout if desired
plt.tight_layout()
Here is another idea. We start by making the plots, then inserting them into a layout. This demonstrates the feasibility of adjusting the positions after the plots are generated. Note: we need to do a bit of gymnastics to get unique plot locations initially.
In [258]:
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(ts, a)
ax.set_ylabel('a')
axs = [ax]
ax = fig.add_subplot(2,1,2)
ax.plot(ts, b)
ax.set_ylabel('b')
axs.append(ax)
ax = fig.add_subplot(3,1,3)
ax.plot(ts, c)
ax.set_ylabel('c')
axs.append(ax)
# Now we know how many plots, we can do the final layout
gs = GridSpec(1, 2)
# In the left half we have some plot:
ax = fig.add_subplot(gs[0,0])
ax.plot(ts, ts)
gs = GridSpecFromSubplotSpec(3, 1, hspace=0.0, subplot_spec=gs[0,1])
#gs = GridSpec(3, 1, hspace=0.0)
ax1, ax2, ax3 = axs
ax2.set_position(gs[1,0].get_position(fig))
ax3.set_position(gs[2,0].get_position(fig))
for n, ax in enumerate(axs):
subplotspec = gs[n,0]
ax.set_position(subplotspec.get_position(fig))
ax.set_subplotspec(subplotspec) # necessary if using tight_layout()
ax.yaxis.tick_right()
ax.yaxis.set_label_position("right")
if n < len(axs) - 1:
ax.xaxis.set_ticklabels([])
ax.tick_params(axis='x', direction='in')
plt.tight_layout()
Now we can encapsulate this in a class. We use the same API as above with the generator, but use some heuristics to determine how to form the final layout.
In [261]:
import attr
class Grid(object):
"""Object for stacking plots.
"""
def __init__(self, fig=None,
subplot_spec=None, right=False, **kw):
if fig is None:
fig = plt.gcf()
if subplot_spec is None:
subplot_spec = GridSpec(1, 1)[0, 0]
self.fig = fig
self.subplot_spec = subplot_spec
self.shape = [0, 0]
self.kw = kw
self.right = right
self.axes = []
def next(self, down=True, share=False):
"""Return the next axis."""
if self.shape == [0, 0]:
# Initialize
self.shape = [1, 1]
ax = self.fig.add_subplot(self.subplot_spec)
self.axes.append(ax)
return ax
if self.shape != [1, 1]:
# Set "down" to be consistent with current shape
down = self.shape[0] > self.shape[1]
if down:
self.shape[0] += 1
assert self.shape[1] == 1
else:
assert self.shape[0] == 1
self.shape[1] += 1
Nx, Ny = self.shape
gs = GridSpecFromSubplotSpec(
Nx, Ny,
subplot_spec=self.subplot_spec,
**self.kw)
N = max(Nx, Ny)
args = {}
if down:
gs = [gs[_n, 0] for _n in range(N)]
if share:
args['sharex'] = self.axes[-1]
else:
gs = [gs[0, _n] for _n in range(N)]
if share:
args['sharey'] = self.axes[-1]
self.axes.append(self.fig.add_subplot(gs[-1], **args))
for ax, subplotspec in zip(self.axes, gs):
ax.set_position(subplotspec.get_position(self.fig))
ax.set_subplotspec(subplotspec) # necessary if using tight_layout()
if self.right:
ax.yaxis.tick_right()
ax.yaxis.set_label_position("right")
if subplotspec is not gs[-1]:
ax.xaxis.set_ticklabels([])
ax.tick_params(axis='x', direction='in')
return ax
In [262]:
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec
# Here is the overall grid layout
gs = GridSpec(1, 2)
# In the left half we have some plot:
plt.subplot(gs[0,0])
plt.plot(ts, ts)
vg = Grid(wspace=0.0, hspace=0.0, right=True, subplot_spec=gs[0, 1])
# Now we construct our desired layout in the right half
ax = vg.next()
ax.plot(ts, a)
ax.set_ylabel('a')
# Plot b
ax = vg.next()
ax.plot(ts, b)
ax.set_ylabel('b')
# Plot c
ax = vg.next()
ax.plot(ts, c)
ax.set_ylabel('c')
ax.set_xlabel('t')
# Adjust layout if desired
plt.tight_layout()
In [ ]:
In [ ]: