text:
from mpi4py import MPI
import unittest

import sys
import os
import numpy as np

import pytools  # runko python tools
import pycorgi
import pyrunko

#from visualize_pic import Particles
#from visualize_pic import plot2dParticles
#from visualize import plotNode
#from visualize import plot2dYee
#from visualize import pytools.visualize.get_yee_2D
#from visualize import saveVisz

try:
    import matplotlib.pyplot as plt
except:
    pass

#make tests deterministic by fixing the RNG seed
np.random.seed(0)

def density_profile(xloc, ispcs, conf):
    return conf.ppc

def filler_no_velocity(xloc, ispcs, conf):

    # perturb position between x0 + RUnif[0,1)
    xx = xloc[0] + np.random.rand(1)
    yy = xloc[1] + np.random.rand(1)
    #zz = xloc[2] + np.random.rand(1)
    zz = 0.0

    x0 = [xx, yy, zz]
    u0 = [0.0, 0.0, 0.0]
    return x0, u0

# random number between [xmin, xmax]
def randab(xmin, xmax):
    return xmin + (xmax-xmin)*np.random.rand(1)

def filler(xloc, ispcs, conf):

    # perturb position between x0 + RUnif[0,1)
    xx = xloc[0] + np.random.rand(1)
    yy = xloc[1] + np.random.rand(1)
    #zz = xloc[2] + np.random.rand(1)
    zz = 0.5

    ur = conf.vel
    uc = randab(0.0, 2.0*np.pi)
    us = randab(0.0, 1.0*np.pi)

    #3D
    #ux = ur*np.sin( uc )*np.sin(us)
    #uy = ur*np.cos( uc )*np.sin(us)
    #uz = ur*np.cos(us)

    #2D
    ux = ur*np.sin( uc )
    uy = ur*np.cos( uc )
    uz = 0.0

    x0 = [xx, yy, zz]
    u0 = [ux, uy, uz]

    return x0, u0

def filler3D(xloc, ispcs, conf):

    # perturb position between x0 + RUnif[0,1)
    xx = xloc[0] + np.random.rand(1)
    yy = xloc[1] + np.random.rand(1)
    zz = xloc[2] + np.random.rand(1)

    ur = conf.vel
    uc = randab(0.0, 2.0*np.pi)
    us = randab(0.0, 1.0*np.pi)

    #3D
    ux = ur*np.sin( uc )*np.sin(us)
    uy = ur*np.cos( uc )*np.sin(us)
    uz = ur*np.cos(us)

    x0 = [xx, yy, zz]
    u0 = [ux, uy, uz]

    return x0, u0

def filler_xvel(xloc, ispcs, conf):

    # perturb position between x0 + RUnif[0,1)
    xx = xloc[0] + np.random.rand(1)
    yy = xloc[1] + np.random.rand(1)
    #zz = xloc[2] + np.random.rand(1)

    #xx = xloc[0] #+ 0.5
    #yy = xloc[1] #+ 0.5
    zz = 0.5

    x0 = [xx, yy, zz]

    #if (conf.NyMesh < yy < 2.0*conf.NyMesh) and (conf.NxMesh < xx < 2.0*conf.NxMesh):
    #    u0 = [0.0, 1.0, 0.0]
    #else:
    #    u0 = [0.0, 0.0, 0.0]

    u0 = [+1.0, +1.0, +1.0]
    #u0 = [-1.0, -1.0, -1.0]

    return x0, u0

def zero_field(x,y,z):
    return 0.0

def const_field(x, y, z):
    return 1.0

def linear_field(x, y, z):
    #print("x = {} y = {} z = {}".format(x,y,z))
    #return 1.0*x
    return 1.0*x + 1.0*y
    #return 10.0*y
    #return 100.0*z
    #return 1.0*x + 10.0*y
    return 1.0*x + 10.0*y + 100.0*z

# insert initial electromagnetic setup (or solve Poisson eq)
def insert_em(grid, conf, ffunc):

    Lx  = conf.Nx*conf.NxMesh #XXX scaled length
    for i in range(grid.get_Nx()):
        for j in range(grid.get_Ny()):
            for k in range(grid.get_Nz()):
                c = grid.get_tile(i,j,k)
                yee = c.get_yee()

                for l in range(conf.NxMesh):
                    for m in range(conf.NyMesh):
                        for n in range(conf.NzMesh):

                            # get x_i,j,k
                            xloc0 = pytools.ind2loc((i,j,k), (l,m,n), conf)

                            #get x_i+1/2, x_j+1/2, x_k+1/2
                            xloc1 = pytools.ind2loc((i,j,k), (l+1,m,  n),   conf)
                            yloc1 = pytools.ind2loc((i,j,k), (l,  m+1,n),   conf)
                            zloc1 = pytools.ind2loc((i,j,k), (l,  m,  n+1), conf)

                            # values in Yee lattice corners
                            xcor = xloc0[0]
                            ycor = xloc0[1]
                            zcor = xloc0[2]

                            # values in Yee lattice mids
                            xmid = 0.5*(xloc0[0] + xloc1[0])
                            ymid = 0.5*(xloc0[1] + yloc1[1])
                            zmid = 0.5*(xloc0[2] + zloc1[2])

                            #val = ffunc(xmid, ymid, zmid)

                            # enforce Yee lattice structure
                            yee.ex[l,m,n] = ffunc(xmid, ycor, zcor)
                            yee.ey[l,m,n] = ffunc(xcor, ymid, zcor)+1.0
                            #yee.ez[l,m,n] = ffunc(xcor, ycor, zmid)+2.0
                            yee.ez[l,m,n] = ffunc(xcor, ycor, zcor)+2.0  #2D hack

                            #yee.bx[l,m,n] = ffunc(xcor, ymid, zmid)+3.0
                            yee.bx[l,m,n] = ffunc(xcor, ymid, zcor)+3.0  #2D hack
                            #yee.by[l,m,n] = ffunc(xmid, ycor, zmid)+4.0 #2D hack
                            yee.by[l,m,n] = ffunc(xmid, ycor, zcor)+4.0
                            yee.bz[l,m,n] = ffunc(xmid, ymid, zcor)+5.0

                            yee.jx[l,m,n] = ffunc(xmid, ymid, zmid)
                            yee.jy[l,m,n] = ffunc(xmid, ymid, zmid)
                            yee.jz[l,m,n] = ffunc(xmid, ymid, zmid)

# basic Conf file/class for PiC simulation testing
class Conf:

    Nx = 1
    Ny = 1
    Nz = 1

    NxMesh = 5
    NyMesh = 5
    NzMesh = 1

    xmin = 0.0
    xmax = 10.0

    ymin = 0.0
    ymax = 10.0

    zmin = 0.0
    zmax = 10.0

    cfl = 0.45
    c_omp = 10.0
    ppc = 1

    gamma_e = 0.0
    gamma_i = 0.0

    dx = 1.0
    dy = 1.0
    dz = 1.0

    me = 1
    mi = 1

    Nspecies = 1

    outdir = "out"

    vel = 0.1

    qe = 1.0
    qi =-1.0

    twoD   = False
    threeD = False

    #def __init__(self):
    #    print("initialized...")

    #update bounding box sizes
    #
    # NOTE: NxMesh = 5 grid looks like this:
    #
    # xmin      xmax
    #  |         |
    #  v         v
    #  |_|_|_|_|_
    #  0 1 2 3 4 5
    #
    def update_bbox(self):
        self.xmin = 0.0
        self.xmax = self.Nx*self.NxMesh

        self.ymin = 0.0
        self.ymax = self.Ny*self.NyMesh

        self.zmin = 0.0
        self.zmax = self.Nz*self.NzMesh

