How to gather arrays of different sizes in the same array with Numpy?-CodePudding

Context: I have 3 arrays. A that is 3x3, B that is 5x2, and C that is a 3D array.

Question: Is there a way to stack arrays of different sizes along the 1st dimension of a 3D array, with Numpy ?

Example: if A and B are stacked in C along its first dimension, and I want to access A, I would type C[0].

Problem: I know I can use Xarray for that, but I wanted to know if there was a Numpy way to do it. So far, I have been artificially extending the arrays with NaNs to match their sizes (see code below).

Code:

# Generate the arrays
A = np.random.rand(3,3)
B = np.random.rand(5,2)
C = np.zeros((2,5,3))

# Resize the arrays to fit C
AR = np.full([C.shape[1], C.shape[2]], np.nan) # Generate an array full of NaNs that fits C
AR[:A.shape[0],:A.shape[1]] = A # Input the previous array in the resized one
BR = np.full([C.shape[1], C.shape[2]], np.nan) # Generate an array full of NaNs that fits C
BR[:B.shape[0],:B.shape[1]] = B # Input the previous array in the resized one

# Stack the resized arrays in C
C[0] = AR
C[1] = BR

CodePudding user response：

You won't be able to slice it as freely, but you can use dtype = 'object' when making a jagged array.

jagged_array = np.array([np.zeros((3, 2)), np.zeros((5, 2)), np.zeros((3, 2, 5))], dtype = 'object')