Skip to content

Commit

Permalink
add plot methods to data objects
Browse files Browse the repository at this point in the history
  • Loading branch information
TomDonoghue committed Apr 9, 2024
1 parent 4d832f1 commit facd87e
Showing 1 changed file with 13 additions and 1 deletion.
14 changes: 13 additions & 1 deletion specparam/objs/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from specparam.core.errors import DataError, InconsistentDataError
from specparam.data import SpectrumMetaData
from specparam.plts.settings import PLT_COLORS
from specparam.plts.spectra import plot_spectra
from specparam.plts.spectra import plot_spectra, plot_spectrogram
from specparam.plts.utils import check_plot_kwargs

###################################################################################################
Expand Down Expand Up @@ -381,6 +381,12 @@ def add_data(self, freqs, spectrogram, freq_range=None):
super().add_data(freqs, spectrogram, freq_range)


def plot(self, **plt_kwargs):
"""Plot the spectrogram."""

plot_spectrogram(self.freqs, self.spectrogram, **plot_kwargs)


class BaseData3D(BaseData2DT):
"""Base object for managing data for spectral parameterization - for 3D data."""

Expand Down Expand Up @@ -440,3 +446,9 @@ def add_data(self, freqs, spectrograms, freq_range=None):
# Otherwise, pass through 2d array to underlying object method
else:
super().add_data(freqs, spectrograms, freq_range)


def plot(self, event_ind):
"""Plot a selected spectrogram."""

plot_spectrogram(self.freqs, self.spectrograms[event_ind, :, :], **plot_kwargs)

0 comments on commit facd87e

Please sign in to comment.