plot_utils.py

from itertools import groupby

import numpy as np
from matplotlib import pyplot as plt, transforms
from matplotlib.figure import Figure
from matplotlib.lines import Line2D
from matplotlib.patches import Ellipse
from matplotlib.ticker import FormatStrFormatter

from dsp import AccelerationData, Segment, SegmentMask, Spectrum, Feature, SegmentType

def __append_acceleration_data_plots(axs: list[plt.Axes], data: AccelerationData):
    axs[0].plot(data.t, data.x, color="#43291f")
    axs[0].set_ylabel("x [m/s^2]")
    axs[1].plot(data.t, data.y, color="#43291f")
    axs[1].set_ylabel("y [m/s^2]")
    axs[2].plot(data.t, data.z, color="#43291f")
    axs[2].set_ylabel("z [m/s^2]")
    axs[3].plot(data.t, data.absolute, color="#43291f")
    axs[3].set_ylabel("abs [m/s^2]")
    return axs


def plot_acceleration_data(data: AccelerationData, fig_size: tuple[int, int] = (7, 6)) \
        -> tuple[Figure, list[plt.Axes]]:
    fig, axs = plt.subplots(4, figsize=fig_size, sharex=True)
    axs = __append_acceleration_data_plots(axs, data)
    axs[3].set_xlabel('t [s]')

    return fig, axs


def plot_resample_preview(
        orig: AccelerationData,
        res: AccelerationData,
        start: int = 0, length: int = 200) -> tuple[plt.Figure, plt.Axes]:
    fig, ax = plt.subplots(3, figsize=(10, 7), sharex=True)
    ratio = orig.size / res.size
    r = length
    s = start

    orig_slice = slice(int(s * ratio), int((r + s) * ratio))
    res_slice = slice(s, int((r + s)))

    color_line = "gray"
    color_marker = "#43291f"

    ax[0].plot(orig.t[orig_slice], orig.x[orig_slice], color=color_line)
    ax[0].plot(res.t[res_slice], res.x[res_slice], "ko", markersize=3, color=color_marker)
    ax[0].set_ylabel("x [m/s^2]")

    ax[1].plot(orig.t[orig_slice], orig.y[orig_slice], color=color_line)
    ax[1].plot(res.t[res_slice], res.y[res_slice], "ko", markersize=3, color=color_marker)
    ax[1].set_ylabel("y [m/s^2]")

    ax[2].plot(orig.t[orig_slice], orig.z[orig_slice], color=color_line)
    ax[2].plot(res.t[res_slice], res.z[res_slice], "ko", markersize=3, color=color_marker)
    ax[2].set_ylabel("z [m/s^2]")

    # ax[3].plot(orig.t[orig_slice], orig.absolute[orig_slice], color=color_line)
    # ax[3].plot(res.t[res_slice], res.absolute[res_slice], "ko", markersize=3, color=color_marker)
    # ax[3].set_ylabel("abs [m/s^2]")
    # ax[3].set_xlabel('t [s]')

    for a in ax:
        a.yaxis.set_major_formatter(FormatStrFormatter('%.1f'))

    fig.align_ylabels(ax)
    fig.subplots_adjust(hspace=0.5)

    return fig, ax


def plot_segments(segments: list[Segment],
                  segments_mask: SegmentMask,
                  fig_size: tuple[int, int] = (15, 7)) -> tuple[Figure, list[plt.Axes]]:
    fig, axs = plt.subplots(5, figsize=fig_size, sharex=True)
    data = segments[0].data

    axs = __append_acceleration_data_plots(axs, data)

    # add threshold line to absolute subplot
    axs[3].axhline(data.segment_threshold(), color="#da2c38", linestyle='--',
                   label=f'threshold={data.segment_threshold():.3f} ~ {data.threshold_percentile:.0f}th perc.')
    axs[3].legend(loc='upper right')

    # add mask subplot
    axs[4].plot(data.t, segments_mask.astype(int), color="#87c38f", linewidth=1)
    axs[4].set_ylabel('segment\nmask')
    axs[4].set_yticks([False, True])
    axs[4].set_xlabel('t [s]')

    # add segments to all subplots
    for idx, segment in enumerate(segments):
        # for x, y, z and absolute subplots
        for i in range(4):
            axs[i].axvspan(segment.start, segment.end, color="#87c38f", alpha=0.15)

        # for mask subplot, also add label
        axs[4].axvspan(segment.start, segment.end, color="#87c38f", alpha=0.3)
        axs[4].text((segment.start + segment.end) / 2, axs[4].get_ylim()[1] * 1 / 2, f'{idx}', ha='center', va='center',
                    fontsize=20,
                    color='black', rotation=0,
                    bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=2.0))

    for a in axs:
        a.yaxis.set_major_formatter(FormatStrFormatter('%.1f'))

    fig.align_ylabels(axs)

    return fig, axs


