text:
#######################################################################
# Collection of functions used in EnKF scripts
# (T. Kent:  amttk@leeds.ac.uk)
#######################################################################
'''
Functions:
> generate_truth(): generates nature run at refined resolution (Nk_tr) and saves run at given observing times (assim_time)
> gasp_cohn(): gaspari cohn taper function
>  analysis_step_enkf(): updates ensemble given observations from truth, returns analysis ensemble and forecast ensemble.
'''
import math as m
import numpy as np
from parameters import *
import os

##################################################################
#------------------ Compute nature run ------------------
##################################################################

def generate_truth(U_tr_array, B_tr, Nk_tr, tr_grid, assim_time, f_path_name):

    from f_modRSW import time_step, step_forward_topog

    Kk_tr = tr_grid[0]
    x_tr = tr_grid[1]

    tn = assim_time[0]
    tmax = assim_time[-1]
    dtmeasure = assim_time[1]-assim_time[0]
    tmeasure = dtmeasure

    U_tr = U_tr_array[:,:,0]

    print ' '
    print 'Integrating forward from t =', tn, 'to', tmax,'...'
    print ' '

    index = 1 # for U_tr_array (start from 1 as 0 contains IC).
    while tn < tmax:
        dt = time_step(U_tr,B_tr,Kk_tr,cfl_tr)
        tn = tn + dt

        if tn > tmeasure:
            dt = dt - (tn - tmeasure) + 1e-12
            tn = tmeasure + 1e-12

        U_tr = step_forward_topog(U_tr,B_tr,dt,tn,Nk_tr,Kk_tr)
#        print 't_tr =',tn

        if tn > tmeasure:
            U_tr_array[:,:,index] = U_tr
            print '*** STORE TRUTH at observing time = ',tmeasure,' ***'
            tmeasure = tmeasure + dtmeasure
            index = index + 1

    np.save(f_path_name,U_tr_array)

    print '* DONE: truth array saved to:', f_path_name, ' with shape:', np.shape(U_tr_array), ' *'

    return U_tr_array

##################################################################
# GASPARI-COHN TAPER FUNCTION FOR COV LOCALISATION
# from Jeff Whitaker's github: https://github.com/jswhit/
################################################################
def gaspcohn(r):
    # Gaspari-Cohn taper function
    # very close to exp(-(r/c)**2), where c = sqrt(0.15)
    # r should be >0 and normalized so taper = 0 at r = 1
    rr = 2.0*r
    rr += 1.e-13 # avoid divide by zero warnings from numpy
    taper = np.where(r<=0.5, \
                     ( ( ( -0.25*rr +0.5 )*rr +0.625 )*rr -5.0/3.0 )*rr**2 + 1.0,\
                     np.zeros(r.shape,r.dtype))

    taper = np.where(np.logical_and(r>0.5,r<1.), \
                    ( ( ( ( rr/12.0 -0.5 )*rr +0.625 )*rr +5.0/3.0 )*rr -5.0 )*rr \
                    + 4.0 - 2.0 / (3.0 * rr), taper)
    return taper

##################################################################
# GASPARI-COHN matrix using taper function
################################################################
def gaspcohn_matrix(loc_rho,Nk_fc):
    # construct localisation matrix rho based on Gaspari Cohn function

    rr = np.arange(0,loc_rho,loc_rho/Nk_fc)
    vec = gaspcohn(rr)

    rho = np.zeros((Nk_fc,Nk_fc))

    for i in range(Nk_fc):
        for j in range(Nk_fc):
            rho[i,j] = vec[np.min([abs(i-j),abs(i+Nk_fc-j),abs(i-Nk_fc-j)])]

    rho = np.tile(rho, (Neq,Neq))

    return rho

##################################################################
#'''------------------ ANALYSIS STEP ------------------'''
##################################################################

# v2: (static) additive inflation applied BEFORE assimilation

def analysis_step_enkf(U_fc, U_tr, tmeasure, dtmeasure, index, pars_ob, pars_enda):
    '''
    (Steps refer to algorithm on page 121 of thesis, as do eq. numbers)

    INPUTS
    U_fc: ensemble trajectories in U space, shape (Neq,Nk_fc,n_ens)
    U_tr: truth trajectory in U space, shape (Neq,Nk_tr,Nmeas+1)
    tmeasure: time of assimilation
    dtmeasure: length of window
    pars_ob: vector of parameter values relating to obs (density and error)
    pars_enda: vector of parameters relating to DA fixes (inflation and localisation)
    '''

    print ' '
    print '----------------------------------------------'
    print '------------ ANALYSIS STEP: START ------------'
    print '----------------------------------------------'
    print ' '

    Nk_fc = np.shape(U_fc)[1] # fc resolution (no. of cells)
    Kk_fc = L/Nk_fc # fc resolution (cell length)
    Nk_tr = np.shape(U_tr)[1] # tr resolution (no. of cells)
    n_ens = np.shape(U_fc)[2]
    n_d = Neq*Nk_fc # total number of variables (dgs of freedom)
    dres = Nk_tr/Nk_fc # ratio of resolutions
    inf = pars_enda[0] # inflation factor
    loc = pars_enda[1]
    add_inf = pars_enda[2]
    obs_dens = pars_ob[0]
    ob_noise = pars_ob[1] # 4-vector of obs noise

    print ' '
    print '--------- ANALYSIS: perturbed obs EnKF ---------'
    print ' '
    print 'Assimilation time = ', tmeasure
    print 'Number of ensembles = ', n_ens

    # project truth onto forecast grid so that U and U_tr are the same dimension
    U_tmp = np.empty([Neq,Nk_fc])
    for i in range(0,Nk_fc):
        U_tmp[:,i] = U_tr[:,i*dres:(i+1)*dres,index+1].mean(axis=-1)
    U_tr = U_tmp

    '''
    step 1.(c)
    '''
    # for assimilation, work with [h,u,r]
    U_fc[1:,:,:] = U_fc[1:,:,:]/U_fc[0,:,:]
    U_tr[1:,:] = U_tr[1:,:]/U_tr[0,:]

    X_tr = U_tr.flatten()
    X_tr = np.matrix(X_tr).T

    # state matrix (flatten the array)
    X = np.empty((n_d,n_ens))
    for N in range(0,n_ens):
        X[:,N] = U_fc[:,:,N].flatten()

    # ADDITIVE INFLATION (Gordon: NEEDS MOVING after analysis...)
    cwd = os.getcwd()
    Q = np.load(str(cwd+'/Q_large.npy'))
    print 'max Q value: ', np.max(Q)
    print 'min Q value: ', np.min(Q)

    q = add_inf*np.random.multivariate_normal(np.zeros(n_d), Q, n_ens)
    q = q.T
    X = X + q # x(t+1) = M(x(t)) + q  (eq. 6.5)

    '''
    Step 2.(a)
    '''
    # make pseudo-obs by defining observation operator and adding perturbations
    n_obs = n_d/obs_dens # no. of observations
    print 'Total no. of obs. =', n_obs
    # observation operator
    H = np.zeros((n_obs,n_d))
    row_vec = range(obs_dens,n_d+1,obs_dens)
    for i in range(0,n_obs):
        H[i,row_vec[i]-1]=1

    # for ensemble of observations
    ob_noise = np.repeat(ob_noise,n_obs/Neq)
    obs_pert = ob_noise[:,None]*np.random.randn(n_obs,n_ens)
    print 'obs_pert shape =', np.shape(obs_pert)

    Y_obs = np.empty([n_obs,n_ens])
    Y_mod = np.dot(H, X_tr)
    print 'Y_mod shape =', np.shape(Y_mod)

    Y_obs = Y_mod + obs_pert #y_o = y_m + e_o (eq. 6.6)

    '''
    Step 2.(b)
    '''
    #### CALCULATE KALMAN GAIN, INNOVATIONS, AND ANALYSIS STATES ####
    ONE = np.ones([n_ens,n_ens])
    ONE = ONE/n_ens # NxN array with elements equal to 1/N
    Xbar = np.dot(X,ONE) # mean
    Xdev = X - Xbar # deviations

    # covariance matrix
    Pf = np.dot(Xdev, Xdev.T)
    Pf = Pf/(n_ens - 1)
    Cf = np.corrcoef(Xdev) # correlation matrix

    '''
    Step 2.(c)
    '''
    # construct localisation matrix rho based on Gaspari Cohn function
    loc_rho = pars_enda[1] # loc_rho is form of lengthscale.
    rho = gaspcohn_matrix(loc_rho,Nk_fc)
    print 'loc matrix rho shape: ', np.shape(rho)

    # construct K
    R = ob_noise*ob_noise*np.identity(n_obs) # obs cov matrix
    K = np.dot(H, np.dot(rho*Pf,H.T)) + R # H B H^T + R
    K = np.linalg.inv(K) # [H B H^T + R]^-1
    K = np.dot(np.dot(rho*Pf,H.T), K) # (rho Pf)H^T [H (rho Pf) H^T + R]^-1

    # compute innovation d = Y-H*X
    D = Y_obs - np.dot(H,X)

    # compute analysis
    Xan = X + np.dot(K,D) # kalman update step
    Xanbar = np.dot(Xan,ONE) # analysis mean
    Xandev = Xan - Xanbar  # analysis deviations

    Pa = np.dot(Xandev,Xandev.T)
    Pa = Pa/(n_ens - 1) # analysis covariance matrix
    Ca = np.corrcoef(Xandev) # analysis correlation matrix

    # masks for locating model variables in state vector
    h_mask = range(0,Nk_fc)
    hu_mask = range(Nk_fc,2*Nk_fc)
    hr_mask = range(2*Nk_fc,3*Nk_fc)

    obs_mask = np.array(row_vec[0:n_obs/Neq])-1
    h_obs_mask = range(0,n_obs/Neq)
    hu_obs_mask = range(n_obs/Neq,2*n_obs/Neq)
    hr_obs_mask = range(2*n_obs/Neq,3*n_obs/Neq)

    # make matrices arrays for manipulating and then saving
    X = np.array(X)
    X_tr = np.array(X_tr)
    Xbar = np.array(Xbar)
    Xan = np.array(Xan)
    Xanbar = np.array(Xanbar)
    Xandev = np.array(Xandev)
    Y_obs = np.array(Y_obs)

    # transform to x = (h,hu,hr) for inflation (to avoid double inflation of hu,hr)
    Xan_new = Xan
    Xan_new[hu_mask,:] = Xan_new[hu_mask,:]*Xan_new[h_mask,:]
    Xan_new[hr_mask,:] = Xan_new[hr_mask,:]*Xan_new[h_mask,:]

    ### MULTIPLICATIVE INFLATION
    if inf != 1.0: # inflate the ensemble
        print 'Covariance (ensemble) inflation factor =', inf
        Xanbar_new = np.dot(Xan_new, ONE)
        Xandev_new = Xan_new - Xanbar_new
        Xandev_new = inf*Xandev_new
        Xan_new = Xandev_new + Xanbar_new # inflated analysis ensemble
        Xan_new = np.array(Xan_new)
    else:
        print 'No multiplicative inflation applied'

    # if hr < 0, set to zero:
    hr = Xan_new[hr_mask,:]
    hr[hr < 0.] = 0.
    Xan_new[hr_mask,:] = hr

    # if h < 0, set to epsilon:
    h = Xan_new[h_mask,:]
    h[h < 0.] = 1e-3
    Xan_new[h_mask,:] = h

    # transform from X to U for next integration:
    U_an = np.empty((Neq,Nk_fc,n_ens))
    for N in range(0,n_ens):
        U_an[:,:,N] = Xan_new[:,N].reshape(Neq,Nk_fc)

    # now inflated, transform back to x = (h,u,r) for saving and later plotting
    Xan[h_mask,:] = Xan_new[h_mask,:]
    Xan[hu_mask,:] = Xan_new[hu_mask,:]/Xan_new[h_mask,:]
    Xan[hr_mask,:] = Xan_new[hr_mask,:]/Xan_new[h_mask,:]

    print ' '
    print '--------- CHECK SHAPE OF MATRICES: ---------'
    print ' '
    print 'U_fc shape   :', np.shape(U_fc)
    print 'U_tr shape   :', np.shape(U_tr)
    print 'X_truth shape:', np.shape(X_tr), '( NOTE: should be', n_d,' by 1 )'
    print 'X shape      :', np.shape(X), '( NOTE: should be', n_d,'by', n_ens,')'
    print 'Xbar shape   :', np.shape(Xbar)
    print 'Xdev shape   :', np.shape(Xdev)
    print 'Pf shape     :', np.shape(Pf), '( NOTE: should be n by n square for n=', n_d,')'
    print 'H shape      :', np.shape(H), '( NOTE: should be', n_obs,'by', n_d,')'
    print 'K shape      :', np.shape(K), '( NOTE: should be', n_d,'by', n_obs,')'
    print 'ob_pert shape:', np.shape(obs_pert)
    print 'Y_mod shape  :', np.shape(Y_mod)
    print 'Y_obs shape  :', np.shape(Y_obs)
    print 'Xan shape    :', np.shape(Xan), '( NOTE: should be the same as X shape)'
    print 'U_an shape   :', np.shape(U_an), '( NOTE: should be the same as U_fc shape)'

    ## observational influence diagnostics
    print ' '
    print '--------- OBSERVATIONAL INFLUENCE DIAGNOSTICS:---------'
    print ' '
    print ' Benchmark: global NWP has an average OI of ~0.18... '
    print ' ... high-res. NWP less clear but should be 0.15 - 0.4'
    print 'Check below: '
    HK = np.dot(H,K)
    HKd = np.diag(HK)
    OI = np.trace(HK)/n_obs
    OI_h = np.sum(HKd[h_obs_mask])/n_obs
    OI_hu = np.sum(HKd[hu_obs_mask])/n_obs
    OI_hr = np.sum(HKd[hr_obs_mask])/n_obs
    OI_vec = np.array([OI , OI_h , OI_hu , OI_hr])

    np.set_printoptions(formatter={'float': '{: 0.4f}'.format})
    print 'trace(HK) =', np.trace(HK)
    print 'shape(HK) =', np.shape(HK), '( NOTE: should be', n_obs,'by', n_obs,')'
    print 'OI =', OI
    print 'OI check = ', np.sum(HKd)/n_obs
    print 'OI_h =', OI_h
    print 'OI_hu =', OI_hu
    print 'OI_hr =', OI_hr
    print ' '
    print ' '
    print '----------------------------------------------'
    print '------------- ANALYSIS STEP: END -------------'
    print '----------------------------------------------'
    print ' '

    return U_an, U_fc, X, X_tr, Xan, Y_obs, OI_vec

