How to fill 3D array in numpy with functional values at the same speed as Java?

Please do not use .vectorize(..), it is not efficient since it will still do the processing at the Python level. .vectorize() should only be used as a last resort if for example the function can not be calculated in "bulk" because its "structure" is too complex.

But you do not need to use .vectorize here, you can implement your function to work over arrays with:

r_xy = np.sqrt(xx ** 2 + yy ** 2)
density = (r_xy <= R0) & (-hz <= zz) & (zz <= hz)

or even a bit faster:

r_xy = xx * xx + yy * yy
density = (r_xy <= R0 * R0) & (-hz <= zz) & (zz <= hz)

This will construct a 2000×2000×20 array of booleans. We can use:

intdens = density.astype(int)

to construct an array of ints.

Printing the array here is quite combersome, but it contains a total of 2'356'047 ones:

>>> density.astype(int).sum()
2356047

Benchmarks: If I run this locally 10 times, I get:

>>> timeit(f, number=10)
18.040479518999973
>>> timeit(f2, number=10)  # f2 is the optimized variant
13.287886952000008

So on average, we calculate this matrix (including casting it to ints) in 1.3-1.8 seconds.

You can also use a compiled version of the function to calculate density. You can use cython or numba for that. I use numba to jit compile the density calculation function in the ans, as it is as easy as putting in a decorator.

Pros :

• You can write if conditions as you mention in your comments
• Slightly faster than the numpy version mentioned in the ans by @Willem Van Onsem, as we have to iterate through the boolean array to calculate the sum in density.astype(int).sum().

Cons:

• Write an ugly three level loop. Looses the beauty of the singlish liner numpy solution.

Code:

import numba as nb
@nb.jit(nopython=True, cache=True)
def calc_density(xx, yy, zz, R0, hz):
    threshold = R0 * R0
    dimensions = xx.shape

    density = 0
    for i in range(dimensions[0]):
        for j in range(dimensions[1]):
            for k in range(dimensions[2]):
                r_xy = xx[i][j][k] ** 2 + yy[i][j][k] ** 2
    
                if(r_xy <= threshold and -hz <= zz[i][j][k] <= hz):
                    density+=1
    return density

Running times:

Willem Van Onsem solution, f2 variant : 1.28s without sum, 2.01 with sum.

Numba solution( calc_density, on second run, to discount the compile time) : 0.48s.

As suggested in the comments, we need not calculate the meshgrid also. We can directly pass the x, y, z to the function. Thus:

@nb.jit(nopython=True, cache=True)
def calc_density2(x, y, z, R0, hz):
    threshold = R0 * R0
    dimensions = len(x), len(y), len(z)

    density = 0
    for i in range(dimensions[0]):
        for j in range(dimensions[1]):
            for k in range(dimensions[2]):
                
                r_xy = x[i] ** 2 + y[j] ** 2
                if(r_xy <= threshold and -hz <= z[k] <= hz):
                    density+=1
    return density

Now, for fair comparison, we also include the time of np.meshgrid in @Willem Van Onsem's ans. Running times:

Willem Van Onsem solution, f2 variant(np.meshgrid time included) : 2.24s

Numba solution( calc_density2, on second run, to discount the compile time) : 0.079s.

This is meant as a lengthy comment on the answer of Deepak Saini.

The main change is to not use the coordinates generated by np.meshgrid which contains unecessary repetitions. This isn't recommandable if you can avoid it (both in terms of memory usage and performance)

Code

import numba as nb
import numpy as np

@nb.jit(nopython=True,parallel=True)
def calc_density_2(x, y, z,R0,hz):
    threshold = R0 * R0

    density = 0
    for i in nb.prange(y.shape[0]):
        for j in range(x.shape[0]):
            r_xy = x[j] ** 2 + y[i] ** 2
            for k in range(z.shape[0]):
                if(r_xy <= threshold and -hz <= z[k] <= hz):
                    density+=1

    return density

Timings

R0 = 500
hz = 1

x = np.arange(-1000, 1000, 1)
y = np.arange(-1000, 1000, 1)
z = np.arange(-10, 10, 1)

xx, yy, zz = np.meshgrid(x, y, z)

#after the first call (compilation overhead)
#calc_density_2          9.7 ms
#calc_density_2 parallel 3.9 ms
#@Deepak Saini           115 ms