updates to python modules

"""

Python interface to DrSPINE

"""

import os.path

import ast

from collections import OrderedDict

def _load_sqt_bin(filename, full_bw=False):

    s = load_txt(filename, delimiter=',', cleanup=['NaN', 'Infinity'])

    sel = s[:,0]==0      # full bw data (it==0)

    if full_bw:

        return s[s[:,0]==0] # only full band width

    return s[s[:,0]>0]      # all but full band width (TOF slices)

        return s[sel] # only full band width

    return s[np.logical_not(sel)]      # all but full band width (TOF slices)

def plot_sqt_bin(filename, save=False, full_bw=False, fmt=None, label=None):

def plot_sqt_bin(filename, **kwargs):

    """

    Plot Q,tau binning

    """

    minsqt  = kwargs.pop('minsqt' ,  0.0)  # min value of sqt to show

    maxsqt  = kwargs.pop('maxsqt' ,  2.0)  # max value of sqt to show

    maxerr  = kwargs.pop('maxerr' ,  0.5)  # max value of abs(var(sqt)/sqt) to show

    maxchi  = kwargs.pop('maxchi' , 100.0) # max chi2 value

    full_bw = kwargs.pop('full_bw', False) #

    debug   = kwargs.pop('debug'  , False) # min

    save    = kwargs.pop('save'   , None)  # save figure to filename

    reverse = kwargs.pop('reverse', False)

    #marker  = kwargs.pop('maxerr', 'o')   # max value of abs(var(sqt)/sqt) to show

    #label   = kwargs.pop('label',  None)  # min

    d = _load_sqt_bin(filename, full_bw=full_bw)

    if full_bw:

        d = _load_sqt_bin(filename, True)

        plt.plot(d[:,3], d[:,4],   fmt or 'ro', label=label or 'Full BW')

    d = _load_sqt_bin(filename, False)

    plt.plot(d[:,3], d[:,4],   fmt or 'k.', label=label or 'TOF slices')

    #

    n0   = d.shape[0]

    w    = d[:,5]  # sqt

    ew   = d[:,6]  # var(sqt)

    chi2 = d[:,16] # chi2(sig)

    sel  = np.logical_and(minsqt<=w, w<=maxsqt)

    sel  = np.logical_and(sel, abs(ew/w)<maxerr)

    #sel  = np.logical_and(sel, chi2<maxchi)

    if reverse:

        sel  = np.logical_not(sel)

    d    = d[sel]

    nsel = d.shape[0]

    t    = d[:,3]

    q    = d[:,4]

    sqt  = d[:,5]

    esqt = d[:,6] #

    chi2 = d[:,16]

    wmin,  wmax  = min(sqt),  max(sqt)

    ewmin, ewmax = min(esqt), max(esqt)

    c  = (sqt - wmin)/(wmax-wmin)

    if debug:

        fig, (ax,ax1) = plt.subplots(1,2, figsize=(12,6))

        print(f"{n0} {nsel}")

        print(f"sqt({wmin},{wmax})")

        print(f"var({ewmin},{ewmax})")

        bins = np.linspace(-1,+3,201)

        ax1.scatter(sqt, abs(esqt/sqt), s=1)

        #ax1.scatter(d[:,15], d[:,16], s=10)

        ax1.set_xlabel('sqt')

        ax1.set_ylabel('var(sqt)/sqt')

        ax1.grid()

    else:

        fig, ax = plt.subplots(1,1, figsize=(8,6))

    ax.scatter(t,q, s=np.ones_like(sqt)*5,c=c, **kwargs)

    ax.set_xscale('log')

    ax.set_xlabel(r'$\tau$ [ns]')

    ax.set_ylabel(r'Q [$\AA$]')

    ax.grid(True, which='both')

    basename  = os.path.basename(filename)

    sel_label = 'Full BW' if full_bw else 'TOF slices'

    fig.suptitle(f"{basename} ({sel_label})")

    plt.xscale('log')

    plt.grid(True, which='both')

    plt.xlabel(r'$\tau$ [ns]')

    plt.ylabel(r'Q [$\AA$]')

    if full_bw:

        plt.legend(loc='best')

    if save:

        plt.savefig(filename+".pdf")

        fig.savefig(save)

