text:
"""
Internal wave functions
"""

import numpy as np
from scipy import linalg, sparse
from scipy.interpolate import interp1d
from scipy.integrate import solve_bvp
from scipy.optimize import least_squares, leastsq

import pdb

GRAV = 9.81
RHO0 = 1020.

###########
# Wave shape
###########
def gaussian(x, a_0, L_w):
    sigma = L_w/4
    return -a_0 * np.exp( - (x/sigma)**2. )

def sine(x, a_0, L_w, x0=0.):

    k = 2*np.pi/L_w
    eta = -a_0/2 - a_0/2 * np.sin(k*x + k*x0 + np.pi/2)
    eta[x>x0+L_w/2] = 0.
    eta[x<x0-L_w/2] = 0.

    return eta

def wave_eta(x, a_0, c_n, L_w, wavefunc=gaussian, **kwargs):
    """
    Initial gaussian wave
    """
    #return -a_0 *c_n* np.exp( - (x/L_w)**2. )
    return wavefunc(x, a_0, L_w, **kwargs)

def wave_init(x, rhoz, dz, d, a_0, L_w, mode=0, wavefunc=gaussian, **kwargs):
    """
    Initialise a wavefield
    """

    phi, cn, drho_dz = iwave_modes(rhoz, dz, d)

    #drho_dz = np.gradient(rhoz, -dz)

    eta = wave_eta(x, a_0, np.real(cn[mode]), L_w, wavefunc=wavefunc, **kwargs)

    phi_n = phi[:,mode].squeeze()
    phi_n /= np.abs(phi_n).max()
    phi_n *= np.sign(phi_n[1])

    rho_pr = eta*drho_dz[:,np.newaxis]*phi_n[:,np.newaxis]

    return rhoz[:,np.newaxis] - rho_pr, phi_n

def wave_init_phi(x, rhoz, drho_dz, phi_n, cn, z, d, a_0, L_w, mode=0):
    """
    Proper way to initialize the wavefield
    """

    #phi, dphi, cn = iwave_modes(rhoz, dz, d)
    Z = z[...,np.newaxis]
    #drho_dz = np.gradient(rhoz, -dz)

    eta = wave_eta(x, a_0, cn, L_w)

    #phi_n = phi[:,mode].squeeze()
    phi = phi_n / np.abs(phi_n).max()
    phi *= np.sign(phi_n.sum())

    #rho_pr = eta*drho_dz[:,np.newaxis]*phi[:,np.newaxis]
    eta_pr = eta*phi[:,np.newaxis]

    #print z.shape, rhoz.shape
    # Interpolation function
    Frho = interp1d(z, rhoz, axis=0)

    eta = z[:,np.newaxis] - eta_pr

    eta[eta>0.] = 0.
    eta[eta<-d] = -d

    # Find rho by interpolating eta
    return Frho(eta), phi
    #return rhoz[:,np.newaxis] - rho_pr, phi

#####
# Nondimensional parameters (Lamb definitions)
#####
def calc_alpha(phi, c, N2, dz):
    """
    Holloway et al 1999 nonlinearity parameter
    """
    phi_z = np.gradient(phi,-np.abs(dz))
    num = 3*c*np.trapz( phi_z**3., dx=dz)
    den = 2*np.trapz( phi_z**2., dx=dz)

    return num/den

def calc_r20(phi, c, N2, dz):
    phi_z = np.gradient(phi,-np.abs(dz))
    S_20 = calc_S20(phi, c, N2, dz)

    num = c*np.trapz( phi*S_20, dx=dz)
    den = 2*np.trapz( phi_z**2., dx=dz)
    return num/den

def calc_alpha_wshear(phi, c, U, dz):
    """
    alpha with shear (see Stastna and Lamb 2002)
    """
    # Stastna and Lamb defn
    E = c/(c-U) * phi
    E_z = np.gradient(E, -np.abs(dz))

    num = 3*np.trapz((c-U)**2. * E_z**3., dx = np.abs(dz))
    den = 2*np.trapz((c-U) * E_z**2., dx = np.abs(dz))

    return num/den

def calc_alpha_wshear_liu(phi, c, U, dz):
    """
    alpha with shear (see Liu et al 1988)
    """
    # Liu et al definition
    phi_z = np.gradient(phi, -np.abs(dz))
    Uz = np.gradient(U, -np.abs(dz))

    num = 3*c*np.trapz( 1/(c-U) * (phi_z - Uz/(U-c)*phi)**3., dx = np.abs(dz))
    den = 2*np.trapz( 1/(c-U) * (phi_z - Uz/(U-c)*phi)**2., dx = np.abs(dz))

    return num/den

def calc_alpha_wshear_old(phi, c, U, dz):
    """
    alpha with shear (see Grimshaw 2004)
    """
    phi_z = np.gradient(phi,-np.abs(dz))
    num = 3*np.trapz((c-U)**2. * phi_z**3., dx=np.abs(dz))
    den = 2*np.trapz( (c-U) * phi_z**2., dx=np.abs(dz))

    return num/den

def calc_beta_wshear(phi, c, U, dz):
    phi_z = np.gradient(phi, -np.abs(dz))
    num = np.trapz( (c-U)**2. * phi**2., dx=np.abs(dz))
    den = 2*np.trapz( (c-U) * phi_z**2., dx=np.abs(dz))

    return num/den

def calc_r10(phi, c, N2, dz):
    """
    alpha in most other papers
    """
    phi_z = np.gradient(phi,-np.abs(dz))
    num = -3*np.trapz( phi_z**3., dx=np.abs(dz))
    den = 4*np.trapz( phi_z**2., dx=np.abs(dz))

    #N2_z = np.gradient(N2,-dz)
    #S10 = N2_z/c**3*phi**2
    #num = c*np.trapz(phi*S10, dx=dz)
    #den = 2*np.trapz(phi_z**2., dx=dz)
    return num/den

def calc_r01(phi, c, dz):
    phi_z = np.gradient(phi, -np.abs(dz))
    num = -c*np.trapz( phi**2., dx=np.abs(dz))
    den = 2*np.trapz( phi_z**2., dx=np.abs(dz))
    #num = -c*np.sum( phi**2. * dz)
    #den = 2*np.sum( phi_z**2. * dz)

    return num/den

def calc_phi10(phi, c, N2, dz):

    kmax = np.argwhere(phi==phi.max())[0][0]
    phi10rhs = calc_phi10_rhs(phi, c, N2, dz)
    phi10 = solve_phi_bvp(phi10rhs, N2, c, dz, kmax=kmax)

    ## Normalize
    #alpha = -phi10[kmax]
    #phi10 += alpha*phi

    return phi10

def calc_phi01(phi, c, N2, dz):

    kmax = np.argwhere(phi==phi.max())[0][0]
    phi01rhs = calc_phi01_rhs(phi, c, N2, dz)
    phi01 = solve_phi_bvp(phi01rhs, N2, c, dz, kmax=kmax)

    # Normalize
    #alpha = -phi01[kmax]
    #phi01 += alpha*phi

    return phi01

def calc_phi20(phi, c, N2, dz):

    kmax = np.argwhere(phi==phi.max())[0][0]
    phi20rhs = calc_phi20_rhs(phi, c, N2, dz)
    phi20 = solve_phi_bvp(phi20rhs, N2, c, dz, kmax=kmax)

    ## Normalize
    #alpha = -phi20[kmax]
    #phi20 += alpha*phi

    return phi20

def calc_phi01_rhs(phi, c, N2, dz):
    r_01 = calc_r01(phi, c, dz)

    RHS = -2*r_01*N2/c**3.*phi - phi

    return RHS

def calc_phi10_rhs(phi, c, N2, dz):
    r_10 = calc_r10(phi, c, N2, dz)

    dN2_dz = np.gradient(N2,-np.abs(dz))
    c3i = 1/c**3.

    RHS = -2*r_10*N2*c3i * phi
    RHS += c3i*dN2_dz*phi**2.

    return RHS

def calc_phi20_rhs(phi, c, N2, dz):
    r_20 = calc_r20(phi, c, N2, dz)
    S_20 = calc_S20(phi, c, N2, dz)

    return -r_20*2*N2/c**3. * phi + S_20

def calc_D01(phi, c, N2, dz):
    """
    Calculates the first order nonlinear term for buoyancy
    """

    r_01 = calc_r01(phi, c, dz)

    phi01 = calc_phi01(phi, c, N2, dz)

    D01 = N2/c*phi01 + r_01*N2/c**2.*phi

    return D01#/N2

def calc_D10(phi, c, N2, dz):
    """
    Calculates the first order nonlinear term for buoyancy
    """

    dN2_dz = np.gradient(N2,-np.abs(dz))
    r_10 = calc_r10(phi, c, N2, dz)

    phi10 = calc_phi10(phi, c, N2, dz)

    D10 = N2/c*phi10 + r_10*N2/c**2.*phi - 1/(2*c**2)*dN2_dz*phi**2.

    return D10 #/N2

def calc_D20(phi, c, N2, dz):
    """
    Calculates the second order nonlinear term for buoyancy
    """
    r_20 = calc_r20(phi, c, N2, dz)

    phi20 = calc_phi20(phi, c, N2, dz)

    T_20 = calc_T20(phi, c, N2, dz)

    D20 = N2/c*phi20 + r_20*N2/c**2.*phi + T_20

    return D20

def calc_S20(phi, c, N2, dz):
    """
    Calculates the second order term
    """

    dN2_dz = np.gradient(N2, -dz)
    d2N2_dz2 = np.gradient(dN2_dz, -dz)

    phi_z = np.gradient(phi, -dz)

    r_10 = calc_r10(phi, c, N2, dz)

    phi10 = calc_phi10(phi, c, N2, dz)

    S20 = 2/c**3.*dN2_dz * phi * phi10\
        - 1/(2*c**4) * d2N2_dz2 * phi**3 \
        + 3*r_10/c**4 * dN2_dz * phi**2 \
        - r_10*N2/c**4 * phi * phi_z \
        -8/3. * r_10 * N2/c**3 * phi10 \
        -4*r_10**2*N2/c**4 * phi

    return S20#/N2

def calc_T20(phi, c, N2, dz):
    """
    Calculates the second order term
    """

    dN2_dz = np.gradient(N2, -dz)
    d2N2_dz2 = np.gradient(dN2_dz, -dz)

    phi_z = np.gradient(phi, -dz)

    r10 = calc_r10(phi, c, N2, dz)

    phi10 = calc_phi10(phi, c, N2, dz)

    return -dN2_dz/c**2.*phi*phi10\
        -r10*dN2_dz/c**3.*phi**2.\
        +1/6.*d2N2_dz2/c**3.*phi**3.\
        +r10/3.*N2/c**3. * phi * phi_z\
        +4/3.*r10*N2/c**2*phi10\
        +4/3.*r10**2.*N2/c**3.*phi

def calc_T10(phi, c, N2, dz):
    """
    Calculates the Grimshaw and co nonlinear correction term
    """

    dN2_dz = np.gradient(N2,-np.abs(dz))
    alpha = calc_alpha(phi, c, N2, dz)

    RHS = alpha*N2/c**4. * phi
    RHS += dN2_dz/c**3. *phi**2.

    kmax = np.argwhere(phi==phi.max())[0][0]
    T10 = solve_phi_bvp(RHS, N2, c, dz, kmax=kmax)

    # Normalize
    #alpha = -T10[kmax]
    #T10 += alpha*phi

    # normalize ??
    #T10 = T10 * phi/np.mean(phi)
    #T10 = T10 * np.mean(phi)/phi

    # Normalize so the max(phi)=1
    #T10 = T10 / np.abs(T10).max()
    #T10 *= np.sign(T10.sum())

    #T10 *= (1-phi)

    return T10

#####
# Nondimensional parameters
#####
def wave_delta_star(phi, dz):
    """
    Nonlinearity parameter - Liu, 1988
    """

    dphi = np.gradient(phi, dz)

    num = np.trapz( np.power(dphi,3.), dx=dz)
    den = np.trapz( np.power(dphi,2.), dx=dz)

    return 1.5*num/den

def wave_delta(phi, dz, a0):
    """
    Nonlinearity parameter
    """
    dphi = np.gradient(phi, dz)

    num = np.trapz( np.power(dphi,3.), dx=dz)
    den = np.trapz( np.power(dphi,2.), dx=dz)

    return -a0 * num/den

def wave_epsilon_star(phi, dz, C0):
    """
    Wave dispersion (nonhydrostasy) parameter - Liu 1988
    """

    dphi = np.gradient(phi, dz)

    num = np.trapz( phi**2., dx=dz)
    den = np.trapz( dphi**2., dx=dz)

    return  C0/2 * num / den

def wave_epsilon(phi, dz, Lw):
    """
    Wave dispersion (nonhydrostasy) parameter
    """

    dphi = np.gradient(phi, dz)

    num = np.trapz( phi**2., dx=dz)
    den = np.trapz( dphi**2., dx=dz)

    return np.sqrt( 3.0/Lw**2 * num / den)

def wave_he(phi, dz):
    """
    Equivalent layer height
    """

    dphi = np.gradient(phi, dz)

    num = np.trapz( phi**2., dx=dz)
    den = np.trapz( dphi**2., dx=dz)

    return np.sqrt(3.0 * num / den)

