Xarray Extensions for Peak Detection and Visualization: This Python code snippet extends xarray DataArrays with custom accessors for detecting and plotting positive and negative peaks in 1D signals, enhancing data analysis and visualization capabilities.

xarray_accessor_peak_detection_and_plotting.md

Xarray Peak Detection and Visualization 📈

TL;DR

Enhance your data analysis with custom Xarray accessors for detecting and visualizing both positive and negative peaks in 1D signals. This Python code snippet makes it easy to identify and annotate peaks directly on your plots. doc xarray

Features

• Find Peaks: Utilize the find_peaks accessor to detect peaks in your DataArray.
• Plot Peaks: With the plot_peaks accessor, visualize positive and negative peaks on your plots, including annotations for peak indices.

Output

# find peaks
return find_peaks(self._da, **kwargs)
# plot
return {
    'pos': peaks_pos,
    'neg': peaks_neg,
    'pos_real': peaks_pos_real,
    'neg_real': peaks_neg_real,
    "neg_values": self._da[peaks_neg],
    "pos_values": self._da[peaks_pos],
    'gca': plt.gca()
}

Usage

1. Detect Peaks: Simply call your DataArray with .find_peaks() to detect peaks based on your specified criteria.
2. Visualize Peaks: Use .plot_peaks(pos_neg=True) on your DataArray to plot and annotate positive and negative peaks.

Example

import xarray as xr
# Assuming `da` is your xarray DataArray
peaks_info = da.plot_peaks(pos_neg=True, height=0.1)

plot = data.u.mean(["t","x"]).diff("y").plot_peaks(pos_neg=True, height=0.5)["gca"]
plot.set_title("u mean profile")

xarray_accessor_peak_detection_and_plotting.py

@xr.register_dataarray_accessor("find_peaks")
class FindPeaksAccessor:
    """Simple accessor to find peaks in a 1D signal using scipy.signal.find_peaks function. works only on 1D signal and exacly like the scipy function.
    """
    def __init__(self, data_array):
        self._da = data_array

    def __call__(self, **kwargs):
        return find_peaks(self._da, **kwargs)

# accessors to find and plot the peaks on the 1d signal of dataarray
@xr.register_dataarray_accessor("plot_peaks")
class PlotPeaksAccessor:
    """
    A custom accessor for xarray DataArray to find and plot peaks.

    This accessor extends xarray DataArrays to allow for easy identification and visualization
    of positive and negative peaks within the data. It provides methods to find peaks based on
    various criteria and plot them with optional annotations.

    Parameters
    ----------
    data_array : xarray.DataArray
        The DataArray to which the accessor is attached.

    Methods
    -------
    pos_neg_peaks(**kwargs)
        Identifies positive and negative peaks in the DataArray.
    __call__(pos_neg=False, anotate=False, vertical_offset=10, height=0.1, threshold=None, distance=None, prominence=None, width=None, wlen=None, rel_height=0.5, plateau_size=None, **kwargs)
        Main method to find and optionally plot and annotate peaks in the DataArray.
    plot_peaks(x_values, y_values, color, label)
        Helper method to plot identified peaks on a graph.
    annotate_peaks(x_values, y_values, indices, vertical_offset=10)
        Annotates identified peaks with their indices on the plot.

    Returns
    -------
    dict
        A dictionary containing indices of positive and negative peaks, their real-world values,
        the values of the DataArray at these peak points, and the current matplotlib Axes object.
    """
    def __init__(self, data_array):
        self._da = data_array

    def pos_neg_peaks(self, **kwargs):
        pos = self._da.where(self._da > 0, 0)
        neg = self._da.where(self._da < 0, 0)
        peaks_pos, _ = find_peaks(pos, **kwargs)
        peaks_neg, _ = find_peaks(-neg, **kwargs)
        return peaks_pos, peaks_neg


    def __call__(self,
                 pos_neg=False,
                 anotate=False,
                 vertical_offset=10,
                 height=0.1,
                 threshold=None,
                 distance=None,
                 prominence=None,
                 width=None,
                 wlen=None,
                 rel_height=0.5,
                 plateau_size=None,
                 **kwargs):
        # Determine peaks based on whether we're looking for positive and negative peaks
        if pos_neg:
            peaks_pos, peaks_neg = self.pos_neg_peaks(height=height)
        else:
            peaks_pos, _ = find_peaks(self._da, height=height)
            peaks_neg = []

        # Convert peak indices to real x-axis values
        dim_ax = self._da.dims[0]  # Dimension axis
        peaks_pos_real = self._da[dim_ax][peaks_pos]
        peaks_neg_real = self._da[dim_ax][peaks_neg]

        # Plot the data array
        self._da.plot(**kwargs)
        # Plot positive and negative peaks
        self.plot_peaks(peaks_pos_real, self._da[peaks_pos], 'red', 'pos peaks')
        self.plot_peaks(peaks_neg_real, self._da[peaks_neg], 'green', 'neg peaks')

        # Annotate peaks with their indices
        if anotate:
            self.annotate_peaks(peaks_pos_real, self._da[peaks_pos], peaks_pos,vertical_offset)
            self.annotate_peaks(peaks_neg_real, self._da[peaks_neg], peaks_neg,vertical_offset)

        return {
            'pos_index': peaks_pos,
            'neg_index': peaks_neg,
            'pos_real': peaks_pos_real,
            'neg_real': peaks_neg_real,
            "neg_values": self._da[peaks_neg],
            "pos_values": self._da[peaks_pos],
            'gca': plt.gca()
        }


    def plot_peaks(self, x_values, y_values, color, label):
        """Plot peaks on the graph."""
        plt.scatter(x_values, y_values, color=color, label=label)

    def annotate_peaks(self, x_values, y_values, indices,vertical_offset=10):
        """Annotate peaks with their indices."""
        for i, index in enumerate(indices):
            plt.annotate(index, (x_values[i], y_values[i]),
                         textcoords="offset points",
                         xytext=(0, vertical_offset),  # Use the fixed offset directly
                         ha='center')