##################################################################
#'''------------------ ANALYSIS STEP v2 ------------------'''
##################################################################

# v2: (static) additive inflation applied AFTER assimilation

def analysis_step_enkf_v2(U_fc, U_tr, tmeasure, dtmeasure, index, pars_ob, pars_enda):
    '''
        (Steps refer to algorithm on page 121 of thesis, as do eq. numbers)

        INPUTS
        U_fc: ensemble trajectories in U space, shape (Neq,Nk_fc,n_ens)
        U_tr: truth trajectory in U space, shape (Neq,Nk_tr,Nmeas+1)
        tmeasure: time of assimilation
        dtmeasure: length of window
        pars_ob: vector of parameter values relating to obs (density and error)
        pars_enda: vector of parameters relating to DA fixes (inflation and localisation)
        '''

    print ' '
    print '----------------------------------------------'
    print '------------ ANALYSIS STEP: START ------------'
    print '----------------------------------------------'
    print ' '

    Nk_fc = np.shape(U_fc)[1] # fc resolution (no. of cells)
    Kk_fc = L/Nk_fc # fc resolution (cell length)
    Nk_tr = np.shape(U_tr)[1] # tr resolution (no. of cells)
    n_ens = np.shape(U_fc)[2]
    n_d = Neq*Nk_fc # total number of variables (dgs of freedom)
    dres = Nk_tr/Nk_fc # ratio of resolutions
    inf = pars_enda[0] # inflation factor
    loc = pars_enda[1]
    add_inf = pars_enda[2]
    obs_dens = pars_ob[0]
    ob_noise = pars_ob[1] # 4-vector of obs noise

    print ' '
    print '--------- ANALYSIS: perturbed obs EnKF ---------'
    print ' '
    print 'Assimilation time = ', tmeasure
    print 'Number of ensembles = ', n_ens

    # project truth onto forecast grid so that U and U_tr are the same dimension
    U_tmp = np.empty([Neq,Nk_fc])
    for i in range(0,Nk_fc):
        U_tmp[:,i] = U_tr[:,i*dres:(i+1)*dres,index+1].mean(axis=-1)
    U_tr = U_tmp

    '''
        step 1.(c)
        '''
    # for assimilation, work with [h,u,r]
    U_fc[1:,:,:] = U_fc[1:,:,:]/U_fc[0,:,:]
    U_tr[1:,:] = U_tr[1:,:]/U_tr[0,:]

    X_tr = U_tr.flatten()
    X_tr = np.matrix(X_tr).T

    # state matrix (flatten the array)
    X = np.empty((n_d,n_ens))
    for N in range(0,n_ens):
        X[:,N] = U_fc[:,:,N].flatten()

    '''
        Step 2.(a)
        '''
    # make pseudo-obs by defining observation operator and adding perturbations
    n_obs = n_d/obs_dens # no. of observations
    print 'Total no. of obs. =', n_obs
    # observation operator
    H = np.zeros((n_obs,n_d))
    row_vec = range(obs_dens,n_d+1,obs_dens)
    for i in range(0,n_obs):
        H[i,row_vec[i]-1]=1

    # for ensemble of observations
    ob_noise = np.repeat(ob_noise,n_obs/Neq)
    obs_pert = ob_noise[:,None]*np.random.randn(n_obs,n_ens)
    print 'obs_pert shape =', np.shape(obs_pert)

    Y_obs = np.empty([n_obs,n_ens])
    Y_mod = np.dot(H, X_tr)
    print 'Y_mod shape =', np.shape(Y_mod)

    Y_obs = Y_mod + obs_pert #y_o = y_m + e_o (eq. 6.6)

    '''
        Step 2.(b)
        '''
    #### CALCULATE KALMAN GAIN, INNOVATIONS, AND ANALYSIS STATES ####
    ONE = np.ones([n_ens,n_ens])
    ONE = ONE/n_ens # NxN array with elements equal to 1/N
    Xbar = np.dot(X,ONE) # mean
    Xdev = X - Xbar # deviations

    # covariance matrix
    Pf = np.dot(Xdev, Xdev.T)
    Pf = Pf/(n_ens - 1)
    Cf = np.corrcoef(Xdev) # correlation matrix

    '''
        Step 2.(c)
        '''
    # construct localisation matrix rho based on Gaspari Cohn function
    loc_rho = pars_enda[1] # loc_rho is form of lengthscale.
    rho = gaspcohn_matrix(loc_rho,Nk_fc)
    print 'loc matrix rho shape: ', np.shape(rho)

    # construct K
    R = ob_noise*ob_noise*np.identity(n_obs) # obs cov matrix
    K = np.dot(H, np.dot(rho*Pf,H.T)) + R # H B H^T + R
    K = np.linalg.inv(K) # [H B H^T + R]^-1
    K = np.dot(np.dot(rho*Pf,H.T), K) # (rho Pf)H^T [H (rho Pf) H^T + R]^-1

    # compute innovation d = Y-H*X
    D = Y_obs - np.dot(H,X)

    # compute analysis
    Xan = X + np.dot(K,D) # kalman update step
    Xanbar = np.dot(Xan,ONE) # analysis mean
    Xandev = Xan - Xanbar  # analysis deviations

    Pa = np.dot(Xandev,Xandev.T)
    Pa = Pa/(n_ens - 1) # analysis covariance matrix
    Ca = np.corrcoef(Xandev) # analysis correlation matrix

    ### ADDITIVE INFLATION ...### (NEEDS MOVING after analysis... here)
    cwd = os.getcwd()
    Q = np.load(str(cwd+'/Q_large.npy'))

    q = add_inf*np.random.multivariate_normal(np.zeros(n_d), Q, n_ens)
    q = q.T
    Xan = Xan + q # x(t+1) = M(x(t)) + q

    # masks for locating model variables in state vector
    h_mask = range(0,Nk_fc)
    hu_mask = range(Nk_fc,2*Nk_fc)
    hr_mask = range(2*Nk_fc,3*Nk_fc)

    obs_mask = np.array(row_vec[0:n_obs/Neq])-1
    h_obs_mask = range(0,n_obs/Neq)
    hu_obs_mask = range(n_obs/Neq,2*n_obs/Neq)
    hr_obs_mask = range(2*n_obs/Neq,3*n_obs/Neq)

    # make matrices arrays for manipulating and then saving
    X = np.array(X)
    X_tr = np.array(X_tr)
    Xbar = np.array(Xbar)
    Xan = np.array(Xan)
    Xanbar = np.array(Xanbar)
    Xandev = np.array(Xandev)
    Y_obs = np.array(Y_obs)

    # transform to x = (h,hu,hr) for inflation (to avoid double inflation of hu,hr)
    Xan_new = Xan
    Xan_new[hu_mask,:] = Xan_new[hu_mask,:]*Xan_new[h_mask,:]
    Xan_new[hr_mask,:] = Xan_new[hr_mask,:]*Xan_new[h_mask,:]

    ### MULTIPLICATIVE INFLATION
    if inf != 1.0: # inflate the ensemble
        print 'Covariance (ensemble) inflation factor =', inf
        Xanbar_new = np.dot(Xan_new, ONE)
        Xandev_new = Xan_new - Xanbar_new
        Xandev_new = inf*Xandev_new
        Xan_new = Xandev_new + Xanbar_new # inflated analysis ensemble
        Xan_new = np.array(Xan_new)
    else:
        print 'No multiplicative inflation applied'

    # if hr < 0, set to zero:
    hr = Xan_new[hr_mask,:]
    hr[hr < 0.] = 0.
    Xan_new[hr_mask,:] = hr

    # if h < 0, set to epsilon:
    h = Xan_new[h_mask,:]
    h[h < 0.] = 1e-3
    Xan_new[h_mask,:] = h

    # transform from X to U for next integration:
    U_an = np.empty((Neq,Nk_fc,n_ens))
    for N in range(0,n_ens):
        U_an[:,:,N] = Xan_new[:,N].reshape(Neq,Nk_fc)

    # now inflated, transform back to x = (h,u,r) for saving and later plotting
    Xan[h_mask,:] = Xan_new[h_mask,:]
    Xan[hu_mask,:] = Xan_new[hu_mask,:]/Xan_new[h_mask,:]
    Xan[hr_mask,:] = Xan_new[hr_mask,:]/Xan_new[h_mask,:]

    print ' '
    print '--------- CHECK SHAPE OF MATRICES: ---------'
    print ' '
    print 'U_fc shape   :', np.shape(U_fc)
    print 'U_tr shape   :', np.shape(U_tr)
    print 'X_truth shape:', np.shape(X_tr), '( NOTE: should be', n_d,' by 1 )'
    print 'X shape      :', np.shape(X), '( NOTE: should be', n_d,'by', n_ens,')'
    print 'Xbar shape   :', np.shape(Xbar)
    print 'Xdev shape   :', np.shape(Xdev)
    print 'Pf shape     :', np.shape(Pf), '( NOTE: should be n by n square for n=', n_d,')'
    print 'H shape      :', np.shape(H), '( NOTE: should be', n_obs,'by', n_d,')'
    print 'K shape      :', np.shape(K), '( NOTE: should be', n_d,'by', n_obs,')'
    print 'ob_pert shape:', np.shape(obs_pert)
    print 'Y_mod shape  :', np.shape(Y_mod)
    print 'Y_obs shape  :', np.shape(Y_obs)
    print 'Xan shape    :', np.shape(Xan), '( NOTE: should be the same as X shape)'
    print 'U_an shape   :', np.shape(U_an), '( NOTE: should be the same as U_fc shape)'

    ## observational influence diagnostics
    print ' '
    print '--------- OBSERVATIONAL INFLUENCE DIAGNOSTICS:---------'
    print ' '
    print ' Benchmark: global NWP has an average OI of ~0.18... '
    print ' ... high-res. NWP less clear but should be 0.15 - 0.4'
    print 'Check below: '
    HK = np.dot(H,K)
    HKd = np.diag(HK)
    OI = np.trace(HK)/n_obs
    OI_h = np.sum(HKd[h_obs_mask])/n_obs
    OI_hu = np.sum(HKd[hu_obs_mask])/n_obs
    OI_hr = np.sum(HKd[hr_obs_mask])/n_obs
    OI_vec = np.array([OI , OI_h , OI_hu , OI_hr])

    np.set_printoptions(formatter={'float': '{: 0.4f}'.format})
    print 'trace(HK) =', np.trace(HK)
    print 'shape(HK) =', np.shape(HK), '( NOTE: should be', n_obs,'by', n_obs,')'
    print 'OI =', OI
    print 'OI check = ', np.sum(HKd)/n_obs
    print 'OI_h =', OI_h
    print 'OI_hu =', OI_hu
    print 'OI_hr =', OI_hr
    print ' '
    print ' '
    print '----------------------------------------------'
    print '------------- ANALYSIS STEP: END -------------'
    print '----------------------------------------------'
    print ' '

    return U_an, U_fc, X, X_tr, Xan, Y_obs, OI_vec

reason: import numpy,from numpy


