Search code examples
pythonpandasdataframesubtraction

How to subtract columns in a multiindex dataframe?


I have a multiindex dataframe like this:

import pandas as pd
import numpy as np

df = pd.DataFrame({'ind1': list('aaaaaaaaabbbbbbbbb'),
                   'ind2': list('cccdddeeecccdddeee'),
                   'ind3': list(range(3))*6,
                   'val1': list(range(100, 118)),
                   'val2': list(range(70, 88))})

df_mult = df.set_index(['ind1', 'ind2', 'ind3'])

                val1  val2
ind1 ind2 ind3            
a    c    0      100    70
          1      101    71
          2      102    72
     d    0      103    73
          1      104    74
          2      105    75
     e    0      106    76
          1      107    77
          2      108    78
b    c    0      109    79
          1      110    80
          2      111    81
     d    0      112    82
          1      113    83
          2      114    84
     e    0      115    85
          1      116    86
          2      117    87

What I want to do is to subtract the values in df_mult.loc['a', 'e', :] and df_mult.loc['b', 'e', :], respectively from the values corresponding to df_mult.loc['a', ['c', 'd'], :] and df_mult.loc['b', ['c', 'd'], :], respectively. The expected outcome would be

                val1  val2
ind1 ind2 ind3            
a    c    0       -6    -6
          1       -6    -6
          2       -6    -6
     d    0       -3    -5
          1       -3    -5
          2       -3    -5
     e    0      106    76
          1      107    77
          2      108    78
b    c    0       -6    -6
          1       -6    -6
          2       -6    -6
     d    0       -3    -3
          1       -3    -3
          2       -3    -3
     e    0      115    85
          1      116    86
          2      117    87

Ideally, something like this would work

df_mult.loc['a', ['c', 'd'], :].subtract(df_mult.loc['a', 'e', :])

but this gives me a lot of NaNs.

How would I do this?


Solution

  • UPDATE2: with kind help of @Divakar:

    def repeat_blocks(a, repeats=2, block_length=None):
        N = a.shape[0]
        if not block_length:
            block_length = N//2
        out = np.repeat(a.reshape(N//block_length,block_length,-1),
                        repeats,
                        axis=0) \
                .reshape(N*repeats,-1)
        return out
    
    In [234]: df_mult.loc[idx[['a','b'], ['c', 'd'], :], :] -= repeat_blocks(df_mult.loc[['a','b'], 'e', :].values)
    
    In [235]: df_mult
    Out[235]:
                    val1  val2
    ind1 ind2 ind3
    a    c    0       -6    -6
              1       -6    -6
              2       -6    -6
         d    0       -3    -3
              1       -3    -3
              2       -3    -3
         e    0      106    76
              1      107    77
              2      108    78
    b    c    0       -6    -6
              1       -6    -6
              2       -6    -6
         d    0       -3    -3
              1       -3    -3
              2       -3    -3
         e    0      115    85
              1      116    86
              2      117    87
    

    UPDATE:

    In [100]: idx = pd.IndexSlice
    
    In [102]: df_mult.loc[idx['a', ['c', 'd'], :], :] -= \
                  np.concatenate([df_mult.loc['a', 'e', :].values] * 2)
    
    In [103]: df_mult
    Out[103]:
                    val1  val2
    ind1 ind2 ind3
    a    c    0       -6    -6
              1       -6    -6
              2       -6    -6
         d    0       -3    -3
              1       -3    -3
              2       -3    -3
         e    0      106    76
              1      107    77
              2      108    78
    b    c    0      109    79
              1      110    80
              2      111    81
         d    0      112    82
              1      113    83
              2      114    84
         e    0      115    85
              1      116    86
              2      117    87
    

    Old (incorrect) answer:

    In [62]: df_mult.loc['a', 'e', :] -= df_mult.loc['b', 'e', :].values
    
    In [63]: df_mult
    Out[63]:
                    val1  val2
    ind1 ind2 ind3
    a    c    0      100    70
              1      101    71
              2      102    72
         d    0      103    73
              1      104    74
              2      105    75
         e    0       -9    -9
              1       -9    -9
              2       -9    -9
    b    c    0      109    79
              1      110    80
              2      111    81
         d    0      112    82
              1      113    83
              2      114    84
         e    0      115    85
              1      116    86
              2      117    87