I’m curious to know how others here handle an issue that I’ve encountered with plotting 2-D NumPy arrays. Matplotlib will accept
plot(x, y) if
(n, m), plotting the columns of the array. If
(m, n), however, it will return an error instead of plotting the rows of the array. The NumPy broadcasting rules and FFT routines are organized around row operations, which forces a choice between rows and columns.
For example, in the code snippet below, I simulate a set of
nw noisy waveforms of length
ns, then plot both the set of waveforms and their power spectral density.
import numpy as np import matplotlib.pyplot as plt from numpy.random import default_rng from numpy.fft import rfftfreq, rfft from scipy.signal import gausspulse ns = 256 nw = 10 dt = 10 / ns t = np.arange(-ns * dt / 2, ns * dt / 2, dt) x = gausspulse(t, fc=1) rng = default_rng(0) # Broadcasting matches rightmost indices xm = x + 0.1 * np.abs(x) * rng.standard_normal((nw, ns)) # Default axis for FFT is -1 f = rfftfreq(ns, t - t) xm_f = rfft(xm) # Need to take the transpose of xm and xm_f to convert rows to columns _, axs = plt.subplots(2, 1) axs.plot(t, xm.T) axs.plot(f, np.abs(xm_f.T) ** 2) plt.show()
To do the same operations with a column orientation, one needs to introduce
newaxis and explicitly apply the FFT along the first dimension.
# Organize around columns xm = x[:, np.newaxis] + 0.1 * np.abs(x[:, np.newaxis]) * rng.standard_normal((ns, nw)) f = rfftfreq(ns, t - t) xm_f = rfft(xm, axis=0) _, axs = plt.subplots(2, 1) axs.plot(t, xm) axs.plot(f, np.abs(xm_f) ** 2) plt.show()
This is not a huge problem, but it is a bit of an annoyance for someone who is used to MATLAB, where the FFT routines are organized around columns and the
plot function accepts both orientations of
y, so long as one of its dimensions is compatible with
x. Has anyone found a better way to deal with this? And are there any plans to implement the more flexible MATLAB-like array compatibility into Matplotlib?