text:
from fewshot_re_kit.data_loader import get_loader, get_loader_pair, get_loader_unsupervised
from fewshot_re_kit.framework import FewShotREFramework
from fewshot_re_kit.sentence_encoder import CNNSentenceEncoder, BERTSentenceEncoder, BERTPAIRSentenceEncoder
import models
from models.regrab import REGRAB
from models.d import Discriminator
import sys
import torch
from torch import optim, nn
import numpy as np
import json
import argparse
import os
import random
import pickle

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--train', default='train_wiki',
            help='train file')
    parser.add_argument('--val', default='val_wiki',
            help='val file')
    parser.add_argument('--test', default='test_wiki',
            help='test file')
    parser.add_argument('--adv', default=None,
            help='adv file')
    parser.add_argument('--trainN', default=10, type=int,
            help='N in train')
    parser.add_argument('--N', default=5, type=int,
            help='N way')
    parser.add_argument('--K', default=5, type=int,
            help='K shot')
    parser.add_argument('--Q', default=5, type=int,
            help='Num of query per class')
    parser.add_argument('--batch_size', default=4, type=int,
            help='batch size')
    parser.add_argument('--train_iter', default=30000, type=int,
            help='num of iters in training')
    parser.add_argument('--val_iter', default=1000, type=int,
            help='num of iters in validation')
    parser.add_argument('--test_iter', default=10000, type=int,
            help='num of iters in testing')
    parser.add_argument('--val_step', default=2000, type=int,
           help='val after training how many iters')
    parser.add_argument('--model', default='regrab',
            help='model name')
    parser.add_argument('--encoder', default='bert',
            help='encoder: cnn or bert')
    parser.add_argument('--max_length', default=128, type=int,
           help='max length')
    parser.add_argument('--lr', default=1e-1, type=float,
           help='learning rate')
    parser.add_argument('--weight_decay', default=1e-5, type=float,
           help='weight decay')
    parser.add_argument('--dropout', default=0.0, type=float,
           help='dropout rate')
    parser.add_argument('--na_rate', default=0, type=int,
           help='NA rate (NA = Q * na_rate)')
    parser.add_argument('--grad_iter', default=1, type=int,
           help='accumulate gradient every x iterations')
    parser.add_argument('--optim', default='sgd',
           help='sgd / adam / bert_adam')
    parser.add_argument('--hidden_size', default=230, type=int,
           help='hidden size')
    parser.add_argument('--load_ckpt', default=None,
           help='load ckpt')
    parser.add_argument('--save_ckpt', default=None,
           help='save ckpt')
    parser.add_argument('--fp16', action='store_true',
           help='use nvidia apex fp16')
    parser.add_argument('--only_test', action='store_true',
           help='only test')
    parser.add_argument('--pair', action='store_true',
           help='use pair model')
    parser.add_argument('--eps', default=0.1, type=float,
           help='step size for SG-MCMC')
    parser.add_argument('--temp', default=10.0, type=float,
           help='temperature for softmax')
    parser.add_argument('--step', default=5, type=int,
           help='steps for SG-MCMC')
    parser.add_argument('--smp', default=10, type=int,
           help='samples for SG-MCMC')
    parser.add_argument('--ratio', default=0.01, type=float,
           help='decay ratio of step size for SG-MCMC')
    parser.add_argument('--wtp', default=0.1, type=float,
           help='weight of the prior term')
    parser.add_argument('--wtn', default=1.0, type=float,
           help='weight of the noise term')
    parser.add_argument('--wtb', default=0.0, type=float,
           help='weight of the background term')
    parser.add_argument('--metric', default='dot',
           help='similarity metric (dot or l2)')
    parser.add_argument('--seed', default=1234, type=int,
           help='random seed')

    opt = parser.parse_args()
    trainN = opt.trainN
    N = opt.N
    K = opt.K
    Q = opt.Q
    batch_size = opt.batch_size
    model_name = opt.model
    encoder_name = opt.encoder
    max_length = opt.max_length

    torch.manual_seed(opt.seed)
    np.random.seed(opt.seed)
    random.seed(opt.seed)
    torch.cuda.manual_seed(opt.seed)

    print("{}-way-{}-shot Few-Shot Relation Classification".format(N, K))
    print("model: {}".format(model_name))
    print("encoder: {}".format(encoder_name))
    print("max_length: {}".format(max_length))

    if encoder_name == 'cnn':
        try:
            glove_mat = np.load('./pretrain/glove/glove_mat.npy')
            glove_word2id = json.load(open('./pretrain/glove/glove_word2id.json'))
        except:
            raise Exception("Cannot find glove files. Run glove/download_glove.sh to download glove files.")
        sentence_encoder = CNNSentenceEncoder(
                glove_mat,
                glove_word2id,
                max_length)
    elif encoder_name == 'bert':
        if opt.pair:
            sentence_encoder = BERTPAIRSentenceEncoder(
                    './pretrain/bert-base-uncased',
                    max_length)
        else:
            sentence_encoder = BERTSentenceEncoder(
                    './pretrain/bert-base-uncased',
                    max_length)
    else:
        raise NotImplementedError

    if opt.pair:
        train_data_loader = get_loader_pair(opt.train, sentence_encoder,
                N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size)
        val_data_loader = get_loader_pair(opt.val, sentence_encoder,
                N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size)
        test_data_loader = get_loader_pair(opt.test, sentence_encoder,
                N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size)
    else:
        train_data_loader = get_loader(opt.train, sentence_encoder,
                N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size)
        val_data_loader = get_loader(opt.val, sentence_encoder,
                N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size)
        test_data_loader = get_loader(opt.test, sentence_encoder,
                N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size)
        if opt.adv:
           adv_data_loader = get_loader_unsupervised(opt.adv, sentence_encoder,
                N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size)

    if opt.optim == 'sgd':
        pytorch_optim = optim.SGD
    elif opt.optim == 'adam':
        pytorch_optim = optim.Adam
    elif opt.optim == 'bert_adam':
        from transformers import AdamW
        pytorch_optim = AdamW
    else:
        raise NotImplementedError
    if opt.adv:
        d = Discriminator(opt.hidden_size)
        framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader, adv_data_loader, adv=opt.adv, d=d)
    else:
        framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader)

    prefix = '-'.join([model_name, encoder_name, opt.train, opt.val, str(N), str(K)])
    if opt.adv is not None:
        prefix += '-adv_' + opt.adv
    if opt.na_rate != 0:
        prefix += '-na{}'.format(opt.na_rate)

    # Loading relation embeddings.
    data = pickle.load(open('./data/embeddings.pkl', 'rb'))
    rellist = data['relations']
    relemb = data['embeddings']
    array0 = np.zeros((1,relemb.shape[1]), dtype=relemb.dtype)
    relemb = np.concatenate([array0, relemb], axis=0)
    rel2id = dict([(rel, k + 1) for k, rel in enumerate(rellist)])
    # Loading relation graphs.
    with open('./data/graph.txt', 'r') as fi:
        us, vs, ws = [], [], []
        for line in fi:
            items = line.strip().split('\t')
            us += [rel2id[items[0]]]
            vs += [rel2id[items[1]]]
            ws += [float(items[2])]
        index = torch.LongTensor([us, vs])
        value = torch.Tensor(ws)
        shape = torch.Size([len(rel2id) + 1, len(rel2id) + 1])
        reladj = torch.sparse.FloatTensor(index, value, shape).cuda()
    # End
    model = REGRAB(sentence_encoder, hidden_size=opt.hidden_size, eps=opt.eps, temp=opt.temp, step=opt.step, smp=opt.smp, ratio=opt.ratio, wtp=opt.wtp, wtn=opt.wtn, wtb=opt.wtb, metric=opt.metric)
    model.set_relemb(rel2id, relemb)
    model.set_reladj(reladj)

    if not os.path.exists('checkpoint'):
        os.mkdir('checkpoint')
    ckpt = 'checkpoint/{}.pth.tar'.format(prefix)
    if opt.save_ckpt:
        ckpt = opt.save_ckpt

    if torch.cuda.is_available():
        model.cuda()

    if not opt.only_test:
        if encoder_name == 'bert':
            bert_optim = True
        else:
            bert_optim = False

        framework.train(model, prefix, batch_size, trainN, N, K, Q,
                pytorch_optim=pytorch_optim, load_ckpt=opt.load_ckpt, save_ckpt=ckpt,
                na_rate=opt.na_rate, val_step=opt.val_step, fp16=opt.fp16, pair=opt.pair,
                train_iter=opt.train_iter, val_iter=opt.val_iter, bert_optim=bert_optim)
    else:
        ckpt = opt.load_ckpt

    acc = framework.eval(model, batch_size, N, K, Q, opt.test_iter, na_rate=opt.na_rate, ckpt=ckpt, pair=opt.pair)
    print("RESULT: %.2f" % (acc * 100))

if __name__ == "__main__":
    main()

ext: py

max_stars_repo_name: DeepGraphLearning/FewShotRE

max_issues_repo_name: DeepGraphLearning/FewShotRE

max_forks_repo_name: DeepGraphLearning/FewShotRE

reason: import numpy


text:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import six.moves
from datetime import datetime
import sys
import math
import time
from data import inputs, standardize_image
import numpy as np
import tensorflow as tf
from detect import *
import re

RESIZE_AOI = 256
RESIZE_FINAL = 227

# Modifed from here
# http://stackoverflow.com/questions/3160699/python-progress-bar#3160819
class ProgressBar(object):
    DEFAULT = 'Progress: %(bar)s %(percent)3d%%'
    FULL = '%(bar)s %(current)d/%(total)d (%(percent)3d%%) %(remaining)d to go'

    def __init__(self, total, width=40, fmt=DEFAULT, symbol='='):
        assert len(symbol) == 1

        self.total = total
        self.width = width
        self.symbol = symbol
        self.fmt = re.sub(r'(?P<name>%\(.+?\))d',
            r'\g<name>%dd' % len(str(total)), fmt)

        self.current = 0

    def update(self, step=1):
        self.current += step
        percent = self.current / float(self.total)
        size = int(self.width * percent)
        remaining = self.total - self.current
        bar = '[' + self.symbol * size + ' ' * (self.width - size) + ']'

        args = {
            'total': self.total,
            'bar': bar,
            'current': self.current,
            'percent': percent * 100,
            'remaining': remaining
        }
        six.print_('\r' + self.fmt % args, end='')

    def done(self):
        self.current = self.total
        self.update(step=0)
        print('')

# Read image files
class ImageCoder(object):

    def __init__(self):
        # Create a single Session to run all image coding calls.
        config = tf.ConfigProto(allow_soft_placement=True)
        self._sess = tf.Session(config=config)

        # Initializes function that converts PNG to JPEG data.
        self._png_data = tf.placeholder(dtype=tf.string)
        image = tf.image.decode_png(self._png_data, channels=3)
        self._png_to_jpeg = tf.image.encode_jpeg(image, format='rgb', quality=100)

        # Initializes function that decodes RGB JPEG data.
        self._decode_jpeg_data = tf.placeholder(dtype=tf.string)
        self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3)
        self.crop = tf.image.resize_images(self._decode_jpeg, (RESIZE_AOI, RESIZE_AOI))

    def png_to_jpeg(self, image_data):
        return self._sess.run(self._png_to_jpeg,
                              feed_dict={self._png_data: image_data})

    def decode_jpeg(self, image_data):
        image = self._sess.run(self.crop, #self._decode_jpeg,
                               feed_dict={self._decode_jpeg_data: image_data})

        assert len(image.shape) == 3
        assert image.shape[2] == 3
        return image

def _is_png(filename):
    """Determine if a file contains a PNG format image.
    Args:
    filename: string, path of the image file.
    Returns:
    boolean indicating if the image is a PNG.
    """
    return '.png' in filename

def make_multi_image_batch(filenames, coder):
    """Process a multi-image batch, each with a single-look
    Args:
    filenames: list of paths
    coder: instance of ImageCoder to provide TensorFlow image coding utils.
    Returns:
    image_buffer: string, JPEG encoding of RGB image.
    """

    images = []

    for filename in filenames:
        with tf.gfile.FastGFile(filename, 'rb') as f:
            image_data = f.read()
        # Convert any PNG to JPEG's for consistency.
        if _is_png(filename):
            print('Converting PNG to JPEG for %s' % filename)
            image_data = coder.png_to_jpeg(image_data)

        image = coder.decode_jpeg(image_data)

        crop = tf.image.resize_images(image, (RESIZE_FINAL, RESIZE_FINAL))
        image = standardize_image(crop)
        images.append(image)
    image_batch = tf.stack(images)
    return image_batch

