Concatenate/stack unpacked numpy array in certain axis to other arrays

Question

I want to create a numpy array variables with a certain structure, as in the example below, consisting of the variables X and Y for a meshgrid and additional parameters params . I need this in the context of plotting a 3D surface plot for a function f(variables) .

This is how the variables array should look like with some arbitrary params 5, 5, 5.

import numpy as np

X = np.linspace(0, 2, num=3)
Y = np.linspace(0, 2, num=3)
X, Y = np.meshgrid(X, Y)
variables = np.array([X, Y, 5, 5, 5])
print(variables)

With the desired output:

array([[[ 0.,  1.,  2.],
        [ 0.,  1.,  2.],
        [ 0.,  1.,  2.]],
       [[ 0.,  0.,  0.],
        [ 1.,  1.,  1.],
        [ 2.,  2.,  2.]],
       5, 5, 5])

And my questions is how can I construct this when having the parameters 5, 5, 5 in a numpy array also. (And explicitly not having something like [X, Y, [5, 5, 5]] .)

params = np.array([5, 5, 5])

I already looked at different numpy functions like stack or concatenate, but could not figure it out. I'm rather new to numpy, so thanks for any help.

Answer 1

You can't get there from here.

Numpy arrays are for storing cuboidal grids of data, where each entry is the same shape as every other entry.

The best you can do here is have a (5,) -grid, containing object s, where each of the first two objects is itself a numpy array:

params = np.array([5, 5, 5])
XY = np.array([X, Y, None])[:-1]  # trick numpy into using a 1d object array
                                  # there might be a better way to do this
variables = np.concatenate((XY, params))
variables

array([array([[ 0.,  1.,  2.],
       [ 0.,  1.,  2.],
       [ 0.,  1.,  2.]]),
       array([[ 0.,  0.,  0.],
       [ 1.,  1.,  1.],
       [ 2.,  2.,  2.]]),
       5, 5, 5], dtype=object)

Are you sure you need a numpy array at all here? Is this enough:

variables = (X, Y) + tuple(params)