repeating a numpy array a specified number of times for itertools

Question

I am trying to write some code that will give me the itertools product, for a varying number of inputs. For example, this works for me.

test = np.array([x for x in itertools.product([0,2],[0,2],[0,2])])

this gives me my desired result:

>>> test
array([[0, 0, 0],
       [0, 0, 2],
       [0, 2, 0],
       [0, 2, 2],
       [2, 0, 0],
       [2, 0, 2],
       [2, 2, 0],
       [2, 2, 2]])

However, I'd like to be able to be able to pass to the product function a varying number of lists. For example:

test = np.array([x for x in itertools.product([0,2],[0,2],[0,2],[0,2])])

or

test = np.array([x for x in itertools.product([0,2],[0,2])])

I have tried

test = np.array([x for x in itertools.product(([0,2],) * 3)])

and

test = np.array([x for x in itertools.product([[0,2]]*3)])

but neither gives me the desired result. Surely there is an easy way to do this. I would appreciate any help.

Answer 1

itertools.product supports another argument called repeat as in itertools.product(*iterables[, repeat]) through which you can manipulate the dimensions of cross product. Note, this argument should be specified explicitly in order to disambiguate from the list content.

So your example extends to

test = np.array([x for x in itertools.product([0,2],[0,2],[0,2],[0,2])])

to

test = np.array([x for x in itertools.product([0,2], repeat = 4)])