3

What would be the fastest and most pythonic way to perform element-wise operations of arrays of different size without oversampling the smaller array?

For example: I have a large array, A 1000x1000 and a small array B 10x10 I want each element in B to respond to 100x100 elements in array B. There is no need for any interpolation, just to use the same element in B for all 10000 operations in A.

I can adjust the size of the two arrays to make the shape of A is a multiple of B. Typically they will be 1:10000 or 1:1000 in all dimensions. The arrays represents data samples with different resolution but same extent.

I know that I can oversample array B, e.g. by using Kronecker product, but It would be nicer to keep array B small, especially as some of my arrays get really big to handle and store. I'm using xarray and dask, but any numpy operation would also work.

I hope this snippet explains what I want to do:

import numpy as np

A = np.random.rand(10,10)
B = np.random.rand(1000,1000)

res = np.shape(B)[0]//np.shape(A)[0]

#I want to add A and B so that each element in A is added to 100x100 elements in B.
#This doesn't work of obvious reasons:
#C = A+B

#This solution sacrifices the resolution of B:
C = A+B[::res,::res]

#These solutions creates an unnecessary large array for the operation(don't they?):
K = np.ones((res,res))
C = np.kron(A, K) + B

C = np.repeat(np.repeat(A,res, axis=0), res, axis=1)+B

I have a feeling that this problem must have been up before, but I couldn't find any answer that works for this particular case.

user2821
  • 1,568
  • 2
  • 12
  • 16
  • 1
    Seems relevant to [`How can I check if one two-dimensional NumPy array contains a specific pattern of values inside it?`](https://stackoverflow.com/questions/32531377/). – Divakar Aug 12 '18 at 13:54

2 Answers2

3

I had a similar problem, and ended up solving it like this :

import numpy as np
import numba as nb
import time

@nb.jit(nb.void(nb.float64[:,:], nb.float64[:,:]), nopython=True, cache=True)
def func(A, B):
    # Assume different resolution along the different axes
    res_row = B.shape[0]//A.shape[0]
    res_col = B.shape[1]//A.shape[1]

    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            # Example to add, can be changed to any other function
            B[i * res_row : (i+1) * res_row, j * res_col : (j+1) * res_col] += A[i,j]


a = np.random.rand(1000,1000)
b = np.random.rand(10000,10000)

t1 = time.time()
func(a,b)
print("Time taken = {time_taken}".format(time_taken = time.time() - t1))
# with jitting on 8GB, 8 cores MacBook Pro = 0.45
# without jitting = 5.46

Though, numba has nothing to do with this question, but this is the approach I took to the problem, with significant performance gain with using JIT compiling.

Deepak Saini
  • 2,810
  • 1
  • 19
  • 26
3

With broadcasting we can 'replicate' a small array to work with a larger one. For example, if

a = np.random.rand(10)
b = np.random.rand(1000).reshape(10,100)
a[:,None]+b

The trick here is to pair each element of A with a (100,100) block of B. To do that we need to reshape and transpose.

In [3]: A = np.random.rand(10,10)
   ...: B = np.random.rand(1000,1000)
   ...: res = np.shape(B)[0]//np.shape(A)[0]

Your target:

In [4]: K = np.ones((res,res))
   ...: C = np.kron(A, K) + B
   ...: 
   ...: C = np.repeat(np.repeat(A,res, axis=0), res, axis=1)+B

In [5]: C.shape
Out[5]: (1000, 1000)

Divide B into these (100,100) chunks:

In [7]: B1 = B.reshape(10,100,10,100).transpose(0,2,1,3)

In [8]: B1.shape
Out[8]: (10, 10, 100, 100)

Now we can add with broadcasting

In [9]: C1 = B1 + A[:,:,None,None]

In [10]: C1.shape
Out[10]: (10, 10, 100, 100)

Reshape back to original B shape:

In [11]: C1=C1.transpose(0,2,1,3).reshape(B.shape)

They match:

In [12]: np.allclose(C,C1)
Out[12]: True
hpaulj
  • 221,503
  • 14
  • 230
  • 353