Question

I have this issue:

in my software I am either dealing with a single array or a list of 3 arrays (they are 1 or 3 components of a pixelized sky map).

If the single array was a list of 1 array, then it would be very easy to iterate over it transparently, regardless the number of elements.

Now, let’s say I want to square these maps:

my_map = np.ones(100) # case of single component
# my_map = [np.ones(100) for c in [0, 1, 2]] # case of 3 components

if isinstance(my_map, list): #this is ugly
    my_map_2 = [m**2 for m in my_map]
else:
    my_map = my_map ** 2

would you have any suggestion on how to improve this?

Answer 1

Why wouldn’t you directly create a 2D array ?

my_array = np.ones((100,3), dtype=float)

That way, you could directly square your ‘three’ arrays at once. You could still access individual elements as:

(x, y, z) = my_array.T

where .T is a shortcut for the .transpose method.

Using this approach would be far more efficient than looping on a list, especially if you apply the same function to the three arrays. Even if you want, say, to square the first array, double the second and take the square root of the third, you could:

my_array[:,0] **=2
my_array[:,1] *=2
my_array[:,2] **=0.5