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

<hr>

<p>Take the case of the following three-dimensional array with lengths of three.</p>

<pre><code>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]]])
</code></pre>

<p>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.</p>

<ul>
<li>1D lines

<ul>
<li>x axis (<code>0, 1, 2</code>)</li>
<li>y axis (<code>0, 3, 6</code>)</li>
<li>z axis (<code>0, 9, 18</code>)</li>
</ul></li>
<li>2D diagonals

<ul>
<li>x/y axes (<code>0, 4, 8</code>, <code>2, 4, 6</code>)</li>
<li>x/z axes (<code>0, 10, 20</code>, <code>2, 10, 18</code>)</li>
<li>y/z axes (<code>0, 12, 24</code>, <code>6, 12, 18</code>)</li>
</ul></li>
<li>3D diagonals

<ul>
<li>x/y/z axes (<code>0, 13, 26</code>, <code>2, 13, 24</code>)</li>
</ul></li>
</ul>

<hr>

<p>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).</p>

<pre><code>In [1]: x=np.arange(27).reshape(3,3,3)
</code></pre>

<p>Selecting individual <code>rows</code> is easy:</p>

<pre><code>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])
</code></pre>

<p>You could iterate over dimensions with an index list:</p>

<pre><code>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])
</code></pre>

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

<p>Reshape and split can also produce a list of <code>rows</code>.  For some dimensions this might require a <code>transpose</code>:</p>

<pre><code>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],
        ...
</code></pre>

<p><code>np.diag</code> can get diagonals from 2d subarrays</p>

<pre><code>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])
</code></pre>

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

<p>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.</p>

<p>==============</p>

<p>All the 1d lines</p>

<pre><code>x.reshape(-1,3)
x.transpose(0,2,1).reshape(-1,3)
x.transpose(1,2,0).reshape(-1,3)
</code></pre>

<p>y/z diagonal and anti-diagonal</p>

<pre><code>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]])
</code></pre>

python arrays numpy diagonal n-dimensional