survBot/utils.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import matplotlib
import numpy as np

from obspy import Stream


def get_bg_color(check_key, status, dt_thresh=None, hex=False):
    message = status.message
    if check_key == 'last active':
        bg_color = get_time_delay_color(message, dt_thresh)
    elif check_key == 'temp':
        bg_color = get_temp_color(message)
    elif check_key == 'mass':
        bg_color = get_mass_color(message)
    else:
        if status.is_warn:
            bg_color = get_warn_color(status.count)
        elif status.is_error:
            bg_color = get_color('FAIL')
        else:
            bg_color = get_color(message)
    if not bg_color:
        bg_color = get_color('undefined')

    if hex:
        bg_color = '#{:02x}{:02x}{:02x}'.format(*bg_color[:3])
    return bg_color


def get_color(key):
    # some GUI default colors
    colors_dict = {'FAIL': (255, 50, 0, 255),
                   'NO DATA': (255, 255, 125, 255),
                   'WARN': (255, 255, 80, 255),
                   'OK': (125, 255, 125, 255),
                   'undefined': (230, 230, 230, 255)}
    return colors_dict.get(key)


def get_color_mpl(key):
    color_tup = get_color(key)
    return np.array([color/255. for color in color_tup])


def get_time_delay_color(dt, dt_thresh):
    """ Set color of time delay after thresholds specified in self.dt_thresh """
    if isinstance(dt, type(dt_thresh[0])):
        if dt < dt_thresh[0]:
            return get_color('OK')
        elif dt_thresh[0] <= dt < dt_thresh[1]:
            return get_color('WARN')
    return get_color('FAIL')


def get_warn_color(count):
    color = (min([255, 200 + count ** 2]), 255, 80, 255)
    return color


def get_mass_color(message):
    # can change this to something else if wanted. This way it always returns get_color (without warn count)
    if isinstance(message, (float, int)):
        return get_color('OK')
    return get_color(message)


def get_temp_color(temp, vmin=-10, vmax=60, cmap='coolwarm'):
    """ Get an rgba temperature value back from specified cmap, linearly interpolated between vmin and vmax. """
    if type(temp) in [str]:
        return get_color('undefined')
    cmap = matplotlib.cm.get_cmap(cmap)
    val = (temp - vmin) / (vmax - vmin)
    rgba = [int(255 * c) for c in cmap(val)]
    return rgba


def modify_stream_for_plot(input_stream, parameters):
    """ copy (if necessary) and modify stream for plotting """

    # make a copy
    st = Stream()

    channels_dict = parameters.get('CHANNELS')

    # iterate over all channels and put them to new stream in order
    for index, ch_tup in enumerate(channels_dict.items()):
        # unpack tuple from items
        channel, channel_dict = ch_tup

        # get correct channel from stream
        st_sel = input_stream.select(channel=channel)
        # in case there are != 1 there is ambiguity
        if not len(st_sel) == 1:
            continue

        # make a copy to not modify original stream!
        tr = st_sel[0].copy()

        # multiply with conversion factor for unit
        unit_factor = channel_dict.get('unit')
        if unit_factor:
            tr.data = tr.data * float(unit_factor)

        # apply transformations if provided
        transform = channel_dict.get('transform')
        if transform:
            tr.data = transform_trace(tr.data, transform)

        # modify trace id to maintain plotting order
        name = channel_dict.get('name')
        tr.id = f'{index + 1}: {name} - {tr.id}'

        st.append(tr)

    return st


def transform_trace(data, transf):
    """
    Transform trace with arithmetic operations in order, specified in transf
    @param data: numpy array
    @param transf: list of lists with arithmetic operations (e.g. [['*', '20'], ] -> multiply data by 20
    """
    # This looks a little bit hardcoded, however it is safer than using e.g. "eval"
    for operator_str, val in transf:
        if operator_str == '+':
            data = data + val
        elif operator_str == '-':
            data = data - val
        elif operator_str == '*':
            data = data * val
        elif operator_str == '/':
            data = data / val
        else:
            raise IOError(f'Unknown arithmethic operator string: {operator_str}')

    return data


def set_axis_ylabels(fig, parameters, verbosity=0):
    """
    Adds channel names to y-axis if defined in parameters.
    """
    names = [channel.get('name') for channel in parameters.get('CHANNELS').values()]
    if not names: # or not len(st.traces):
        return
    if not len(names) == len(fig.axes):
        if verbosity:
            print('Mismatch in axis and label lengths. Not adding plot labels')
        return
    for channel_name, ax in zip(names, fig.axes):
        if channel_name:
            ax.set_ylabel(channel_name)


def set_axis_color(fig, color='0.8'):
    """
    Set all axes of figure to specific color
    """
    for ax in fig.axes:
        for key in ['bottom', 'top', 'right', 'left']:
            ax.spines[key].set_color(color)


def set_axis_yticks(fig, parameters, verbosity=0):
    """
    Adds channel names to y-axis if defined in parameters.
    """
    ticks = [channel.get('ticks') for channel in parameters.get('CHANNELS').values()]
    if not ticks:
        return
    if not len(ticks) == len(fig.axes):
        if verbosity:
            print('Mismatch in axis tick and label lengths. Not changing plot ticks.')
        return
    for ytick_tripple, ax in zip(ticks, fig.axes):
        if not ytick_tripple:
            continue
        ymin, ymax, step = ytick_tripple

        yticks = list(np.arange(ymin, ymax + step, step))
        ax.set_yticks(yticks)
        ax.set_ylim(ymin - 0.33 * step, ymax + 0.33 * step)


def plot_axis_thresholds(fig, parameters, verbosity=0):
    """
    Adds channel thresholds (warn, fail) to y-axis if defined in parameters.
    """
    if verbosity > 0:
        print('Plotting trace thresholds')

    keys_colors = {'warn': dict(color=0.8 * get_color_mpl('WARN'), linestyle=(0, (5, 10)), alpha=0.5, linewidth=0.7),
                   'fail': dict(color=0.8 * get_color_mpl('FAIL'), linestyle='solid', alpha=0.5, linewidth=0.7)}

    for key, kwargs in keys_colors.items():
        channel_threshold_list = [channel.get(key) for channel in parameters.get('CHANNELS').values()]
        if not channel_threshold_list:
            continue
        plot_threshold_lines(fig, channel_threshold_list, parameters, **kwargs)


def plot_threshold_lines(fig, channel_threshold_list, parameters, **kwargs):
    for channel_thresholds, ax in zip(channel_threshold_list, fig.axes):
        if not channel_thresholds:
            continue

        if not isinstance(channel_thresholds, (list, tuple)):
            channel_thresholds = [channel_thresholds]

        for warn_thresh in channel_thresholds:
            if isinstance(warn_thresh, str):
                warn_thresh = parameters.get('THRESHOLDS').get(warn_thresh)
            if type(warn_thresh in (float, int)):
                ax.axhline(warn_thresh, **kwargs)