Given a numpy array of dimension n with each direction having length m, I would like to iterate through all 1-dimensional arrays of length m.
For example, consider:
import numpy as np
x = np.identity(4)
array([[1., 0., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]])
then I would like to find all all 1-dimensional arrays with length 4. So the result should include all 4 rows, all 4 columns, and the 2 diagonals:
[x[i,:] for i in range(4)] + [x[:,i] for i in range(4)] + [np.array([x[i,i] for i in range(4)])] + [np.array([x[3-i,i] for i in range(4)])]
It's unclear to me how to generalize this to higher dimensional arrays since the position of the ":" in the slice needs to iterate as well. In a higher dimensional analogue with
slices = [[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2], [2, 0], [2, 1], [2, 2]]
we can get
[x[i,j,:] for (i,j) in slices]
but then I'm not sure how to proceed to iterate through the permutations of [i,j,:].
Although the comprehensions are readable, for large arrays you probably want to use what numpy gives you:
import numpy as np
n = 4
# random n*n numpy array
arr = np.random.rand(n, n)
print(arr)
# this has all the data you need, relatively efficiently - but not in 1D shape
result = (
np.vsplit(arr, n) +
np.hsplit(arr, n) +
[arr.diagonal()] +
[arr[np.arange(n), np.arange(n - 1, -1, -1)]]
)
# you can just flatten them as you used them:
for xs in result:
print(xs.ravel())
# or flatten them into a new result:
result_1d = [xs.ravel() for xs in result]
Edit: user @Matt correctly pointed out in the comments that this solution only works for the case of a 2-dimensional array.
Things get a bit more complicated for an arbitrary number of dimensions n with size m across all dimensions. This works, but given the complexity, can probably be improved upon for simplicity:
import numpy as np
import itertools as it
# random n*n numpy array
m = 2
n = 3
arr = np.random.rand(*([m] * n))
print(arr)
def get_all_lines(arr):
ndim = arr.ndim
size = arr.shape[0] # assuming the same size for all dimensions
# generate each 1d slice along and across each axis
for fixed in it.product(range(size), repeat=ndim - 1):
for axis in range(ndim):
yield arr[fixed[:axis] + (slice(None, ),) + fixed[axis:]]
# generate each 1d diagonal for each combination of axes
for d_dim in range(2, ndim+1): # d_dim is the number of varying axes
for fixed in it.product(range(size), repeat=(ndim - d_dim)): # fixed indices for the other axes
of, od = 0, 0 # offsets for accessing fixed values and directions
# each varying axis can be traversed in one of two directions
for d_tail in it.product((0, 1), repeat=d_dim - 1): # dir is the direction for each varying axis
d = (1, *d_tail)[::-1] # deduplicate and reverse the direction
for axes in it.combinations(range(ndim), d_dim): # axes to vary
fm = d_dim * (d_dim + 1) // 2 - sum(axes) # first dimension with a fixed index
dm = min(axes) # first dimension with a varying index
yield [
arr[*[fixed[of := 0 if j == fm else of + 1]
if j not in axes else
(i if d[od := 0 if j == dm else od + 1] else size - (i + 1))
for j in range(ndim)]] for i in range(size)
]
lines = get_all_lines(arr)
for line in lines:
print(line)
The mentioned "deduplication" avoids including each diagonal twice (once in both directions).
Also note that this yields 1d arrays as well as lists of numbers, you can of course cast these appropriately.