class PIC(unittest.TestCase):

    def test_communication(self):

        #plt.fig = plt.figure(1, figsize=(3,3))
        #plt.rc('font', family='serif', size=12)
        #plt.rc('xtick')
        #plt.rc('ytick')
        #
        #gs = plt.GridSpec(1, 1)
        #
        #axs = []
        #for ai in range(1):
        #    axs.append( plt.subplot(gs[ai]) )

        conf = Conf()
        conf.NxMesh = 3
        conf.NyMesh = 3

        conf.twoD = True
        conf.Nx = 3
        conf.Ny = 3
        conf.Ny = 1
        conf.update_bbox()

        conf.vel = 0.3

        grid = pycorgi.twoD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax)

        pytools.pic.load_tiles(grid, conf)
        insert_em(grid, conf, const_field)
        pytools.pic.inject(grid, filler, density_profile, conf) #pytools.pic.injecting plasma particles

        # push particles couple of times to make them leak into neighboring tiles
        pusher   = pyrunko.pic.twoD.BorisPusher()

        for lap in range(40):
            #plot2dParticles(axs[0], grid, conf)
            #saveVisz(lap, grid, conf)

            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                pusher.solve(tile)

            ##################################################
            # communication

            #update particle boundaries
            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                tile.check_outgoing_particles()

            # global mpi exchange (independent)
            for cid in grid.get_boundary_tiles():
                tile = grid.get_tile(cid)
                tile.pack_outgoing_particles()

            # MPI global exchange
            # transfer primary and extra data
            grid.send_data(0) #(indepdendent)
            grid.send_data(1) #(indepdendent)

            grid.recv_data(0) #(indepdendent)
            grid.recv_data(1) #(indepdendent)

            grid.wait_data(0) #(indepdendent)
            grid.wait_data(1) #(indepdendent)

            # global unpacking (independent)
            for cid in grid.get_virtual_tiles():
                tile = grid.get_tile(cid)
                tile.unpack_incoming_particles()
                tile.check_outgoing_particles()

            # transfer local + global
            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                tile.get_incoming_particles(grid)

            # delete local transferred particles
            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                tile.delete_transferred_particles()

            for cid in grid.get_virtual_tiles():
                tile = grid.get_tile(cid)
                tile.delete_all_particles()

        # count how many particles we now have
        n_particles = 0
        for i in range(conf.Nx):
            for j in range(conf.Ny):
                for k in range(conf.Nz):
                    cid = grid.id(i,j)
                    c = grid.get_tile(cid)

                    container = c.get_container(0)
                    #print("({},{},{}) has {}".format(i,j,k,len(container.loc(0))))
                    n_particles += len(container.loc(0))

                    #self.assertTrue( 0.0 <= container.loc(0) <= conf.xmax )
                    #self.assertTrue( 0.0 <= container.loc(1) <= conf.ymax )
                    #self.assertTrue( 0.0 <= container.loc(2) <= conf.zmax )

                    for prtcl in range(len(container.loc(0))):
                        #print("{} {} {} maxs {} {} {} id {}/{}".format(
                        #container.loc(0)[prtcl],
                        #container.loc(1)[prtcl],
                        #container.loc(2)[prtcl],
                        #conf.xmax, conf.ymax, conf.zmax,
                        #container.id(0)[prtcl],
                        #container.id(1)[prtcl],
                        #))

                        #print("prtcl {} x={} y={} z={} vx={} vy={} vz={}".format(
                        #    prtcl,
                        #    container.loc(0)[prtcl],
                        #    container.loc(1)[prtcl],
                        #    container.loc(2)[prtcl],
                        #    container.vel(0)[prtcl],
                        #    container.vel(1)[prtcl],
                        #    container.vel(2)[prtcl]))

                        # check location
                        self.assertTrue( 0.0 <= container.loc(0)[prtcl] <= conf.xmax )
                        self.assertTrue( 0.0 <= container.loc(1)[prtcl] <= conf.ymax )
                        self.assertTrue( 0.0 <= container.loc(2)[prtcl] <= conf.zmax )

                        # check velocity
                        velx = container.vel(0)[prtcl]
                        vely = container.vel(1)[prtcl]
                        velz = container.vel(2)[prtcl]
                        vel = np.sqrt( velx*velx + vely*vely + velz*velz )
                        self.assertAlmostEqual( vel, conf.vel, places=5 )

        tot_particles = (conf.Nx*conf.NxMesh *
                        conf.Ny*conf.NyMesh *
                        conf.Nz*conf.NzMesh *
                        conf.ppc)

        #tot_particles =(conf.NxMesh *
        #                conf.NyMesh *
        #                conf.NzMesh *
        #                conf.ppc)

        # assert that there is equal number of particles as we began with
        self.assertEqual( tot_particles, n_particles )

    def test_communication3D(self):

        #plt.fig = plt.figure(1, figsize=(3,3))
        #plt.rc('font', family='serif', size=12)
        #plt.rc('xtick')
        #plt.rc('ytick')
        #
        #gs = plt.GridSpec(1, 1)
        #
        #axs = []
        #for ai in range(1):
        #    axs.append( plt.subplot(gs[ai]) )

        conf = Conf()
        conf.threeD = True
        conf.NxMesh = 3
        conf.NyMesh = 3
        conf.NzMesh = 3

        conf.Nx = 3
        conf.Ny = 3
        conf.Nz = 3
        conf.update_bbox()

        conf.vel = 0.3

        grid = pycorgi.threeD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax, conf.zmin, conf.zmax)

        pytools.pic.load_tiles(grid, conf)
        insert_em(grid, conf, const_field)
        pytools.pic.inject(grid, filler3D, density_profile, conf) #pytools.pic.injecting plasma particles

        # push particles couple of times to make them leak into neighboring tiles
        pusher = pyrunko.pic.threeD.BorisPusher()

        for lap in range(40):
            #plot2dParticles(axs[0], grid, conf)
            #saveVisz(lap, grid, conf)

            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                pusher.solve(tile)

            ##################################################
            # communication

            #update particle boundaries
            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                tile.check_outgoing_particles()

            # global mpi exchange (independent)
            for cid in grid.get_boundary_tiles():
                tile = grid.get_tile(cid)
                tile.pack_outgoing_particles()

            # MPI global exchange
            # transfer primary and extra data
            grid.send_data(0) #(indepdendent)
            grid.send_data(1) #(indepdendent)

            grid.recv_data(0) #(indepdendent)
            grid.recv_data(1) #(indepdendent)

            grid.wait_data(0) #(indepdendent)
            grid.wait_data(1) #(indepdendent)

            # global unpacking (independent)
            for cid in grid.get_virtual_tiles():
                tile = grid.get_tile(cid)
                tile.unpack_incoming_particles()
                tile.check_outgoing_particles()

            # transfer local + global
            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                tile.get_incoming_particles(grid)

            # delete local transferred particles
            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                tile.delete_transferred_particles()

            for cid in grid.get_virtual_tiles():
                tile = grid.get_tile(cid)
                tile.delete_all_particles()

        # count how many particles we now have
        n_particles = 0
        for i in range(conf.Nx):
            for j in range(conf.Ny):
                for k in range(conf.Nz):
                    cid = grid.id(i,j,k)
                    c = grid.get_tile(cid)

                    container = c.get_container(0)
                    #print("({},{},{}) has {}".format(i,j,k,len(container.loc(0))))
                    n_particles += len(container.loc(0))

                    #self.assertTrue( 0.0 <= container.loc(0) <= conf.xmax )
                    #self.assertTrue( 0.0 <= container.loc(1) <= conf.ymax )
                    #self.assertTrue( 0.0 <= container.loc(2) <= conf.zmax )

                    for prtcl in range(len(container.loc(0))):
                        #print("{} {} {} maxs {} {} {} id {}/{}".format(
                        #container.loc(0)[prtcl],
                        #container.loc(1)[prtcl],
                        #container.loc(2)[prtcl],
                        #conf.xmax, conf.ymax, conf.zmax,
                        #container.id(0)[prtcl],
                        #container.id(1)[prtcl],
                        #))

                        #print("prtcl {} x={} y={} z={} vx={} vy={} vz={}".format(
                        #    prtcl,
                        #    container.loc(0)[prtcl],
                        #    container.loc(1)[prtcl],
                        #    container.loc(2)[prtcl],
                        #    container.vel(0)[prtcl],
                        #    container.vel(1)[prtcl],
                        #    container.vel(2)[prtcl]))

                        # check location
                        self.assertTrue( 0.0 <= container.loc(0)[prtcl] <= conf.xmax )
                        self.assertTrue( 0.0 <= container.loc(1)[prtcl] <= conf.ymax )
                        self.assertTrue( 0.0 <= container.loc(2)[prtcl] <= conf.zmax )

                        # check velocity
                        velx = container.vel(0)[prtcl]
                        vely = container.vel(1)[prtcl]
                        velz = container.vel(2)[prtcl]
                        vel = np.sqrt( velx*velx + vely*vely + velz*velz )
                        self.assertAlmostEqual( vel, conf.vel, places=5 )

        tot_particles = (conf.Nx*conf.NxMesh *
                         conf.Ny*conf.NyMesh *
                         conf.Nz*conf.NzMesh *
                        conf.ppc)

        #tot_particles =(conf.NxMesh *
        #                conf.NyMesh *
        #                conf.NzMesh *
        #                conf.ppc)

        # assert that there is equal number of particles as we began with
        self.assertEqual( tot_particles, n_particles )

    def test_const_field_interpolation(self):

        conf = Conf()
        conf.twoD = True
        grid = pycorgi.twoD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax)
        pytools.pic.load_tiles(grid, conf)
        insert_em(grid, conf, const_field)
        pytools.pic.inject(grid, filler_no_velocity, density_profile, conf) #pytools.pic.injecting plasma particles

        ##update boundaries
        for j in range(grid.get_Ny()):
            for i in range(grid.get_Nx()):
                tile = grid.get_tile(i,j)
                tile.update_boundaries(grid)

        #interpolate fields
        fintp = pyrunko.pic.twoD.LinearInterpolator()
        for j in range(grid.get_Ny()):
            for i in range(grid.get_Nx()):
                tile = grid.get_tile(i,j)
                fintp.solve(tile)

        #test results
        for i in range(conf.Nx):
            for j in range(conf.Ny):

                cid = grid.id(i,j)
                c = grid.get_tile(cid)
                container = c.get_container(0)

                xx = container.loc(0)
                yy = container.loc(1)
                zz = container.loc(2)

                ux = container.vel(0)
                uy = container.vel(1)
                uz = container.vel(2)

                for i, x in enumerate(xx):
                    #print(i)
                    ex_ref = 1.0
                    ey_ref = 2.0
                    ez_ref = 3.0
                    self.assertEqual(container.ex(i), ex_ref)
                    self.assertEqual(container.ey(i), ey_ref)
                    self.assertEqual(container.ez(i), ez_ref)

                    bx_ref = 4.0
                    by_ref = 5.0
                    bz_ref = 6.0
                    self.assertEqual(container.bx(i), bx_ref)
                    self.assertEqual(container.by(i), by_ref)
                    self.assertEqual(container.bz(i), bz_ref)

    def test_linear_field_interpolation(self):

        conf = Conf()
        conf.twoD = True
        conf.Nx = 3
        conf.Ny = 3

        grid = pycorgi.twoD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax)
        pytools.pic.load_tiles(grid, conf)
        insert_em(grid, conf, linear_field)
        pytools.pic.inject(grid, filler_no_velocity, density_profile, conf) #pytools.pic.injecting plasma particles

        ##update boundaries
        for j in range(grid.get_Ny()):
            for i in range(grid.get_Nx()):
                tile = grid.get_tile(i,j)
                tile.update_boundaries(grid)

        #interpolate fields
        fintp = pyrunko.pic.twoD.LinearInterpolator()
        for j in range(grid.get_Ny()):
            for i in range(grid.get_Nx()):
                tile = grid.get_tile(i,j)
                fintp.solve(tile)

        #test results; avoid boundaries because they are cyclic
        for i in range(1,conf.Nx-1):
            for j in range(1,conf.Ny-1):

                cid = grid.id(i,j)
                c = grid.get_tile(cid)
                container = c.get_container(0)

                xx = container.loc(0)
                yy = container.loc(1)
                zz = container.loc(2)

                for i, x in enumerate(xx):
                    #print(i)

                    ref = linear_field(xx[i], yy[i], zz[i]) #exact/true solution
                    ex_ref = ref + 0.0
                    ey_ref = ref + 1.0
                    ez_ref = ref + 2.0

                    #print("asserting {} {} {} vs {}".format(xx[i], yy[i], zz[i], ref))
                    self.assertAlmostEqual(container.ex(i), ex_ref, places=5)
                    self.assertAlmostEqual(container.ey(i), ey_ref, places=5)
                    self.assertAlmostEqual(container.ez(i), ez_ref, places=5)

                    bx_ref = ref + 3.0
                    by_ref = ref + 4.0
                    bz_ref = ref + 5.0
                    self.assertAlmostEqual(container.bx(i), bx_ref, places=5)
                    self.assertAlmostEqual(container.by(i), by_ref, places=5)
                    self.assertAlmostEqual(container.bz(i), bz_ref, places=5)

    def skip_test_filters(self):
        """ filter integration test with rest of the PIC functions"""

        try:
            plt.fig = plt.figure(1, figsize=(5,7))
            plt.rc('font', family='serif', size=12)
            plt.rc('xtick')
            plt.rc('ytick')

            gs = plt.GridSpec(8, 1)

            axs = []
            for ai in range(8):
                axs.append( plt.subplot(gs[ai]) )
        except:
            pass

        conf = Conf()
        conf.twoD = True
        conf.Nx = 3
        conf.Ny = 3
        conf.Nz = 1
        conf.NxMesh = 3
        conf.NyMesh = 3
        conf.NzMesh = 1
        conf.ppc = 10
        conf.vel = 0.1
        conf.update_bbox()

        grid = pycorgi.twoD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax)
        pytools.pic.load_tiles(grid, conf)
        #insert_em(grid, conf, linear_field)
        insert_em(grid, conf, zero_field)
        pytools.pic.inject(grid, filler, density_profile, conf) #pytools.pic.injecting plasma particles
        #pytools.pic.inject(grid, filler_xvel, conf) #pytools.pic.injecting plasma particles

        #pusher   = pyrunko.pic.twoD.BorisPusher()
        #fintp    = pyrunko.pic.twoD.LinearInterpolator()
        currint  = pyrunko.pic.twoD.ZigZag()

        flt =  pytools.Filter(conf.NxMesh, conf.NyMesh)
        flt.init_gaussian_kernel(1.0, 1.0)

        #for lap in range(0, conf.Nt):
        for lap in range(1):

            #update boundaries
            for j in range(grid.get_Ny()):
                for i in range(grid.get_Nx()):
                    tile = grid.get_tile(i,j)
                    tile.update_boundaries(grid)

            #deposit current
            for j in range(grid.get_Ny()):
                for i in range(grid.get_Nx()):
                    tile = grid.get_tile(i,j)
                    currint.solve(tile)

            #exchange currents
            for j in range(grid.get_Ny()):
                for i in range(grid.get_Nx()):
                    tile = grid.get_tile(i,j)
                    tile.exchange_currents(grid)

            #try:
            #    #plotNode(axs[0], grid, conf)
            #    #plot2dParticles(axs[0], grid, conf, downsample=0.1)
            #    #plot2dYee(axs[1], grid, conf, 'rho')
            #    #plot2dYee(axs[2], grid, conf, 'jx')
            #    #plot2dYee(axs[3], grid, conf, 'jy')
            #    #plot2dYee(axs[4], grid, conf, 'jz')
            #except:
            #    pass

            yee_ref = pytools.visualize.get_yee_2D(grid, conf)

            #filter
            for j in range(grid.get_Ny()):
                for i in range(grid.get_Nx()):
                    #print(" i j ({},{})".format(i,j))
                    tile = grid.get_tile(i,j)
                    flt.get_padded_current(tile, grid)

                    # fourier space filtering
                    flt.fft_image_forward()
                    flt.apply_kernel()
                    flt.fft_image_backward()

                    # direct filtering
                    #for fj in range(1):
                    #    flt.direct_convolve_3point()
                    flt.set_current(tile)

            #cycle new and temporary currents
            for j in range(grid.get_Ny()):
                for i in range(grid.get_Nx()):
                    tile = grid.get_tile(i,j)
                    tile.cycle_current()

            yee = pytools.visualize.get_yee_2D(grid, conf)

            #try:
            #    plot2dYee(axs[5], grid, conf, 'jx')
            #    plot2dYee(axs[6], grid, conf, 'jy')
            #    plot2dYee(axs[7], grid, conf, 'jz')
            #    #saveVisz(lap, grid, conf)
            #except:
            #    pass

            for j in range(conf.Ny*conf.NyMesh):
                for i in range(conf.Nx*conf.NxMesh):
                    #print("({},{})".format(i,j))
                    self.assertAlmostEqual( yee_ref['jx'][i,j], yee['jx'][i,j], places=5 )
                    self.assertAlmostEqual( yee_ref['jy'][i,j], yee['jy'][i,j], places=5 )
                    self.assertAlmostEqual( yee_ref['jz'][i,j], yee['jz'][i,j], places=5 )

    def test_current_exchange(self):
        """test that current exchange is producing hand-build array"""

        #plt.fig = plt.figure(1, figsize=(5,7))
        #plt.rc('font', family='serif', size=12)
        #plt.rc('xtick')
        #plt.rc('ytick')
        #
        #gs = plt.GridSpec(5, 1)
        #
        #axs = []
        #for ai in range(5):
        #    axs.append( plt.subplot(gs[ai]) )

        conf = Conf()
        conf.twoD = True
        conf.Nx = 3
        conf.Ny = 3
        conf.Nz = 1
        conf.NxMesh = 10
        conf.NyMesh = 10
        conf.NzMesh = 1
        conf.ppc = 10
        conf.vel = 0.1
        conf.Nspecies = 2
        conf.me =-1.0        #electron mass-to-charge
        conf.mi = 1.0        #ion mass-to-charge

        conf.update_bbox()

        grid = pycorgi.twoD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax)
        pytools.pic.load_tiles(grid, conf)
        #insert_em(grid, conf, linear_field)
        insert_em(grid, conf, const_field)
        pytools.pic.inject(grid, filler_xvel, density_profile, conf) #pytools.pic.injecting plasma particles

        #pusher   = pyrunko.pic.twoD.BorisPusher()
        #fintp    = pyrunko.pic.twoD.LinearInterpolator()
        currint  = pyrunko.pic.twoD.ZigZag()

        for j in range(grid.get_Ny()):
            for i in range(grid.get_Nx()):
                c = grid.get_tile(i,j)
                yee = c.get_yee(0)
                for l in range(-3, conf.NxMesh+3):
                    for m in range(-3,conf.NyMesh+3):
                        for n in range(-3,conf.NzMesh+3):
                            yee.jx[l,m,n] = 1.0
                            yee.jy[l,m,n] = 1.0
                            yee.jz[l,m,n] = 1.0

        #deposit current
        #for j in range(grid.get_Ny()):
        #    for i in range(grid.get_Nx()):
        #for j in [1]:
        #    for i in [1]:
        #        tile = grid.get_tile(i,j)
        #        currint.solve(tile)

        #exchange currents for the middle one only
        for j in [1]:
            for i in [1]:
                tile = grid.get_tile(i,j)
                tile.exchange_currents(grid)

        #plotNode(axs[0], grid, conf)
        ##plot2dParticles(axs[0], grid, conf, downsample=0.1)
        #plot2dYee(axs[1], grid, conf, 'rho')
        #plot2dYee(axs[2], grid, conf, 'jx')
        #plot2dYee(axs[3], grid, conf, 'jy')
        #plot2dYee(axs[4], grid, conf, 'jz')
        #saveVisz(-1, grid, conf)

        # create reference 3 halo width array by hand
        ref = np.ones((conf.NxMesh, conf.NyMesh))
        ref[0:3,  : ]   = 2
        ref[-3:10,: ]   = 2
        ref[:, 0:3, ]   = 2
        ref[:, -3:10]   = 2

        ref[0:3,  0:3 ] = 4
        ref[7:10, 7:10] = 4
        ref[0:3,  7:10] = 4
        ref[7:10, 0:3 ] = 4
        #print(ref)

        for j in [1]:
            for i in [1]:
                c = grid.get_tile(i,j)
                yee = c.get_yee(0)
                for l in range(conf.NxMesh):
                    for m in range(conf.NyMesh):
                        self.assertEqual(ref[l,m], yee.jx[l,m,0] )
                        self.assertEqual(ref[l,m], yee.jy[l,m,0] )
                        self.assertEqual(ref[l,m], yee.jz[l,m,0] )

    def test_current_deposit(self):
        """
        test that current deposit of + and - particles with same x and v
        produce zero current in total.
        """

        try:
            plt.fig = plt.figure(1, figsize=(5,7))
            plt.rc('font', family='serif', size=12)
            plt.rc('xtick')
            plt.rc('ytick')
            gs = plt.GridSpec(5, 1)

            axs = []
            for ai in range(5):
                axs.append( plt.subplot(gs[ai]) )
        except:
            pass

        conf = Conf()
        conf.twoD = True
        conf.Nx = 3
        conf.Ny = 3
        conf.Nz = 1
        conf.NxMesh = 10
        conf.NyMesh = 10
        conf.NzMesh = 1
        conf.ppc = 1
        conf.vel = 0.1
        conf.Nspecies = 2
        conf.me = -1.0        #electron mass-to-charge
        conf.mi =  1.0        #ion mass-to-charge

        conf.update_bbox()

        grid = pycorgi.twoD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax)
        pytools.pic.load_tiles(grid, conf)
        #insert_em(grid, conf, const_field)
        pytools.pic.inject(grid, filler_xvel, density_profile, conf) #pytools.pic.injecting plasma particles

        #pusher   = pyrunko.pic.twoD.BorisPusher()
        #fintp    = pyrunko.pic.twoD.LinearInterpolator()
        currint  = pyrunko.pic.twoD.ZigZag()

        #deposit current
        for j in range(grid.get_Ny()):
            for i in range(grid.get_Nx()):
                tile = grid.get_tile(i,j)
                currint.solve(tile)

        #exchange currents for the middle one only
        #for j in [1]:
        #    for i in [1]:
        for j in range(grid.get_Ny()):
            for i in range(grid.get_Nx()):
                tile = grid.get_tile(i,j)
                tile.exchange_currents(grid)

        #try:
        #    #plotNode(axs[0], grid, conf)
        #    #plot2dParticles(axs[0], grid, conf, downsample=0.1)
        #    plot2dYee(axs[1], grid, conf, 'rho')
        #    plot2dYee(axs[2], grid, conf, 'jx')
        #    plot2dYee(axs[3], grid, conf, 'jy')
        #    plot2dYee(axs[4], grid, conf, 'jz')
        #    #saveVisz(-2, grid, conf)
        #except:
        #    pass

        for j in range(grid.get_Ny()):
            for i in range(grid.get_Nx()):
                c = grid.get_tile(i,j)
                yee = c.get_yee(0)
                for l in range(conf.NxMesh):
                    for m in range(conf.NyMesh):
                        self.assertAlmostEqual(yee.jx[l,m,0], 0.0, places=5 )
                        self.assertAlmostEqual(yee.jy[l,m,0], 0.0, places=5 )
                        self.assertAlmostEqual(yee.jz[l,m,0], 0.0, places=5 )
                        #self.assertEqual(yee.jx[l,m,0], 0.0 )
                        #self.assertEqual(yee.jy[l,m,0], 0.0 )
                        #self.assertEqual(yee.jz[l,m,0], 0.0 )

    def test_test_particle_initialization(self):

        conf = Conf()
        conf.twoD = True

        conf.NxMesh = 3
        conf.NyMesh = 3
        conf.Nx = 3
        conf.Ny = 3
        conf.Ny = 1
        conf.ppc = 1
        conf.update_bbox()

        grid = pycorgi.twoD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax)

        # this calls internally test particle addition
        pytools.pic.load_tiles(grid, conf)

    def test_current_array_aligning(self):
        #test current deposition + interpolation to see array aligning in practise
        #this is also closest test we have for a full integration test (non-mpi)

        do_plots = False
        try:
            if do_plots:
                plt.fig = plt.figure(1, figsize=(8,10))
                plt.rc('font', family='serif', size=8)
                plt.rc('xtick')
                plt.rc('ytick')

                gs = plt.GridSpec(4, 3)
                gs.update(hspace = 0.0)

                axs = []
                for ai in range(12):
                    axs.append( plt.subplot(gs[ai]) )
        except:
            #print()
            pass

        conf = Conf()
        conf.twoD = True

        conf.outdir = "align_test"
        conf.NxMesh = 3
        conf.NyMesh = 3
        conf.Nx = 3
        conf.Ny = 3
        conf.Nz = 1
        conf.ppc = 1
        conf.npasses = 0
        conf.Nt = 2
        conf.update_bbox()

        if not os.path.exists( conf.outdir ):
            os.makedirs(conf.outdir)

        grid = pycorgi.twoD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax)

        pytools.pic.load_tiles(grid, conf)

        #pytools.pic.inject(grid, filler_xvel, conf) #pytools.pic.injecting plasma particles
        #pytools.pic.inject by hand
        #--------------------------------------------------
        cid    = grid.id(0,0)
        c      = grid.get_tile(cid) #get cell ptr

        container = c.get_container(0) #ispcs
        container.set_keygen_state(0, 0) #number, rank

        #pytools.pic.inject random tricky points

        #x moving
        x0 = [0.0, 0.0, 0.5]
        u0 = [+0.9, 0.0, 0.0]
        container.add_particle(x0, u0, 1.0)

        x0 = [0.0, 2.99, 0.5]
        u0 = [-0.9, 0.0, 0.0]
        container.add_particle(x0, u0, 1.0)

        #y moving
        x0 = [2.99, 2.99, 0.5]
        u0 = [0.0,-0.9, 0.0]
        container.add_particle(x0, u0, 1.0)

        x0 = [0.0, 0.00, 0.5]
        u0 = [0.0,+0.9, 0.0]
        container.add_particle(x0, u0, 1.0)

        #z moving
        x0 = [2.99, 2.99, 0.5]
        u0 = [0.0, 0.0,-0.9]
        container.add_particle(x0, u0, 1.0)

        x0 = [2.0, 0.0, 0.5]
        u0 = [0.0, 0.0,+0.9]
        container.add_particle(x0, u0, 1.0)

        #x0 = [0.5, 0.5, 0.5]
        #for u0 in [
        #        [ 0.9, 0.0, 0.0],
        #        [-0.9, 0.0, 0.0],
        #        [ 0.0, 0.9, 0.0],
        #        [ 0.0,-0.9, 0.0]]:
        #    container.add_particle(x0, u0, 1.0)

        #--------------------------------------------------

        pusher   = pyrunko.pic.twoD.BorisPusher()
        fldprop  = pyrunko.fields.twoD.FDTD2()
        fintp    = pyrunko.pic.twoD.LinearInterpolator()
        currint  = pyrunko.pic.twoD.ZigZag()
        flt      = pyrunko.fields.twoD.Binomial2(conf.NxMesh, conf.NyMesh, conf.NzMesh)

        lap = 0
        for lap in range(lap, conf.Nt):

            #update tile boundaries
            for cid in grid.get_tile_ids():
                tile = grid.get_tile(cid)
                tile.update_boundaries(grid)

            #interpolate fields
            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                fintp.solve(tile)

            #push
            #for cid in grid.get_local_tiles():
            #    tile = grid.get_tile(cid)
            #    pusher.solve(tile)

            #push half B
            for cid in grid.get_tile_ids():
                tile = grid.get_tile(cid)
                fldprop.push_half_b(tile)

            for cid in grid.get_tile_ids():
                tile = grid.get_tile(cid)
                tile.update_boundaries(grid)

            #push full E
            for cid in grid.get_tile_ids():
                tile = grid.get_tile(cid)
                fldprop.push_e(tile)

            #current
            for cid in grid.get_local_tiles():
                tile = grid.get_tile(cid)
                currint.solve(tile)

            for cid in grid.get_tile_ids():
                tile = grid.get_tile(cid)
                tile.exchange_currents(grid)

            #particle movement
            #for cid in grid.get_local_tiles():
            #    tile = grid.get_tile(cid)
            #    tile.check_outgoing_particles()

            #for cid in grid.get_local_tiles():
            #    tile = grid.get_tile(cid)
            #    tile.get_incoming_particles(grid)

            #for cid in grid.get_local_tiles():
            #    tile = grid.get_tile(cid)
            #    tile.delete_transferred_particles()

            #filtering
            #for fj in range(conf.npasses):

            #    #filter each tile
            #    for cid in grid.get_local_tiles():
            #        tile = grid.get_tile(cid)
            #        flt.solve(tile)

            #    #get halo boundaries
            #    for cid in grid.get_local_tiles():
            #        tile = grid.get_tile(cid)
            #        tile.update_boundaries(grid)

            for cid in grid.get_tile_ids():
                tile = grid.get_tile(cid)
                tile.deposit_current()

            #--------------------------------------------------
            # end of cycle // analyzing next

            # TODO can this be removed?
            #build analysis statistics
            #for cid in grid.get_local_tiles():
            #    tile = grid.get_tile(cid)

            xx = container.loc(0)
            if do_plots:
                for i in range(len(xx)):
                    print("i = ", i)
                    print("container ex:", container.ex(i))
                    print("container ey:", container.ey(i))
                    print("container ez:", container.ez(i))
                    print("container bx:", container.bx(i))
                    print("container by:", container.by(i))
                    print("container bz:", container.bz(i))

            #plot
            yee = pytools.visualize.get_yee_2D(grid, conf)

            if do_plots:
                print("jx:", yee['jx'])
                print("jy:", yee['jy'])
                print("jz:", yee['jz'])
                print("sum of current jx:", np.sum(yee['jx']))
                print("sum of current jy:", np.sum(yee['jy']))
                print("sum of current jz:", np.sum(yee['jz']))
                #plotNode(axs[0], grid, conf)
                #plot2dYee(axs[1],  yee, grid, conf, 'rho', label_title=True)

                #plot2dYee(axs[3],  yee, grid, conf, 'jx' , label_title=True)
                #plot2dYee(axs[4],  yee, grid, conf, 'jy' , label_title=True)
                #plot2dYee(axs[5],  yee, grid, conf, 'jz' , label_title=True)
                #plot2dYee(axs[6],  yee, grid, conf, 'ex' , label_title=True)
                #plot2dYee(axs[7],  yee, grid, conf, 'ey' , label_title=True)
                #plot2dYee(axs[8],  yee, grid, conf, 'ez' , label_title=True)
                #plot2dYee(axs[9],  yee, grid, conf, 'bx' , label_title=True)
                #plot2dYee(axs[10], yee, grid, conf, 'by' , label_title=True)
                #plot2dYee(axs[11], yee, grid, conf, 'bz' , label_title=True)
                #saveVisz(lap, grid, conf)

            #assert that arrays are zero (i.e., scheme is charge conserving)
            self.assertTrue( np.abs(np.sum(yee['jx'])) < 1.e7)
            self.assertTrue( np.abs(np.sum(yee['jy'])) < 1.e7)
            self.assertTrue( np.abs(np.sum(yee['jz'])) < 1.e7)

    def test_problematic_prtcls(self):

        # test pic loop behavior when particle is located in wrong container

        conf = Conf()
        conf.twoD = False
        conf.threeD = True
        conf.Nx = 2
        conf.Ny = 2
        conf.Nz = 2

        conf.NxMesh = 3
        conf.NyMesh = 3
        conf.NzMesh = 3

        conf.Nt = 3
        conf.update_bbox()

        grid = pycorgi.threeD.Grid(conf.Nx, conf.Ny, conf.Nz)
        grid.set_grid_lims(conf.xmin, conf.xmax, conf.ymin, conf.ymax, conf.zmin, conf.zmax)

        pytools.pic.load_tiles(grid, conf)
        insert_em(grid, conf, zero_field)

        print("============================================================")

        #--------------------------------------------------
        cid    = grid.id(0,0,0)
        c      = grid.get_tile(cid) #get cell ptr

        container = c.get_container(0) #ispcs
        container.set_keygen_state(0, 0) #number, rank

        u0 = [0.0, 0.0, 0.0]

        # problematic prtcl set 1 (at minimum boundary)
        #x0 = [0.0, 1.0, 1.0]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [1.0, 0.0, 1.0]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [1.0, 1.0, 0.0]
        #container.add_particle(x0, u0, 1.0)

        # problematic prtcl set 2 (at maximum boundary)
        #x0 = [3.0, 1.0, 1.0]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [1.0, 3.0, 1.0]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [1.0, 1.0, 3.0]
        #container.add_particle(x0, u0, 1.0)

        # problematic prtcl set 3 (completely outside)
        #x0 = [3.1, 1.0, 1.0]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [1.0, 3.1, 1.0]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [1.0, 1.0, 3.1]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [-0.1, 1.0, 1.0]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [1.0, -0.1, 1.0]
        #container.add_particle(x0, u0, 1.0)

        #x0 = [1.0, 1.0, -0.1]
        #container.add_particle(x0, u0, 1.0)

        # last tile
        cid    = grid.id(1,1,1)
        c      = grid.get_tile(cid) #get cell ptr
        container = c.get_container(0) #ispcs
        container.set_keygen_state(0, 0) #number, rank

        x0 = [6.0, 5.0, 5.0]
        container.add_particle(x0, u0, 1.0)

        x0 = [0.0, 5.0, 5.0]
        container.add_particle(x0, u0, 1.0)

        #--------------------------------------------------
        pusher   = pyrunko.pic.threeD.BorisPusher()
        fldprop  = pyrunko.fields.threeD.FDTD2()
        fintp    = pyrunko.pic.threeD.LinearInterpolator()
        currint  = pyrunko.pic.threeD.ZigZag()
        flt      = pyrunko.fields.threeD.Binomial2(conf.NxMesh, conf.NyMesh, conf.NzMesh)

        lap = 0
        for lap in range(lap, conf.Nt):

            # --------------------------------------------------
            # push B half
            for tile in pytools.tiles_all(grid):
                fldprop.push_half_b(tile)

            # --------------------------------------------------
            # comm B
            grid.send_data(2)
            grid.recv_data(2)
            grid.wait_data(2)

            # --------------------------------------------------
            # update boundaries
            for tile in pytools.tiles_all(grid):
                tile.update_boundaries(grid)

            ##################################################
            # move particles (only locals tiles)

            # --------------------------------------------------
            # interpolate fields
            for tile in pytools.tiles_local(grid):
                fintp.solve(tile)
            #print('successful interp')

            # --------------------------------------------------
            # push particles in x and u
            #for tile in pytools.tiles_local(grid):
            #    pusher.solve(tile)

            ##################################################
            # advance B half

            # --------------------------------------------------
            # push B half
            for tile in pytools.tiles_all(grid):
                fldprop.push_half_b(tile)

            # --------------------------------------------------
            # comm B
            grid.send_data(1)
            grid.recv_data(1)
            grid.wait_data(1)

            # --------------------------------------------------
            # update boundaries
            for tile in pytools.tiles_all(grid):
                tile.update_boundaries(grid)

            ##################################################
            # advance E

            # --------------------------------------------------
            # push E
            for tile in pytools.tiles_all(grid):
                fldprop.push_e(tile)

            # --------------------------------------------------
            # current calculation; charge conserving current deposition
            for tile in pytools.tiles_local(grid):
                currint.solve(tile)
            #print('successful currint')

            # --------------------------------------------------
            # clear virtual current arrays for boundary addition after mpi
            for tile in pytools.tiles_virtual(grid):
                tile.clear_current()

            # --------------------------------------------------
            # mpi send currents
            grid.send_data(0)
            grid.recv_data(0)
            grid.wait_data(0)

            # --------------------------------------------------
            # exchange currents
            for tile in pytools.tiles_all(grid):
                tile.exchange_currents(grid)

            ##################################################
            # particle communication (only local/boundary tiles)

            # --------------------------------------------------
            # local particle exchange (independent)
            for tile in pytools.tiles_local(grid):
                tile.check_outgoing_particles()

            # --------------------------------------------------
            # global mpi exchange (independent)
            for tile in pytools.tiles_boundary(grid):
                tile.pack_outgoing_particles()

            # --------------------------------------------------
            # MPI global particle exchange
            # transfer primary and extra data
            grid.send_data(3)
            grid.recv_data(3)
            grid.wait_data(3)

            # orig just after send3
            grid.send_data(4)
            grid.recv_data(4)
            grid.wait_data(4)

            # --------------------------------------------------
            # global unpacking (independent)
            for tile in pytools.tiles_virtual(grid):
                tile.unpack_incoming_particles()
                tile.check_outgoing_particles()

            # --------------------------------------------------
            # transfer local + global
            for tile in pytools.tiles_local(grid):
                tile.get_incoming_particles(grid)

            # --------------------------------------------------
            # delete local transferred particles
            for tile in pytools.tiles_local(grid):
                tile.delete_transferred_particles()

            # --------------------------------------------------
            # delete all virtual particles (because new prtcls will come)
            for tile in pytools.tiles_virtual(grid):
                tile.delete_all_particles()

            # --------------------------------------------------
            # add current to E
            for tile in pytools.tiles_all(grid):
                tile.deposit_current()

            # comm E
            grid.send_data(1)
            grid.recv_data(1)
            grid.wait_data(1)

            # --------------------------------------------------
            # comm B
            grid.send_data(2)
            grid.recv_data(2)
            grid.wait_data(2)

            # --------------------------------------------------
            # update boundaries
            for tile in pytools.tiles_all(grid):
                tile.update_boundaries(grid)

        # --------------------------------------------------
        ip = 0
        prtcls = {}
        for tile in pytools.tiles_all(grid):
            container = tile.get_container(0) #ispcs

            xp = container.loc(0)
            yp = container.loc(1)
            zp = container.loc(2)

            for i in range(len(xp)):
                prtcls[ip] = (xp[i],yp[i],zp[i], tile.cid)
                ip += 1

        #for p in prtcls.keys():
        #    print(prtcls[p])


