How to make a multidimension numpy array with a varying row size?

In numpy 1.14.3, using append:

d = []                 # initialize an empty list
a = np.arange(3)       # array([0, 1, 2])
d.append(a)            # [array([0, 1, 2])]
b = np.arange(3,-1,-1) #array([3, 2, 1, 0])
d.append(b)            #[array([0, 1, 2]), array([3, 2, 1, 0])]

what you get an list of arrays (that can be of different lengths) and you can do operations like d[0].mean(). On the other hand,

cells = numpy.array([[0,1,2,3], [2,3,4]])

results in an array of lists.

You may want to do this:

a1 = np.array([1,2,3])
a2 = np.array([3,4])
a3 = np.array([a1,a2])
a3 # array([array([1, 2, 3]), array([3, 4])], dtype=object)
type(a3) # numpy.ndarray
type(a2) # numpy.ndarray

Slightly off-topic, but not as much as one would think because of eager mode which is now the default: If you are using Tensorflow, you can do:

a = tf.ragged.constant([[0, 1, 2, 3]])
b = tf.ragged.constant([[2, 3, 4]])
c = tf.concat([a, b], axis=0)

And you can then do all the mathematical operations still, like tf.math.reduce_mean, etc.

np.array([[0,1,2,3], [2,3,4]], dtype=object) returns an "array" of lists.

a = np.array([np.array([0,1,2,3]), np.array([2,3,4])], dtype=object) returns an array of arrays. It allows already for operations such as a+1.

Building up on this, the functionality can be enhanced by subclassing.

import numpy as np

class Arrays(np.ndarray):
    def __new__(cls, input_array, dims=None):
        obj = np.array(list(map(np.array, input_array))).view(cls)
        return obj
    def __getitem__(self, ij):
        if isinstance(ij, tuple) and len(ij) > 1:
            # handle twodimensional slicing
            if isinstance(ij[0],slice) or hasattr(ij[0], '__iter__'):
                # [1:4,:] or [[1,2,3],[1,2]]
                return Arrays(arr[ij[1]] for arr in self[ij[0]])
            return self[ij[0]][ij[1]] # [1,:] np.array
        return super(Arrays, self).__getitem__(ij)
    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        axis = kwargs.pop('axis', None)
        dimk = [len(arg) if hasattr(arg, '__iter__') else 1 for arg in inputs]
        dim = max(dimk)
        pad_inputs = [([i]*dim if (d<dim) else i) for d,i in zip(dimk, inputs)]
        result = [np.ndarray.__array_ufunc__(self, ufunc, method, *x, **kwargs) for x in zip(*pad_inputs)]
        if method == 'reduce':
            # handle sum, min, max, etc.
            if axis == 1:
                return np.array(result)
            else:
                # repeat over remaining axis
                return np.ndarray.__array_ufunc__(self, ufunc, method, result, **kwargs)
        return Arrays(result)

Now this works:

a = Arrays([[0,1,2,3], [2,3,4]])
a[0:1,0:-1]
# Arrays([[0, 1, 2]])
np.sin(a)
# Arrays([array([0.        , 0.84147098, 0.90929743, 0.14112001]),
#        array([ 0.90929743,  0.14112001, -0.7568025 ])], dtype=object)
a + 2*a
# Arrays([array([0, 3, 6, 9]), array([ 6,  9, 12])], dtype=object)

To get nanfunctions working, this can be done

# patch for nanfunction that cannot handle the object-ndarrays along with second axis=-1
def nanpatch(func):
    def wrapper(a, axis=None, **kwargs):
        if isinstance(a, Arrays):
            rowresult = [func(x, **kwargs) for x in a]
            if axis == 1:
                return np.array(rowresult)
            else:
                # repeat over remaining axis
                return func(rowresult)
        # otherwise keep the original version
        return func(a, axis=axis, **kwargs)
    return wrapper

np.nanmean = nanpatch(np.nanmean)
np.nansum = nanpatch(np.nansum)
np.nanmin = nanpatch(np.nanmin)
np.nanmax = nanpatch(np.nanmax)

np.nansum(a)
# 15
np.nansum(a, axis=1)
# array([6, 9])