#####
# Boundary value problem solvers
#####
def solve_phi_bvp_newton(B, N2, c, dz):
    """
    Use the scipy integration function (Uses Newton method)
    """
    nx = N2.shape[0]
    x = np.linspace(0, nx*dz, nx)
    k = -N2/c**2.
    Fk = interp1d(x, k, kind=2)
    Fb = interp1d(x, B, kind=2)
    def fun(x, y, p):
        return np.vstack([y[1], Fk(x) * y[0] + Fb(x)])

    def bc(ya, yb, p):
        # BC: y(0) = y(H) = 0 and y'(0) = 0
        #return np.array([ ya[0], yb[0], ya[1] ])
        # BC: y(0) = y(H) = 0
        return np.array([ ya[0], yb[0], yb[1] ])

    y = np.zeros((2,x.size))

    res_a = solve_bvp(fun, bc, x, y, p=[0], tol=1e-10)

    return res_a.sol(x)[0,:]

def solve_phi_bvp(B, N2, c, dz, kmax=None):
    """
    Finite-difference  solver for the higher-order vertical structure
    functions

    """
    # Interpolate onto the mid-point
    Bmid = 0.5*(B[1:] + B[:-1])
    N2mid = 0.5*(N2[1:] + N2[:-1])
    nz = Bmid.shape[0]
    dz2 = 1/dz**2

    ### Construct the LHS matrix, A
    # (use a dense matrix for now)
    A = np.diag(dz2*np.ones((nz-1)),-1) +\
        np.diag(-2*dz2*np.ones((nz,)) +\
        N2mid/c**2*np.ones((nz,)) ,0) + \
        np.diag(dz2*np.ones((nz-1)),1)

    ## BC's
    A[0,0] = -3*dz2 + N2mid[0]/c**2 # Dirichlet BC
    A[-1,-1] = -3*dz2 + N2mid[-1]/c**2 # Dirichlet BC

    # Apply normalization condition at kmax
    if kmax is not None:
        A[kmax, kmax] = 0.

    #soln = linalg.solve(A, Bmid)

    # Solve with optimization routines
    def minfun(x0):
        return A.dot(x0) - Bmid

    soln, covc, info, mesg, status = leastsq(minfun, np.zeros_like(Bmid),\
        xtol=1e-12, full_output=True)

    #if status != 2:
    #    print mesg
    #    # These are bad eggs... zero them as its also a solution??
    #    soln *= 0.

    # Interpolate back onto the grid points
    phi = np.zeros((nz+1,))
    soln_mid = 0.5*(soln[1:] + soln[:-1])
    # Zeros are BCs
    phi[1:-1] = soln_mid

    return phi

def _oldsolve_phi_bvp_fd(B, N2, c, dz):
    """
    Finite-difference  solver for the higher-order vertical structure
    functions

    !!!!DOES NOT WORK PROPERLY!!!!
    """

    nz = B.shape[0]
    dz2 = 1/dz**2

    ### Construct the LHS matrix, A
    # (use a dense matrix for now)
    A = np.diag(dz2*np.ones((nz-1)),-1) +\
        np.diag(-2*dz2*np.ones((nz,)) +\
        N2/c**2*np.ones((nz,)) ,0) + \
        np.diag(dz2*np.ones((nz-1)),1)

    ## BC's
    #eps = 1e-10
    #A[0,0] = -1
    #A[0,1] = eps
    #A[-1,-1] = -1
    #A[-1,-2] = eps

    #B[0] = eps
    #B[-1] = eps

    return linalg.solve(A, B)

    #nz = B.shape[0] - 2
    #dz2 = 1/dz**2

    #### Construct the LHS matrix, A
    ## (use a dense matrix for now)
    #A = np.diag(dz2*np.ones((nz-1)),-1) +\
    #    np.diag(-2*dz2*np.ones((nz,)) +\
    #    N2[1:-1]/c**2*np.ones((nz,)) ,0) + \
    #    np.diag(dz2*np.ones((nz-1)),1)

    ## BC's
    ##eps = 1e-10
    ##A[0,0] = eps
    ##A[0,1] = eps
    ##A[-1,-1] = eps
    ##A[-1,-2] = eps

    #soln = np.zeros((nz+2,))
    #soln[1:-1] = linalg.solve(A, B[1:-1])
    #return soln

    #### Use the direct banded solver
    #A = np.zeros((3,nz))
    #A[0,:] = dz2*np.ones(nz)
    #A[1,:] = -2*dz2*np.ones(nz)
    #A[1,:] += N2[1:-1]/c**2*np.ones(nz)
    #A[2,:] = dz2*np.ones(nz)
    #
    ##A[0,0]=eps
    ##A[1,0]=eps
    ##A[2,-1] = eps
    ##A[1,-1] = eps
    #
    #soln = np.zeros((nz+2,))
    #soln[1:-1] = linalg.solve_banded((1,1), A, B[1:-1])
    #return soln

#####
# Eigenvalue solver functions
#####
def iwave_modes_banded(N2, dz, k=None):
    """
    !!! DOES NOT WORK!!!
    Calculates the eigenvalues and eigenfunctions to the internal wave eigenvalue problem:

    $$
    \left[ \frac{d^2}{dz^2} - \frac{1}{c_0} \bar{\rho}_z \right] \phi = 0
    $$

    with boundary conditions
    """

    nz = N2.shape[0] # Remove the surface values
    if k is None:
        k = nz-2

    dz2 = 1/dz**2

    # Construct the LHS matrix, A
    A = np.vstack([-1*dz2*np.ones((nz,)),\
        2*dz2*np.ones((nz,)),\
        -1*dz2*np.ones((nz,)),\
    ])

    # BC's
    #A[0,0] = -1.
    #A[0,1] = 0.
    #A[-1,-1] = -1.
    #A[-1,-2] = 0.
    A[1,0] = -1.
    A[2,0] = 0.
    A[1,-1] = -1.
    A[0,-1] = 0.

    # Now convert from a generalized eigenvalue problem to
    #       A.v = lambda.B.v
    # a standard problem
    #       A.v = lambda.v
    # By multiply the LHS by inverse of B
    #       (B^-1.A).v = lambda.v
    # B^-1 = 1/N2 since B is diagonal
    A[0,:] /= N2
    A[1,:] /= N2
    A[2,:] /= N2

    w, phi = linalg.eig_banded(A)

    pdb.set_trace()

    ## Main diagonal
    #dd = 2*dz2*np.ones((nz,))

    #dd /= N2

    #dd[0] = -1
    #dd[-1] = -1

    ## Off diagonal
    #ee = -1*dz2*np.ones((nz-1,))
    #ee /= N2[0:-1]

    #ee[0] = 0
    #ee[-1] = 0

    ## Solve... (use scipy not numpy)
    #w, phi = linalg.eigh_tridiagonal(dd, ee )

    #####

    c = 1. / np.power(w, 0.5) # since term is ... + N^2/c^2 \phi

    # Sort by the eigenvalues
    idx = np.argsort(c)[::-1] # descending order

    ## Calculate the actual phase speed
    cn = np.real( c[idx] )

    idxgood = ~np.isnan(cn)
    phisort = phi[:,idx]

    return np.real(phisort[:,idxgood]), np.real(cn[idxgood])

def iwave_modes_sparse(N2, dz, k=None, Asparse=None, return_A=False):
    """
    Calculates the eigenvalues and eigenfunctions to the internal wave eigenvalue problem:

    $$
    \left[ \frac{d^2}{dz^2} - \frac{1}{c_0} \bar{\rho}_z \right] \phi = 0
    $$

    with boundary conditions
    """

    nz = N2.shape[0] # Remove the surface values
    if k is None:
        k = nz-2

    if Asparse is None:
        dz2 = 1/dz**2

        # Construct the LHS matrix, A
        A = np.vstack([-1*dz2*np.ones((nz,)),\
                2*dz2*np.ones((nz,)),\
                -1*dz2*np.ones((nz,)),\
        ])

        # BC's
        eps = 1e-10
        #A[0,0] = -1.
        #A[0,1] = 0.
        #A[-1,-1] = -1.
        #A[-1,-2] = 0.
        A[1,0] = -1.
        A[2,0] = 0.
        A[1,-1] = -1.
        A[0,-1] = 0.

        Asparse = sparse.spdiags(A,[-1,0,1],nz,nz, format='csc')

    # Construct the RHS matrix i.e. put N^2 along diagonals
    #B = np.diag(N2,0)
    B = sparse.spdiags(N2,[0],nz,nz, format='csc')
    Binv = sparse.spdiags(1/N2,[0],nz,nz, format='csc')
    w, phi = sparse.linalg.eigsh(Asparse, M=B, Minv=Binv, which='SM', k=k)

    # Solve... (use scipy not numpy)
    #w, phi = linalg.eig(A, b=B)

    c = 1. / np.power(w, 0.5) # since term is ... + N^2/c^2 \phi

    # Sort by the eigenvalues
    #idx = np.argsort(c)[::-1] # descending order

    ## Calculate the actual phase speed
    #cn = np.real( c[idx] )

    if return_A:
        return np.real(phi), np.real(c), Asparse
    else:
        return np.real(phi), np.real(c)

def iwave_modes(N2, dz, k=None):
    """
    Calculates the eigenvalues and eigenfunctions to the internal wave eigenvalue problem:

    $$
    \left[ \frac{d^2}{dz^2} - \frac{1}{c_0} \bar{\rho}_z \right] \phi = 0
    $$

    with boundary conditions
    """

    nz = N2.shape[0] # Remove the surface values
    dz2 = 1/dz**2

    # Construct the LHS matrix, A
    A = np.diag(-1*dz2*np.ones((nz-1)),-1) + \
        np.diag(2*dz2*np.ones((nz,)),0) + \
        np.diag(-1*dz2*np.ones((nz-1)),1)

    # BC's
    A[0,0] = -1.
    A[0,1] = 0.
    A[-1,-1] = -1.
    A[-1,-2] = 0.

    # Construct the RHS matrix i.e. put N^2 along diagonals
    B = np.diag(N2,0)

    # Solve... (use scipy not numpy)
    w, phi = linalg.eig(A, b=B)

    c = 1. / np.power(w, 0.5) # since term is ... + N^2/c^2 \phi

    # Sort by the eigenvalues
    idx = np.argsort(c)[::-1] # descending order

    # Calculate the actual phase speed
    cn = np.real( c[idx] )

    return phi[:,idx], cn

def iwave_modes_uneven(N2, z):
    """
    Calculates the eigenvalues and eigenfunctions to the internal wave eigenvalue problem:

    $$
    \left[ \frac{d^2}{dz^2} - \frac{1}{c_0} \bar{\rho}_z \right] \phi = 0
    $$

    with boundary conditions
    """

    nz = N2.shape[0]

    dz = np.zeros((nz,))
    zm = np.zeros((nz,))
    dzm = np.zeros((nz,))

    dz[0:-1] = z[0:-1] - z[1:]
    zm[0:-1] = z[0:-1] - 0.5*dz[0:-1]

    dzm[1:-1] = zm[0:-2] - zm[1:-1]
    dzm[0] = dzm[1]
    dzm[-1] = dzm[-2]

    A = np.zeros((nz,nz))
    for i in range(1,nz-1):
        A[i,i] = 1/ (dz[i-1]*dzm[i]) + 1/(dz[i]*dzm[i])
        A[i,i-1] = -1/(dz[i-1]*dzm[i])
        A[i,i+1] = -1/(dz[i]*dzm[i])

    # BC's
    eps = 1e-10
    A[0,0] = -1.
    A[0,1] = 0.
    A[-1,-1] = -1.
    A[-1,-2] = 0.

    # Construct the RHS matrix i.e. put N^2 along diagonals
    B = np.diag(N2,0)

    # Solve... (use scipy not numpy)
    w, phi = linalg.eig(A, b=B)

    c = 1. / np.power(w, 0.5) # since term is ... + N^2/c^2 \phi

    # Sort by the eigenvalues
    idx = np.argsort(c)[::-1] # descending order

    # Calculate the actual phase speed
    cn = np.real( c[idx] )

    phiall = phi[:,idx]

    # Normalize so the max(phi)=1
    for ii in range(nz):
        phi_1 = phiall[:,ii]
        phi_1 = phi_1 / np.abs(phi_1).max()
        phi_1 *= np.sign(phi_1.sum())
        phiall[:,ii] = phi_1

    return phiall, cn

def iwave_modes_nondim(N2, dz, d, nondim=True):
    """
    Calculates the eigenvalues and eigenfunctions to the internal wave eigenvalue problem:

    $$
    \left[ \frac{d^2}{dz^2} - \frac{1}{c_0} \bar{\rho}_z \right] \phi = 0
    $$

    with boundary conditions
    """

    nz = N2.shape[0] # Remove the surface values
    dz2 = 1/dz**2

    # Construct the LHS matrix, A
    A = np.diag(-1*dz2*np.ones((nz-1)),-1) + \
        np.diag(2*dz2*np.ones((nz,)),0) + \
        np.diag(-1*dz2*np.ones((nz-1)),1)

    # BC's
    eps = 1e-10
    A[0,0] = -1.
    A[0,1] = 0.
    A[-1,-1] = -1.
    A[-1,-2] = 0.

    # Construct the RHS matrix i.e. put N^2 along diagonals
    B = np.diag(N2,0)

    # Solve... (use scipy not numpy)
    w, phi = linalg.eig(A, b=B)

    c = 1. / np.power(w, 0.5) # since term is ... + N^2/c^2 \phi

    # Sort by the eigenvalues
    idx = np.argsort(c)[::-1] # descending order

    # Calculate the actual phase speed
    cn = np.real( c[idx] )

    return phi[:,idx], cn

max_stars_repo_name: mrayson/iwaves

max_issues_repo_name: mrayson/iwaves

max_forks_repo_name: iosonobert/iwaves


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Provide the various kinematic constraints used in QP.