def plot_spectra(spectra: list[Spectrum],
                 fig_size: tuple[int, int] = (7, 5)) -> tuple[Figure, list[plt.Axes]]:
    fig, axs = plt.subplots(3, figsize=fig_size, sharex=True)

    freq: np.ndarray = spectra[0].freq
    combined_x: np.ndarray = spectra[0].x
    combined_y: np.ndarray = spectra[0].y
    combined_z: np.ndarray = spectra[0].z

    for spectrum in spectra[1:]:
        combined_x += spectrum.x
        combined_y += spectrum.y
        combined_z += spectrum.z

    axs[0].plot(freq, combined_x, color="#43291f")
    axs[0].set_ylabel("x")
    axs[1].plot(freq, combined_y, color="#43291f")
    axs[1].set_ylabel("y")
    axs[2].plot(freq, combined_z, color="#43291f")
    axs[2].set_ylabel("z")

    axs[2].set_xlabel('frequency [Hz]')

    for a in axs:
        a.yaxis.set_major_formatter(FormatStrFormatter('%.1f'))

    fig.align_ylabels(axs)
    fig.subplots_adjust(hspace=0.5)

    return fig, axs


class FinalPlot:

    def __init__(self, features: list[Feature]):
        self.features = features

    @staticmethod
    def get_color_by_type(segment_type: SegmentType):
        return {
            SegmentType.FLAT: "#43291f",
            SegmentType.DOWN: "#da2c38",
            SegmentType.UP: "#226f54"
        }[segment_type]

    def __selected_features(self, features: list[Feature] = None):
        if features is None:
            features = self.features

        f1 = [f.re_x for f in features]
        f2 = [f.re_z for f in features]

        return (
            np.array(f1) * 100,
            np.array(f2) * 100,
            f"relative energy [%]\ndirection x | <{features[0].f_low}, {features[0].f_high}> Hz",
            f"relative energy [%]\ndirection z |  <{features[0].f_low}, {features[0].f_high}> Hz",
            # "Standard deviation [m s^-2] | direction y"
        )

    def plot_features(self, fig_size: tuple[int, int] = (7, 5)) -> tuple[Figure, plt.Axes]:
        fig, axs = plt.subplots(1, figsize=fig_size)

        f1, f2, f1_label, f2_label = self.__selected_features()
        colors = [self.get_color_by_type(f.segment_type) for f in self.features]
        labels = map(lambda f: f.segment_type.name, self.features)

        axs.scatter(f1, f2, color=colors, sizes=0.5 * 10 * np.ones(len(f1)), label=labels)
        axs.set_xlabel(f1_label)
        axs.set_ylabel(f2_label)

        legend_elements = [
            Line2D([0], [0], marker='o', color='w', label='UP',
                   markerfacecolor=self.get_color_by_type(SegmentType.UP), markersize=8),
            Line2D([0], [0], marker='o', color='w', label='FLAT',
                   markerfacecolor=self.get_color_by_type(SegmentType.FLAT), markersize=8),
            Line2D([0], [0], marker='o', color='w', label='DOWN',
                   markerfacecolor=self.get_color_by_type(SegmentType.DOWN), markersize=8)
        ]

        axs.legend(handles=legend_elements, loc='best')

        return fig, axs

    @staticmethod
    def __add_confidence_ellipse(ax, x, y, n_std: float = 1.0, edge_color: str = 'black',
                                 center_point_size: float = 50.0, **kwargs):
        if x.size < 2:
            return None

        # calculate covariance and Pearson correlation coefficient
        cov = np.cov(x, y)
        pearson = cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])

        # calculate major and minor axes for the unit ellipse (pre-scaling/rotation)
        ell_radius_x: float = float(np.sqrt(1 + pearson))
        ell_radius_y: float = float(np.sqrt(1 - pearson))
        ellipse = Ellipse((0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2,
                          edgecolor=edge_color, facecolor='none', **kwargs)

        # Scaling factors based on standard deviations and the desired n_std
        scale_x: float = float(np.sqrt(cov[0, 0])) * n_std
        scale_y: float = float(np.sqrt(cov[1, 1])) * n_std

        mean_x: float = float(np.mean(x))
        mean_y: float = float(np.mean(y))

        ax.scatter(mean_x, mean_y, color=edge_color, sizes=[center_point_size])

        # Transformation to rotate, scale, and translate the ellipse
        transform = transforms.Affine2D() \
            .rotate_deg(45) \
            .scale(scale_x, scale_y) \
            .translate(mean_x, mean_y)

        ellipse.set_transform(transform + ax.transData)
        return ax.add_patch(ellipse)

    def add_confidence_ellipses(self, ax, n_std: list[float] = None):
        if n_std is None:
            n_std = [0.5, 1.0, 1.5]

        features_grouped = groupby(self.features, lambda f: f.segment_type)

        for segment_type, features in features_grouped:
            features = list(features)

            color = self.get_color_by_type(segment_type)
            f1, f2, _, _ = self.__selected_features(features)

            for std in n_std:
                self.__add_confidence_ellipse(ax, f1, f2, n_std=std, edge_color=color)

        pass