Find all n-dimensional lines and diagonals with NumPy

Question

Using NumPy, I would like to produce a list of all lines and diagonals of an n-dimensional array with lengths of k.

---

Take the case of the following three-dimensional array with lengths of three.

array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])

For this case, I would like to obtain all of the following types of sequences. For any given case, I would like to obtain all of the possible sequences of each type. Examples of desired sequences are given in parentheses below, for each case.

• 1D lines
  • x axis (0, 1, 2)
  • y axis (0, 3, 6)
  • z axis (0, 9, 18)
• 2D diagonals
  • x/y axes (0, 4, 8, 2, 4, 6)
  • x/z axes (0, 10, 20, 2, 10, 18)
  • y/z axes (0, 12, 24, 6, 12, 18)
• 3D diagonals
  • x/y/z axes (0, 13, 26, 2, 13, 24)

---

The solution should be generalized, so that it will generate all lines and diagonals for an array, regardless of the array's number of dimensions or length (which is constant across all dimensions).

Answer 1

This solution generalized over n

Lets rephrase this problem as "find the list of indices".

We're looking for all of the 2d index arrays of the form

array[i[0], i[1], i[2], ..., i[n-1]]

Let n = arr.ndim

Where i is an array of shape (n, k)

Each of i[j] can be one of:

• The same index repeated n times, ri[j] = [j, ..., j]
• The forward sequence, fi = [0, 1, ..., k-1]
• The backward sequence, bi = [k-1, ..., 1, 0]

With the requirements that each sequence is of the form ^(ri)*(fi)(fi|bi|ri)*$ (using regex to summarize it). This is because:

• there must be at least one fi so the "line" is not a point selected repeatedly
• no bis come before fis, to avoid getting reversed lines

---

def product_slices(n):
    for i in range(n):
        yield (
            np.index_exp[np.newaxis] * i +
            np.index_exp[:] +
            np.index_exp[np.newaxis] * (n - i - 1)
        )

def get_lines(n, k):
    """
    Returns:
        index (tuple):   an object suitable for advanced indexing to get all possible lines
        mask (ndarray):  a boolean mask to apply to the result of the above
    """
    fi = np.arange(k)
    bi = fi[::-1]
    ri = fi[:,None].repeat(k, axis=1)

    all_i = np.concatenate((fi[None], bi[None], ri), axis=0)

    # inedx which look up every possible line, some of which are not valid
    index = tuple(all_i[s] for s in product_slices(n))

    # We incrementally allow lines that start with some number of `ri`s, and an `fi`
    #  [0]  here means we chose fi for that index
    #  [2:] here means we chose an ri for that index
    mask = np.zeros((all_i.shape[0],)*n, dtype=np.bool)
    sl = np.index_exp[0]
    for i in range(n):
        mask[sl] = True
        sl = np.index_exp[2:] + sl

    return index, mask

---

Applied to your example:

# construct your example array
n = 3
k = 3
data = np.arange(k**n).reshape((k,)*n)

# apply my index_creating function
index, mask = get_lines(n, k)

# apply the index to your array
lines = data[index][mask]
print(lines)

array([[ 0, 13, 26],
       [ 2, 13, 24],
       [ 0, 12, 24],
       [ 1, 13, 25],
       [ 2, 14, 26],
       [ 6, 13, 20],
       [ 8, 13, 18],
       [ 6, 12, 18],
       [ 7, 13, 19],
       [ 8, 14, 20],
       [ 0, 10, 20],
       [ 2, 10, 18],
       [ 0,  9, 18],
       [ 1, 10, 19],
       [ 2, 11, 20],
       [ 3, 13, 23],
       [ 5, 13, 21],
       [ 3, 12, 21],
       [ 4, 13, 22],
       [ 5, 14, 23],
       [ 6, 16, 26],
       [ 8, 16, 24],
       [ 6, 15, 24],
       [ 7, 16, 25],
       [ 8, 17, 26],
       [ 0,  4,  8],
       [ 2,  4,  6],
       [ 0,  3,  6],
       [ 1,  4,  7],
       [ 2,  5,  8],
       [ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 13, 17],
       [11, 13, 15],
       [ 9, 12, 15],
       [10, 13, 16],
       [11, 14, 17],
       [ 9, 10, 11],
       [12, 13, 14],
       [15, 16, 17],
       [18, 22, 26],
       [20, 22, 24],
       [18, 21, 24],
       [19, 22, 25],
       [20, 23, 26],
       [18, 19, 20],
       [21, 22, 23],
       [24, 25, 26]])

---

Another good set of test data is np.moveaxis(np.indices((k,)*n), 0, -1), which gives an array where every value is its own index

---

I've solved this problem before to implement a higher dimensional tic-tac-toe

Answer 2

In [1]: x=np.arange(27).reshape(3,3,3)

Selecting individual rows is easy:

In [2]: x[0,0,:]
Out[2]: array([0, 1, 2])
In [3]: x[0,:,0]
Out[3]: array([0, 3, 6])
In [4]: x[:,0,0]
Out[4]: array([ 0,  9, 18])

You could iterate over dimensions with an index list:

In [10]: idx=[slice(None),0,0]
In [11]: x[idx]
Out[11]: array([ 0,  9, 18])
In [12]: idx[2]+=1
In [13]: x[idx]
Out[13]: array([ 1, 10, 19])

Look at the code for np.apply_along_axis to see how it implements this sort of iteration.

Reshape and split can also produce a list of rows. For some dimensions this might require a transpose:

In [20]: np.split(x.reshape(x.shape[0],-1),9,axis=1)
Out[20]: 
[array([[ 0],
        [ 9],
        [18]]), array([[ 1],
        [10],
        [19]]), array([[ 2],
        [11],
        ...

np.diag can get diagonals from 2d subarrays

In [21]: np.diag(x[0,:,:])
Out[21]: array([0, 4, 8])
In [22]: np.diag(x[1,:,:])
Out[22]: array([ 9, 13, 17])
In [23]: np.diag?
In [24]: np.diag(x[1,:,:],1)
Out[24]: array([10, 14])
In [25]: np.diag(x[1,:,:],-1)
Out[25]: array([12, 16])

And explore np.diagonal for direct application to the 3d. It's also easy to index the array directly, with range and arange, x[0,range(3),range(3)].

As far as I know there isn't a function to step through all these alternatives. Since dimensions of the returned arrays can differ, there's little point to producing such a function in compiled numpy code. So even if there was a function, it would step through the alternatives as I outlined.

==============

All the 1d lines

x.reshape(-1,3)
x.transpose(0,2,1).reshape(-1,3)
x.transpose(1,2,0).reshape(-1,3)

y/z diagonal and anti-diagonal

In [154]: i=np.arange(3)
In [155]: j=np.arange(2,-1,-1)
In [156]: np.concatenate((x[:,i,i],x[:,i,j]),axis=1)
Out[156]: 
array([[ 0,  4,  8,  2,  4,  6],
       [ 9, 13, 17, 11, 13, 15],
       [18, 22, 26, 20, 22, 24]])

Answer 3

np.einsum can be used to build all these kind of expressions; for instance:

# 3d diagonals
print(np.einsum('iii->i', a))
# 2d diagonals
print(np.einsum('iij->ij', a))
print(np.einsum('iji->ij', a))