pysmo.tools.plotutils

def plotseis(
    *seismograms: Seismogram,
    outfile: str = "",
    title: str = "",
    **kwargs: Any,
) -> matplotlib.figure.Figure:
    """Plot Seismogram objects.

    Returns the [`Figure`][matplotlib.figure.Figure] without displaying it;
    call `fig.show()` or integrate it into your own application.

    Args:
        seismograms: One or more seismogram objects. A `label` attribute, if
            present, labels that trace in the plot.
        outfile: Optionally save the figure to this filename.
        title: Optionally set the figure title.
        kwargs: Keyword arguments passed through to `matplotlib.pyplot.plot`.

    Returns:
        The matplotlib [`Figure`][matplotlib.figure.Figure] containing the plot.

    Raises:
        TypeError: If the unsupported `showfig` keyword is passed.

    Examples:
        ```python
        >>> from pysmo.classes import SAC
        >>> from pysmo.tools.plotutils import plotseis
        >>> seis = SAC.from_file("example.sac").seismogram
        >>> fig = plotseis(seis)
        >>>
        ```
    """
    if "showfig" in kwargs:
        raise TypeError(
            "plotseis() no longer accepts 'showfig'; it always returns the "
            + "Figure without displaying it. Call fig.show() yourself."
        )
    fig = plt.figure()
    any_labelled = False
    for seis in seismograms:
        time = time_array(seis)
        plot_kwargs = dict(kwargs)
        if "label" not in plot_kwargs:
            plot_kwargs["label"] = getattr(seis, "label", None)
        any_labelled = any_labelled or bool(plot_kwargs["label"])
        plt.plot(time, seis.data, scalex=True, scaley=True, **plot_kwargs)
    plt.xlabel("Time")
    plt.gcf().autofmt_xdate()
    fmt = mdates.DateFormatter("%H:%M:%S")
    plt.gca().xaxis.set_major_formatter(fmt)
    if not title:
        left, _ = plt.xlim()
        title = mdates.num2date(left).strftime("%Y-%m-%d %H:%M:%S")
    plt.title(title)
    if any_labelled:
        plt.legend()
    if outfile:
        plt.savefig(outfile)
    return fig

def relative_time_array(
    seismogram: Seismogram, reference: pd.Timestamp
) -> npt.NDArray[np.floating]:
    """Create an array of elapsed seconds relative to a reference time.

    Args:
        seismogram: Seismogram object.
        reference: Reference time.

    Returns:
        Array containing the elapsed time (in seconds) of each point in the
        seismogram data, relative to `reference`. Values are negative for
        points before `reference`.

    Examples:
        ```python
        >>> from pysmo.tools.plotutils import relative_time_array
        >>> from pysmo.classes import SAC
        >>> seis = SAC.from_file("example.sac").seismogram
        >>> reference = seis.begin_time + (seis.end_time - seis.begin_time) / 2
        >>> rel_times = relative_time_array(seis, reference)
        >>> bool(rel_times[0] < 0 < rel_times[-1])
        True
        >>>
        ```
    """
    start = (seismogram.begin_time - reference).total_seconds()
    end = (seismogram.end_time - reference).total_seconds()
    return np.linspace(start, end, len(seismogram.data))

def time_array(seismogram: Seismogram) -> npt.NDArray[np.floating]:
    """Create an array containing Matplotlib dates.

    Args:
        seismogram: Seismogram object.

    Returns:
        Array containing the Matplotlib dates (number of days since the
        Matplotlib epoch, default 1970-01-01) of each point in the
        seismogram data.

    Examples:
        ```python
        >>> from pysmo.tools.plotutils import time_array
        >>> from pysmo.classes import SAC
        >>> seis = SAC.from_file("example.sac").seismogram
        >>> seis_data = seis.data
        >>> seis_times = time_array(seis)
        >>> for t, v in zip(seis_times, seis_data):
        ...     print(t,v)
        ...
        14667.280625804839 -47201.0
        14667.280626383543 -47361.0
        14667.280626962245 -47511.0
        14667.28062754095 -47666.0
        14667.280628119654 -47826.0
        14667.280628698358 -47993.0
        ...
        >>>
        ```
    """
    start = mdates.date2num(seismogram.begin_time)
    end = mdates.date2num(seismogram.end_time)
    return np.linspace(start, end, len(seismogram.data))

def unix_time_array(seismogram: Seismogram) -> npt.NDArray[np.floating]:
    """Create an array containing Unix epoch times.

    Args:
        seismogram: Seismogram object.

    Returns:
        Array containing the Unix epoch time (seconds since 1970) of each
        point in the seismogram data.

    Examples:
        ```python
        >>> from pysmo.classes import SAC
        >>> from pysmo.tools.plotutils import unix_time_array
        >>> seis = SAC.from_file("example.sac").seismogram
        >>> seis_data = seis.data
        >>> seis_times = unix_time_array(seis)
        >>> for t, v in zip(seis_times, seis_data):
        ...     print(t,v)
        ...
        1267253046.069538 -47201.0
        1267253046.119538 -47361.0
        1267253046.169538 -47511.0
        1267253046.2195382 -47666.0
        1267253046.2695382 -47826.0
        1267253046.319538 -47993.0
        ...
        >>>
        ```
    """
    start = seismogram.begin_time.timestamp()
    end = seismogram.end_time.timestamp()
    return np.linspace(start, end, len(seismogram.data))