Concatenating views in NumPy

Viewed 336

Indexing NumPy array through slices/indexes creates a view that is lightweight (doesn't copy data) and allows assigning to elements of original array. I.e.

import numpy as np
a = np.array([1, 2, 3, 4, 5])
a[2:4] = [6, 7]
print(a)
# [1 2 6 7 5]

But how about multiple views, how do I concatenate them to create a bigger view that still assigns to original first array. E.g. for imaginary function concatenate_views(...):

a = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
concatenate_views((a[1:3], a[4:6], a[7:9])) = [11, 12, 13, 14, 15, 16]
print(a)
# should print [1 11 12 4 13 14 7 15 16 10]

Of course, I can create a list of indexes for each view that it views to, just by converting slices to indexes and then concatenating these indexes. This way I'll get all indexes of concatenated views and can use these indexes to create a combined view. But this is not what I want. I want NumPy to keep the notion of slices representation, because all slices can be very long and it will be inefficient to convert and store these slices as indexes. I want NumPy to be aware of underlying slices of all concatenated views to make internally just simple looping of slices ranges.

Also would be nice to generalize the problem. Not only concatenate views, but also allow to form any arbitrary tree of slicing/indexing operations, e.g. concatenate views, then apply some slicing, then indexing, then slicing, then concatenate again. Also N-dimensional slicing/indexing. I.e. all fancy stuff that can be done with single un-concatenated view.

The main point of concatenated views is only efficiency. Of cause we can represent any view or slicing operation by N-D array of integer indexes (coordinates, like meshgrid) and then can use this array to make a view of source array. But if numpy can keep a notion of source set of slices instead of array of integers, then first it will be lightweight (much less memory consumption), second instead of reading indexes from memory numpy internally can loop (iterate) through each slice more efficiently in C++ loops.

By having concatenated view I wanted to be able to apply any numpy operation like np.mean(...) to the combined view in efficient way.

Full procedure of concatenating views of N-D slicing based on 2D example is described down below:

    1 Step described below:
    
    2D array slicing using 3 slices for each axis
    
    a,b,c - sizes of "slices" along axis 0
    d,e,f - sizes of "slices" along axis 1
    
    Each "slice" - is either slice(start, stop, step) or 1D array of integer indexes
            
      d e f
     .......
    a.0.1.2.
     .......
    b.3.4.5.
     .......
    c.6.7.8.
     .......

    Above 0 1 2 3 4 5 6 7 8 mean not a single integer but some 2D sub-array.
    Dots (`.`) also mean some 2D sub-arrays.
    
    Sub-views shapes:
    0:(a, d), 1:(a, e), 2:(a, f)
    3:(b, d), 4:(b, e), 5:(b, f)
    6:(c, d), 7:(c, e), 8:(c, f)

    Final aggregated (concatenated) view shape:
    ((a + b + c), (d + e + f))
    containing 2D array
    012
    345
    678
    
    There can be more than one Steps, each next Step applies new sequence of slicing
    to the final view obtained on previous Step. Each Step has different set of sizes
    of slices and different amount of slices per each dimension.
    In general each next Step reduces number of total elements, except the case
    when slices or indexes overlap then you may get more elements but with duplicates.
2 Answers

You can use np.r_ to concatenate slice objects and assign back to the indexed array:

a[np.r_[1:3, 4:6, 7:9]] = [11, 12, 13, 14, 15, 16]

print(a)
array([ 1, 11, 12,  4, 13, 14,  7, 15, 16, 10])

Update

Based on your update, I think you might want something like:

from itertools import islice

it = iter([11, 12, 13, 14, 15, 16])
for s in slice(1,3), slice(4,6), slice(7,9):
    a[s] = list(islice(it, s.stop-s.start))

print(a)
array([ 1, 11, 12,  4, 13, 14,  7, 15, 16, 10])

You can only concatenate views if they are contiguous in terms of dtypes, strides and offsets. Here is one way to check. This way is likely incomplete, but it illustrates the gist of it. Basically, if the views share a base, and the strides and offsets are aligned so that they are on the same grid, you can concatenate.

In the spirit of TDD, I will work with the following example:

x = np.arange(24).reshape(4, 6)

We (or at least I) want the following to be concatenatable:

a, b = x[:, :4], x[:, 4:]        # Basic case
a, b = x[:, :4:2], x[:, 4::2]    # Strided
a, b = x[:, :4:2], x[:, 2::2]    # Strided overlapping
a, b = x[1:2, 1:4], x[2:4, 1:4]  # Stacked

# Completely reshaped:
a, b = x.ravel()[:12].reshape(3, 4), x.ravel()[12:].reshape(3, 4)
# Equivalent to
a, b = x[:2, :].reshape(3, 4), x[2:, :].reshape(3, 4)

We do not want the following to be concatenatable:

a, b = x, np.arange(12).reshape(2, 6)   # Buffer mismatch
a, b = x[0, :].view(np.uint), x[1:, :]  # Dtype mismatch
a, b = x[:, ::2], x[:, ::3]             # Stride mismatch
a, b = x[:, :4], x[:, 4::2]             # Stride mismatch
a, b = x[:, :3], x[:, 4:]               # Overlap mismatch
a, b = x[:, :4:2], x[:, 3::2]           # Overlap mismatch
a, b = x[:-1, :-1], x[1:, 1:]           # Overlap mismatch
a, b = x[:-1, :4], x[:, 4:]             # Shape mismatch

The following could be interpreted as concatenatable, but won't be in this case:

a, b = x, x[1:-1, 1:-1]

The idea is that everything (dtype, strides, offsets) has to match exactly. Only one axis offset is allowed to be different between the views, as long as it is no more than one stride away from the edge of the other view. The only possible exception is when one view is fully contained in another, but we will ignore this scenario here. Generalizing to multiple dimensions should be pretty simple if we use array operations on the offsets and strides.

def cat_slices(a, b):
    if a.base is not b.base:
        raise ValueError('Buffer mismatch')
    if a.dtype != b.dtype:  # I don't thing you can use `is` here in general
        raise ValueError('Dtype mismatch')

    sa = np.array(a.strides)
    sb = np.array(b.strides)

    if (sa != sb).any():
        raise ValueError('Stride mismatch')

    oa = np.byte_bounds(a)[0]
    ob = np.byte_bounds(b)[0]

    if oa > ob:
        a, b = b, a
        oa, ob = ob, oa

    offset = ob - oa

    # Check if you can get to `b` from a by moving along exactly one axis
    # This part works consistently for arrays with internal overlap
    div = np.zeros_like(sa)
    mod = np.ones_like(sa)  # Use ones to auto-flag divide-by zero
    np.divmod(offset, sa, where=sa.astype(bool), out=(div, mod))

    zeros = np.flatnonzero((mod == 0) & (div >= 0) & (div <= a.shape))

    if not zeros.size:
        raise ValueError('Overlap mismatch')

    axis = zeros[0]

    check_shape = np.equal(a.shape, b.shape)
    check_shape[axis] = True

    if not check_shape.all():
        raise ValueError('Shape mismatch')

    shape = list(a.shape)
    shape[axis] = b.shape[axis] + div[axis]

    start = np.byte_bounds(a)[0] - np.byte_bounds(a.base)[0]

    return np.ndarray(shape, dtype=a.dtype, buffer=a.base, offset=start, strides=a.strides)

Some things that this function does not handle:

  • Merging flags
  • Broadcasting
  • Handling arrays that are fully contained within each other but with multi-axis offsets
  • Negative strides

You can, however, check that it returns the expected views (and errors) for all the cases shown above. In a more production-y version, I could envision this enhancing np.concatenate, so for failed cases, it would just copy data instead of raising an error.

Related