text:
# Description: Local ancillary functions.
#
# Author:      André Palóczy Filho
# E-mail:      paloczy@gmail.com

import numpy as np
import matplotlib.pyplot as plt
from scipy.special import erfinv
from netCDF4 import num2date
from mpl_toolkits.basemap import Basemap
from pygamman import gamma_n as gmn
from pygeodesy import Datums, VincentyError
from pygeodesy.ellipsoidalVincenty import LatLon as LatLon
from pygeodesy.sphericalNvector import LatLon as LatLon_sphere
from pandas import Series

def seasonal_avg(t, F):
    tmo = np.array([ti.month for ti in t])
    ftmo = [tmo==mo for mo in range(1, 13)]

    return np.array([F[ft].mean() for ft in ftmo])

def seasonal_std(t, F):
    """
    USAGE
    -----
    F_seasonal = seasonal_std(t, F)

    Calculates the seasonal standard deviation of variable F(t).
    Assumes 't' is a 'datetime.datetime' object.
    """
    tmo = np.array([ti.month for ti in t])
    ftmo = [tmo==mo for mo in range(1, 13)]

    return np.array([F[ft].std() for ft in ftmo])

def deseason(t, F):
    Fssn = seasonal_avg(t, F)
    nyears = int(t.size/12)
    aux = np.array([])
    for n in range(nyears):
        aux = np.concatenate((aux, Fssn))

    return F - aux

def blkavgt(t, x, every=2):
    """
    Block-averages a variable x(t). Returns its block average
    and standard deviation and new t axis.
    """
    nt = t.size
    t = np.array([tt.toordinal() for tt in t])
    tblk, xblk, xblkstd = np.array([]), np.array([]), np.array([])
    for i in range(every, nt+every, every):
        xi = x[i-every:i]
        tblk = np.append(tblk, t[i-every:i].mean())
        xblk = np.append(xblk, xi.mean())
        xblkstd = np.append(xblkstd, xi.std())

    tblk = num2date(tblk, units='days since 01-01-01')
    return tblk, xblk, xblkstd

def blkavg(x, y, every=2):
    """
    Block-averages a variable y(x). Returns its block average
    and standard deviation and new x axis.
    """
    nx = x.size
    xblk, yblk, yblkstd = np.array([]), np.array([]), np.array([])
    for i in range(every, nx+every, every):
        yi = y[i-every:i]
        xblk = np.append(xblk, x[i-every:i].mean())
        yblk = np.append(yblk, yi.mean())
        yblkstd = np.append(yblkstd, yi.std())

    return xblk, yblk, yblkstd

def blksum(x, y, every=2):
    """
    Sums a variable y(x) in blocks. Returns its block average
    and new x axis.
    """
    nx = x.size
    xblk, yblksum = np.array([]), np.array([])
    for i in range(every, nx+every, every):
        xblk = np.append(xblk, x[i-every:i].mean())
        yblksum = np.append(yblksum, y[i-every:i].sum())

    return xblk, yblksum

def stripmsk(arr, mask_invalid=False):
    if mask_invalid:
        arr = np.ma.mask_invalid(arr)
    if np.ma.isMA(arr):
        msk = arr.mask
        arr = arr.data
        arr[msk] = np.nan

    return arr

def rcoeff(x, y):
        """
        USAGE
        -----
        r = rcoeff(x, y)

        Computes the Pearson correlation coefficient r between series x and y.

        References
        ----------
        e.g., Thomson and Emery (2014),
        Data analysis methods in physical oceanography,
        p. 257, equation 3.97a.
        """
        x,y = map(np.asanyarray, (x,y))

        # Sample size.
        assert x.size==y.size
        N = x.size

        # Demeaned series.
        x = x - x.mean()
        y = y - y.mean()

        # Standard deviations.
        sx = x.std()
        sy = y.std()

        ## Covariance between series. Choosing unbiased normalization (N-1).
        Cxy = np.sum(x*y)/(N-1)

        ## Pearson correlation coefficient r.
        r = Cxy/(sx*sy)

        return r

def near(x, x0, npts=1, return_index=False):
    """
    USAGE
    -----
    xnear = near(x, x0, npts=1, return_index=False)

    Finds 'npts' points (defaults to 1) in array 'x'
    that are closest to a specified 'x0' point.
    If 'return_index' is True (defauts to False),
    then the indices of the closest points are
    returned. The indices are ordered in order of
    closeness.
    """
    x = list(x)
    xnear = []
    xidxs = []
    for n in range(npts):
        idx = np.nanargmin(np.abs(np.array(x)-x0))
        xnear.append(x.pop(idx))
        if return_index:
            xidxs.append(idx)
    if return_index: # Sort indices according to the proximity of wanted points.
        xidxs = [xidxs[i] for i in np.argsort(xnear).tolist()]
    xnear.sort()

    if npts==1:
        xnear = xnear[0]
        if return_index:
            xidxs = xidxs[0]
    else:
        xnear = np.array(xnear)

    if return_index:
        return xidxs
    else:
        return xnear

def near2(x, y, x0, y0, npts=1, return_index=False):
    """
    USAGE
    -----
    xnear, ynear = near2(x, y, x0, y0, npts=1, return_index=False)

    Finds 'npts' points (defaults to 1) in arrays 'x' and 'y'
    that are closest to a specified '(x0, y0)' point. If
    'return_index' is True (defauts to False), then the
    indices of the closest point(s) are returned.

    Example
    -------
    >>> x = np.arange(0., 100., 0.25)
    >>> y = np.arange(0., 100., 0.25)
    >>> x, y = np.meshgrid(x, y)
    >>> x0, y0 = 44.1, 30.9
    >>> xn, yn = near2(x, y, x0, y0, npts=1)
    >>> print("(x0, y0) = (%f, %f)"%(x0, y0))
    >>> print("(xn, yn) = (%f, %f)"%(xn, yn))
    """
    x, y = map(np.array, (x, y))
    shp = x.shape

    xynear = []
    xyidxs = []
    dx = x - x0
    dy = y - y0
    dr = dx**2 + dy**2
    for n in range(npts):
        xyidx = np.unravel_index(np.nanargmin(dr), dims=shp)
        if return_index:
            xyidxs.append(xyidx)
        xyn = (x[xyidx], y[xyidx])
        xynear.append(xyn)
        dr[xyidx] = np.nan

    if npts==1:
        xynear = xynear[0]
        if return_index:
            xyidxs = xyidxs[0]

    if return_index:
        return xyidxs
    else:
        return xynear

def xy2dist(x, y, cyclic=False, datum='WGS84'):
    """
    USAGE
    -----
    d = xy2dist(x, y, cyclic=False, datum='WGS84')
    """
    if datum is not 'Sphere':
        xy = [LatLon(y0, x0, datum=Datums[datum]) for x0, y0 in zip(x, y)]
    else:
        xy = [LatLon_sphere(y0, x0) for x0, y0 in zip(x, y)]
    d = np.array([xy[n].distanceTo(xy[n+1]) for n in range(len(xy)-1)])

    return np.append(0, np.cumsum(d))

def lon180to360(lon):
        """
        Converts longitude values in the range [-180,+180]
        to longitude values in the range [0,360].
        """
        lon = np.asanyarray(lon)
        return (lon + 360.0) % 360.0

def lon360to180(lon):
        """
        Converts longitude values in the range [0,360]
        to longitude values in the range [-180,+180].
        """
        lon = np.asanyarray(lon)
        return ((lon + 180.) % 360.) - 180.

def rot_vec(u, v, angle=-45, degrees=True):
        """
        USAGE
        -----
        u_rot,v_rot = rot_vec(u,v,angle=-45.,degrees=True)

        Returns the rotated vector components (`u_rot`,`v_rot`)
        from the zonal-meridional input vector components (`u`,`v`).
        The rotation is done using the angle `angle` positive counterclockwise
        (trigonometric convention). If `degrees` is set to `True``(default),
        then `angle` is converted to radians.
        is

        Example
        -------
        >>> from matplotlib.pyplot import quiver
        >>> from ap_tools.utils import rot_vec
        >>> u = -1.
        >>> v = -1.
        >>> u2,v2 = rot_vec(u,v, angle=-30.)
        """
        u,v = map(np.asanyarray, (u,v))
        if degrees:
                angle = angle*np.pi/180. # Degrees to radians.

        u_rot = +u*np.cos(angle) + v*np.sin(angle) # Usually the across-shore component.
        v_rot = -u*np.sin(angle) + v*np.cos(angle) # Usually the along-shore component.

        return u_rot, v_rot

def angle_isobath(xiso, yiso, cyclic=True):
    R = 6371000.0 # Mean radius of the earth in meters (6371 km), from gsw.constants.earth_radius.
    deg2rad = np.pi/180. # [rad/deg]

    if cyclic: # Add cyclic point.
        xiso = np.append(xiso, xiso[0])
        yiso = np.append(yiso, yiso[0])

    # From the coordinates of the isobath, find the angle it forms with the
    # zonal axis, using points k+1 and k.
    shth = yiso.size-1
    theta = np.zeros(shth)
    for k in range(shth):
        dyk = R*(yiso[k+1]-yiso[k])
        dxk = R*(xiso[k+1]-xiso[k])*np.cos(yiso[k]*deg2rad)
        theta[k] = np.arctan2(dyk,dxk)

    xisom = 0.5*(xiso[1:] + xiso[:-1])
    yisom = 0.5*(yiso[1:] + yiso[:-1])

    return xisom, yisom, theta

def fmt_isobath(cs, fontsize=8, fmt='%g', inline=True, inline_spacing=7, manual=True, **kw):
	"""
	Formats the labels of isobath contours. `manual` is set to `True` by default,
	but can be `False`, or a tuple/list of tuples with the coordinates of the labels.
	All options are passed to plt.clabel().
	"""
	isobstrH = plt.clabel(cs, fontsize=fontsize, fmt=fmt, inline=inline, \
                          inline_spacing=inline_spacing, manual=manual, **kw)
	for ih in range(0, len(isobstrH)): # Appends 'm' for meters at the end of the label.
		isobstrh = isobstrH[ih]
		isobstr = isobstrh.get_text()
		isobstr = isobstr.replace('-','') + ' m'
		isobstrh.set_text(isobstr)

def gamman(Sp, T, p, x, y):
    assert Sp.shape==T.shape
    n = np.size(p)
    skel = Sp*np.nan
    Spmsk = np.ma.masked_invalid(Sp).mask
    Tmsk = np.ma.masked_invalid(T).mask
    gammamsk = np.logical_or(Tmsk, Spmsk)
    Sp[Spmsk] = 35
    T[Tmsk] = 0
    gm, dgl, dgh = gmn(Sp, T, p, n, x, y)
    # gm, dgl, dgh = gmn.gamma_n(Sp, T, p, n, x, y)
    gm[gammamsk] = np.nan

    return gm, dgl, dgh

def bmap_antarctica(ax, resolution='h'):
	"""
	Full Antartica basemap (Polar Stereographic Projection).
	"""
	m = Basemap(boundinglat=-60,
		        lon_0=60,
	            projection='spstere',
	            resolution=resolution,
	            ax=ax)

	m.fillcontinents(color='0.9', zorder=9)
	m.drawcoastlines(zorder=10)
	m.drawmapboundary(zorder=-9999)
	m.drawmeridians(np.arange(-180, 180, 20), linewidth=0.15, labels=[1, 1, 1, 1], zorder=12)
	m.drawparallels(np.arange(-90, -50, 5), linewidth=0.15, labels=[0, 0, 0, 0], zorder=12)

	return m