text:
import numpy as np
print(np.matrix([1,2])+np.matrix([2,2]))

list1 = [[1,2],[1,1],[3,3]]
list2 = [6,6,6]
print(list1+list2)
list3 = [[6],[6],[6]]
print(list1+list3)

print(list(map(lambda x:x+['a']  , list1)))

print(list(zip(list2,list1)))

print(dict(zip(list2,list1)))

list4 = [9,8,7]
print(zip(list4,list1))

b = list(zip(list4,list1))
print(b)

def sortFun(item):
    return item[0]

print(sorted(b, key = sortFun))

print(sorted(zip(list4,list1), key = lambda x:x[0]))

list5 = [8,8,6]
c = sorted(zip(list5,list1), key = lambda x:x[0])

print("BBBBBBBBB")
print(sum(np.matrix([1,2])[0]))

print("C:" + str(c))
d = {}
from itertools import *
from functools import *
from numpy import *
for k, g in groupby(c, lambda x:x[0]):
    gl = list(g)
    print("========")
    t = tuple((k,list([0,0])))
    print(t[1])
    gl.append(t)
    print(gl)
    def lam( x,y):
        print("AAAAAAAAAAAA")
        print(x,y)
        return (x[0]+np.matrix(y[1]), x[1] + sum(y[1]))
    d[k] = reduce(lam,gl,(np.matrix(ones(2)),2))

