Syntactic sugar for creating a pandas series with a MultiIndex

Question

The following code generates a Pandas series with a hierarchical MultiIndex:

import pandas as pd
number_of_classes = 3
number_of_price_brackets = 10
survival_table = pd.Series(
  index = pd.MultiIndex.from_tuples(
    [(gender,klass,fare)
     for gender in range(2)
     for klass in range(number_of_classes)
     for fare in range(number_of_price_brackets)],
     names=['Gender','Class','Price Bracket']
     ))

But this approach seems a bit "manual". Is there any prettier way of generating the survival_table variable?

Answer 1

Seems fine to me.

You could use itertools.product to create the tuples without the three for sections in the comprehension. But I think this is just as a good and will be readable for novices, whereas approaches with itertools are less so. The itertools way would be like this:

list(itertools.product(range(2), 
                       range(number_of_classes), 
                       range(number_of_price_brackets)))

You could also define a helper function if you find that you do this often.

def make_category_multiindex(categories, names):
    from itertools import product
    return pd.MultiIndex.from_tuples(list(product(*categories)), names=names)

Then do:

categories = [range(2), 
              range(number_of_classes), 
              range(number_of_price_brackets)]

names = ['Gender','Class','Price Bracket']

survival_table = pd.Series(index=make_category_multiindex(categories, names))