pythonnumpyloopsphysics

# Iterate over a n-dimensional numpy array which contains unknown number of numpy array as its values

I have a numpy array of the following form

``````[[array([ 1.  ,  3.25,  5.5 ,  7.75, 10.  ]) 0 0]
[0 array([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.]) 0]
[0 0 0]]
``````

The array will be N dimensional, and at specific positions, I will have arrays over which I need to iterate. And I do not know the final dimension, or how many arrays I will have.

How do I write a program/function so that I get all the possible N dimensional tensors/matrices with only a single value at each position? (i.e., in positions where there are arrays, I need to iterate over them and I need to generate a list of tensor/matrices having only singles values at each position.

I do not know if it is important, but the maximum length in each direction is 3.

(Physical meaning of the problem: basically, I want to create input files for another software. I need to iterate over variation of electric field in the directions which I will specify. So, in 1 dimension, I will have only X, Y and Z, for 2 dimensions I will have XX, YY, ZZ, XZ, etc. And in each position I indicate the arrays of values over which I want to iterate, i.e. for example a linspace(10, 100, 10) etc).

I do not even know how to approach it, or how to google it. I was thinking maybe creating some additional function to help me somehow iterate etc.

So, the desired result is that after each iteration I get a numpy array with only values from arrays inside it, so I can submit it to another function. For example, the resulting np.arrays from the np.array above would be

``````First,

[[1 0 0]
[0 1 0]
[0 0 0]]
Second,

[[1 0 0]
[0 2 0]
[0 0 0]]

...

Last,

[10 0 0]
[0 10 0]
[0 0 0]]

``````

Solution

• The builtin itertools module has the perfect function for this: product

Then it's just a matter of handling the numpy array. I found the key to that was creating a mask of boolean values indicating where the nested arrays are.

I've only tested it on a 3x3 matrix but I think it should work on any array.

``````from itertools import product

import numpy as np

array = np.array([[np.array([1, 2]), 0.5, 0.75],
[-0.5, np.array([10, 20]), -0.75],
[np.array([-1, -2]), 3.5, 3.75]], dtype=object)

# create a vectorized function to find nested arrays
is_cell_an_array = np.vectorize(lambda cell: isinstance(cell, np.ndarray))
# create an array of booleans indicating where the nested cells are