def make_multi_crop_batch(filename, coder):
    """Process a single image file.
    Args:
    filename: string, path to an image file e.g., '/path/to/example.JPG'.
    coder: instance of ImageCoder to provide TensorFlow image coding utils.
    Returns:
    image_buffer: string, JPEG encoding of RGB image.
    """
    # Read the image file.
    with tf.gfile.FastGFile(filename, 'rb') as f:
        image_data = f.read()

    # Convert any PNG to JPEG's for consistency.
    if _is_png(filename):
        print('Converting PNG to JPEG for %s' % filename)
        image_data = coder.png_to_jpeg(image_data)

    image = coder.decode_jpeg(image_data)

    crops = []
    print('Running multi-cropped image')
    h = image.shape[0]
    w = image.shape[1]
    hl = h - RESIZE_FINAL
    wl = w - RESIZE_FINAL

    crop = tf.image.resize_images(image, (RESIZE_FINAL, RESIZE_FINAL))
    crops.append(standardize_image(crop))
    crops.append(standardize_image(tf.image.flip_left_right(crop)))

    corners = [ (0, 0), (0, wl), (hl, 0), (hl, wl), (int(hl/2), int(wl/2))]
    for corner in corners:
        ch, cw = corner
        cropped = tf.image.crop_to_bounding_box(image, ch, cw, RESIZE_FINAL, RESIZE_FINAL)
        crops.append(standardize_image(cropped))
        flipped = standardize_image(tf.image.flip_left_right(cropped))
        crops.append(standardize_image(flipped))

    image_batch = tf.stack(crops)
    return image_batch

def face_detection_model(model_type, model_path):
    model_type_lc = model_type.lower()
    if model_type_lc == 'yolo_tiny':
        from yolodetect import PersonDetectorYOLOTiny
        return PersonDetectorYOLOTiny(model_path)
    elif model_type_lc == 'yolo_face':
        from yolodetect import FaceDetectorYOLO
        return FaceDetectorYOLO(model_path)
    elif model_type == 'dlib':
        from dlibdetect import FaceDetectorDlib
        return FaceDetectorDlib(model_path)
    return ObjectDetectorCascadeOpenCV(model_path)

max_stars_repo_name: pengbohua/age-gender

max_issues_repo_name: pengbohua/age-gender

max_forks_repo_name: pengbohua/age-gender


text:
"""
Creates small figure that shows H_row_EC gives an unbiased estimate of H (in contrast to just H_row)
Uses variant 1 for with and without EC
For each length l, a different entry from the first row of H_rows are chosen (H[0][(l+1) % 2] which corresponds to the max.
(Just using max will sometimes choose another entry and thus not give an unbiased estimate)

First version: Nov 3, 2016
This version: March 3, 2020
Author: Wolfgang Gatterbauer
"""

import numpy as np
import time
import datetime
import random
# import os                       # for displaying created PDF
import sys
sys.path.append('./../sslh')
from fileInteraction import save_csv_record
from utils import from_dictionary_beliefs, create_parameterized_H, replace_fraction_of_rows, showfig
from graphGenerator import (planted_distribution_model_H,
                            calculate_average_outdegree_from_graph)
from estimation import M_observed, H_observed
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import pandas as pd
pd.set_option('display.max_columns', None)      # show all columns

# -- Determine path to data *irrespective* of where the file is run from
from os.path import abspath, dirname, join
from inspect import getfile, currentframe
current_path = dirname(abspath(getfile(currentframe())))
figure_directory = join(current_path, 'figs')
data_directory = join(current_path, 'datacache')

# %% -- main method
def run(choice, variant, create_data=False, show_plot=False, create_pdf=False, show_pdf=False):
    """main parameterized method to produce all figures.
    Can be run from external jupyther notebook or method to produce all figures in PDF
    """

    # %% -- Setup
    CREATE_DATA = create_data
    CHOICE = choice
    VARIANT = variant
    SHOW_PLOT = show_plot
    CREATE_PDF = create_pdf
    SHOW_PDF = show_pdf
    SHOW_TITLE = True
    LEGEND_MATCH_COLORS = False
    SHOW_DISTRIBUTION_IN_TITLE = True

    SHOW_BACKTRACK_ESTIMATE = True
    SHOW_NONBACKTRACK_ESTIMATE = True
    plot_colors = ['darkgreen', 'darkorange', 'blue']
    label_vec = [r'$\mathbf{H}^{\ell}\,\,\,\,$', r'$\mathbf{\hat P}^{(\ell)}$', r'$\mathbf{\hat P}_{\mathrm{NB}}^{(\ell)}$']

    csv_filename = 'Fig_Backtracking_Advantage_{}.csv'.format(CHOICE)
    fig_filename = 'Fig_Backtracking_Advantage_{}-{}.pdf'.format(CHOICE, VARIANT)

    header = ['currenttime',
              'choice',     # H, Hrow, HrowEC
              'l',
              'valueH',    # maximal values in first row of H
              'valueM']    # average value across entries in M
    if CREATE_DATA:
        save_csv_record(join(data_directory, csv_filename), header, append=False)

    # %% -- Default parameters
    ymin=0.3
    ymax=1
    exponent = None

    # %% -- CHOICES and VARIANTS
    if CHOICE == 1:     # n=1000, shows NB to be slight lower for l=2: probably due to sampling issues (d=3, thus very few points available)
        n = 1000
        h = 8
        d = 3
        f = 0.1
        distribution = 'uniform'
        rep = 10000
        length = 8

    elif CHOICE == 2:
        n = 1000
        h = 8
        d = 10
        f = 0.1
        distribution = 'uniform'
        rep = 10000
        length = 8

    elif CHOICE == 3: # nice: shows nicely that difference is even bigger for smaller h
        n = 1000
        h = 3
        d = 10
        f = 0.1
        distribution = 'uniform'
        rep = 10000
        length = 8
        ymax = 0.8

    elif CHOICE == 4:
        n = 10000
        h = 3
        d = 10
        f = 0.1
        distribution = 'uniform'
        rep = 100
        length = 8
        ymin = 0.333
        ymax = 0.65

    elif CHOICE == 5:
        n = 10000
        h = 3
        d = 3
        f = 0.1
        distribution = 'uniform'
        rep = 1000
        length = 8

    elif CHOICE == 6:  # n=1000, the powerlaw problem with small graphs and high exponent
        n = 1000
        h = 8
        d = 3
        f = 0.1
        distribution = 'powerlaw'
        exponent = -0.5
        rep = 10000
        length = 8

    elif CHOICE == 7:
        n = 10000
        h = 8
        d = 3
        f = 0.1
        distribution = 'uniform'
        rep = 1000
        length = 8
        # ymin = 0.4
        ymax = 1

    elif CHOICE == 8:
        n = 10000
        h = 8
        d = 10
        f = 0.1
        distribution = 'uniform'
        rep = 1000
        length = 8
        # ymin = 0.4
        ymax = 1

    elif CHOICE == 9: # shows lower NB due to problem with sampling from high powerlaw -0.5
        n = 10000
        h = 8
        d = 10
        f = 0.1
        distribution = 'powerlaw'
        exponent = -0.5
        rep = 1000
        length = 8

    elif CHOICE == 10:
        n = 10000
        h = 8
        d = 3
        f = 0.1
        distribution = 'powerlaw'
        exponent = -0.5
        rep = 1000
        length = 8

    elif CHOICE == 11:    # problem: shows that NB is too low (probably because of problem with sampling from -0.5 factor)
        n = 1000
        h = 8
        d = 10
        f = 0.1
        distribution = 'powerlaw'
        exponent = -0.5
        rep = 1000
        length = 8

    elif CHOICE == 12:  # problem: shows no problem with NB (probably because no problem with sampling from -0.2 factor)
        n = 1000
        h = 8
        d = 10
        f = 0.1
        distribution = 'powerlaw'
        exponent = -0.2
        rep = 1000
        length = 8

    elif CHOICE == 20:
        n = 10000
        h = 3
        d = 10
        f = 0.1
        distribution = 'powerlaw'
        exponent = -0.3
        rep = 1000
        length = 8
        ymin = 0.333
        ymax = 0.65

    elif CHOICE == 21:          # originally used before color change
        n = 10000
        h = 3
        d = 25
        f = 0.1
        distribution = 'powerlaw'
        exponent = -0.3
        rep = 1000
        length = 8
        ymin = 0.333
        ymax = 0.65

        if VARIANT == 1:
            SHOW_TITLE = False
            plot_colors = ['red', 'blue', 'darkorange']
            label_vec = [r'$\mathbf{H}^{\ell}\quad\quad$', 'naive', 'better']
            LEGEND_MATCH_COLORS = True

        if VARIANT == 2:
            SHOW_TITLE = False
            plot_colors = ['red', 'blue', 'darkorange']
            label_vec = [r'$\mathbf{H}^{\ell}\quad\quad$', 'naive', 'better']
            SHOW_NONBACKTRACK_ESTIMATE = False
            LEGEND_MATCH_COLORS = True

        if VARIANT == 3:
            SHOW_TITLE = False
            plot_colors = ['red', 'blue', 'darkorange']
            label_vec = [r'$\mathbf{H}^{\ell}\quad\quad$', 'naive', 'better']
            SHOW_BACKTRACK_ESTIMATE = False
            SHOW_NONBACKTRACK_ESTIMATE = False
            LEGEND_MATCH_COLORS = True

        if VARIANT == 4:
            plot_colors = ['red', 'blue', 'darkorange']
            LEGEND_MATCH_COLORS = True

    elif CHOICE == 25:
        n = 10000
        h = 8
        d = 5
        f = 0.1
        distribution = 'uniform'
        rep = 1000
        length = 8

    elif CHOICE == 26:
        n = 10000
        h = 8
        d = 25
        f = 0.1
        distribution = 'uniform'
        rep = 1000
        length = 8
        ymax = 0.9
        ymin = 0.4

    elif CHOICE == 27:
        n = 10000
        h = 8
        d = 10
        f = 0.1
        distribution = 'powerlaw'
        exponent = -0.3
        rep = 1000
        length = 8
        ymax = 0.9
        ymin = 0.33

    elif CHOICE == 31:
        n = 10000
        h = 3
        d = 10
        f = 0.1
        distribution = 'uniform'
        length = 8
        ymin = 0.333
        ymax = 0.65
        SHOW_DISTRIBUTION_IN_TITLE = False
        plot_colors = ['red', 'blue', 'darkorange']
        LEGEND_MATCH_COLORS = True

        if VARIANT == 0:
            rep = 1000

        if VARIANT == 1:
            rep = 20

    else:
        raise Warning("Incorrect choice!")

    k = 3
    a = 1
    alpha0 = np.array([a, 1., 1.])
    alpha0 = alpha0 / np.sum(alpha0)
    H0 = create_parameterized_H(k, h, symmetric=True)
    RANDOMSEED = None  # For repeatability
    random.seed(RANDOMSEED)  # seeds some other python random generator
    np.random.seed(seed=RANDOMSEED)  # seeds the actually used numpy random generator; both are used and thus needed
    # print("CHOICE: {}".format(CHOICE))

    # %% -- Create data
    if CREATE_DATA:

        # Calculations H
        print("Max entry of first rows of powers of H0:")
        for l in range(1, length + 1):
            valueH = np.max(np.linalg.matrix_power(H0, l)[0])

            tuple = [str(datetime.datetime.now())]
            text = ['H',
                    l,
                    valueH,
                    '']
            text = np.asarray(text)  # without np, entries get ugly format
            tuple.extend(text)
            print("{}: {}".format(l, valueH))
            save_csv_record(join(data_directory, csv_filename), tuple)

        # Calculations Hrow and HrowEC
        for r in range(rep):
            print('Repetition {}'.format(r))

            # Create graph
            start = time.time()
            W, Xd = planted_distribution_model_H(n, alpha=alpha0, H=H0, d_out=d,    # notice that for undirected graphs, actual degree = 2*d
                                                      distribution=distribution,
                                                      exponent=exponent,
                                                      directed=False,
                                                      debug=False)
            X0 = from_dictionary_beliefs(Xd)
            X1, ind = replace_fraction_of_rows(X0, 1 - f)
            time_calc = time.time() - start
            # print("\nTime for graph:{}".format(time_calc))

            print("Average outdegree: {}".format(calculate_average_outdegree_from_graph(W)))

            # Calculate H_vec and M_vec versions (M_vec to calculate the average number of entries in M)
            H_vec = H_observed(W, X1, distance=length, NB=False, variant=1)
            H_vec_EC = H_observed(W, X1, distance=length, NB=True, variant=1)
            M_vec = M_observed(W, X1, distance=length, NB=False)
            M_vec_EC = M_observed(W, X1, distance=length, NB=True)

            # Calculation H_vec
            # print("Max entry of first rows of H_vec")
            for l, H in enumerate(H_vec):
                valueH = H[0][(l + 1) % 2]       # better than 'value = np.max(H[0])', otherwise sometimes chooses another higher entry -> biased estimate
                valueM = np.average(M_vec[l+1])
                # print(M_vec[l+1])
                # print(valueM)

                tuple = [str(datetime.datetime.now())]
                text = ['Hrow',
                        l + 1,
                        valueH,
                        valueM]
                text = np.asarray(text)         # without np, entries get ugly format
                tuple.extend(text)
                # print("{}: {}".format(l + 1, value))
                save_csv_record(join(data_directory, csv_filename), tuple)

            # Calculation H_vec_EC
            # print("Max entry of first rows of H_vec_EC")
            for l, H in enumerate(H_vec_EC):
                valueH = H[0][(l + 1) % 2]
                valueM = np.average(M_vec_EC[l+1])
                # print(M_vec_EC[l+1])
                # print(valueM)

                tuple = [str(datetime.datetime.now())]
                text = ['HrowEC',
                        l + 1,
                        valueH,
                        valueM]
                text = np.asarray(text)  # without np, entries get ugly format
                tuple.extend(text)
                # print("{}: {}".format(l + 1, value))
                save_csv_record(join(data_directory, csv_filename), tuple)

    #%% -- Read, aggregate, and pivot data
    df1 = pd.read_csv(join(data_directory, csv_filename))
    # print("\n-- df1 (length {}):\n{}".format(len(df1.index), df1.head(15)))
    df2 = df1.groupby(['choice', 'l']).agg \
        ({'valueH': [np.mean, np.std, np.size],  # Multiple Aggregates
          'valueM': [np.mean],
          })
    df2.columns = ['_'.join(col).strip() for col in df2.columns.values]     # flatten the column hierarchy
    df2.reset_index(inplace=True)                                           # remove the index hierarchy
    df2.rename(columns={'valueH_size': 'count'}, inplace=True)
    # print("\n-- df2 (length {}):\n{}".format(len(df2.index), df2.head(30)))
    df3 = pd.pivot_table(df2, index=['l'], columns=['choice'], values=['valueH_mean', 'valueH_std', 'valueM_mean'])  # Pivot
    # print("\n-- df3 (length {}):\n{}".format(len(df3.index), df3.head(30)))
    df3.columns = ['_'.join(col).strip() for col in df3.columns.values]     # flatten the column hierarchy
    # print("\n-- df3 (length {}):\n{}".format(len(df3.index), df3.head(30)))
    # df3.drop(['valueM_mean_H', 'valueH_std_H'], axis=1, inplace=True)
    # print("\n-- df3 (length {}):\n{}".format(len(df3.index), df3.head(30)))
    df3.reset_index(level=0, inplace=True)                                  # get l into columns
    # print("\n-- df3 (length {}):\n{}".format(len(df3.index), df3.head(30)))

    #%% -- Setup figure
    mpl.rcParams['backend'] = 'pdf'
    mpl.rcParams['lines.linewidth'] = 3
    mpl.rcParams['font.size'] = 16
    mpl.rcParams['axes.labelsize'] = 20
    mpl.rcParams['axes.titlesize'] = 16
    mpl.rcParams['xtick.labelsize'] = 16
    mpl.rcParams['ytick.labelsize'] = 16
    mpl.rcParams['legend.fontsize'] = 20
    mpl.rcParams['axes.edgecolor'] = '111111'   # axes edge color
    mpl.rcParams['grid.color'] = '777777'   # grid color
    mpl.rcParams['figure.figsize'] = [4, 4]
    mpl.rcParams['xtick.major.pad'] = 4     # padding of tick labels: default = 4
    mpl.rcParams['ytick.major.pad'] = 4         # padding of tick labels: default = 4

    fig = plt.figure()
    ax = fig.add_axes([0.13, 0.17, 0.8, 0.8])

    #%% -- Extract values into columns (plotting dataframew with bars plus error lines and lines gave troubles)
    l_vec = df3['l'].values                   # .tolist() does not work with bar plot
    mean_H_vec = df3['valueH_mean_H'].values
    mean_Hrow_vec = df3['valueH_mean_Hrow'].values
    mean_Hrow_vecEC = df3['valueH_mean_HrowEC'].values
    std_Hrow_vec = df3['valueH_std_Hrow'].values
    std_Hrow_vecEC = df3['valueH_std_HrowEC'].values

    #%% -- Draw the plot and annotate
    width = 0.3       # the width of the bars
    if SHOW_BACKTRACK_ESTIMATE:
        left_vec = l_vec
        if SHOW_NONBACKTRACK_ESTIMATE:
            left_vec = left_vec-width
        bar1 = ax.bar(left_vec, mean_Hrow_vec, width, color=plot_colors[1],
                      yerr=std_Hrow_vec, error_kw={'ecolor':'black', 'linewidth':2},    # error-bars colour
                      label=label_vec[1])
    if SHOW_NONBACKTRACK_ESTIMATE:
        bar2 = ax.bar(l_vec, mean_Hrow_vecEC, width, color=plot_colors[2],
                      yerr=std_Hrow_vecEC, error_kw={'ecolor':'black', 'linewidth':2},  # error-bars colour
                      label=label_vec[2])
    gt = ax.plot(l_vec, mean_H_vec, color=plot_colors[0], linestyle ='solid', linewidth = 2,
                 marker='o', markersize=10, markeredgewidth=2, markerfacecolor='None', markeredgecolor=plot_colors[0],
                 label=label_vec[0])

    if CHOICE == 4 or CHOICE == 20:
        ax.annotate(np.round(mean_Hrow_vec[1], 2), xy=(2.15, 0.65), xytext=(2.1, 0.60),
                    arrowprops=dict(facecolor='black', arrowstyle="->"),)

    #%% -- Legend
    if distribution == 'uniform' and SHOW_DISTRIBUTION_IN_TITLE:
        distribution_label = ',$uniform'
    else:
        distribution_label = '$'
    if SHOW_TITLE:
        plt.title(r'$\!\!\!\!n\!=\!{}\mathrm{{k}}, d\!=\!{}, h\!=\!{}, f\!=\!{}{}'.format(int(n / 1000), 2*d, h, f, distribution_label))    # notice that actual d is double than in one direction

    handles, labels = ax.get_legend_handles_labels()
    legend = plt.legend(handles, labels,
                        loc='upper right',
                        handlelength=1.5,
                        labelspacing=0,             # distance between label entries
                        handletextpad=0.3,          # distance between label and the line representation
                        # title='Iterations'
                        borderaxespad=0.1,        # distance between legend and the outer axes
                        borderpad=0.1,                # padding inside legend box
                        numpoints=1,    # put the marker only once
                        )

    if LEGEND_MATCH_COLORS:     # TODO: how to get back the nicer line spacing defined in legend above after changing the legend text colors
       legend.get_texts()[0].set_color(plot_colors[0])
       if SHOW_BACKTRACK_ESTIMATE:
           legend.get_texts()[1].set_color(plot_colors[1])
       if SHOW_NONBACKTRACK_ESTIMATE:
           legend.get_texts()[2].set_color(plot_colors[2])

    frame=legend.get_frame()
    frame.set_linewidth(0.0)
    frame.set_alpha(0.8)        # 0.8

    # %% -- Figure settings & plot
    ax.set_xticks(range(10))
    plt.grid(b=True, which='both', alpha=0.2, linestyle='solid', axis='y', linewidth=0.5)  # linestyle='dashed', which='minor'
    plt.xlabel(r'Path length ($\ell$)', labelpad=0)
    plt.ylim(ymin,ymax)                 # placed after yticks
    plt.xlim(0.5,5.5)
    plt.tick_params(
        axis='x',          # changes apply to the x-axis
        which='both',      # both major and minor ticks are affected
        bottom='off',      # ticks along the bottom edge are off        TODO: Paul, this does not work anymore :(    1/26/2020
        top='off',         # ticks along the top edge are off
        # labelbottom='off',    # labels along the bottom edge are off
    )

    if CREATE_PDF:
        plt.savefig(join(figure_directory, fig_filename), format='pdf',
            dpi=None,
            edgecolor='w',
            orientation='portrait',
            transparent=False,
            bbox_inches='tight',
            pad_inches=0.05,
            # frameon=None
            )
    if SHOW_PDF:
        showfig(join(figure_directory, fig_filename))
    if SHOW_PLOT:
        plt.show()

