lax.scan over 3D array

[dui1234]

I am tring to calculate a spatial average over a set of information.

The information I want to calcualte contains 10^5 sets of information of 21 numbers. It comes with laber on a spatial space ijk. Both data are shown below.

The spatial space I talked about is a huge cubic box with dimention Lx x Ly x Lz, i.e. 25 x 26 x 2**5 in the below figure, and now you slice it to a cube with dimension 1x1x1. I then use a vanilar for loop to perform a summation of each 21 number from each set in each i,j,k label in order to get the average of each 1x1x1 cube, I tried only for the first 5000 sets of information:

stresses = np.zeros((2**5, 2**6, 2**5 ,21))
count = np.zeros((2**5, 2**6, 2**5))
num = 0
for i,j,k in ijk[:5000]:
    for c in range(21):
        stresses[i,j,k,c] += S_t[num][c]
    count[i,j,k] += 1
    num+=1
    
for i,j,k in ijk[:5000]:
    for c in range(21): stresses[i,j,k,c]/= count[i,j,k]

This takes so much time (even for just first 5000 points), I thus, tried to use the lax.scan instead but it shows the following error.

stresses_m = np.zeros((2**5, 2**6, 2**5 ,21))
count_m = np.zeros((2**5, 2**6, 2**5))
def lax_stress(s_c,ijk_sr):
    stresses, count = s_c
    ijk,sr,num = ijk_sr
    for c in range(21): 
        stresses[ijk[0],ijk[1],ijk[2],c] += sr[num][c]
    count[i,j,k] += 1
    s_c = (stresses, count)
    return s_c,s_c`

`n = len(ijk[:5000])
num = np.arange(n)
ijk_sr = (ijk[:5000],S_t[:5000])
s_ci = (stresses_m,count_m)
final, result = lax.scan(lax_stress, init=s_ci, xs=num)

How should I fix it?

[jakevdp]

We should be able to help you with this, but can you edit your question to paste code as text within triple-tick code blocks rather than including code by screenshot? It helps to be able to copy-paste your for-loop rather than having to transcribe it from the image. Thanks!

[dui1234]

Thank you for your prompt reply. So sorry for causing the inconvenience. I have just edited the post including the code blocks just below each images. Very much obliged for your help.

[hericks]

0. Setup. I used the following code as setup.

   import jax
   import jax.numpy as jnp
   import numpy as np
   
   key= jax.random.PRNGKey(0)
   key, key_ijk, key_S_t = jax.random.split(key, 3)
   
   ijk = jax.random.randint(key_ijk, (159600, 3), 0, 2**5)
   S_t = jax.random.normal(key_S_t, (159600, 21))
   

1. Fixing bug in provided solution. I believe that your code does not solve the problem that you'd actually like to solve, i.e. computing a mean. For this, note that when a tuple of indices i, j, k appears multiple times in the array ijk, then the entries in stresses[i, j, k, :] are divided by count[i, j, k] multiple times - once for each occurrence of the indices i, j, k in ijk.

   This can be fixed dividing all non-zero entries in stresses exactly once.

   n = 500
   num = 0
   for i,j,k in ijk[:n]:
       for c in range(21):
           stresses[i,j,k,c] += S_t[num][c]
       count[i,j,k] += 1
       num+=1
   
   for (i, j, k) in np.argwhere(count):
       for c in range(21):
           stresses[i,j,k,c] /= count[i,j,k]
   

2. Exploiting NumPy vectorisation. Next, we can optimise the code by using the vectorisation capabilities of NumPy to eliminate the inner loop over the last array axis. On my machine, this already yields a 10x speed-up.

   for idx in range(n):
       i, j, k = ijk[idx]
       stresses[i,j,k,:] += S_t[idx,:]
       count[i,j,k] += 1
   
   for (i, j, k) in np.argwhere(count):
       stresses[i,j,k,:] /= count[i,j,k]
   

3. Rewrite in JAX using jax.lax.fori_loop. Finally, this code can be rewritten in JAX using jax.lax.fori_loop. Again, this gives a further 3.6x speed-up on my machine.

   def accumulate_fun(idx, val):
       stresses, count, ijk, S_t = val
       i, j, k = ijk.at[idx,:].get()
       stresses = stresses.at[i,j,k,:].add(S_t.at[idx,:].get())
       count = count.at[i,j,k].add(1.)
       return (stresses, count, ijk, S_t)
   
   stresses, count, _, _ = jax.lax.fori_loop(0, n, accumulate_fun, (stresses, count, ijk, S_t))
   
   def divide_fun(idx, val):
       stresses, count, indices = val
       i, j, k = indices.at[idx,:].get()
       stresses = stresses.at[i,j,k,:].divide(count.at[i,j,k].get())
       return (stresses, count, indices)
   
   indices = jnp.argwhere(count)
   stresses, _, _ = jax.lax.fori_loop(0, indices.shape[0], divide_fun, (stresses, count, indices))
   

   I am sure that you can even wrap this block into a jittable function. For this, you would make sure that jnp.argwhere(count) is of a size only depending on the input shapes - take a look at the size and fill_value arguments of jnp.argwhere. In this case, you may need a jax.lax.cond in divide_fun to check for the fill_value.

   Still, the given JAX solution for me already runs in <10s using the full array.

4. Edit. Rewrite in JAX using jax.lax.scan.

   stresses = np.zeros((2**5, 2**6, 2**5, 21))
   count = np.zeros((2**5, 2**6, 2**5))
   
   def fun_accumulate(carry, ijk):
       (i, j, k) = ijk
       stresses, count, S_t, idx = carry
       stresses = stresses.at[i,j,k,:].add(S_t.at[idx,:].get())
       count = count.at[i,j,k].add(1.)
       return (stresses, count, S_t, idx + 1), 0
   
   (stresses, count, _, _), _ = jax.lax.scan(fun_accumulate, init=(stresses, count, S_t, 0), xs=ijk[:n])
   
   def fun_divide(carry, ijk):
       (i, j, k) = ijk
       stresses, count = carry
       stresses = stresses.at[i,j,k,:].divide(count.at[i,j,k].get())
       return (stresses, count), 0
   
   (stresses, count), _ = jax.lax.scan(fun_divide, init=(stresses, count), xs=jnp.unique(ijk[:n], axis=0))
   
   

[dui1234]

That's really splendid! Thank you for these solution.