print("RRRRRRRRRRRRRR")
print(d)

aa = {}
bb = {1:1}
cc = aa.update(bb)
print(aa)

print("map dict")
def mapdict(x):
    print(x)
    return {x:d[x][0]/d[x][1]}
mapdict = map(mapdict, d)
#print(list(mapdict))
def lam2(x,y):
    print(x,y)
    x.update(y)
    return x
dict1 = reduce(lam2, mapdict, {})
print(dict1)

print ([0]*3*3)

print([[0 for i in range(3)] for k in range(3)])

print(np.array([2,2])*np.array(np.matrix([2,3])[0]))

ext: py

max_stars_repo_name: simonwanggh/MachineLearning

max_issues_repo_name: simonwanggh/MachineLearning

max_forks_repo_name: simonwanggh/MachineLearning

reason: import numpy,from numpy


text:
#%%
import os
import sys
import requests

import dgl
import torch
import numpy as np

from torch.utils.data import DataLoader, Dataset
from Bio.PDB import PDBParser
from Bio.PDB.NeighborSearch import NeighborSearch

import warnings
warnings.filterwarnings("ignore")
os.path.abspath('.')

DTYPE = np.float32
IDTYPE = np.int32

data_dir = '/home/flower/projects/def-laurence/flower'
# data_dir = '../data'
residue2idx = torch.load('../data/res2idx_dict_core.pt')
# residue2count = torch.load('../data/res2count_dict.pt')

