How can I interpolate based on index values when using a pandas MultiIndex?

Question

I have demographic panel data, where each data point is categorized by a country, sex, year, and age. For a given country, sex, and year, my age-pattern has missing data, and I want to interpolate it based on the value of the age. For example, if 5 year-olds have a value of 5, and 10 year-olds have a value of 10, 6.3 year-olds should have a value of 6.3. I cannot use the default pandas 'linear' interpolation method because my age groups are not spaced linearly. My data look something like this:

iso3s = ['USA', 'CAN']
age_start_in_years = [0, 0.01, 0.1, 1]
years = [1990, 1991]
sexes = [1,2]
multi_index = pd.MultiIndex.from_product([iso3s,sexes,years,age_start_in_years],
                                          names = ['iso3','sex','year','age_start'])

frame_length = len(iso3s)*len(age_start_in_years)*len(years)*len(sexes)
test_df = pd.DataFrame({'value':range(frame_length)},index=multi_index)
test_df=test_df.sortlevel()

# Insert missingness to practice interpolating
idx = pd.IndexSlice
test_df.loc[idx[:,:,:,[0.01,0.1]],:] = np.NaN
test_df

                                value
iso3    sex year    age_start   
CAN     1   1990    0.00        0
                    0.01        NaN
                    0.10        NaN
                    1.00        3
            1991    0.00        4
                    0.01        NaN
                    0.10        NaN
                    1.00        7
       2    1990    0.00        8
...

However, when I try to use test_df.interpolate(method='index') , I get this error:

ValueError: Only `method=linear` interpolation is supported on MultiIndexes.

Surely there must be some way to interpolate based on the index values.

Answer 1

I found this hacky work-around that gets rid of the MultiIndex and uses a combination of groupby and transform:

def multiindex_interp(x, interp_col, step_col):

    valid = ~pd.isnull(x[interp_col])
    invalid = ~valid

    x['last_valid_value'] = x[interp_col].ffill()
    x['next_valid_value'] = x[interp_col].bfill()

    # Generate a new Series filled with NaN's
    x['last_valid_step'] =  np.NaN
    # Copy the step value where we have a valid value
    x['last_valid_step'][valid] = x[step_col][valid]
    x['last_valid_step'] = x['last_valid_step'].ffill()

    x['next_valid_step'] =  np.NaN
    x['next_valid_step'][valid] = x[step_col][valid]
    x['next_valid_step'] = x['next_valid_step'].bfill()

    # Simple linear interpolation= distance from last step / (range between closest valid steps) *
    #                              difference between closest values + last value
    x[interp_col][invalid] = (x[step_col]-x['last_valid_step'])/(x['next_valid_step'] - x['last_valid_step']) \
                             * (x['next_valid_value']-x['last_valid_value']) \
                             + x['last_valid_value']
    return x

test_df = test_df.reset_index(drop=False)
grouped = test_df.groupby(['iso3','sex','year'])
interpolated = grouped.transform(multiindex_interp,'value','age_start')
test_df['value'] = interpolated['value']
test_df
    iso3    sex year    age_start   value
0   CAN     1   1990    0.00        16.00
1   CAN     1   1990    0.01        16.03
2   CAN     1   1990    0.10        16.30
3   CAN     1   1990    1.00        19.00
4   CAN     1   1991    0.00        20.00
5   CAN     1   1991    0.01        20.03
6   CAN     1   1991    0.10        20.30
7   CAN     1   1991    1.00        23.00
8   CAN     2   1990    0.00        24.00
9   CAN     2   1990    0.01        24.03
10  CAN     2   1990    0.10        24.30
11  CAN     2   1990    1.00        27.00
...

Answer 2

This might come a little late, but I ran into the same problem today. What I came up with is also just a workaround, but it uses pandas built-ins at least. My approach was to reset the index, then group by the first subset of index columns (ie all but age_start ). These sub-DataFrames can then be interpolated with the method='index' parameter and put back together into a whole frame with pd.concat . The resulting DataFrame then gets its original index reassigned.

idx_names = test_df.index.names
test_df = test_df.reset_index()
concat_list = [grp.set_index('age_start').interpolate(method='index') for _, grp in test_df.groupby(['iso3', 'sex', 'year'])]
test_df = pd.concat(concat_list)
test_df = test_df.reset_index().set_index(idx_names)
test_df
                         value
iso3 sex year age_start       
CAN  1   1990 0.00       16.00
              0.01       16.03
              0.10       16.30
              1.00       19.00
         1991 0.00       20.00
              0.01       20.03
              0.10       20.30
              1.00       23.00
     2   1990 0.00       24.00

EDIT

I got back to this problem today and found a bug in my originally proposed solution. When the multi-index is not ordered as it is in your example, the above code sorts your DataFrame by index values. To get around this, I joined the result back into a DataFrame with the original index so that index order is preserved. I've also put it inside a function.

def interp_multiindex(df, interp_idx_name):
    """
    Provides index-based interpolation for pd.Multiindex which usually only support linear
    interpolation. Interpolates full DataFrame.

    Parameters
    ----------
    df : pd.DataFrame
        The DataFrame with NaN values
    interp_idx_name : str
        The name of the multiindex level on which index-based interpolation should take place

    Returns
    -------
    df : pd.DataFrame
        The DataFrame with index-based interpolated values
    """
    # Get all index level names in order
    existing_multiidx = df.index
    # Remove the name on which interpolation will take place
    noninterp_idx_names = [idx_name for idx_name in existing_multiidx.names 
                           if idx_name != interp_idx_name]
    df = df.reset_index()
    concat_list = [grp.set_index(interp_idx_name).interpolate(method='index') 
                   for _, grp in df.groupby(noninterp_idx_names)]
    df = pd.concat(concat_list)
    df = df.reset_index().set_index(existing_multiidx.names)
    df = pd.DataFrame(index=existing_multiidx).join(df)
    return df

Answer 3

You can try something like this:

test_df.groupby(level=[0,1,2])\
       .apply(lambda g: g.reset_index(level=[0,1,2], drop=True)
                         .interpolate(method='index'))

Output:

                         value
iso3 sex year age_start       
CAN  1   1990 0.00       16.00
              0.01       16.03
              0.10       16.30
              1.00       19.00
         1991 0.00       20.00
              0.01       20.03
              0.10       20.30
              1.00       23.00
     2   1990 0.00       24.00
              0.01       24.03
              0.10       24.30
              1.00       27.00
         1991 0.00       28.00
              0.01       28.03
              0.10       28.30
              1.00       31.00
USA  1   1990 0.00        0.00
              0.01        0.03
              0.10        0.30
              1.00        3.00
         1991 0.00        4.00
              0.01        4.03
              0.10        4.30
              1.00        7.00
     2   1990 0.00        8.00
              0.01        8.03
              0.10        8.30
              1.00       11.00
         1991 0.00       12.00
              0.01       12.03
              0.10       12.30
              1.00       15.00