Fill a multidimensional array efficiently that have many if else statements

Question

I want to fill an 4dim numpy array in a specific and efficient way. Because I don't know better I startet to write the code with if else statements, but that doesn't look nice, is probably slow and I also can not be really sure if I thought about every combination. Here is the code which I stopped writing down:

sercnew2 = numpy.zeros((gn, gn, gn, gn))
for x1 in range(gn):
    for x2 in range(gn):
        for x3 in range(gn):
            for x4 in range(gn):
                if x1 == x2 == x3 == x4: 
                    sercnew2[x1, x2, x3, x4] = ewp[x1]
                elif x1 == x2 == x3 != x4:
                    sercnew2[x1, x2, x3, x4] = ewp[x1] * ewp[x4]
                elif x1 == x2 == x4 != x3:
                    sercnew2[x1, x2, x3, x4] = ewp[x1] * ewp[x3]
                elif x1 == x3 == x4 != x2:
                    sercnew2[x1, x2, x3, x4] = ewp[x1] * ewp[x2]
                elif x2 == x3 == x4 != x1:
                    sercnew2[x1, x2, x3, x4] = ewp[x2] * ewp[x1]
                elif x1 == x2 != x3 == x4:
                    sercnew2[x1, x2, x3, x4] = ewp[x1] * ewp[x3]
                elif ... many more combinations which have to be considered

So basically what should happen is, that if all variables (x1, x2, x3, x4) are different from each other, the entry would be:

sercnew2[x1, x2, x3, x4] = ewp[x1]* ewp[x2] * ewp[x3] * ewp[x4]

Now if lets say the variable x2 and x4 is the same then:

sercnew2[x1, x2, x3, x4] = ewp[x1]* ewp[x2] * ewp[x3]

Others examples can be seen in the code above. Basically if two or more variables are the same, then I only consider on of them. I hope the pattern is clear. Otherwise please let me note and I will try to express my problem better. I am pretty sure, that there is a much more intelligent way to do it. Hope you know better and thanks in advance :)

Answer 1

I don't really know what you mean by saying that those variables are the same, but if indeed they are, than all you have to do is just use set().

from functools import reduce
from operator import mul
sercnew2 = numpy.zeros((gn, gn, gn, gn))
for x1 in range(gn):
    for x2 in range(x1, gn):
        for x3 in range(x2, gn):
            for x4 in range(x3, gn):
                set_ = [ewp[n] for n in set([x1, x2, x3, x4])]
                sercnew2[x1, x2, x3, x4] = reduce(mul, set_, 1)

The way it works is that it creates a set() which delete duplicates, and later with reduce function I pick first number from the set_, multiplies it with 1 (the initializer value), and the result of this is going to be passed to reduce as first argument, and second will be the second item from set_. Sorry for my bad explanation.

Answer 2

Really hoping that I have got it! Here's a vectorized approach -

from itertools import product

n_dims = 4 # Number of dims

# Create 2D array of all possible combinations of X's as rows
idx = np.sort(np.array(list(product(np.arange(gn), repeat=n_dims))),axis=1)

# Get all X's indexed values from ewp array
vals = ewp[idx]

# Set the duplicates along each row as 1s. With the np.prod coming up next,
#these 1s would not affect the result, which is the expected pattern here.
vals[:,1:][idx[:,1:] == idx[:,:-1]] = 1

# Perform product along each row and reshape into multi-dim array
out = vals.prod(1).reshape([gn]*n_dims)

Answer 3

You can do it in a single for loop too. Building on Divakar's Trick for a list of indexes, the first thing we need to do is figure out how to extract only the unique indices of a given element in the 4d array sercnew2.

One of the fastest ways to do this (Refer: https://www.peterbe.com/plog/uniqifiers-benchmark) is using sets. Then, we simply have to initialize sercnew2 as an array of ones, rather than zeros.

from itertools import product
import numpy as np

sercnew2 = np.ones((gn, gn, gn, gn))
n_dims=4
idx = list(product(np.arange(gn), repeat=n_dims))

for i,j,k,l in idx:
    unique_items = set((i,j,k,l))
    for ele in unique_items:
        sercnew2[i,j,k,l] *= ewp[ele]

EDIT: As @unutbu suggested, we could also use the cartesian_product function from https://stackoverflow.com/a/11146645/5714445 to speed up the initialization of idx

Edit2: In case you are having difficulty understanding what product from itertools does, it provides all permutations. For instance, suppose gn=2, with repeat dimension set to 4, you get

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