tracer_analysis.py

from __future__ import division, print_function
import os
import sys
import time
import inspect
from copy import deepcopy
from collections import namedtuple, defaultdict
import math
import numpy as np
from numpy import linalg
import networkx as nx
import community as community_louvain
from matplotlib import cm
import matplotlib.pyplot as plt
from matplotlib import rcParams
from matplotlib.patches import Patch
from scipy.spatial.distance import pdist
from scipy.cluster.hierarchy import linkage
import seaborn as sns
import pandas as pd
from roi_groups import ROI_GROUPS, ROI_MAPPINGS
from custom_atlas_config import CustomAtlasConfig
from tracer_connectivity import TracerConnectivity
rcParams.update({'figure.autolayout': True})


class TracerAnalysis:
    """
    TracerAnalysis instance is constructed from a list of TracerConnectivity
    objects, which are expected to have the same tracer directionality, cell
    quantification mode and custom atlas

    self.output_dir:
    directory to store output. default to current working directory

    self.ara_level_subset:
    a subset of ara levels to include in the analysis. defaults to None, where
    all ara levels encountered are included.
    this attribute should only be modified through self.set_ara_level_subset,
    which calls each TracerConnectivity instances's set_ara_level_subset method,
    and rebuild self.data.
    when building its tracer_connectivity attribute, TracerConnectivity
    instances treat ara levels not in the subset as if their csv files contains
    all 0s in their overlap and atlas_only fields. actual data for all levels
    are still maintained in  TracerConnectivity instances's level_connectivities
    attribute. the fraction normalization for TracerAnalysis is performed
    against total overlap counts from ara levels within the subset.

    self.roi_exclusion_subset:
    not implemented

    self.data:
    dict with keys count', 'area', 'fraction' and density'.
    values are dataframes with MultiIndex (injection_site, case_tracer) and
    roi names as columns. the dataframes are constructed from TracerConnectivity
    instances. whenever self.ara_level_subset is set, self.data is rebuilt by
    self.build_data()

    self.selection_parameters
    set by self.set_selection_parameters as named tuple
    the fields in self.selection_parameters defines how to apply selection and
    grouping on self.data to generate self.selected data

    self.selected_data
    dict with keys count', 'area', 'fraction' and density'.
    values are dataframes in self.data with self.selection_parameters applied

    outputs are generated to self.output_dir with an autogenerated name
    including a time stamp for disambiguity. the attributes of the instance
    that are relevant to the output is written to a meta info txt file with
    matching name.
    """
    def __init__(self, tracer_connectivities, description=None):
        # analysis is initialized from a list of TracerConnectivity objects
        self.tracer_connectivities = tracer_connectivities
        self.description = 'description: {}'.format(description)
        # attributes that need to be identical in all TracerConnectivity instances
        self.anterograde = self.cell_count = self.custom_atlas = None
        self._validate_input()

        self.analysis_name = 'None'
        self.output_dir = os.getcwd()

        self.ara_level_subset = None
        # placeholder
        self.roi_exclusion_subset = None

        self.custom_atlas_config = CustomAtlasConfig(custom_atlas=self.custom_atlas)

        self.data = {}
        self._max_density = None
        self.multi_index = self.get_indices()
        self.build_data()

        self.selected_data = {}
        self.selection_parameters = None

        self.color_map = {}
        self.set_cmap('injection_site')
        self.distance_type = None

    def _validate_input(self):
        for tracer_connectivity in self.tracer_connectivities:
            if not isinstance(tracer_connectivity, TracerConnectivity):
                raise TypeError("TracerConnectivity object expected, "
                                "{} encountered".format(type(tracer_connectivity)))
        self.anterograde = self.tracer_connectivities[0].anterograde
        self.cell_count = self.tracer_connectivities[0].cell_count
        self.custom_atlas = self.tracer_connectivities[0].custom_atlas
        roi_names_hemi = self.tracer_connectivities[0].tracer_connectivity.index
        for tracer_connectivity in self.tracer_connectivities[1:]:
            try:
                assert tracer_connectivity.custom_atlas == self.custom_atlas
                assert tracer_connectivity.anterograde == self.anterograde
                assert tracer_connectivity.cell_count == self.cell_count
                assert (tracer_connectivity.tracer_connectivity.index == roi_names_hemi).all()
            except AssertionError as e:
                print(e.message)
                print('inconsistent meta data fields between {} channel {} and {} channel {}'
                      .format(self.tracer_connectivities[0].case_name, self.tracer_connectivities[0].channel_number,
                              tracer_connectivity.case_name, tracer_connectivity.channel_number))

    def unique_injection_sites(self):
        return self.multi_index.get_level_values('injection_site').unique()

    def case_tracers(self, series_insensitive=False):
        if not series_insensitive:
            return self.multi_index.get_level_values('case_tracer')
        else:
            return list(set([tracer_connectivity.case_tracer(series_insensitive=True)
                             for tracer_connectivity in self.tracer_connectivities]))

    def get_indices(self):
        """
        multiple index (injection_site, case_tracer)
        the index uses tracer instead of channel number, since tracer more
        meaningfully identifies data from different series of a same case.
        it's possible for phal data in A series to be imaged as ch2, while
        phal data in B series to be imaged as ch3
        """
        injection_sites = [tracer_connectivity.injection_site
                           for tracer_connectivity in self.tracer_connectivities]
        case_tracers = ['{}_{}'.format(tracer_connectivity.case_name, tracer_connectivity.tracer)
                         for tracer_connectivity in self.tracer_connectivities]
        return pd.MultiIndex.from_arrays([injection_sites, case_tracers],
                                         names=['injection_site', 'case_tracer'])

    def _get_series_insensitive_indices(self):
        """
        index (injection_site:case_tracer_series_insensitive)
        """
        injection_sites = [tracer_connectivity.injection_site
                           for tracer_connectivity in self.tracer_connectivities]
        case_tracers = [tracer_connectivity.case_tracer(series_insensitive=True)
                        for tracer_connectivity in self.tracer_connectivities]
        return pd.Index(['{}:{}'.format(injection_site, case_tracer) for
                         injection_site, case_tracer in zip(injection_sites, case_tracers)],
                        name='injection_site:case_tracer')

    def build_data(self):
        tracer_count_dfs = []
        tracer_area_dfs = []
        for tracer_connectivity in self.tracer_connectivities:
            tracer_df = tracer_connectivity.tracer_connectivity
            tracer_count_dfs.append(tracer_df.loc[:, ['overlap']].transpose())
            tracer_area_dfs.append(tracer_df.loc[:, ['area']].transpose())
        self.data['count'] = pd.concat(tracer_count_dfs).set_index(self.multi_index)
        self.data['area'] = pd.concat(tracer_area_dfs).set_index(self.multi_index)
        # divide count by row sum
        self.data['fraction'] = self.data['count'].div(self.data['count'].sum(axis=1), axis=0)
        self.data['density'] = self.data['fraction'] / (self.data['area'] + 1e-9)
        self._max_density = np.max(self.data['density'].values)
        self.data['density'] /= self._max_density

    # returns hemisphere roi names
    def _selected_roi_names(self):
        # if rois and roi_groups are both None, return all roi names
        if not self.selection_parameters.roi_names and not self.selection_parameters.roi_group_names:
            return set(self.data['count'].columns)
        elif not self.selection_parameters.roi_group_names:
            roi_names = self.selection_parameters.roi_names
        # if roi_groups is not None, return a set of grey matter region names based on roi_groups
        else:
            roi_names = set()
            for roi_group_name in self.selection_parameters.roi_group_names:
                assert roi_group_name in ROI_GROUPS
                roi_group = ROI_GROUPS[roi_group_name]
                if isinstance(roi_group, set):
                    roi_names |= roi_group
                else:
                    roi_names |= set([key for key in roi_group.keys()])
        roi_names_hemi = set()
        for roi_name in roi_names:
            roi_names_hemi |= {'{}_i'.format(roi_name), '{}_c'.format(roi_name)}
        return sorted(list(roi_names_hemi))

    # apply roi selection and aggregation
    def _merge_rois(self):
        if len(self.selected_data) == 0:
            raise ValueError('data selection should have been performed by _merge_case_tracers')
        selected_roi_names = self._selected_roi_names()
        for data_type in ['count', 'area', 'fraction']:
            self.selected_data[data_type] = self.selected_data[data_type].loc[:, list(selected_roi_names)]
            # merge roi
            if self.selection_parameters.roi_aggregate_rule is not None:
                if isinstance(self.selection_parameters.roi_aggregate_rule, str):
                    roi_aggregate_rule = ROI_MAPPINGS[self.selection_parameters.roi_aggregate_rule]
                else:
                    assert isinstance(self.selection_parameters.roi_aggregate_rule, dict)
                    roi_aggregate_rule = self.selection_parameters.roi_aggregate_rule
                roi_mapping = {}
                for roi_name, mapped_roi_name in roi_aggregate_rule.items():
                    roi_mapping['{}_i'.format(roi_name)] = '{}_i'.format(mapped_roi_name)
                    roi_mapping['{}_c'.format(roi_name)] = '{}_c'.format(mapped_roi_name)
                self.selected_data[data_type] = self.selected_data[data_type].groupby(by=roi_mapping, axis='columns').sum()
        self.selected_data['density'] = self.selected_data['fraction'] / (self.selected_data['area'] + 1e-9)
        self.selected_data['density'] /= self._max_density

    # this function performs dataframe row merge based on case_aggregate_rule
    # case merging should occur before roi merging to ensure correct mle estimates
    # when pooling data from multiple rows. only allow specifying index ordering
    # under injection site merging rule
    def _merge_case_tracers(self, aggregated_index_orderings=None):
        # no row merging occuring
        if self.selection_parameters.case_aggregate_rule is None:
            self.selected_data = deepcopy(self.data)
        # if merging rule is known, merge accordingly
        else:
            if self.selection_parameters.case_aggregate_rule == 'injection_site':
                for data_type in ['count', 'area']:
                    self.selected_data[data_type] = self.data[data_type].groupby(level='injection_site').sum()
                    if aggregated_index_orderings:
                        assert set(self.selected_data[data_type].index.values) == set(aggregated_index_orderings)
                        self.selected_data[data_type] = self.selected_data[data_type].reindex(index=aggregated_index_orderings)
            if self.selection_parameters.case_aggregate_rule == 'combine_series':
                for data_type in ['count', 'area']:
                    # exchange self.multi_index with series insensitive index
                    # make the index a column and group on it
                    # split to multi index with levels injection_site and case_tracer
                    self.selected_data[data_type] = \
                        self.data[data_type].set_index(self._get_series_insensitive_indices())\
                                            .reset_index()\
                                            .groupby(by='injection_site:case_tracer').sum()
                    self.selected_data[data_type].index = self.selected_data[data_type].index.str.split(pat=':', expand=True)
                    self.selected_data[data_type].index.names = ['injection_site', 'case_tracer']
            self.selected_data['fraction'] = self.selected_data['count'].div(self.selected_data['count'].sum(axis=1), axis=0)
            self.selected_data['density'] = self.selected_data['fraction'] / (self.selected_data['area'] + 1e-9)
            self.selected_data['density'] /= self._max_density

    def _select_hemispheres(self):
        if self.selection_parameters.hemispheres == 'i+c':
            return
        hemisphere_cols = [col for col in self.selected_data['fraction'].columns if col.endswith(self.selection_parameters.hemispheres)]
        for data_type in ['fraction', 'density', 'count', 'area']:
            self.selected_data[data_type] = self.selected_data[data_type].loc[:, hemisphere_cols]

    def select_data(self, hemispheres='i+c', roi_names=None, roi_group_names=None, roi_aggregate_rule=None,
                    case_aggregate_rule=None, aggregated_index_orderings=None):
        """
        populate self.selected_data
        key: 'count', 'fraction', density' and 'area'
        value: dataframe
        :param hemispheres:
               'i+c': select data from both hemisphere. 'i': select data from ipsilateral. 'c': select data from contralateral
        :param roi_names:
               the rois to include in the selected data, applying to both hemisphere
        :param roi_group_names:
               if given, roi_names argument is ignored. all rois within the groups are included
        :param roi_aggregate_rule:
               dictionary that map roi names to destination name
        :param case_aggregate_rule: None, 'injection_site' or 'combine_series'
        :param aggregated_index_orderings:
               the row index for selected_data dataframes. should be a reordering of result of case aggregation
        :return: None
        """
        self.set_selection_parameters(hemispheres=hemispheres, roi_names=roi_names, roi_group_names=roi_group_names,
                                      roi_aggregate_rule=roi_aggregate_rule, case_aggregate_rule=case_aggregate_rule)
        self._merge_case_tracers(aggregated_index_orderings=aggregated_index_orderings)
        self._merge_rois()
        self._select_hemispheres()

    def filter_data(self, data_type, threshold=None, threshold_quantile=None):
        """
        filter self.selected_data[data_type]
        if threshold and threshold_quantile are both not None, threshold is used
        :param data_type: 'count', 'fraction' or 'density'
        :param threshold: if all values in a column is less than threshold, the column is discarded
        :param threshold_quantile: set threshold to the given quantile of the entire data frame
        :return:
        """
        if len(self.selected_data) == 0:
            raise ValueError('no data selected')
        df = self.selected_data[data_type]
        if not threshold and not threshold_quantile:
            return df
        elif not threshold:
            assert 0 <= threshold_quantile <= 1
            threshold = np.quantile(df.values, threshold_quantile)
        # columns where all entries are less than threshold
        roi_is_sparse = (df < threshold).all(axis=0)
        # remove sparse columns
        return df.loc[:, ~roi_is_sparse]

    def append_total_columns(self):
        """
        calculate TOTAL_i and/or TOTAL_c columns and append to self.selected_data
        """
        if len(self.selected_data) == 0:
            raise ValueError('no data selected')
        # prevent adding total columns multiple times
        self.drop_total_columns()
        for data_type in ['count', 'area', 'fraction']:
            df = self.selected_data[data_type]
            df_total = df.groupby(lambda x: 'TOTAL_i' if x.endswith('i') else 'TOTAL_c', axis=1).sum()
            self.selected_data[data_type] = pd.concat([df, df_total], axis=1)
        df_density_total = self.selected_data['fraction'].iloc[:, [-2, -1]] / self.selected_data['area'].iloc[:, [-2, -1]]
        self.selected_data['density'] = pd.concat([self.selected_data['density'], df_density_total], axis=1)
        self.selected_data['density'] /= self.selected_data['density'].max()

    def drop_total_columns(self):
        if len(self.selected_data) == 0:
            raise ValueError('no data selected')
        for data_type in ['count', 'area', 'fraction', 'density']:
            if 'TOTAL_c' in self.selected_data[data_type].columns:
                self.selected_data[data_type] = self.selected_data[data_type].drop(labels='TOTAL_c', axis=1)
            if 'TOTAL_i' in self.selected_data[data_type].columns:
                self.selected_data[data_type] = self.selected_data[data_type].drop(labels='TOTAL_i', axis=1)

    # return tidy dataframe for self.selected_data[data_type]
    # with columns (case_tracer), injection_site, roi_name, hemisphere, data_type
    # the last column (named data_type) contains value of fraction/density/count data
    def tidy_data(self, data_type):
        if len(self.selected_data) == 0:
            raise ValueError('no data selected')
        df = self.selected_data[data_type].reset_index()
        if self.selection_parameters.case_aggregate_rule == 'injection_site':
            id_vars = ['injection_site']
        elif self.selection_parameters.case_aggregate_rule == 'combine_series' or self.selection_parameters.case_aggregate_rule is None:
            id_vars = ['injection_site', 'case_tracer']
        else:
            raise ValueError('illegal case aggregate rule')
        df = pd.melt(df, id_vars=id_vars, var_name='roi_name', value_name=data_type)
        df = df.assign(hemisphere=df.roi_name.apply(lambda x: 'ipsilateral' if x[-1] == 'i' else 'contralateral'),
                       roi_name=df.roi_name.str[0:-2])
        return df

    # return Graph built from tidy data frame. note that at the moment contralateral data is discarded
    def get_graph(self, data_type):
        self.drop_total_columns()
        # using tidy_data to build a NetworkX.Graph
        df = self.tidy_data(data_type)
        # drop contralateral data
        df = df.loc[df['hemisphere'] == 'ipsilateral', :]
        # drop columns that do not denote nodes
        df = df.drop(labels='hemisphere', axis=1)
        if not self.selection_parameters.case_aggregate_rule == 'injection_site':
            df['case_tracer_injection_site'] = df.case_tracer.str.cat(df.injection_site, sep='_')
            df = df.drop(labels='injection_site', axis=1)
            id_var = 'case_tracer_injection_site'
        else:
            id_var = 'injection_site'
        source = id_var if self.anterograde else 'roi_name'
        target = 'roi_name' if self.anterograde else id_var
        graph = nx.from_pandas_edgelist(df, source, target, edge_attr=[data_type])
        return graph

    def community_detection(self, data_type):
        graph = self.get_graph(data_type)
        return community_louvain.best_partition(graph, weight=data_type)

    @staticmethod
    def _heatmap_row_labels(df):
        # create row labels for the plot
        if isinstance(df.index, pd.Index):
            row_labels = df.index.values
        # MultiIndex
        else:
            case_tracers = df.index.get_level_values('case_tracer')
            injection_sites = df.index.get_level_values('injection_site')
            row_labels = ['{}:{}'.format(case_tracer, injection_site)
                          for case_tracer, injection_site in zip(case_tracers, injection_sites)]
        return row_labels

    @staticmethod
    def _clip_df(df, vis_ceil_val):
        if vis_ceil_val is None:
            return df
        if vis_ceil_val < 0:
            vis_ceil_val = df.values.max() * np.abs(vis_ceil_val)
        return df.clip(lower=0, upper=vis_ceil_val)

    @staticmethod
    def save_fig_maximize(figure, save_path, cg=None):
        # retrieve the gui manager of the figure, has it maximize the gui display
        # use a non blocking show() call, flush this gui event back to the figure.
        # cleanup and exit display with close()
        plot_backend = plt.get_backend()
        mng = getattr(figure.canvas, 'manager')
        if plot_backend == 'TkAgg':
            mng.resize(*mng.window.maxsize())
        elif plot_backend == 'wxAgg':
            mng.frame.Maximize(True)
        elif plot_backend == 'Qt4Agg':
            mng.window.showMaximized()
        plt.show(block=False)
        figure.canvas.flush_events()
        if not cg:
            figure.savefig(save_path)
        else:
            cg.savefig(save_path)
        plt.close()

    def plot_data_matrix(self, data_type, threshold=None, threshold_quantile=None,
                         vis_ceil_val=None, title=None):
        """
        plot previously selected data as a 2d color mesh. columns of the mesh
        are hemisphere roi names, rows of the mesh are case_tracer:injection_site
        :param data_type: 'count', 'fraction' or 'density'
        :param threshold: see filter_data
        :param threshold_quantile: see filter_data
        :param vis_ceil_val: if positive, clip data with upper bound as vis_ceil_val.
                             if negative, calculate upper bound as df.max() * abs(vis_ceil_val)
        :param title: title of the plot
        :param show_fig:
        :param save_dir:
        :return:
        """
        if len(self.selected_data) == 0:
            raise ValueError('no data selected')

        df = self.filter_data(data_type, threshold=threshold, threshold_quantile=threshold_quantile)
        df = TracerAnalysis._clip_df(df, vis_ceil_val)

        # note: pcolor plots last row first.
        # reverse dataframe rows before plotting
        df = df.reindex(index=df.index[::-1])
        fig, ax = plt.subplots()
        plt.pcolormesh(df, cmap="Oranges", edgecolors='powderblue', linewidth=0.01)

        # set column labels for the plot
        xticks = range(0, df.shape[1])
        xticks = [xtick + 0.5 for xtick in xticks]
        ax.set_xticks(xticks, minor=False)
        ax.set_xticklabels(df.columns)
        plt.xticks(rotation=90)

        # create row labels for the plot
        row_labels = TracerAnalysis._heatmap_row_labels(df)
        yticks = range(0, df.shape[0])
        yticks = [ytick + 0.5 for ytick in yticks]
        ax.set_yticks(yticks)
        ax.set_yticklabels(row_labels)

        if title is None:
            title = data_type
        plt.title(title)

        plt.colorbar(fraction=0.01)

        output_basename = self.output_basename()
        save_path = os.path.join(self.output_dir, '{}.png'.format(output_basename))
        TracerAnalysis.save_fig_maximize(plt.gcf(), save_path)
        self.write_meta_info(output_basename)

    def _pair_distance(self, v0, v1):
        """ modified cosine for fractions, cosine for densities """
        epsilon = 1e-6
        cosine_d = 1 - np.dot(v0, v1) / (linalg.norm(v0, 2) * linalg.norm(v1, 2) + epsilon)
        r_d = max(np.sum(v0), np.sum(v1)) / (min(np.sum(v0), np.sum(v1)) + epsilon)
        if self.distance_type == 'modified_cosine':
            return cosine_d * r_d
        else:
            return cosine_d

    def cluster_2d(self, data_type, threshold=None, threshold_quantile=None, vis_ceil_val=None, title=None):
        if len(self.selected_data) == 0:
            raise ValueError('no data selected')

        self.distance_type = 'modified_cosine' if data_type == 'fraction' else 'cosine'
        df = self.filter_data(data_type, threshold=threshold, threshold_quantile=threshold_quantile)

        # distances between cases and regions
        case_distance_vec = pdist(df, self._pair_distance)
        case_distance_vec = np.clip(case_distance_vec, a_min=0, a_max=np.iinfo('i').max)
        case_distance_vec[np.isnan(case_distance_vec)] = 0
        Z_injections = linkage(case_distance_vec, method='complete')
        Z_injections = np.clip(Z_injections, a_min=0, a_max=np.inf)
        # transpose to calculate distance between rois
        region_distance_vec = pdist(df.transpose(), self._pair_distance)
        region_distance_vec = np.clip(region_distance_vec, a_min=0, a_max=np.iinfo('i').max)
        region_distance_vec[np.isnan(region_distance_vec)] = 0
        Z_rois = linkage(region_distance_vec, method='complete')
        Z_rois = np.clip(Z_rois, a_min=0, a_max=np.inf)

        df = TracerAnalysis._clip_df(df, vis_ceil_val)
        cg = sns.clustermap(df, row_linkage=Z_injections, col_linkage=Z_rois, xticklabels=True, cmap='copper', robust=False)
        if title is None:
            title = data_type
        cg.fig.suptitle('2D clustermap: {}'.format(title))
        plt.setp(cg.ax_heatmap, autoscalex_on=True, rasterized=False)
        plt.setp(cg.ax_heatmap.yaxis.get_majorticklabels(), rotation=0)
        plt.setp(cg.ax_heatmap.xaxis.get_ticklabels(), rotation=90, fontsize=9)

        output_basename = self.output_basename()
        save_path = os.path.join(self.output_dir, '{}.png'.format(output_basename))
        # save with the savefig method of the cluster grid (from seaborn documentation)
        # this avoids incorrectly clipping figure content
        TracerAnalysis.save_fig_maximize(cg.fig, save_path, cg=cg)
        self.write_meta_info(output_basename)

    def plot_data_scatter(self, data_type, free_ymax=False, use_global_ymax=False, hue_order=None):
        # get tidy data
        df = self.tidy_data(data_type)
        # rank the order for each roi to appear in the scatter plot
        # roi's are ranked by the maximum value across all rows
        ranked_max_values = df.groupby(by='roi_name')[data_type].max().sort_values(ascending=False)
        ranked_rois = ranked_max_values.index

        # set cmap to per injection site
        self.set_cmap('injection_site')

        n_regions = len(ranked_rois)
        n_cols = n_rows = 3
        global_ymax = df[data_type].max()
        if self.selection_parameters.hemispheres == 'i+c':
            order = ['ipsilateral', 'contralateral']
        else:
            order = ['ipsilateral'] if self.selection_parameters.hemispheres == 'i' else ['contralateral']
        for roi_rank in range(0, n_regions, n_rows * n_cols):
            subplot_id = 0
            fig, axes = plt.subplots(nrows=n_rows, ncols=n_cols, sharey=False, sharex=False)
            # retrieve max y value appearing in the figure
            fig_ymax = ranked_max_values[roi_rank]
            ymax = global_ymax if use_global_ymax else fig_ymax
            axes = axes.flatten()
            for i, roi in enumerate(ranked_rois[roi_rank:]):
                df_roi = df.loc[df['roi_name'] == roi]
                axes[i].set_xlim([0, 5])
                axes[i].set_xticks([1, 4])
                axes[i].axvline(x=3, ymin=0, ymax=1, color='black', linestyle='dashed')
                # for seaborn 0.9.1, if y is passed as series instead of numpy array, the legend is incorrectly generated
                sns.stripplot(x='hemisphere', y=df_roi[data_type].values, hue='injection_site',
                              data=df_roi, jitter=False, ax=axes[i], palette=self.color_map, dodge=True, size=7.5,
                              order=order, hue_order=hue_order)
                if not free_ymax:
                    axes[i].set_ylim([0, ymax * 1.1])
                tracer_preamble = 'Projection to' if self.anterograde else 'Input from'
                axes[i].set_ylabel('{} {}'.format(tracer_preamble, roi), fontsize=14)
                axes[i].set_xlabel('')
                axes[i].set_xticklabels(order, fontsize=14)
                subplot_id += 1
                if i == n_rows * n_cols - 1:
                    break
            if subplot_id < n_rows * n_cols:
                for j in range(subplot_id, n_rows * n_cols):
                    axes[j].set_visible(False)
            figure_id = roi_rank // (n_rows * n_cols)
            output_basename = self.output_basename()
            save_path = os.path.join(self.output_dir, '{}_{}.png'.format(output_basename, str(figure_id).zfill(3)))
            TracerAnalysis.save_fig_maximize(plt.gcf(), save_path)
            self.write_meta_info(output_basename)

    def plot_fraction_by_ara_level(self):
        """
        plot connectivity fraction distribution by ara level. this function does
        not use selected data. it does combine TracerConnectivity data with
        identical series insensitive case_tracer
        :param injection_color_codes:
        """
        n_injection_sites = len(self.unique_injection_sites())
        n_figure_rows = int(math.ceil(n_injection_sites / 2))
        n_figure_cols = 2
        fig, axes = plt.subplots(nrows=n_figure_rows, ncols=n_figure_cols, sharex=True, sharey=True)
        if n_figure_rows * n_figure_cols == 1:
            axes = [axes]
        else:
            axes = axes.flatten()

        # set cmap to case_tracer_series_insensitive
        self.set_cmap('case_tracer_series_insensitive')

        # case_tracer: (injection_site, injection_site_level)
        case_tracer_injection_sites = {}
        # injection_site: subplot_id
        injection_subplot_ids = dict(zip(self.unique_injection_sites(),
                                         range(n_injection_sites)))
        legend_elements = defaultdict(list)

        df_ara_levels = []
        for tracer_connectivity in self.tracer_connectivities:
            case_tracer = tracer_connectivity.case_tracer(series_insensitive=True)
            case_tracer_injection_sites[case_tracer] = \
                (tracer_connectivity.injection_site, tracer_connectivity.injection_site_level)
            # sum of overlap at each ara level in the TracerConnectivity instance
            # columns: area_level, overlap, case_tracer
            df_ara_level = tracer_connectivity.tracer_connectivity_by_ara_levels()\
                                              .groupby('ara_level')['overlap'].sum().reset_index()
            df_ara_level['case_tracer'] = case_tracer
            df_ara_levels.append(df_ara_level)
        df_ara_levels = pd.concat(df_ara_levels)
        # do not display ara levels not in the subset
        if self.ara_level_subset is not None:
            df_ara_levels = df_ara_levels.loc[df_ara_levels.ara_level.isin(self.ara_level_subset)]

        fraction_max = 0
        for case_tracer in self.case_tracers(series_insensitive=True):
            injection_site = case_tracer_injection_sites[case_tracer][0]
            injection_site_level = case_tracer_injection_sites[case_tracer][1]
            color = self.color_map[case_tracer]
            ax_id = injection_subplot_ids[injection_site]
            df_case_tracer = df_ara_levels.loc[df_ara_levels['case_tracer'] == case_tracer]
            df_case_tracer = df_case_tracer.assign(fraction=lambda df: df.overlap / df.overlap.sum())
            fraction_max = max(fraction_max, df_case_tracer.fraction.max())
            axes[ax_id].scatter(df_case_tracer.ara_level, df_case_tracer.fraction, color=color, edgecolors=None)
            if injection_site_level:
                axes[ax_id].axvline(x=injection_site_level, color=color, alpha=0.6, linewidth=3, linestyle='--')
            legend_elements[ax_id].append(Patch(facecolor=color, edgecolor=color, label=injection_site))

        for ax_id in range(n_injection_sites):
            if ax_id % 2 == 0:
                axes[ax_id].set_ylabel('Fractions', fontsize=16)
            if ax_id // n_figure_cols == n_figure_rows - 1:
                axes[ax_id].set_xlabel('ARA Levels', fontsize=16)
            axes[ax_id].set_ylim([-0.01, 1.1 * fraction_max])
            axes[ax_id].legend(handles=legend_elements[ax_id])
        for ax_id in range(n_injection_sites, n_figure_rows * n_figure_cols):
            axes[ax_id].set_visible(False)
        output_basename = self.output_basename()
        save_path = os.path.join(self.output_dir, '{}.png'.format(output_basename))
        TracerAnalysis.save_fig_maximize(plt.gcf(), save_path)
        self.write_meta_info(output_basename)

    def plot_ara_levels(self, group_by='injection_site'):
        if group_by not in { 'injection_site', 'case_tracer', 'case_tracer_series_insensitive' }:
            raise ValueError('group_by argument {} not understood'.format(group_by))

        self.set_cmap(group_by)
        group_by_column = 'injection_site' if group_by == 'injection_site' else 'case_tracer'

        df_levels = []
        for tracer_connectivity in self.tracer_connectivities:
            if group_by == 'injection_site':
                group_by_var = tracer_connectivity.injection_site
            elif group_by == 'case_tracer':
                group_by_var = tracer_connectivity.case_tracer()
            else:
                group_by_var = tracer_connectivity.case_tracer(series_insensitive=True)
            df_levels.append(pd.DataFrame(data={
                'ara_level': tracer_connectivity.ara_levels(),
                group_by_column: group_by_var,
                'hue': tracer_connectivity.injection_site,
            }))
        df_levels = pd.concat(df_levels)
        sns.stripplot(x='ara_level', y=group_by_column, hue='hue', data=df_levels)
        output_basename = self.output_basename()
        save_path = os.path.join(self.output_dir, '{}.png'.format(output_basename))
        TracerAnalysis.save_fig_maximize(plt.gcf(), save_path)
        self.write_meta_info(output_basename)

    def write_data_long(self, data_type):
        if len(self.selected_data) == 0:
            raise ValueError('no data selected')
        df = self.selected_data[data_type]
        output_basename = '{}_{}_long_{}'.format(self.analysis_name, data_type, TracerAnalysis.time_stamp())
        out_csv_path = os.path.join(self.output_dir, '{}.csv'.format(output_basename))
        df.to_csv(out_csv_path)
        self.write_meta_info(output_basename)

    def write_data_tidy(self, data_type):
        df = self.tidy_data(data_type)\
                 .reindex(columns=['injection_site', 'case_tracer', 'roi_name', 'hemisphere', 'fraction'])\
                 .sort_values(by=['injection_site', 'case_tracer', 'roi_name', 'hemisphere'])
        output_basename = '{}_{}_tidy_{}'.format(self.analysis_name, data_type, TracerAnalysis.time_stamp())
        out_csv_path = os.path.join(self.output_dir, '{}.csv'.format(output_basename))
        df.to_csv(out_csv_path)
        self.write_meta_info(output_basename)

    ############################################################################
    #                                  setters                                 #
    ############################################################################
    def set_analysis_name(self, analysis_name):
        self.analysis_name = analysis_name

    def set_output_dir(self, output_dir):
        self.output_dir = output_dir
        if not os.path.isdir(self.output_dir):
            os.makedirs(self.output_dir)
        assert os.path.isdir(self.output_dir)

    def set_cmap(self, color_map_by):
        """
        sets the colors to be used in plots. if number of colors required is
        low, use qualitative cmaps tab10. otherwise use gist_ncar
        :return:
        """
        assert color_map_by in { 'injection_site', 'case_tracer', 'case_tracer_series_insensitive' }
        if color_map_by == 'injection_site':
            id_vars = self.unique_injection_sites()
        elif color_map_by == 'case_tracer':
            id_vars = self.case_tracers()
        else:
            id_vars = self.case_tracers(series_insensitive=True)
        n_colors = len(id_vars)
        if n_colors <= 10:
            cmap = cm.get_cmap('tab10', 10)
        else:
            cmap = cm.get_cmap('gist_ncar', n_colors)
        self.color_map.clear()
        # https://matplotlib.org/3.1.1/tutorials/colors/colormap-manipulation.html#sphx-glr-tutorials-colors-colormap-manipulation-py
        # https://matplotlib.org/_modules/ matplotlib/colors.html#LinearSegmentedColormap
        # ColorMap instances are callable (with [0, 1) floating number or integer)
        for i in range(n_colors):
            self.color_map[id_vars[i]] = cmap(i)

    def set_ara_level_subset(self, ara_level_subset):
        if self.ara_level_subset != ara_level_subset:
            self.ara_level_subset = ara_level_subset
        for tracer_connectivity in self.tracer_connectivities:
            tracer_connectivity.set_ara_level_subset(ara_level_subset=self.ara_level_subset)
        self.build_data()
        self.selected_data.clear()

    def set_roi_exclusion_subset(self, roi_exclusion_subset):
        """
        not implemented for now. unclear whether this should be allowed to
        differ for each TracerConnectivity, since one of the use case is to
        exclude rois very near injection site, and each case may have different
        exclusions
        """

    def set_selection_parameters(self, hemispheres='i+c', roi_names=None, roi_group_names=None,
                                 roi_aggregate_rule=None, case_aggregate_rule=None):
        SelectionParameters = namedtuple('SelectionParameters', ['hemispheres', 'roi_names', 'roi_group_names', 'roi_aggregate_rule', 'case_aggregate_rule'])
        if hemispheres not in {'i+c', 'i', 'c'}:
            raise ValueError('hemispheres: {} not understood'.format(hemispheres))
        if isinstance(roi_names, str):
            roi_names = [roi_names]
        elif isinstance(roi_group_names, str):
            roi_group_names = [roi_group_names]
        if roi_names is not None and not isinstance(roi_names, (list, set)):
            raise ValueError('roi_names parameter not understood')
        if roi_group_names is not None and not isinstance(roi_group_names, list):
            raise ValueError('roi_group_names parameter not understood')
        if isinstance(roi_aggregate_rule, str):
            if roi_aggregate_rule not in ROI_MAPPINGS:
                raise ValueError('roi_aggregate_rule {} not understood'.format(roi_aggregate_rule))
        elif roi_aggregate_rule is not None:
            if not isinstance(roi_aggregate_rule, dict):
                raise ValueError('roi_aggregate_rule must be string or a dict')
        if case_aggregate_rule not in {None, 'injection_site', 'combine_series'}:
            raise ValueError('case_aggregate_rule: {} not understood'.format(case_aggregate_rule))
        self.selection_parameters = SelectionParameters(hemispheres, roi_names, roi_group_names, roi_aggregate_rule, case_aggregate_rule)

    ############################################################################
    #                            utility functions                             #
    ############################################################################
    # string representation of selection parameters. None parameters are skipped
    def selection_parameter_str(self):
        hemispheres_str, roi_names_str, roi_group_names_str, \
        roi_aggregate_rule_str, case_aggregate_rule_str = '', '', '', '', ''
        if self.selection_parameters.hemispheres == 'i + c':
            hemispheres_str = '# hemispheres = ipsilateral + contralateral'
        else:
            hemispheres_str = '# hemispheres = {}'.format('ipsilateral' if self.selection_parameters.hemispheres == 'i' else 'contralateral')
        if self.selection_parameters.roi_names:
            roi_names_str = '# roi_names={}'.format('+'.join(self.selection_parameters.roi_names))
        if self.selection_parameters.roi_group_names:
            roi_group_names_str = '# roi_group_names = {}'.format('+'.join(self.selection_parameters.roi_group_names))
        if not self.selection_parameters.roi_names and not self.selection_parameters.roi_group_names:
            roi_names_str = '# roi_names=all'
        if self.selection_parameters.roi_aggregate_rule:
            roi_aggregate_rule_str = '# roi_aggregate_rule = {}'.format(self.selection_parameters.roi_aggregate_rule)
        if self.selection_parameters.case_aggregate_rule:
            case_aggregate_rule_str = '# case_aggregate_rule = {}'.format(self.selection_parameters.case_aggregate_rule)
        parameter_strs = [hemispheres_str, roi_names_str, roi_group_names_str, roi_aggregate_rule_str, case_aggregate_rule_str]
        return '\n'.join([parameter_str for parameter_str in parameter_strs if parameter_str])

    def ara_level_subset_str(self):
        if self.ara_level_subset is None:
            return '# ara_level_subset = None'
        else:
            return '# ara_level_subset = {}'.format(list(self.ara_level_subset))

    def analysis_name_str(self):
        return '# analysis_name = {}'.format(self.analysis_name)

    def hemispheres_str(self):
        return '# hemispheres = {}'.format(self.selection_parameters.hemispheres)

    def class_name_str(self):
        return '# {}'.format(self.__class__.__name__)

    def __str__(self):
        return '\n'.join([self.class_name_str(), self.analysis_name_str(),
                          self.selection_parameter_str(), self.ara_level_subset_str(), self.description])

    @staticmethod
    # A calls B, B calls get_caller_name. get_caller_name returns A
    def _get_caller_name():
        caller_frame_info = inspect.stack()[2]
        # retrive calling function's name
        # caller_frame_info differ between python major versions
        if sys.version < '3.5':
            caller_name = caller_frame_info[3]
        else:
            caller_name = caller_frame_info.function
        return caller_name

    @staticmethod
    def time_stamp():
        return time.strftime("%y%m%d-%H_%M_%S", time.localtime())

    # return output basename and its associated time stamp
    def output_basename(self):
        caller_name = TracerAnalysis._get_caller_name()
        time_stamp = TracerAnalysis.time_stamp()
        return '_'.join([self.analysis_name, caller_name, time_stamp])

    def write_meta_info(self, basename):
        meta_info_path = os.path.join(self.output_dir, '{}_meta_info.txt'.format(basename))
        with open(meta_info_path, 'w') as f:
            f.write('{}\n'.format(self.__str__()))