def _plot_chi2_histo(bins, chi2, color='r', fmt='.', label=None):

    db   =  bins[1:]-bins[:-1]      # bin width

    Plot Q,tau binning

    """

    d  = load_sqt_bin(filename, False)

    w = abs(d[:,10]/d[:, 9]) # SQT

    #d = d[w<2]

    print(d.shape)

    t  = d[:,3]

    q  = d[:,4]

    #w = d[:,11] + d[:,13]

    w = d[:, 7] # SQT

    #w = abs(d[:,10]/d[:, 9]) # SQT

    w  = d[:,5]  # SQT

    ew = d[:,6] #

    wr  = d[:,7]  # R(SQT)

    #ewr = d[:,8] #

    ws  = d[:,9]  # F(SQT)

    #ewr = d[:,10] #

    print(w.shape)

    #nrow, ncol = d.shape

    #count = 0

    #for i,r in enumerate(d):

    #    err = abs(r[6]/r[5])

    #    if err>=1.0:

    #        count = count + 1

    #        print(f"{i:6d}/{count:6d} t={r[3]:6.2f} q={r[4]:7.4f} w={r[5]:8.3f} e={r[6]:8.3f} rel={100*err:8.2f}")

    #print(nrow, ncol, count)

    sel1 = abs(ew/w)<1.0 # relative error less than 100%

    sel2 = np.logical_and(w>=0, w<2)

    sel = np.logical_and(sel1, sel2)

    #sel = np.logical_not(sel)

    t = t[sel]

    q = q[sel]

    w = w[sel]

    wr = wr[sel]

    ws = ws[sel]

    print(min(w),  max(w))

    print(min(wr), max(wr))

    print(min(ws), max(ws))

    print(w.shape)

    wmin, wmax = min(w), max(w)

    print(wmin, wmax)

    _, axx = plt.subplots(1,1)

    bins = np.linspace(-1,+3,201)

    #axx.hist(w,  bins=bins, histtype='step', label='sqt', lw=2)

    axx.hist(wr, bins=bins, histtype='step', label='res')

    axx.hist(ws, bins=bins, histtype='step', label='sig')

    #axx.hist(ws/wr, bins=bins, histtype='step', label='ratio', color='k', lw=0.5, ls='--')

    axx.set_yscale('log')

    axx.legend()

    wmin, wmax = min(w), max(w)

    wn  = (w - wmin)/(wmax-wmin)*20

    #print(min(wn), max(wn))

    print(wmin, wmax)

    c  = (w - wmin)/(wmax-wmin)

    print(min(c), max(c))

    ax.scatter(t, q, s=wn, **kwargs)

    sqt1 = np.ma.masked_where(w <  0, np.ones_like(w)*5)

    sqt2 = np.ma.masked_where(w >= 0, np.ones_like(w)*100)

    ax.scatter(t, q, s=sqt1, marker='o', c=c, cmap='jet')

    #ax.scatter(t, q, s=sqt2, marker='v', c='k')

    ## positive

    #p = ax.scatter(t, q, s=wn, **kwargs)

    #color

    #print(help(p.set_color))

    #print(p.get_edgecolor())

    #print(p.get_edgecolors())

    #print(p.get_facecolor())

    #print(p.get_facecolors())

    ax.set_xscale('log')

    ax.grid(True, which='both')

    ax.set_xlabel(r'$\tau$ [ns]')

    ax.set_ylabel(r'Q [$\AA$]')

    d = d[sel]

    print(d.shape)

def q_auto_bin(ax, filename, **kwargs):

    "auto q binning"

def autobin(x, w, **kwargs):

    "autobin"

    xmin   = kwargs.pop('xmin', 0)

    xmax   = kwargs.pop('xmax', np.inf)

    levels = kwargs.pop('levels', None)

    smooth = kwargs.pop('smooth', False)

    edges  = np.linspace(xmin, xmax, max(len(x)//10,1001) )

    #

    if smooth:

        kde = gaussian_kde(x, weights=w)

        cdf = kde((edges[:-1]+edges[1:])/2)

    else:

        cdf, edges = np.histogram(x, weights=w, bins=edges)

    cdf = np.cumsum(cdf)/np.sum(cdf)

    res = []

    for y in levels:

        iy = np.where(cdf<y, cdf, 0)

        iy = np.argmax(iy) if len(iy) else 0

        xq = edges[iy]

        if smooth:

            if y<=0:

                xq = xmin

            elif y>=1:

                xq = xmax

            else:

                def func(xu, y0):

                    return kde.integrate_box_1d(0, xu)-y0

                xq = min(fsolve(func, xq, args=(y,)))

        res.append(xq)

    res[-1] = edges[-1]

    return res, ( edges, cdf)

def make_qbins(ax, filename, **kwargs):

    """make q-bins

    auto   = False - uniform binning between (qmin, qmax)

    auto   = True  - build a CDF so that q bins have the same statistics

    smooth = True (implies auto=True) create a "smooth" CDF

    returs edges of q bins

    """

    #pylint: disable=too-many-locals

    histtype = kwargs.pop('histtype', 'step')

    bins     = kwargs.pop('bins', 11)

    auto     = kwargs.pop('auto', False)

    qmin     = kwargs.pop('qmin', None)

    qmax     = kwargs.pop('qmax', None)

    auto     = kwargs.pop('auto', False) # auto

    smooth   = kwargs.pop('smooth',  False)

    weights  = kwargs.pop('weights', None) #'sqt')

    #

    histtype = kwargs.pop('histtype', 'step')

    orientation= kwargs.pop('orientation','horizontal')

    d = load_sqt_bin(filename, False)

    q = d[:, 4]

        qmax = max(q)

    if auto:

        tmpbin=len(q)//10

        ax1 = ax.twiny()

        if smooth:

            edges = np.linspace(qmin, qmax, tmpbin)

            kde   = gaussian_kde(q, weights=w)

            hist  = kde(edges[:-1])

            def xcdf(x):

                return kde.integrate_box_1d(-np.inf, x)

        else:

            hist, edges = np.histogram(q, weights=w, bins=np.linspace(qmin, qmax, tmpbin))

        cdf = np.cumsum(hist)/np.sum(hist)

        levels=np.linspace(0,1,bins+1)

        bins, (edges,cdf) = autobin(q, w, xmin=qmin, xmax=qmax, levels=levels, smooth=smooth)

        ax1.plot(cdf, edges[:-1], lw=2, color='m')

        q_edges = []

        qave = np.average(q)

        for x in np.linspace(0,1,bins+1):

            ix = np.where(cdf<x, cdf,0)

            if len(ix):

                ix = np.argmax(ix)

            else:

                ix = 0

            xq = edges[ix]

            q_edges.append(xq)

            #if auto:

            #    if 0<x and x<1:

            #        def func(_q, _x0):

            #            return abs(xcdf(_q)-_x0)

            #        res = fsolve(func, qave, args=(x,))

            #        print(x, xq, *res)

            ax.axhline(xq, lw=1, ls='--', color='k')

            ax1.axvline(x, lw=1, ls='--', color='k')

        q_edges[-1] = edges[-1]

        bins = q_edges

        for x in levels:

            ax1.axvline(x, lw=0.5, ls='--', color='k')

        ax1.set_xticklabels([])

    else:

        bins = np.linspace(qmin, qmax, bins+1)

    _, q_edges, _ = ax.hist(q, weights=w, bins=bins, histtype=histtype, **kwargs)

    _, q_edges, _ = ax.hist(q, weights=w, bins=bins,

                    orientation=orientation, histtype=histtype, **kwargs)

    ax.grid(True, which='both')

    ax.set_xlabel('Intensity (arb)')

    return q_edges

#def tau_auto_bin(ax, filename, **kwargs):

#    "auto tau binning"

#    #pylint: disable=too-many-locals

#    histtype = kwargs.pop('histtype', 'step')

#    bins     = kwargs.pop('bins', 11)

#    auto     = kwargs.pop('auto', False)

def main():

    "the main"

    parser = argparse.ArgumentParser(description='q binning')

    parser.set_defaults(filename='./last_sqt_bin.dat', nbins=5, uniform=False,

                        qmin=None, qmax=None, weights='sqt', smooth=False)

    parser.set_defaults(filename='./last_sqt_bin.dat', nbins=11,

                        qmin=None, qmax=None, weights='sqt',

                        auto=False,smooth=False, show=True)

    parser.add_argument('filename', metavar='file', help='sqt_bin filename')

    parser.add_argument('--q-min', dest='qmin', type=float)

    parser.add_argument('--q-max', dest='qmax', type=float)

    parser.add_argument('--num-bins', '-n', dest='nbins', type=int)

    parser.add_argument('--uniform',  '-u',  dest='uniform', action='store_true',

                        help='create uniform prob. distribution')

    parser.add_argument('--auto',     '-a', dest='auto', action='store_true',

                        help='create uniform division based on a cumulative distribution function')

    parser.add_argument('--smooth'  ,  '-s', dest='smooth',  action='store_true',

                        help='"smooth" prob. distribution')

                        help='make a "smooth" cumulative distribution functio ')

    parser.add_argument('--no-show',  '-N', dest='show',  action='store_false',

                        help='do not show any plots')

    args = parser.parse_args()

    if args.smooth:

    plot_sqt_map(ax1, args.filename)

    q_edges = q_auto_bin(ax2, args.filename, auto=args.uniform,

                weights=args.weights, smooth=args.smooth,

    q_edges = make_qbins(ax2, args.filename,

                auto=args.auto, smooth=args.smooth, weights=args.weights,

                qmin=args.qmin, qmax=args.qmax, bins=args.nbins,

                color='b', linewidth=2, orientation='horizontal')

                color='b', linewidth=2)

    print("histo q custom", end=' ')

    for q in q_edges:

        ax1.axhline(q, color='b', ls='--', lw=1)

        ax2.axhline(q, color='b', ls='--', lw=1)

        print(f"{q:.4f}/A", end=' ')

    if args.uniform:

        print(" ! uniform binning")

    if args.auto:

        print(" ! auto calculated")

    else:

        print(" ! auto calc ")

        print(" ! uniform binning")

    plt.tight_layout()

    if args.show:

        plt.show()

if __name__ == "__main__":