Hello,
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 x.shape is (n,) and y.shape is (n, m), plotting the columns of the array. If y.shape is (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[1] - t[0])
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[0].plot(t, xm.T)
axs[1].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[1] - t[0])
xm_f = rfft(xm, axis=0)
_, axs = plt.subplots(2, 1)
axs[0].plot(t, xm)
axs[1].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?