#%%
# def blockPrinting(func):
#     def func_wrapper(*args, **kwargs):
#         # block all printing to the console
#         sys.stdout = open(os.devnull, 'w')
#         # call the method in question
#         value = func(*args, **kwargs)
#         # enable all printing to the console
#         sys.stdout = sys.__stdout__
#         # pass the return value of the method back
#         return value

#     return func_wrapper

class ProtProcess:
    @staticmethod
    def download_pdb(pdb_id, outfile):
        # check if the pdb file has been downloaded in the outfile
        if os.path.exists(outfile):
            # print(f"The file has been downloaded...")
            return 1

        page = 'http://files.rcsb.org/view/{}.pdb'.format(pdb_id)
        req = requests.get(page)
        if req.status_code == 200:
            response = req.text
            if outfile:
                with open(outfile, 'w') as f:
                    f.write(response)
                print('downloaded')
            return 1
        else:
            return 0

    @staticmethod
    def get_residue_feature(chain):
        # init
        res_list = []
        coo_list = []

        # record sequential info
        for res in chain:
            if res.get_resname() == 'HOH':
                continue

            # record residue sequence
            res_name = res.get_resname()
            if residue2idx.get(res_name) is None:
                continue                            # with res2idx_dict_core.pt
                # residue2idx[res_name] = len(residue2idx)
                # residue2count[res_name] = 0                 # TODO: remove later

            res_list.append(residue2idx[res_name])
            # residue2count[res_name] += 1                    # TODO: remove later

            # compute coo of residue by averaging its atoms' coos
            coo = np.concatenate([i.get_coord() for i in res]).reshape((-1, 3)).mean(axis=0)
            coo_list.append(coo)

        return np.array(res_list).astype(IDTYPE), np.concatenate(coo_list).reshape((-1, 3)).astype(DTYPE)

    @staticmethod
    def get_edge_set(dis_cut, l, atom_list, level='R'):
        ns = NeighborSearch(atom_list)
        ns_list = ns.search_all(dis_cut, level=level)

        edge = set()
        for ai, aj in ns_list:
            i, j = ai.get_id()[1]-1, aj.get_id()[1]-1
            if i < j-1 and j < l and i > 0:
                edge.add((i, j))

        return edge

class RandomRotation(object):
    def __init__(self):
        pass

    def __call__(self, x):
        M = np.random.randn(3,3)
        Q, __ = np.linalg.qr(M)
        return x @ Q

#%%
class ProtFunctDataset(Dataset):
    atom_feature_size = len(residue2idx)

    def __init__(self, file_path, mode: str='train', if_transform: bool=True, dis_cut: list=[3.0, 3.5]):
        """Create a dataset object

        Args:
            file_path (str): path to data
            mode (str, optional): {train/test/valid}. Defaults to 'train'.
            if_transform (bool, optional): if applying data augmentation function. Defaults to True.
        """
        self.file_path = file_path
        self.mode = mode

        self.dis_cut = dis_cut
        self.num_bonds = len(dis_cut) + 1
        self.bond2idx = {'covalent':0, 'neighbor<{dis_cut[0]}':1, 'neighbor<{dis_cut[1]}': 2}

        self.transform = RandomRotation() if if_transform else None

        self.__load_data()
        self.len = len(self.targets)

    def __load_data(self):
        """Load preprocessed dataset and parser for input protein structure."""
        data = torch.load(self.file_path)[self.mode]
        # s, e = 100, 500
        self.inputs = data['input_list'][80:]               # TODO: remove
        self.targets = data['target_list'][80:]

        self.parser = PDBParser()

    def __len__(self):
        return self.len

    def to_one_hot(self, data, num_classes):
        one_hot = np.zeros(list(data.shape) + [num_classes])
        one_hot[np.arange(len(data)),data] = 1
        return one_hot

    def __getitem__(self, idx):
        pdb, y = self.inputs[idx], self.targets[idx]

        # parse protein structure
        p, c = pdb.split('.')           # pdb ID and chain ID
        ProtProcess.download_pdb(p, f'{data_dir}/pdb/{p}.pdb')
        # structure = self.parser.get_structure('a', f'{data_dir}/pdb/{p}.pdb')
        # chain = structure[0][c]

        # # generate node features
        # try:
        #     res, x = ProtProcess.get_residue_feature(chain)
        # except:
        #     print('error pdb: ', pdb)
        # num_residues = res.shape[0]
        # res = self.to_one_hot(res, len(residue2idx))[...,None]

        # # augmentation on the coordinates(
        # if self.transform:
        #     x = self.transform(x).astype(DTYPE)

        # # generate edge features
        # src, dst, w = self.__connect_partially([i for i in chain.get_atoms()], num_residues)
        # w = self.to_one_hot(w, self.num_bonds).astype(DTYPE)

        # # create protein representation graph
        # G = dgl.DGLGraph((src, dst))
        # # add node feature
        # x = torch.Tensor(x)
        # G.ndata['x'] = x
        # G.ndata['f'] = torch.Tensor(res)
        # # add edge feature
        # G.edata['d'] = x[dst] - x[src]
        # G.edata['w'] = torch.Tensor(w)

        return y
        # return torch.Tensor([0, 0])

    def norm2units(self, x, denormalize=True, center=True):
        # Convert from normalized to QM9 representation
        if denormalize:
            x = x * self.std
            # Add the mean: not necessary for error computations
            if not center:
                x += self.mean
        # x = self.unit_conversion[self.task] * x
        return x

    def __connect_partially(self, atom_list, num_residues):
        # initialize edges satisfy the different distance cutoffs
        adjacency = {}
        for c in range(1, self.num_bonds):
            for i, j in ProtProcess.get_edge_set(self.dis_cut[-c], num_residues, atom_list):
                adjacency[(i, j)] = self.num_bonds - c
                adjacency[(j, i)] = self.num_bonds - c

        # add covalent bonds
        for i in range(1, num_residues):
            adjacency[(i-1, i)] = 0
            adjacency[(i, i-1)] = 0

        # convert to numpy arrays
        src, dst, w = [], [], []
        for edge, bond in adjacency.items():
            src.append(edge[0])
            dst.append(edge[1])
            w.append(bond)

        return np.array(src).astype(IDTYPE), np.array(dst).astype(IDTYPE), np.array(w)

# def collate(samples):
    # graphs, y, pdb = map(list, zip(*samples))
    # batched_graph = dgl.batch(graphs)
    # return batched_graph, torch.tensor(y), pdb

def to_np(x):
    return x.cpu().detach().numpy()

#%%
# test
if __name__ == '__main__':
    try:
        for mode in ['valid', 'test', 'train']:
            dataset = ProtFunctDataset('../data/ProtFunct.pt', mode=mode)
            dataloader = DataLoader(dataset, batch_size=1, shuffle=False)

            for i, data in enumerate(dataloader):
                if not i % 1000:
                    print(i)
                continue
                # print(f'{mode}: {i}')
    finally:
        torch.save(residue2idx, 'res2idx_dict_tmp.pt')
        # torch.save(residue2count, 'res2count_dict.pt')

reason: import numpy


text:
from mlbase.lazy import numpy as np

def tiled_image(images: "np.ndarray"):
    n, h, w, ch = images.shape
    sq = int(np.ceil(np.sqrt(n)))
    output = np.zeros([sq * h, sq * w, ch], dtype=np.uint8)
    for i, image in enumerate(images):
        x = i % sq
        y = i // sq
        output[h * y:h * (y + 1), w * x:w * (x + 1)] = image
    return output

max_stars_repo_name: n-kats/mlbase

max_issues_repo_name: n-kats/mlbase

max_forks_repo_name: n-kats/mlbase


text:
import gzip

import boto3
import brotli
import numpy as np
import pytest
from pytest import raises
from requests import exceptions

import boss_export.libs.py_compressed_segmentation as csegpy
from boss_export.libs import bosslib, mortonxyz, ngprecomputed