def near(x, x0, npts=1, return_index=False):
    """
    USAGE
    -----
    xnear = near(x, x0, npts=1, return_index=False)

    Finds 'npts' points (defaults to 1) in array 'x'
    that are closest to a specified 'x0' point.
    If 'return_index' is True (defauts to False),
    then the indices of the closest points are
    returned. The indices are ordered in order of
    closeness.
    """
    x = list(x)
    xnear = []
    xidxs = []
    for n in range(npts):
        idx = np.nanargmin(np.abs(np.array(x)-x0))
        xnear.append(x.pop(idx))
        if return_index:
            xidxs.append(idx)
    if return_index: # Sort indices according to the proximity of wanted points.
        xidxs = [xidxs[i] for i in np.argsort(xnear).tolist()]
    xnear.sort()

    if npts==1:
        xnear = xnear[0]
        if return_index:
            xidxs = xidxs[0]
    else:
        xnear = np.array(xnear)

    if return_index:
        return xidxs
    else:
        return xnear

def UVz2iso(U, V, hisopyc, z):
	ny, nx = hisopyc.shape
	Uisopyc = np.nan*np.ones((ny,nx))
	Visopyc = np.nan*np.ones((ny,nx))
	for j in range(ny):
		print("Row %d of %d"%(j+1,ny))
		for i in range(nx):
			if np.isnan(hisopyc[j,i]):
				continue
			else:
				Uisopyc[j,i] = np.interp(hisopyc[j,i], z, U[:,j,i])
				Visopyc[j,i] = np.interp(hisopyc[j,i], z, V[:,j,i])

	return Uisopyc, Visopyc

def isopyc_depth(dens0, z, isopyc=1027.75, dzref=1.):
    """
    USAGE
    -----
    hisopyc = isopyc_depth(z, dens0, isopyc=1027.75)

    Calculates the spatial distribution of the depth of a specified isopycnal 'isopyc'
    (defaults to 1027.75 kg/m3) from a 2D density section rho0 (in kg/m3) with shape
    (nz,ny,nx) and a 1D depth array 'z' (in m) with shape (nz).

    'dzref' is the desired resolution for the refined depth array (defaults to 1 m) which
    is generated for calculating the depth of the isopycnal. The smaller 'dzref', the smoother
    the resolution of the returned isopycnal depth array 'hisopyc'.
    """
    dens0, z = map(np.array, (dens0, z))
    if not np.all(np.diff(z)>0):
        z = np.flipud(z)
        dens0 = np.flipud(dens0)
    if dens0.ndim==2:
        nz, nx = dens0.shape
    else:
        nz = dens0.size
        nx = 1
    zref = np.arange(z.min(), z.max()+dzref, dzref)

    if np.ma.isMaskedArray(dens0):
        dens0 = np.ma.filled(dens0, np.nan)

    hisopyc = np.nan*np.ones((nx))
    for i in range(nx):
        if nx==1:
            dens0i = dens0
        else:
            dens0i = dens0[:,i]

        cond1 = np.logical_or(isopyc<np.nanmin(dens0i), np.nanmax(dens0i)<isopyc)
        if np.logical_or(cond1, np.isnan(dens0i).all()):
            continue
        else:
            dens0ref = np.interp(zref, z, dens0i) # Refined density profile.
            fz = near(dens0ref, isopyc, return_index=True)
            try:
                hisopyc[i] = zref[fz]
            except ValueError:
                print("Warning: More than 1 (%d) nearest depths found. Using the median of the depths for point (i=%d)."%(fz.sum(), i))
                hisopyc[i] = np.nanmedian(zref[fz])

    return hisopyc

def isopyc_depth2(z, dens0, isopyc=1027.75, dzref=1.):
    """
    USAGE
    -----
    hisopyc = isopyc_depth(z, dens0, isopyc=1027.75)

    Calculates the spatial distribution of the depth of a specified isopycnal 'isopyc'
    (defaults to 1027.75 kg/m3) from a 3D density array rho0 (in kg/m3) with shape
    (nz,ny,nx) and a 1D depth array 'z' (in m) with shape (nz).

    'dzref' is the desired resolution for the refined depth array (defaults to 1 m) which
    is generated for calculating the depth of the isopycnal. The smaller 'dzref', the smoother
    the resolution of the returned isopycnal depth array 'hisopyc'.
    """
    z, dens0 = map(np.asanyarray, (z, dens0))
    ny, nx = dens0.shape[1:]
    if not np.all(np.diff(z>0)):
        z = np.flipud(z)
        dens0 = np.flipud(dens0)
    zref = np.arange(z.min(), z.max(), dzref)

    if np.ma.isMaskedArray(dens0):
        dens0 = np.ma.filled(dens0, np.nan)

    hisopyc = np.nan*np.ones((ny,nx))
    for j in range(ny):
        print("Row %d of %d"%(j+1,ny))
        for i in range(nx):
            dens0ij = dens0[:,j,i]
            if np.logical_or(np.logical_or(isopyc<np.nanmin(dens0ij), np.nanmax(dens0ij)<isopyc), np.isnan(dens0ij).all()):
                continue
            else:
                dens0ref = np.interp(zref, z, dens0ij) # Refined density profile.
                dens0refn = near(dens0ref, isopyc)
                fz=dens0ref==dens0refn
                try:
                    hisopyc[j,i] = zref[fz]
                except ValueError:
                    print("Warning: More than 1 (%d) nearest depths found. Using the median of the depths for point (j=%d,i=%d)."%(fz.sum(), j, i))
                    hisopyc[j,i] = np.nanmedian(zref[fz])

    return hisopyc

def montecarlo_gamman(sp_mean, spSE, t_mean, tSE, p0, z0, x0, y0, wanted_isoneutrals, dzref,
                      nmc=1e4, alpha=0.95, nbins=100, plot_cdf=True, verbose=True):
    """Using block-averaged lon, lat (x0, y0)."""
    nmc = int(nmc)
    ngamma = len(wanted_isoneutrals)
    nz = z0.size
    if not np.all(np.diff(z0)>=0):
        z0 = np.flipud(z0)
        p0 = np.flipud(p0)
        sp_mean = np.flipud(sp_mean)
        spSE = np.flipud(spSE)
        t_mean = np.flipud(t_mean)
        tSE = np.flipud(tSE)
    z00 = z0.copy()
    # Mean neutral density profile for this segment.
    gamman_MEANSEG, _, _ = gamman(sp_mean, t_mean, p0, x0, y0)
    # Calculate derived variable (the isoneutral depth) with the synthetic random data.
    hgamman_mc = np.empty((ngamma, nmc))*np.nan
    hgamman_errs = np.empty(ngamma)*np.nan
    for n in range(nmc):
        if verbose:
            print("Simulated Monte Carlo profile %d of %d"%(n+1,nmc))
        SP_mc = spSE*np.random.randn(nz) + sp_mean
        T_mc = tSE*np.random.randn(nz) + t_mean
        gamman_mc, _, _ = gamman(SP_mc, T_mc, p0, x0, y0)
        # gamman_mc = np.ma.masked_invalid(gamman_mc)
        fgud = ~np.isnan(gamman_mc)
        gamman_mc = gamman_mc[fgud]
        z0 = z00[fgud]
        for ngm in range(ngamma):
            ison = wanted_isoneutrals[ngm]
            isodpth = isopyc_depth(gamman_mc, z0, isopyc=ison, dzref=dzref)
            # print(isodpth)
            hgamman_mc[ngm, n] = isodpth

    # Calculate PDF and CDF of the simulated data.
    hgamman_MEANSEG = np.empty(ngamma)*np.nan
    gamman_MEANSEG = gamman_MEANSEG[fgud]
    for ngm in range(ngamma):
        ison = wanted_isoneutrals[ngm]
        hgamman_mcm = hgamman_mc[ngm, :]
        if np.all(np.isnan(hgamman_mcm)):
            print("Skipping %.2f kg/m3 isoneutral (Not found in any of the simulated profiles)."%ison)
            continue
        # Mean depth of this isoneutral for this segment.
        hgamman_MEANSEG[ngm] = isopyc_depth(gamman_MEANSEG, z0, isopyc=ison, dzref=dzref)
        fig, ax = plt.subplots()
        # Remove the mean from the absolute value of the simulated depths
        # so the error is a +- (depth interval).
        hgamman_mcm = hgamman_mcm[~np.isnan(hgamman_mcm)]
        hgamman_mcm = hgamman_mcm - hgamman_MEANSEG[ngm]
        hgamman_mcm = np.abs(hgamman_mcm)
        hlb, hrb = hgamman_mcm.min(), hgamman_mcm.max()
        hgamman_hist, bins, _ = ax.hist(hgamman_mcm, bins=nbins, density=True,
                                        range=(hlb, hrb), color='b', histtype='bar')

        # Calculate and plot CDF.
        dbin = bins[1:] - bins[:-1]
        cdf = np.cumsum(hgamman_hist*dbin) # CDF [unitless].
        cdf = np.insert(cdf, 0, 0)
        fci_alpha = near(cdf, alpha, return_index=True)

        # Value of the Monte Carlo derived variable associated with the alpha*100 percentile.
        hgamman_err = bins[fci_alpha]

        if plot_cdf:
                ax2 = ax.twinx()
                ax2.plot(bins, cdf, 'k', linewidth=3.0)
                # ax2.plot(binm, cdf, 'k', linewidth=3.0)
                ax2.axhline(y=alpha, linewidth=1.5, linestyle='dashed', color='grey')
                ax2.axvline(x=hgamman_err, linewidth=1.5, linestyle='dashed', color='grey')
                ax.set_ylabel(r'Probability density', fontsize=18, fontweight='black')
                ax2.set_ylabel(r'Cumulative probability', fontsize=18, fontweight='black')
                ax.set_xlabel(r'Departure from mean isoneutral depth [m]', fontsize=18, fontweight='black')
                ytks = ax2.get_yticks()
                ytks = np.append(ytks, alpha)
                ytks.sort()
                ytks = ax2.set_yticks(ytks)
                fig.canvas.draw()
                plt.show()
        else:
                plt.close()

        # Error is symmetric about zero, because the
        # derived variable is assumed to be normal.
        hgamman_err = hgamman_err/2.
        if verbose:
            print("The Monte Carlo error for the derived variable is +- %.2f (alpha=%.2f)"%(hgamman_err, alpha))
        hgamman_errs[ngm] = hgamman_err
    # hgamman_err is the envelope containing 100*'alpha' % of the values
    # of the simulated derived variable. So the error bars
    # are +-hgamman_err/2 (symmetric about each data point).

    return hgamman_MEANSEG, hgamman_err, hgamman_mc