if __name__ == "__main__":
    run(choice=31, variant=0, show_plot=False, create_pdf=True, show_pdf=True, create_data=False)

max_stars_repo_name: northeastern-datalab/factorized-graphs

max_issues_repo_name: northeastern-datalab/factorized-graphs

max_forks_repo_name: northeastern-datalab/factorized-graphs


text:
import argparse
import sys
import time
import math
import numpy as np
import datetime
from numpy.random.mtrand import RandomState
from gym import wrappers, logger
from gym_DC.simulator.DataParameters import Data_Parameters
#from gym_DC.DCgym_Env import DCGymEnv
import matplotlib.pyplot as plt
import numpy as np
from gym_DC.simulator.Simulator import Simulator

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--env', default = 'DCgym-v0')
    parser.add_argument('--num_packets', type = int, default = 10000)
    parser.add_argument('--num_sp', type = int, default = 10)
    parser.add_argument('--num_storage', type = int, default = 1)
    parser.add_argument('--num_days', type = int, default = 100)
    parser.add_argument('--num_departures_sp', type = int, default = 80)
    parser.add_argument('--num_destination', type = int, default = 6)
    parser.add_argument('--Budget', type = int, default = 2000)
    parser.add_argument('--Penalty', type = int, default = 0)
    parser.add_argument('--num_actions', type=int, default=10)
    parser.add_argument('--HP', type = int, default = 3 )
    parser.add_argument('--start_date', type = datetime.date, default = datetime.date(2020, 1, 1))
    parser.add_argument('--mean_packet_weight', type = int, default = 8)
    parser.add_argument('--stddev_packet_weight', type = int, default = 2)
    parser.add_argument('--mean_packet_budget', type = int, default= 8)
    parser.add_argument('--stddev_packet_budget', type = int, default = 3)
    parser.add_argument('--mean_packet_penalty_late', type = int, default = 6)
    parser.add_argument('--stddev_packet_penalty_late', type = int, default = 2)
    parser.add_argument('--mean_packet_penalty_unmatched', type = int, default = 2)
    parser.add_argument('--stddev_packet_penalty_unmatched', type = int, default = 1)
    parser.add_argument('--mean_service_provider_price_per_kg', type = int, default = 10)
    parser.add_argument('--stddev_service_provider_price_per_kg', type = int, default = 2)
    parser.add_argument('--loss_service_provider_max', type = float, default = 0.6)
    parser.add_argument('--loss_service_provider_min', type = float, default = 0.1)
    parser.add_argument('--mean_service_provider_capacity', type = int, default = 200)
    parser.add_argument('--stddev_service_provider_capacity', type = int, default = 15)
    parser.add_argument('--mean_storage_price_per_kg', type = int, default = 1)
    parser.add_argument('--stddev_storage_price_per_kg', type = int, default = 1)
    parser.add_argument('--mean_storage_capacity', type = int, default = 5000)
    parser.add_argument('--stddev_storage_capacity', type = int, default = 50)
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--epsilon', type=float, default = 0.4)

    args = parser.parse_args()
    env = args.env
    num_packets = args.num_packets
    num_sp = args.num_sp
    num_storage = args.num_storage
    num_days = args.num_days
    num_departures_sp = args.num_departures_sp
    num_destination = args.num_destination
    Budget = args.Budget
    Penalty = args.Penalty
    num_actions = args.num_actions
    HP = args.HP
    Data_Params = Data_Parameters(args.start_date, args.mean_packet_weight, args.stddev_packet_weight, args.mean_packet_budget, args.stddev_packet_budget, args.mean_packet_penalty_late, args.stddev_packet_penalty_late, args.mean_packet_penalty_unmatched, args.stddev_packet_penalty_unmatched, args.mean_service_provider_price_per_kg, args.stddev_service_provider_price_per_kg, args.loss_service_provider_max, args.loss_service_provider_min, args.mean_service_provider_capacity, args.stddev_service_provider_capacity, args.mean_storage_price_per_kg, args.stddev_storage_price_per_kg, args.mean_storage_capacity, args.stddev_storage_capacity)

    data_x_axis = []
    data_y_axis = []

    HP=0
    for index in range (100):
        simulator = Simulator(num_packets, num_sp, num_storage, num_days, num_departures_sp, num_destination, Budget, Penalty, HP, Data_Params)
        profit, penalty, rate_of_successful, rate_of_delayed, rate_of_unmatched = simulator.Simulation_Period()
        reward = 0
        if profit > 0:
            reward = int(rate_of_successful * profit / 100)
        state = (reward, rate_of_successful, rate_of_delayed, rate_of_unmatched)

max_stars_repo_name: Saeid-Rezaei-projects/Gym-DC

max_issues_repo_name: Saeid-Rezaei-projects/Gym-DC

max_forks_repo_name: Saeid-Rezaei-projects/Gym-DC


text:
from .initial_condition import InitialCondition
from ..mesh import Mesh
import numpy as np
from ..eos.stiffened_gas import StiffenedGas

