Feature request: advanced slicing (on left hand side)

[ZisIsNotZis]

Introduction

numpy is an important tool in scientific computation (e.g. optimization), and visualization is an important subject in this field. Sometimes, after done some numerical optimization, people would want to visualize the result of how things get optimized. In order to do so, a common way is to draw some raw video (by manipulating array) and play that later using tools such as ffmpeg.

With advanced indexing, point drawing can be extremely simple, but line/box drawing is getting inconvenient. What people usually have is indexes of start/end point. Since they can't put the indexes in slice, neither can they unite the tiny slices to a big slice, they have to explicitly draw using (nested) for loop + tiny-slice assigning. This is not taking advantage of numpy's vectorization, and is slow due to huge python overhead.

Definition

One way to solve this is allowing advanced slicing (i.e. slice array using array). A straightforward example would be

import numpy as np
a = np.arange(8)
a[[0,1,3]:[1,3,6]]  # "should" be [[0], [1,2], [3,4,5]]

The problem is, when applied on RHS, advanced slicing have to generate a "physical" array, which might make the result non-rectangular and unsupported by numpy.

On the contrary, indexing on LHS is fully "imaginary". Therefore, advanced slicing becomes possible:

import numpy as np
a = np.arange(8)
a[[0,1,3]:[1,3,6]] = [0,1,2]
a # should be [0, 1, 1, 2, 2, 2, 6, 7]  

Since this makes ndarray much more powerful, I would strongly request such an enhancement. There are two set of detailed rules on my mind:

Possible Rule 1

• On RHS: advanced slicing is disallowed.
• On LHS: advanced slicing broadcasts with regular advanced indexing. All axes sliced such way will be invisible (in terms of broadcasting) to RHS, and all values/operations from RHS will be broadcasted to all elements in these invisible axes.

Possible Rule 2

• Generally: advanced slicing broadcasts with advanced index.
• On RHS: advanced slicing should produce a valid array, otherwise an exception is raised.
• On LHS: advanced slicing might produce array with invalid axes (i.e. axes that have different shapes inside). Those axes will have a shape of ?, which is only broadcastable against shape 1.

Example

Draw 10 random rectangles on a 256x256 black area with random grayscale

import numpy as np
a = np.zeros((256, 256))
xBeg,xEnd = np.random.randint(0, 256, (2,10)).sort(0)
yBeg,yEnd = np.random.randint(0, 256, (2,10)).sort(0)
colors = np.random.rand(10)
if POSSIBLE_RULE == 1:
    a[xBeg:xEnd,yBeg:yEnd] = colors
else:
    # `a[xBeg:xEnd,yBeg:yEnd]` will have shape `10,?,?,3`
    a[xBeg:xEnd,yBeg:yEnd] = colors[:,None,None]

Given info of 10 rectangles in 1000 frames, and color of each rectangle, draw them into a raw video of 256x256

import numpy as np
def draw_advanced_slice(xBeg, xEnd, yBeg, yEnd, colors):
    assert xBeg.shape == xEnd.shape == yBeg.shape == yEnd.shape == (1000,10)
    assert colors.shape == (10,3)
    a = np.memmap('video.raw', mode='w+', shape=(1000,256,256,3))
    if POSSIBLE_RULE == 1:
        a[np.arange(1000)[:,None], xBeg:xEnd, yBeg:yEnd] = colors
    else:
        # LHS will have shape `1000,10,?,?,3`
        a[np.arange(1000)[:,None], xBeg:xEnd, yBeg:yEnd] = colors[:,None,None]

Compared to the current implementation

import numpy as np
def draw_without_advanced_slice(xBeg, xEnd, yBeg, yEnd, colors):
    assert xBeg.shape == xEnd.shape == yBeg.shape == yEnd.shape == (1000,10)
    assert colors.shape == (10,3)
    a = np.memmap('video.raw', mode='w+', shape=(1000,256,256,3))
    for a, xBeg, xEnd, yBeg, yEnd in zip(a, xBeg, xEnd, yBeg, yEnd):
        for xBeg, xEnd, yBeg, yEnd, colors in zip(xBeg, xEnd, yBeg, yEnd, colors):
            a[xBeg:xEnd, yBeg:yEnd] = colors

Things other than assigning. This might be more tricky because it cannot be translated to a[xs:ys] = a[xs:ys] * b

import numpy as np
a = np.arange(8)
if POSSIBLE_RULE == 1:
    a[[0,1,3]:[1,3,6]] *= [0,1,2]