def crosscorr(x, y, nblks, maxlags=0, overlap=0, onesided=False, verbose=False):
    """
    Lag-N cross correlation averaged with Welch's Method.
    Parameters
    ----------
    x, y     : Arrays of equal length.
    nblks    : Number of blocks to average cross-correlation.
    maxlags  : int, default (0) calculates te largest possible number of lags,
               i.e., the number of points in each chunk.
    overlap  : float, fraction of overlap between consecutive chunks. Default 0.
    onesided : Whether to calculate the cross-correlation only at
               positive lags (default False). Has no effect if
               x and y are the same array, in which case the
               one-sided autocorrelation function is calculated.

    Returns
    ----------
    crosscorr : float array.
    """
    if x is y:
        auto = True
    else:
        auto = False
    x, y = np.array(x), np.array(y)
    nx, ny = x.size, y.size
    assert x.size==y.size, "The series must have the same length"

    nblks, maxlags = int(nblks), int(maxlags)
    ni = int(nx/nblks)               # Number of data points in each chunk.
    dn = int(round(ni - overlap*ni)) # How many indices to move forward with
                                     # each chunk (depends on the % overlap).

    if maxlags==0:
        if verbose:
            print("Maximum lag was not specified. Accomodating it to block size (%d)."%ni)
        maxlags = ni
    elif maxlags>ni:
        if verbose:
            print("Maximum lag is too large. Accomodating it to block size (%d)."%ni)
        maxlags = ni

    if onesided:
        lags = range(maxlags+1)
    else:
        lags = range(-maxlags, maxlags+1)

    # Array that will receive cross-correlation of each block.
    xycorr = np.zeros(len(lags))

    n=0
    il, ir = 0, ni
    while ir<=nx:
        xn = x[il:ir]
        yn = y[il:ir]

        # Calculate cross-correlation for current block up to desired maximum lag - 1.
        xn, yn = map(Series, (xn, yn))
        xycorr += np.array([xn.corr(yn.shift(periods=lagn)) for lagn in lags])

        il+=dn; ir+=dn
        n+=1

    # pandas.Series.corr(method='pearson') -> pandas.nanops.nancorr() ...
    # -> pandas.nanops.get_corr_function() -> np.corrcoef -> numpy.cov(bias=False as default).
    # So np.corrcoef() returns the UNbiased correlation coefficient by default
    # (i.e., normalized by N-k instead of N).

    xycorr /= n    # Divide by number of blocks actually used.
    ncap = nx - il # Number of points left out at the end of array.

    if verbose:
        print("")
        if ncap==0:
            print("No data points were left out.")
        else:
            print("Left last %d data points out (%.1f %% of all points)."%(ncap,100*ncap/nx))
        print("Averaged %d blocks, each with %d lags."%(n,maxlags))
        if overlap>0:
            print("Intended %d blocks, but could fit %d blocks, with"%(nblks,n))
            print('overlap of %.1f %%, %d points per block.'%(100*overlap,dn))
        print("")

    lags = np.array(lags)
    if auto:
        fo = np.where(lags==0)[0][0]
        xycorr[fo+1:] = xycorr[fo+1:] + xycorr[:fo]
        lags = lags[fo:]
        xycorr = xycorr[fo:]

    fgud=~np.isnan(xycorr)

    return lags[fgud], xycorr[fgud]

def Tdecorr(Rxx, M=None, dtau=1., verbose=False):
    """
    USAGE
    -----
    Td = Tdecorr(Rxx)

    Computes the integral scale Td (AKA decorrelation scale, independence scale)
    for a data sequence with autocorrelation function Rxx. 'M' is the number of
    lags to incorporate in the summation (defaults to all lags) and 'dtau' is the
    lag time step (defaults to 1).

    The formal definition of the integral scale is the total area under the
    autocorrelation curve Rxx(tau):

    /+inf
    Td = 2 * |     Rxx(tau) dtau
    /0

    In practice, however, Td may become unrealistic if all of Rxx is summed
    (e.g., often goes to zero for data dominated by periodic signals); a
    different approach is to instead change M in the summation and use the
    maximum value of the integral Td(t):

    /t
    Td(t) = 2 * |     Rxx(tau) dtau
    /0

    References
    ----------
    e.g., Thomson and Emery (2014),
    Data analysis methods in physical oceanography,
    p. 274, equation 3.137a.

    Gille lecture notes on data analysis, available
    at http://www-pord.ucsd.edu/~sgille/mae127/lecture10.pdf
    """
    Rxx = np.asanyarray(Rxx)
    C0 = Rxx[0]
    N = Rxx.size # Sequence size.

    # Number of lags 'M' to incorporate in the summation.
    # Sum over all of the sequence if M is not chosen.
    if not M:
        M = N

    # Integrate the autocorrelation function.
    Td = np.zeros(M)
    for m in range(M):
        Tdaux = 0.
        for k in range(m-1):
            Rm = (Rxx[k] + Rxx[k+1])/2. # Midpoint value of the autocorrelation function.
            Tdaux = Tdaux + Rm*dtau # Riemann-summing Rxx.

        Td[m] = Tdaux

    # Normalize the integral function by the autocorrelation at zero lag
    # and double it to include the contribution of the side with
    # negative lags (C is symmetric about zero).
    Td = (2./C0)*Td

    if verbose:
        print("")
        print("Theoretical integral scale --> 2 * int 0...+inf [Rxx(tau)] dtau: %.2f."%Td[-1])
        print("")
        print("Maximum value of the cumulative sum: %.2f."%Td.max())

    return Td

def Tdecorrw(x, nblks=30, ret_median=True, verbose=True):
    """
    USAGE
    -----
    Ti = Tdecorrw(x, nblks=30, ret_median=True, verbose=True)

    'Ti' is the integral timescale calculated from the
    autocorrelation function calculated for variable 'x'
    block-averaged in 'nblks' chunks.
    """
    x = np.array(x)
    dnblkslr = round(nblks/2)

    tis = [Tdecorr(crosscorr(x, x, nblks=n, verbose=verbose)[1]).max() for n in range(nblks-dnblkslr, nblks+dnblkslr+1)]
    tis = np.ma.masked_invalid(tis)

    if verbose:
        print("========================")
        print(tis)
        print("========================")
        p1, p2, p3, p4, p5 = map(np.percentile, [tis]*5, (10, 25, 50, 75, 90))
        print("--> 10 %%, 25 %%, 50 %%, 75 %%, 90 %% percentiles for Ti:  %.2f,  %.2f,  %.2f,  %.2f,  %.2f."%(p1, p2, p3, p4, p5))
        print("------------------------")

    if ret_median:
        return np.median(tis)
    else:
        return tis

def rsig(ndof_eff, alpha=0.95):
	"""
	USAGE
	-----
	Rsig = rsig(ndof_eff, alpha=0.95)

	Computes the minimum (absolute) threshold value 'rsig' that
	the Pearson correlation coefficient r between two normally-distributed
	data sequences with 'ndof_eff' effective degrees of freedom has to have
	to be statistically significant at the 'alpha' (defaults to 0.95)
	confidence level.

	For example, if rsig(ndof_eff, alpha=0.95) = 0.2 for a given pair of
	NORMALLY-DISTRIBUTED samples with a correlation coefficient r>0.7, there
	is a 95 % chance that the r estimated from the samples is significantly
	different from zero. In other words, there is a 5 % chance that two random
	sequences would have a correlation coefficient higher than 0.7.

	OBS: This assumes that the two data series have a normal distribution.

	Translated to Python from the original matlab code by Prof. Sarah Gille
	(significance.m), available at http://www-pord.ucsd.edu/~sgille/sio221c/

	References
	----------
	Gille lecture notes on data analysis, available
	at http://www-pord.ucsd.edu/~sgille/mae127/lecture10.pdf

	Example
	-------
	TODO
	"""
	rcrit_z = erfinv(alpha)*np.sqrt(2./ndof_eff)

	return rcrit_z

def arsig(r0, Ndt, T1, T2, verbose=True):
    """
    USAGE
    -----
    alpha_rsig = arsig(r0, Ndt, T1, T2, verbose=True)

    The returned 'alpha_sig' is the CL at which the
    correlation coefficient 'r0' between two variables
    with integral timescales 'T1' and 'T2' and length 'Ndt'
    is significant.
    """
    Tslow = np.maximum(T1, T2) # The effective number of
    edof = Ndt/Tslow           # DoFs is constrained by the
                               # slower-decorrelating variable.
    alphai = 0.4
    issig = True
    while issig and alphai<1.0:
        rsigi = rsig(edof, alpha=alphai)
        issig=r0>=rsigi
        alphai+=0.01

    if verbose:
        print("Queried r = %.3f with %.1f EDoF. It is significant at **%.2f** CL."%(r0, edof, alphai))

    return alphai


text:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from vae_train.vae_utils import *

class Predictor(nn.Module):
    def __init__(self, n_class, out_class, z_size ,hard = True, hidden_size=64, num_layers=1, bidirectional=False):
        super(Predictor, self).__init__()
        self.z_size = z_size
        self.bidirectional = 2 if bidirectional else 1
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.hard = hard
        if hard:
            self.embedding = nn.Linear(n_class, hidden_size, bias=False)
            self.lstm = nn.LSTM(input_size=n_class + z_size, hidden_size=hidden_size,
                                num_layers=num_layers, batch_first=True, bidirectional=bidirectional)
        else:
            self.lstm = nn.LSTM(input_size=n_class + z_size, hidden_size=hidden_size,
                                num_layers=num_layers, batch_first=True, bidirectional=bidirectional)

        self.predict = nn.Linear(hidden_size, out_class)

    def forward(self, x, z, doc_len):
        batch_size = x.shape[0]
        seq_len = x.shape[1]

        if self.hard:
            x = torch.cat([x, z], dim=2)
            output, (h_n, c_n) = self.lstm(x)
            hidden = output[torch.arange(output.shape[0]), doc_len.view(-1)]
            return self.predict(hidden), self.embedding.weight
        else:
            x = torch.cat([x, z], dim=2)
            output, (h_n, c_n) = self.lstm(x)
            hidden = output[torch.arange(output.shape[0]), doc_len.view(-1)]
            return self.predict(hidden), None

ext: py

max_stars_repo_name: GT-SALT/Persuasive-Orderings

max_issues_repo_name: GT-SALT/Persuasive-Orderings

max_forks_repo_name: GT-SALT/Persuasive-Orderings

reason: import numpy


text:
import argparse
import sys
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

import numpy as np
from nnmnkwii.frontend import merlin as fe
from nnmnkwii.io import hts
from tqdm import tqdm

