neurodsp-tools / neurodsp

Digital signal processing for neural time series.
https://neurodsp-tools.github.io/
Apache License 2.0
281 stars 61 forks source link

[ENH] Plot spectra accepts 2d arrays #269

Closed ryanhammonds closed 3 years ago

ryanhammonds commented 3 years ago

Before, plot_power_spectra only accepted 1d array or list of 1d array for freqs and powers. This adds support for 2d arrays. The requirement to call .tolist() on a 2d power array for plotting to work felt strange - espcially since compute_spectrum returns 2d arrays for 2d inputs.

import numpy as np
from neurodsp.sim import sim_oscillation
from neurodsp.spectral import compute_spectrum
from neurodsp.plts import plot_power_spectra

n_seconds = 10
fs = 1000
freqs = [10, 20]

# Simulate 2d array
sigs = np.zeros((2, n_seconds * fs))

for ind, freq in enumerate(freqs):
    sigs[ind] = sim_oscillation(n_seconds, fs, freq)

freqs, powers = compute_spectrum(sigs, fs)

# Freqs: list of 1d array
plot_power_spectra([freqs, freqs], powers)

# Freqs: 2d array
plot_power_spectra(np.array([freqs, freqs]), powers)

# Powers: list of 1d array
plot_power_spectra(freqs, [powers[0], powers[1]])

# Powers: 2d array
plot_power_spectra(freqs, powers)
TomDonoghue commented 3 years ago

This looks like a super sensible and nice update! I left one small comment on the docs, otherwise, LGTM!