#! failing because offset not aligned with the cube size
@pytest.mark.xfail
def test_ngmorton():
    cube_size = 512, 512, 16
    offset = 0, 0, 2917

    boss_xyz_idx = [o // c for o, c in zip(offset, cube_size)]

    mortonid_actual = mortonxyz.XYZMorton(*boss_xyz_idx)

    ng_morton = ngprecomputed.ngmorton(mortonid_actual, cube_size, offset)
    assert ng_morton == 0

def test_get_scales_ngpath():
    ngpath = "s3://nd-precomputed-volumes/bock11/image"
    info = ngprecomputed.get_scales_ngpath(ngpath)
    assert info["data_type"] == "uint8"
    assert info["num_channels"] == 1
    assert info["type"] == "image"
    assert "scales" in info

    ngpath_broken = "gobblygook"
    assert raises(
        exceptions.MissingSchema, ngprecomputed.get_scales_ngpath, ngpath_broken
    )

def test_ng_key():
    dataset = "bock11_test"
    layer = "image_test"
    chunk_name = "32_32_40/0-512_0-512_2917-2933"

    s3_key = ngprecomputed.get_ng_key(dataset, layer, chunk_name)

    assert s3_key == "bock11_test/image_test/32_32_40/0-512_0-512_2917-2933"

def test_get_chunk_name():
    basescale = 4, 4, 40
    res = 3
    cube_shape = 512, 512, 16
    array_shape = 16, 512, 512
    offset = 0, 0, 2917

    mortonid = 0
    chunk_name = ngprecomputed.get_chunk_name(
        mortonid, basescale, res, cube_shape, array_shape, offset
    )
    assert "32_32_40/0-512_0-512_2917-2933" == chunk_name

    mortonid = 14
    # xyz = mortonxyz.MortonXYZ(mortonid)
    # xyz = [2,1,1]

    res = 0
    chunk_name = ngprecomputed.get_chunk_name(
        mortonid, basescale, res, cube_shape, array_shape, offset
    )
    assert f"4_4_40/{512*2}-{512*3}_{512*1}-{512*2}_{2917+16}-{2917+16*2}" == chunk_name

def test_get_chunk_name_crop():
    # ,295,[512, 512, 11cube_shaped63f2fa28ecd8ef1a1884a787&36&116&869&0&0&295&0,[798.0, 798.0, 2000.0],(16, 109, 512)
    # extent = [1406, 621, 138]

    mortonid = 295
    basescale = [798.0, 798.0, 2000.0]
    array_shape = (16, 109, 512)
    cube_shape = (512, 512, 16)
    offset = [0, 0, 0]
    res = 0
    chunk_name = ngprecomputed.get_chunk_name(
        mortonid, basescale, res, cube_shape, array_shape, offset
    )

    chunk_name_correct = "798.0_798.0_2000.0/512-1024_512-621_112-128"

    assert chunk_name == chunk_name_correct

def test_limit_to_extent():
    # bock11 extents
    extent = 135424, 119808, 4156  # xyz order

    xyz = 135168, 119296, 4144  # xyz order

    # this is the shape it comes back from when reading it out of boss
    data_full = np.zeros((16, 512, 512), dtype="uint8", order="F")  # zyx order

    data_limit = ngprecomputed.crop_to_extent(data_full, xyz, extent)

    limited_shape = tuple(e - i for e, i in zip(extent, xyz))

    assert data_limit.shape == limited_shape[::-1]

def test_numpy_chunk():
    data_array = np.random.randint(0, 2500, (16, 512, 512), "uint16")

    # test no compression
    bstring = ngprecomputed.numpy_chunk(data_array, compression="")
    assert bstring == np.transpose(data_array, (2, 1, 0)).tobytes("F")

    with raises(NotImplementedError):
        ngprecomputed.numpy_chunk(data_array, compression="lz4")

    bstring_gzip = ngprecomputed.numpy_chunk(data_array, compression="gzip")
    bstring_gzip_test = gzip.compress(np.transpose(data_array, (2, 1, 0)).tobytes("F"))
    assert gzip.decompress(bstring_gzip) == gzip.decompress(bstring_gzip_test)

    bstring_br = ngprecomputed.numpy_chunk(data_array, compression="br")
    bstring_br_test = brotli.compress(
        np.transpose(data_array, (2, 1, 0)).tobytes("F"), quality=6
    )
    assert brotli.decompress(bstring_br) == brotli.decompress(bstring_br_test)

#! failing because offset not aligned with the cube size
@pytest.mark.xfail
def test_save_obj():
    # chan_id:1005
    # col_id:51
    # digest:'89bb785630a9446b6a564c8779b3678d'
    # exp_id:174
    # mortonid:12282054
    # parent_iso:None
    # res:0
    # s3key:'89bb785630a9446b6a564c8779b3678d&51&174&1005&0&0&12282054&0'
    # t:0
    # version:0

    # mortonxyz.MortonXYZ(12282054)
    # [132, 117, 249]

    session = boto3.Session(profile_name="boss-s3")
    boss_s3resource = session.resource("s3")

    boss_s3Bucket = "cuboids.production.neurodata"
    boss_s3Key = "89bb785630a9446b6a564c8779b3678d&51&174&1005&0&0&12282054&0"
    dtype = "uint8"
    cube_size = 512, 512, 16
    offset = 0, 0, 2917

    basescale = 4, 4, 40
    dest_dataset = "bock11_test"
    dest_layer = "image_test"

    bosskey = bosslib.parts_from_bosskey(boss_s3Key)

    ngmorton = ngprecomputed.ngmorton(bosskey.mortonid, cube_size, offset)

    chunk_name = ngprecomputed.get_chunk_name(
        ngmorton, basescale, bosskey.res, cube_size, offset
    )
    ngkey = ngprecomputed.get_ng_key(dest_dataset, dest_layer, chunk_name)

    data = bosslib.get_boss_data(
        boss_s3resource, boss_s3Bucket, boss_s3Key, dtype, cube_size
    )

    # save it to a test folder
    ng_session = boto3.Session(profile_name="NGuser")
    ng_s3resource = ng_session.resource("s3")
    ng_s3Bucket = "nd-precomputed-volumes"

    resp = ngprecomputed.save_obj(ng_s3resource, ng_s3Bucket, ngkey, data)

    # attempt to read it back in?

    # examine dict to determine what params we can check

def test_get_scale_at_res():
    def test_res_scale(res_scales, iso=False):
        base_scale = res_scales[0][1]
        for res, scale in res_scales:
            scale_at_res = ngprecomputed.get_scale_at_res(base_scale, res, iso)
            assert all([str(s) == str(t) for s, t in zip(scale_at_res, scale)])

    res_scales = [(0, [2, 2, 50]), (1, [4, 4, 50])]
    test_res_scale(res_scales)

    res_scales = [(0, [3.8, 3.8, 50]), (1, [7.6, 7.6, 50]), (2, [15.2, 15.2, 50])]
    test_res_scale(res_scales)

    res_scales = [(0, [9.252, 741.5, 25]), (2, [37.008, 2966, 25])]
    test_res_scale(res_scales)

    res_scales = [(0, [9.252, 741.5, 25]), (2, [37.008, 2966, 25])]
    test_res_scale(res_scales)

    with raises(ValueError, match="voxel resolution is < 1"):
        ngprecomputed.get_scale_at_res([1, 0.5, 50], 1)

    res_scales = [(0, [20, 20, 20]), (1, [40, 40, 40]), (3, [160, 160, 160])]
    test_res_scale(res_scales, iso=True)

    scales = [
        [4, 4, 40],
        [8, 8, 40],
        [16, 16, 40],
        [32, 32, 40],
        [64, 64, 40],
        [128, 128, 40],
        [256, 256, 40],
        [512, 512, 40],
        [1024, 1024, 40],
        [2048, 2048, 40],
    ]

    for res in range(0, 10):
        assert scales[res] == ngprecomputed.get_scale_at_res(scales[0], res)

    for res in range(0, 10):
        assert scales[res][0:2] + [
            scales[res][2] * 2 ** res
        ] == ngprecomputed.get_scale_at_res(scales[0], res, iso=True)

def test_numpy_chunk_segmentation():
    data_array = np.random.randint(0, 10, (16, 512, 512), "uint64")

    bstring_br = ngprecomputed.numpy_chunk(data_array, compression="")

    block_size = (8, 8, 8)
    bstring_br_test = csegpy.encode_chunk(
        np.expand_dims(data_array, axis=0), block_size
    )
    assert bstring_br == bstring_br_test

@pytest.mark.xfail
def test_assume_role():
    # will need to specify the IAM to get this to pass
    # iam = "arn:aws:iam::222222222222:role/role-on-source-account"
    iam = ""

    session = boto3.Session(profile_name="ben-boss-dev")

    s3_write_resource = ngprecomputed.assume_role_resource(iam, session)

    bucket = "colm-data-from-boss"
    key = "test1"
    data = "a string of data"

    ngprecomputed.save_obj(
        s3_write_resource, bucket, key, data, storage_class="STANDARD", public=False
    )

max_stars_repo_name: neurodata/boss-export

max_issues_repo_name: neurodata/boss-export

max_forks_repo_name: neurodata/boss-export


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

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import shutil
import tempfile

import numpy as np

from tensorflow.python.keras._impl import keras
from tensorflow.python.platform import test
from tensorflow.python.training import training as training_module

try:
  import h5py  # pylint:disable=g-import-not-at-top
except ImportError:
  h5py = None

class TestModelSaving(test.TestCase):

  def test_sequential_model_saving(self):
    if h5py is None:
      return  # Skip test if models cannot be saved.

    with self.test_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(2, input_shape=(3,)))
      model.add(keras.layers.RepeatVector(3))
      model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))
      model.compile(loss=keras.losses.MSE,
                    optimizer=keras.optimizers.RMSprop(lr=0.0001),
                    metrics=[keras.metrics.categorical_accuracy],
                    sample_weight_mode='temporal')
      x = np.random.random((1, 3))
      y = np.random.random((1, 3, 3))
      model.train_on_batch(x, y)

      out = model.predict(x)
      _, fname = tempfile.mkstemp('.h5')
      keras.models.save_model(model, fname)

      new_model = keras.models.load_model(fname)
      os.remove(fname)

      out2 = new_model.predict(x)
      self.assertAllClose(out, out2, atol=1e-05)

      # test that new updates are the same with both models
      x = np.random.random((1, 3))
      y = np.random.random((1, 3, 3))
      model.train_on_batch(x, y)
      new_model.train_on_batch(x, y)
      out = model.predict(x)
      out2 = new_model.predict(x)
      self.assertAllClose(out, out2, atol=1e-05)

  def test_sequential_model_saving_2(self):
    if h5py is None:
      return  # Skip test if models cannot be saved.

    with self.test_session():
      # test with custom optimizer, loss

      class CustomOp(keras.optimizers.RMSprop):
        pass

      def custom_loss(y_true, y_pred):
        return keras.losses.mse(y_true, y_pred)

      model = keras.models.Sequential()
      model.add(keras.layers.Dense(2, input_shape=(3,)))
      model.add(keras.layers.Dense(3))
      model.compile(loss=custom_loss, optimizer=CustomOp(), metrics=['acc'])

      x = np.random.random((1, 3))
      y = np.random.random((1, 3))
      model.train_on_batch(x, y)

      out = model.predict(x)
      _, fname = tempfile.mkstemp('.h5')
      keras.models.save_model(model, fname)

      model = keras.models.load_model(
          fname,
          custom_objects={'CustomOp': CustomOp,
                          'custom_loss': custom_loss})
      os.remove(fname)

      out2 = model.predict(x)
      self.assertAllClose(out, out2, atol=1e-05)

  def test_functional_model_saving(self):
    if h5py is None:
      return  # Skip test if models cannot be saved.

    with self.test_session():
      inputs = keras.layers.Input(shape=(3,))
      x = keras.layers.Dense(2)(inputs)
      output = keras.layers.Dense(3)(x)

      model = keras.models.Model(inputs, output)
      model.compile(loss=keras.losses.MSE,
                    optimizer=keras.optimizers.RMSprop(lr=0.0001),
                    metrics=[keras.metrics.categorical_accuracy])
      x = np.random.random((1, 3))
      y = np.random.random((1, 3))
      model.train_on_batch(x, y)

      out = model.predict(x)
      _, fname = tempfile.mkstemp('.h5')
      keras.models.save_model(model, fname)

      model = keras.models.load_model(fname)
      os.remove(fname)

      out2 = model.predict(x)
      self.assertAllClose(out, out2, atol=1e-05)

  def test_saving_without_compilation(self):
    if h5py is None:
      return  # Skip test if models cannot be saved.

    with self.test_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(2, input_shape=(3,)))
      model.add(keras.layers.Dense(3))
      model.compile(loss='mse', optimizer='sgd', metrics=['acc'])

      _, fname = tempfile.mkstemp('.h5')
      keras.models.save_model(model, fname)
      model = keras.models.load_model(fname)
      os.remove(fname)

  def test_saving_with_tf_optimizer(self):
    if h5py is None:
      return  # Skip test if models cannot be saved.

    with self.test_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(2, input_shape=(3,)))
      model.add(keras.layers.Dense(3))
      model.compile(loss='mse',
                    optimizer=training_module.AdadeltaOptimizer(0.1),
                    metrics=['acc'])

      _, fname = tempfile.mkstemp('.h5')
      keras.models.save_model(model, fname)
      model = keras.models.load_model(fname)
      os.remove(fname)

  def test_saving_right_after_compilation(self):
    if h5py is None:
      return  # Skip test if models cannot be saved.

    with self.test_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(2, input_shape=(3,)))
      model.add(keras.layers.Dense(3))
      model.compile(loss='mse', optimizer='sgd', metrics=['acc'])
      model.model._make_train_function()

      _, fname = tempfile.mkstemp('.h5')
      keras.models.save_model(model, fname)
      model = keras.models.load_model(fname)
      os.remove(fname)

  def test_saving_lambda_numpy_array_arguments(self):
    if h5py is None:
      return  # Skip test if models cannot be saved.

    mean = np.random.random((4, 2, 3))
    std = np.abs(np.random.random((4, 2, 3))) + 1e-5
    inputs = keras.layers.Input(shape=(4, 2, 3))
    output = keras.layers.Lambda(lambda image, mu, std: (image - mu) / std,
                                 arguments={'mu': mean, 'std': std})(inputs)
    model = keras.models.Model(inputs, output)
    model.compile(loss='mse', optimizer='sgd', metrics=['acc'])

    _, fname = tempfile.mkstemp('.h5')
    keras.models.save_model(model, fname)

    model = keras.models.load_model(fname)
    os.remove(fname)

    self.assertAllClose(mean, model.layers[1].arguments['mu'])
    self.assertAllClose(std, model.layers[1].arguments['std'])

