Find indices using for loop

I should estimate the effect of B fields on particle deflection using $$\F=V \times B \times Force\$$.

I have a 2D mesh with $$\n_x\$$ and $$\n_y\$$ which show the number of points in each direction. $$\B\$$ is the magnetic field which is a 2D mesh arrays (with dimension $$\n_x \times n_y= 1600 \times 900 = 1440000\$$) and $$\V\$$ is the velocity of particles and is a 1D array (with dimension 431605462). I also have two 1D arrays (with the same dimension as the $$\V\$$ array) which show the x and y coordinates of each particles. I called them grid_pro_x and grid_pro_y.

Since $$\V\$$ is a 1D array and $$\B\$$ is a 2D array I can not directly multiply $$\V\$$ with $$\B\$$. To solve this issue I created a 2D zero array for the $$\V\$$ and using 2 for-loops I find particles which stand in each cell of the mesh. I consider the average velocity of all particles as a velocity component of each cell. Because of the large dimensionality of the mesh, if I keep the original dimension, the code needs several days to complete the $$\V\$$ array since it needs to be repeated 1440000 times. As a solution, I compressed $$\V\$$ with an optional variable called divider. If I use divider = 10 I guess the code needs 2 days to be finished and if I use divider = 100 the code will be finished in half an hour but now the precision of the results reduces significantly.

You can see the code in the following.

nx,ny= 1600, 900

xmin = -30.0e-6
xmax =  50.0e-6

ymin = -45.0e-6
ymax = 45.0e-6

#### create a 2D  velocity array

divider = 100
m_p= 1.6726219e-27   #  (kg unit)  proton mass

nx_new = nx/divider
ny_new = ny/divider

Vx = np.zeros((ny_new, nx_new))
Vy = np.zeros((ny_new, nx_new))

res_x = (xmax - xmin)/nx_new   # resolution in x direction
res_y = (ymax - ymin)/ny_new   # resolution in y direction

#   Find corresponding indices

for i in range( 0,ny_new ):
for j in range(  0,nx_new):

xx = i*res_x
yy = j*res_y

##### all particles with different Px confined in a specific cell

Px_pro_lim = Px_pro[(grid_pro_x >=  xx ) & (grid_pro_x < xx+res_x ) & (grid_pro_y >=  yy ) & (grid_pro_y < yy+res_y ) ]

Vx[i,j] = np.mean(Px_pro_lim)/m_p
#print 'Vx[i,j]= ' ,Vx[i,j]

• I think an example or something would go a long way in explaining what your code does. Aug 3 '19 at 23:34
• As it sits, your code will throw NameErrors against Px_pro, grid_pro_x, and gird_pro_y. Can you provide context into what these are? Oct 3 '19 at 18:31

If you are trying to make it fast, the last thing to do is drop into Python for-loops. Keep it all in numpy as much as possible.

The code below maps the particles to the mesh and calculates the mean v for all the particles in each mesh square. With 400M particles having random x,y, and v on a 1600x900 size mesh, it runs in about 100 seconds on my laptop.

xi = np.digitize(x, np.linspace(xmin, xmax, nx, endpoint=False)) - 1
yi = np.digitize(y, np.linspace(ymin, ymax, ny, endpoint=False)) - 1

flat_index = nx*yi + xi

counts = np.bincount(flat_index)

vsums = np.bincount(flat_index, weights=v)

means = (vsums / counts).reshape(nx,ny)


np.digitize(x, bins) maps values the values in the array x, to ints based on the values in bins. x < bin[0] maps to 0; bin[0] <= x < bin[1] maps to 1; and so on. The endpoints=False makes sure the indices go from 1 to 1440000 instead of 1 to 143999; and the -1 shifts it to 0 to 1439999.

np.linspace(start, end, nsteps) divides the range [start, end) in to nstep even intervals. This is used to define the bins for np.digitize().

nx*yi + xi basically converts the x and y indices to a 1D index from 0 to 1439999, because bincount only works on 1D arrays

np.bincount(z, weights) if weights is omitted, this counts the number of occurrences of each integer in z (like a histogram). If provided, weights is an array of the same dim as z, the values of the corresponding weights are summed instead of just counting the integers. So setting weights to v, sums the velocity of each particle in a mesh square.

(sums / counts).reshape(nx,ny) calculates the means and reshapes the result into the 1600x900 mesh.