If we have a quantum system that has n_sites and n_levels, itertools.product provides a convenient way to enumerate the basis_states:
1 2 3 4 |
import itertools n_sites=2 # 2 lattice sites n_levels=3 # 3 energy levels per lattice site basis_states = list(itertools.product(range(n_levels), repeat=n_sites)) |
1 2 |
In [19]: list(itertools.product(range(n_levels), repeat=n_sites)) Out[19]: [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)] |