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