References:
    [1] "Quadratic Programming in Python" (https://scaron.info/blog/quadratic-programming-in-python.html), Caron, 2017
    [2] "OpenSoT: A whole-body control library for the compliant humanoid robot COMAN", Rocchi et al., 2015
    [3] "Robot Control for Dummies: Insights and Examples using OpenSoT", Hoffman et al., 2017
"""

import rbdl
import numpy as np

from pyrobolearn.robots.robot import Robot
from pyrobolearn.priorities.constraints.constraint import Constraint

__author__ = "Brian Delhaisse"
__copyright__ = "Copyright 2018, PyRoboLearn"
__credits__ = ["OpenSoT (Enrico Mingo Hoffman and Alessio Rocchi)", "Songyan Xin"]
__license__ = "GNU GPLv3"
__version__ = "1.0.0"
__maintainer__ = "Brian Delhaisse"
__email__ = "briandelhaisse@gmail.com"
__status__ = "Development"

class KinematicConstraint(Constraint):
    r"""Kinematic constraint

    """
    pass

class CartesianPose(KinematicConstraint):
    r"""Cartesian Pose constraint.

    """
    pass

class CartesianVelocity(KinematicConstraint):
    r"""Cartesian velocity constraint.

    """
    pass

class CoMVelocity(KinematicConstraint):
    r"""Center-of-mass velocity constraint.

    """
    pass

class JointLimits(KinematicConstraint):
    r"""Joint limits constraint.

    """
    pass

class JointVelocityLimits(KinematicConstraint):
    r"""Velocity limits constraint.

    """
    pass

class SelfCollisionAvoidance(KinematicConstraint):
    r"""Self-collision avoidance constraint.

    """
    pass

class ConvexHull(Constraint):
    r"""Convex Hull constraint

    """
    pass

ext: py

max_stars_repo_name: Pandinosaurus/pyrobolearn

max_issues_repo_name: Pandinosaurus/pyrobolearn

max_forks_repo_name: Pandinosaurus/pyrobolearn

reason: import numpy


text:
import numpy as np
from scipy.ndimage import gaussian_filter
import sys
from scipy.spatial import Delaunay
import os
import subprocess
from matplotlib import image as mpimg
from PIL import Image, ImageDraw
import shutil
import cv2
import time

GAUSS = 2
OUTDIR = 'temp4/'
PERSISTENCE_THRESHOLD = 512
l = 512
w = 512

MAX_INT_PERC = 95

def output_edge_scores(paths, st):
    output_edge = []
    for i in range(len(paths)):
        p = paths[i]
        for j in range(len(p) - 1):
            v1 = p[j]
            v2 = p[j + 1]
            output_edge.append([v1, v2, st[i]])
    return output_edge

def dist(p1, p2):
    return (((p2[0] - p1[0]) ** 2) + ((p2[1] - p1[1]) ** 2)) ** .5

def line_function(val, cap):
    capped_val = min(val, cap)
    return capped_val

def length(p):
    sum = 0
    for i in range(len(p) - 1):
        i_dist = dist(p[i], p[i + 1])
        sum += i_dist
    return sum

def cmp_dm_img_tri2D(output_dir, input_filename):
    # print('running morse')
    t_start = time.time()
    subprocess.check_call([r"src/a.out", output_dir + 'input.png', output_dir, str(PERSISTENCE_THRESHOLD)])
    t_end = time.time()
    print("only morse:", t_end - t_start)
    # print('process done')
    o_vert = np.loadtxt(output_dir + "/vert.txt")
    o_edge = np.loadtxt(output_dir + "/edge.txt")

    if len(o_edge) == 0:
        return [], []

    o_vert = o_vert[:, :2]

    # visualize result
    stable_vert = o_vert.copy()
    stable_vert[:, 0] = o_vert[:, 1]
    stable_vert[:, 1] = o_vert[:, 0]
    o_vert = stable_vert

    return o_vert, o_edge

def make_png(verts, edges, path, l, w, linestroke=1):
    im = Image.new('L', (l, w), color=0)
    draw = ImageDraw.Draw(im)
    for e in edges:
        bright = int(e[2] * 255)
        u, v = verts[e[0]], verts[e[1]]
        draw.line((u[0], u[1], v[0], v[1]), fill=bright, width=linestroke)

    im.save(path, format='png')

def dm_cal(tile, id):
    start = time.time()
    print(PERSISTENCE_THRESHOLD)
# def dm_cal():
    #print('begin testing')
    img = tile
    '''
    img = []
    for i in range(512):
        row = []
        for j in range(512):
            row.append(0)
        img.append(row)
    '''
    # img = mpimg.imread('imgs/test1.tif')
    # id = 3
    img = np.asarray(img)
    # print('img max:', np.max(img))
    nx, ny = img.shape
    '''
    for i in range(nx):
        for j in range(ny):
            print('py vert:', i, j, img[i, j])

    sys.exit()
    '''
    if GAUSS > 0:
        input_img = gaussian_filter(img, GAUSS)
    else:
        input_img = img

    # print('gauss img max:', np.max(input_img))

    input_img.astype('uint16')

    # print('creating dir')
    if not os.path.exists(OUTDIR):
        os.makedirs(OUTDIR)

    id_dir = OUTDIR + str(id) + '/'
    if not os.path.exists(id_dir):
        os.makedirs(id_dir)

    input_filename = id_dir + '/input.png'

    # print('outputting img for morse')
    cv2.imwrite(input_filename, input_img)

    threshold = PERSISTENCE_THRESHOLD
    verts, edges = cmp_dm_img_tri2D(id_dir, input_filename)
    verts = [(int(v[0]), int(v[1])) for v in verts]

    end = time.time()
    print("old:", end - start)

    no_dup_edges = []
    for edge in edges:
        if edge[0] < edge[1]:
            e = [int(edge[0]), int(edge[1])]
        else:
            e = [int(edge[1]), int(edge[0])]
        if e not in no_dup_edges:
            no_dup_edges.append(e)

    no_dup_edge_filename = id_dir + "no_dup_edge.txt"
    with open(no_dup_edge_filename, 'w') as output_file:
        for e in no_dup_edges:
            output_file.write(str(e[0]) + ' ' + str(e[1]) + '\n')
        output_file.close()

    # sys.exit()

    subprocess.check_call([r"paths_src/a.out", id_dir])

    perc = np.percentile(img, MAX_INT_PERC)
    intensity = {}

    for v in verts:
        x = v[0]
        y = v[1]
        f = img[y, x]
        # print('f', f)
        intensity[v] = f

    # print('max int:', intensity.items())

    above = []
    for i in range(nx):
        for j in range(ny):
            f = img[i, j]
            if f >= perc:
                above.append(f)
    MAX_INT = int(sum(above) / len(above)) + 8

    raw_paths = []
    # sys.stdout.flush()
    # print('reading in paths...')
    # sys.stdout.flush()
    lines = [line.rstrip('\n').split(' ') for line in open(id_dir + 'paths.txt')]

    for i in range(len(lines)):
        lines[i] = lines[i][:len(lines[i]) - 1]

    for line in lines:
        raw_paths.append([int(x) for x in line])

    paths = []
    # sys.stdout.flush()
    # print('computing valid paths...')
    # sys.stdout.flush()
    for p in raw_paths:
        if len(p) <= 1:
            # print('path of len 1 or 0')
            sys.exit()

        edges = [(p[i], p[i + 1]) for i in range(len(p) - 1)]

        '''
        if test_edge not in edges and reflect_edge not in edges:
            continue

        print('edge:', test_edge)
        print('path', p)
        '''

        bad_ind = []
        for i in range(len(edges)):
            edge = edges[i]
            # print('edge:', edge)
            e_len = dist(verts[edge[0]], verts[edge[1]])
            if float(20) < e_len:
                bad_ind.append(i)
        num_bad = len(bad_ind)
        # print('bad:', num_bad)
        if num_bad == 0:
            paths.append(p)
            continue
        if num_bad == len(edges):
            continue
        if num_bad == 1:
            b = bad_ind[0]
            p1 = p[:b + 1]
            p2 = p[b + 1:]
            if len(p1) > 1:
                paths.append(p1)
            if len(p2) > 1:
                paths.append(p2)
            continue

        # print('bad inds', bad_ind)

        b0 = bad_ind[0]
        bl = bad_ind[num_bad - 1]
        p1 = p[:b0 + 1]
        p2 = p[bl + 1:]

        if len(p1) > 1:
            paths.append(p1)
        if len(p2) > 1:
            paths.append(p2)

        if num_bad == 2:
            if len(p[b0 + 1:bl + 1]) > 1:
                paths.append(p[b0 + 1:bl + 1])

        # if num bad == 2 then range is empty
        # print('range:', [i for i in range(1,len(bad_ind) - 1)])
        for i in range(1, len(bad_ind) - 1):
            bim1 = bad_ind[i - 1]
            bi = bad_ind[i]
            bip1 = bad_ind[i + 1]
            p1 = p[bim1 + 1:bi + 1]
            p2 = p[bi + 1: bip1 + 1]
            if len(p1) > 1:
                paths.append(p1)
            if len(p2) > 1:
                paths.append(p2)

    # sys.stdout.flush()
    # print(len(paths), 'valid paths')

    # sys.stdout.flush()
    # print('ready to go!')
    # sys.stdout.flush()
    scores = []
    lengths = []
    degree_dict = {}
    for i in range(len(verts)):
        degree_dict[i] = 0

    p_index = 0
    for path in paths:
        p_index += 1
        # print('path:', p_index,'/',len(paths))
        v_path = [verts[i] for i in path]
        # print(v_path)
        p_len = length(v_path)

        degree_dict[path[0]] += 1
        degree_dict[path[len(path) - 1]] += 1

        lengths.append(p_len)
        # print(p_len)
        f_vals = [intensity[v] for v in v_path]

        score = 0
        vec_score = 0
        for i in range(len(v_path) - 1):
            b1 = line_function(f_vals[i], MAX_INT)
            b2 = line_function(f_vals[i + 1], MAX_INT)
            h = dist(v_path[i], v_path[i + 1])
            area = h * (b1 + b2) / 2

            # print(h, b1, b2)

            score += area
        # print('score for path:', score/p_len)
        scores.append(score / p_len)

    # sys.stdout.flush()
    # print('outputting...')
    # sys.stdout.flush()

    if len(scores) == 0:
        # print('empty morse!')
        empty = []
        for i in range(nx):
            row = []
            for j in range(ny):
                row.append(0)
            empty.append(row)
        empty = np.asarray(empty)
        return empty

    min_score = min(scores)
    max_score = max(scores)
    # sys.stdout.flush()

    print('min max', min_score, max_score)
    # sys.stdout.flush()

    if min_score == 0 and max_score == 0:
        normalized_scores =    [0 for s in scores]
    else:
        normalized_scores = [s / MAX_INT for s in scores]

    # print('scores:', normalized_scores)
    for s in normalized_scores:
        assert (0 <= s <= 1)

    # pe_filename = id_dir + 'simplify_edges.txt'
    output_edges = output_edge_scores(paths, normalized_scores)
    output_path = id_dir + 'g_dimo.png'
    make_png(verts, output_edges, output_path, nx, ny, linestroke=1)
    morse_tile = mpimg.imread(output_path)
    return morse_tile

'''
if __name__ == '__main__':
    dm_cal()
'''

reason: import numpy,from scipy


text:
import triton

import os
import sys
import pdb
# sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)))
import math
import argparse

import torch
import torch.nn.functional as F
import pytest
import numpy
from unittest import mock
from megatron.model.transformer import ParallelAttention, BSMMParallelAttention
from megatron import get_args
from megatron.initialize import initialize_megatron

from megatron.model.enums import AttnMaskType, AttnType
from megatron.model import LayerNorm
from megatron.model.utils import init_method_normal

def assert_almost_equal(x, y, decimal=2, err_msg=''):
    import numpy.testing as npt
    if isinstance(x, torch.Tensor):
        x = x.cpu().detach().numpy()
    if isinstance(y, torch.Tensor):
        y = y.cpu().detach().numpy()
    npt.assert_array_almost_equal(x, y, err_msg=err_msg, decimal=decimal)

def test_attn():
    initialize_megatron(extra_args_provider=None, args_defaults={'tokenizer_type': 'GPT2BPETokenizer','no_load_rng': True, 'no_load_optim': True})
    args = get_args()

    init_method = init_method_normal(args.init_method_std)
    output_layer_init_method = init_method_normal(args.init_method_std)
    layer_number = 1
    #######################
    # self attention test #
    #######################
    self_attn = ParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.self_attn,
            attn_mask_type=AttnMaskType.padding
            )
    torch.save(self_attn.state_dict(), "./self_attn.pt")
    # print(self_attn.state_dict())
    # pdb.set_trace()

    self_attn_bsmm = BSMMParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.self_attn,
            attn_mask_type=AttnMaskType.padding)
    self_attn_bsmm.load_state_dict(torch.load("./self_attn.pt"))

    # hidden_states: [sq, b, h]... 그래서 이게 key야 query야 value야??
    hidden_states = torch.randn(args.seq_length, args.micro_batch_size, args.hidden_size).float().cuda().contiguous()
    attention_mask = torch.tril(torch.ones(args.micro_batch_size, args.seq_length, args.seq_length)).view(args.micro_batch_size, 1, args.seq_length, args.seq_length).cuda().contiguous()
    attention_mask = (attention_mask < 0.5).cuda().contiguous()
    attention_mask_bsmm = torch.tril(torch.ones(args.seq_length, args.seq_length)).cuda()

    # save query_layer,
    encoder_output, encoder_bias = self_attn(hidden_states, attention_mask, layer_past=None, get_key_value=False)
    encoder_output_bsmm, encoder_bias_bsmm = self_attn_bsmm(hidden_states, attention_mask_bsmm, layer_past=None, get_key_value=False)

    assert_almost_equal(encoder_output, encoder_output_bsmm)
    assert_almost_equal(encoder_bias, encoder_bias_bsmm)
    # ########################
    # # cross attention test #
    ########################
    cross_attn = ParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.cross_attn,
            attn_mask_type=AttnMaskType.causal)
    torch.save(cross_attn.state_dict(), "./cross_attn.pt")

    cross_attn_bsmm = BSMMParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.cross_attn,
            attn_mask_type=AttnMaskType.causal)
    cross_attn_bsmm.load_state_dict(torch.load("./cross_attn.pt"))

    hidden_states = torch.randn(args.seq_length, args.micro_batch_size, args.hidden_size).float().cuda().contiguous()
    attention_mask = torch.tril(torch.ones(args.micro_batch_size, args.seq_length, args.seq_length)).view(args.micro_batch_size, 1, args.seq_length, args.seq_length).cuda().contiguous()
    attention_mask = (attention_mask < 0.5).cuda().contiguous()
    attention_mask_bsmm = torch.tril(torch.ones(args.seq_length, args.seq_length)).cuda()

    decoder_output, decoder_bias = cross_attn(hidden_states, attention_mask, layer_past=None, get_key_value=False, encoder_output=encoder_output)

    decoder_output_bsmm, decoder_bias_bsmm = cross_attn_bsmm(hidden_states,
            attention_mask_bsmm,
            layer_past=None,
            get_key_value=False,
            encoder_output=encoder_output_bsmm)

    assert_almost_equal(decoder_output, decoder_output_bsmm)
    assert_almost_equal(decoder_bias, decoder_bias_bsmm)

if __name__ == "__main__":
        test_attn()

max_stars_repo_name: t-hyun/Megatron-LM

max_issues_repo_name: t-hyun/Megatron-LM

max_forks_repo_name: t-hyun/Megatron-LM


text:
import numpy as np
import json
from PIL import Image
from pathlib import Path

'''record_###.json
{"user/throttle": 0.7424248049084544, "timestamp": "2019-03-06 21:04:48.552623", "user/mode": "user", "cam/image_array": "219_cam-image_array_.jpg", "user/angle": 0.12798827723777417, "is-simulated" : false}
'''
THROTTLE="user/throttle"
STEERING="user/angle"
IS_SIM="is-simulated"
IMAGE="cam/image_array"

def norm_image(image):
    return (image / 255.) - 0.5

def load_image(img):
    img = np.asarray(Image.open(img), dtype=np.float32)
    img = norm_image(img)
    return img

def batcher(seq, size):
    return (seq[pos:pos + size] for pos in range(0, len(seq), size))

def get_gen(data, batch_size, mode):
    while True:
        np.random.shuffle(data)
        for batch in batcher(data, batch_size):
            x = []
            steering = []
            throttle = []
            is_sim   = []
            smoosh   = []
            for ele in batch:

                # Load image the first time then hold it in memory
                # this will slow down the first epoch, but will
                # speed up subsiquent epochs
                if isinstance(ele[IMAGE], str):
                    ele[IMAGE] = load_image((ele[IMAGE]))
                x.append(ele[IMAGE])
                steering.append(ele[STEERING])
                throttle.append(ele[THROTTLE])

                if mode in ["L", "C"]:
                    if mode == "C":
                        smoosh.append([0.5,0.5])
                        is_sim.append([ele[IS_SIM], not ele[IS_SIM]])
                    elif mode == "L":
                        smoosh.append(0.5)
                        is_sim.append(ele[IS_SIM])
                    else:
                        assert False, "POOP"
                    y = [steering, throttle, is_sim, smoosh]
                else:
                    y = [steering, throttle]

            yield [x], y

def get_gens(tub_paths, batch_size=32, train_frac=0.8, seed=42,
             mode  = "POOP"):
    np.random.seed(seed)
    records = []
    for path in tub_paths:
        for r_path in Path(path).glob("record*.json"):
            with open(r_path, "r") as f:
                data = json.load(f)
                data[IMAGE] = str(Path(path) / data[IMAGE])
                records.append(data)
    np.random.shuffle(records)
    train_split = int(len(records) * train_frac)
    train_data = records[:train_split]
    valid_data = records[train_split:]

    train_gen = get_gen(train_data, batch_size, mode)
    valid_gen = get_gen(valid_data, batch_size, mode)

    return (train_gen, int(len(train_data)/batch_size),
            valid_gen, int(len(valid_data)/batch_size))

max_stars_repo_name: tall-josh/dingocar-research

max_issues_repo_name: tall-josh/dingocar-research

max_forks_repo_name: tall-josh/dingocar-research


text:
"""
    Create graph from dmatrix
!!!!!!!!!!!!!!!! Please use CreateGraphNumpy.py, this file is obsolete !!!!!!!!!!!!!!!
"""
import networkx as nx
import numpy as np
import os
import pickle
from random import choice
from FileParser import GetAllinfo
from BuildTree import BuildTree
# Global variable so I can change easily
__SAVEDIR__ = "./Data/"

def createGraph( dmatrix, criteria, moldict, typeofbinding ):
    newGraph = nx.Graph()
    row, col = dmatrix.shape
    for eachrow in range(row):
        # in case, isolated nodes exist in graph
        if not IsTypeofBinding(eachrow, moldict, typeofbinding):
            continue
        newGraph.add_node( eachrow )
        for eachcol in range( eachrow + 1, col ):
            if not IsTypeofBinding(eachcol, moldict, typeofbinding):
                continue
            eachweight = dmatrix[eachrow, eachcol]
            if eachweight < criteria:
                newGraph.add_edge( eachrow, eachcol, weight = eachweight )
    return newGraph

def RandomPickFromList( alist ):
    return choice( alist )

def BindingTypeFilter( alist, moldict, bindingType = None ):
    if bindingType is None:
        return alist
    newlist = []
    for eachIndex in alist:
        if IsTypeofBinding( eachIndex, moldict, bindingType):
            newlist.append( eachIndex )
    if len(newlist) == 0:
        raise RuntimeError("Empty list after filtering!!")
    return newlist

def LeaderInCluster( graphObj, moldict ):
    leaderList = []
    for eachSubGraph in nx.connected_component_subgraphs( graphObj ):
        centerlist = nx.center(eachSubGraph)
        if len(centerlist) < 1:
            raise "A subgraph with less than 1 center"
        leaderID   = RandomPickFromList(centerlist)
        leaderList.append( leaderID )
        graphSize  = len(eachSubGraph)
        moldict[ leaderID ][ "size" ] = graphSize
    return leaderList

def IsTypeofBinding( index, moldict, typeofbinding ):
    if typeofbinding == "all":
        return True
    return moldict[index]["typeofbinding"] == typeofbinding

def MoleculeDictionary( infile ):
    all_info = GetAllinfo( infile )
    totalrow = len( all_info[ "ligandid" ] )
    molDict  = dict()
    for index in range( totalrow ):
        ligandid = all_info[ "ligandid" ][ index ]
        bindingtype = all_info[ "typeofbinding" ][ index ]
        molDict[index] = { "ligandid": ligandid, "typeofbinding": bindingtype }
    return molDict

def CheckExistingLeaderlist( typeofbinding, criteria ):
    criString = str(criteria)
    expectFile = typeofbinding + "_" + criString
    numpyarray = None
    moldictfile = None
    for eachfile in os.listdir(__SAVEDIR__):
        if typeofbinding in eachfile and criString in eachfile:
            filedir = __SAVEDIR__ + eachfile
            if eachfile == expectFile + ".npz":
                numpyarray = filedir
            if eachfile == expectFile + ".p":
                moldictfile = filedir
    if not numpyarray is None and not moldictfile is None:
        return (numpyarray, moldictfile)
    return None

def SizeHistogram( moldict ):
    sizeList = []
    for eachID in moldict:
        try:
            sizeList.append(moldict[eachID]["size"])
        except:
            continue
    value, binedge = np.histogram(sizeList)
    print "number:", value
    print "value:", binedge

def SaveLeaderAndMolDict(leaderlist, moldict, typeofbinding, criteria):
    criString = str(criteria)
    filename = "_".join([ typeofbinding, criString ])
    filedir  = __SAVEDIR__ + filename
    np.savez( filedir, leaderlist )
    pickle.dump( moldict, open( filedir + ".p", "wb" ))

def main( bindingtype, minDistance, dmatrix ):
    #minDistance = 0.75
    infile      = "./Data/ligand_5_7_ppilot_modified.txt"
    #leaderAndmol = CheckExistingLeaderlist( bindingtype, minDistance )
    leaderAndmol = None
    if leaderAndmol:
        leaderfile, moldictfile = leaderAndmol
        with np.load(leaderfile) as leader_moldict:
            print leader_moldict.files
            leaderlist     = leader_moldict["arr_0"]
            print "leaderlist length:", len(leaderlist)
        with open(moldictfile, "rb") as moldictobj:
            moldict        = pickle.load(moldictobj)
        SizeHistogram( moldict )
    else:
        moldict = MoleculeDictionary( infile )
        SizeHistogram( moldict )
        newgraph = createGraph( dmatrix, minDistance, moldict, bindingtype)
        leaderlist = LeaderInCluster( newgraph, moldict )
        SaveLeaderAndMolDict(leaderlist, moldict, bindingtype, minDistance)
    leaderlist = BindingTypeFilter( leaderlist, moldict, bindingtype)
    BuildTree( leaderlist, dmatrix, moldict, str(minDistance) + "_" + bindingtype )

if __name__ == "__main__":
    smatrixfile = "./Data/similarityMatrix.npy"
    dmatrix = 1 - np.load(smatrixfile)
    print dmatrix.shape
    #distanceList = [ 0.6, 0.65, 0.7, 0.8 ]
    distanceList = [ 0.1 ]
    #for each in ["allosteric", "competitive"]:
    for each in ["all"]:
        for distance in distanceList:
            main(each, distance, dmatrix)

max_stars_repo_name: ajing/clusterVis

max_issues_repo_name: ajing/clusterVis

max_forks_repo_name: ajing/clusterVis

reason: import numpy,import networkx


text:
from utils.misc import get_logger
log = get_logger()

import threading, time
import utils.kaldi_io
import numpy as np

def load_and_prep_dev_set_tar(egs_dir, nnet3_copy_egs_to_feats):
    dev_feat_rsp        = "ark:" + nnet3_copy_egs_to_feats + " ark:" + egs_dir + "/valid_egs.1.ark ark:- |"
    dev_feats_generator = utils.kaldi_io.read_mat_ark(dev_feat_rsp)
    dev_set             = list(dev_feats_generator)
    dev_lab             = np.array([int(dev_set[i][0].split("-")[-1]) for i in range(len(dev_set))])
    dev_feat            = np.vstack([dev_set[i][1][np.newaxis,:,:] for i in range(len(dev_set))])
    dev_idx             = list(range(0, dev_feat.shape[0]*(dev_feat.shape[1]+1), dev_feat.shape[1]))

    return dev_lab, dev_feat.reshape(1,-1,dev_feat.shape[2]), dev_idx

class egsBatchQue(object):

    def __init__(self, nnet3_copy_egs_to_feats, archive_path, b_size, n_archives, ark_que_length, feat_size=40, do_shuffle=False):
        self.delete = False

        self.nnet3_copy_egs_to_feats = nnet3_copy_egs_to_feats
        self.archive_path = archive_path
        self.b_size = b_size
        self.n_archives = n_archives
        self.ark_que_length = ark_que_length
        self.feat_size = feat_size
        self.do_shuffle = do_shuffle

        if self.do_shuffle:
            self.archive_list =  np.random.permutation( len( self.n_archives ) ) +1 # +1 Because egs indices starts from 1
            log.debug( "Shuffled the archive list " )
        else:
            self.archive_list =  np.arange( self.n_archives ) +1                    # np.random.permutation( len( self.n_archives ) )

        self.qued_archives = []
        self.archive_idx = 0       # Index for which archive to process
        self.batch_idx_ark = 0     # Index for where to start the batch within the current archive

        self.batch_number = 1

        self.batch_thread = threading.Thread( target =self.prep_archives )
        self.batch_thread.daemon = True # This will make the process die if the main process dies I THINK...???
        self.batch_thread.start()

    def prep_archives( self ):
        while not self.delete:
            if ( len(self.qued_archives ) < self.ark_que_length ):
                log.info( "Loading new archive" ) # self.qued_archives
                # If we have reached the last archive.
                if self.archive_idx == len( self.archive_list ) -1:
                    self.archive_idx = 0
                    if self.do_shuffle:
                        self.archive_list =  np.random.permutation( len( self.n_archives ) ) + 1
                        log.debug( "Shuffled the archive list " )
                feat_rsp="ark:" + self.nnet3_copy_egs_to_feats + " ark:" + self.archive_path + "/egs." + str( self.archive_list[ self.archive_idx ] ) +  ".ark ark:- |"
                feats_generator=utils.kaldi_io.read_mat_ark(feat_rsp)

                a=list(feats_generator)
                lab=np.array([int(a[i][0].split("-")[-1]) for i in range(len(a))])
                feat=np.vstack([a[i][1][np.newaxis,:,:] for i in range(len(a))])
                log.debug("loading archive done.")

                self.archive_idx += 1

                if self.do_shuffle:
                    idx=np.random.permutation( len(lab) )
                    feat = feat[idx]
                    lab = lab[idx]
                    log.debug( "Shuffled the loaded archive." )

                self.qued_archives.append( [lab,feat] )

    def get_batch(self):
        X=[]
        U=[]
        bad_tr_files = []
        tr_idx = None
        control_nb = 0

        log.debug ( "Retrieving batch" )

        while len(self.qued_archives) < 1:
            time.sleep(1)

        print("  " + str(len( self.qued_archives[0][0] )))
        print("  " + str(self.batch_idx_ark + self.b_size))

        if len( self.qued_archives[0][0] ) >= self.batch_idx_ark + self.b_size:
            assert ( len( self.qued_archives[0][1] ) == len( self.qued_archives[0][0] ) )

            start = self.batch_idx_ark
            end   = self.batch_idx_ark + self.b_size
            tr_feats = self.qued_archives[0][1][start:end]
            Y        = self.qued_archives[0][0][start:end]

            self.batch_idx_ark = end

        else:

            assert ( len( self.qued_archives[0][1] ) == len( self.qued_archives[0][0] ) )
            start = self.batch_idx_ark
            end   = self.batch_idx_ark +  self.b_size        # Will be beyond last index but this is OK
            tr_feats = self.qued_archives[0][1][start:end]
            Y        = self.qued_archives[0][0][start:end]
            self.qued_archives.pop(0)

            n_needed = self.b_size - Y.shape[0]

            while len(self.qued_archives) < 1 :
                time.sleep(1)

            assert ( len( self.qued_archives[0][1] ) == len( self.qued_archives[0][0] ) )
            log.debug( tr_feats.shape )
            log.debug( self.qued_archives[0][1][0:n_needed].shape )
            log.debug( Y.shape )
            log.debug( self.qued_archives[0][0][0:n_needed].shape )

            tr_feats = np.vstack( (tr_feats, self.qued_archives[0][1][0:n_needed]) )
            Y        = np.hstack( (Y, self.qued_archives[0][0][0:n_needed]) )

            self.batch_idx_ark = n_needed

        self.batch_number += 1
        return [[X, Y, U], [bad_tr_files], [tr_feats.astype('float32'), tr_idx],  self.batch_number, control_nb]

max_stars_repo_name: gulamungon/SEQUENS

max_issues_repo_name: gulamungon/SEQUENS

max_forks_repo_name: gulamungon/SEQUENS


text:
import cv2
import numpy as np
import os

from time import time
from PyQt5.QtWidgets import QProgressBar

from img import Tiff

# Image Registration using
# Enhanced Correlation Coefficient (ECC) Maximization
def images_alignement(one, two, outpath):
    _h, _w = one.shape

    # Define 2x3 and initialize the matrix to identity
    warp_matrix = np.eye(2, 3, dtype=np.float32)

    # Specify the number of iterations
    number_of_iterations = 5000

    # Specify the threshold of the increment
    # in the correlation coefficient between two iterations
    termination_eps = 1e-10

    # Define termination criteria
    criteria = (
        cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
        number_of_iterations,
        termination_eps
    )

    # Run the ECC algorithm. The results are stored in warp_matrix.
    (cc, warp_matrix) = cv2.findTransformECC (
        one, two,
        warp_matrix,
        cv2.MOTION_TRANSLATION,
        criteria
    )

    # Use warpAffine for Translation, Euclidean and Affine
    three = cv2.warpAffine(
        two,
        warp_matrix, (_w, _h),
        flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP
    )

    # Save image
    cv2.imwrite(outpath, three.astype('uint16'))
    return outpath

def alignement(paths, pbar: QProgressBar):
    timestamp = str(int(time()))
    spath = "../.cache/alignment/" + timestamp + "/"

    pbar.setMaximum(len(paths))
    pbar.setEnabled(True)

    if not os.path.exists(spath):
        os.makedirs(spath)

    new_paths = []

    # Algorithm need to work with 8bits images
    one = Tiff(paths[0]).to_8bits()
    sname = spath + os.path.basename(paths[0])
    new_paths.append(sname)
    cv2.imwrite(sname, one.astype('uint16'))

    for i in range(1, len(paths)):
        two = Tiff(paths[i]).to_8bits()

        sname = spath + os.path.basename(paths[i])

        new_paths.append(
            images_alignement(one, two, sname)
        )

        one = two

        pbar.setValue( pbar.value() + 1 )

    pbar.setValue(0)
    pbar.setEnabled(False)

    return new_paths

max_stars_repo_name: Schrut/PRT

max_issues_repo_name: Schrut/PRT

max_forks_repo_name: Schrut/PRT


text:
import os
import cv2
import argparse
import uuid
import numpy as np
import pathlib

from filters.noise.noise_removal import NoiseRemoval
from filters.binarization import otsu, sauvola
from utils.helpers import load_image, debug_display_image, debug_plot_array
from utils.char_mapper import vocab_letter_to_class

from ocr_image import OCRImage
from text_image_v3 import TextImageBaseline
from scipy.ndimage import interpolation as inter

INVALID_WORD_FOLDER = "/Users/filipgulan/ds/invalid_words"

def draw_box(image, ocr_image):
    b_box = ocr_image.get_bounding_box()
    cv2.rectangle(image,
                  (b_box['x'] - 1, b_box['y'] - 1),
                  (b_box['x'] + b_box['width'] - 1,
                   b_box['y'] + b_box['height'] - 1),
                  (255, 0, 0), 1)

def process_image(image_path):
    input_image = load_image(image_path)

    binarizer = otsu.OtsuBinarization()
    binarized_img = binarizer.process(input_image)
    binarized_img = cv2.bitwise_not(binarized_img)

    denoiser = NoiseRemoval()
    input_image = denoiser.process(binarized_img)

    height, width = input_image.shape[:2]
    ocr_image = OCRImage(input_image, width, height)
    angle = ocr_image.fix_skew()

    rotated = inter.rotate(input_image, angle, reshape=False, order=0)
    original_image = cv2.cvtColor(rotated, cv2.COLOR_GRAY2RGB)

    roi_image, width, height, min_x, min_y = ocr_image.get_segments()
    text_image = TextImageBaseline(roi_image, width, height, min_x, min_y)
    return text_image, original_image

def process_text_file(file_path):
    lines = []
    with open(file_path) as text_file:
        lines = text_file.read().splitlines()
    lines = list(map(lambda line: line.strip(), lines))
    lines = list(filter(None, lines))
    return lines

def get_words(file_line):
    words = file_line.split(" ")
    words = list(map(lambda text: text.strip(), words))
    words = list(filter(None, words))
    return words

def process_chars(ocr_chars, file_chars, root_output_folder, avg_line_height):
    for ocr_char, file_char in zip(ocr_chars, file_chars):
        scaled_image = ocr_char.get_scaled_image(avg_line_height)

        file_name = str(uuid.uuid4()) + ".jpg"
        letter_class = vocab_letter_to_class[file_char]
        output_folder = os.path.join(root_output_folder, letter_class)
        output_file = os.path.join(output_folder, file_name)
        pathlib.Path(output_folder).mkdir(parents=True, exist_ok=True)
        cv2.imwrite(output_file, scaled_image)

def process_words(ocr_words, file_words, root_output_folder, line_index, avg_line_height):
    for ocr_word, file_word in zip(ocr_words, file_words):
        ocr_chars = ocr_word.get_segments()
        file_chars = list(file_word)

        if len(ocr_chars) != len(file_chars):
            # print("Neispravan broj znakova u rijeci: " + file_word, "Linija:", str(line_index + 1))
            # print("ocr_chars:", len(ocr_chars), "file_chars", len(file_chars))
            file_name = str(uuid.uuid4()) + ".jpg"
            output_path = os.path.join(INVALID_WORD_FOLDER, file_name)
            ocr_word.save(output_path)
            continue

        process_chars(ocr_chars, file_chars,
                      root_output_folder, avg_line_height)

def process(args):
    file_lines = process_text_file(args.text_file)
    text_image, original_image = process_image(args.image)

    text_lines = text_image.get_segments()
    if len(text_lines) != len(file_lines):
        print("Neispravan broj linija!")
        return

    avg_line_height = 0
    for line in text_lines:
        avg_line_height += line.get_height()
    avg_line_height /= len(text_lines)
    print(avg_line_height)

    all_words = []
    for index, (ocr_line, file_line) in enumerate(zip(text_lines, file_lines)):
        ocr_words = ocr_line.get_segments()
        file_words = get_words(file_line)

        all_words.extend(ocr_words)
        if len(ocr_words) == len(file_words):
            process_words(ocr_words, file_words, args.output,
                          index, avg_line_height)
        else:
            print("Neispravan broj rijeci na liniji " + str(index + 1))
            continue

    export_words_image(original_image, all_words)

def export_words_image(out_image, ocr_words):
    for ocr_word in ocr_words:
        draw_box(out_image, ocr_word)

    cv2.imwrite("sample_words.jpg", out_image)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--image', help="Input text image", required=True, type=str)
    parser.add_argument(
        '--text-file', help="Input text file", required=True, type=str)
    parser.add_argument(
        '--output', help="Input image is grayscale", required=True, type=str)
    args = parser.parse_args()
    process(args)

if __name__ == '__main__':
    main()

max_stars_repo_name: fgulan/PyOCR

max_issues_repo_name: fgulan/PyOCR

max_forks_repo_name: fgulan/PyOCR


text:
import numpy as np
import onnx
import onnxruntime as ort
x=np.random.rand(1,3,32,256)
session = ort.InferenceSession("ch_ppocr_mobile_v2.0_det_infer.onnx.new")
onnx_inputs = { session.get_inputs()[0].name: x.astype(np.float32)}
onnx_outputs=session.run(None, onnx_inputs)
print(onnx_outputs)

max_stars_repo_name: Darcy-C/RapidOCR

max_issues_repo_name: maxpark/RapidOCR

max_forks_repo_name: maxpark/RapidOCR


text:
import time
import random
import numpy as np
from pylsl import StreamInfo, StreamInlet, StreamOutlet, local_clock, resolve_byprop,\
    resolve_bypred, proc_clocksync, proc_dejitter, proc_monotonize
try:
    from pyfftw.interfaces.numpy_fft import rfft, irfft       # Performs much better than numpy's fftpack
except ImportError:
    from numpy.fft import rfft, irfft
try:
    import pyqtgraph as pg
    import sys
    haspyqtgraph = True
except ImportError:
    haspyqtgraph = False

# The code for pink noise generation is taken from
# https://github.com/python-acoustics/python-acoustics/blob/master/acoustics/generator.py
# which is distributed under the BSD license.
def ms(x):
    """Mean value of signal `x` squared.
    :param x: Dynamic quantity.
    :returns: Mean squared of `x`.
    """
    return (np.abs(x) ** 2.0).mean()

def normalize(y, x=None):
    """normalize power in y to a (standard normal) white noise signal.
    Optionally normalize to power in signal `x`.
    #The mean power of a Gaussian with :math:`\\mu=0` and :math:`\\sigma=1` is 1.
    """
    #return y * np.sqrt( (np.abs(x)**2.0).mean() / (np.abs(y)**2.0).mean() )
    if x is not None:
        x = ms(x)
    else:
        x = 1.0
    return y * np.sqrt( x / ms(y) )
    #return y * np.sqrt( 1.0 / (np.abs(y)**2.0).mean() )

def pink(N):
    """
    Pink noise.

    :param N: Amount of samples.

    Pink noise has equal power in bands that are proportionally wide.
    Power density decreases with 3 dB per octave.

    """
    # This method uses the filter with the following coefficients.
    # b = np.array([0.049922035, -0.095993537, 0.050612699, -0.004408786])
    # a = np.array([1, -2.494956002, 2.017265875, -0.522189400])
    # return lfilter(B, A, np.random.randn(N))
    # Another way would be using the FFT
    # x = np.random.randn(N)
    # X = rfft(x) / N
    uneven = N % 2
    X = np.random.randn(N // 2 + 1 + uneven) + 1j * np.random.randn(N // 2 + 1 + uneven)
    S = np.sqrt(np.arange(len(X)) + 1.)  # +1 to avoid divide by zero
    y = (irfft(X / S)).real
    if uneven:
        y = y[:-1]
    return normalize(y)

class PinkNoiseGenerator(object):
    def __init__(self, nSampsPerBlock=1024):
        self.N = nSampsPerBlock
        self.uneven = self.N % 2
        lenX = self.N//2 + 1 + self.uneven
        self.S = np.sqrt(np.arange(lenX) + 1.)

    def generate(self):
        X = np.random.randn(self.N // 2 + 1 + self.uneven) + 1j * np.random.randn(self.N // 2 + 1 + self.uneven)
        y = (irfft(X / self.S)).real
        if self.uneven:
            y = y[:-1]
        return normalize(y)

class BetaGeneratorOutlet(object):
    def __init__(self, Fs=2**14, FreqBeta=20.0, AmpBeta=100.0, AmpNoise=20.0, NCyclesPerChunk=4,
                 channels=["RAW1", "SPK1", "RAW2", "SPK2", "RAW3", "SPK3"]):
        """
        :param Fs:              Sampling rate
        :param FreqBeta:        Central frequency of beta band
        :param AmpBeta:         Amplitude of beta (uV)
        :param AmpNoise:        Amplitude of pink noise (uV)
        :param NCyclesPerChunk: Minimum number of cycles of beta in a chunk.
        :param channels:        List of channel names
        """
        # Saved arguments
        self.FreqBeta = FreqBeta
        self.AmpBeta = AmpBeta                                          # Amplitude of Beta (uV)
        self.AmpNoise = AmpNoise                                        # Amplitude of pink noise
        self.channels = channels
        # Derived variables
        chunk_dur = NCyclesPerChunk / self.FreqBeta           # Duration, in sec, of one chunk
        chunk_len = int(Fs * chunk_dur)                  # Number of samples in a chunk
        self.tvec = 1.0 * (np.arange(chunk_len) + 1) / Fs     # time vector for chunk (sec)
        # Pink noise generator
        self.pinkNoiseGen = PinkNoiseGenerator(nSampsPerBlock=chunk_len)

        # Create a stream of fake 'raw' data
        raw_info = StreamInfo(name='BetaGen', type='EEG',
                           channel_count=len(self.channels), nominal_srate=Fs,
                           channel_format='float32', source_id='betagen1234')
        raw_xml = raw_info.desc()
        chans = raw_xml.append_child("channels")
        for channame in self.channels:
            chn = chans.append_child("channel")
            chn.append_child_value("label", channame)
            chn.append_child_value("unit", "microvolts")
            chn.append_child_value("type", "generated")
        self.eeg_outlet = StreamOutlet(raw_info)
        print("Created outlet with name BetaGen and type EEG")

        self.last_time = local_clock()

    def update(self, task={'phase': 'precue', 'class': 1}):
        # Convert phase and class_id into beta_amp
        if task['phase'] in ['cue', 'go']:
            beta_amp = 0 if task['class'] == 3 else self.AmpBeta
        else:
            beta_amp = self.AmpBeta / 5.0

        this_tvec = self.tvec + self.last_time  # Sample times
        # Put the signal together
        this_sig = self.AmpNoise * np.asarray(self.pinkNoiseGen.generate(), dtype=np.float32)  # Start with some pink noise
        this_sig += beta_amp * np.sin(this_tvec * 2 * np.pi * self.FreqBeta)  # Add our beta signal
        this_sig = np.atleast_2d(this_sig).T * np.ones((1, len(self.channels)), dtype=np.float32)  # Tile across channels

        time_to_sleep = max(0, this_tvec[-1] - local_clock())
        time.sleep(time_to_sleep)

        print("Beta outlet pushing signal with shape {},{} and Beta amp {}".format(this_sig.shape[0], this_sig.shape[1],
                                                                                beta_amp))
        self.eeg_outlet.push_chunk(this_sig, timestamp=this_tvec[-1])

        self.last_time = local_clock()

class BetaInlet(object):
    def __init__(self):
        print("looking for an EEG stream...")
        streams = resolve_byprop("type", "EEG")

        # create a new inlet to read from the stream
        proc_flags = proc_clocksync | proc_dejitter | proc_monotonize
        self.inlet = StreamInlet(streams[0], processing_flags=proc_flags)

        # The following is an example of how to read stream info
        stream_info = self.inlet.info()
        stream_Fs = stream_info.nominal_srate()
        stream_xml = stream_info.desc()
        chans_xml = stream_xml.child("channels")
        chan_xml_list = []
        ch = chans_xml.child("channel")
        while ch.name() == "channel":
            chan_xml_list.append(ch)
            ch = ch.next_sibling("channel")
        self.channel_names = [ch_xml.child_value("label") for ch_xml in chan_xml_list]
        print("Reading from inlet named {} with channels {} sending data at {} Hz".format(stream_info.name(),
                                                                                self.channel_names, stream_Fs))

    def update(self):
        max_samps = 3276*2
        data = np.nan * np.ones((max_samps, len(self.channel_names)), dtype=np.float32)
        _, timestamps = self.inlet.pull_chunk(max_samples=max_samps, dest_obj=data)
        data = data[:len(timestamps), :]
        print("Beta inlet retrieved {} samples.".format(len(timestamps)))
        return data, np.asarray(timestamps)

class MarkersGeneratorOutlet(object):
    phases = {
        'precue': {'next': 'cue', 'duration': 1.0},
        'cue': {'next': 'go', 'duration': 0.5},
        'go': {'next': 'evaluate', 'duration': 5.0},
        'evaluate': {'next': 'precue', 'duration': 0.1}
    }

    def __init__(self, class_list=[1, 3], classes_rand=True, target_list=[1, 2], targets_rand=True):
        """

        :param class_list:  A list of integers comprising different class ids. Default: [1, 3]
        :param classes_rand: If True, classes are chosen randomly from list. If False, the list is cycled. Default: True
        :param target_list: A list of integers comprising different target ids. Default: [1, 2]
        :param targets_rand: If True, targets are chosen randomly from list. If False, the list is cycled. Default: True
        """
        stream_name = 'GeneratedCentreOutMarkers'
        stream_type = 'Markers'
        outlet_info = StreamInfo(name=stream_name, type=stream_type,
                   channel_count=1, nominal_srate=0,
                   channel_format='string',
                   source_id='centreoutmarkergen1234')
        outlet_xml = outlet_info.desc()
        channels_xml = outlet_xml.append_child("channels")
        chan_xml = channels_xml.append_child("channel")
        chan_xml.append_child_value("label", "EventMarkers")
        chan_xml.append_child_value("type", "generated")
        self.outlet = StreamOutlet(outlet_info)
        print("Created outlet with name {} and type {}".format(stream_name, stream_type))

        self.class_list = class_list
        self.classes_rand = classes_rand
        self.target_list = target_list
        self.targets_rand = targets_rand
        self.next_transition = -1
        self.in_phase = 'evaluate'
        self.trial_ix = 0
        self.class_id = self.class_list[0]
        self.target_id = self.target_list[0]

    def update(self):
        now = local_clock()
        if (now > self.next_transition):

            # Transition phase
            self.in_phase = self.phases[self.in_phase]['next']
            self.next_transition = now + self.phases[self.in_phase]['duration']

            # Send markers
            out_string = "undefined"
            if self.in_phase == 'precue':
                # transition from evaluate to precue
                # print("Previous class_id: {}, target_id: {}".format(self.class_id, self.target_id))
                self.trial_ix += 1
                self.target_id = random.choice(self.target_list) if self.targets_rand else self.target_list[
                    (self.target_list.index(self.target_id) + 1) % len(self.target_list)]
                self.class_id = random.choice(self.class_list) if self.classes_rand else self.class_list[
                    (self.class_list.index(self.class_id) + 1) % len(self.class_list)]
                # print("New class_id: {}, target_id: {}".format(self.class_id, self.target_id))
                out_string = "NewTrial {}, Class {}, Target {}".format(self.trial_ix, self.class_id, self.target_id)
            elif self.in_phase == 'cue':
                # transition from precue to cue
                out_string = "TargetCue, Class {}, Target {}".format(self.class_id, self.target_id)
            elif self.in_phase == 'go':
                # transition from cue to go
                out_string = "GoCue, Class {}, Target {}".format(self.class_id, self.target_id)
            elif self.in_phase == 'evaluate':
                # transition from go to evaluate
                hit_string = "Hit" if random.randint(0, 1) == 1 else "Miss"
                out_string = "{}, Class {}, Target {}".format(hit_string, self.class_id, self.target_id)
            print("Marker outlet pushing string: {}".format(out_string))
            self.outlet.push_sample([out_string,])

            return True
        return False

class MarkerInlet(object):
    def __init__(self):
        self.task = {'phase':'precue', 'class':1, 'target':1}
        print("Looking for stream with type Markers")
        streams = resolve_bypred("type='Markers'", minimum=1)
        proc_flags = 0  # Marker events are relatively rare. No need to post-process.
        self.inlet = StreamInlet(streams[0], processing_flags=proc_flags)
        # The following is an example of how to read stream info
        stream_info = self.inlet.info()
        stream_Fs = stream_info.nominal_srate()
        stream_xml = stream_info.desc()
        chans_xml = stream_xml.child("channels")
        chan_xml_list = []
        ch = chans_xml.child("channel")
        while ch.name() == "channel":
            chan_xml_list.append(ch)
            ch = ch.next_sibling("channel")
        stream_ch_names = [ch_xml.child_value("label") for ch_xml in chan_xml_list]
        print("Reading from inlet named {} with channels {}".format(stream_info.name(), stream_ch_names))

    def update(self):
        marker_samples, marker_timestamps = self.inlet.pull_chunk(timeout=0.0)
        if (marker_timestamps):
            [phase_str, class_str, targ_str] = marker_samples[-1][0].split(', ')
            if phase_str in ['TargetCue']:
                self.task['phase'] = 'cue'
            elif phase_str in ['GoCue']:
                self.task['phase'] = 'go'
            elif phase_str in ['Miss', 'Hit']:
                self.task['phase'] = 'evaluate'
            elif phase_str[:8] == 'NewTrial':
                self.task['phase'] = 'precue'
            else:
                print(phase_str)
            self.task['class'] = int(class_str.split(' ')[1])
            self.task['target'] = int(targ_str.split(' ')[1])
            print("Marker inlet updated with task {}".format(self.task))

betaGen = BetaGeneratorOutlet()
markerGen = MarkersGeneratorOutlet()
betaIn = BetaInlet()
markerIn = MarkerInlet()

if haspyqtgraph:
    qapp = pg.QtGui.QApplication(sys.argv)
    qwindow = pg.plot()
    qwindow.clear()
    qwindow.parent().setWindowTitle("pylsl PerformanceTest")

def update():
    markerGen.update()
    markerIn.update()
    betaGen.update(task=markerIn.task)  # Rate-limiting step. Will time.sleep as needed.
    signal, tvec = betaIn.update()

    if haspyqtgraph:
        plot = qwindow.getPlotItem()
        graphs = plot.listDataItems()
        if not graphs:
            # create graphs
            for i in range(signal.shape[1]):
                plot.plot(tvec, signal[:, i])
        else:
            # update graphs
            for i in range(signal.shape[1]):
                graphs[i].setData(signal[:, i], x=tvec)

if __name__ == '__main__':
    """
    python3 -m cProfile -o pylsl.cprof PerformanceTest.py
    gprof2dot -f pstats pylsl.cprof | dot -Tpng -o pylsl_prof.png
    """
    try:
        if haspyqtgraph:
            timer = pg.QtCore.QTimer()
            timer.timeout.connect(update)
            timer.start(1)  # Delay not needed because update has time.sleep
            if (sys.flags.interactive != 1) or not hasattr(pg.QtCore, 'PYQT_VERSION'):
                sys.exit(pg.QtGui.QApplication.instance().exec_())
        else:
            while True:
                update()

    except KeyboardInterrupt:
        # No cleanup necessary?
        pass

max_stars_repo_name: stfnrpplngr/liblsl-Python

max_issues_repo_name: stfnrpplngr/liblsl-Python

max_forks_repo_name: stfnrpplngr/liblsl-Python

reason: import numpy,from numpy


text:
#!/usr/bin/env python3
import numpy as np
import gym

from stable_baselines_custom.common.cmd_util import mujoco_arg_parser
from stable_baselines_custom import bench, logger
from stable_baselines_custom.common import set_global_seeds
from stable_baselines_custom.common.vec_env.vec_normalize import VecNormalize
from stable_baselines_custom.ppo2 import PPO2
from stable_baselines_custom.common.policies import MlpPolicy
from stable_baselines_custom.common.vec_env.dummy_vec_env import DummyVecEnv

def train(env_id, num_timesteps, seed):
    """
    Train PPO2 model for Mujoco environment, for testing purposes

    :param env_id: (str) the environment id string
    :param num_timesteps: (int) the number of timesteps to run
    :param seed: (int) Used to seed the random generator.
    """
    def make_env():
        env_out = gym.make(env_id)
        env_out = bench.Monitor(env_out, logger.get_dir(), allow_early_resets=True)
        return env_out

    env = DummyVecEnv([make_env])
    env = VecNormalize(env)

    set_global_seeds(seed)
    policy = MlpPolicy
    model = PPO2(policy=policy, env=env, n_steps=2048, nminibatches=32, lam=0.95, gamma=0.99, noptepochs=10,
                 ent_coef=0.0, learning_rate=3e-4, cliprange=0.2)
    model.learn(total_timesteps=num_timesteps)

    return model, env

def main():
    """
    Runs the test
    """
    args = mujoco_arg_parser().parse_args()
    logger.configure()
    model, env = train(args.env, num_timesteps=args.num_timesteps, seed=args.seed)

    if args.play:
        logger.log("Running trained model")
        obs = np.zeros((env.num_envs,) + env.observation_space.shape)
        obs[:] = env.reset()
        while True:
            actions = model.step(obs)[0]
            obs[:] = env.step(actions)[0]
            env.render()

if __name__ == '__main__':
    main()

max_stars_repo_name: iamlab-cmu/stable-baselines

max_issues_repo_name: iamlab-cmu/stable-baselines

max_forks_repo_name: iamlab-cmu/stable-baselines


text:
import numpy as np
from torch.utils.data import Dataset
import torch
import warnings
warnings.filterwarnings("ignore")

# https://arxiv.org/pdf/1811.07490.pdf

def get_binary_mask(x):
    error = torch.abs(x-x[:,0:1]).sum(dim=1)
    binary_mask = error<0.05
    y = binary_mask[0]
    for t in range(1,binary_mask.shape[0]):
        y = y&binary_mask[t]
    return y

class TrafficDataset(Dataset):
    def __init__(self, X,Y,mean=None, std=None, require_back=False):
        super(TrafficDataset,self).__init__()
        self.X = (X+1)/2
        self.Y = (Y+1)/2
        self.mean = -1
        self.std = 2

        self.require_back=require_back
        # if require_back:
        #     backgrounds = torch.load('/usr/data/gzy/Weather_Forecast/ex_sota/backgrounds/traffic.pt')
        #     backgrounds = backgrounds.swapaxes(2,3).swapaxes(1,2)

        #     N1 = self.X.shape[0]
        #     x = self.X[:,0,::].reshape(N1,-1)
        #     x = x/np.linalg.norm(x,axis=1,keepdims=True)

        #     N2 = backgrounds.shape[0]
        #     b = backgrounds.reshape(N2,-1)
        #     b = b/np.linalg.norm(b,axis=1,keepdims=True)

        #     S = np.matmul(x,b.T)
        #     self.back_idx = np.argmax(S,axis=1)
        #     self.backgrounds = backgrounds[:,np.newaxis,::]

    def __len__(self):
        return self.X.shape[0]

    def __getitem__(self, index):
        data = torch.tensor(self.X[index,::]).float()
        labels = torch.tensor(self.Y[index,::]).float()
        if self.require_back:
            b_mask = get_binary_mask(data)
            H,W = b_mask.shape
            return data, labels, b_mask.view(1,1,H,W)
        else:
            return data, labels

def load_data(
        batch_size, val_batch_size,
        data_root,require_back=False):

    # X_train, Y_train, X_test, Y_test, mmn, timestamp_train, timestamp_test=load_traffic_data(data_root,len_closeness=4,len_pred=4,len_test=500)
    dataset = np.load(data_root+'TaxiBJ/dataset.npz')
    X_train, Y_train, X_test, Y_test = dataset['X_train'], dataset['Y_train'], dataset['X_test'], dataset['Y_test']
    train_set = TrafficDataset(X=X_train,
                              Y=Y_train,
                              require_back=require_back)

    test_set = TrafficDataset( X=X_test,
                              Y=Y_test,
                              mean=train_set.mean,
                              std=train_set.std,
                              require_back=require_back)

    dataloader_train = torch.utils.data.DataLoader(
        train_set, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=20)

    dataloader_test = torch.utils.data.DataLoader(
        test_set, batch_size=val_batch_size, shuffle=False, pin_memory=True, num_workers=20)

    return dataloader_train, None, dataloader_test, train_set.mean, train_set.std

if __name__ == '__main__':
    dataloader_train, _, _, _, _ = load_data(batch_size=128, val_batch_size=128, data_root='/usr/data/video_dataset/data/')
    for item in dataloader_train:
        print(item[0].shape)
        break

max_stars_repo_name: gaozhangyang/DecST

max_issues_repo_name: gaozhangyang/DecST

max_forks_repo_name: gaozhangyang/DecST


text:
##################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
#######################################################################
# Network Pruning via Transformable Architecture Search, NeurIPS 2019 #
#######################################################################
import sys, time, torch, random, argparse
from PIL     import ImageFile
from os      import path as osp
ImageFile.LOAD_TRUNCATED_IMAGES = True
import numpy as np
from copy    import deepcopy
from pathlib import Path

lib_dir = (Path(__file__).parent / '..' / 'lib').resolve()
if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir))
from config_utils import load_config, configure2str, obtain_search_args as obtain_args
from procedures   import prepare_seed, prepare_logger, save_checkpoint, copy_checkpoint
from procedures   import get_optim_scheduler, get_procedures
from datasets     import get_datasets, SearchDataset
from models       import obtain_search_model, obtain_model, change_key
from utils        import get_model_infos
from log_utils    import AverageMeter, time_string, convert_secs2time

def main(args):
  assert torch.cuda.is_available(), 'CUDA is not available.'
  torch.backends.cudnn.enabled   = True
  torch.backends.cudnn.benchmark = True
  #torch.backends.cudnn.deterministic = True
  torch.set_num_threads( args.workers )

  prepare_seed(args.rand_seed)
  logger = prepare_logger(args)

  # prepare dataset
  train_data, valid_data, xshape, class_num = get_datasets(args.dataset, args.data_path, args.cutout_length)
  #train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True , num_workers=args.workers, pin_memory=True)
  valid_loader = torch.utils.data.DataLoader(valid_data, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True)

  split_file_path = Path(args.split_path)
  assert split_file_path.exists(), '{:} does not exist'.format(split_file_path)
  split_info      = torch.load(split_file_path)

  train_split, valid_split = split_info['train'], split_info['valid']
  assert len( set(train_split).intersection( set(valid_split) ) ) == 0, 'There should be 0 element that belongs to both train and valid'
  assert len(train_split) + len(valid_split) == len(train_data), '{:} + {:} vs {:}'.format(len(train_split), len(valid_split), len(train_data))
  search_dataset  = SearchDataset(args.dataset, train_data, train_split, valid_split)

  search_train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size,
                      sampler=torch.utils.data.sampler.SubsetRandomSampler(train_split), pin_memory=True, num_workers=args.workers)
  search_valid_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size,
                      sampler=torch.utils.data.sampler.SubsetRandomSampler(valid_split), pin_memory=True, num_workers=args.workers)
  search_loader       = torch.utils.data.DataLoader(search_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, sampler=None)
  # get configures
  if args.ablation_num_select is None or args.ablation_num_select <= 0:
    model_config = load_config(args.model_config, {'class_num': class_num, 'search_mode': 'shape'}, logger)
  else:
    model_config = load_config(args.model_config, {'class_num': class_num, 'search_mode': 'ablation', 'num_random_select': args.ablation_num_select}, logger)

  # obtain the model
  search_model = obtain_search_model(model_config)
  MAX_FLOP, param  = get_model_infos(search_model, xshape)
  optim_config = load_config(args.optim_config, {'class_num': class_num, 'FLOP': MAX_FLOP}, logger)
  logger.log('Model Information : {:}'.format(search_model.get_message()))
  logger.log('MAX_FLOP = {:} M'.format(MAX_FLOP))
  logger.log('Params   = {:} M'.format(param))
  logger.log('train_data : {:}'.format(train_data))
  logger.log('search-data: {:}'.format(search_dataset))
  logger.log('search_train_loader : {:} samples'.format( len(train_split) ))
  logger.log('search_valid_loader : {:} samples'.format( len(valid_split) ))
  base_optimizer, scheduler, criterion = get_optim_scheduler(search_model.base_parameters(), optim_config)
  arch_optimizer = torch.optim.Adam(search_model.arch_parameters(optim_config.arch_LR), lr=optim_config.arch_LR, betas=(0.5, 0.999), weight_decay=optim_config.arch_decay)
  logger.log('base-optimizer : {:}'.format(base_optimizer))
  logger.log('arch-optimizer : {:}'.format(arch_optimizer))
  logger.log('scheduler      : {:}'.format(scheduler))
  logger.log('criterion      : {:}'.format(criterion))

  last_info, model_base_path, model_best_path = logger.path('info'), logger.path('model'), logger.path('best')
  network, criterion = torch.nn.DataParallel(search_model).cuda(), criterion.cuda()

  # load checkpoint
  if last_info.exists() or (args.resume is not None and osp.isfile(args.resume)): # automatically resume from previous checkpoint
    if args.resume is not None and osp.isfile(args.resume):
      resume_path = Path(args.resume)
    elif last_info.exists():
      resume_path = last_info
    else: raise ValueError('Something is wrong.')
    logger.log("=> loading checkpoint of the last-info '{:}' start".format(resume_path))
    checkpoint  = torch.load(resume_path)
    if 'last_checkpoint' in checkpoint:
      last_checkpoint_path = checkpoint['last_checkpoint']
      if not last_checkpoint_path.exists():
        logger.log('Does not find {:}, try another path'.format(last_checkpoint_path))
        last_checkpoint_path = resume_path.parent / last_checkpoint_path.parent.name / last_checkpoint_path.name
      assert last_checkpoint_path.exists(), 'can not find the checkpoint from {:}'.format(last_checkpoint_path)
      checkpoint = torch.load( last_checkpoint_path )
    start_epoch = checkpoint['epoch'] + 1
    #for key, value in checkpoint['search_model'].items():
    #  print('K {:} = Shape={:}'.format(key, value.shape))
    search_model.load_state_dict( checkpoint['search_model'] )
    scheduler.load_state_dict ( checkpoint['scheduler'] )
    base_optimizer.load_state_dict ( checkpoint['base_optimizer'] )
    arch_optimizer.load_state_dict ( checkpoint['arch_optimizer'] )
    valid_accuracies = checkpoint['valid_accuracies']
    arch_genotypes   = checkpoint['arch_genotypes']
    discrepancies    = checkpoint['discrepancies']
    max_bytes        = checkpoint['max_bytes']
    logger.log("=> loading checkpoint of the last-info '{:}' start with {:}-th epoch.".format(resume_path, start_epoch))
  else:
    logger.log("=> do not find the last-info file : {:} or resume : {:}".format(last_info, args.resume))
    start_epoch, valid_accuracies, arch_genotypes, discrepancies, max_bytes = 0, {'best': -1}, {}, {}, {}

  # main procedure
  train_func, valid_func = get_procedures(args.procedure)
  total_epoch = optim_config.epochs + optim_config.warmup
  start_time, epoch_time = time.time(), AverageMeter()
  for epoch in range(start_epoch, total_epoch):
    scheduler.update(epoch, 0.0)
    search_model.set_tau(args.gumbel_tau_max, args.gumbel_tau_min, epoch*1.0/total_epoch)
    need_time = 'Time Left: {:}'.format( convert_secs2time(epoch_time.avg * (total_epoch-epoch), True) )
    epoch_str = 'epoch={:03d}/{:03d}'.format(epoch, total_epoch)
    LRs       = scheduler.get_lr()
    find_best = False

    logger.log('\n***{:s}*** start {:s} {:s}, LR=[{:.6f} ~ {:.6f}], scheduler={:}, tau={:}, FLOP={:.2f}'.format(time_string(), epoch_str, need_time, min(LRs), max(LRs), scheduler, search_model.tau, MAX_FLOP))

    # train for one epoch
    train_base_loss, train_arch_loss, train_acc1, train_acc5 = train_func(search_loader, network, criterion, scheduler, base_optimizer, arch_optimizer, optim_config, \
                                                                                {'epoch-str'  : epoch_str,        'FLOP-exp': MAX_FLOP * args.FLOP_ratio,
                                                                                'FLOP-weight': args.FLOP_weight, 'FLOP-tolerant': MAX_FLOP * args.FLOP_tolerant}, args.print_freq, logger)
    # log the results
    logger.log('***{:s}*** TRAIN [{:}] base-loss = {:.6f}, arch-loss = {:.6f}, accuracy-1 = {:.2f}, accuracy-5 = {:.2f}'.format(time_string(), epoch_str, train_base_loss, train_arch_loss, train_acc1, train_acc5))
    cur_FLOP, genotype = search_model.get_flop('genotype', model_config._asdict(), None)
    arch_genotypes[epoch]  = genotype
    arch_genotypes['last'] = genotype
    logger.log('[{:}] genotype : {:}'.format(epoch_str, genotype))
    # save the configuration
    configure2str(genotype, str( logger.path('log') / 'seed-{:}-temp.config'.format(args.rand_seed) ))
    arch_info, discrepancy = search_model.get_arch_info()
    logger.log(arch_info)
    discrepancies[epoch]   = discrepancy
    logger.log('[{:}] FLOP : {:.2f} MB, ratio : {:.4f}, Expected-ratio : {:.4f}, Discrepancy : {:.3f}'.format(epoch_str, cur_FLOP, cur_FLOP/MAX_FLOP, args.FLOP_ratio, np.mean(discrepancy)))

    #if cur_FLOP/MAX_FLOP > args.FLOP_ratio:
    #  init_flop_weight = init_flop_weight * args.FLOP_decay
    #else:
    #  init_flop_weight = init_flop_weight / args.FLOP_decay

    # evaluate the performance
    if (epoch % args.eval_frequency == 0) or (epoch + 1 == total_epoch):
      logger.log('-'*150)
      valid_loss, valid_acc1, valid_acc5 = valid_func(search_valid_loader, network, criterion, epoch_str, args.print_freq_eval, logger)
      valid_accuracies[epoch] = valid_acc1
      logger.log('***{:s}*** VALID [{:}] loss = {:.6f}, accuracy@1 = {:.2f}, accuracy@5 = {:.2f} | Best-Valid-Acc@1={:.2f}, Error@1={:.2f}'.format(time_string(), epoch_str, valid_loss, valid_acc1, valid_acc5, valid_accuracies['best'], 100-valid_accuracies['best']))
      if valid_acc1 > valid_accuracies['best']:
        valid_accuracies['best'] = valid_acc1
        arch_genotypes['best']   = genotype
        find_best                = True
        logger.log('Currently, the best validation accuracy found at {:03d}-epoch :: acc@1={:.2f}, acc@5={:.2f}, error@1={:.2f}, error@5={:.2f}, save into {:}.'.format(epoch, valid_acc1, valid_acc5, 100-valid_acc1, 100-valid_acc5, model_best_path))
      # log the GPU memory usage
      #num_bytes = torch.cuda.max_memory_allocated( next(network.parameters()).device ) * 1.0
      num_bytes = torch.cuda.max_memory_cached( next(network.parameters()).device ) * 1.0
      logger.log('[GPU-Memory-Usage on {:} is {:} bytes, {:.2f} KB, {:.2f} MB, {:.2f} GB.]'.format(next(network.parameters()).device, int(num_bytes), num_bytes / 1e3, num_bytes / 1e6, num_bytes / 1e9))
      max_bytes[epoch] = num_bytes

    # save checkpoint
    save_path = save_checkpoint({
          'epoch'        : epoch,
          'args'         : deepcopy(args),
          'max_bytes'    : deepcopy(max_bytes),
          'valid_accuracies': deepcopy(valid_accuracies),
          'model-config' : model_config._asdict(),
          'optim-config' : optim_config._asdict(),
          'search_model' : search_model.state_dict(),
          'scheduler'    : scheduler.state_dict(),
          'base_optimizer': base_optimizer.state_dict(),
          'arch_optimizer': arch_optimizer.state_dict(),
          'arch_genotypes': arch_genotypes,
          'discrepancies' : discrepancies,
          }, model_base_path, logger)
    if find_best: copy_checkpoint(model_base_path, model_best_path, logger)
    last_info = save_checkpoint({
          'epoch': epoch,
          'args' : deepcopy(args),
          'last_checkpoint': save_path,
          }, logger.path('info'), logger)

    # measure elapsed time
    epoch_time.update(time.time() - start_time)
    start_time = time.time()

  logger.log('')
  logger.log('-'*100)
  last_config_path = logger.path('log') / 'seed-{:}-last.config'.format(args.rand_seed)
  configure2str(arch_genotypes['last'], str(last_config_path))
  logger.log('save the last config int {:} :\n{:}'.format(last_config_path, arch_genotypes['last']))

  best_arch, valid_acc = arch_genotypes['best'], valid_accuracies['best']
  for key, config in arch_genotypes.items():
    if key == 'last': continue
    FLOP_ratio = config['estimated_FLOP'] / MAX_FLOP
    if abs(FLOP_ratio - args.FLOP_ratio) <= args.FLOP_tolerant:
      if valid_acc <= valid_accuracies[key]:
        best_arch, valid_acc = config, valid_accuracies[key]
  print('Best-Arch : {:}\nRatio={:}, Valid-ACC={:}'.format(best_arch, best_arch['estimated_FLOP'] / MAX_FLOP, valid_acc))
  best_config_path = logger.path('log') / 'seed-{:}-best.config'.format(args.rand_seed)
  configure2str(best_arch, str(best_config_path))
  logger.log('save the last config int {:} :\n{:}'.format(best_config_path, best_arch))
  logger.log('\n' + '-'*200)
  logger.log('Finish training/validation in {:} with Max-GPU-Memory of {:.2f} GB, and save final checkpoint into {:}'.format(convert_secs2time(epoch_time.sum, True), max(v for k, v in max_bytes.items()) / 1e9, logger.path('info')))
  logger.close()

if __name__ == '__main__':
  args = obtain_args()
  main(args)


text:
from typing import List

import numpy as np

__author__ = "@de-code"

"""
Utility class, from:
https://github.com/elifesciences/sciencebeam-trainer-delft/blob/develop/sciencebeam_trainer_delft/utils/numpy.py
"""

def concatenate_or_none(arrays: List[np.array], **kwargs) -> np.array:
    if arrays[0] is None:
        return None
    return np.concatenate(arrays, **kwargs)

# https://stackoverflow.com/a/51526109/8676953
def shuffle_arrays(arrays: List[np.array], random_seed: int = None) -> List[np.array]:
    """Shuffles arrays in-place, in the same order, along axis=0

    Parameters:
    -----------
    arrays : List of NumPy arrays.
    random_seed : Seed value if not None, else seed is random.
    """
    assert all(len(arr) == len(arrays[0]) for arr in arrays)
    if random_seed is None:
        random_seed = np.random.randint(0, 2**(32 - 1) - 1)

    for arr in arrays:
        rstate = np.random.RandomState(random_seed)  # pylint: disable=no-member
        rstate.shuffle(arr)

def shuffle_triple_with_view(a, b=None, c=None):
    """
    Shuffle a pair/triple with view, without modifying the initial elements.
    The assumption is that if b is None, c is also None.
    """
    if b is not None and c is not None:
        assert "Cannot shuffle with view if the arrays have different dimensions: " + str(len(a)) + " vs " \
               + str(len(b)), len(a) == len(b) == len(c)
    elif b is not None and c is None:
        assert "Cannot shuffle with view if the arrays have different dimensions: " + str(len(a)) + " vs " \
               + str(len(b)), len(a) == len(b)
    elif c is not None and b is None:
        raise Exception("b is None but c is not.")

    # generate permutation index array
    permutation = np.random.permutation(a.shape[0])

    # shuffle the arrays
    if c is not None and b is not None:
        return a[permutation], b[permutation], c[permutation]
    if b is not None and c is None:
        return a[permutation], b[permutation], None
    else:
        return a[permutation], None, None

max_stars_repo_name: tantikristanti/delft

max_issues_repo_name: tantikristanti/delft

max_forks_repo_name: tantikristanti/delft


text:
1#!/usr/bin/python
# -*- coding: UTF-8 -*-

import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt
import cv2

import scipy.misc

def plot_value_array(i, predictions_array, true_label):
  predictions_array, true_label = predictions_array, true_label[i]
  plt.grid(False)
  plt.xticks(range(10))
  plt.yticks([])
  thisplot = plt.bar(range(10), predictions_array, color="#777777")
  plt.ylim([0, 1])
  predicted_label = np.argmax(predictions_array)

  thisplot[predicted_label].set_color('red')
  thisplot[true_label].set_color('blue')

#注意不要太多空白，否则缩放之后，文字信息不够大
#注意要黑底白字，否则识别不行的
#需要二值化，切黑色是0-1的浮点数
#识别：5  0  2  3  6
#未识别：1->6   4->3    7->2    8->3   9->1
img = cv2.imread('test_data/9_2.png')  # 手写数字图像所在位置

img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 转换图像为单通道(灰度图)

ret, thresh_img = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY) #二值化  必须要有ret 两个参数来接受返回值

re_img = cv2.resize(thresh_img, (28,28))

cv2.imshow('image', re_img) #图片很小，在左上角。。。。注意仔细找

re_img = re_img / 255.0

scipy.misc.imsave('xx.jpg', re_img)#保存图片

print(re_img.shape)

cv2.waitKey(0)

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()

plt.figure()
plt.imshow(x_train[7])
plt.colorbar()
plt.grid(False)
plt.show()

x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=6)

model.evaluate(x_test,  y_test, verbose=2)

# 将整个模型保存为HDF5文件
model.save('my_model.h5')

# 重新创建完全相同的模型，包括其权重和优化程序
new_model = keras.models.load_model('my_model.h5')

# 显示网络结构
new_model.summary()

loss, acc = new_model.evaluate(x_test,  y_test, verbose=2)
print("Restored model, accuracy: {:5.2f}%".format(100*acc))

print(x_test[0].shape)

predictions = model.predict(re_img.reshape(1,28,28), batch_size=1)

print(predictions[0])

class_names = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

#用柱状图显示预测结果，第一个参数表示“索引”，在评估时根据对错来显示颜色，这里做预测，所以这个参数没有意义
plot_value_array(4, predictions[0], y_test)
_ = plt.xticks(range(10), class_names, rotation=0)
plt.show()

max_stars_repo_name: asysbang/tensorflow

max_issues_repo_name: asysbang/tensorflow

max_forks_repo_name: asysbang/tensorflow

reason: import numpy,import scipy


text:
import cv2 as cv
import numpy as np

def gaussian_kernel_2d(shape, sigma=0.5) -> np.ndarray:
    if len(shape) != 2:
        raise Exception("Wrong shape")
    x: np.ndarray = cv.getGaussianKernel(shape[0], sigma)
    y: np.ndarray = cv.getGaussianKernel(shape[1], sigma)

    return x.dot(y.T)

UNSHARP_MASKING_5x5_KERNEL = np.array([
    [1, 4, 6, 4, 1],
    [4, 16, 24, 16, 4],
    [6, 24, -476, 24, 6],
    [4, 16, 24, 16, 4],
    [1, 4, 6, 4, 1]
]) / -256


text:
import numpy as np
from os import path as osp
import pickle
from experiments.image_retrieval.base import BaseRetrievalAgent
from experiments.image_retrieval.download_radenovic_data import download_datasets
from experiments.image_retrieval.utils import compute_map_and_print

class RadenovicRetrievalAgent(BaseRetrievalAgent):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.available_datasets = ['oxford5k', 'paris6k', 'roxford5k', 'rparis6k']
        self.data = self._create_retrieval_data()

    @staticmethod
    def config_imname(cfg, i):
        return osp.join(cfg['dir_images'], cfg['imlist'][i] + cfg['ext'])

    @staticmethod
    def config_qimname(cfg, i):
        return osp.join(cfg['dir_images'], cfg['qimlist'][i] + cfg['qext'])

    def _config_dataset(self, dir_main, dataset):
        dataset = dataset.lower()

        if dataset not in self.available_datasets:
            raise ValueError('Unknown dataset: {}!'.format(dataset))

        # loading imlist, qimlist, and gnd, in cfg as a dict
        gnd_fname = osp.join(dir_main, dataset, 'gnd_{}.pkl'.format(dataset))
        with open(gnd_fname, 'rb') as f:
            cfg = pickle.load(f)
        cfg['gnd_fname'] = gnd_fname

        cfg['ext'] = '.jpg'
        cfg['qext'] = '.jpg'
        cfg['dir_data'] = osp.join(dir_main, dataset)
        cfg['dir_images'] = osp.join(cfg['dir_data'], 'jpg')

        cfg['n'] = len(cfg['imlist'])
        cfg['nq'] = len(cfg['qimlist'])

        cfg['im_fname'] = RadenovicRetrievalAgent.config_imname
        cfg['qim_fname'] = RadenovicRetrievalAgent.config_qimname

        cfg['dataset'] = dataset
        return cfg

    def _create_retrieval_data(self):
        cfg = self._config_dataset(self.cfg.paths.datasets_home_dir,
                                   self.retrieval_agent_cfg.dataset)
        q_images = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])]
        db_images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])]
        try:
            bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
        except:
            bbxs = None  # for holidaysmanrot and copydays

        return {'qs': q_images,
                'dbs': db_images,
                'bbxs': bbxs}

    def retrieve(self):
        print(f'Downloading the revisited datasets (roxford5k, rparis6k)...')
        download_datasets(self.cfg.paths.datasets_home_dir)
        print(f'Downloading the revisited datasets (roxford5k, rparis6k)... Done!')

        print(f'Extract keypoints and local descriptors from the database images...')
        self.ldd_model.evaluate(self.data['dbs'])

        print(f'Extract keypoints and local descriptors from the query images...')
        self.ldd_model.evaluate(self.data['qs'], bbxs=self.data['bbxs'])

        # Load GND data
        with open(osp.join(self.retrieval_agent_cfg.paths.img_home_dir,
                           f'gnd_{self.retrieval_agent_cfg.dataset}.pkl'), "rb") as f:
            cfg_gnd = pickle.load(f)

        # Let us load Radenovic image retrieval resutls
        with open(self.retrieval_agent_cfg.paths.radenovic_dict, "rb") as f:
            data_rad = pickle.load(f)

        db_rank_orig_dict = dict(zip(data_rad["db_images"], range(len(data_rad["db_images"]))))
        new_ranks = np.zeros(data_rad["ranks"].shape)
        db_feats_dict = {}
        for i_q, query in enumerate(data_rad["q_images"]):
            ranks_rad = data_rad["ranks"][:, i_q]
            db_rerank_rad = np.array(data_rad["db_images"])[ranks_rad]
            orig_top_k_fnames = db_rerank_rad[:self.retrieval_agent_cfg.topk]

            pkl_fname = query.split('/')[-1][:-3] + 'pkl'
            with open(osp.join(self.retrieval_agent_cfg.output.precomputed_feats_dir, 'jpg', pkl_fname), 'rb') as f:
                q_kpt_data = pickle.load(f)

            inls_count = []
            for db_fname in orig_top_k_fnames:
                pkl_fname = db_fname.split('/')[-1][:-3] + 'pkl'
                if pkl_fname not in db_feats_dict:
                    with open(osp.join(self.retrieval_agent_cfg.output.precomputed_feats_dir, 'jpg', pkl_fname), 'rb') as f:
                        d_kpt_data = pickle.load(f)
                    db_feats_dict[pkl_fname] = d_kpt_data
                else:
                    d_kpt_data = db_feats_dict[pkl_fname]

                n_inliers, _, _ = self.matcher.get_inliers_count(q_kpt_data, d_kpt_data)
                inls_count.append(n_inliers)
            rerank_top_k = [orig_top_k_fnames[idx] for idx in (-np.asarray(inls_count)).argsort()]

            for item in db_rerank_rad[self.retrieval_agent_cfg.topk:]:
                rerank_top_k.append(item)

            rerank_full = rerank_top_k
            for i_d, db_rerank in enumerate(rerank_full):
                new_ranks[i_d, i_q] = db_rank_orig_dict[db_rerank]

            print("Query ", i_q, " out of ", len(data_rad["q_images"]), " processed")

        output = compute_map_and_print(self.retrieval_agent_cfg.dataset,
                                       new_ranks,
                                       cfg_gnd["gnd"])

        # write results to the file
        with open(self.cfg.task.task_params.output.res_fname, "w") as f:
            f.write(f"mAP benchmark:\n")
            mapE, mapM, mapH = output['mAP']
            f.write(f'E: {mapE*100:04.2f}, M: {mapM*100:04.2f}, H: {mapH*100:04.2f}\n')
            f.write(f"mP@k [1, 5, 10] benchmark:\n")
            mprE, mprM, mprH = output['mP@k']
            for type_, res in zip(['E', 'M', 'H'], [mprE, mprM, mprH]):
                f.write(f'{type_}: [{res[0]*100:04.2f}, {res[1]*100:04.2f}, {res[2]*100:04.2f}]\n')

        print('Done')

max_stars_repo_name: imelekhov/HNDesc

max_issues_repo_name: imelekhov/HNDesc

max_forks_repo_name: imelekhov/HNDesc