How can I add new dimension to a Numpy array to perform multithreading?

Viewed 45

Let's keep it simple: calculation of a*x+b for many x values and two set of parameters a and b.

import numpy as np

def myFoo1 (par,x):
    return par[:,0,None]*x+par[:,1,None]

x=np.array(np.arange(0,10,2.0))
par=np.array([[i*10.0,i*4.0] for i in range(1,3)])
myFoo1(par,x)

I need to add an extra dimension (par[:,0,None] because par[:,0] returns a 1D array). This code works.

But, when I tried to do the same thing with @jit from numba, I could not do it.

First, I tested 3 methods to add an extra dimension, only one works. I do not why. Any explanation would be welcomed even if it is not really the question.

@jit(('float64[:,:], float64[:]'), nopython=True)
def myFoo2(par,x):
    # print(par[:,0,None]) # It does not compile
    #print(par[:,0][:,None]) # It does not compile
    print(par[:,0:1]) # It compiles and return a 2D array
    return par[:,0:1]*x+par[:,1:2]

x=np.array(np.arange(0,10,2.0))
par=np.array([[i*10.0,i*4.0] for i in range(1,3)])
myFoo2(par,x)

I thought I had found the solution to do multithreading. But now, I am facing a new problem. It does not calculate. With the following code,

@jit(('float64[:,:], float64[:]'), nopython=True, parallel=True)
def myFoo3(par,x):
    length=x.shape[0]
    numthreads=2
    result = np.empty((par.shape[0],x.shape[0]), dtype=np.float64)
    chunklen = (length + numthreads - 1) // numthreads    
    for i in prange(numthreads):
        # It seems good 
        print(i,":",par[:,0:1],x[i * chunklen:(i + 1) * chunklen])
        # But it does not calculate
        print(i,":",par[:,0:1]*x[i * chunklen:(i + 1) * chunklen]) 
    return result

x=np.array(np.arange(0,10,2.0))
par=np.array([[i*10.0,i*4.0] for i in range(1,3)])
myFoo3(par,x)

I get an Assertion Error.

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
/tmp/ipykernel_64043/2473996700.py in <module>
     12 x=np.array(np.arange(0,10,2.0))
     13 par=np.array([[i*10.0,i*4.0] for i in range(1,3)])
---> 14 myFoo3(par,x)

AssertionError: Sizes of $148binary_subscr.42, $170binary_subscr.54 do not match on /tmp/ipykernel_64043/2473996700.py (9)

I do not understand where this problem comes from. Any help would be appreciated.

0 Answers
Related