class ShockTubeInitialCondition(InitialCondition):
    def __init__(self, mesh: Mesh, rho, vx, p):
        super().__init__()
        self.rho = rho
        self.vx = vx
        self.p = p
        self.rhoU = rho * vx
        self.E = StiffenedGas().get_energy(rho, self.rhoU, p)
        self.rhoE = rho * self.E

    def __str__(self):
        return f"Initial Conditon: \n" \
               f"- Primitive Variables: \n" \
               f"\tDensity:\t{self.rho} \n" \
               f"\tVelocityX:\t{self.vx} \n" \
               f"\tPressure:\t{self.p} \n" \
               f"- Conservative Variables: \n" \
               f"\tDensity:\t{self.rho} \n" \
               f"\tMomentumX:\t{self.rhoU} \n" \
               f"\tTotal E:\t{self.rhoE} \n"

    def set_rho(self, density):
        self.rho = density

    def set_vx(self, velocityX):
        self.vx = velocityX

    def set_p(self, pressure):
        self.p = pressure

    def set_initial_condition(self, rho, vx, p):
        self.rho = rho
        self.vx = vx
        self.p = p

max_stars_repo_name: fengyiqi/riemannsolver

max_issues_repo_name: fengyiqi/riemannsolver

max_forks_repo_name: fengyiqi/riemannsolver


text:
import matplotlib.pyplot as plt
import numpy as np
from typing import *

class Plotter:
    nx: np.ndarray
    fx: np.ndarray
    px: np.ndarray
    px1: np.ndarray

    def __init__(self, nx: np.ndarray, fx: np.ndarray, px1: np.ndarray, px: np.ndarray):
        self.nx = nx
        self.fx = fx
        self.px = px
        self.px1 = px1

    def plot_fx_px1_histograms(self, bins=10):
        fig, ax = plt.subplots(1, 2, figsize=(11, 4))
        # fX histogram
        fx_c1 = self.fx[self.fx[:, -1] == 0][:, :-1].flatten()
        fx_c2 = self.fx[self.fx[:, -1] == 1][:, :-1].flatten()
        ax[0].hist([fx_c1, fx_c2], stacked=True, color=["tab:blue", "tab:orange"], bins=bins)
        ax[0].set_title("fX Histogram")
        ax[0].set_xlabel("fX")
        ax[0].set_ylabel("Count")
        # ax[0].margins(0.1)  # 1% padding in all directions
        ax[0].legend({"Class 1": "tab:blue", "Class 2": "tab:orange"})
        ax[0].grid(True)
        # fX histogram
        px1_c1 = self.px1[self.px1[:, -1] == 0][:, :-1].flatten()
        px1_c2 = self.px1[self.px1[:, -1] == 1][:, :-1].flatten()
        ax[1].hist([px1_c1, px1_c2], stacked=True, color=["tab:blue", "tab:orange"], bins=bins)
        # Annotate Plot
        ax[1].set_title("pX1 Histogram")
        ax[1].set_xlabel("pX1")
        ax[1].set_ylabel("Count")
        # ax[1].margins(0.1)  # 1% padding in all directions
        ax[1].legend({"Class 1": "tab:blue", "Class 2": "tab:orange"})
        ax[1].grid(True)
        # Fig config
        fig.tight_layout()

    @staticmethod
    def plot_roc(confusion_matrix_data: List[Dict]):
        cm_data_sorted = sorted(confusion_matrix_data, key=lambda row: row['fpr'])
        x = [cm_row['fpr'] for cm_row in cm_data_sorted]
        y = [cm_row['tpr'] for cm_row in cm_data_sorted]
        point_labels = [cm_row['priors'] for cm_row in cm_data_sorted]
        fig, ax = plt.subplots(1, 1, figsize=(11, 11))
        ax.plot(x, y, '-', color='tab:orange')
        previous_point = 0
        y_step = 0.3
        x_step = 0.01
        for px, py, pl in zip(x, y, point_labels):
            # Don't annotate everything
            if abs(previous_point - px) > 0.1:
                previous_point = px
                pxl = px - x_step
                pyl = py - y_step
                y_step -= 0.032
                x_step += 0.008
                pl_round = f'Priors: ({pl[0]:.2f}, {pl[1]:.2f})'
                ax.annotate(pl_round, xy=(px, py), xytext=(pxl, pyl),
                            bbox=dict(boxstyle="round", fc="none", ec="gray"),
                            arrowprops=dict(facecolor='black', arrowstyle="fancy",
                                            fc="0.6", ec="none",
                                            connectionstyle="angle3,angleA=0,angleB=-90"))
        # Annotate Plot
        ax.set_title('ROC for pX using Case 3 for different priors')
        ax.set_xlabel('FPR')
        ax.set_ylabel('TPR')
        ax.set_xticks(np.arange(0, 1.1, 0.1))
        ax.set_yticks(np.arange(0, 1.1, 0.1))
        ax.set_xlim([0, 1])
        ax.set_ylim([0, 1])
        ax.set_xbound(0, 1)
        ax.set_ybound(0, 1)
        ax.grid(True)
        ax.set_aspect('equal')
        # Fig Config
        fig.tight_layout()


text:
import math
import copy
import os

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

from matplotlib.collections import LineCollection

VIDEO_PATH = '../videos/1.mp4'
OUT_VIDEO_PATH = '../out_videos/result_1.avi'
NUMBER_LINES_CUTOFF = 5
AXIS_SIZE = 3
LENGTH_PERCENTAGE_THRESH = 70

def find_median_axis(frame):
	indices = np.argwhere(edges > 0)

	averaged_median_axis = copy.deepcopy(edges)

	variance_nan = [np.var(indices[indices[:, 0] == i][:,1]) for i in range(frame.shape[0])]
	variance = [ variance_nan[i] if not math.isnan(variance_nan[i]) else 0 for i in range(len(variance_nan))]

	for i in range(edges.shape[0]):
		if len(indices) > 0 and variance[i]<2000 and variance[i]!=0:
			averaged_median_axis[i, :] = 0
			pixel = np.average(indices[indices[:, 0] == i], axis=0)[1]
			if not math.isnan(pixel):
				pixel = np.int16(pixel).item()
				for x in range(max(0, pixel-AXIS_SIZE), min(np.int16(averaged_median_axis.shape[1]).item(), pixel+AXIS_SIZE)):
					averaged_median_axis[i][x] = 255
		else:
			averaged_median_axis[i, :] = 0

	return averaged_median_axis

# Get lines from Rho and Theta returned by Hough Transform using the bound y2
def get_line(lines,y2):
	avg = [0, 0, 0, 0]

	num = min(len(lines), NUMBER_LINES_CUTOFF)

	# Averaging over top NUMBER_LINES_CUTOFF most voted lines
	for rho, theta in lines[0:min(len(lines), NUMBER_LINES_CUTOFF)]:
		a = np.cos(theta)
		b = np.sin(theta)
		x0 = a*rho
		y0 = b*rho

		x1 = x0 + (b*y0/a)
		x2 = x1 - (b*y2/a)
		y1 = 0

		avg[0] += x1
		avg[1] += y1
		avg[2] += x2
		avg[3] += y2

	return (int(avg[0]/num), int(avg[1]/num)), (int(avg[2]/num), int(avg[3]/num))

# Finding out the length of the line to be a median axis and not overflow
def approximate_split(frame):
	frame_copy = copy.deepcopy(frame)
	indices = np.argwhere(frame_copy>0)

	variance_nan = [np.var(indices[indices[:, 0] == i][:,1]) for i in range(frame.shape[0])]
	variance = [ variance_nan[i] if not math.isnan(variance_nan[i]) else 0 for i in range(len(variance_nan))]

	sum_ = np.sum(frame_copy>0,axis=1)
	cum_sum = np.cumsum(sum_)
	percentages = 100*cum_sum/cum_sum[-1]

	for i in range(percentages.shape[0]):
		if(percentages[i]>LENGTH_PERCENTAGE_THRESH):
			return i

	return 0

backSub = cv.createBackgroundSubtractorKNN()
backSub1 = cv.createBackgroundSubtractorMOG2()

vids = os.listdir("../Testcase1/")
for video in vids:
	VIDEO_PATH = '../Testcase1/'+video
	OUT_VIDEO_PATH = '../out_videos/result_'+ video.split('.')[0] +'.avi'
	capture = cv.VideoCapture(cv.samples.findFileOrKeep(VIDEO_PATH))
	fourcc = cv.VideoWriter_fourcc(*'XVID')
	out = cv.VideoWriter(OUT_VIDEO_PATH,fourcc, 30, (int(capture.get(3)),int(capture.get(4))))

	if not capture.isOpened:
		print('ERROR: Unable to open video path ' + VIDEO_PATH)
		exit(0)

	# cv.namedWindow("AveragedAxis", cv.WINDOW_NORMAL)
	cv.namedWindow("Final", cv.WINDOW_NORMAL)
	# cv.namedWindow("Edges", cv.WINDOW_NORMAL)
	x_axis_intersection = []
	angle = []
	length = []

	count = 0
	# prev_line = None

	# Until frames keep coming
	while True:
		_, frame = capture.read()
		if frame is None:
			break

		fgMask = backSub.apply(frame)
		fgMask1 = backSub1.apply(frame)

		combined_fgMask = cv.bitwise_and(fgMask,fgMask1)

		edges = cv.Canny(combined_fgMask, 50, 100, apertureSize=3)
		edges_RGB = cv.cvtColor(edges,cv.COLOR_GRAY2RGB)

		denoised_edges = edges

		middle_axis = find_median_axis(denoised_edges)

		y_split = approximate_split(middle_axis)

		lines = cv.HoughLines(middle_axis, 1, np.pi/180, 40)

		if lines is not None:
			pt1, pt2 = get_line(lines[0],y_split)
			for r,theta in lines[0]:
				a = np.cos(theta)
				b = np.sin(theta)
				x0 = a*r
				y0 = b*r
				x1 = int(x0 + 1000*(-b))
				y1 = int(y0 + 1000*(a))
				x2 = int(x0 - 1000*(-b))
				y2 = int(y0 - 1000*(a))
				cv.line(edges_RGB,(x1,y1), (x2,y2), (0,255,0),2)
			cv.line(edges_RGB, pt1, pt2, (0, 0, 255), 6)
			# prev_line = (pt1, pt2)
			# cv.line(frame, pt1, pt2, (0, 0, 255), 6)
		# else:
		# 	if prev_line is not None:
		   #      cv.line(frame, prev_line[0], prev_line[1], (0, 0, 255), 6)

		# out.write(frame)
		# cv.imshow('Final', frame)
		# cv.imshow('AveragedAxis', middle_axis)
		cv.imshow('Final', edges_RGB)
		out.write(edges_RGB)

		keyboard = cv.waitKey(30)
		if keyboard == 'q' or keyboard == 27:
			print('LOG: Exit from keyboard')
			break
		count += 1

	out.release()
	capture.release()
	cv.destroyAllWindows()


text:
import os
import sys
import numpy as np
from six.moves import cPickle as pickle

def load_mnist_dataset():
    if sys.version_info.major == 3:
        dataset = pickle.load(open('dataset/mnist.pkl', 'rb'), encoding="bytes")
    else:
        dataset = pickle.load(open('dataset/mnist.pkl', 'rb'))
    train_set_x = np.concatenate((dataset[0][0], dataset[1][0]), axis=0).astype("float32")
    train_set_y = np.concatenate((dataset[0][1], dataset[1][1]), axis=0)
    return (train_set_x, train_set_y), (dataset[2][0], dataset[2][1])

def load_mnist_for_semi_sup(n_l=1000, n_v=1000, keep=False):
    dataset = load_mnist_dataset()

    _train_set_x, _train_set_y = dataset[0]
    test_set_x, test_set_y = dataset[1]
    test_set_x = test_set_x.reshape(test_set_x.shape[0], 1, 28, 28)

    rand_ind = np.random.permutation(_train_set_x.shape[0])
    _train_set_x = _train_set_x[rand_ind]
    _train_set_y = _train_set_y[rand_ind]

    s_c = int(n_l / 10.0)
    train_set_x = np.zeros((n_l, 28 ** 2))
    train_set_y = np.zeros(n_l)
    for i in range(10):
        ind = np.where(_train_set_y == i)[0]
        train_set_x[i * s_c:(i + 1) * s_c, :] = _train_set_x[ind[0:s_c], :]
        train_set_y[i * s_c:(i + 1) * s_c] = _train_set_y[ind[0:s_c]]
        if keep is False:
            _train_set_x = np.delete(_train_set_x, ind[0:s_c], 0)
            _train_set_y = np.delete(_train_set_y, ind[0:s_c])

    l_rand_ind = np.random.permutation(train_set_x.shape[0])
    train_set_x = train_set_x[l_rand_ind].astype("float32").reshape(l_rand_ind.shape[0], 1, 28, 28)
    train_set_y = train_set_y[l_rand_ind]
    valid_set_x = _train_set_x[:n_v]
    valid_set_x = valid_set_x.reshape(n_v, 1, 28, 28)
    valid_set_y = _train_set_y[:n_v]
    train_set_ul_x = _train_set_x[n_v:]
    train_set_ul_x = train_set_ul_x.reshape(train_set_ul_x.shape[0], 1, 28, 28)
    train_set_ul_y = _train_set_y[n_v:]
    # train_set_ul_x = np.concatenate((train_set_x, _train_set_x[n_v:]), axis=0)
    # train_set_ul_x = train_set_ul_x[np.random.permutation(train_set_ul_x.shape[0])]
    # train_set_ul_y = train_set_ul_y[np.random.permutation(train_set_ul_x.shape[0])]

    return (train_set_x, train_set_y), (train_set_ul_x, train_set_ul_y), (valid_set_x, valid_set_y), (test_set_x, test_set_y)