def get_parser():
    parser = argparse.ArgumentParser(
        description="Preprocess for duration models",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("utt_list", type=str, help="utternace list")
    parser.add_argument("lab_root", type=str, help="label directory")
    parser.add_argument("qst_file", type=str, help="HTS style question file")
    parser.add_argument("out_dir", type=str, help="out directory")
    parser.add_argument("--n_jobs", type=int, default=1, help="Number of jobs")
    return parser

def preprocess(lab_file, binary_dict, numeric_dict, in_dir, out_dir):
    labels = hts.load(lab_file)
    in_feats = fe.linguistic_features(labels, binary_dict, numeric_dict)
    out_feats = fe.duration_features(labels)

    utt_id = lab_file.stem
    np.save(
        in_dir / f"{utt_id}-feats.npy", in_feats.astype(np.float32), allow_pickle=False
    )
    np.save(
        out_dir / f"{utt_id}-feats.npy",
        out_feats.astype(np.float32),
        allow_pickle=False,
    )

if __name__ == "__main__":
    args = get_parser().parse_args(sys.argv[1:])

    with open(args.utt_list) as f:
        utt_ids = [utt_id.strip() for utt_id in f]
    lab_files = [Path(args.lab_root) / f"{utt_id}.lab" for utt_id in utt_ids]
    binary_dict, numeric_dict = hts.load_question_set(args.qst_file)

    in_dir = Path(args.out_dir) / "in_duration"
    out_dir = Path(args.out_dir) / "out_duration"
    in_dir.mkdir(parents=True, exist_ok=True)
    out_dir.mkdir(parents=True, exist_ok=True)

    with ProcessPoolExecutor(args.n_jobs) as executor:
        futures = [
            executor.submit(
                preprocess, lab_file, binary_dict, numeric_dict, in_dir, out_dir
            )
            for lab_file in lab_files
        ]
        for future in tqdm(futures):
            future.result()

max_stars_repo_name: kunosato-mado/ttslearn

max_issues_repo_name: kunosato-mado/ttslearn

max_forks_repo_name: kunosato-mado/ttslearn


text:
#################################################################################
#   Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.          #
#                                                                               #
#   Licensed under the Apache License, Version 2.0 (the "License").             #
#   You may not use this file except in compliance with the License.            #
#   You may obtain a copy of the License at                                     #
#                                                                               #
#       http://www.apache.org/licenses/LICENSE-2.0                              #
#                                                                               #
#   Unless required by applicable law or agreed to in writing, software         #
#   distributed under the License is distributed on an "AS IS" BASIS,           #
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.    #
#   See the License for the specific language governing permissions and         #
#   limitations under the License.                                              #
#################################################################################

"""
lidar_preprocess.py

This module defines the interface for the lidar preprocessing and different
preprocessing methods.
"""

import abc
import numpy as np

class LidarPreprocessingInterface():
    """ This class defines an interface for LiDAR preprocessing, it defines
        the basic functionality required by all preprocess types.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def preprocess_data(self, lidar_data):
        """Abstract method that has to be implemented in the classes that implement
           this interface with the details of preprocessing the interpolated LiDAR
           data passed as input.

        Args:
            lidar_data (np.array): This function expects a interpolated np.array
                                   with the distance information.

        Raises:
            NotImplementedError: Abstract method to force override this method when
                                 implementing thisinterface.
        """
        raise NotImplementedError("Lidar preprocessing object must implement a way "
                                  "to preprocess data")

class DefaultPreprocessing(LidarPreprocessingInterface):
    """Default preprocessing method that returns the interpolated LiDAR data as is.

    Args:
        LidarPreprocessingInterface: Implements the LidarPreprocessingInterface.
    """
    def preprocess_data(self, lidar_data):
        """Return the interpolated LiDAR data as is without any modification.

        Args:
            lidar_data (np.array): LiDAR distance data.

        Returns:
            np.array: LiDAR distance data.
        """
        return lidar_data

class SectorPreprocessing(LidarPreprocessingInterface):
    """Preprocessing method that is used for SECTOR_LIDAR models where the interpolated LiDAR
       distance information is further preprocessed into binary sector values to indicate there
       was any obstacle present in that sector.

    Args:
        LidarPreprocessingInterface: Implements the LidarPreprocessingInterface.
    """
    def __init__(self, num_sectors, max_lidar_dist):
        """Create the SectorPreprocessing object.

        Args:
            num_sectors (int): Number of sectors to split the LiDAR data.
            max_lidar_dist (float): Maximum distance until which we look for an
                                    obstacle in the sector.
        """
        self.num_sectors = num_sectors
        self.max_lidar_dist = max_lidar_dist

    def preprocess_data(self, lidar_data):
        """Split the LiDAR distance array into sectors and flag 1 if there is an distance
           value less than the maximum distance value.

        Args:
            lidar_data (np.array): Interpolated LiDAR distance data.

        Returns:
            np.array: Array of shape [num_sectors] with values 1/0 indicating if there is
                      an obstacle in that sector.
        """
        lidar_data = np.min(lidar_data.reshape(-1, self.num_sectors), axis=1)
        return (lidar_data < self.max_lidar_dist).astype(float)

max_stars_repo_name: jsspric/aws-deepracer-sensor-fusion-pkg

max_issues_repo_name: jsspric/aws-deepracer-sensor-fusion-pkg

max_forks_repo_name: jsspric/aws-deepracer-sensor-fusion-pkg


text:
"""
This module implements a Tensorflow neural network as a Scikit learn estimator.
"""

from __future__ import print_function
import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.utils.validation import check_X_y, check_array
from sklearn.metrics import mean_squared_error
from sklearn.metrics import mean_absolute_error
import tensorflow as tf
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score
import seaborn as sns
import pandas as pd

class MLPRegFlow(BaseEstimator, ClassifierMixin):
    """
    Neural-network with multiple hidden layers to do regression.
    This model optimises the squared error function using the Adam optimiser.

    :hidden_layer_sizes: Tuple, length = number of hidden layers, default (5,).

        The ith element represents the number of neurons in the ith
        hidden layer.

    :alpha: float, default 0.0001

        L2 penalty (regularization term) parameter.

    :batch_size: int, default 'auto'.

        Size of minibatches for stochastic optimizers.
        If the solver is 'lbfgs', the classifier will not use minibatch.
        When set to "auto", `batch_size=min(200, n_samples)`

    :learning_rate_init: double, default 0.001.

        The value of the learning rate in the numerical minimisation.

    :max_iter: int, default 200.

        Total number of iterations that will be carried out during the training process.

    """

    def __init__(self, hidden_layer_sizes=(5,), alpha=0.0001, batch_size='auto', learning_rate_init=0.001,
                 max_iter=80):

        # Initialising the parameters
        self.alpha = alpha
        self.batch_size = batch_size
        self.learning_rate_init = learning_rate_init
        self.max_iter = max_iter
        self.hidden_layer_sizes = hidden_layer_sizes

        # Initialising parameters needed for the Tensorflow part
        self.alreadyInitialised = False
        self.trainCost = []
        self.testCost = []
        self.isVisReady = False

    def fit(self, X, y, *test):
        """
        Fit the model to data matrix X and target y.

        :X: array of shape (n_samples, n_features).

            This contains the input data with samples in the rows and features in the columns.

        :y: array of shape (n_samples,).

            This contains the target values for each sample in the X matrix.

        :test: list with 1st element an array of shape (n_samples, n_features) and 2nd element an array of shape (n_samples, )

            This is a test set to visualise whether the model is overfitting.

        """

        print( "Starting the fitting process ... \n")

        # Check that X and y have correct shape
        X, y = check_X_y(X, y, multi_output=True)
        # Modification of the y data, because tensorflow wants a column vector, while scikit learn uses a row vector
        if y.ndim == 1:
            y = np.atleast_2d(y).T

        # Checking if a test set has been passed
        if test:
            if len(test) > 2:
                raise TypeError("foo() expected 2 arguments, got %d" % (len(test)))
            X_test = test[0]
            y_test = test[1]
            check_X_y(X_test, y_test, multi_output=True)
            if y.ndim == 1:
                y_test = np.atleast_2d(y_test).T

        self.n_feat = X.shape[1]
        self.n_samples = X.shape[0]
        self.n_output = y.shape[1]

        # Check the value of the batch size
        self.batch_size = self.checkBatchSize()

        # Initial set up of the NN
        X_train = tf.placeholder(tf.float32, [None, self.n_feat])
        Y_train = tf.placeholder(tf.float32, [None, self.n_output])

        # This part either randomly initialises the weights and biases or restarts training from wherever it was stopped
        if self.alreadyInitialised == False:
            weights, biases = self.__generate_weights()
            self.alreadyInitialised = True
        else:
            weights = []
            biases = []

            for ii in range(len(self.all_weights)):
                weights.append(tf.Variable(self.all_weights[ii]))
                biases.append(tf.Variable(self.all_biases[ii]))

        model = self.modelNN(X_train, weights, biases)
        cost = self.costReg(model, Y_train, weights, self.alpha)
        optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate_init).minimize(cost)

        # Initialisation of the variables
        init = tf.global_variables_initializer()

        # Running the graph
        with tf.Session() as sess:
            sess.run(init)

            for iter in range(self.max_iter):
                # This is the total number of batches in which the training set is divided
                n_batches = int(self.n_samples / self.batch_size)
                # This will be used to calculate the average cost per iteration
                avg_cost = 0
                # Learning over the batches of data
                for i in range(n_batches):
                    batch_x = X[i * self.batch_size:(i + 1) * self.batch_size, :]
                    batch_y = y[i * self.batch_size:(i + 1) * self.batch_size, :]
                    opt, c = sess.run([optimizer, cost], feed_dict={X_train: batch_x, Y_train: batch_y})
                    avg_cost += c / n_batches
                if test and iter % 50 == 0:
                    optTest, cTest = sess.run([optimizer, cost], feed_dict={X_train: X_test, Y_train: y_test})
                    self.testCost.append(cTest)
                self.trainCost.append(avg_cost)
                if iter % 500 == 0:
                    print( "Completed " + str(iter) + " iterations. \n")

            # Saving the weights for later re-use
            self.all_weights = []
            self.all_biases = []
            for ii in range(len(weights)):
                self.all_weights.append(sess.run(weights[ii]))
                self.all_biases.append(sess.run(biases[ii]))

    def modelNN(self, X, weights, biases):
        """
        This function evaluates the output of the neural network. It takes as input a data set, the weights and the
        biases.

        :X: tf.placeholder of shape (n_samples, n_features)
        :weights: list of tf.Variables of length len(hidden_layer_sizes) + 1
        :biases: list of tf.Variables of length len(hidden_layer_sizes) + 1
        :return: tf.Variable of size (n_samples, 1)
        """

        # Calculating the activation of the first hidden layer
        z = tf.add(tf.matmul(X, tf.transpose(weights[0])), biases[0])
        h = tf.nn.sigmoid(z)

        # Calculating the activation of all the hidden layers
        for ii in range(len(self.hidden_layer_sizes)-1):
            z = tf.add(tf.matmul(h, tf.transpose(weights[ii+1])), biases[ii+1])
            h = tf.nn.sigmoid(z)

        # Calculating the output of the last layer
        z = tf.add(tf.matmul(h, tf.transpose(weights[-1])), biases[-1])

        return z

    def __generate_weights(self):
        """
        This function generates the weights and the biases. It does so by looking at the size of the hidden layers and
        the number of features in the descriptor. The weights are initialised randomly.

        :return: lists (of length n_hidden_layers + 1) of tensorflow variables
        """

        weights = []
        biases = []

        # Weights from input layer to first hidden layer
        weights.append(tf.Variable(tf.truncated_normal([self.hidden_layer_sizes[0], self.n_feat], stddev=0.01),
                                   name='weight_in'))
        biases.append(tf.Variable(tf.zeros([self.hidden_layer_sizes[0]]), name='bias_in'))

        # Weights from one hidden layer to the next
        for ii in range(len(self.hidden_layer_sizes) - 1):
            weights.append(tf.Variable(
                tf.truncated_normal([self.hidden_layer_sizes[ii + 1], self.hidden_layer_sizes[ii]], stddev=0.01),
                name='weight_hidden'))
            biases.append(tf.Variable(tf.zeros([self.hidden_layer_sizes[ii + 1]]), name='bias_hidden'))

        # Weights from lat hidden layer to output layers
        weights.append(
            tf.Variable(tf.truncated_normal([self.n_output, self.hidden_layer_sizes[-1]], stddev=0.01), name='weight_out'))
        biases.append(tf.Variable(tf.zeros([self.n_output]), name='bias_out'))

        return weights, biases

    def costReg(self, qm_data, nn_data, weights, regu):
        """
        This function calculates the squared error cost function with L2 regularisation.

        :nn_data: tensor

            This tensor contains the neural network model.

        :qm_data: TensorFlow Place holder

            This tensor contains the y part of the data once the graph is initialised.

        :weights: array of TensorFlow variables of shape (len(hidden_layer_sizes+1), )

            It contains the weights for each hidden layer and the output layer.

        :regu: float

            The parameter that tunes the amount of regularisation.

        :return: tensor
        """
        err = tf.square(tf.subtract(qm_data, nn_data))
        cost = tf.reduce_mean(err, name="unreg_cost") # scalar

        # Calculating the regularisation term
        reg_l2 = tf.Variable(tf.zeros([1]), name="reg_term")
        for ii in range(len(weights)):
            reg_l2 = tf.add(reg_l2, tf.nn.l2_loss(weights[ii]))

        # Multiplying the regularisation term by the regularisation parameter
        reg_l2 = tf.scalar_mul(regu, reg_l2)

        # Regularised cost
        cost_reg = tf.add(cost, reg_l2, name="reg_cost")

        return cost_reg

    def plotLearningCurve(self):
        """
        This function plots the cost versus the number of iterations for the training set and the test set in the
        same plot. The cost on the train set is calculated every 50 iterations.
        """
        fig, ax = plt.subplots(figsize=(6, 6))
        ax.plot(self.trainCost, label="Train set", color="b")
        iterTest = range(0, self.max_iter, 50)
        ax.plot(iterTest, self.testCost, label="Test set", color="red")
        ax.set_xlabel('Number of iterations')
        ax.set_ylabel('Cost Value')
        ax.legend()
        plt.yscale("log")
        plt.show()

    def checkBatchSize(self):
        """
        This function is called to check if the batch size has to take the default value or a user-set value.
        If it is a user set value, it checks whether it is a reasonable value.

        :return: int

            The default is 100 or to the total number of samples present if this is smaller than 100. Otherwise it is
            checked whether it is smaller than 1 or larger than the total number of samples.
        """
        if self.batch_size == 'auto':
            batch_size = min(100, self.n_samples)
        else:
            if self.batch_size < 1 or self.batch_size > self.n_samples:
                print( "Warning: Got `batch_size` less than 1 or larger than sample size. It is going to be clipped")
                batch_size = np.clip(self.batch_size, 1, self.n_samples)
            else:
                batch_size = self.batch_size

        return batch_size

    def checkIsFitted(self):
        """
        This function checks whether the weights and biases have been changed from their initial values.

        :return: True if the weights and biases are not all zero.
        """
        if self.alreadyInitialised == False:
            raise StandardError("The fit function has not been called yet")
        else:
            return True

    def predict(self, X):
        """
        This function uses the X data and plugs it into the model and then returns the predicted y

        :X: array of shape (n_samples, n_features)

            This contains the input data with samples in the rows and features in the columns.

        :return: array of size (n_samples, n_output)

            This contains the predictions for the target values corresponding to the samples contained in X.

        """
        print( "Calculating the predictions. \n")

        if self.checkIsFitted():
            check_array(X)

            X_test = tf.placeholder(tf.float32, [None, self.n_feat])

            weights = []
            biases = []

            for ii in range(len(self.all_weights)):
                weights.append(tf.Variable(self.all_weights[ii]))
                biases.append(tf.Variable(self.all_biases[ii]))

            model = self.modelNN(X_test, weights, biases)

            init = tf.global_variables_initializer()

            with tf.Session() as sess:
                sess.run(init)
                predictions = sess.run(model, feed_dict={X_test: X})
                predictions = np.reshape(predictions,(predictions.shape[0],self.n_output))

            return predictions
        else:
            raise StandardError("The fit function has not been called yet, so the model has not been trained yet.")

    def score(self, X, y, sample_weight=None):
        """
        Returns the mean accuracy on the given test data and labels. It calculates the R^2 value. It is used during the
        training of the model.

        :X: array of shape (n_samples, n_features)

            This contains the input data with samples in the rows and features in the columns.

        :y: array of shape (n_samples, n_output)

            This contains the target values for each sample in the X matrix.

        :sample_weight: array of shape (n_samples,)

            Sample weights (not sure what this is, but i need it for inheritance from the BaseEstimator)

        :return: double
            This is a score between -inf and 1 (best value is 1) that tells how good the correlation plot is.
        """

        y_pred = self.predict(X)
        r2 = r2_score(y, y_pred)
        return r2

    def scoreFull(self, X, y):
        """
        This scores the predictions more thouroughly than the function 'score'. It calculates the r2, the root mean
        square error, the mean absolute error and the largest positive/negative outliers. They are all in the units of
        the data passed.

        :X: array of shape (n_samples, n_features)

            This contains the input data with samples in the rows and features in the columns.

        :y: array of shape (n_samples,)

            This contains the target values for each sample in the X matrix.

        :return:
        :r2: double

            This is a score between -inf and 1 (best value is 1) that tells how good the correlation plot is.

        :rmse: double

            This is the root mean square error

        :mae: double

            This is the mean absolute error

        :lpo: double

            This is the largest positive outlier.

        :lno: double

            This is the largest negative outlier.

        """

        y_pred = self.predict(X)
        r2 = r2_score(y, y_pred)
        rmse = np.sqrt(mean_squared_error(y, y_pred))
        mae = mean_absolute_error(y, y_pred)
        lpo, lno = self.largestOutliers(y, y_pred)

        return r2, rmse, mae, lpo, lno

    def largestOutliers(self, y_true, y_pred):
        """
        This function calculates the larges positive and negative outliers from the predictions of the neural net.

        :y_true: array of shape (n_samples,)

            This contains the target values for each sample.

        :y_pred: array of shape (n_samples,)

            This contains the neural network predictions of the target values for each sample.

        :return:

        :lpo: double

            This is the largest positive outlier.

        :lno: double

            This is the largest negative outlier.
        """
        diff = y_pred - y_true
        lpo = np.amax(diff)
        lno = - np.amin(diff)

        return lpo, lno

    def errorDistribution(self, X, y):
        """
        This function plots histograms of how many predictions have an error in a certain range.

        :X: array of shape (n_samples, n_features)

            This contains the input data with samples in the rows and features in the columns.

        :y: array of shape (n_samples,)

            This contains the target values for each sample in the X matrix.
        """
        y_pred = self.predict(X)
        diff_kJmol = (y - y_pred)*2625.50
        df = pd.Series(diff_kJmol, name="Error (kJ/mol)")
        # sns.set_style(style='white')
        # sns.distplot(df, color="#f1ad1e")
        # sns.plt.savefig("ErrorDist.png", transparent=True, dpi=800)
        plt.show()

    def correlationPlot(self, X, y, ylim=(1.90, 1.78), xlim=(1.90, 1.78)):
        """
        This function plots a correlation plot of the values that are in the data set and the NN predictions. It expects
        the target values to be in Hartrees.

        :X: array of shape (n_samples, n_features)

            This contains the input data with samples in the rows and features in the columns.

        :y: array of shape (n_samples,)

            This contains the target values for each sample in the X matrix.

        :ylim: tuple of shape (2,) containing doubles

            These are the limits of the y values for the plot.

        :xlim: tuple of shape (2,) containing doubles

            These are the limits of the x values for the plot.
        """
        y_pred = self.predict(X)
        df = pd.DataFrame()
        df['High level calculated energies (Ha)'] = y
        df['NN predicted energies (Ha)'] = y_pred
        lm = sns.lmplot('High level calculated energies (Ha)', 'NN predicted energies (Ha)', data=df,
                        scatter_kws={"s": 20, "alpha": 0.6}, line_kws={"alpha": 0.5})
        lm.set(ylim=ylim)
        lm.set(xlim=xlim)
        plt.show()

    def plotWeights(self):
        """
        This function plots the weights of the first layer of the neural network as a heat map.
        """

        w1_square_tot = []

        for i in range(self.hidden_layer_sizes[0]):
            w1_square = self.reshape_triang(self.w1[i], 7)
            w1_square_tot.append(w1_square)

        n = int(np.ceil(np.sqrt(self.hidden_layer_sizes)))
        additional = n**2 - self.hidden_layer_sizes[0]

        fig, axn = plt.subplots(n, n, sharex=True, sharey=True)
        fig.set_size_inches(11.7, 8.27)
        cbar_ax = fig.add_axes([.91, .3, .03, .4])

        for i, ax in enumerate(axn.flat):
            if i >= self.hidden_layer_sizes[0]:
                break
            df = pd.DataFrame(w1_square_tot[i])
            sns.heatmap(df,
                        ax=ax,
                        cbar=i == 0,
                        vmin=-0.2, vmax=0.2,
                        cbar_ax=None if i else cbar_ax, cmap="PiYG")

        fig.tight_layout(rect=[0, 0, 0.9, 1])
        sns.plt.savefig("weights_l1.png", transparent=False, dpi=600)
        # sns.plt.show()

    def reshape_triang(self, X, dim):
        """
        This function reshapes a single flattened triangular matrix back to a square diagonal matrix.

        :X: array of shape (n_atoms*(n_atoms+1)/2, )

            This contains a sample of the Coulomb matrix trimmed down so that it contains only the a triangular matrix.

        :dim: int

            The triangular matrix X will be reshaped to a matrix that has size dim by dim.

        :return: array of shape (n_atoms, n_atoms)

            This contains the square diagonal matrix.
        """

        x_square = np.zeros((dim, dim))
        counter = 0
        for i in range(dim):
            for j in range(i, dim):
                x_square[i][j] = X[counter]
                x_square[j][i] = X[counter]
                counter = counter + 1

        return x_square

    def __vis_input(self, initial_guess):
        """
        This function does gradient ascent to generate an input that gives the highest activation for each neuron of
        the first hidden layer.

        :initial_guess: array of shape (n_features,)

            A coulomb matrix to use as the initial guess to the gradient ascent in the hope that the closest local
            maximum will be found.

        :return: list of arrays of shape (num_atoms, num_atoms)

            each numpy array is the input for a particular neuron that gives the highest activation.

        """

        self.isVisReady = True
        initial_guess = np.reshape(initial_guess, newshape=(1, initial_guess.shape[0]))
        input_x = tf.Variable(initial_guess, dtype=tf.float32)
        activations = []
        iterations = 7000
        lambda_reg = 0.0002
        self.x_square_tot = []

        for node in range(self.hidden_layer_sizes[0]):

            # Calculating the activation of the first layer
            w1_node = tf.constant(self.w1[node], shape=(1,self.n_feat))
            b1_node = tf.constant(self.b1[node])
            z1 = tf.add(tf.matmul(tf.abs(input_x), tf.transpose(w1_node)), b1_node)
            a1 = tf.nn.sigmoid(z1)
            a1_reg = a1 - lambda_reg * tf.tensordot(input_x, tf.transpose(input_x), axes=1)

            # Function to maximise a1
            optimiser = tf.train.AdamOptimizer(learning_rate=0.01).minimize(-a1_reg)

            # Initialising the model
            init = tf.global_variables_initializer()

            # Running the graph
            with tf.Session() as sess:
                sess.run(init)

                for i in range(iterations):
                    sess.run(optimiser)

                temp_a1 = sess.run(a1)
                activations.append(temp_a1)     # Calculating the activation for checking later if a node has converged
                final_x = sess.run(input_x)     # Storing the best input

            x_square = self.reshape_triang(final_x[0,:], 7)
            self.x_square_tot.append(x_square)
        print("The activations at the end of the optimisations are:")
        print(activations)

        return self.x_square_tot

    def vis_input_matrix(self, initial_guess, write_plot=False):
        """
        This function calculates the inputs that would give the highest activations of the neurons in the first hidden
        layer of the neural network. It then plots them as a heat map.

        :initial_guess: array of shape (n_features,)

            A coulomb matrix to use as the initial guess to the gradient ascent in the hope that the closest local
            maximum will be found.

        :write_plot: boolean, default False

            If this is true, the plot is written to a png file.
        """

        if self.isVisReady == False:
            self.x_square_tot = self.__vis_input(initial_guess)

        n = int(np.ceil(np.sqrt(self.hidden_layer_sizes)))
        additional = n ** 2 - self.hidden_layer_sizes[0]

        fig, axn = plt.subplots(n, n, sharex=True, sharey=True)
        fig.set_size_inches(11.7, 8.27)
        cbar_ax = fig.add_axes([.91, .3, .03, .4])
        counter = 0

        for i, ax in enumerate(axn.flat):
            df = pd.DataFrame(self.x_square_tot[counter])
            ax.set(xticks=[], yticks=[])
            sns.heatmap(df, ax=ax, cbar=i == 0, cmap='RdYlGn',
                        vmax=8, vmin=-8,
                        cbar_ax=None if i else cbar_ax)
            counter = counter + 1
            if counter >= self.hidden_layer_sizes[0]:
                break

        fig.tight_layout(rect=[0, 0, 0.9, 1])
        if write_plot==True:
            sns.plt.savefig("high_a1_input.png", transparent=False, dpi=600)
        sns.plt.show()

    def vis_input_network(self, initial_guess, write_plot=False):
        """
        This function calculates the inputs that would give the highest activations of the neurons in the first hidden
        layer of the neural network. It then plots them as a netwrok graph.

        :initial_guess: array of shape (n_features,)

            A coulomb matrix to use as the initial guess to the gradient ascent in the hope that the closest local
            maximum will be found.

        :write_plot: boolean, default False

            If this is true, the plot is written to a png file.
        """
        import networkx as nx

        if self.isVisReady == False:
            self.x_square_tot = self.__vis_input(initial_guess)

        n = int(np.ceil(np.sqrt(self.hidden_layer_sizes)))

        fig = plt.figure(figsize=(10, 8))
        for i in range(n**2):
            if i >= self.hidden_layer_sizes[0]:
                break
            fig.add_subplot(n,n,1+i)
            A = np.matrix(self.x_square_tot[i])
            graph2 = nx.from_numpy_matrix(A, parallel_edges=False)
            # nodes and their label
            # pos = {0: np.array([0.46887886, 0.06939788]), 1: np.array([0, 0.26694294]),
            #        2: np.array([0.3, 0.56225267]),
            #        3: np.array([0.13972517, 0.]), 4: np.array([0.6, 0.9]), 5: np.array([0.27685853, 0.31976436]),
            #        6: np.array([0.72, 0.9])}
            pos = {}
            for i in range(7):
                x_point = 0.6*np.cos((i+1)*2*np.pi/7)
                y_point = 0.6*np.sin((i+1)*2*np.pi/7)
                pos[i] = np.array([x_point, y_point])
            labels = {}
            labels[0] = 'H'
            labels[1] = 'H'
            labels[2] = 'H'
            labels[3] = 'H'
            labels[4] = 'C'
            labels[5] = 'C'
            labels[6] = 'N'
            node_size = np.zeros(7)
            for i in range(7):
                node_size[i] =  abs(graph2[i][i]['weight'])*10
            nx.draw_networkx_nodes(graph2, pos, node_size=node_size)
            nx.draw_networkx_labels(graph2, pos, labels=labels, font_size=15, font_family='sans-serif', font_color='blue')
            # edges
            edgewidth = [d['weight'] for (u, v, d) in graph2.edges(data=True)]
            nx.draw_networkx_edges(graph2, pos, width=edgewidth)
            plt.axis('off')

        if write_plot==True:
            plt.savefig("high_a1_network.png")  # save as png

        plt.show()  # display

# This example tests the module on fitting a simple cubic function and then plots the results

if __name__ == "__main__":

    estimator = MLPRegFlow(hidden_layer_sizes=(5, 5, 5), learning_rate_init=0.01, max_iter=5000, alpha=0)
    x = np.arange(-2.0, 2.0, 0.05)
    X = np.atleast_2d(x).T
    y = X ** 3

    estimator.fit(X, y)
    y_pred = estimator.predict(X)

    print("RMSD error:", np.sqrt(np.sum((y_pred - y)**2)/x.size))

    #  Visualisation of predictions
    fig2, ax2 = plt.subplots(figsize=(6,6))
    ax2.scatter(x, y, label="original", marker="o", c="r")
    ax2.scatter(x, y_pred, label="predictions", marker="o", c='b')
    ax2.set_xlabel('x')
    ax2.set_ylabel('y')
    ax2.legend()

    # Correlation plot
    fig3, ax3 = plt.subplots(figsize=(6,6))
    ax3.scatter(y, y_pred, marker="o", c="r")
    ax3.set_xlabel('original y')
    ax3.set_ylabel('prediction y')
    plt.show()

    # estimator.errorDistribution(X, y)

max_stars_repo_name: larsbratholm/YAMLP

max_issues_repo_name: larsbratholm/YAMLP

max_forks_repo_name: larsbratholm/YAMLP

reason: import numpy,import networkx


text:
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import bruges as b
from scipy.interpolate import interp1d
import os

topbox = dict(boxstyle='round', ec='none', fc='w', alpha=0.6)
format_tops={'fontsize':10, 'color':'blue', 'ha':'right', 'bbox':topbox}
format_title={'fontsize':14, 'weight':'bold'}

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def contactcement(K0, G0, phi, phi_c=0.4, Cn=8.6, Kc=37, Gc=45, scheme=2):
    PR0=(3*K0-2*G0)/(6*K0+2*G0)
    PRc = (3*Kc-2*Gc)/(6*Kc+2*Gc)
    if scheme == 1: # scheme 1: cement deposited at grain contacts
        alpha = ((phi_c-phi)/(3*Cn*(1-phi_c))) ** (1/4)
    else: # scheme 2: cement evenly deposited on grain surface
        alpha = ((2*(phi_c-phi))/(3*(1-phi_c)))**(1/2)
    LambdaN = (2*Gc*(1-PR0)*(1-PRc)) / (np.pi*G0*(1-2*PRc))
    N1 = -0.024153*LambdaN**-1.3646
    N2 = 0.20405*LambdaN**-0.89008
    N3 = 0.00024649*LambdaN**-1.9864
    Sn = N1*alpha**2 + N2*alpha + N3
    LambdaT = Gc/(np.pi*G0)
    T1 = -10**-2*(2.26*PR0**2+2.07*PR0+2.3)*LambdaT**(0.079*PR0**2+0.1754*PR0-1.342)
    T2 = (0.0573*PR0**2+0.0937*PR0+0.202)*LambdaT**(0.0274*PR0**2+0.0529*PR0-0.8765)
    T3 = 10**-4*(9.654*PR0**2+4.945*PR0+3.1)*LambdaT**(0.01867*PR0**2+0.4011*PR0-1.8186)
    St = T1*alpha**2 + T2*alpha + T3
    K_DRY = 1/6*Cn*(1-phi_c)*(Kc+(4/3)*Gc)*Sn
    G_DRY = 3/5*K_DRY+3/20*Cn*(1-phi_c)*Gc*St
    return K_DRY, G_DRY

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def constantcement(K0, G0, phi, phi_cem=0.38, phi_c=0.4, Cn=8.6, Kc=37, Gc=45, scheme=2):
    # contact cement model
    K_HI, G_HI = contactcement(K0, G0, phi, phi_c=phi_c, Cn=Cn, Kc=Kc, Gc=Gc, scheme=scheme)
    # lower bound Hashin-Shtrikman starting from phi_cem
    Kcc, Gcc = contactcement(K0, G0, phi_cem, phi_c=phi_c, Cn=Cn, Kc=Kc, Gc=Gc, scheme=scheme)
    K_LO = -4/3*Gcc + (((phi/phi_cem)/(Kcc+4/3*Gcc)) + ((1-phi/phi_cem)/(K0+4/3*Gcc)))**-1
    tmp = Gcc/6*((9*Kcc+8*Gcc) / (Kcc+2*Gcc))
    G_LO= -tmp + ((phi/phi_cem)/(Gcc+tmp) + ((1-phi/phi_cem)/(G0+tmp)))**-1
    # initialize empty vectors for K and G dry
    K_DRY, G_DRY=(np.full(phi.size, np.nan) for _ in range(2))
    # for porosities>phi_cem use [K,G]_HI = contact cement model
    # for porosities<=phi_cem use [K,G]_LO = constant cement model
    K_DRY[phi>phi_cem]=K_HI[phi>phi_cem]
    K_DRY[phi<=phi_cem]=K_LO[phi<=phi_cem]
    G_DRY[phi>phi_cem]=G_HI[phi>phi_cem]
    G_DRY[phi<=phi_cem]=G_LO[phi<=phi_cem]
    return K_DRY, G_DRY

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def inccement(K0, G0, phi, phi_cem=0.38, phi_c=0.4, Cn=8.6, Kc=37, Gc=45, scheme=2):
    Kcc, Gcc = contactcement(K0, G0, phi_cem, phi_c=phi_c, Cn=Cn, Kc=Kc, Gc=Gc, scheme=scheme)
    K_DRY = -4/3*G0 + (((phi/phi_cem)/(Kcc+4/3*G0)) + ((1-phi/phi_cem)/(K0+4/3*G0)))**-1
    tmp = G0/6*((9*K0+8*G0) / (K0+2*G0))
    G_DRY = -tmp + ((phi/phi_cem)/(Gcc+tmp) + ((1-phi/phi_cem)/(G0+tmp)))**-1
    return K_DRY, G_DRY

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def vels(k_dry,mu_dry,k_min,rho_min,k_fl,rho_fl,phi):
    # converts all inputs to SI (density in kg/m3 and moduli in Pa)
    KD = k_dry*1e9
    GD = mu_dry*1e9
    K0 = k_min*1e9
    D0 = rho_min*1e3
    Kf = k_fl*1e9
    Df = rho_fl*1e3
    rho = D0*(1-phi)+Df*phi
    K = KD + (1-KD/K0)**2 / ( (phi/Kf) + ((1-phi)/K0) - (KD/K0**2) )
    vp = np.sqrt((K+4/3*GD)/rho)
    vs = np.sqrt(GD/rho)
    return vp, vs, rho/1e3, K/1e9

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def blok_plot(REF,BLK,ztop=None,zbot=None):
    if ztop is None: ztop = BLK.index.min()
    if zbot is None: zbot = BLK.index.max()
    REF=REF[(REF.index>=ztop) & (REF.index<=zbot)]
    REF.columns=REF.columns.str.upper()
    BLK.columns=BLK.columns.str.upper()

    f, ax = plt.subplots(nrows=1,ncols=5,sharey=True,figsize=(8,5))
    ax[0].plot(REF.VSH, REF.index, color='.5')
    ax[0].set_xlabel('Vsh', color='.5')
    ax[1].plot(REF.PHIE, REF.index, color='.5')
    ax[1].set_xlabel('phi', color='.5')
    ax[2].plot(REF.VP_FRMB, REF.index, color='black')
    ax[2].plot(BLK.VP, BLK.index, color='red', lw=4)
    ax[2].set_xlabel('Vp [m/s]')
    ax[3].plot(REF.VS_FRMB, REF.index, color='black')
    ax[3].plot(BLK.VS, BLK.index, color='red', lw=4)
    ax[3].set_xlabel('Vs [m/s]')
    ax[4].plot(REF.RHO_FRMB, REF.index, color='black')
    ax[4].plot(BLK.RHO, BLK.index, color='red', lw=4)
    ax[4].set_xlabel('Density [g/cc]')
    ax[0].set_ylim(zbot,ztop)
    plt.tight_layout()

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def generate_wells(WW,top_blocks,name,output_dir):
    # top_blocks is an array with 4 values defining 3 blocks:
    # block 1 = shallow block, caprock
    # block 2 = central block, reservoir
    # block 3 = deep block
    z = WW.index.values
    nblks = np.size(top_blocks)-1

    # compute and store average vp,vs,rho for fluids 1 (brine) and 2 (oil)
    vp1_k,vs1_k,rho1_k,vp2_k,vs2_k,rho2_k,phi_k = (np.zeros(nblks) for _ in range(7))
    for nn in range(nblks):
        rr = (WW.index>=top_blocks[nn]) & (WW.index<top_blocks[nn+1])
        vp1_k[nn] = WW.VP_FRMB[rr].mean()
        vs1_k[nn] = WW.VS_FRMB[rr].mean()
        rho1_k[nn] = WW.RHO_FRMB[rr].mean()
        vp2_k[nn] = WW.VP_FRMO[rr].mean()
        vs2_k[nn] = WW.VS_FRMO[rr].mean()
        rho2_k[nn] = WW.RHO_FRMO[rr].mean()

    # range of variations for thickness (vTH), water saturation (vSW) and porosity (vPH)
    vTH= np.linspace(0,60,4)
    vSW= np.linspace(0,1,5)
    vPH = [-.1, -.05, 0, .05]

    # output depths will be within top block 1 and base block 3
    z_out = z[(z>=top_blocks[0]) & (z<=top_blocks[3])]
    # mm0 = to select depths above reservoir
    mm0 = (z_out>=top_blocks[0]) & (z_out<top_blocks[1])

    # loop over thickness variations
    for tt in vTH:
        # initialize vp,vs,rho for both fluids with values taken from block 3
        vp1_TH  = np.ones(z_out.size)*vp1_k[2]
        vs1_TH  = np.ones(z_out.size)*vs1_k[2]
        rho1_TH = np.ones(z_out.size)*rho1_k[2]
        vp2_TH  = np.ones(z_out.size)*vp2_k[2]
        vs2_TH  = np.ones(z_out.size)*vs2_k[2]
        rho2_TH = np.ones(z_out.size)*rho2_k[2]

        # sets vp,vs,rho for block 1
        vp1_TH[mm0] = vp1_k[0]
        vs1_TH[mm0] = vs1_k[0]
        rho1_TH[mm0] = rho1_k[0]
        vp2_TH[mm0] = vp2_k[0]
        vs2_TH[mm0] = vs2_k[0]
        rho2_TH[mm0] = rho2_k[0]

        # mm1 = to select depths in block 2=reservoir (base varies with thickness variations)
        mm1 = (z_out>=top_blocks[1]) & (z_out<top_blocks[1]+tt)

        # sets vp,vs,rho for block 2
        vp1_TH[mm1] = vp1_k[1]
        vs1_TH[mm1] = vs1_k[1]
        rho1_TH[mm1] = rho1_k[1]
        vp2_TH[mm1] = vp2_k[1]
        vs2_TH[mm1] = vs2_k[1]
        rho2_TH[mm1] = rho2_k[1]
        # loop over water saturation variations
        for ss in vSW:
            # modify velocities and density according to saturation changes
            # assume a linear change from brine to oil filled rock
            vp_SW = ss*vp1_TH + (1-ss)*vp2_TH
            vs_SW = ss*vs1_TH + (1-ss)*vs2_TH
            rho_SW = ss*rho1_TH + (1-ss)*rho2_TH

            # loop over porosity variations
            for pp in vPH:
                # modify velocities and density according to porosity changes
                # assume linear porosity-velocity ratios in the limited range of study
                # relations are given by rock physics study
                dvp_b = -pp*4500
                dvp_o = -pp*4500
                dvs_b = -pp*3100
                dvs_o = -pp*3050
                dvs   = ss*dvs_b + (1-ss)*dvs_o
                dvp   = ss*dvp_b + (1-ss)*dvp_o
                drho_b = (2.65-1.0)*-pp
                drho_o = (2.65-0.8)*-pp
                drho   = ss*drho_b + (1-ss)*drho_o
                # calculate new parameters with varying saturations
                vp_out=np.copy(vp_SW); vs_out=np.copy(vs_SW); rho_out=np.copy(rho_SW);
                vp_out[mm1] = vp_SW[mm1] + dvp
                vs_out[mm1] = vs_SW[mm1] + dvs
                rho_out[mm1] = rho_SW[mm1] + drho
                d_out = {'DEPTH':z_out, 'VP': vp_out, 'VS': vs_out, 'RHO': rho_out}
                df_out = pd.DataFrame(data=d_out)
                filename = '{:s}_Z{:02d}_Sw{:03d}_Por{:02d}'.format(name,int(np.round(tt)),int(100*ss),int((0.3+pp)*100))
                df_out.to_csv(output_dir+'/'+filename, index=False)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_well_files(wells_dir, name):
   well_files = []
   for dirpath, _, filenames in os.walk(wells_dir):
       well_files += [os.path.join(dirpath, f) for f in filenames if f.startswith(name)]
   return well_files

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_nears_fars(df):
   sample_size = 64
   number_of_splits = abs(df['NEAR'].size/64)
   nears = np.array_split(df['NEAR'].values, number_of_splits)
   fars = np.array_split(df['FAR'].values, number_of_splits)
   nears = np.asarray(nears).transpose()
   fars = np.asarray(fars).transpose()
   return nears, fars

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_twt(tdr,z):
    tt=tdr[:,1]
    zz=tdr[:,0]
    d2t = interp1d(zz, tt, kind='linear', bounds_error=False, fill_value='extrapolate')
    return d2t(z)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def make_synt(WW,ang,wavelet,method='shuey'):
    '''
    WW: Pandas dataframe with VP, VS, RHO (optionally EPSILON and DELTA)
    ang: angle range, define with ang=np.arange(0,50,1)
    wavelet
    method: 'shuey' (Shuey 2-terms), 'shuey3' (Shuey 3-terms),
            'aki' (Aki-Richards)
    '''
    WW.columns=WW.columns.str.upper()
    uvp, lvp   = WW.VP.values[:-1], WW.VP.values[1:]
    uvs, lvs   = WW.VS.values[:-1], WW.VS.values[1:]
    urho, lrho = WW.RHO.values[:-1], WW.RHO.values[1:]
    z=WW.index.values  # z is two-way-time
    synt  = np.zeros((z.size,ang.size))
    #--> calculate reflectivities with AVO equation,
    #--> convolve with input wavelet and fill in traces of synthetic seismogram
    for i,alpha in enumerate(ang):
        if method is 'shuey':
            RC = shuey(uvp,uvs,urho,lvp,lvs,lrho,alpha)
        elif method is 'shuey3':
            RC = shuey(uvp,uvs,urho,lvp,lvs,lrho,alpha,approx=False)
        else:
            RC = akirichards(uvp,uvs,urho,lvp,lvs,lrho,alpha)
        RC = np.append(np.nan, RC)
        RC = np.nan_to_num(RC)
        synt[:,i] = np.convolve(RC, wavelet, mode='same')
    return RC, synt

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def plot_synt(WW,synt,ztop,zbot,gain=10):
    '''
    WW: Pandas dataframe with VP, VS, RHO (in time)
    synt: synthetic seismogram computed with make_synt (in time)
    ztop,zbot: display window
    gain: multiplier to be applied to wiggles (default=5)
    method: 'shuey' (Shuey 2-terms), 'shuey3' (Shuey 3-terms),
            'aki' (Aki-Richards)
    '''
    WW.columns=WW.columns.str.upper()
    it1=np.abs(WW.index-ztop).argmin()
    it2=np.abs(WW.index-zbot).argmin()
    ss = synt[it1:it2,:]
    clip=np.abs(synt.max())
    f,ax=plt.subplots(nrows=1,ncols=5)
    opz1={'color':'k','linewidth':.5}
    opz2={'linewidth':0, 'alpha':0.6}
    ax[0].plot(WW.VP*WW.RHO,WW.index,'-k')
    ax[1].plot(WW.VP/WW.VS,WW.index,'-k')
    ax[2].plot(synt[:,0],WW.index,'-k')
    ax[3].plot(synt[:,1],WW.index,'-k')
    im=ax[4].imshow(ss,interpolation='none', cmap='Greys',aspect='auto')
    cbar=plt.colorbar(im, ax=ax[4])
    ax[0].set_xlabel('AI [m/s*g/cc]')
    ax[0].set_ylabel('TWT [s]')
    ax[1].set_xlabel('Vp/Vs')
    ax[2].set_title('Near')
    ax[3].set_title('Far')
    ax[4].set_title('Near|Far')
    for aa in ax[:4]:
        aa.set_ylim(zbot,ztop)
        aa.grid()
    for aa in ax[1:]:
        aa.set_yticklabels([])
    for aa in ax[2:4]:
        aa.set_xlim(-clip,+clip)
        aa.set_xticklabels([])
    for aa in ax[:2]:
        aa.xaxis.tick_top()
        plt.setp(aa.xaxis.get_majorticklabels(), rotation=90)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def td(WW,sr=0.1524,KB=0,WD=0,repl_vel=1600):
    '''
    td (C) aadm 2016
    Calculates time-depth relation by sonic integration.

    INPUT
    WW: Pandas dataframe
    sr: depth sampling rate (m, default 0.1524 m = half-foot)
    KB: kelly bushing elevation (m, default 0)
    WD: water depth (m, default 0)
    repl_vel: replacement velocity for overburden i.e. interval between seafloor and beginning of the logs
              (m/s, default 1600)

    OUTPUT
    numpy array with 2 columns: 0=depth, 1=twt (secs)
    '''
    WW.columns=WW.columns.str.upper()
    if 'TVD' in WW.columns:
        depth = WW.TVD.values
    else:
        depth = WW.index.values
    WW.VP.interpolate(inplace=True)
    sonic=1/WW.VP.values # VP in m/s
    start = depth.min()
    water_vel = 1480
    wb_twt = 2.0*WD/water_vel
    sonic_start=depth[np.isfinite(sonic)].min()
    sonic_start_twt=2*(sonic_start-KB-WD)/repl_vel + wb_twt
    scaled_sonic = sr*sonic[depth>=sonic_start]
    twt = 2*np.cumsum(scaled_sonic) + sonic_start_twt
    print('[TD] water bottom two-way-time: {:.3f} [s]'.format(wb_twt))
    print('[TD] sonic log start: {:.3f} [m] = {:.3f} [s]'.format(sonic_start, sonic_start_twt))
    print('[TD] computed twt scale range: {:.3f}-{:.3f} [s]'.format(twt.min(),twt.max()))
    return np.column_stack((depth[depth>=sonic_start],twt))

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def welltime(WW,tdr,dt=0.001,ztop=None,zbot=None,name=None,tops=None,qcplot=True):
    '''
    welltime (C) aadm 2016
    Converts logs sampled in depth to time using a reference time-depth function.
    Load existing t-d with tdr=np.loadtxt('TD.dat', delimiter=',')
    or use squit.well.td to create one.

    INPUT
    WW: Pandas dataframe
    tdr: time-depth table: numpy array shape Nx2, column 0 = MD [m], column 1 = TWT [s]
    dt: sample rate in seconds
    ztop,zbot: display window (defaults to min,max depth)
    name: well name (or anything else) to print
    tops: dictionary containing stratigraphic tops, e.g.: tops={'Trias': 1200,'Permian': 2310}
          or Pandas Series, e.g: tops=pd.Series({'Trias': 1200,'Permian': 2310})

    OUTPUT
    Pandas dataframe with logs sampled in time
    ztop, zbot converted to TWT
    '''
    WW.columns=WW.columns.str.upper()
    flagfrm=True if 'VP_FRMB' in WW.columns else False
    z = WW.index
    if ztop is None: ztop = z.min()
    if zbot is None: zbot = z.max()
    flagtops=False if tops is None else True
    if flagtops:
        if not isinstance(tops, pd.Series):
            tops=pd.Series(tops)
        tops=tops.dropna().sort_values()

    #--> load depth-time relationship (depth is MD)
    tt=tdr[:,1]
    zz=tdr[:,0]
    #-->  twt reference log sampled like depth reference log
    # twt_log = np.interp(z, zz, tt, left=np.nan, right=np.nan)
    ff=(z>=zz.min()) & (z<=zz.max())
    twt_log = np.interp(z[ff], zz, tt, left=np.nan, right=np.nan)
    #-->  interpolant to convert depths to times on the fly (e.g., tops)
    d2t = interp1d(zz, tt, kind='linear', bounds_error=False, fill_value='extrapolate')
    if qcplot:
        print('[WELLTIME] plot window top, bottom [m]:{:.0f}-{:.0f}, [s]:{:.4f}-{:.4f}'.format(ztop,zbot,float(d2t(ztop)),float(d2t(zbot))))
    #-->  regularly-sampled twt scale and its depth (MD) equivalent on the basis of depth-time rel.
    twt = np.arange(0, tt.max(), dt)
    zt = np.interp(x=twt, xp=tt, fp=zz, left=np.nan, right=np.nan)

    #-->  resample logs to twt
    WWt=pd.DataFrame(data=zt, columns=['DEPTH'], index=twt)
    WWt.index.rename('TWT',inplace=True)
    loglist=WW.columns
    for i in loglist:
        tmp = np.interp(x=twt, xp=twt_log, fp=WW[i][ff].values, left=np.NaN, right=np.NaN)
        WWt=pd.concat([WWt,pd.Series(tmp, index=twt, name=i)],axis=1)
        WWt.interpolate(inplace=True)
        WWt.fillna(method = 'bfill',inplace=True)

    #--> QC plot with IP in depth and time
    if qcplot:
        tmp_IP = WW['VP']*WW['RHO']
        tmp_IPt = WWt['VP']*WWt['RHO']
        plotmax = tmp_IP[(z>=ztop) & (z<=zbot)].max()
        plotmin = tmp_IP[(z>=ztop) & (z<=zbot)].min()
        plotmax += plotmax*.1
        plotmin -= plotmin*.1

        f, ax = plt.subplots(nrows=1,ncols=2,figsize=(5,5), facecolor='w')
        ax[0].plot(tmp_IP, z, '-k')
        ax[1].plot(tmp_IPt, WWt.index, '-k')
        ax[0].set_xlabel('AI [m/s*g/cc]'), ax[0].set_ylabel('MD [m]')
        ax[1].set_xlabel('AI [m/s*g/cc]'), ax[1].set_ylabel('TWT [s]')
        ax[1].yaxis.set_label_position('right')
        ax[1].yaxis.set_ticks_position('right')
        ax[0].set_ylim(ztop,zbot)
        ax[1].set_ylim(d2t(ztop),d2t(zbot))
        for aa in ax.flatten():
            aa.invert_yaxis()
            aa.grid()
            aa.xaxis.tick_top()
            plt.setp(aa.xaxis.get_majorticklabels(), rotation=90, fontsize=8)
            # aa.set_xlim(plotmin,plotmax)
        if flagtops: # plot top markers on all columns
            for topn,topz in tops.iteritems():
                if (topz>=ztop) & (topz<=zbot):
                    ax[0].axhline(y=topz,color=color_top,alpha=alpha_top)
                    ax[0].text(x=plotmax,y=topz,s=topn,**format_tops)
                    ax[1].axhline(y=d2t(topz),color=color_top,alpha=alpha_top)
                    ax[1].text(x=plotmax,y=d2t(topz),s=topn,**format_tops)
        if name is not None:
            plt.suptitle(name, **format_title)
        plt.tight_layout()
    return WWt,float(d2t(ztop)),float(d2t(zbot))

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def shuey(vp1, vs1, rho1, vp2, vs2, rho2, theta, approx=True, terms=False):
    '''
    shuey (C) aadm 2016
    Calculates P-wave reflectivity with Shuey's equation

    reference:
    Avseth et al. (2005), Quantitative Seismic Interpretation, Cambridge University Press (p.182)

    INPUT
    vp1, vs1, rho1: P-, S-wave velocity (m/s) and density (g/cm3) of upper medium
    vp2, vs2, rho2: P-, S-wave velocity (m/s) and density (g/cm3) of lower medium
    theta: angle of incidence (degree)
    approx: returns approximate (2-term) form (default: True)
    terms: returns reflectivity, intercept and gradient (default: False)

    OUTPUT
    reflectivity (and optionally intercept, gradient; see terms option) at angle theta
    '''
    a = np.radians(theta)
    dvp = vp2-vp1
    dvs = vs2-vs1
    drho = rho2-rho1
    vp  = np.mean([vp1,vp2])
    vs  = np.mean([vs1,vs2])
    rho = np.mean([rho1,rho2])
    R0 = 0.5*(dvp/vp + drho/rho)
    G  = 0.5*(dvp/vp) - 2*(vs**2/vp**2)*(drho/rho+2*(dvs/vs))
    F =  0.5*(dvp/vp)
    if approx:
        R = R0 + G*np.sin(a)**2
    else:
        R = R0 + G*np.sin(a)**2 + F*(np.tan(a)**2-np.sin(a)**2)
    if terms:
        return R,R0,G
    else:
        return R

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def akirichards(vp1, vs1, rho1, vp2, vs2, rho2, theta):
    '''
    Aki-Richards (C) aadm 2017
    Calculates P-wave reflectivity with Aki-Richards approximate equation
    only valid for small layer contrasts.

    reference:
    Mavko et al. (2009), The Rock Physics Handbook, Cambridge University Press (p.182)

    INPUT
    vp1, vs1, rho1: P-, S-wave velocity (m/s) and density (g/cm3) of upper medium
    vp2, vs2, rho2: P-, S-wave velocity (m/s) and density (g/cm3) of lower medium
    theta: angle of incidence (degree)

    OUTPUT
    reflectivity at angle theta
    '''
    a = np.radians(theta)
    p = np.sin(a)/vp1
    dvp = vp2-vp1
    dvs = vs2-vs1
    drho = rho2-rho1
    vp  = np.mean([vp1,vp2])
    vs  = np.mean([vs1,vs2])
    rho = np.mean([rho1,rho2])
    A = 0.5*(1-4*p**2*vs**2)*drho/rho
    B = 1/(2*np.cos(a)**2) * dvp/vp
    C = 4*p**2*vs**2*dvs/vs
    R = A + B - C
    return R

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def quicklook(WW,ztop=None,zbot=None,name=None,tops=None):
    '''
    quicklook (C) aadm 2015-2018
    Summary well plot with raw and processed logs.

    INPUT
    WW: Pandas dataframe with VP, [VS], RHO, IP, [VPVS or PR], [SWE], PHIE, VSH
    ztop,zbot: depth range to plot (defaults to min,max depth)
    name: well name (or anything else) to print
    tops: dictionary containing stratigraphic tops, e.g.: tops={'Trias': 1200,'Permian': 2310}
          or Pandas Series, e.g: tops=pd.Series({'Trias': 1200,'Permian': 2310})
    '''
    if ztop is None: ztop = WW.index.min()
    if zbot is None: zbot = WW.index.max()
    WW=WW[(WW.index>=ztop) & (WW.index<=zbot)]
    WW.columns=WW.columns.str.upper()

    flagtops=False if tops is None else True
    if flagtops:
        if not isinstance(tops, pd.Series):
            tops=pd.Series(tops)
        tops=tops.dropna().sort_values()

    f, ax = plt.subplots(nrows=1,ncols=4,sharey=True,figsize=(8,5))
    ax[0].plot(WW.VSH, WW.index, color='.5')
    ax[0].set_xlabel('Vsh', color='.5')
    ax[1].plot(WW.PHIE, WW.index, color='.5')
    ax[1].set_xlabel('phi', color='.5')
    ax[2].plot(WW.VP*WW.RHO, WW.index, 'black')
    ax[2].set_xlabel('AI [m/s*g/cc]')
    ax[3].plot(WW.VP/WW.VS, WW.index, 'black')
    ax[3].set_xlabel('Vp/Vs')
    rlims = ax[3].get_xlim()
    for i,aa in enumerate(ax):
        if flagtops: # plot top markers on all columns
            for topn,topz in tops.iteritems():
                if (topz>=ztop) & (topz<=zbot):
                    aa.axhline(y=topz,color='blue',alpha=.8)
                    if i is 3: # plot top name also on last column
                        aa.text(x=rlims[1]+rlims[1]*.01,y=topz,s=topn,**format_tops)
    if name is not None:
        plt.suptitle(name, **format_title)
    ax[0].set_ylim(zbot,ztop)

max_stars_repo_name: aadm/ASAP

max_issues_repo_name: aadm/ASAP

max_forks_repo_name: aadm/ASAP

reason: import numpy,from scipy