else:
    # LHS will have shape `3,?`
    a[[0,1,3]:[1,3,6]] *= np.array([0,1,2])[:,None]
a # should be [0, 1, 2, 6, 8, 10, 6, 7]  

Note

Another way to solve non-rectangular broadcast assigning problem is using a masked-array. But generating such a mask is inefficient, and sometimes non-trivial. It might take huge memory (and long time) to generate the mask in the first place

[seberg]

Masked arrays or advanced indexing in general can help you do this. To be honest, I am a bit cautious about such an addition. I admit, it is probably fairly straight forward in a sense, but it also seems to introduce a hole new complexity to indexing (which already is complicated enough that most users do not really know the full power of advanced indexing).

I see two main possibilities:

1. Your main concern is with overheads for many such (small to mid sized) slices, in which case advanced-indexing should be an OK solution (you should probably create the array of indexes instead of using a mask).
2. The concern is that advanced-indexing is slow, in which case you probably have largish areas, so that the overhead of iterating the slices is not very bad.

Of course I am handwaving a bit here, your approach can obviously be better optimized, and there will be some in-between use cases where that may be more substantial.

The main question is whether the use-case (and common of use) is worth the additional complexity (both in code and teaching users).

Another alternative could be to not do this within indexing, but rather its own function, although it may not change things much. Could you write more about your use-case, and maybe float this idea to the mailing list, since it is a fair big API change proposal it needs to go there anyway. Once there is some agreement that this should be tried, someone will have to pick it up and try to create a PR for it.

[ZisIsNotZis]

Thanks for the reply, I'll definitely think about the use case of this more and probably send something to the mailing list later!

Another thing, actually I believe advanced indexing won't help in this case:

The height/width of different rectangles are different. Therefore, those advanced indices cannot broadcast together due to different length. I'm not sure is there any way to generate the indices using less than O(n_rectangles) python-level function calls.

This is the most efficient code I've came up with (draw 20 rectangles):

import numpy as np
xy = np.random.randint(0, 256, (20,2,2), 'B')
xy.sort(2)
c = np.random.randint(0, 256, (20,3), 'B')

a = np.zeros((256,256,3), 'B')
vindex = np.concatenate([np.broadcast_to(np.arange(i,j,1,'B')[:,None],(j-i,J-I))for (i,j),(I,J) in xy],None)
hindex = np.concatenate([np.broadcast_to(np.arange(I,J,1,'B')        ,(j-i,J-I))for (i,j),(I,J) in xy],None)
a[vindex,hindex] = c.repeat((xy[...,1]-xy[...,0]).prod(1,'H'), 0)

Essentially it still uses python loop of size O(n_rectangles). The slicing method uses for loop of the same size, and does not have overhead of creating index. Therefore, advanced index might only be slower.

---

I also implemented a tricky solution using masked array:

# xy and c same as above
N = np.arange(0, 256, 1, 'B')[:,None]

a = np.zeros((256,256,3), 'B')
m = (xy[:,1,0]>N) | (N>=xy[:,1,1]) | ((xy[:,0,0]>N) | (N>=xy[:,0,1]))[:,None]
A,m = np.broadcast_arrays(a[:,:,None], m[...,None])
A = np.ma.masked_array(A, m, hard_mask=True)
A[:] = c

This is fully vectorized without python loop, but creating such "broadcasted mask" is very memory costly

[seberg]

Yeah, you will have to repeat c. With the advanced indexing I was mostly thinking about a use-case where you draw the same rectangles multiple times (so that building the index is at least not a performance issue). But that is probably just not what you need...

Of course, I am sure it could be a small factor faster, but note that:

def draw2(): 
    for ((i, j), (I, J)), color in zip(xy, c): 
        a[i:j, I:J, :] = color 

is almost 5 times as fast as your function for that example (although all the cast to 'B' will be counterproductive for speed at least unless the arrays are large enough for memory to take an effect, maybe even then.)

There is always the possibility of writing something like this yourself. E.g. cython is probably pretty nice to do.

[eric-wieser]

One thing that might be neat is allowing some kind of hook for slice objects, such that regular_ndarray[MyUnusualSliceObj(...)] calls MyUnusualSliceObj(...).__array_slice__(regular_ndarray) or something. That would give @ZisIsNotZis ways to get the syntax they want.

Probably something we should not think about before nailing down the OIndex / VIndex ideas.