Correct usage of numpy.vstack with python’s numba

I am currently trying to speed up some python code using numba. According to the documentation of numba, numpy.meshgrid is not supported but numpy.vstack is. So I replaced the meshgrid call by vstack which works fine when not using numba. It does not work, however, when using numba.

Here is the code:

import numpy as np
import timeit
from numba import njit, prange

@njit(parallel=True)
def method1_par1(n):
    # create some data
    x   = np.linspace(0,2*np.pi,n)
    y   = np.linspace(0,2*np.pi,n)
    # np.meshgrid not supported by numba, therefore trying to use np.vstack
    X   = np.vstack( n*[x] )
    Y   = np.hstack( n*[y[:,np.newaxis]] )    # tranform row vector into column vector, then duplicate
    Z   = np.sin(X) * np.sin(Y)

    # calculate centered finite difference using for loop
    Z_diff1 = Z*.0
    for ii in prange(1,Z.shape[0]-1,2):
        for jj in range(1,Z.shape[1]-1,2):
            Z_diff1[ii,jj]  = Z[ii+1, jj] - Z[ii-1,jj]

runs    = 1000
print( min(timeit.repeat( "method1_par1(50)", "from __main__ import method1_par1",
           number=runs )) )

And here is the error message:

No implementation of function Function(<built-in function getitem>) found for signature:

>>> getitem(array(float64, 1d, C), Tuple(slice<a:b>, none))

There are 22 candidate implementations:
  - Of which 20 did not match due to:
  Overload of function 'getitem': File: <numerous>: Line N/A.
    With argument(s): '(array(float64, 1d, C), Tuple(slice<a:b>, none))':
   No match.
  - Of which 2 did not match due to:
  Overload in function 'GetItemBuffer.generic': File: numba/core/typing/arraydecl.py: Line 162.
    With argument(s): '(array(float64, 1d, C), Tuple(slice<a:b>, none))':
  Rejected as the implementation raised a specific error:
    TypeError: unsupported array index type none in Tuple(slice<a:b>, none)
 raised from /usr/local/lib/python3.8/dist-packages/numba/core/typing/arraydecl.py:68

During: typing of intrinsic-call at forloop_slicing_parallel_q.py (12)
During: typing of static-get-item at forloop_slicing_parallel_q.py (12)

File "forloop_slicing_parallel_q.py", line 12:
def method1_par1(n):
    <source elided>
    # np.meshgrid not supported by numba, therefore trying to use np.vstack
    Y   = np.hstack( n*[y[:,np.newaxis]] )
    ^

It sounds to me that the way I am using vstack is not supported, is that correct?

The versions I am using:

numpy:  1.20.3
numba:  0.54.1
python: 3.8.10

Update1: Using np.concatenate( n*[[x]] ) results in a similar error:

 >>> mul(int64, list(list(array(float64, 1d, C))<iv=None>)<iv=None>)
There are 12 candidate implementations:
    - Of which 10 did not match due to:
    Overload of function 'mul': File: <numerous>: Line N/A.
      With argument(s): '(int64, list(list(array(float64, 1d, C))<iv=None>)<iv=None>)':
     No match.
    - Of which 2 did not match due to:
    Operator Overload in function 'mul': File: unknown: Line unknown.
      With argument(s): '(int64, list(list(array(float64, 1d, C))<iv=None>)<iv=None>)':
    No match for registered cases:
     * (int64, int64) -> int64
     * (int64, uint64) -> int64
     * (uint64, int64) -> int64
     * (uint64, uint64) -> uint64
     * (float32, float32) -> float32
     * (float64, float64) -> float64
     * (complex64, complex64) -> complex64
     * (complex128, complex128) -> complex128

Update2: Using np.hstack( n*[y[:,np.newaxis]] ) results in the same error as the original error.

Using y.reshape(-1,1).repeat(n,1) results in the following error:

No implementation of function Function(<function array_repeat at 0x7f0dac54ee50>) found for signature:

 >>> array_repeat(array(float64, 2d, C), int64, Literal[int](1))

There are 2 candidate implementations:
    - Of which 2 did not match due to:
    Overload in function 'array_repeat': File: numba/np/arrayobj.py: Line 2066.
      With argument(s): '(array(float64, 2d, C), int64, int64)':
     Rejected as the implementation raised a specific error:
       TypeError: array_repeat() takes 2 positional arguments but 3 were given
    raised from /usr/local/lib/python3.8/dist-packages/numba/core/typing/templates.py:775

- Resolution failure for non-literal arguments:
No implementation of function Function(<function array_repeat at 0x7f0dac54ee50>) found for signature:

 >>> array_repeat(array(float64, 2d, C), int64, int64)

There are 2 candidate implementations:
   - Of which 2 did not match due to:
   Overload in function 'array_repeat': File: numba/np/arrayobj.py: Line 2066.
     With argument(s): '(array(float64, 2d, C), int64, int64)':
    Rejected as the implementation raised a specific error:
      TypeError: array_repeat() takes 2 positional arguments but 3 were given
   raised from /usr/local/lib/python3.8/dist-packages/numba/core/typing/templates.py:775

During: resolving callee type: BoundFunction((<class 'numba.core.types.npytypes.Array'>, 'repeat') for array(float64, 2d, C))

Answer

If you look at the error message carefully, you will see that it says

No implementation of function Function(<built-in function getitem>) found for signature:

>>> getitem(array(float64, 1d, C), Tuple(slice<a:b>, none))

getitem is how numba compiles the [] operator. The signature shows that numba does not support a call like array[slice, None]. Specifically, the problem is

y[:, np.newaxis]

Numba does support the reshape operation and repeat, but without an axis argument, so you can change that line to

Y = y.repeat(n).reshape(n, n)