def load_npz_as_dict(path):
    data = np.load(path)
    return {key: data[key] for key in data}

def augmentation(images, random_crop=True, random_flip=True):
    # random crop and random flip
    h, w = images.shape[2], images.shape[3]
    pad_size = 2
    aug_images = []
    padded_images = np.pad(images, ((0, 0), (0, 0), (pad_size, pad_size), (pad_size, pad_size)), 'reflect')
    for image in padded_images:
        if random_flip:
            image = image[:, :, ::-1] if np.random.uniform() > 0.5 else image
        if random_crop:
            offset_h = np.random.randint(0, 2 * pad_size)
            offset_w = np.random.randint(0, 2 * pad_size)
            image = image[:, offset_h:offset_h + h, offset_w:offset_w + w]
        else:
            image = image[:, pad_size:pad_size + h, pad_size:pad_size + w]
        aug_images.append(image)
    ret = np.stack(aug_images)
    assert ret.shape == images.shape
    return ret

class Data:
    def __init__(self, data, label):
        # Store in numpy.array, CPU load < 100%, less GPU memory usage, takes longer time to delivery to GPU
        # Store in torch.Tensor in CPU, CPU load > 100%, similar GPU memory usage and time consuming as numpy
        # Store in torch.cuda.Tensor in GPU, GPU load < 100%, more GPU memory usage, less running time
        self.data = data
        self.label = label

    @property
    def size(self):
        return len(self.data)

    def get(self, n=None, shuffle=True, aug_trans=False, aug_flip=False, get_idx=False):
        if shuffle:
            ind = np.random.permutation(self.data.shape[0])
        else:
            ind = np.arange(self.data.shape[0])
        if n is None:
            n = self.data.shape[0]
        index = ind[:n]
        batch_data = self.data[index]
        batch_label = self.label[index]
        if aug_trans or aug_flip:
            assert batch_data.ndim == 4
            # shape of `image' [N, K, W, H]
            batch_data = augmentation(batch_data, aug_trans, aug_flip)
        if get_idx:
            return batch_data, batch_label, index
        return batch_data, batch_label

def load_dataset(dir_path, size=0, valid_size=0, valid=False, dataset_seed=1, keep=False):
    if dataset_seed > 5:
        # 1-5 -> [1, 5],  >5 -> [1, 5]
        dataset_seed = dataset_seed % 5 + 1

    if 'mnist' in dir_path:
        train_l, train_ul, val_set, test_set = load_mnist_for_semi_sup(n_l=size, n_v=valid_size, keep=keep)
        train_l = {'images': train_l[0], 'labels': train_l[1]}
        train_ul = {'images': train_ul[0], 'labels': train_ul[1]}
        if valid:
            test = {'images': test_set[0], 'labels': test_set[1]}
        else:
            test = {'images': val_set[0], 'labels': val_set[1]}
        c, h, w = 1, 28, 28
    else:
        if valid:
            train_l = load_npz_as_dict(os.path.join(dir_path, 'seed' + str(dataset_seed), 'labeled_train_valid.npz'))
            train_ul = load_npz_as_dict(os.path.join(dir_path, 'seed' + str(dataset_seed), 'unlabeled_train_valid.npz'))
            test = load_npz_as_dict(os.path.join(dir_path, 'seed' + str(dataset_seed), 'test_valid.npz'))
        else:
            train_l = load_npz_as_dict(os.path.join(dir_path, 'seed' + str(dataset_seed), 'labeled_train.npz'))
            train_ul = load_npz_as_dict(os.path.join(dir_path, 'seed' + str(dataset_seed), 'unlabeled_train.npz'))
            test = load_npz_as_dict(os.path.join(dir_path, 'seed' + str(dataset_seed), 'test.npz'))
        c, h, w = 3, 32, 32
    train_l['images'] = train_l['images'].reshape(train_l['images'].shape[0], c, h, w).astype(np.float32)
    train_ul['images'] = train_ul['images'].reshape(train_ul['images'].shape[0], c, h, w).astype(np.float32)
    test['images'] = test['images'].reshape(test['images'].shape[0], c, h, w).astype(np.float32)
    train_l['labels'] = train_l['labels'].astype(np.int32)
    train_ul['labels'] = train_ul['labels'].astype(np.int32)
    test['labels'] = test['labels'].astype(np.int32)
    return Data(train_l['images'], train_l['labels']), Data(train_ul['images'], train_ul['labels']), Data(test['images'], test['labels'])


text:
import Mask_RCNN.utils as utils
import Mask_RCNN.model as modellib
import Mask_RCNN.functions as f

from Mask_RCNN.bowl_config import bowl_config
from Mask_RCNN.bowl_dataset import BowlDataset
from Mask_RCNN.model import log
from Mask_RCNN.inference_config import inference_config
from Mask_RCNN.utils import rle_encode, rle_decode, rle_to_string

from glob import glob

import pandas as pd
import cv2
import numpy as np
from tqdm import tqdm

import os
import time

ROOT_DIR = os.getcwd()

MODEL_DIR = os.path.join(ROOT_DIR, "model/checkpoint")
CHECKPOINT_PATH = os.path.join(MODEL_DIR, "RCNN_checkpoint.h5")
SUBMISSION_FILE = "./stage2_sample_submission_final.csv"
FILE_PREFIX = "./sub-dsbowl2018-RCNN_"
TEST_DIR = "./test_2/"
TRAIN_DIR = "./test_2/"
init_with = "checkpoint"

gen_predictions = True
train_model = False

model = modellib.MaskRCNN(mode="inference",
                          config=inference_config,
                          model_dir=MODEL_DIR)

# Test dataset
dataset_test = BowlDataset()
dataset_test.load_bowl(TEST_DIR)
dataset_test.prepare()

if init_with == "checkpoint":
    model.load_weights(CHECKPOINT_PATH, by_name=True)
else:
    # Recreate the model in inference mode
    model = modellib.MaskRCNN(mode="inference",
                              config=inference_config,
                              model_dir=MODEL_DIR)
    model_path = model.find_last()[1]
    # Load trained weights (fill in path to trained weights here)
    assert model_path != "", "Provide path to trained weights"
    print("Loading weights from ",
          model_path)
    model.load_weights(model_path,
                       by_name=True)

def generate_RCNN_prediction(model = None,
                             SUBMISSION_FILE = None,
                             TEST_DIR = None,
                             OUTPUT_FILE = None):

    output = []
    sample_submission = pd.read_csv(SUBMISSION_FILE)

    ImageId = []
    EncodedPixels = []

    for image_id in tqdm(sample_submission.ImageId):
        image_path = os.path.join(TEST_DIR,
                                  image_id, "images",
                                  image_id + ".png")

        original_image = cv2.imread(image_path)
        results = model.detect([original_image],
                               verbose=0)
        r = results[0]

        masks = r["masks"]

        ImageId_batch, EncodedPixels_batch = f.numpy2encoding_no_overlap2(masks,
                                                                          image_id,
                                                                          r["scores"])
        ImageId += ImageId_batch
        EncodedPixels += EncodedPixels_batch

    f.write2csv(OUTPUT_FILE,
                ImageId,
                EncodedPixels)

def generate_timestamp():
    timestring = time.strftime("%Y_%m_%d-%H_%M_%S")
    print ("Time stamp generated: " + timestring)
    return timestring

timestr = generate_timestamp()

generate_RCNN_prediction(model = model,
                         SUBMISSION_FILE = SUBMISSION_FILE,
                         TEST_DIR = TEST_DIR,
                         OUTPUT_FILE = os.path.join(FILE_PREFIX + timestr + ".csv"))

max_stars_repo_name: rahulremanan/HIMA

max_issues_repo_name: rahulremanan/HIMA

max_forks_repo_name: rahulremanan/HIMA


text:
#!/usr/bin/env python
# Written by Kristoffer Rakstad Solberg, Student
# Copyright (c) 2020 Manta AUV, Vortex NTNU.
# All rights reserved.

import os
import rospy
import numpy as np
import yaml
from nav_msgs.msg import Odometry
from geometry_msgs.msg import Pose
from tf.transformations import quaternion_from_euler

        #   "atan2"
        #            y 90
        #     180    |     0         ^
        #        <---*----> x   rot+ |
        #    -180    |     0
        #           -90
        #

class PrepareWaypoints:

    def __init__(self, inertial_frame_id='world'):

        # Waypoints INIT
        assert inertial_frame_id in ['world', 'world_enu']
        self._inertial_frame_id = inertial_frame_id
        self.waypoints = []

    def read_from_file(self, filename):
        if not os.path.isfile(filename):
            print 'Invalid waypoint filename, file', filename
            return False
        try:
            with open(filename, 'r') as wp_file:
                wps = yaml.load(wp_file)

                # check if .yaml has a list
                if isinstance(wps['waypoints'], list):

                    # pick every waypoint an insert into array
                    for wp in wps['waypoints']:
                        x = wp['point'][0]
                        y = wp['point'][1]
                        z = wp['point'][2]
                        #yaw = wp['R']
                        #self.waypoints.append((x,y,z,yaw))
                        R = wp['RPY'][0]
                        P = wp['RPY'][1]
                        Y = wp['RPY'][2]
			self.waypoints.append((x,y,z,R,P,Y))
                    #print(self.waypoints)

        except Exception, e:
            print'Error when loading the waypoint file'
            print str(e)
            return False
        return True

    def delete_waypoint(self,waypoints):
        # pop first waypoint in list
        waypoints.pop(0)

"""
# ROS spin
if __name__ == '__main__':

    try:

    	# ROS INIT
        rospy.init_node('load_waypoints')
        rate = rospy.Rate(20) #20hz
        # get file
        if rospy.is_shutdown():
                rospy.logerr('ROS master not running!')
                sys.exit(-1)
                rospy.logerr('Found waypoint file!')
        if rospy.has_param('~filename'):

                filename = rospy.get_param('~filename')
        else:
                raise rospy.ROSException('No filename found')

        # Start tracking by reading waypoints from file
        load = PrepareWaypoints()
        load.read_from_file(filename)
        rospy.spin()
    except rospy.ROSInterruptException:
        print('caught exeption')
    print('exiting')

"""

max_stars_repo_name: theBadMusician/Vortex-AUV

max_issues_repo_name: oyssolbo/Vortex-AUV

max_forks_repo_name: oyssolbo/Vortex-AUV


text:
import json
import time

import numpy as np
import tensorflow as tf
from scipy.optimize import BFGS, basinhopping, minimize

from .fit_improve import Cached_FG
from .fit_improve import minimize as my_minimize
from .utils import time_print

class LargeNumberError(ValueError):
    pass

def fit_minuit(fcn, bounds_dict={}, hesse=True, minos=False, **kwargs):
    try:
        import iminuit
    except ImportError:
        raise RuntimeError(
            "You haven't installed iminuit so you can't use Minuit to fit."
        )
    """

    :param fcn:
    :param bounds_dict:
    :param hesse:
    :param minos:
    :return:
    """
    if int(iminuit.__version__[0]) < 2:
        return fit_minuit_v1(
            fcn, bounds_dict=bounds_dict, hesse=hesse, minos=minos, **kwargs
        )
    return fit_minuit_v2(
        fcn, bounds_dict=bounds_dict, hesse=hesse, minos=minos, **kwargs
    )

def fit_minuit_v1(fcn, bounds_dict={}, hesse=True, minos=False, **kwargs):

    """

    :param fcn:
    :param bounds_dict:
    :param hesse:
    :param minos:
    :return:
    """
    from iminuit import Minuit

    var_args = {}
    var_names = fcn.vm.trainable_vars
    for i in var_names:
        var_args[i] = fcn.vm.get(i)
        if i in bounds_dict:
            var_args["limit_{}".format(i)] = bounds_dict[i]
        var_args["error_" + i] = 0.1

    f_g = Cached_FG(fcn.nll_grad)
    m = Minuit(
        f_g.fun,
        name=var_names,
        errordef=0.5,
        grad=f_g.grad,
        print_level=2,
        use_array_call=True,
        **var_args,
    )
    print("########## begin MIGRAD")
    now = time.time()
    m.migrad()  # (ncall=10000))#,precision=5e-7))
    print("MIGRAD Time", time.time() - now)
    if hesse:
        print("########## begin HESSE")
        now = time.time()
        m.hesse()
        print("HESSE Time", time.time() - now)
    if minos:
        print("########## begin MINOS")
        now = time.time()
        m.minos()  # (var="")
        print("MINOS Time", time.time() - now)
    ndf = len(m.list_of_vary_param())
    ret = FitResult(
        dict(m.values), fcn, m.fval, ndf=ndf, success=m.migrad_ok()
    )
    ret.set_error(dict(m.errors))
    return ret

