hladnik_grewe_heterogeneity/code/plots/info_surface.py

###############################################################################
##                   information as function of delay and kernel             ##
from argparse import ArgumentError
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import matplotlib.patches as mpatches

from matplotlib import cm

from .figure_style import subfig_labelsize, subfig_labelweight, despine
from ..util import DelayType

def plot_surface(df, axis, args):
    kernels = df.kernel_sigma.unique()
    delays = df.delay.unique()
    delays = delays[delays >= 0.0]

    if args.pop_size not in df.pop_size.unique():
        raise ValueError("Requested population size %i was not analyzed! valid values are: " % args.pop_size,  df.pop_size.unique())

    if args.delay_type not in df.delay_type.unique():
        raise ArgumentError(f"Delay type argument ({args.delay_type}) is invalid in data frame {args.heterogeneous_data}! Valid options are: {list(df.delay_type.unique())}")

    # create the surface plot
    log_delays = delays.copy()
    log_delays[0] += 0.00005
    log_delays = np.log10(log_delays)
    log_kernels = np.log10(kernels)
    X, Y = np.meshgrid(log_delays, log_kernels)
    Z = np.zeros_like(X)
    E = np.zeros_like(X)
    for i, d in enumerate(delays):
        for j, k in enumerate(kernels):
            data_df = df[(df.kernel_sigma == k) & (df.pop_size == args.pop_size) & (df.delay == d)]
            Z[j, i] = np.mean(data_df.mi)
            E[j, i] = np.std(data_df.mi)
    CS = axis.contourf(X, Y, Z, cmap=cm.coolwarm, antialiased=True, alpha=0.85, levels=20)
    axis.set_xlabel(r"$\sigma_{delay}$ [ms]")
    xticklabels = list(map(lambda s: str(np.round(s*1000, 1)), delays))
    x2ticklabels = list(map(lambda s: str(np.round((s*2)/np.sqrt(12)*1000, 1)), delays))
    for i in [4, 5, 7, 8, 9, 10, 12, 13]: 
        xticklabels[i] = "" 
        x2ticklabels[i] = ""
    axis.plot([np.log10(0.000125), np.log10(0.015)], [np.log10(0.000125), np.log10(0.015)], ls="--", lw=0.5, color="k")
    axis.set_xlim([log_delays[0], log_delays[-1]])
    axis.set_xticks(log_delays, minor=True)
    axis.set_xticks([log_delays[0], np.log10(0.001), np.log10(0.01)])
    axis.set_xticklabels([0.0, 1.0, 10.0], fontsize=7, rotation=90, ha="center")
    axis.set_ylim([min(log_kernels), max(log_kernels)])
    axis.set_yticks(log_kernels)
    axis.set_yticklabels(kernels*1000, fontsize=7)
    axis.set_ylabel(r"$\sigma_{kernel}$ [ms]")
    axis.xaxis.grid(b=True, which="major", color="white", ls="--", lw=0.1)
    axis.yaxis.grid(b=True, which="major", color="white", ls="--", lw=0.1)
    cbar = plt.gcf().colorbar(CS, ax=axis)
    cbar.set_ticks(np.arange(200, 751, 100))
    cbar.ax.set_ylabel('mutual information [bit/s]')
    
    # add a rectangle for LS delays
    rect = mpatches.Rectangle([np.log10(0.0002), log_kernels[0]],
                              np.log10(0.0007) - np.log10(0.0002),
                              log_kernels[-1] - log_kernels[0], lw=1., ls="--", ec="silver", fc="None", alpha=1.0)
    axis.add_patch(rect)
    axis.text(np.log10(0.0002), log_kernels[-2], r"receptive field size lower bound", rotation=90, fontsize=5, va="top", ha="right")
    axis.text(np.log10(0.0007), log_kernels[-2], r"receptive field size upper bound", rotation=90, fontsize=5, va="top", ha="right")

    rect = mpatches.Rectangle([log_delays[0], np.log10(0.0005)],
                              log_delays[-1] - log_delays[0],
                              np.log10(0.0015) - np.log10(0.0005), lw=1., ls="--", ec="silver", fc="None", alpha=1.0)
    axis.add_patch(rect)
    axis.text(log_delays[-2], np.log10(0.0005), r"fast EPSP component", rotation=0, fontsize=5, va="top", ha="right")
    axis.text(log_delays[-2], np.log10(0.0015), r"slow EPSP component", rotation=0, fontsize=5, va="top", ha="right")


def layout_figure():
    fig = plt.figure(figsize=(3.42, 2.5))
    axis = fig.add_subplot()
    fig.subplots_adjust(left=0.175, top=0.95, right=0.95, bottom=0.2)
    return fig, axis


def performance_surface(args):
    """ The postsynaptic pyramidal cells in the ELL integrate over different population size, i.e. 
    receptive fields of different sizes (Maler 2009) and do show different EPSP time courses.

    This function creates a surface plot     

    Args:
        args (argparse.arguments): the command line arguments

    Returns:
        [type]: [description]
    """
    if not os.path.exists(args.heterogeneous_data):
        raise ValueError(f"Results DataFrame does not exist! {args.heterogeneous_data}")
    het_df = pd.read_csv(args.heterogeneous_data, sep=";", index_col=0)
    fig, axis = layout_figure()
    if "delay_type" in het_df.columns:
        plot_surface(het_df[het_df.delay_type=="gaussian"], axis, args)
    else:
        plot_surface(het_df, axis, args)
    if args.nosave:
        plt.show()
    else:
        fig.savefig(args.outfile, dpi=500)
        plt.close()


def command_line_parser(subparsers):
    parser = subparsers.add_parser("info_surface", help="Plots the mutual information as function of delay and (synpatic) kernel width.")
    parser.add_argument("-hetdf", "--heterogeneous_data", default=os.path.join("derived_data", "heterogeneous_populationcoding.csv"))
    parser.add_argument("-o", "--outfile", type=str, default=os.path.join("figures", "performance_surface.pdf"), help="The filename of the figure")
    parser.add_argument("-ps", "--pop_size", type=int, default=16, help="The population size")
    parser.add_argument("-dt", "--delay_type", type=str, default=str(DelayType.Gaussian), help="The type of distribution from which delays have been drawn. Usually a Gaussian normal distribution.")

    parser.add_argument("-n", "--nosave", action='store_true', help="No saving of the figure, just showing.")
    parser.set_defaults(func=performance_surface)