<commit_before>"""Plots for time series."""

from itertools import repeat, cycle

import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt

from neurodsp.plts.style import style_plot
from neurodsp.plts.utils import check_ax, savefig

###################################################################################################
###################################################################################################

@savefig
@style_plot
def plot_time_series(times, sigs, labels=None, colors=None, ax=None, **kwargs):
    """Plot a time series.

    Parameters
    ----------
    times : 1d array or list of 1d array
        Time definition(s) for the time series to be plotted.
    sigs : 1d array or list of 1d array
        Time series to plot.
    labels : list of str, optional
        Labels for each time series.
    colors : str or list of str
        Colors to use to plot lines.
    ax : matplotlib.Axes, optional
        Figure axes upon which to plot.
    **kwargs
        Keyword arguments for customizing the plot.

    Examples
    --------
    Create a time series plot:

    >>> from neurodsp.sim import sim_combined
    >>> from neurodsp.utils import create_times
    >>> sig = sim_combined(n_seconds=10, fs=500,
    ...                    components={'sim_powerlaw': {'exponent': -1.5, 'f_range': (2, None)},
    ...                                'sim_oscillation' : {'freq': 10}})
    >>> times = create_times(n_seconds=10, fs=500)
    >>> plot_time_series(times, sig)
    """

    ax = check_ax(ax, (15, 3))

    times = repeat(times) if isinstance(times, np.ndarray) else times
    sigs = [sigs] if isinstance(sigs, np.ndarray) else sigs

    if labels is not None:
        labels = [labels] if not isinstance(labels, list) else labels
    else:
        labels = repeat(labels)

    # If not provided, default colors for up to two signals to be black & red
    if not colors and len(sigs) <= 2:
        colors = ['k', 'r']
    colors = repeat(colors) if not isinstance(colors, list) else cycle(colors)

    for time, sig, color, label in zip(times, sigs, colors, labels):
        ax.plot(time, sig, color=color, label=label)

    ax.set_xlabel('Time (s)')
    ax.set_ylabel('Voltage (uV)')


@savefig
@style_plot
def plot_instantaneous_measure(times, sigs, measure='phase', ax=None, **kwargs):
    """Plot an instantaneous measure, of phase, amplitude or frequency.

    Parameters
    ----------<commit_msg>allow plotting of 2d arrays<commit_after>"""Plots for time series."""

from itertools import repeat, cycle

import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt

from neurodsp.plts.style import style_plot
from neurodsp.plts.utils import check_ax, savefig

###################################################################################################
###################################################################################################

@savefig
@style_plot
def plot_time_series(times, sigs, labels=None, colors=None, ax=None, **kwargs):
    """Plot a time series.

    Parameters
    ----------
    times : 1d array, list of 1d array, or 2d array
        Time definition(s) for the time series to be plotted.
    sigs : 1d array, list of 1d array, or 2d array
        Time series to plot.
    labels : list of str, optional
        Labels for each time series.
    colors : str or list of str
        Colors to use to plot lines.
    ax : matplotlib.Axes, optional
        Figure axes upon which to plot.
    **kwargs
        Keyword arguments for customizing the plot.

    Examples
    --------
    Create a time series plot:

    >>> from neurodsp.sim import sim_combined
    >>> from neurodsp.utils import create_times
    >>> sig = sim_combined(n_seconds=10, fs=500,
    ...                    components={'sim_powerlaw': {'exponent': -1.5, 'f_range': (2, None)},
    ...                                'sim_oscillation' : {'freq': 10}})
    >>> times = create_times(n_seconds=10, fs=500)
    >>> plot_time_series(times, sig)
    """

    ax = check_ax(ax, (15, 3))

    n_repeats = len(sigs) if isinstance(sigs, list) or sigs.ndim == 2 else 1

    # Repeat times if needed
    if isinstance(times, np.ndarray) and times.ndim != 2:
        times = np.tile(times, (n_repeats, 1))

    # Make sigs iterable if 1D
    sigs = np.reshape(sigs, (1, -1)) if not isinstance(sigs, list) and sigs.ndim == 1 else sigs

    if labels is not None:
        labels = [labels] if not isinstance(labels, list) else labels
    else:
        labels = repeat(labels, n_repeats)

    # If not provided, default colors for up to two signals to be black & red
    if not colors and len(sigs) <= 2:
        colors = ['k', 'r']
    colors = repeat(colors, n_repeats) if not isinstance(colors, list) else cycle(colors)

    for time, sig, color, label in zip(times, sigs, colors, labels):
        ax.plot(time, sig, color=color, label=label)

    ax.set_xlabel('Time (s)')
    ax.set_ylabel('Voltage (uV)')


@savefig
@style_plot
def plot_instantaneous_measure(times, sigs, measure='phase', ax=None, **kwargs):
    """Plot an instantaneous measure, of phase, amplitude or frequency.

    Parameters
    ----------
{'fraction_non_alphanumeric': 0.18391689111196613, 'fraction_numerical': 0.010965756060023086, 'mean_word_length': 2.993087557603687, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '942215', 'n_tokens_mistral': 1595, 'n_tokens_neox': 1501, 'n_words': 524}