mrestimator_triallength/plt/utility.py

"""
    @Author:        F. Paul Spitzner
    @Email:         paul.spitzner@ds.mpg.de
    @Created:       2020-01-21
    @Last Modified: 2020-01-21
    @Description:
        Helper functions to load and plot the preprocessed hdf5 files containing
        correlation coefficients.
"""


import numpy as np
import mrestimator as mre
import h5py


def h5_load(filename, dsetname, raise_ex=False):
    try:
        file = h5py.File(filename, "r")
        try:
            res = file[dsetname][:]
        # maybe it is a scalar
        except ValueError:
            res = file[dsetname][()]
        file.close()
        return res
    except Exception as e:
        print(f"failed to load {dsetname} from {filename}: {e}")
        if raise_ex:
            raise e

# draw an arrow at specified side
def arrowed_spines(ax, side="bottom", **ap):

    # ap = dict(arrowstyle="->", fc="k", ec="k", connectionstyle="arc3",
    # shrinkA=0, shrinkB=0, lw=1, mutation_scale=5)
    ap.setdefault("arrowstyle", "->")
    ap.setdefault("fc", "k")
    ap.setdefault("ec", "k")
    ap.setdefault("connectionstyle", "arc3")
    ap.setdefault("shrinkA", 0)
    ap.setdefault("shrinkB", 0)
    ap.setdefault("lw", 1)
    ap.setdefault("mutation_scale", 7)

    cs = dict(xycoords="axes fraction", textcoords="axes fraction", zorder=20)

    for s in side:
        if s == "bottom":
            ax.annotate("", xy=(0, 1), xytext=(0, 0), **cs, arrowprops=ap)
        elif s == "left":
            ax.annotate("", xy=(1, 0), xytext=(0, 0), **cs, arrowprops=ap)
        elif s == "top":
            ax.annotate("", xy=(1, 1), xytext=(0, 1), **cs, arrowprops=ap)
        elif s == "right":
            ax.annotate("", xy=(1, 1), xytext=(1, 0), **cs, arrowprops=ap)

    # other option, way more complicated, modified:
    # https://stackoverflow.com/questions/33737736/matplotlib-axis-arrow-tip
    # hl = 1./20.*(xmax-xmin) / width


def plot_stdd(ax, x, y_mean, y_stdd, mode="bar", eb_ls='-', **kwargs):
    if mode=="fill":
        if "alpha" not in kwargs:
            kwargs = dict(kwargs, alpha=0.1)
        if "lw" not in kwargs:
            kwargs = dict(kwargs, lw=0)

        ax.fill_between(
            x,
            (y_mean - y_stdd),
            (y_mean + y_stdd),
            **kwargs
        )
    elif mode=="bar":
        eb = ax.errorbar(x, y_mean, yerr=y_stdd, fmt='o', markersize=0, capsize=2, capthick=0.5, lw=0.5, **kwargs)
        try:
            eb[-1][0].set_linestyle(eb_ls)
        except Exceptions as e:
            print(e)

# find_prec(arr, "312.62", ".2f")
def find_prec(arr, expr, ptrn):
    for idx, i in enumerate(arr):
        if "{i:{ptrn}}".format(i=i, ptrn=ptrn) == expr:
            return idx
    return -1


# bias corrected solution by Marriott and Pope
def rk_analytic_408(k_max, m, numsteps):
    k_range = np.arange(1, int(k_max + 1))
    return np.power(m, k_range) - 2 * k_range * np.power(m, k_range) / numsteps


def rk_analytic_407(k_max, m, numsteps):
    k_range = np.arange(1, int(k_max + 1))
    geom = np.zeros(k_max)
    for idx, k in enumerate(k_range):
        # print(f"{idx} {k}")
        # print(f"\t{np.power(m, k_range[:k]-1)}")
        geom[idx] = np.sum(np.power(m, k_range[:k] - 1))

    res = np.power(m, k_range) - 1.0 / numsteps * (
        (1 + m) * geom - 2 * k_range * np.power(m, k_range) / numsteps
    )

    return res


def tau_est_over_tau_407(tau, targetlength, k_max, fdx=0):
    # targetlength will be multiplied by tau
    m = np.exp(-1.0 / tau)
    k_range = np.arange(1, k_max + 1)
    res = np.zeros(len(targetlength))
    for idx, T in enumerate(targetlength):
        print(f"{T}/{targetlength[-1]}")
        rks = rk_analytic_407(k_max, m, T * tau)
        rks = mre.CoefficientResult(steps=k_range, coefficients=rks)
        res[idx] = mre.fit(rks, fitfunc=fitfuncs[fdx]).tau / tau
    return res


def tau_est_over_tau_408(tau, targetlength, k_max, fdx=0):
    # targetlength will be multiplied by tau
    m = np.exp(-1.0 / tau)
    k_range = np.arange(1, k_max + 1)
    res = np.zeros(len(targetlength))
    for idx, T in enumerate(targetlength):
        print(f"{T}/{targetlength[-1]}")
        rks = rk_analytic_408(k_max, m, T * tau)
        rks = mre.CoefficientResult(steps=k_range, coefficients=rks)
        res[idx] = mre.fit(rks, fitfunc=fitfuncs[fdx]).tau / tau
    return res