bat-studio/modules/graphs/linear_target.py

import numpy as np

from modules.graphs.style_settings import DEFAULT_STYLE_SETTINGS


class LinearTarget:
    def __init__(self, fig, subplot_spec, max_dist, step, threshold, color_settings=DEFAULT_STYLE_SETTINGS):
        self._plt = fig.add_subplot(subplot_spec)
        self._max_dist = max_dist
        self._step = step
        self._threshold = threshold
        self._color_settings = color_settings

    def plot(self, z, resize_by_data=True):
        if resize_by_data and len(z) > 0:
            actual_max_dist = max(np.max(np.abs(z)), self._max_dist)
        else:
            actual_max_dist = self._max_dist

        x = np.zeros(z.size)

        self._show_scale(actual_max_dist + self._step)


        self._show_ticks(actual_max_dist + self._step)

        self._plt.scatter(
            x,z,
            c=self._color_settings.color_data_fill,
            s=self._color_settings.size_data_point,
            edgecolors=self._color_settings.color_data_border,
            linewidths=self._color_settings.size_data_border,
            zorder=3,
            alpha=self._color_settings.alpha_data
        )

        # Vyznačenie nuly a úprava osí
        self._plt.axhline(0, color=self._color_settings.color_contour, linestyle='-', alpha=1.0, linewidth=1.5, zorder=2)
       # self._plt.set_xlim(-0.15, 0.15)
        self._plt.set_xticks([])

        self._plt.set_ylim(-actual_max_dist, actual_max_dist)
        #self._plt.set_ylabel("Odchýlka hĺbky Z (mm)", fontweight='bold', color=self._color_settings.color_contour)
        #self._plt.set_title("Jednoosová hĺbka", fontsize=11, fontweight='bold', color=self._color_settings.color_contour, pad=10)

        # Mriežka a rámčeky
        self._plt.grid(True, axis='y', linestyle='-', color=self._color_settings.color_contour, alpha=0.8)
        for spine in self._plt.spines.values():
            spine.set_edgecolor(self._color_settings.color_contour)
            spine.set_alpha(self._color_settings.alpha_contour)
        self._plt.tick_params(colors=self._color_settings.color_contour, grid_color=self._color_settings.color_contour)

    def _show_scale(self, y_limit_val):
        for val in range(0, int(y_limit_val) + self._step, self._step):
            idx = val // self._step

            # show positive
            if val + self._step <= self._threshold:
                c_pos = self._color_settings.color_ok_dark if idx % 2 == 0 else self._color_settings.color_ok_light
            else:
                c_pos = self._color_settings.color_err_dark if idx % 2 == 0 else self._color_settings.color_err_light
            self._plt.axhspan(val, val + self._step, color=c_pos, alpha=self._color_settings.alpha_scale, zorder=0)

            # Záporný smer Z
            if val + self._step <= self._threshold:
                c_neg = self._color_settings.color_ok_dark if idx % 2 == 0 else self._color_settings.color_ok_light
            else:
                c_neg = self._color_settings.color_err_dark if idx % 2 == 0 else self._color_settings.color_err_light
            self._plt.axhspan(-val - self._step, -val, color=c_neg, alpha=self._color_settings.alpha_scale, zorder=0)

    def _show_ticks(self, max_dist):
        self._plt.set_xticks([])
        self._plt.set_xticklabels([])

        self._plt.set_yticks(np.arange(-max_dist, max_dist+1, self._step))
        self._plt.grid(True, linestyle='-', color=self._color_settings.color_contour, alpha=self._color_settings.alpha_contour)
        #self._plt.spines['polar'].set_color(self._color_settings.color_contour)
        #self._plt.spines['polar'].set_alpha(self._color_settings.alpha_contour)