Search code examples
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).

Thanks in advance.

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
    nested_mask = is_cell_an_array(array)
    
    # use itertools.product to generate all the combinations of the nested arrays
    for permutation in product(*array[nested_mask]):
        # create a new array
        new_array = np.zeros(array.shape, dtype=float)
        # copy in the unique combination of the nested arrays for this loop
        new_array[nested_mask] = permutation
        # copy in the rest of the values from the original array
        new_array[~nested_mask] = array[~nested_mask]
    
        print(new_array)