Source code for lab_driver.process_plots

import numpy as np
from os import makedirs
from os.path import join
from matplotlib import pyplot as plt




def get_plot_color(idx: int) -> str:
    """Getting the color string"""
    sel_color = ['k', 'r', 'b', 'g', 'y', 'c', 'm', 'gray']
    return sel_color[idx % len(sel_color)]





def get_plot_marker(idx: int) -> str:
    """Getting the marker for plotting"""
    sel_marker = '.+x_'
    return sel_marker[idx % len(sel_marker)]





def save_figure(fig, path: str, name: str, formats: list=('pdf', 'svg')) -> None:
    """Saving figure in given format
    Args:
        fig:        Matplot which will be saved
        path:       Path for saving the figure
        name:       Name of the plot
        formats:    List with data formats for saving the figures
    Returns:
        None
    """
    makedirs(path, exist_ok=True)
    path2fig = join(path, name)
    for idx, form in enumerate(formats):
        fig.savefig(f"{path2fig}.{form}", format=form)





def scale_auto_value(data: np.ndarray | float) -> [float, str]:
    """Getting the scaling value and corresponding string notation for unit scaling in plots
    Args:
        data:   Array or value for calculating the SI scaling value
    Returns:
        Tuple with [0] = scaling value and [1] = SI pre-unit
    """
    ref_dict = {'T': -4, 'G': -3, 'M': -2, 'k': -1, '': 0, 'm': 1, 'µ': 2, 'n': 3, 'p': 4, 'f': 5}

    value = np.max(np.abs(data)) if isinstance(data, np.ndarray) else data
    str_value = str(value).split('.')
    digit = 0
    if 'e' not in str_value[1]:
        if not str_value[0] == '0':
            # --- Bigger Representation
            sys = -np.floor(len(str_value[0]) / 3)
        else:
            # --- Smaller Representation
            for digit, val in enumerate(str_value[1], start=1):
                if '0' not in val:
                    break
            sys = np.ceil(digit / 3)
    else:
        val = int(str_value[1].split('e')[-1])
        sys = -np.floor(abs(val) / 3) if np.sign(val) == 1 else np.ceil(abs(val) / 3)

    scale = 10 ** (sys * 3)
    units = [key for key, div in ref_dict.items() if sys == div][0]
    return scale, units





def plot_transfer_function_norm(data: dict, path2save: str='',
                                xlabel: str='Stimulus Input', ylabel: str='Stimulus Output',
                                title: str='Transfer Function', file_name: str='') -> None:
    """Function for plotting the transfer function
    :param data:        Dictionary with extracted values from measurement data
    :param path2save:   Path for saving the figure
    :param xlabel:      Text Label for x-axis
    :param ylabel:      Text Label for y-axis
    :param title:       Text Label for title
    :param file_name:   File name of the saved figure
    :return:            None
    """
    val_input = data['stim']
    xaxis = np.linspace(start=val_input[0], stop=val_input[-1], num=9, endpoint=True)

    val_output = np.array([data[key]['mean'] for key in data.keys() if not key == 'stim'])
    yaxis = np.linspace(start=val_output.min(), stop=val_output.max(), num=9, endpoint=True)
    dy = np.diff(yaxis).max()

    plt.figure()
    for idx, key in enumerate(data.keys()):
        if not key == 'stim':
            plt.step(val_input, data[key]['mean'], where='mid', marker='.', c=get_plot_color(idx), label=key)
            plt.fill_between(val_input, data[key]['mean'] - data[key]['std'], data[key]['mean'] + data[key]['std'],
                             step='mid', alpha=0.3, color='gray')

    plt.xticks(xaxis)
    plt.xlim([val_input[0], val_input[-1]])
    plt.yticks(yaxis)
    plt.ylim([val_output.min()-dy, val_output.max()+dy])

    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)

    plt.legend(loc='upper left')
    plt.grid()
    plt.tight_layout()
    if path2save and file_name:
        save_figure(plt, path2save, f'{file_name.lower()}')

    plt.show()





def plot_transfer_function_metric(data: dict, func: object, path2save: str='',
                                  xlabel: str='Stimulus Input', ylabel: str='Stimulus Output',
                                  title: str='Transfer Function', file_name: str='') -> None:
    """Function for plotting the metric, extracted from the transfer function
    :param data:        Dictionary with pre-processed data from measurement with keys: ['stim', 'ch<x>': {'mean', 'std'}}
    :param func:        Function for calculating the metric
    :param path2save:   Path for saving the figure
    :param xlabel:      Text Label for x-axis
    :param ylabel:      Text Label for y-axis
    :param title:       Text Label for title
    :param file_name:   File name of the saved figure
    :return:            None
    """
    data_lsb = {'stim': data['stim']}
    for key in data.keys():
        if not key == 'stim':
            scale_val = 1.0
            metric = func(data['stim'], data[key]['mean'])
            if not metric.size == data['stim'].size:
                metric = np.concatenate((np.array((metric[0], )), metric), axis=0)
            data_lsb.update({key: {'mean': metric,
                                   'std': scale_val * data[key]['std']}})

    plot_transfer_function_norm(
        data=data_lsb,
        path2save=path2save,
        xlabel=xlabel,
        ylabel=ylabel,
        title=title,
        file_name=file_name
    )