I am working on a script that uses itertools to generate desired combinations of given parameters. I am having trouble generating the appropriate 'sets' however when some variables are linked, i.e. excluding certain combinations. Consider the following:
import itertools
A = ['a1','a2','a3']
B = ['b1','b2','b3']
C = ['c1','c2']
If I want to generate all possible combinations of these elements, I can simply use itertools.product()
all_combinations = list(itertools.product(A,B,C))
Which gives the expected
[('a1', 'b1', 'c1'), ('a1', 'b1', 'c2'), ('a1', 'b2', 'c1'), ...
('a3', 'b2', 'c2'), ('a3', 'b3', 'c1'), ('a3', 'b3', 'c2')]
For 18 combinations (3*3*2)
However, how can I 'link' parameters A and B such that each returned set contains only 'an','bn' elements? That is, I have tried:
ABprecombine = zip(A,B)
limited_combinations = list(itertools.product(ABprecombine,C))
Which returns
[(('a1', 'b1'), 'c1'), (('a1', 'b1'), 'c2'), (('a2', 'b2'), 'c1'),
(('a2', 'b2'), 'c2'), (('a3', 'b3'), 'c1'), (('a3', 'b3'), 'c2')]
This is the six (3*1*2) desired products, but obviously due to the way I created it I now have an extra tuple.
Of course I could generate all combinations and then filter out given ones, but is there a smart way to 'link' parameters as above?
Here, zipping A and B is the right way to go. You can flatten the tuples pretty easily if you want:
limited_combinations = [(a, b, c) for ((a, b), c) in itertools.product(zip(A, B), C)]
If you want more detailed control of what combinations get produced, things can rapidly get more complicated, up to the difficulty of needing to solve NP-hard problems like boolean satisfiability. If that happens, look into existing libraries for that kind of thing.