tuple as output of @vectorize function (numba): what signature so I have to use?

I am trying to create a vectorized function

@vectorize()
def rhou(f):
    rho = numpy.sum(f, axis=-1)
    u = 1.0
    v = 2.0
    w = 3.0
    velocity = numpy.array([u,v,w])
    return rho, velocity

rho, f, u, v, w, and velocity are all of the type float64. f is an numpy array of the form f[:,:,:,:]. So the input array (being f) can be an array of any size regarding the first three indexes. The sum is taken over the fourth in the function rhou.

My question is: what signature should I put in @vectorize()?

I have tried @vectorize((float 64, float64)(float64)) (and many other probably stupid attempts with Tuple for example), but that does not work. The help page of numba deals with s single output, but in this case the function returns two results.

Thanks!

Tried many things, but nothing helps. The numba help page is not clear about multiple outputs

The vectorized decorator does not support multiple outputs (afaik), but you can use guvectorize in order to achieve this. The guvectorize is more flexible, but that also means you'll have to specify a bit more information regarding shapes & types.

But your code is also unclear at the moment, and you don't provide any sample input making it impossible to run. Applying vectorized on any array, will as the name suggests, vectorize the operation, meaning it will perform an element-wise calculation on each entry in the input individually. So that makes your sum(f) statement meaningless, since it will always take the sum of a single value. So I expect you either don't want to vectorize (element-wise), or don't want to take the sum, but this is unclear from your post.

The example below using guvectorize shows something similar as what I think you're trying to achieve, but I'm not sure. With guvectorize you'll have to specify the dimensions you want to keep ("see" in the function), and everything else will be vectorized. The default assumption is to vecorize over remaining dimensions at the "end", but the function will automatically get an axis keyword (similar as Numpy ufuncs) that allows you to change that.

@guvectorize([
    "void(float64[:], float64[:], float64[:], float64[:])"
], "(n),(m)->(),(m)")
def rhou(f, dummy, rho, velocity):
    rho[0] = np.sum(f)
    velocity[0] = 1.0
    velocity[1] = 2.0
    velocity[2] = 3.0

f = np.random.randn(3)
velocity = np.empty(3, dtype=np.float64)
rho = np.empty(1, dtype=np.float64)

rhou(f, velocity, rho, velocity)

Note how there is no return statement, the output arrays are arguments to the function, and the results as set on those! This is a similar pattern as the out= keyword on most Numpy functions.

I don't know anyway to avoid having to pass a dummy input to let Numba know what the size of "m" is (shape of velocity), this is necessary even if you provide the return array explicitly as shown above.

An alternative way of writing this is having Numba create/allocate the output arrays. It can do this because both the shape (from the signature) and the datatype is known. But I often see people who try this not realizing you have to set the result by indexing the array (even if you output a scalar!). Especially at first it's probably more intuitive to create the outputs yourself, and pass the explicitly to the function.

But that would look something like:

import numba; import numpy as np

@numba.guvectorize([
    "void(float64[:], float64[:], float64[:], float64[:])",
], "(n),(m)->(),(m)")
def rhou(f, _, rho, velocity):
    rho[0] = np.sum(f)
    velocity[0] = 1.0
    velocity[1] = 2.0
    velocity[2] = 3.0

f = np.random.randn(3)
dummy = np.empty(3, dtype=np.float64)

rho, velocity = rhou(f, dummy)