Suppose I have three different colors of paint and I want to test what happens when I mix them in various ratios. How do I get a list of those ratios?
I can do it with nested for-loops like this:
for i in np.linspace(0, 1, 6):
for j in np.linspace(0, 1, 6):
for k in np.linspace(0, 1, 6):
if np.isclose(i+j+k, 1.0): ## catch floating point errors
print(i, j, k)
This outputs the following:
0.0 0.0 1.0
0.0 0.2 0.8
0.0 0.4 0.6
0.0 0.6 0.4
0.0 0.8 0.2
0.0 1.0 0.0
0.2 0.0 0.8
0.2 0.2 0.6
0.2 0.4 0.4
0.2 0.6 0.2
0.2 0.8 0.0
0.4 0.0 0.6
0.4 0.2 0.4
0.4 0.4 0.2
0.4 0.6 0.0
0.6 0.0 0.4
0.6 0.2 0.2
0.6 0.4 0.0
0.8 0.0 0.2
0.8 0.2 0.0
1.0 0.0 0.0
Each line has three numbers and they all add up to one, and I have every possible combination for my specified degree of resolution. It works, but it's an awkward construction, and it requires a new loop for every new color of paint I add. If I wanted to mix 10 types of paint, I'd need 10 loops!
Is there a more elegant/pythonic way of doing this? A built-in numpy function maybe?
(Also I found it hard to summarise the question in the title, can anyone think of a better concise way of describing the problem?)
You can use itertools.product for this; effectively this is combinations with replacement. You can construct a list comprehension which filters for sum equal to 1.
from itertools import product
res = np.array([i for i in product(np.linspace(0, 1, 6), repeat=3) if sum(i) == 1])
# array([[ 0. , 0. , 1. ],
# [ 0. , 0.2, 0.8],
# [ 0. , 0.4, 0.6],
# [ 0. , 0.6, 0.4],
# [ 0. , 0.8, 0.2],
# [ 0. , 1. , 0. ],
# [ 0.2, 0. , 0.8],
# [ 0.2, 0.2, 0.6],
# [ 0.2, 0.4, 0.4],
# [ 0.2, 0.6, 0.2],
# [ 0.2, 0.8, 0. ],
# [ 0.4, 0. , 0.6],
# [ 0.4, 0.2, 0.4],
# [ 0.4, 0.4, 0.2],
# [ 0.4, 0.6, 0. ],
# [ 0.6, 0. , 0.4],
# [ 0.6, 0.2, 0.2],
# [ 0.6, 0.4, 0. ],
# [ 0.8, 0. , 0.2],
# [ 0.8, 0.2, 0. ],
# [ 1. , 0. , 0. ]])