def fit_minuit_v2(fcn, bounds_dict={}, hesse=True, minos=False, **kwargs):

    """

    :param fcn:
    :param bounds_dict:
    :param hesse:
    :param minos:
    :return:
    """
    from iminuit import Minuit

    var_args = {}
    var_names = fcn.vm.trainable_vars
    x0 = []
    for i in var_names:
        x0.append(fcn.vm.get(i))
        var_args[i] = fcn.vm.get(i)
        if i in bounds_dict:
            var_args["limit_{}".format(i)] = bounds_dict[i]
        # var_args["error_" + i] = 0.1

    f_g = Cached_FG(fcn.nll_grad)
    m = Minuit(
        f_g.fun,
        np.array(x0),
        name=var_names,
        grad=f_g.grad,
    )
    m.strategy = 0
    for i in var_names:
        if i in bounds_dict:
            m.limits[i] = bounds_dict[i]
    m.errordef = 0.5
    m.print_level = 2
    print("########## begin MIGRAD")
    now = time.time()
    m.migrad()  # (ncall=10000))#,precision=5e-7))
    print("MIGRAD Time", time.time() - now)
    if hesse:
        print("########## begin HESSE")
        now = time.time()
        m.hesse()
        print("HESSE Time", time.time() - now)
    if minos:
        print("########## begin MINOS")
        now = time.time()
        m.minos()  # (var="")
        print("MINOS Time", time.time() - now)
    ndf = len(var_names)
    ret = FitResult(
        dict(zip(var_names, m.values)), fcn, m.fval, ndf=ndf, success=m.valid
    )
    # print(m.errors)
    ret.set_error(dict(zip(var_names, m.errors)))
    return ret

def fit_root_fitter(fcn):
    from array import array

    import ROOT

    var_names = fcn.vm.trainable_vars
    x0 = []
    for i in var_names:
        x0.append(fcn.vm.get(i))
    f_g = Cached_FG(fcn.nll_grad)

    class MyMultiGenFCN(ROOT.Math.IMultiGenFunction):
        def NDim(self):
            return len(x0)

        def DoEval(self, x):
            return f_g.fun(x)

        def Clone(self):
            x = MyMultiGenFCN()
            ROOT.SetOwnership(x, False)
            return x

    fitter = ROOT.Fit.Fitter()
    myMultiGenFCN = MyMultiGenFCN()
    params = array("d", x0)
    fitter.FitFCN(myMultiGenFCN, params)
    fit_result = fitter.Result()
    x = dict(zip(var_names, [fit_result.Parameter(i) for i in range(len(x0))]))
    xerr = dict(zip(var_names, [fit_result.Error(i) for i in range(len(x0))]))
    ndf = len(var_names)
    ret = FitResult(
        x, fcn, fit_result.MinFcnValue(), ndf=ndf, success=fit_result.IsValid()
    )
    ret.set_error(xerr)
    return ret

def fit_scipy(
    fcn,
    method="BFGS",
    bounds_dict={},
    check_grad=False,
    improve=False,
    maxiter=None,
):
    """

    :param fcn:
    :param method:
    :param bounds_dict:
    :param kwargs:
    :return:
    """
    args_name = fcn.vm.trainable_vars
    x0 = []
    bnds = []
    for name, i in zip(args_name, fcn.vm.trainable_variables):
        x0.append(i.numpy())
        if name in bounds_dict:
            bnds.append(bounds_dict[name])
        else:
            bnds.append((None, None))

    points = []
    nlls = []
    hess_inv = None
    now = time.time()
    if maxiter is None:
        maxiter = max(100 * len(x0), 2000)
    min_nll = 0.0
    ndf = 0
    # maxiter = 0
    def v_g2(x0):
        f_g = fcn.vm.trans_fcn_grad(fcn.nll_grad)
        nll, gs0 = f_g(x0)
        gs = []
        for i, name in enumerate(args_name):
            x0[i] += 1e-5
            nll0, _ = f_g(x0)
            x0[i] -= 2e-5
            nll1, _ = f_g(x0)
            x0[i] += 1e-5
            gs.append((nll0 - nll1) / 2e-5)
        return nll, np.array(gs)

    if check_grad:
        print("checking gradients ...")
        f_g = fcn.vm.trans_fcn_grad(fcn.nll_grad)
        nll, gs0 = f_g(x0)
        _, gs = v_g2(x0)
        for i, name in enumerate(args_name):
            print(args_name[i], gs[i], gs0[i])

    if method in ["BFGS", "CG", "Nelder-Mead", "test"]:

        def callback(x):
            if np.fabs(x).sum() > 1e7:
                x_p = dict(zip(args_name, x))
                raise LargeNumberError("x too large: {}".format(x_p))
            points.append(fcn.vm.get_all_val())
            nlls.append(float(fcn.cached_nll))
            # if len(nlls) > maxiter:
            #    with open("fit_curve.json", "w") as f:
            #        json.dump({"points": points, "nlls": nlls}, f, indent=2)
            #    pass  # raise Exception("Reached the largest iterations: {}".format(maxiter))
            print(fcn.cached_nll)

        # bd = Bounds(bnds)
        fcn.vm.set_bound(bounds_dict)

        f_g = fcn.vm.trans_fcn_grad(fcn.nll_grad)
        x0 = np.array(fcn.vm.get_all_val(True))
        # print(x0, fcn.vm.get_all_dic())
        # s = minimize(f_g, x0, method='trust-constr', jac=True, hess=BFGS(), options={'gtol': 1e-4, 'disp': True})
        if method == "test":
            try:
                s = my_minimize(
                    f_g,
                    x0,
                    method=method,
                    jac=True,
                    callback=callback,
                    options={"disp": 1, "gtol": 1e-3, "maxiter": maxiter},
                )
            except LargeNumberError:
                return except_result(fcn, x0.shape[0])
        else:
            try:
                s = minimize(
                    f_g,
                    x0,
                    method=method,
                    jac=True,
                    callback=callback,
                    options={"disp": 1, "gtol": 1e-3, "maxiter": maxiter},
                )
            except LargeNumberError:
                return except_result(fcn, x0.shape[0])

        while improve and not s.success:
            min_nll = s.fun
            maxiter -= s.nit
            s = minimize(
                f_g,
                s.x,
                method=method,
                jac=True,
                callback=callback,
                options={"disp": 1, "gtol": 1e-2, "maxiter": maxiter},
            )
            if hasattr(s, "hess_inv"):
                edm = np.dot(np.dot(s.hess_inv, s.jac), s.jac)
            else:
                break
            if edm < 1e-5:
                print("edm: ", edm)
                s.message = "Edm allowed"
                break
            if abs(s.fun - min_nll) < 1e-3:
                break
        fcn.vm.set_trans_var(s.x)  # make sure fit results same as variable
        print(s)
        # xn = s.x  # fcn.vm.get_all_val()  # bd.get_y(s.x)

        # fcn.vm.set_all(s.x, True)
        ndf = s.x.shape[0]
        min_nll = s.fun
        success = s.success
        hess_inv = fcn.vm.trans_error_matrix(s.hess_inv, s.x)
        fcn.vm.remove_bound()

        xn = fcn.vm.get_all_val()
    elif method in ["L-BFGS-B"]:

        def callback(x):
            if np.fabs(x).sum() > 1e7:
                x_p = dict(zip(args_name, x))
                raise LargeNumberError("x too large: {}".format(x_p))
            points.append([float(i) for i in x])
            nlls.append(float(fcn.cached_nll))

        try:

            s = minimize(
                fcn.nll_grad,
                np.array(x0),
                method=method,
                jac=True,
                bounds=bnds,
                callback=callback,
                options={
                    "disp": 1,
                    "maxcor": 50,
                    "ftol": 1e-15,
                    "maxiter": maxiter,
                },
            )
        except LargeNumberError:
            return except_result(fcn, len(x0))
        xn = s.x
        fcn.vm.set_var(xn)
        print(s)
        ndf = s.x.shape[0]
        min_nll = s.fun
        success = s.success
    elif method in ["Newton-CG", "trust-krylov", "trust-ncg", "trust-exact"]:
        fcn.vm.set_bound(bounds_dict)
        return fit_newton_cg(fcn, method, False)
    elif method in ["Newton-CG-p", "trust-krylov-p", "trust-ncg-p"]:
        fcn.vm.set_bound(bounds_dict)
        return fit_newton_cg(fcn, method[:-2], True)
    elif method in ["iminuit"]:
        m = fit_minuit(fcn)
        return m
    elif method in ["root"]:
        m = fit_root_fitter(fcn)
        return m
    else:
        raise Exception("unknown method")
    if check_grad:
        print("checking gradients ...")
        f_g = fcn.vm.trans_fcn_grad(fcn.nll_grad)
        _, gs0 = f_g(xn)
        gs = []
        for i, name in enumerate(args_name):
            xn[i] += 1e-5
            nll0, _ = f_g(xn)
            xn[i] -= 2e-5
            nll1, _ = f_g(xn)
            xn[i] += 1e-5
            gs.append((nll0 - nll1) / 2e-5)
            print(args_name[i], gs[i], gs0[i])
    params = fcn.get_params()  # vm.get_all_dic()
    return FitResult(
        params, fcn, min_nll, ndf=ndf, success=success, hess_inv=hess_inv
    )

def fit_newton_cg(
    fcn, method="Newton-CG", use_hessp=False, check_hess=False, gtol=1e-4
):
    vm = fcn.vm

    points = []

    def callback(x):
        points.append(fcn.vm.get_all_val())
        # if len(nlls) > maxiter:
        #    with open("fit_curve.json", "w") as f:
        #        json.dump({"points": points, "nlls": nlls}, f, indent=2)
        #    pass  # raise Exception("Reached the largest iterations: {}".format(maxiter))
        print(fcn.cached_nll)

    # bd = Bounds(bnds)

    f_g = vm.trans_fcn_grad(fcn.nll_grad)
    if use_hessp:
        hessp = vm.trans_grad_hessp(fcn.grad_hessp)
    else:
        hess = vm.trans_f_grad_hess(fcn.nll_grad_hessian)
        hess = time_print(hess)
    f_g = Cached_FG(f_g)

    x0 = np.array(vm.get_all_val(True))

    if check_hess:
        # check if hessp works well
        # hess(x0, x0)
        gs = []
        for i, _ in enumerate(x0):
            x0[i] += 1e-3
            _, g1 = f_g(x0)
            x0[i] -= 2e-3
            _, g2 = f_g(x0)
            x0[i] += 1e-3
            gs.append((g1 - g2) / 2e-3)
        gs = np.array(gs)

        # print(gs)
        if use_hessp:
            p = np.random.random(x0.shape)
            print("check", hessp(x0, p)[1], "==", np.dot(gs, p))
            p = np.random.random(x0.shape)
            print("check", hessp(x0, p)[1], "==", np.dot(gs, p))
        else:
            print(hess(x0)[2] - gs)

    if use_hessp:
        s = minimize(
            f_g, x0, jac=True, hessp=lambda x, p: hessp(x, p)[1], method=method
        )
    else:
        s = minimize(
            f_g, x0, jac=True, hess=lambda x: hess(x)[2], method=method
        )
    fcn.vm.set_trans_var(s.x)
    xn = s.x
    ndf = s.x.shape[0]
    min_nll = s.fun
    if not s.success:
        if np.min(np.abs(s.jac)) < gtol:
            s.success = True
            s.message = s.message + "\n But gradients allow"
    success = s.success
    print(s)
    # fcn.vm.set_all(xn)
    params = fcn.get_params()
    return FitResult(
        params, fcn, min_nll, ndf=ndf, success=success, hess_inv=None
    )

def except_result(fcn, ndf):
    params = fcn.vm.get_all_dic()
    return FitResult(
        params, fcn, float(fcn.cached_nll), ndf=ndf, success=False
    )

# import pymultinest
def fit_multinest(fcn):
    pass

class FitResult(object):
    def __init__(
        self, params, fcn, min_nll, ndf=0, success=True, hess_inv=None
    ):
        self.params = params
        self.error = {}
        self.fcn = fcn
        self.min_nll = float(min_nll)
        self.ndf = int(ndf)
        self.success = success
        self.hess_inv = hess_inv

    def save_as(self, file_name, save_hess=False):
        s = {
            "value": self.params,
            "error": self.error,
            "status": {
                "success": self.success,
                "NLL": self.min_nll,
                "Ndf": self.ndf,
            },
        }
        if save_hess and self.hess_inv is not None:
            s["free_params"] = [str(i) for i in self.error]
            s["hess_inv"] = [[float(j) for j in i] for i in self.hess_inv]
        with open(file_name, "w") as f:
            json.dump(s, f, indent=2)

    def set_error(self, error):
        self.error = error.copy()

max_stars_repo_name: tdayu/tf-pwa

max_issues_repo_name: tdayu/tf-pwa

max_forks_repo_name: tdayu/tf-pwa

reason: import numpy,from scipy