ENH: Generalize split function to split matrix into sub-matrices.

[KoningR]

Closes #15380

Co-authored-by: KoningR r.m.koning@student.tudelft.nl
Co-authored-by: erwinvanthiel e.l.vanthiel@student.tudelft.nl

[eric-wieser]

Thanks for the contribution.

This is a good start, but in the context of the "0, 1, infinity rule", I don't think we should take this patch in its current form.
Essentially, if we're going to generalize a function that works on 1d arrays, we should generalize it to nd arrays, and not intermediately generalize to 2d.

If possible, do you think you could implement the suggestion in #15380 (comment)? Modifying one of your existing examples, that would give:

    >>> xs = np.array_split(x, [3, 3], axis=(0, 1))
    >>> xs.shape
    (3, 3)
    >>> xs
    array([[array([[0, 1], [4, 5]),
            array([[2], [6]]),
            array([[3], [7]])],
           [array([[8, 9]]),
            array([[10]]),
            array([[11]])],
           [array([[12, 13]]),
            array([[14]]),
            array([[15]])]], dtype=object)

[eric-wieser]

Suggested change

	if len(indices_or_sections) is not len(axis):
	if len(indices_or_sections) != len(axis):

Comparing integers with is is not guaranteed to work

[charris]

Also needs a release note if it goes in.

[eric-wieser]

Formatting is off here:

Suggested change

	>>> np.array_split(x, 3)
	[array([0., 1., 2.]), array([3., 4.]), array([5., 6.])]
	>>> x = np.reshape(np.arange(16), (4, 4))
	>>> np.array_split(x, [3, 3], (0, 1))
	[array([[0, 1],
	[4, 5]]),
	array([[2], [6]]),
	array([[3], [7]]),
	array([[8, 9]]),
	array([[10]]),
	array([[11]]),
	array([[12, 13]]),
	array([[14]]),
	array([[15]])]
	>>> x = np.reshape(np.arange(8), (2, 2, 2))
	>>> np.array_split(x, [2, 2, 2], (0, 1, 2))
	[array([[[0]]]),
	array([[[1]]]),
	array([[[2]]]),
	array([[[3]]]),
	array([[[4]]]),
	array([[[5]]]),
	array([[[6]]]),
	array([[[7]]])]
	>>> np.array_split(x, 3)
	[array([0., 1., 2.]), array([3., 4.]), array([5., 6.])]
	>>> x = np.reshape(np.arange(16), (4, 4))
	>>> np.array_split(x, [3, 3], (0, 1))
	[array([[0, 1],
	[4, 5]]),
	array([[2], [6]]),
	array([[3], [7]]),
	array([[8, 9]]),
	array([[10]]),
	array([[11]]),
	array([[12, 13]]),
	array([[14]]),
	array([[15]])]
	>>> x = np.reshape(np.arange(8), (2, 2, 2))
	>>> np.array_split(x, [2, 2, 2], (0, 1, 2))
	[array([[[0]]]),
	array([[[1]]]),
	array([[[2]]]),
	array([[[3]]]),
	array([[[4]]]),
	array([[[5]]]),
	array([[[6]]]),
	array([[[7]]])]

[eric-wieser]

Suggested change

	>>> np.split(x, [4, 2], (0, 1))
	[array([[0, 1]]),
	array([[2, 3]]),
	array([[4, 5]]),
	array([[6, 7]]),
	array([[8, 9]]),
	array([[10, 11]]),
	array([[12, 13]]),
	array([[14, 15]])]
	[array([[0, 1]]),
	array([[2, 3]]),
	array([[4, 5]]),
	array([[6, 7]]),
	array([[8, 9]]),
	array([[10, 11]]),
	array([[12, 13]]),
	array([[14, 15]])]

[eric-wieser]

Suggested change

	[[0, 1],
	[4, 5]],
	[[0, 1],
	[4, 5]],

[eric-wieser]

or

Suggested change

	[[0, 1],
	[4, 5]],
	[[0, 1], [4, 5]],

[eric-wieser]

It's my opinion that blocks = np.split(arr, (X, Y), (0, 1)) should return an object array of shape (X, Y), not a list of length X*Y. The latter is easy to construct from the former as blocks.ravel().tolist(), but the former is much harder to construct from the latter.

The advantage would be that things like blocks[0,:] could be used to select the first row of blocks, etc.

For compatilibity, you'd only return an object array if a tuple is passed for the number of splits.

[KoningR]

@eric-wieser

Hi,
We stumbled upon the problem that the sub-matrices within the array do not necessarily have the same dimensions.
This is the case quite often when splitting on indices instead of scalars. For example:

matrix = np.reshape(np.arange(16), (4, 4))
res = array_split(matrix, [[2, 3], [2, 3]], (0, 1))

In this case, res will be:
res = [
[[0, 1], [4, 5]],
[[2], [6]],
[[3], [7]],
[[8, 9]],
[[10]],
[[11]],
[[12, 13]],
[[14]],
[[15]]
]

The elements in res are clearly not of the same size.

What would be the best way to convert this to an N-dimensional array?
We already tried to reshape the res array but that seems to require all elements be of the same dimensions.

[eric-wieser]

Just a tip - use triple backticks for multi-line codeblocks in comments.

[eric-wieser]

One way to transform that would be:

# as you have it now
res = [
    np.array([[0, 1], [4, 5]]),
    np.array([[2], [6]]),
    np.array([[3], [7]]),
    np.array([[8, 9]]),
    np.array([[10]]),
    np.array([[11]]),
    np.array([[12, 13]]),
    np.array([[14]]),
    np.array([[15]]),
]

# convert list to object array.
# It's possible you can assemble it this way in the first place,
# but doesn't really matter if you do it this way
res_arr = np.empty(len(res), dtype=object)
for i, subarr in enumerate(res):
    res_arr[i] = subarr

# I assume you already have a way to compute this shape
return res.reshape((3, 3))

[KoningR]

@eric-wieser

Hi,
Currently, two tests are failing the pipeline. Both tests appear to execute the commented examples on lines 761 and 922 of shape_base.py. The only differences between the expected and received outputs are whitespace characters. The output on lines 48-54 of the error log appears to be formatted in a rather odd fashion, though we cannot figure out what causes this as the object itself seems to behave correctly. Therefore, it appears that only the string representation of the object is odd. Should we reformat the example as to match the object array’s string formatting? Or did we overlook anything and are the contents of the object array wrong, causing the string to be wrong as well?

The string representation of the object array:

    array([[array([[0, 1],
           [4, 5]]), array([[2],
           [6]]),
            array([[3],
           [7]])],
           [array([[8, 9]]), array([[10]]), array([[11]])],
           [array([[12, 13]]), array([[14]]), array([[15]])]], dtype=object)

[eric-wieser]

Unfortunately you're running up against #15869.

[KoningR]

@eric-wieser

Hi. It appears PR #15997 was delayed because it might break doctests already in the code base. However, we are now dependent on this PR. If doctests will be refactored sometime in the future anyhow, can we push our doctests such that they pass the build, and be done with this issue?

[eric-wieser]

The PR you reference is now merged, so that question is moot :)

