cobrawap_publication_code/plot_figures/scripts/plotting_utils.py

import numpy as np
import pandas as pd
from pathlib import Path
import matplotlib as mpl

int_types = ['wavefronts_id', 'wavemodes_id', 'wavemode', 'channel_id', 'dim_x',
             'dim_y', 'number_of_triggers', 'x_coords', 'y_coords']
float_types = ['spatial_scale', 'sampling_rate', 'recording_length', 'duration',
               'inter_wave_interval', 'inter_wave_interval_std',
               'inter_wave_interval_local', 'planarity', 'velocity_planar',
               'velocity_planar_std', 'velocity_local', 'direction_x',
               'direction_y', 'direction_x_std', 'direction_y_std',
               'direction_local_x', 'direction_local_y', 'time_stamp',
               'flow_direction_local_x', 'flow_direction_local_y']
df_dtypes = dict.fromkeys(int_types, 'Int64') | dict.fromkeys(float_types, float) \
                                              | {'is_planar': bool}


colormap = {'WT': '#35618f',
            'KO': '#72be3e',
            'WBS': '#35618f',
            'FXS': '#72be3e',
            'mix': '#f56808',
            'ketamine': '#dd4601',
            'isoflurane': '#ebd111',
            'light': '#f8765c',
            'medium': '#d3436e',
            'deep': '#982d80'}


def filter_velocity_local(df, vmin=0.0, vmax=120):
    # minimum velocity
    df = df[df.velocity_local > vmin]
    # maximum velocity
    df = df[df.velocity_local < vmax]
    # nan values
    df = df.loc[pd.notnull(df.velocity_local)]
    return df

def load_df(filename, dtype=None):
    filename = Path(filename)
    if dtype is None:
        if 'avg' in str(filename.name):
            dtype = df_dtypes | dict(wavefronts_id=float,
                                     number_of_triggers=float,
                                     is_planar=float)
        else:
            dtype = df_dtypes

    keep = ['wavefronts_id', 'wavemode', 'wavemodes_id', 
            'anesthetic', 'profile', 'technique', 'disease_model', 'model_type']
    local_measures = ['velocity_local', 'x_coords', 'y_coords',
                      'direction_local_x', 'direction_local_y', 'inter_wave_interval_local', 
                      'flow_direction_local_x', 'flow_direction_local_y',
                      'angle_local', 'flow_angle_local']
    global_measures = ['velocity_planar', 'direction_x', 'direction_y',
                       'planarity', 'inter_wave_interval', 'number_of_triggers',
                       'angle']
    
    usecols = keep + local_measures if 'channel-wise' in str(filename.name) \
         else keep + global_measures
    
    df = pd.read_csv(filename, usecols=lambda c: c in usecols, 
                     dtype=dtype, low_memory=False)

    for measure in ['direction_local', 'flow_direction_local', 'direction']:
        if f'{measure}_x' in df.columns and f'{measure}_y' in df.columns:
            df = add_angle_column(df, measure)

    for measure in ['wavemode', 'wavemodes_id', 'wavefronts_id']:
        if measure in df.columns:
            df.drop(df[df[measure] == -1].index, inplace=True)

    df = simplify_anesthetic_names(df)
    return df

def simplify_anesthetic_names(df):
    df['anesthetic'] = ['ketamine' if 'ketamine' in name else name
                          for name in df.anesthetic]
    return df

def direction_to_angle(dx, dy):
    return np.angle(dx + 1j*dy)

def add_angle_column(df, dir_key='direction', angle_key=None):
    if angle_key is None:
        angle_key = dir_key.replace('direction', 'angle')
    df[angle_key] = df.apply(lambda row: direction_to_angle(row[f'{dir_key}_x'],
                                                            row[f'{dir_key}_y']),
                                                            axis=1)
    return df

def color_to_map(hex_color, N):
    h = hex_color.lstrip('#')
    r, g, b = tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
    M = N+2
    vals = np.ones((M, 4))
    vals[:, 0] = np.linspace(r/256, 1, M)
    vals[:, 1] = np.linspace(g/256, 1, M)
    vals[:, 2] = np.linspace(b/256, 1, M)
    return mpl.colors.ListedColormap(vals).colors[:N]