text:
# -*- coding: utf-8 -*-
"""
Created on Sat May 15 15:38:35 2021

@author: fionnlagh
"""

import sys
import matplotlib
#matplotlib.use('Agg')
matplotlib.use('TkAgg')  # revert above
import matplotlib.pyplot as plt
import os
import numpy as np
import glob
from pathlib import Path
import matplotlib.pyplot as plt
import pickle
import pandas as pd
import math
from scipy.interpolate import griddata
from scipy.optimize import curve_fit
sys.path.append("/home/fionnlagh/forked_amrvac/amrvac/tools/python")
import amrvac_pytools as apt

def dis_2_grid(slice_height, physical_length, resoltion):
    # translate ypts height into index value
    convert = resoltion/physical_length
    return int(round(slice_height*convert))

def grid_2_dis(physical_length, resoltion, clip_domain, positon):
    # convert index position to physical length
    # physical_length of domain - [x,y]
    # resoultion of domian - [x,y]
    # clip_domain resolution of smaller regoin to returen correct x vbalue
    # positon of interest for conveting into units - [x,y]
    physical_length = np.asarray(physical_length)
    resoltion = np.asarray(resoltion)
    positon = np.asarray(positon)
    mid_pt_x = round(clip_domain[0]/2)
    # redefine zero poitn to jet centre
    positon[0] -= mid_pt_x
    convert = physical_length/resoltion
    return positon*convert

def side_pts_of_jet_dt(clipped_data, slice_height, DOMIAN, shape):
    # gets the side points of the jet at spefified height for td plot
    # clipped_data - data set to scan
    # slice_height - phyiscal value of where to take slice in x dir
    # DOMIAN - phyiscal length of whole domain (before clipping)
    # shape - shape of whole domain before clipping
    clip_domain = np.shape(clipped_data)

    # pick up the sides of jets
    xslice_idex = dis_2_grid(slice_height, DOMIAN[1], shape[1])

    x_slice = clipped_data[:, xslice_idex]

    if np.sum(x_slice) == 0:
        jet_sides_index1 = None
        jet_sides_index2 = None

        side_values1 = None
        side_values2 = None
    else:
        indexs_x = np.nonzero(x_slice)

        jet_sides_index1 = [min(indexs_x[0]), xslice_idex]
        jet_sides_index2 = [max(indexs_x[0]), xslice_idex]

        side_values1 = grid_2_dis(DOMIAN, shape, clip_domain, jet_sides_index1)
        side_values2 = grid_2_dis(DOMIAN, shape, clip_domain, jet_sides_index2)
    return jet_sides_index1, jet_sides_index2, side_values1, side_values2

def side_pts(clipped_data, slice_height, DOMIAN, shape):
    # gets the side points of the jet at spefified height for td plot
    # clipped_data - data set to scan
    # slice_height - index value of where to take slice in x dir
    # DOMIAN - phyiscal length of whole domain (before clipping)
    # shape - shape of whole domain before clipping
    clip_domain = np.shape(clipped_data)

    x_slice = clipped_data[:, slice_height]

    if np.sum(x_slice) == 0:
        jet_sides_index1 = None
        jet_sides_index2 = None

        side_values1 = None
        side_values2 = None
    else:
        indexs_x = np.nonzero(x_slice)

        jet_sides_index1 = [min(indexs_x[0]), slice_height]
        jet_sides_index2 = [max(indexs_x[0]), slice_height]

        side_values1 = grid_2_dis(DOMIAN, shape, clip_domain, jet_sides_index1)
        side_values2 = grid_2_dis(DOMIAN, shape, clip_domain, jet_sides_index2)
    return jet_sides_index1, jet_sides_index2, side_values1, side_values2

def side_pts_of_jet(clipped_data, jet_height, nb_pts, DOMIAN, shape):
    # gets multiple pts on jet side at one instance of time
    # clipped_data - data set to scan
    # jet_height - index value of the jet height
    # DOMIAN - phyiscal length of whole domain (before clipping)
    # shape - shape of whole domain before clipping
    indexs_of_slices = np.linspace(0, jet_height[1], nb_pts, dtype=int)
    js_idx_x = []
    js_val_x = []
    js_idx_y = []
    js_val_y = []
    for idx in indexs_of_slices:
        dumvar1,dumvar2,dumvar3,dumvar4 = side_pts(clipped_data, idx, DOMIAN, shape)
        # restructiong data
        # all x values
        js_idx_x.append(dumvar1[0])
        js_idx_x.append(dumvar2[0])
        js_val_x.append(dumvar3[0])
        js_val_x.append(dumvar4[0])
        # all y data
        js_idx_y.append(dumvar1[1])
        js_idx_y.append(dumvar2[1])
        js_val_y.append(dumvar3[1])
        js_val_y.append(dumvar4[1])
    return js_idx_x, js_idx_y, js_val_x, js_idx_y

def data_slice(data, xranges, yranges, yh_index):
    # will produce slices for dt plotting
    clipped_data = data[xranges[0]:xranges[-1],
                        yranges[0]:yranges[-1]]
    slice_data = clipped_data[:,yh_index]
    return slice_data

def angle_cal(A,B):
    # A & B are (x_n,y_n)
    Y1 = max(A[1],B[1])
    Y2 = min(A[1],B[1])
    x1, x2 = A[0], B[0]
    m = A-B
    if m[1]/m[0] < 0:
        theta = -np.arccos((Y1-Y2)/(np.sqrt((x1-x2)**2+(Y1-Y2)**2)))+2*np.pi
    else:
        theta = np.arccos((Y1-Y2)/(np.sqrt((x1-x2)**2+(Y1-Y2)**2)))
    return theta

def func(m, x, c):
    return m*x + c

def LoBf(xy_data):
    # reutrns a linear line of best fit
    x = xy_data[:,0]
    y = xy_data[:,1]
    # best fit will only return single val, crashin angle func
    if min(x)-max(x) == 0:
        angle=0
    else:
        xline = np.linspace(min(x),max(x),100)
        popt, pcov = curve_fit(func, x, y)
        yline = func(xline, *popt)
        if max(yline)-min(yline) < 1e-6:
            angle = 0
        else:
            start_pt, end_pt = np.asarray((xline[0], yline[0])), \
                                           np.asarray((xline[-1], yline[-1]))
            angle = angle_cal(start_pt, end_pt)
#            print(x, y, yline, start_pt, end_pt)
#        print('angle', angle)
    return angle

def vec_angle(A,B,C):
    a = B-A
    b = C-B
    unit_vec_A = a/np.linalg.norm(a)
    unit_vec_C = b/np.linalg.norm(b)
#    print(A,B,C, a, b, np.linalg.norm(a),np.linalg.norm(b))
    dot_product = np.dot(unit_vec_A, unit_vec_C)
    theta = np.arccos(dot_product)
#    print(dot_product, np.arccos(dot_product), math.degrees(np.arccos(dot_product)))
    return theta

def distance_cal(A, B):
    # A = (x1, y1)
    # B = (x2, y2)
#    print('NUMBERS!!')
#    print(A[0]-B[0], B[1]-B[1],(A[0]-B[0])**2+(B[1]-B[1])**2)
    return np.sqrt((A[0]-B[0])**2+(A[1]-B[1])**2)

SMALL_SIZE = 12
MEDIUM_SIZE = 14
BIGGER_SIZE = 16

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

path_2_shared_drive = '/run/user/1001/gvfs/smb-share:server=uosfstore.shef.ac.uk,share=shared/mhd_jet1/User/smp16fm/j'
dir_paths =  glob.glob('../T/P300/B60/A60/T30*')
#dir_paths =  glob.glob('../T/P300/B80/A80/T5*')

#dir_paths =  glob.glob('../T/P300/B60/A60/T15*')

#dir_paths = dir_paths[-8:]
#dir_paths = [dir_paths[-7]]
#dir_paths = dir_paths[-8:-4]

#dir_paths =  glob.glob('../T/P*/B50/A*/T*')
#dir_paths =  glob.glob('../2D/P*/T*')
#dir_paths =  glob.glob('../hight_dt/P*/B*/A*')
#dir_paths = [dir_paths[-3]]

# constants
unit_length = 1e9  # cm
DOMIAN = [5*unit_length, 3*unit_length]
unit_temperature = 1e6  # K
unit_numberdensity = 1e9  # cm^-3

g_cm3_to_kg_m3 = 1e3
dyne_cm2_to_Pa = 1e-1
cm_to_km = 1e-5
cm_to_Mm = 1e-8
s_to_min = 1/60

unit_density = 2.3416704877999998E-015
unit_velocity = 11645084.295622544
unit_pressure = 0.31754922400000002
unit_magenticfield = 1.9976088799077159
unit_time = unit_length/unit_velocity

dt = unit_time/20

unit_mass = unit_density*unit_length**3
unit_specific_energy = (unit_length/unit_time)**2

# otpions
testing = True
plotting_on = False
data_save = False
# NOTE: if name already there it will append to file
max_h_data_fname = 'max_h_data_sj_p2.dat'
big_data_fname = 'big_data_set_sj_p2.dat'
td_plotting = False
td_plot_root_folder = 'td_plots_data_sj/'
td_file_name = 'data.csv'
stop_height_condition = 1e8  #5e6 # 5 cells high
stop_indx_condition = 120  # 24 # ~50s which is min driver time
thresh_hold = 20  #0.4
central_axis_tracking = True
c_data_root = 'c_data/'
central_axis_step_size = 0.1  # Mm
# abdandon methods
method_1 = False
method_2 = False
# chosen method
method_3 = True
method_4 = False
x_pad = 1/2  # Mm
y_pad = 0.75/2  # Mm
dummy_dis = 0

#t0 = 1
#t1 = 80
#nb_steps = 80

t0 = 10
t1 = 20
nb_steps = 2

xres = 4096
yres = 2944

physical_grid_size_xy = DOMIAN/np.array([xres, yres])*cm_to_Mm

peak_hi = 0

big_data = []
data = []
big_data_indexs = []
for path in dir_paths:
    FIRST = True
    data_c_first = True
    JL_data_first = True
    h_check = 0
    path_parts = path.split('/')
    path_parts = path_parts[2:]
    path_numerics = np.zeros(len(path_parts))

    for j, item in enumerate(path_parts):
#        print(item)
        path_numerics[j] = float(item[1:])

    #old style
#    full_paths = glob.glob(path+'/jet_'+path_parts[1]+'_'+path_parts[-1]+'_*.vtu')
    full_paths = glob.glob(path + '/jet_'+path_parts[0] + '_'+path_parts[1] +
                           '_' + path_parts[2] + '_' + path_parts[3] +
                           '_*.vtu')
    # skip first step as no value
#    full_paths = full_paths[1:]

    # testing breaks in code
    full_paths = full_paths[68:69]