[mattip]

needs some rewrapping to fit into 80 columns

Edit: ... and in the rest of the PR as well

[eric-wieser]

Probably clearer to explain this with a kwarg.

Suggested change

	To use this new functionality, write ``np.split(array, [4, 2], (0, 1))`` or ``np.array_split(array, [4, 2], (0, 1))``
	To use this new functionality, write ``np.split(array, [4, 2], axis=(0, 1))`` or ``np.array_split(array, [4, 2], axis=(0, 1))``

[eric-wieser]

Do you think we should allow

Suggested change

	>>> np.array_split(x, [3, 3], (0, 1))
	>>> np.array_split(x, 3, (0, 1))

to mean the same thing?

[KoningR]

I like the idea and have added the functionality. Based on your proposed test_zero_dimensional_split, I assumed that ret = split(matrix, 3, axis=()) should operate the same as ret = split(matrix, [], axis=()) and therefore do nothing.

[eric-wieser]

Yes, "split into three parts on no axes" seems like a legal instruction to do nothing, and an error would just be annoying.

[KoningR]

Hi. I think this PR is done. Would anyone mind having a look at it?

[seberg]

We discussed it today, we think this can be pushed forward API wise. We haven't reviewed the implementation though, so we will need to find/prioritize some review time actually push it forward.