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.