class TestSequential(test.TestCase):
  """Most Sequential model API tests are covered in `training_test.py`.
  """

  def test_basic_methods(self):
    model = keras.models.Sequential()
    model.add(keras.layers.Dense(1, input_dim=2))
    model.add(keras.layers.Dropout(0.3, name='dp'))
    model.add(keras.layers.Dense(2, kernel_regularizer='l2',
                                 kernel_constraint='max_norm'))
    model.build()
    self.assertEqual(model.state_updates, model.model.state_updates)
    self.assertEqual(model.get_layer(name='dp').name, 'dp')

  def test_sequential_pop(self):
    num_hidden = 5
    input_dim = 3
    batch_size = 5
    num_classes = 2
    with self.test_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(num_hidden, input_dim=input_dim))
      model.add(keras.layers.Dense(num_classes))
      model.compile(loss='mse', optimizer='sgd')
      x = np.random.random((batch_size, input_dim))
      y = np.random.random((batch_size, num_classes))
      model.fit(x, y, epochs=1)
      model.pop()
      self.assertEqual(len(model.layers), 1)
      self.assertEqual(model.output_shape, (None, num_hidden))
      model.compile(loss='mse', optimizer='sgd')
      y = np.random.random((batch_size, num_hidden))
      model.fit(x, y, epochs=1)

      # Test popping single-layer model
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(num_hidden, input_dim=input_dim))
      model.pop()
      self.assertEqual(len(model.layers), 0)
      self.assertEqual(len(model.outputs), 0)

      # Invalid use case
      model = keras.models.Sequential()
      with self.assertRaises(TypeError):
        model.pop()

  def test_sequential_weight_loading(self):
    if h5py is None:
      return

    temp_dir = self.get_temp_dir()
    self.addCleanup(shutil.rmtree, temp_dir)
    h5_path = os.path.join(temp_dir, 'test.h5')

    num_hidden = 5
    input_dim = 3
    batch_size = 5
    num_classes = 2

    with self.test_session():
      model = keras.models.Sequential()
      model.add(keras.layers.Dense(num_hidden, input_dim=input_dim))
      model.add(keras.layers.Dense(num_classes))

      x = np.random.random((batch_size, input_dim))
      ref_y = model.predict(x)

      model.save_weights(h5_path)

      model = keras.models.Sequential()
      model.add(keras.layers.Dense(num_hidden, input_dim=input_dim))
      model.add(keras.layers.Dense(num_classes))
      model.load_weights(h5_path)
      y = model.predict(x)

      self.assertAllClose(y, ref_y)

  def test_invalid_use_cases(self):
    with self.test_session():
      # Added objects must be layer instances
      with self.assertRaises(TypeError):
        model = keras.models.Sequential()
        model.add(None)

      # Added layers must have an inputs shape
      with self.assertRaises(ValueError):
        model = keras.models.Sequential()
        model.add(keras.layers.Dense(1))

      # Added layers cannot have multiple outputs
      class MyLayer(keras.layers.Layer):

        def call(self, inputs):
          return [3 * inputs, 2 * inputs]

        def _compute_output_shape(self, input_shape):
          return [input_shape, input_shape]

      with self.assertRaises(ValueError):
        model = keras.models.Sequential()
        model.add(MyLayer(input_shape=(3,)))
      with self.assertRaises(TypeError):
        model = keras.models.Sequential()
        model.add(keras.layers.Dense(1, input_dim=1))
        model.add(MyLayer())

      # Building empty model
      model = keras.models.Sequential()
      with self.assertRaises(TypeError):
        model.build()

  def test_nested_sequential_trainability(self):
    input_dim = 20
    num_units = 10
    num_classes = 2

    inner_model = keras.models.Sequential()
    inner_model.add(keras.layers.Dense(num_units, input_shape=(input_dim,)))

    model = keras.models.Sequential()
    model.add(inner_model)
    model.add(keras.layers.Dense(num_classes))

    self.assertEqual(len(model.trainable_weights), 4)
    inner_model.trainable = False
    self.assertEqual(len(model.trainable_weights), 2)
    inner_model.trainable = True
    self.assertEqual(len(model.trainable_weights), 4)

class TestModelCloning(test.TestCase):

  def test_clone_sequential_model(self):
    with self.test_session():
      val_a = np.random.random((10, 4))
      val_out = np.random.random((10, 4))

      model = keras.models.Sequential()
      model.add(keras.layers.Dense(4, input_shape=(4,)))
      model.add(keras.layers.Dropout(0.5))
      model.add(keras.layers.Dense(4))

    # Everything should work in a new session.
    keras.backend.clear_session()

    with self.test_session():
      # With placeholder creation
      new_model = keras.models.clone_model(model)
      new_model.compile('rmsprop', 'mse')
      new_model.train_on_batch(val_a, val_out)

      # On top of new tensor
      input_a = keras.Input(shape=(4,))
      new_model = keras.models.clone_model(
          model, input_tensors=input_a)
      new_model.compile('rmsprop', 'mse')
      new_model.train_on_batch(val_a, val_out)

      # On top of new, non-Keras tensor
      input_a = keras.backend.variable(val_a)
      new_model = keras.models.clone_model(
          model, input_tensors=input_a)
      new_model.compile('rmsprop', 'mse')
      new_model.train_on_batch(None, val_out)

  def test_clone_functional_model(self):
    with self.test_session():
      val_a = np.random.random((10, 4))
      val_b = np.random.random((10, 4))
      val_out = np.random.random((10, 4))

      input_a = keras.Input(shape=(4,))
      input_b = keras.Input(shape=(4,))
      dense_1 = keras.layers.Dense(4,)
      dense_2 = keras.layers.Dense(4,)

      x_a = dense_1(input_a)
      x_a = keras.layers.Dropout(0.5)(x_a)
      x_b = dense_1(input_b)
      x_a = dense_2(x_a)
      outputs = keras.layers.add([x_a, x_b])
      model = keras.models.Model([input_a, input_b], outputs)

    # Everything should work in a new session.
    keras.backend.clear_session()

    with self.test_session():
      # With placeholder creation
      new_model = keras.models.clone_model(model)
      new_model.compile('rmsprop', 'mse')
      new_model.train_on_batch([val_a, val_b], val_out)

      # On top of new tensors
      input_a = keras.Input(shape=(4,), name='a')
      input_b = keras.Input(shape=(4,), name='b')
      new_model = keras.models.clone_model(
          model, input_tensors=[input_a, input_b])
      new_model.compile('rmsprop', 'mse')
      new_model.train_on_batch([val_a, val_b], val_out)

      # On top of new, non-Keras tensors
      input_a = keras.backend.variable(val_a)
      input_b = keras.backend.variable(val_b)
      new_model = keras.models.clone_model(
          model, input_tensors=[input_a, input_b])
      new_model.compile('rmsprop', 'mse')
      new_model.train_on_batch(None, val_out)

  def test_model_cloning_invalid_use_cases(self):
    seq_model = keras.models.Sequential()
    seq_model.add(keras.layers.Dense(4, input_shape=(4,)))

    x = keras.Input((4,))
    y = keras.layers.Dense(4)(x)
    fn_model = keras.models.Model(x, y)

    with self.assertRaises(ValueError):
      keras.models._clone_functional_model(seq_model)
    with self.assertRaises(ValueError):
      keras.models._clone_functional_model(None)
    with self.assertRaises(ValueError):
      keras.models._clone_sequential_model(fn_model)

    with self.assertRaises(ValueError):
      keras.models._clone_sequential_model(seq_model, input_tensors=[x, x])
    with self.assertRaises(ValueError):
      keras.models._clone_sequential_model(seq_model, input_tensors=y)

if __name__ == '__main__':
  test.main()

max_stars_repo_name: yasunakacho/tensorflow

max_issues_repo_name: yasunakacho/tensorflow

max_forks_repo_name: yasunakacho/tensorflow


text:
import time
import os
from dateutil.parser import parse
import simplejson as json

import pandas as pd
import numpy as np

f = open('settings.json', 'r')
SETTINGS = json.loads(f.read().replace('\n', ''))
f.close()

def get_paths():
    """
    Redefine data_path and submissions_path here to run the benchmarks on your machine
    """
    data_path = SETTINGS['data_path']
    submission_path = SETTINGS['submission_path']
    return data_path, submission_path

def get_train_df(data_path = None):
    if data_path is None:
        data_path, submission_path = get_paths()

    train = pd.read_csv(os.path.join(data_path, SETTINGS['train']),
        converters={"saledate": parse})
    return train

def get_machines_df(data_path = None):
    if data_path is None:
        data_path, submission_path = get_paths()

    machines = pd.read_csv(os.path.join(data_path, SETTINGS['machines']))
    return machines

def get_test_df(data_path = None):
    if data_path is None:
        data_path, submission_path = get_paths()

    test = pd.read_csv(os.path.join(data_path, SETTINGS['test']),
        converters={"saledate": parse})
    return test

def get_train_test_df(data_path = None):
    return get_train_df(data_path), get_test_df(data_path)

def write_submission(submission_name, predictions, \
        IDs=None, submission_path=None):

    if submission_path is None:
        data_path, submission_path = get_paths()
    if IDs is None:
        test = get_test_df()
        test = test.join(pd.DataFrame({"SalePrice": predictions}))
    else:
        test = pd.DataFrame({'SalesID': IDs, 'SalePrice': predictions})

    test[["SalesID", "SalePrice"]].to_csv(os.path.join(submission_path,
            submission_name), index=False)