#    full_paths = full_paths[36:40]

    sub_data_1 = []
    sub_data_2 = []
    physical_time = []

    for ind, path_element in enumerate(full_paths):
        Full_path = path_2_shared_drive + path_element[2:-8]
        ti = int(path_element[-8:-4])
        # Reading vtu file, allows to set custom grid poitns
        ds0 = apt.load_vtkfile(ti, file=Full_path, type='vtu')
        data0 = apt.vtkfiles.rgplot(ds0.trp1, data=ds0, cmap='hot')
        plt.close()
        var_tr_data, x_grid0, y_grid0 = data0.get_data(xres=xres, yres=yres)
        bin_data = np.where(var_tr_data < 15, 0, 1)
        # dims in [y,x]
        shape = np.shape(bin_data)
        # This mid point doesnt corospond to jet centre
        if FIRST == True:
            # These don't work for first time step
            indexs_x = np.nonzero(bin_data[:, 0])[0]
            mid_pt_x = int(round((min(indexs_x) + (max(indexs_x) -
                           min(indexs_x))/2)))
            clip_range_x = round(0.1*shape[0])
            scan_range_x = [mid_pt_x-clip_range_x, mid_pt_x+clip_range_x]

            mid_pt_y = round(shape[1]/2)
            clip_range_y = round(0.2*shape[1])
            scan_range_y = [0, mid_pt_y+clip_range_y]
            x_extent = np.asarray(scan_range_x)*physical_grid_size_xy[0] - \
                       2.547205e+09*cm_to_Mm
            cf = sum(x_extent)/2
            x_extent -= cf
            y_extent = np.asarray(scan_range_y)*physical_grid_size_xy[1]
            x_pad_dex_size = int(np.ceil(x_pad/physical_grid_size_xy[0]))
            y_pad_dex_size = int(np.ceil(y_pad//physical_grid_size_xy[0]))

            FIRST = False

        # clips data around jet
        sorted_data = bin_data[scan_range_x[0]:scan_range_x[-1],
                               scan_range_y[0]:scan_range_y[-1]]
        # All indexs that belong to jet bc
        indexs = np.nonzero(sorted_data)
        # index for top of jet
        jet_top_index = np.argmax(indexs[1])
        jet_top_pixel_pos = [indexs[0][jet_top_index],
                             indexs[1][jet_top_index]]
        # need to fix x postion as its zero point is not at jet centre
        values = grid_2_dis(DOMIAN, shape, np.shape(sorted_data),
                            jet_top_pixel_pos)
        # top
        height_x = values[0]*cm_to_Mm
        height_y = values[1]*cm_to_Mm
        physical_time = ti*dt
        sub_data_1.append((physical_time, height_y))

        # stops the loop for creating data
        if (height_y < stop_height_condition*cm_to_Mm and
            ind > stop_indx_condition):
            break

        # plot side values every 1 Mm interval
        for hi in range(1, int(np.floor(height_y))+1):
#            print(hi, h_check, h_check-hi, h_check-hi<0)
            if h_check-hi < 0:
                td_first = True
                h_check = hi
            if hi > peak_hi:
                peak_hi = hi

            slice_height = hi/cm_to_Mm
            jet_sides_index1, jet_sides_index2, val1, val2 = side_pts_of_jet_dt(sorted_data, slice_height, DOMIAN, shape)

            if jet_sides_index1 == None:
                pass
            else:
                dis_x = (val2[0]-val1[0])*cm_to_km
                dis_y = slice_height*cm_to_Mm

                side_xL = val1[0]*cm_to_km
                side_xR = val2[0]*cm_to_km
#                side_y.append(slice_height*cm_to_Mm)
#                side_y.append(slice_height*cm_to_Mm)
                side_time = ti*dt
                sub_data_2.append((side_time, dis_x, dis_y,
                                   side_xL, side_xR))
                # This wont work for tilted jets or any asymetries, its a quick fix
                if testing == True:
                    plt.scatter((jet_sides_index1[0]+scan_range_x[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf,
                                 jet_sides_index1[1]*physical_grid_size_xy[1],
                                 s=40, color='blue')
                    plt.scatter((jet_sides_index2[0]+scan_range_x[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf,
                                 jet_sides_index2[1]*physical_grid_size_xy[1],
                                 s=40, color='blue')

                if td_plotting == True:
                    td_save_path = td_plot_root_folder + \
                                   full_paths[ind].split('/')[-1][:-9] + \
                                   '/'+str(hi)+'Mm'
                    Path(td_save_path).mkdir(parents=True, exist_ok=True)
                    height_km = dis_y
                    height_index = jet_sides_index2[1]

                    datarho = apt.vtkfiles.rgplot(ds0.rho, data=ds0,
                                                  cmap='hot')
                    plt.close()
                    datarho, dummy_x, dummy_y = datarho.get_data(xres=xres,
                                                                 yres=yres)
                    rho_slice = data_slice(datarho, scan_range_x, scan_range_y,
                                           height_index)*g_cm3_to_kg_m3
                    datarho = []

                    dataTe = apt.vtkfiles.rgplot(ds0.T, data=ds0, cmap='hot')
                    plt.close()
                    dataTe, dummy_x, dummy_y = dataTe.get_data(xres=xres,
                                                               yres=yres)
                    Te_slice = data_slice(dataTe, scan_range_x, scan_range_y,
                                          height_index)
                    dataTe = []

                    dataVx = apt.vtkfiles.rgplot(ds0.v1, data=ds0, cmap='hot')
                    plt.close()
                    dataVx, dummy_x, dummy_y = dataVx.get_data(xres=xres,
                                                               yres=yres)
                    Vx_slice = data_slice(dataVx, scan_range_x, scan_range_y,
                                          height_index)*cm_to_km
                    dataVx = []

                    dataVy = apt.vtkfiles.rgplot(ds0.v2, data=ds0, cmap='hot')
                    dataVy, dummy_x, dummy_y = dataVy.get_data(xres=xres,
                                                               yres=yres)
                    plt.close()
                    Vy_slice = data_slice(dataVy, scan_range_x, scan_range_y,
                                          height_index)*cm_to_km
                    dataVy = []

                    td_xvales = [scan_range_x[0]*physical_grid_size_xy[0] -
                                 2.547205e+09*cm_to_Mm - cf,
                                 scan_range_x[1]*physical_grid_size_xy[0] -
                                 2.547205e+09*cm_to_Mm-cf]
                    td_xranges = np.linspace(td_xvales[0], td_xvales[1],
                                             len(Vy_slice))

                    df_td = pd.DataFrame(np.transpose(np.array([np.ones(len(Vy_slice))*ti*dt, td_xranges, rho_slice, Te_slice, Vx_slice, Vy_slice])),
                                                                columns=['time [s]', 'x [Mm]', 'density [kg m-3]', 'Te [k]', 'vx [km s-1]', 'vy [km s-1]'])
                    if td_first == True:
#                        print('writting')
                        df_td.to_csv(td_save_path+'/'+td_file_name,
                                     index=False,
                                     columns=['time [s]', 'x [Mm]',
                                              'density [kg m-3]',
                                              'Te [k]', 'vx [km s-1]',
                                              'vy [km s-1]'])
                        td_first = False
                    else:
                        df_td.to_csv(td_save_path+'/'+td_file_name, mode='a',
                                     index=False, header=None)

# test putting side tracking here
        # need to add more points to using above
        if central_axis_tracking == True:
            if height_y < 1:
                pass
            else:
                # +1 ensures endpts remain
                nb_step = int((height_y/central_axis_step_size)+1)
                hi_locs = np.linspace(0, height_y, nb_step,
                                      endpoint=True)/cm_to_Mm
                central_pts = []
                central_sides = []
                for c_pts in hi_locs:
                    cjet_sides_index1, cjet_sides_index2, cval1, cval2 = side_pts_of_jet_dt(sorted_data, c_pts, DOMIAN, shape)
                    if cjet_sides_index1 is not None:
                        central_sides.append((cjet_sides_index1,
                                              cjet_sides_index2))
                        central_pts.append(np.add(cjet_sides_index1,
                                                  cjet_sides_index2)//2)
                    else:
                        print('Cenrtal axis pt missed')
                        continue

                central_pts = np.reshape(central_pts, np.shape(central_pts))
                if testing == True:
    #                        plt.plot(central_pts[:,0], central_pts[:,1])
                    plt.scatter((central_pts[:, 0]+scan_range_x[0])*physical_grid_size_xy[0]
                                 -2.547205e+09*cm_to_Mm-cf,
                                 central_pts[:, 1]*physical_grid_size_xy[1],
                                 s=40, color='yellow', marker="*", zorder=3)

                # need to calc length of jet
                central_pts_phy = np.zeros(np.shape(central_pts))
                central_pts_phy[:, :1] = (central_pts[:, :1]+scan_range_x[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf
                central_pts_phy[:, -1:] = central_pts[:, -1:]*physical_grid_size_xy[1]
                # vectorization of dis_Cal func
                p2p_dis = np.sqrt(np.sum((central_pts_phy[:-1]-central_pts_phy[1:])**2,
                                         axis=1))
                p2p_dis_array = np.zeros([np.size(p2p_dis), 3])
                p2p_dis_array[:, 1:] = central_pts_phy[1:]
                for i in range(1, len(p2p_dis)):
                    p2p_dis_array[i, 0] = sum(p2p_dis[:i+1])
    #            jet_length = sum(p2p_dis)
                jet_length = p2p_dis_array[-1][0]
    #            print(jet_length)
                if data_save == True:
                    df_JL_data = pd.DataFrame([[jet_length, physical_time]],
                                              columns=['Jet length [Mm]',
                                                       'Time [s]'])
                    if JL_data_first:
                        data_c_save_path = c_data_root+full_paths[ind].split('/')[-1][:-9]
                        Path(data_c_save_path).mkdir(parents=True,
                                                     exist_ok=True)
                        df_JL_data.to_csv(data_c_save_path+'/'+full_paths[ind].split('/')[-1][:-9]+'_'+'df_jl.csv',
                                          index = False,
                                          columns=['Jet length [Mm]',
                                                   'Time [s]'])
                        JL_data_first = False
                    else:
                        df_JL_data.to_csv(data_c_save_path+'/'+full_paths[ind].split('/')[-1][:-9]+'_'+'df_jl.csv',
                                          mode='a', index=False, header=None)
                #-------------------------------------------
                if method_1 == True:
                    # trying method of avg angles
                    for hi_indx in range(1,len(central_pts)-1):
                        p1,p2 = angle_cal(central_pts[hi_indx-1],
                                          central_pts[hi_indx]), \
                                          angle_cal(central_pts[hi_indx],
                                                    central_pts[hi_indx+1])
                        perp_avg_tilt = np.mean([p1,p2])-np.pi/2
                        m_grad = 1/np.tan(perp_avg_tilt)
                        const = central_pts[hi_indx][1] - \
                                m_grad*central_pts[hi_indx][0]
                        x_slit = np.linspace(0, clip_range_x*2, 50)
                        line = m_grad*x_slit+const
                        if testing == True:
                            plt.plot((x_slit+scan_range_x[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf,
                                     line*physical_grid_size_xy[1], 'r-')
                # ------------------------------------------------------
                # method 2: middle angles
                if method_2 == True:
                    for hi_indx in range(1, len(central_pts)-1):
                        vec_A = central_pts[hi_indx-1]
                        vec_B = central_pts[hi_indx]
                        vec_C = central_pts[hi_indx+1]
                        vec_ang = vec_angle(vec_A, vec_B, vec_C)
                        width_angle = vec_ang/2
                        m_grad = np.tan(width_angle)
                        const = vec_B[1]-m_grad*vec_B[0]
                        x_slit = np.linspace(0, clip_range_x*2, 50)
                        line = m_grad*x_slit+const
                        if testing == True:
                            plt.plot((x_slit+scan_range_x[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf,
                                     line*physical_grid_size_xy[1], 'b--')
                # ------------------------------------------------------
                # method 3: top angles
                if method_3 == True:
                    for hi_indx in range(1, int(np.floor(height_y))+1):
                        current_x_pad_dex_size = x_pad_dex_size
                        current_y_pad_dex_size = y_pad_dex_size
                        # +1 matches it with central_pts as 1 element is lost with dis calc
                        c_index = np.argmin(abs(p2p_dis_array[:, 0]-hi_indx))+1
#                         if value fall at top of arry angle cant be calc
                        if c_index+1 >= len(p2p_dis_array):
                            pass
                        else:
                            p1 = angle_cal(central_pts[c_index], central_pts[c_index+1])
                            perp_avg_tilt = p1-np.pi/2
                            m_grad = 1/np.tan(perp_avg_tilt)
        #                   current method
                            const = central_pts[c_index][1] - \
                                    m_grad*central_pts[c_index][0]
                            z_line_switches = [0]
                            # makes sure that more than 1 edge is detected
                            while_count = 0
                            while sum(np.abs(z_line_switches)) < 2:
                                print(while_count)
                                while_count += 1
                                # defines search region
                                x_search = (central_sides[c_index][0][0] -
                                            current_x_pad_dex_size,
                                            central_sides[c_index][1][0] +
                                            current_x_pad_dex_size)
                                y_search = (central_sides[c_index][0][1] -
                                            current_y_pad_dex_size,
                                            central_sides[c_index][0][1] +
                                            current_y_pad_dex_size)
                                # grid in phy units
                                points = np.array((y_grid0[scan_range_x[0]+x_search[0]:scan_range_x[0]+x_search[1],
                                                           y_search[0]:y_search[1]].flatten(),
                                                   x_grid0[scan_range_x[0]+x_search[0]:scan_range_x[0]+x_search[1],
                                                           y_search[0]:y_search[1]].flatten())).T*cm_to_Mm
                                values = (bin_data[scan_range_x[0]+x_search[0]:scan_range_x[0]+x_search[1],
                                                      y_search[0]:y_search[1]]).flatten()

                                line_dis_phy = np.sqrt(((x_search[0]-x_search[-1])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf)**2+((y_search[0]-y_search[-1])*physical_grid_size_xy[1])**2)
                                nb_pts_for_line = int(line_dis_phy//0.05)
                                x_slit = np.linspace(x_search[0], x_search[1],
                                                     nb_pts_for_line)
                                x_slit_phy = (x_slit+scan_range_x[0])*physical_grid_size_xy[0] - \
                                             2.547205e+09*cm_to_Mm-cf

                                line = m_grad*x_slit+const
                                line_phy = line*physical_grid_size_xy[1]
                                xi = np.array(list(zip(line_phy, x_slit_phy)))

                                z_line_vale = griddata(points, values, xi)
    #                            z_line_vale = z_line_vale[~np.isnan(z_line_vale)]
                                z_line_vale[np.where(np.isnan(z_line_vale))] = 0
                                z_line_vale = np.where(z_line_vale < 1, 0, 1)
                                z_line_switches = np.diff(z_line_vale)
                                # expand search area
                                if sum(np.abs(z_line_switches)) < 2:
#                                    print('while not broken', sum(np.abs(z_line_switches)))
                                    current_x_pad_dex_size += 5
                                    current_y_pad_dex_size += 5
                                    continue
#                                print('while will be broken', sum(np.abs(z_line_switches)))
                                # make sure only 2 pts are sleceted
                                LR_edge_fix = np.argwhere(abs(z_line_switches) > 0)
                                LR_edge_fix_index = [np.min(LR_edge_fix), np.max(LR_edge_fix)]
                                spatial_locs_widths = xi[LR_edge_fix_index]
                                # old method resultes in > 2pts
    #                            z_line_side_indexs = np.argwhere(abs(z_line_switches)>0)
    #                            spatial_locs_widths = xi[z_line_side_indexs][:,0]

                                # xi = [[y1,x1],[y2,x2]]
                                # Will be give Mm
                                tilt_widths = distance_cal(spatial_locs_widths[0],
                                                           spatial_locs_widths[1])
                                data_c = np.asarray((float(path_parts[0][1:]),
                                                     float(path_parts[1][1:]),
                                                     float(path_parts[2][1:]),
                                                     float(path_parts[3][1:]),
                                                     p2p_dis_array[c_index][0],
                                                     p2p_dis_array[c_index][-1],
                                                     tilt_widths,
                                                     physical_time,
                                                     p2p_dis_array[-1][0],
                                                     p2p_dis_array[-1][-1]))
                                if data_save == True:
                                    df_dc = pd.DataFrame([data_c],
                                                         columns=['Driver time [s]',
                                                                  'Magnetic field strength [B]',
                                                                  'Amplitude [km/s]',
                                                                  'Tilt angle [degree]',
                                                                  'Jet length [Mm]',
                                                                  'Jet height [Mm]',
                                                                  'Jet width [Mm]',
                                                                  'Time [s]',
                                                                  'Max len [Mm]',
                                                                  'Max height [Mm]'])
                                    if data_c_first == True:
                #                        print('writting')
                                        data_c_save_path = c_data_root+full_paths[ind].split('/')[-1][:-9]
                                        Path(data_c_save_path).mkdir(parents=True, exist_ok=True)
                                        df_dc.to_csv(data_c_save_path+'/'+full_paths[ind].split('/')[-1][:-9]+'_'+td_file_name,
                                                     index=False, columns=['Driver time [s]',
                                                                             'Magnetic field strength [B]',
                                                                             'Amplitude [km/s]',
                                                                             'Tilt angle [degree]',
                                                                             'Jet length [Mm]',
                                                                             'Jet height [Mm]',
                                                                             'Jet width [Mm]',
                                                                             'Time [s]',
                                                                             'Max len [Mm]',
                                                                             'Max height [Mm]'])

                                        data_c_first = False
                                    else:
                                        df_dc.to_csv(data_c_save_path+'/'+full_paths[ind].split('/')[-1][:-9]+'_'+td_file_name,
                                                     mode='a', index=False, header=None)

                                if testing == True:
                                    # Physical grid checking
                                    # Issue with grid aligment due to how yt written data, most likely cause by the sterech grids.
                                    # width are correctly measure but are shift leftward due to difference in physical value for index pts of the grid and line
                                    extra_cf = (x_grid0[:,0][scan_range_x[0]+x_search[0]])*cm_to_Mm-min(x_slit_phy)
                                    plt.scatter(spatial_locs_widths[:,1:]-extra_cf,spatial_locs_widths[:,:-1], color='red', marker='s', zorder=2)
#                                    plt.plot((x_slit+scan_range_x[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf,line*physical_grid_size_xy[1], 'g-o')
                                    cmap = 'gray'
#                                    plt.imshow(np.rot90(var_tr_data[scan_range_x[0]:scan_range_x[-1], scan_range_y[0]:scan_range_y[-1]]), cmap=cmap, extent = [x_extent[0], x_extent[1], y_extent[0],y_extent[1]])
#                                    plt.plot(x_slit_phy,line_phy, 'g-', zorder=1)
                                    # test to purely size slice area
                                    plt.plot(x_slit_phy-extra_cf,line_phy, 'g-', zorder=1)
                                    plt.imshow(np.rot90(bin_data[scan_range_x[0]:scan_range_x[-1], scan_range_y[0]:scan_range_y[-1]]),
                                                        cmap=cmap, extent=[x_extent[0], x_extent[1], y_extent[0],y_extent[1]])
                                    plt.imshow(np.rot90(bin_data[scan_range_x[0]+x_search[0]:scan_range_x[0]+x_search[1], y_search[0]:y_search[1]]),
                                                        cmap='Spectral', extent=[(scan_range_x[0]+x_search[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf,
                                                                             (scan_range_x[0]+x_search[1])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf,
                                                                             y_search[0]*physical_grid_size_xy[1],y_search[1]*physical_grid_size_xy[1]])
#
                                    plt.xlim(-1,1)
                                    plt.ylim(0,2)
                                    plt.xlabel('Y [Mm]')
                                    plt.ylabel('X [Mm]')
                                    plt.savefig('sharc_run/fig_for_paper/example_of_tilt_jet_code.png', dpi=200, bbox_inches='tight')
                                    plt.show()
                if method_4 == True:
                    for hi_indx in range(1,int(np.floor(height_y))+1):
                        current_x_pad_dex_size = x_pad_dex_size
                        current_y_pad_dex_size = y_pad_dex_size
                        # +1 matches it with central_pts as 1 element is lost with dis calc
                        c_index = np.argmin(abs(p2p_dis_array[:,0]-hi_indx))+1
                        pts_of_influence = 3 # need to be moved but here for convenice
                        if (c_index + pts_of_influence >= len(p2p_dis_array)) and (c_index-pts_of_influence<0):
                            pass
                        else:
                            p1 = LoBf(central_pts[c_index-pts_of_influence:c_index+pts_of_influence])
                            perp_avg_tilt = p1-np.pi/2
                            m_grad = 1/np.tan(perp_avg_tilt)
        #                   current method
                            const = central_pts[c_index][1]-m_grad*central_pts[c_index][0]
                            z_line_switches = [0]
                            # makes sure that more than 1 edge is detected
                            while_count = 0
                            while sum(np.abs(z_line_switches)) < 2:
                                print(while_count)
                                while_count += 1
                                # defines search region
                                x_search = (central_sides[c_index][0][0]-current_x_pad_dex_size,
                                            central_sides[c_index][1][0]+current_x_pad_dex_size)
                                y_search = (central_sides[c_index][0][1]-current_y_pad_dex_size,
                                            central_sides[c_index][0][1]+current_y_pad_dex_size)
                                # grid in phy units
                                points = np.array((y_grid0[scan_range_x[0]+x_search[0]:scan_range_x[0]+x_search[1],
                                                           y_search[0]:y_search[1]].flatten(),
                                                   x_grid0[scan_range_x[0]+x_search[0]:scan_range_x[0]+x_search[1],
                                                           y_search[0]:y_search[1]].flatten())).T*cm_to_Mm
                                values = (bin_data[scan_range_x[0]+x_search[0]:scan_range_x[0]+x_search[1],
                                                      y_search[0]:y_search[1]]).flatten()

                                line_dis_phy = np.sqrt(((x_search[0]-x_search[-1])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf)**2+((y_search[0]-y_search[-1])*physical_grid_size_xy[1])**2)
                                nb_pts_for_line = int(line_dis_phy//0.05)
                                x_slit = np.linspace(x_search[0], x_search[1],
                                                     nb_pts_for_line)
                                x_slit_phy = (x_slit+scan_range_x[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf

                                line = m_grad*x_slit + const
                                line_phy = line*physical_grid_size_xy[1]
                                xi = np.array(list(zip(line_phy, x_slit_phy)))

                                z_line_vale = griddata(points, values, xi)
                                z_line_vale[np.where(np.isnan(z_line_vale))] = 0
                                z_line_vale = np.where(z_line_vale<1, 0, 1)
                                z_line_switches = np.diff(z_line_vale)
                                # expand search area
                                if sum(np.abs(z_line_switches)) < 2:
#                                    print('while not broken', sum(np.abs(z_line_switches)))
                                    current_x_pad_dex_size += 5
                                    current_y_pad_dex_size += 5
                                    continue
#                                print('while will be broken', sum(np.abs(z_line_switches)))

                                # make sure only 2 pts are sleceted
                                LR_edge_fix = np.argwhere(abs(z_line_switches) > 0)
                                LR_edge_fix_index = [np.min(LR_edge_fix), np.max(LR_edge_fix)]
                                spatial_locs_widths = xi[LR_edge_fix_index]
                                # Will be give Mm
                                tilt_widths = distance_cal(spatial_locs_widths[0],
                                                           spatial_locs_widths[1])

                                if testing == True:
                                    # Physical grid checking
                                    # Issue with grid aligment due to how yt written data, most likely cause by the sterech grids.
                                    # width are correctly measure but are shift leftward due to difference in physical value for index pts of the grid and line
                                    extra_cf = (x_grid0[:,0][scan_range_x[0]+x_search[0]])*cm_to_Mm-min(x_slit_phy)
                                    plt.scatter(spatial_locs_widths[:,1:]-extra_cf,spatial_locs_widths[:,:-1], color='pink', marker='P', zorder=2)
                                    # test to purely size slice area
                                    plt.plot(x_slit_phy-extra_cf, line_phy, 'c:', zorder=2)
        if testing == True:
            # testing
            cmap = 'gray'
            plt.scatter((jet_top_pixel_pos[0]+scan_range_x[0])*physical_grid_size_xy[0]-2.547205e+09*cm_to_Mm-cf,
                        jet_top_pixel_pos[1]*physical_grid_size_xy[1], s=40, color='red')
            # image
#            plt.imshow(sorted_data, cmap=cmap)
            plt.imshow(np.rot90(var_tr_data[scan_range_x[0]:scan_range_x[-1], scan_range_y[0]:scan_range_y[-1]]),
                       cmap=cmap, extent=[x_extent[0], x_extent[1],
                                          y_extent[0], y_extent[1]])
            plt.xlim(-1, 1)
#            plt.xlim(-1,1)
            plt.ylim(0, 8)
            plt.gca().set_aspect(0.5, adjustable='box')
            plt.xlabel('x (Mm)')
            plt.ylabel('y (Mm)')
            plt.colorbar()
            plt.show()
            #saves testing image
#            Path('image_check/'+full_paths[ind].split('/')[-1][:-9]).mkdir(parents=True, exist_ok=True)
#            plt.savefig('image_check/'+full_paths[ind].split('/')[-1][:-9]+'/jet_P'+str(int(path_numerics[0]))+'_B' +
#                        str(int(path_numerics[1])) +
#                        'A_' + str(int(path_numerics[2])) +
#                        'T_'+str(ti).zfill(4) + '.png',
#                        format='png', dpi=500)
#            plt.clf()

    # data frame to nest data in
    df_sub1 = (pd.DataFrame(sub_data_1, columns=['time [s]', 'Height [Mm]'],
                            index = [i for i in range(len(sub_data_1))]))

    df_sub2 = (pd.DataFrame(sub_data_2,
                                 columns=['side time [s]', 'jet Width [km]',
                                'height [Mm]', 'jet side left [km]',
                                'jet side right [km]'],
                                 index = [i for i in range(len(sub_data_2))]))

#    big_data.append(pd.concat([df_sub1, df_sub2], axis=1))
#    big_data_indexs.append(path_numerics.astype(int))    # first data set
#    data.append(np.hstack([path_numerics,max(sub_data_1, key=lambda x: float(x[1]))[1]]))
#    df = pd.DataFrame(data, columns=['driver time [s]',
#                                'magnetic field strength [G]',
#                                'amplitude [km s-1]',
#                                'Tilt [deg]',
#                                'max height [Mm]'],
#                                index = [i for i in range(np.shape(data)[0])])
    big_data = pd.concat([df_sub1, df_sub2], axis=1)
    big_data_indexs = path_numerics.astype(int)    # first data set
    df_collect = pd.DataFrame([{'idx': big_data_indexs, 'dfs': big_data}])

    data = np.hstack([path_numerics, max(sub_data_1, key=lambda x: float(x[1]))[1]])
    df = pd.DataFrame([data], columns=['driver time [s]',
                                       'magnetic field strength [G]',
                                       'amplitude [km s-1]',
                                       'Tilt [deg]',
                                       'max height [Mm]'])
    if data_save == True:
        if os.path.exists(max_h_data_fname):
            data_max_h_t0 = pd.read_pickle(max_h_data_fname)
#            print('I add to file ' + max_h_data_fname + '!!!!!!!!!!' )
            dummy_max_h = data_max_h_t0.append(df, ignore_index=True)
            dummy_max_h.to_pickle(max_h_data_fname)
        else:
#            print('I made file ' + max_h_data_fname + '!!!!!!!!!!')
            df.to_pickle(max_h_data_fname)
        if os.path.exists(big_data_fname):
            big_data_t0 = pd.read_pickle(big_data_fname)
#            print('I add to file ' + big_data_fname + '!!!!!!!!!!' )
            dummy_big_data = big_data_t0.append(df_collect, ignore_index=True)
            dummy_big_data.to_pickle(big_data_fname)
        else:
#            print('I made file ' + big_data_fname + '!!!!!!!!!!')
            df_collect.to_pickle(big_data_fname)


text:
"""
Trialwise Decoding on BCIC IV 2a Dataset
========================================

This tutorial shows you how to train and test deep learning models with
Braindecode in a classical EEG setting: you have trials of data with
labels (e.g., Right Hand, Left Hand, etc.).

"""

######################################################################
# Loading and preprocessing the dataset
# -------------------------------------
#

######################################################################
# Loading
# ~~~~~~~
#

######################################################################
# First, we load the data. In this tutorial, we use the functionality of
# braindecode to load datasets through
# `MOABB <https://github.com/NeuroTechX/moabb>`__ to load the BCI
# Competition IV 2a data.
#
# .. note::
#    To load your own datasets either via mne or from
#    preprocessed X/y numpy arrays, see `MNE Dataset
#    Tutorial <./plot_mne_dataset_example.html>`__ and `Numpy Dataset
#    Tutorial <./plot_custom_dataset_example.html>`__.
#

from braindecode.datasets.moabb import MOABBDataset
import mne

def exp(subject_id):
    dataset = MOABBDataset(dataset_name="BNCI2014001", subject_ids=[subject_id])

    ######################################################################
    # Preprocessing
    # ~~~~~~~~~~~~~
    #

    ######################################################################
    # Now we apply preprocessing like bandpass filtering to our dataset. You
    # can either apply functions provided by
    # `mne.Raw <https://mne.tools/stable/generated/mne.io.Raw.html>`__ or
    # `mne.Epochs <https://mne.tools/0.11/generated/mne.Epochs.html#mne.Epochs>`__
    # or apply your own functions, either to the MNE object or the underlying
    # numpy array.
    #
    # .. note::
    #    These prepocessings are now directly applied to the loaded
    #    data, and not on-the-fly applied as transformations in
    #    PyTorch-libraries like
    #    `torchvision <https://pytorch.org/docs/stable/torchvision/index.html>`__.
    #

    from braindecode.datautil.preprocess import exponential_moving_standardize
    from braindecode.datautil.preprocess import MNEPreproc, NumpyPreproc, preprocess

    low_cut_hz = 0.  # low cut frequency for filtering
    high_cut_hz = 38.  # high cut frequency for filtering
    # Parameters for exponential moving standardization
    factor_new = 1e-3
    init_block_size = 1000

    preprocessors = [
        # keep only EEG sensors
        MNEPreproc(fn='pick_types', eeg=True, meg=False, stim=False),
        # convert from volt to microvolt, directly modifying the numpy array
        NumpyPreproc(fn=lambda x: x * 1e6),
        # # bandpass filter
        MNEPreproc(fn='filter', l_freq=low_cut_hz, h_freq=high_cut_hz),
        # # exponential moving standardization
        NumpyPreproc(fn=exponential_moving_standardize, factor_new=factor_new,
            init_block_size=init_block_size)
    ]

    # Transform the data
    preprocess(dataset, preprocessors)

    ######################################################################
    # Cut Compute Windows
    # ~~~~~~~~~~~~~~~~~~~
    #

    ######################################################################
    # Now we cut out compute windows, the inputs for the deep networks during
    # training. In the case of trialwise decoding, we just have to decide if
    # we want to cut out some part before and/or after the trial. For this
    # dataset, in our work, it often was beneficial to also cut out 500 ms
    # before the trial.
    #

    import numpy as np
    from braindecode.datautil.windowers import create_windows_from_events

    trial_start_offset_seconds = -0.5
    # Extract sampling frequency, check that they are same in all datasets
    sfreq = dataset.datasets[0].raw.info['sfreq']
    assert all([ds.raw.info['sfreq'] == sfreq for ds in dataset.datasets])
    # Calculate the trial start offset in samples.
    trial_start_offset_samples = int(trial_start_offset_seconds * sfreq)

    # Create windows using braindecode function for this. It needs parameters to define how
    # trials should be used.
    # windows_dataset = create_windows_from_events(
    #     dataset,
    #     trial_start_offset_samples=trial_start_offset_samples,
    #     trial_stop_offset_samples=0,
    #     preload=True,
    # )

    #sliding window aug
    windows_dataset = create_windows_from_events(
        dataset,
        trial_start_offset_samples=trial_start_offset_samples,
        trial_stop_offset_samples=0,
        # window_size_samples=1000,
        # window_stride_samples=74,
        drop_last_window=False,
        preload=True,
    )

    ######################################################################
    # Split dataset into train and valid
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    #

    ######################################################################
    # We can easily split the dataset using additional info stored in the
    # description attribute, in this case ``session`` column. We select
    # ``session_T`` for training and ``session_E`` for validation.
    #

    splitted = windows_dataset.split('session')
    train_set = splitted['session_T']
    valid_set = splitted['session_E']

    ######################################################################
    # Create model
    # ------------
    #

    ######################################################################
    # Now we create the deep learning model! Braindecode comes with some
    # predefined convolutional neural network architectures for raw
    # time-domain EEG. Here, we use the shallow ConvNet model from `Deep
    # learning with convolutional neural networks for EEG decoding and
    # visualization <https://arxiv.org/abs/1703.05051>`__. These models are
    # pure `PyTorch <https://pytorch.org>`__ deep learning models, therefore
    # to use your own model, it just has to be a normal PyTorch
    # `nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__.
    #

    import torch
    from braindecode.util import set_random_seeds
    from braindecode.models import ShallowFBCSPNet,Deep4Net, EEGNetv4, EEGResNet
    from mynetworks import Deep4Net_origin, ConvClfNet, FcClfNet, EEGNet_v2, ShallowNet

    cuda = torch.cuda.is_available()  # check if GPU is available, if True chooses to use it
    device = 'cuda' if cuda else 'cpu'
    if cuda:
        torch.backends.cudnn.benchmark = True
    seed = 20200220  # random seed to make results reproducible
    # Set random seed to be able to reproduce results
    set_random_seeds(seed=seed, cuda=cuda)

    n_classes=4
    # Extract number of chans and time steps from dataset
    n_chans = train_set[0][0].shape[0]
    input_window_samples = train_set[0][0].shape[1]

    # embedding_net = ShallowNet(4, 22, input_window_samples)
    # model = FcClfNet(embedding_net)

    model = EEGResNet(
        n_chans,
        n_classes,
        input_window_samples=input_window_samples
    )

    # Send model to GPU
    if cuda:
        model.cuda()

    print(model)
    ######################################################################
    # Training
    # --------
    #

    ######################################################################
    # Now we train the network! EEGClassifier is a Braindecode object
    # responsible for managing the training of neural networks. It inherits
    # from skorch.NeuralNetClassifier, so the training logic is the same as in
    # `Skorch <https://skorch.readthedocs.io/en/stable/>`__.
    #

    ######################################################################
    #    **Note**: In this tutorial, we use some default parameters that we
    #    have found to work well for motor decoding, however we strongly
    #    encourage you to perform your own hyperparameter optimization using
    #    cross validation on your training data.
    #

    from skorch.callbacks import LRScheduler
    from skorch.helper import predefined_split

    from braindecode import EEGClassifier
    # # These values we found good for shallow network:
    lr = 0.0625 * 0.01
    weight_decay = 0

    # # For deep4 they should be:
    # lr = 1 * 0.01
    # weight_decay = 0.5 * 0.001

    batch_size = 60
    n_epochs = 200

    clf = EEGClassifier(
        model,
        criterion=torch.nn.NLLLoss,
        optimizer=torch.optim.AdamW,
        train_split=predefined_split(valid_set),  # using valid_set for validation
        optimizer__lr=lr,
        optimizer__weight_decay=weight_decay,
        batch_size=batch_size,
        callbacks=[
            "accuracy", ("lr_scheduler", LRScheduler('CosineAnnealingLR', T_max=n_epochs - 1)),
        ],
        device=device,
    )

    # Model training for a specified number of epochs. `y` is None as it is already supplied
    # in the dataset.
    clf.fit(train_set, y=None, epochs=n_epochs)

    ######################################################################
    # Plot Results
    # ------------
    #

    ######################################################################
    # Now we use the history stored by Skorch throughout training to plot
    # accuracy and loss curves.
    #

    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D
    import pandas as pd
    # Extract loss and accuracy values for plotting from history object
    results_columns = ['train_loss', 'valid_loss', 'train_accuracy', 'valid_accuracy']
    df = pd.DataFrame(clf.history[:, results_columns], columns=results_columns,
                      index=clf.history[:, 'epoch'])

    # get percent of misclass for better visual comparison to loss
    df = df.assign(train_misclass=100 - 100 * df.train_accuracy,
                   valid_misclass=100 - 100 * df.valid_accuracy)

    plt.style.use('seaborn')
    fig, ax1 = plt.subplots(figsize=(8, 3))
    df.loc[:, ['train_loss', 'valid_loss']].plot(
        ax=ax1, style=['-', ':'], marker='o', color='tab:blue', legend=False, fontsize=14)

    ax1.tick_params(axis='y', labelcolor='tab:blue', labelsize=14)
    ax1.set_ylabel("Loss", color='tab:blue', fontsize=14)

    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

    df.loc[:, ['train_misclass', 'valid_misclass']].plot(
        ax=ax2, style=['-', ':'], marker='o', color='tab:red', legend=False)
    ax2.tick_params(axis='y', labelcolor='tab:red', labelsize=14)
    ax2.set_ylabel("Misclassification Rate [%]", color='tab:red', fontsize=14)
    ax2.set_ylim(ax2.get_ylim()[0], 85)  # make some room for legend
    ax1.set_xlabel("Epoch", fontsize=14)

    # where some data has already been plotted to ax
    handles = []
    handles.append(Line2D([0], [0], color='black', linewidth=1, linestyle='-', label='Train'))
    handles.append(Line2D([0], [0], color='black', linewidth=1, linestyle=':', label='Valid'))
    plt.legend(handles, [h.get_label() for h in handles], fontsize=14)
    plt.tight_layout()

    return df

if __name__ == '__main__':
    import pandas as pd
    df_all = pd.DataFrame()
    for id in range(1,10):
        df = exp(id)
        df_all = pd.concat([df_all, df], axis=1)
        df_all.to_csv("EEGnet_trial.csv",mode='w')

ext: py

max_stars_repo_name: PatternRecognition/OpenBMI

max_issues_repo_name: deep-bci-g/OpenBMI

max_forks_repo_name: deep-bci-g/OpenBMI

reason: import numpy


text:
import numpy as np
import random
from collections import namedtuple, deque
import sys

from model import QNetwork

import tensorflow as tf

BUFFER_SIZE = int(1e5)  # replay buffer size
BATCH_SIZE = 64         # minibatch size
GAMMA = 0.99            # discount factor
TAU = 1e-3              # for soft update of target parameters
LR = 5e-4               # learning rate
UPDATE_EVERY = 25        # how often to update the network

class Agent():
    """Interacts with and learns from the environment."""

    def __init__(self, state_size, action_size, seed):
        """Initialize an Agent object.
        Params
        ======
            state_size (int): dimension of each state
            action_size (int): dimension of each action
            seed (int): random seed
        """
        self.state_size = state_size
        self.action_size = action_size
        self.seed = random.seed(seed)

        # Q-Network
        #optimizer = tf.train.RMSPropOptimizer(learning_rate= LR)
        optimizer = tf.train.AdamOptimizer(learning_rate = LR)
        self.Qnetwork = QNetwork(state_size = state_size,
                                 action_size = action_size,
                                 optimizer = optimizer,
                                 gamma=GAMMA, tau = TAU, minibatch_size = BATCH_SIZE)
        # Replay memory
        self.memory = ReplayBuffer(action_size, BUFFER_SIZE, BATCH_SIZE, seed)
        # Initialize time step (for updating every UPDATE_EVERY steps)
        self.t_step = 0

    def step(self, state, action, reward, next_state, done):
        # Save experience in replay memory
        self.memory.add(state, action, reward, next_state, done)

        if len(self.memory) > BATCH_SIZE:
            experiences = self.memory.sample()
            self.learn(experiences)

        # Learn every UPDATE_EVERY time steps.
        self.t_step = (self.t_step + 1) % UPDATE_EVERY
        if self.t_step == 0:
            # ------------------- update target network ------------------- #
            self.Qnetwork.update_target_network()

    def act(self, state, eps=0.):
        """Returns actions for given state as per current policy.

        Params
        ======
            state (array_like): current state
            eps (float): epsilon, for epsilon-greedy action selection
        """
        if len(state.shape) == 1:
            # make it batch-like
            state = state[np.newaxis, :]

        # Epsilon-greedy action selection
        if random.random() > eps:
            action_values = self.Qnetwork.get_action(state)
            return np.argmax(action_values)
        else:
            return random.choice(np.arange(self.action_size))

    def learn(self, experiences):
        """Update value parameters using given batch of experience tuples.

        Params
        ======
            experiences (Tuple[torch.Variable]): tuple of (s, a, r, s', done) tuples
            gamma (float): discount factor
        """
        #states, actions, rewards, next_states, dones = experiences

        ## TODO: compute and minimize the loss
        "*** YOUR CODE HERE ***"
        current_loss = self.Qnetwork.train(experiences)

class ReplayBuffer:
    """Fixed-size buffer to store experience tuples."""

    def __init__(self, action_size, buffer_size, batch_size, seed):
        """Initialize a ReplayBuffer object.

        Params
        ======
            action_size (int): dimension of each action
            buffer_size (int): maximum size of buffer
            batch_size (int): size of each training batch
            seed (int): random seed
        """
        self.action_size = action_size
        self.memory = deque(maxlen=buffer_size)
        self.batch_size = batch_size
        self.experience = namedtuple("Experience", field_names=["state", "action", "reward", "next_state", "done"])
        self.seed = random.seed(seed)

    def add(self, state, action, reward, next_state, done):
        """Add a new experience to memory."""
        e = self.experience(state, action, reward, next_state, done)
        self.memory.append(e)

    def sample(self):
        """Randomly sample a batch of experiences from memory."""
        experiences = random.sample(self.memory, k=self.batch_size)

        states = np.vstack([e.state for e in experiences if e is not None])
        actions = np.vstack([e.action for e in experiences if e is not None])
        rewards = np.vstack([e.reward for e in experiences if e is not None])
        next_states = np.vstack([e.next_state for e in experiences if e is not None])
        dones = np.vstack([e.done for e in experiences if e is not None])

        return (states, actions, rewards, next_states, dones)

    def __len__(self):
        """Return the current size of internal memory."""
        return len(self.memory)


text:
#!/usr/bin/env python3

"""
Tool for automated capturing of EM traces. EMcap can send commands to the target device for starting and stopping
operations using a simple communication protocol over either a serial connection or over TCP.
"""

import numpy as np
import sys
import socket
import os
import signal
import logging
import struct
import binascii
import osmosdr
import argparse
import serial
import subprocess
from gnuradio import blocks
from gnuradio import gr
from gnuradio import uhd
from time import sleep
from datetime import datetime
from emma.utils.socketwrapper import SocketWrapper
from scipy.signal import hilbert
from scipy import fftpack
from emma.emcap.online_client import EMCapOnlineClient
from collections import defaultdict
from emma.emcap.sdr import SDR
from emma.emcap.types import *
from emma.emcap.ttywrapper import TTYWrapper
from emma.utils.utils import binary_to_hex

logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
logger = logging.getLogger(__name__)

hilbert3 = lambda x: hilbert(x, fftpack.next_fast_len(len(x)))[:len(x)]

def handler(signum, frame):
    print("Got CTRL+C")
    exit(0)

signal.signal(signal.SIGINT, handler)

def reset_usrp():
    print("Resetting USRP")
    p = subprocess.Popen(["/usr/lib/uhd/utils/b2xx_fx3_utils", "--reset-device"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    print(p.communicate())

# EMCap class: wait for signal and start capturing using a SDR
class EMCap:
    def __init__(self, args):
        # Determine ctrl socket type
        self.ctrl_socket_type = None
        if args.ctrl == 'serial':
            self.ctrl_socket_type = CtrlType.SERIAL
        elif args.ctrl == 'udp':
            self.ctrl_socket_type = CtrlType.UDP

        # Set up data socket
        self.data_socket = SocketWrapper(socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM), ('127.0.0.1', 3884), self.cb_data)
        self.online = args.online

        # Set up sockets
        if self.ctrl_socket_type == CtrlType.DOMAIN:
            unix_domain_socket = '/tmp/emma.socket'
            self.clear_domain_socket(unix_domain_socket)
            self.ctrl_socket = SocketWrapper(socket.socket(family=socket.AF_UNIX, type=socket.SOCK_STREAM), unix_domain_socket, self.cb_ctrl)
        elif self.ctrl_socket_type == CtrlType.UDP:
            self.ctrl_socket = SocketWrapper(socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM), ('10.0.0.1', 3884), self.cb_ctrl)
        elif self.ctrl_socket_type == CtrlType.SERIAL:
            self.ctrl_socket = TTYWrapper("/dev/ttyUSB0", self.cb_ctrl)
        else:
            logger.error("Unknown ctrl_socket_type")
            exit(1)

        if self.online is not None:
            try:
                self.emma_client = EMCapOnlineClient()
                self.emma_client.connect(self.online, 3885)
            except Exception as e:
                print(e)
                exit(1)

        self.sdr_args = {'hw': args.hw, 'samp_rate': args.sample_rate, 'freq': args.frequency, 'gain': args.gain, 'ds_mode': args.ds_mode, 'agc': args.agc}
        self.sdr = SDR(**self.sdr_args)
        self.store = False
        self.stored_plaintext = []
        self.stored_key = []
        self.stored_data = []
        self.trace_set = []
        self.plaintexts = []
        self.keys = []
        self.preprocessed = []
        self.preprocessed_keys = []
        self.preprocessed_plaintexts = []
        self.limit_counter = 0
        self.limit = args.limit
        self.compress = args.compress
        self.args = args

        if self.sdr.hw == 'usrp':
            self.wait_num_chunks = 0
        else:
            self.wait_num_chunks = 50  # Bug in rtl-sdr?

        self.global_meta = {
            "core:datatype": "cf32_le",
            "core:version": "0.0.1",
            "core:license": "CC0",
            "core:hw": self.sdr.hw,
            "core:sample_rate": self.sdr.samp_rate,
            "core:author": "Pieter Robyns"
        }

        self.capture_meta = {
            "core:sample_start": 0,
            "core:frequency": self.sdr.freq,
            "core:datetime": str(datetime.utcnow()),
        }

    def clear_domain_socket(self, address):
        try:
            os.unlink(address)
        except OSError:
            if os.path.exists(address):
                raise

    def cb_timeout(self):
        logger.warning("Timeout on capture, skipping...")
        self.sdr.stop()

    def cb_data(self, client_socket, client_address, data):
        self.stored_data.append(data)
        return len(data)

    def cb_ctrl(self, client_socket, client_address, data):
        logger.log(logging.NOTSET, "Control packet: %s" % binary_to_hex(data))
        if len(data) < 5:
            # Not enough for TLV
            return 0
        else:
            pkt_type, payload_len = struct.unpack(">BI", data[0:5])
            payload = data[5:]
            if len(payload) < payload_len:
                return 0  # Not enough for payload
            else:
                self.process_ctrl_packet(pkt_type, payload)
                # Send ack
                if self.ctrl_socket_type == CtrlType.SERIAL:
                    client_socket.write(b"k")
                else:
                    client_socket.sendall("k")
                return payload_len + 5

    def parse_ies(self, payload):
        while len(payload) >= 5:
            # Extract IE header
            ie_type, ie_len = struct.unpack(">BI", payload[0:5])
            payload = payload[5:]

            # Extract IE data
            ie = payload[0:ie_len]
            payload = payload[ie_len:]
            logger.debug("IE type %d of len %d: %s" % (ie_type, ie_len, binary_to_hex(ie)))

            # Determine what to do with IE
            if ie_type == InformationElementType.PLAINTEXT:
                self.stored_plaintext = [byte_value for byte_value in ie]
            elif ie_type == InformationElementType.KEY:
                self.stored_key = [byte_value for byte_value in ie]
            else:
                logger.warning("Unknown IE type: %d" % ie_type)

    def preprocess(self, trace_set, plaintexts, keys):
        all_traces = []
        key_set = defaultdict(set)
        pt_set = defaultdict(set)

        for i, trace in enumerate(trace_set):
            if len(trace) < 16384:
                print("--skip")
                continue
            trace = trace[8192:16384]
            trace = np.square(np.abs(np.fft.fft(trace)))
            all_traces.append(trace)

            # Check keys and plaintexts
            num_key_bytes = keys.shape[1]
            for j in range(0, num_key_bytes):
                key_set[j].add(keys[i][j])
                pt_set[j].add(plaintexts[i][j])

                if len(key_set[j]) != 1 or len(pt_set[j]) != 1:
                    print("Keys or plaintexts not equal at index %d" % j)
                    print(key_set)
                    print(pt_set)
                    exit(1)

        all_traces = np.array(all_traces)
        self.preprocessed.append(np.mean(all_traces, axis=0))
        self.preprocessed_plaintexts.append(plaintexts[0])
        self.preprocessed_keys.append(keys[0])

    def save(self, trace_set, plaintexts, keys, ciphertexts=None):
        filename = str(datetime.utcnow()).replace(" ", "_").replace(".", "_").replace(":", "-")
        output_dir = self.args.output_dir

        if self.args.preprocess:
            self.preprocess(trace_set, plaintexts, keys)
            if len(self.preprocessed) >= self.args.traces_per_set:
                logger.info("Dumping %d preprocessed traces to file" % len(self.preprocessed))
                np.save(os.path.join(output_dir, "%s_traces.npy" % filename), np.array(self.preprocessed))
                np.save(os.path.join(output_dir, "%s_textin.npy" % filename), np.array(self.preprocessed_plaintexts))
                np.save(os.path.join(output_dir, "%s_knownkey.npy" % filename), np.array(self.preprocessed_keys))
                self.preprocessed = []
                self.preprocessed_plaintexts = []
                self.preprocessed_keys = []
        else:
            logger.info("Dumping %d traces to file" % len(self.trace_set))
            np.save(os.path.join(output_dir, "%s_traces.npy" % filename), trace_set)
            np.save(os.path.join(output_dir, "%s_textin.npy" % filename), plaintexts)
            np.save(os.path.join(output_dir, "%s_knownkey.npy" % filename), keys)
            if self.compress:
                logger.info("Calling emcap-compress...")
                subprocess.call(['/usr/bin/python', 'emcap-compress.py', os.path.join(output_dir, "%s_traces.npy" % filename)])

    def process_ctrl_packet(self, pkt_type, payload):
        if pkt_type == CtrlPacketType.SIGNAL_START:
            logger.debug("Starting for payload: %s" % binary_to_hex(payload))
            self.parse_ies(payload)
            self.sdr.start()

            # Spinlock until data
            timeout = 3
            current_time = 0.0
            while len(self.stored_data) <= self.wait_num_chunks:
                sleep(0.0001)
                current_time += 0.0001
                if current_time >= timeout:
                    logger.warning("Timeout while waiting for data. Did the SDR crash? Reinstantiating...")
                    del self.sdr
                    self.data_socket.socket.close()
                    self.data_socket = SocketWrapper(socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM), ('127.0.0.1', 3884), self.cb_data)
                    self.data_socket.start()
                    self.sdr = SDR(**self.sdr_args)
                    self.process_ctrl_packet(pkt_type, payload)
        elif pkt_type == CtrlPacketType.SIGNAL_END:
            # self.sdr.sdr_source.stop()
            self.sdr.stop()
            self.sdr.wait()

            logger.debug("Stopped after receiving %d chunks" % len(self.stored_data))
            # sleep(0.5)
            # logger.debug("After sleep we have %d chunks" % len(self.stored_data))

            # Successful capture (no errors or timeouts)
            if len(self.stored_data) > 0:  # We have more than 1 chunk
                # Data to file
                np_data = np.fromstring(b"".join(self.stored_data), dtype=np.complex64)
                self.trace_set.append(np_data)
                self.plaintexts.append(self.stored_plaintext)
                self.keys.append(self.stored_key)

                if len(self.trace_set) >= self.args.traces_per_set:
                    assert(len(self.trace_set) == len(self.plaintexts))
                    assert(len(self.trace_set) == len(self.keys))

                    np_trace_set = np.array(self.trace_set)
                    np_plaintexts = np.array(self.plaintexts, dtype=np.uint8)
                    np_keys = np.array(self.keys, dtype=np.uint8)

                    if self.online is not None:  # Stream online
                        self.emma_client.send(np_trace_set, np_plaintexts, None, np_keys, None)
                    else:  # Save to disk
                        if not self.args.dry:
                            # Write metadata to sigmf file
                            # if sigmf
                            #with open(test_meta_path, 'w') as f:
                            #    test_sigmf = SigMFFile(data_file=test_data_path, global_info=copy.deepcopy(self.global_meta))
                            #    test_sigmf.add_capture(0, metadata=capture_meta)
                            #    test_sigmf.dump(f, pretty=True)
                            # elif chipwhisperer:
                            self.save(np_trace_set, np_plaintexts, np_keys)
                        else:
                            print("Dry run! Not saving.")

                        self.limit_counter += len(self.trace_set)
                        if self.limit_counter >= self.limit:
                            print("Done")
                            exit(0)

                    # Clear results
                    self.trace_set = []
                    self.plaintexts = []
                    self.keys = []

                # Clear
                self.stored_data = []
                self.stored_plaintext = []

    def capture(self, to_skip=0, timeout=1.0):
        # Start listening for signals
        self.data_socket.start()
        self.ctrl_socket.start()

        # Wait until supplicant signals end of acquisition
        while self.ctrl_socket.is_alive():
            self.ctrl_socket.join(timeout=1.0)

        logging.info("Supplicant disconnected on control channel. Stopping...")

def main():
    parser = argparse.ArgumentParser(description='EMCAP')
    parser.add_argument('hw', type=str, choices=['usrp', 'hackrf', 'rtlsdr'], help='SDR capture hardware')
    parser.add_argument('ctrl', type=str, choices=['serial', 'udp'], help='Controller type')
    parser.add_argument('--sample-rate', type=int, default=4000000, help='Sample rate')
    parser.add_argument('--frequency', type=float, default=64e6, help='Capture frequency')
    parser.add_argument('--gain', type=float, default=50, help='RX gain')
    parser.add_argument('--traces-per-set', type=int, default=256, help='Number of traces per set')
    parser.add_argument('--limit', type=int, default=256*400, help='Limit number of traces')
    parser.add_argument('--output-dir', dest="output_dir", type=str, default="/tmp/", help='Output directory to store samples')
    parser.add_argument('--online', type=str, default=None, help='Stream samples to remote EMMA instance at <IP address> for online processing.')
    parser.add_argument('--dry', default=False, action='store_true', help='Do not save to disk.')
    parser.add_argument('--ds-mode', default=False, action='store_true', help='Direct sampling mode.')
    parser.add_argument('--agc', default=False, action='store_true', help='Automatic Gain Control.')
    parser.add_argument('--compress', default=False, action='store_true', help='Compress using emcap-compress.')
    parser.add_argument('--preprocess', default=False, action='store_true', help='Preprocess before storing')  # TODO integrate into emcap.py
    args, unknown = parser.parse_known_args()

    e = EMCap(args)
    e.capture()

if __name__ == '__main__':
    main()

reason: import numpy,from scipy


text:
import torch
import torch.nn as nn
import torch.nn.functional as F

from .res_stack import ResStack
# from res_stack import ResStack

import numpy as np

from .melgan_wav_images import render_histogram
from .melgan_wav_images import rescale_mel

#from wav_images import congrid
import skimage
from skimage.transform import resize

MAX_WAV_VALUE = 32768.0

class Generator(nn.Module):
    def __init__(self, mel_channel):
        super(Generator, self).__init__()
        self.mel_channel = mel_channel

        self.generator = nn.Sequential(
            nn.ReflectionPad1d(3),
            nn.utils.weight_norm(nn.Conv1d(mel_channel, 512, kernel_size=7, stride=1)),

            nn.LeakyReLU(0.2),
            nn.utils.weight_norm(nn.ConvTranspose1d(512, 256, kernel_size=16, stride=8, padding=4)),

            ResStack(256),

            nn.LeakyReLU(0.2),
            nn.utils.weight_norm(nn.ConvTranspose1d(256, 128, kernel_size=16, stride=8, padding=4)),

            ResStack(128),

            nn.LeakyReLU(0.2),
            nn.utils.weight_norm(nn.ConvTranspose1d(128, 64, kernel_size=4, stride=2, padding=1)),

            ResStack(64),

            nn.LeakyReLU(0.2),
            nn.utils.weight_norm(nn.ConvTranspose1d(64, 32, kernel_size=4, stride=2, padding=1)),

            ResStack(32),

            nn.LeakyReLU(0.2),
            nn.ReflectionPad1d(3),
            nn.utils.weight_norm(nn.Conv1d(32, 1, kernel_size=7, stride=1)),
            nn.Tanh(),
        )

    def forward(self, mel):
        mel = (mel + 5.0) / 5.0 # roughly normalize spectrogram
        return self.generator(mel)

    def eval(self, inference=False):
        super(Generator, self).eval()

        # don't remove weight norm while validation in training loop
        if inference:
            self.remove_weight_norm()

    def remove_weight_norm(self):
        for idx, layer in enumerate(self.generator):
            if len(layer.state_dict()) != 0:
                try:
                    nn.utils.remove_weight_norm(layer)
                except:
                    layer.remove_weight_norm()

    def inference(self, mel):
        # TODO/NB: The actual shapes expected
        # eg, torch.Size([1, 80, 972]), where last dim is dynamic
        # but got: torch.Size([1, 1000, 1025])

        print('>>> Original Tacotron Mel Size:')
        print(mel.shape)

        mel = mel.cpu()

        # Swap
        #mel = np.swapaxes(mel, 1, 2)
        #print('>>> Swapped:')
        print(mel.shape)
        #render_histogram(mel.cpu(), 'histograms/swapped.png', row_index=1, col_index=2)

        # Resize
        new_dims = (1, 80, mel.shape[2])
        #new_dims = (1, mel.shape[1], 80)
        #TODO TODO TODOmel = resize(mel.cpu(), new_dims)
        #mel = np.resize(mel.cpu(), (1, 80, mel.shape[2])) # works, but physical appearance is wrong
        #mel = np.resize(mel.cpu(), (1, mel.shape[1], 80))

        #print('>>> Mel Size (resized):')
        #print(mel.shape)
        #render_histogram(mel, 'histograms/resized.png', row_index=1, col_index=2)

        # RESCALE
        #print('>>> RESCALE MEL:')
        #rescale_mel(mel, -14.0, -0.1)

        # Swap 2
        #mel = np.swapaxes(mel, 1, 2)
        #print('>>> Swapped:')
        #print(mel.shape)
        #render_histogram(mel, 'histograms/swapped2.png', row_index=1, col_index=2)

        hop_length = 256
        # pad input mel with zeros to cut artifact
        # see https://github.com/seungwonpark/melgan/issues/8
        #zero = torch.full((1, self.mel_channel, 10), -11.5129).to(mel.device)
        zero = torch.full((1, self.mel_channel, 10), -11.5129) # TODO Removed device

        #TODO TODO mel = torch.from_numpy(mel)

        #print('>>> Zero Size:')
        #print(zero.shape)

        mel = torch.cat((mel, zero), dim=2)

        #mel = mel.cuda()

        audio = self.forward(mel)
        audio = audio.squeeze() # collapse all dimension except time axis
        audio = audio[:-(hop_length*10)]
        # NB(bt, 2020-02-01): I think this is causing it to sound overblown, pinned to const 1 signal
        #audio = MAX_WAV_VALUE * audio
        #audio = MAX_WAV_VALUE * 0.25 * audio
        print('>>> MELGAN <<<')
        print(audio)
        #audio = audio * 20.0
        print('>>> MELGAN AMP <<<')
        print(audio)
        #audio = audio.clamp(min=-MAX_WAV_VALUE, max=MAX_WAV_VALUE-1)
        #audio = audio.short()

        return audio

'''
    to run this, fix
    from . import ResStack
    into
    from res_stack import ResStack
'''
if __name__ == '__main__':
    model = Generator(80)

    x = torch.randn(3, 80, 10)
    print(x.shape)

    y = model(x)
    print(y.shape)
    assert y.shape == torch.Size([3, 1, 2560])

    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(pytorch_total_params)


text:
import os, sys
curr_path = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, os.path.join(curr_path, "../../python"))
sys.path.insert(0, os.path.join(curr_path, "../../example/image-classification"))
import mxnet as mx
import logging
import argparse
import time
import numpy as np

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def parse_args():
    parser = argparse.ArgumentParser(description="command for benchmark kv-store")
    parser.add_argument('--network', type=str, default="resnet",
                        help='the neural network to test')
    parser.add_argument('--gpus', type=str, default='0,1',
                        help='the gpus to be used, e.g "0,1,2,3"')
    parser.add_argument('--depth', type=int, default=152,
                        help='the depth of network, only valid for resnet')
    parser.add_argument('--kv-store', type=str, default='device',
                        help='the kvstore type')
    parser.add_argument('--batch-size', type=int, default=128,
                        help='batch size. should not affect the results')
    parser.add_argument('--num-batches', type=int, default=5,
                        help='number of batches to run')
    parser.add_argument('--disp-batches', type=int, default=1,
                        help='show averaged results for every n batches')
    parser.add_argument('--test-results', type=int, default=1,
                        help='if or not evalute the results correctness')
    parser.add_argument('--data-shape', type=str, default='128,3,224,224',
                        help='input data shape')
    parser.add_argument('--num-classes', type=int, default=1000,
                        help='number of classes')
    parser.add_argument('--optimizer', type=str, default='None',
                        help='the optimizer set to kvstore. None means no optimizer')
    args = parser.parse_args()
    logging.info(args)
    return args

def get_resnet(args):
    resnet_path = os.path.join(curr_path, "./ResNet")
    if not os.path.isdir(resnet_path):
        os.system("git clone https://github.com/tornadomeet/ResNet")
    sys.path.insert(0, resnet_path)
    from symbol_resnet import resnet
    if args.depth == 18:
        units = [2, 2, 2, 2]
    elif args.depth == 34:
        units = [3, 4, 6, 3]
    elif args.depth == 50:
        units = [3, 4, 6, 3]
    elif args.depth == 101:
        units = [3, 4, 23, 3]
    elif args.depth == 152:
        units = [3, 8, 36, 3]
    elif args.depth == 200:
        units = [3, 24, 36, 3]
    else:
        raise ValueError("no experiments done on detph {}, you can do it youself".format(args.depth))

    filter_list=[64, 256, 512, 1024, 2048] if args.depth >=50 else [64, 64, 128, 256, 512]
    bottle_neck = True if args.depth >= 50 else False
    symbol = resnet(units=units, num_stage=4, filter_list=filter_list,
                    num_class=args.num_classes, data_type="imagenet", bottle_neck=bottle_neck, bn_mom=.9, workspace=512)
    return symbol

def get_shapes(symbol, data_shape):
    arg_name = symbol.list_arguments()
    arg_shape, _, _ = symbol.infer_shape(data=data_shape)
    shapes = [s for n,s in zip(arg_name, arg_shape) if 'weight' in n or 'bias' in n]
    return shapes

def diff(a, b):
    return np.sum(np.abs(a.asnumpy() - b.asnumpy()))

def error(gpu_res, cpu_res):
    res = sum([sum([diff(a, b) for a in w]) for w, b in zip(gpu_res, cpu_res)])
    res /= sum([np.sum(np.abs(g.asnumpy())) for g in cpu_res])
    return res

def run():
    args = parse_args();
    # create kvstore and optimizer
    devs = [mx.gpu(int(i)) for i in args.gpus.split(',')]
    kv = mx.kv.create(args.kv_store)
    if args.optimizer == 'None':
        optimizer = None
    else:
        optimizer = mx.optimizer.Optimizer.create_optimizer(args.optimizer)
        updater = mx.optimizer.get_updater(mx.optimizer.Optimizer.create_optimizer(args.optimizer))
        kv.set_optimizer(optimizer)

    # create network
    if args.network == 'resnet':
        symbol = get_resnet(args)
    else:
        import importlib
        symbol = importlib.import_module('symbol_' + args.network).get_symbol(args.num_classes)
    data_shape = tuple([int(s) for s in args.data_shape.split(',')])
    shapes = get_shapes(symbol, data_shape)

    size = float(sum([reduce(lambda x,y : x*y, s, 1) for s in shapes])) * 4 / 1e6
    logging.info('num of arrays = %d, total size = %f MB' % (len(shapes), size))

    for i, s in enumerate(shapes):
        kv.init(i, mx.nd.zeros(s))

    grads_val = [[mx.random.uniform(-1,1,shape=s) for d in devs] for s in shapes]
    grads = [[g.as_in_context(d) for g, d in zip(gs, devs)] for gs in grads_val]
    weights = [[mx.nd.zeros(s, d) for d in devs] for s in shapes]

    cpu_grads = [mx.nd.array(sum([g.asnumpy() for g in gs]))*kv.num_workers for gs in grads_val]
    cpu_weights = [mx.nd.zeros(s) for s in shapes]
    toc = 0
    for b in range(0, args.num_batches+1):
        tic = time.time()
        for i,g in enumerate(grads):
            kv.push(i, g, i)

        for i,w in enumerate(weights):
            kv.pull(i, w, i)
        for ws in weights:
            for w in ws:
                w.wait_to_read()
        toc += time.time() - tic
        if args.test_results:
            if optimizer == None:
                err = error(weights, cpu_grads)
            else:
                for i, wg in enumerate(zip(cpu_weights, cpu_grads)):
                    updater(i, wg[1], wg[0])
                err = error(weights, cpu_weights)
        else:
            err = -1

        if b % args.disp_batches == 0:
            toc /= args.disp_batches
            if b != 0:
                # 0 is used for warmup, ignored
                logging.info('iter %d, %f sec, %f GB/sec per gpu, error %f' % (
                    b, toc, size*2*(len(devs)-1)/len(devs)/toc/1e3, err))
            toc = 0

if __name__ == "__main__":
    run()

max_stars_repo_name: xcgoner/ps-mxnet-new

max_issues_repo_name: xcgoner/ps-mxnet-new

max_forks_repo_name: xcgoner/ps-mxnet-new


text:
#!/usr/bin/python
#
# Converts a LCM log to a "structured" format that is easier to work with
# external tools such as Matlab or Python. The set of messages on a given
# channel can be represented as a structure preserving the original lcm message
# structure.
#
# lcm-log2smat is based on libbot2 script bot-log2mat.
# Modified by G.Troni

import os
import sys
import binascii
import types
import numpy
import re
import getopt

# Python 2.7 and 3
if sys.version_info >= (3, 0):
    import pickle
    xrange  = range
    long    = int
    unicode = int
else:
    import cPickle as pickle

# check which version for mio location
if sys.version_info < (2, 6):
    import scipy.io.mio
else:
    import scipy.io.matlab.mio

from lcm import EventLog
from .scan_for_lcmtypes import *

longOpts = ["help", "print", "pickle", "format", "separator", "channelsToProcess", "ignore", "outfile", "lcm_packages"]

def usage():
    pname, sname = os.path.split(sys.argv[0])
    sys.stderr.write("usage: % s %s < filename > \n" % (sname, str(longOpts)))
    print( """
    -h --help                 print this message
    -p --print                Output log data to stdout instead of to .mat or .pkl
    -k --pickle               Output log data to python pickle .pkl instead of to .mat
    -f --format               print the data format to stderr
    -s --seperator=sep        print data with separator [sep] instead of default to ["" ""]
    -c --channelsToProcess=chan        Parse channelsToProcess that match Python regex [chan] defaults to [".*"]
    -i --ignore=chan          Ignore channelsToProcess that match Python regex [chan]
                              ignores take precedence over includes!
    -o --outfile=ofname       output data to [ofname] instead of default [filename.mat or filename.pkd or stdout]
    -l --lcmtype_pkgs=pkgs    load python modules from comma seperated list of packages [pkgs] defaults to ["botlcm"]
    -v                        Verbose

    """ )
    sys.exit()

def msg_getfields (lcm_msg):
    return lcm_msg.__slots__

def msg_getconstants (lcm_msg):
    # Get full list of valid attributes
    fulllist = dir(lcm_msg)
    # Get constants
    constantslist = [x for x in fulllist if not(x[0]=='_')
                    if not(x=='decode')
                    if not(x=='encode')
                    if x not in msg_getfields (lcm_msg)]
    return constantslist

def msg_to_dict (data, e_channel, msg, statusMsg, verbose=False, lcm_timestamp=-1):

    # Initializing channel
    if e_channel not in data:
        data[e_channel] = dict()

        # Iterate each constant of the LCM message
        constants = msg_getconstants (msg)
        for i in xrange(len(constants)):
            myValue = None
            myValue = eval('msg.' + constants[i])
            data[e_channel][constants[i][:31]] = myValue

    # Get lcm fields and constants
    fields = msg_getfields (msg)

    # Iterate each field of the LCM message
    for i in xrange(len(fields)):
        myValue = None
        myValue = eval(' msg.' + fields[i])
        if (isinstance(myValue,int)     or
            isinstance(myValue,long)    or
            isinstance(myValue,float)   or
            isinstance(myValue,tuple)   or
            isinstance(myValue,unicode) or
            isinstance(myValue,str)):
            try:
                data[e_channel][fields[i][:31]].append(myValue)
            except KeyError as AttributeError:
                data[e_channel][fields[i][:31]] = [(myValue)]

        elif (hasattr(myValue,'__slots__')):
            submsg = eval('msg.' + fields[i])
            msg_to_dict (data[e_channel], fields[i][:31], submsg, statusMsg, verbose)

        else:
            if verbose:
                statusMsg = deleteStatusMsg(statusMsg)
                sys.stderr.write("ignoring field %s from channel %s. \n" %(fields[i], e_channel))
            continue

    # Add extra field with lcm_timestamp
    if lcm_timestamp > 0:
        try:
            data[e_channel]['lcm_timestamp'].append(lcm_timestamp)
        except KeyError as AttributeError:
            data[e_channel]['lcm_timestamp'] = [(lcm_timestamp)]

def deleteStatusMsg(statMsg):
    if statMsg:
        sys.stderr.write("\r")
        sys.stderr.write(" " * (len(statMsg)))
        sys.stderr.write("\r")
    return ""

def parse_and_save (args, opts={}):

    if isinstance(args, (list, tuple)):
        fname = args[0]
    elif isinstance(args, str):
        fname = args
    else:
        usage()

    #default options
    lcm_packages = [ "botlcm"]

    printFname = "stdout"
    printFile = sys.stdout
    verbose = False
    printOutput = False
    savePickle  = False
    printFormat = False
    channelsToIgnore = ""
    checkIgnore = False
    channelsToProcess = ".*"
    separator = ' '
    for o, a in opts:
        if o == "-v":
            verbose = True
        elif o in ("-h", "--help"):
            usage()
        elif o in ("-p", "--print"):
            printOutput = True
        elif o in ("-k", "--pickle"):
            savePickle = True
        elif o in ("-f", "--format"):
            printFormat = True
        elif o in ("-s", "--separator="):
            separator = a
        elif o in ("-o", "--outfile="):
            outFname = a
            printFname = a
        elif o in ("-c", "--channelsToProcess="):
            channelsToProcess = a
        elif o in ("-i", "--ignore="):
            channelsToIgnore = a
            checkIgnore = True
        elif o in ("-l", "--lcm_packages="):
            lcm_packages = a.split(",")
        else:
            assert False, "unhandled option"

    try:
        outFname
    except NameError:
        outDir, outFname = os.path.split(os.path.abspath(fname))

        if savePickle:
            outFname = outDir + "/" + outFname + ".pkl"
        else:
            outFname = outFname.replace(".", "_")
            outFname = outFname.replace("-", "_")
            outFname = outDir + "/" + outFname + ".mat"

    fullPathName = os.path.abspath(outFname)
    dirname = os.path.dirname(fullPathName)
    outBaseName = ".".join(os.path.basename(outFname).split(".")[0:-1])
    fullBaseName = dirname + "/" + outBaseName

    data = {}

    print("Searching for LCM types...")
    type_db = make_lcmtype_dictionary()

    channelsToProcess = re.compile(channelsToProcess)
    channelsToIgnore = re.compile(channelsToIgnore)
    log = EventLog(fname, "r")

    if printOutput:
        sys.stderr.write("opened % s, printing output to %s \n" % (fname, printFname))
        if printFname == "stdout":
            printFile = sys.stdout
        else:
            printFile = open(printFname, "w")
    else:
        sys.stderr.write("opened % s, outputing to % s\n" % (fname, outFname))

    ignored_channels = []
    msgCount = 0
    statusMsg = ""
    startTime = 0

    # Iterate LCM log file
    for e in log:
        if msgCount == 0:
            startTime = e.timestamp

        if e.channel in ignored_channels:
            continue
        if ((checkIgnore and channelsToIgnore.match(e.channel) and len(channelsToIgnore.match(e.channel).group())==len(e.channel)) \
             or (not channelsToProcess.match(e.channel))):
            if verbose:
                statusMsg = deleteStatusMsg(statusMsg)
                sys.stderr.write("ignoring channel %s\n" % e.channel)
            ignored_channels.append(e.channel)
            continue

        packed_fingerprint = e.data[:8]
        lcmtype = type_db.get(packed_fingerprint, None)
        if not lcmtype:
            if verbose:
                statusMsg = deleteStatusMsg(statusMsg)
                sys.stderr.write("ignoring channel %s -not a known LCM type\n" % e.channel)
            ignored_channels.append(e.channel)
            continue
        try:
            msg = lcmtype.decode(e.data)
        except:
            statusMsg = deleteStatusMsg(statusMsg)
            sys.stderr.write("error: couldn't decode msg on channel %s\n" % e.channel)
            continue

        msgCount = msgCount + 1
        if (msgCount % 5000) == 0:
            statusMsg = deleteStatusMsg(statusMsg)
            statusMsg = "read % d messages, % d %% done" % (msgCount, log.tell() / float(log.size())*100)
            sys.stderr.write(statusMsg)
            sys.stderr.flush()

        msg_to_dict (data, e.channel, msg, statusMsg, verbose, (e.timestamp - startTime) / 1e6)

    deleteStatusMsg(statusMsg)
    if not printOutput:

        sys.stderr.write("loaded all %d messages, saving to % s\n" % (msgCount, outFname))

        if savePickle:

            # Pickle the list/dictonary using the highest protocol available.
            output = open(outFname, 'wb')
            pickle.dump(data, output, -1)
            output.close()
        else:
            # Matlab format using scipy
            if sys.version_info < (2, 6):
                scipy.io.mio.savemat(outFname, data)
            else:
                scipy.io.matlab.mio.savemat(outFname, data, oned_as='row')

if __name__ == "__main__":
    # Parse command line arguments
    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], "hpktvfs:c:i:o:l:", longOpts)
    except getopt.GetoptError as err:
        # print help information and exit:
        print (str(err)) # will print something like "option -a not recognized"
        usage()
    if len(args) != 1:
        usage()

    # Run main parser
    parse_and_save (args,opts)

max_issues_repo_name: WatsonZhouAnda/Cheetah-Software

max_forks_repo_name: WatsonZhouAnda/Cheetah-Software

reason: import numpy,import scipy


text:
import pytest

import numpy as np

import astropy.units as u
from astropy.modeling import InputParameterError

from .helpers import _invalid_x_range
from .helpers import (
    all_models,
    param_ave_models_Rv,
    param_ave_models_Rv_fA,
    param_ave_models,
    ave_models,
    grain_models,
)

@pytest.mark.parametrize("model", all_models)
def test_nounits_warning(model):
    ext = model()
    x = np.arange(ext.x_range[0], ext.x_range[1], 0.1)

    with pytest.warns(
        UserWarning, match="x has no units, assuming x units are inverse microns"
    ):
        ext(x)

@pytest.mark.parametrize("model", all_models)
def test_units_nowarning_expected(model):
    ext = model()
    x = [0.5 * (ext.x_range[0] + ext.x_range[1])] / u.micron

    with pytest.warns(None) as record:
        ext(x)
    assert len(record) == 0

@pytest.mark.parametrize("model", param_ave_models_Rv)
@pytest.mark.parametrize("Rv_invalid", [-1.0, 0.0, 1.9, 6.1, 10.0])
def test_invalid_Rv_input(model, Rv_invalid):
    with pytest.raises(InputParameterError) as exc:
        model(Rv=Rv_invalid)
    assert exc.value.args[0] == "parameter Rv must be between 2.0 and 6.0"

@pytest.mark.parametrize("RvA_invalid", [-1.0, 0.0, 1.9, 6.1, 10.0])
@pytest.mark.parametrize("model", param_ave_models_Rv_fA)
def test_invalid_RvA_input(model, RvA_invalid):
    with pytest.raises(InputParameterError) as exc:
        model(RvA=RvA_invalid)
    assert exc.value.args[0] == "parameter RvA must be between 2.0 and 6.0"

@pytest.mark.parametrize("fA_invalid", [-1.0, -0.1, 1.1, 10.0])
@pytest.mark.parametrize("model", param_ave_models_Rv_fA)
def test_invalid_fA_input(model, fA_invalid):
    with pytest.raises(InputParameterError) as exc:
        model(fA=fA_invalid)
    assert exc.value.args[0] == "parameter fA must be between 0.0 and 1.0"

@pytest.mark.parametrize("model", all_models)
def test_invalid_wavenumbers(model):
    tmodel = model()
    x_invalid_all = [-1.0, 0.9 * tmodel.x_range[0], 1.1 * tmodel.x_range[1]]
    for x_invalid in x_invalid_all:
        _invalid_x_range(x_invalid, tmodel, tmodel.__class__.__name__)
        _invalid_x_range(x_invalid / u.micron, tmodel, tmodel.__class__.__name__)
        _invalid_x_range(u.micron / x_invalid, tmodel, tmodel.__class__.__name__)
        _invalid_x_range(
            u.angstrom * 1e4 / x_invalid, tmodel, tmodel.__class__.__name__
        )

@pytest.mark.parametrize("model", param_ave_models + ave_models)
def test_extinguish_no_av_or_ebv(model):
    ext = model()
    with pytest.raises(InputParameterError) as exc:
        ext.extinguish(ext.x_range[0])
    assert exc.value.args[0] == "neither Av or Ebv passed, one required"

@pytest.mark.parametrize("model", grain_models)
def test_possible_grain_model(model):
    with pytest.raises(InputParameterError) as exc:
        model("badmodename")
    assert exc.value.args[0] == "modelname not recognized"

reason: import numpy,import astropy,from astropy