hladnik_grewe_heterogeneity/code/plots/intro_figure.py

###############################################################################
##                            introductory plot                              ##
import os
import numpy as np
import pandas as pd
import scipy.stats as stats
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

from .figure_style import subfig_labelsize, subfig_labelweight, despine


def plot_linregress(ax, x, y, bonferroni_factor, color="tab:red", label="", feature=""):
    xRange = np.asarray([x.min(), x.max()])

    slope, intercept, *params = stats.linregress(x, y)
    p_star = params[1] * bonferroni_factor
    if p_star > 1:
        p_star = 1
    l = r'%s r: %.2f, %s: %.2f' % ("%s, " % label if len(label) > 0 else "", params[0], "p" if bonferroni_factor == 1 else r"p", params[1] if bonferroni_factor == 1 else p_star)
    ax.plot(xRange, intercept + slope * xRange, color=color, label=l, ls="-")
    ax.legend(loc='upper left', frameon=False, ncol=1,fontsize=6, bbox_to_anchor=(0., 1.25))
    print("%s: minimum: %.2f, maximum: %.2f, mean: %.2f, std: %.2f" % (feature, np.min(y), np.max(y), np.mean(y), np.std(y)))
    return slope, intercept


def plotMeanProperty(ax, bins, x, y):
    my = np.zeros(len(bins)-1)
    sy = np.zeros(len(bins)-1)
    for i, (start, end) in enumerate(zip(bins[:-1], bins[1:])):
        yy = y[(x >= start) & (x < end)]
        if len(yy) > 0:
            my[i] = np.mean(yy)
            sy[i] = np.std(yy)
        else:
            my[i] = np.nan
            sy[i] = np.nan
    centers = bins[:-1]+(bins[1]-bins[0])/2

    ax.plot(centers, my, '--', lw=0.5, color='black')
    ax.fill_between(centers, my+sy, my-sy, color='black', alpha=0.1)

    return my


def add_model_sketch(axis):
    if os.path.exists(os.path.join("figures", "delay_problem_sketch.png")):
        img = mpimg.imread(os.path.join("figures", "delay_problem_sketch.png"))
        axis.imshow(img)
        axis.set_xticklabels([])
        axis.set_yticklabels([])
        despine(axis, ["top", "bottom", "left", "right"], True)


def add_fish_sketch(axis):
    if os.path.exists(os.path.join("figures", "fishsketch.png")):
        img = mpimg.imread(os.path.join("figures", "fishsketch_tight.png"))
        axis.imshow(img)
        despine(axis, ["top", "bottom", "left", "right"], True)
    else:
        raise FileNotFoundError("Missing fishsketch.png file.")


def plot_delay(axis, args):
    if not os.path.exists(args.baseline_frame):
        raise ValueError(f"Results data frame not found! {args.baseline_frame})")
    df = pd.read_csv(args.baseline_frame, sep=";", index_col=0)
    absolute_positions = df.receptor_pos_absolute.values
    delay_times = df.phase_time.values * 1000
    axis.scatter(absolute_positions, delay_times, color="tab:blue", s=1)
    m, n = plot_linregress(axis, absolute_positions, delay_times, bonferroni_factor=4, feature="phase")
    # plotMeanProperty(axis, np.linspace(0, 120, 12), absolute_positions, delay_times)
    axis.set_yticks([0.0, 1.0, 2.0])
    axis.set_xlim(0.0, 120.0)
    despine(axis, ["top", "right"], False)
    axis.set_xlabel("receptor position [mm]", fontsize=8)
    axis.set_ylabel("delay [ms]", fontsize=8)


def layout_figure():
    fig = plt.figure(figsize=(3.42, 1.4)) # , constrained_layout=True)
    axes = []
    fig_grid = (8, 8)
    axes.append(plt.subplot2grid(fig_grid, (0, 0), 8, 4))
    axes.append(plt.subplot2grid(fig_grid, (0, 5), 2, 3))
    axes.append(plt.subplot2grid(fig_grid, (2, 5), 3, 3))

    axes[0].text(0.05, 1.00, "A", transform=axes[0].transAxes, ha="center",fontsize=subfig_labelsize, fontweight=subfig_labelweight)
    axes[1].text(-0.35, 0.9, "B", transform=axes[1].transAxes, ha="center",fontsize=subfig_labelsize, fontweight=subfig_labelweight)

    fig.subplots_adjust(left=-0.0, top=0.95, bottom=0.00, right=0.95)
    return fig, axes


def introductory_figure(args):
    fig, axes = layout_figure()
    add_model_sketch(axes[0])
    #add_distributions(axes[1], axes[2], axes[3])
    add_fish_sketch(axes[1])
    plot_delay(axes[-1], args)
    if args.nosave:
            plt.show()
    else:
        fig.savefig(args.outfile, dpi=500)
        plt.close()


def command_line_parser(subparsers):
    default_bf = os.path.join("derived_data", "figure2_baseline_properties.csv")
    parser = subparsers.add_parser("intro_figure", help="Introductory plot (figure 1).")
    parser.add_argument("-bf", "--baseline_frame", default=default_bf,
                        help=f"Full file name of a  CSV table readable with pandas that holds the baseline properties and positions (defaults to {default_bf}).")
    parser.add_argument("-o", "--outfile", type=str, default=os.path.join("figures","delay_problem.pdf"), help="The filename of the figure")
    parser.add_argument("-n", "--nosave", action='store_true', help="no saving of the figure, just showing") 

    parser.set_defaults(func=introductory_figure)