def get_rmse(Y1, Y2):
    Y1 = np.array(Y1)
    Y2 = np.array(Y2)
    sq = (Y1 - Y2) ** 2
    return np.sqrt(np.mean(sq))

class Logger():
    def __init__(self, count, step, tag=None):
        self.count = count
        self.step_size = step
        self.time = time.time()
        self.tag = tag

    def step(self):
        self.count -= 1
        if self.count % self.step_size == 0:
            s = '%s - %s' % (self.count, time.time() - self.time)
            if self.tag:
                s = '[%s] %s' % (self.tag, s)
            print(s)
            self.time = time.time()

max_stars_repo_name: Keesiu/meta-kaggle

max_issues_repo_name: Keesiu/meta-kaggle

max_forks_repo_name: Keesiu/meta-kaggle


text:
#!/usr/bin/env python
#
# Copyright 2019 DFKI GmbH.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the
# following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.
import numpy as np
from PySignal import Signal
from ..graphics.renderer.lines import DebugLineRenderer
from .animation_controller import AnimationController
from ..graphics.renderer import SphereRenderer
from ..scene.components import ComponentBase
from ..graphics.utils import get_translation_matrix
from ..io import load_json_file

DEFAULT_COLOR = [0, 0, 1]

class PointCloudAnimationController(ComponentBase, AnimationController):
    updated_animation_frame = Signal()
    reached_end_of_animation = Signal()

    def __init__(self, scene_object, color=DEFAULT_COLOR, visualize=True):
        ComponentBase.__init__(self, scene_object)
        AnimationController.__init__(self)
        self.animated_joints = []
        self.visualize = visualize
        self.skeleton = None
        if visualize:
            self._sphere = SphereRenderer(10, 10, 1, color=color)
            a = [0, 0, 0]
            b = [1, 0, 0]
            self._line = DebugLineRenderer(a, b, color)
        self.draw_bone = False
        self._semantic_annotation = None
        self.frameTime = 1.0/30
        self.motion_data = []
        self.skeleton_model = None
        self.target_skeleton = None

    def toggle_animation_loop(self):
        self.loopAnimation = not self.loopAnimation

    def isLoadedCorrectly(self):
        return len(self.motion_data) > 0

    def set_data(self, data):
        self.motion_data = []
        for idx, frame in enumerate(data["motion_data"]):
            pc_frame = []
            for p in frame:
                pc_frame.append(list(map(float, p)))
            self.motion_data.append(pc_frame)
        self.motion_data = np.array(self.motion_data)

        if 'skeleton' in list(data.keys()):
            self.skeleton = data["skeleton"]
        else:
            self.skeleton = None

        if 'has_skeleton' in list(data.keys()) and 'skeleton' in list(data.keys()):
            self.draw_bone = False#data['has_skeleton']

    def update(self, dt):
        dt *= self.animationSpeed
        if self.isLoadedCorrectly():
            if self.playAnimation:
                self.animationTime += dt
                self.currentFrameNumber = int(self.animationTime / self.getFrameTime())

                # update gui
                if self.currentFrameNumber > self.getNumberOfFrames():
                    self.resetAnimationTime()
                else:
                    self.updated_animation_frame.emit(self.currentFrameNumber)

    def draw(self, modelMatrix, viewMatrix, projectionMatrix, lightSources):
        if self._sphere is None:
                return

        if self.currentFrameNumber < 0 or self.currentFrameNumber >= self.getNumberOfFrames():
            return

        for position in self.motion_data[self.currentFrameNumber]:
            m = get_translation_matrix(position[:3])
            m = np.dot(m, modelMatrix)
            self._sphere.draw(m, viewMatrix, projectionMatrix, lightSources)

        if self.draw_bone:
            for joint, value in list(self.skeleton.items()):
                if value['parent'] is not None:
                    joint_idx = value['index']
                    joint_parent_idx = self.skeleton[value['parent']]['index']
                    start_point = self.motion_data[self.currentFrameNumber][joint_parent_idx]
                    end_point = self.motion_data[self.currentFrameNumber][joint_idx]
                    self._line.set_line(start_point, end_point)
                    self._line.draw(modelMatrix, viewMatrix, projectionMatrix)

    def getNumberOfFrames(self):
        return len(self.motion_data)

    def getFrameTime(self):
        return self.frameTime

    def updateTransformation(self, frame_number=None):
        if frame_number is not None:
            self.currentFrameNumber = frame_number
        self.animationTime = self.getFrameTime() * self.currentFrameNumber

    def setCurrentFrameNumber(self, frame_number=None):
        if frame_number is not None:
            self.currentFrameNumber = frame_number
        self.animationTime = self.getFrameTime() * self.currentFrameNumber

    def apply_scale(self, scale):
        self.motion_data[:, :, :3] *= scale

    def apply_transform(self, m):
        for i, frame in enumerate(self.motion_data):
            for j, p in enumerate(frame):
                p = self.motion_data[i, j, :3]
                self.motion_data[i, j, :3] = np.dot(m, p)

    def setColor(self, color):
        self._sphere.technique.material.diffuse_color = color
        self._sphere.technique.material.ambient_color = color * 0.1

    def getColor(self):
        return self._sphere.technique.material.diffuse_color

    def replace_skeleton_model(self, filename):
        data = load_json_file(filename)
        self.set_skeleton_model(data)

    def set_skeleton_model(self, model):
        self.skeleton_model = model
        #self.skeleton = SkeletonBuilder().load_from_json_data(data)
        #self.skeleton.joints = self._joints

    def get_semantic_annotation(self):
        return dict()

    def get_current_frame(self):
        if self.currentFrameNumber < 0 or self.currentFrameNumber >= self.getNumberOfFrames():
            return None
        return self.motion_data[self.currentFrameNumber]

    def get_skeleton(self):
        return self.skeleton

    def get_frame_time(self):
        return self.getFrameTime()

    def get_label_color_map(self):
        return dict()

    def set_frame_time(self, frame_time):
        self.frameTime = frame_time


text:
"""
Control Amplitude Noise
=======================

This example demonstrates how to add Gaussian noise to the control pulse.
"""
import numpy as np
import matplotlib.pyplot as plt
from qutip.qip.device import Processor
from qutip.qip.noise import RandomNoise
from qutip.operators import sigmaz, sigmay

# add control Hamiltonians
processor = Processor(N=1)
processor.add_control(sigmaz(), targets=0)

# define pulse coefficients and tlist for all pulses
processor.pulses[0].coeff = np.array([0.3, 0.5, 0. ])
processor.set_all_tlist(np.array([0., np.pi/2., 2*np.pi/2, 3*np.pi/2]))

# define noise, loc and scale are keyword arguments for np.random.normal
gaussnoise = RandomNoise(
            dt=0.01, rand_gen=np.random.normal, loc=0.00, scale=0.02)
processor.add_noise(gaussnoise)

# Plot the ideal pulse
processor.plot_pulses(title="Original control amplitude", figsize=(5,3))

# Plot the noisy pulse
qobjevo, _ = processor.get_qobjevo(noisy=True)
noisy_coeff = qobjevo.to_list()[1][1] + qobjevo.to_list()[2][1]
fig2, ax2 = processor.plot_pulses(title="Noisy control amplitude", figsize=(5,3))
ax2[0].step(qobjevo.tlist, noisy_coeff)

max_stars_repo_name: xinyufei/Quantum-Control-qutip

max_issues_repo_name: xinyufei/Quantum-Control-qutip

max_forks_repo_name: xinyufei/Quantum-Control-qutip


text:
import torch
import torch.nn as nn

import numpy as np
from collections import defaultdict
from functools import partial
import copy
import logging

from flopco_keras.compute_layer_flops import *

log=logging.getLogger(__name__)
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)
log.addHandler(ch)

class FlopCoKeras():

    def __init__(self, model):
        '''
        instances: list of layer types,
            supported types are [nn.Conv2d, nn.Linear,
            nn.BatchNorm2d, nn.ReLU, nn.MaxPool2d, nn.AvgPool2d, nn.Softmax]
        '''
        self.model = model

        self.flops = []
        self.macs = []
        self.get_flops = {
                    'ReLU': compute_relu_flops,
                    'InputLayer': compute_input_flops,
                    'Conv2D': compute_conv2d_flops,
                    'ZeroPadding2D': compute_padding_flops,
                    'Activation': compute_activation_flops,
                    'Dense': compute_fc_flops,
                    'BatchNormalization': compute_bn2d_flops,
                    'TensorFlowOpLayer': compute_tfop_flops,
                    'MaxPooling2D': compute_pool2d_flops,
                    'Add': compute_add_flops,
                    'Flatten': compute_flatten_flops,
                    'Dropout': compute_dropout_flops,
                    'GlobalAveragePooling2D': compute_globalavgpool2d_flops,
                    'GlobalMaxPooling2D': compute_globalmaxpool2d_flops,
                    'DepthwiseConv2D': compute_depthwiseconv2d_flops,
                    'Reshape': compute_reshape_flops,
        }

        self.total_flops =None
        self.total_macs = None
        self.relative_flops = None
        self.relative_macs = None

        try:
            self.get_stats( flops = True, macs = True)
        except Exception as e:
            log.error(f'Could not compute flops: caught {e} (probably cannot compute this type of layer)')
            del self.model
            return
        self.total_flops = sum(self.flops)
        self.total_macs = sum(self.macs)
        # self.total_params = sum(self.params) #TO DO

        self.relative_flops = [k/self.total_flops for k in self.flops]

        self.relative_macs = [k/self.total_macs for k in self.macs]

        # self.relative_params = [k/self.total_params for k in self.params] #TO DO

        del self.model

    def __str__(self):
        print_info = "\n".join([str({k:v}) for k,v in self.__dict__.items()])

        return str(self.__class__) + ": \n" + print_info

    # def count_params(self):
    #     self.params = [0]
        # self.params = defaultdict(int)

        # for mname, m in self.model.named_modules():
        #     if m.__class__ in self.instances:

        #         self.params[mname] = 0

        #         for p in m.parameters():
        #             self.params[mname] += p.numel()

    def _save_flops(self, layer, macs=False):
        """ Computes flops for a layer and optionally saves the MACs for that layer

        :param layer: the layer
        :param macs: set True to save layer MACs to self.flops list
        """
        flops = self.get_flops[layer.__class__.__name__](layer, macs)
        if macs:
            self.macs.append(flops)
        else:
            self.flops.append(flops)

    def get_stats(self, flops = False, macs = False):
        """ Gets stats for a layer and optionally saves the MACs for that layer
        Use it by constucting f=FlopCo(model), calling f.get_stats(), and then accessing f.total_flops, etc

         :param layer: the layer
         :param macs: set True to save layer MACs to self.flops list

         :returns: None
         """

        # if params:
        #     self.count_params()

        if flops:
            self.flops = []

        if macs:
            self.macs = []

        layers = None
        if hasattr(self.model, 'layers'):
            log.info('found model.layers')
            layers = self.model.layers
        elif isinstance(self.model, tuple) and hasattr(self.model[0], 'layers'):
            log.info('found model.layers')
            layers = self.model[0].layers
        if flops:
            for layer in layers:
                self._save_flops(layer)
        if macs:
            for layer in layers:
                self._save_flops(layer, macs=True)

max_stars_repo_name: tobidelbruck/flopco-keras

max_issues_repo_name: tobidelbruck/flopco-keras

max_forks_repo_name: tobidelbruck/flopco-keras