How to speed up nested loops for groupby multiindex

Viewed 141

I have two Multiindex dataframes, namely, panel1 and panel2: both have the same level 0 index-the dates, but different level 1 index; see sample code below:

# panel1:
idx1 = pd.MultiIndex.from_product([['2017-05-02', '2017-05-03', '2017-05-04'],['id1', 'id2', 'id3']],names=['Dates', 'id'])
panel1=pd.DataFrame(np.random.randn(9,2), index=idx1,columns=['ytm','mat'])
# panel2:
idx2 = pd.MultiIndex.from_product([['2017-05-02', '2017-05-03', '2017-05-04'],['0.5', '1.5', '2.5']],names=['Dates', 'yr'])
panel2=pd.DataFrame(np.random.randn(9), index=idx2,columns=['curve'])

I want to loop over the two panels by the Dates (level 0 index). So for each day (e.g. '2017-05-02'), I search the mat of each id/row (of panel1) in the yr column (of panel2), if there is a match, I want to get the corresponding curve values (of panel2) and add it as a new column (named CDB) in panel1.

My current code as following:

group1=panel1.groupby(level=0)
group2=panel2.groupby(level=0)

lst=[]
for ytm in group1:              # loop over each day
    for yr in group2:           # loop over each day
        df_ytm=ytm[1]           # get df of id, yt & mat
        df_ytm=df_ytm.assign(CDB=np.nan)      # add a col of nan, later will be replaced by matched curve values
        df_curve=yr[1].reset_index()          # need get rid of index to match yr with t_mat
        df_curve.yr=df_curve.yr.astype(float) 
        for i in range(df_ytm.shape[0]):      # loop over each row
            if (df_ytm.iloc[i,1]==df_curve.yr).any()==True:      # search if each 'mat' value in 'yr' column
                df_ytm.iloc[i,2]=df_curve[df_curve.yr.isin([df_ytm.t_mat[i]])].curve.values   # if matched, set 'CDB' as curve value
    lst.append(df_ytm)      # need get modified 'df_ytm' (with matched 'CDB')  

The code works as I tried with a small sample, but I have a huge panel 1 (sized 800 days times 10000 ids) and big panel 2 as well. So the code has been running for more than 24 hours.

I wonder how could I rewrite the code (use possible vectorization) to speed it up?

Any comments would be much appreciated!

2 Answers

If I understand you correctly you need to construct new MultiIndex from Dates index and mat column and get values of curve for this index.

import pandas as pd
import numpy as np

np.random.seed(12)
idx1 = pd.MultiIndex.from_product(
    [["2017-05-02", "2017-05-03", "2017-05-04"], ["id1", "id2", "id3"]],
    names=["Dates", "id"],
)
panel1 = pd.DataFrame(
    np.random.randint(3, size=(9, 2)), index=idx1, columns=["ytm", "mat"]
)
idx2 = pd.MultiIndex.from_product(
    [["2017-05-02", "2017-05-03", "2017-05-04"], ["0", "1", "2"]], names=["Dates", "yr"]
)
panel2 = pd.DataFrame(np.random.randint(3, size=9), index=idx2, columns=["curve"])
print(panel1)
#                 ytm  mat
# Dates      id
# 2017-05-02 id1    2    1
#            id2    1    2
#            id3    0    0
# 2017-05-03 id1    2    1
#            id2    0    1
#            id3    1    1
# 2017-05-04 id1    2    2
#            id2    2    0
#            id3    1    0
print(panel2)
#                curve
# Dates      yr
# 2017-05-02 0       0
#            1       1
#            2       2
# 2017-05-03 0       1
#            1       2
#            2       0
# 2017-05-04 0       1
#            1       2
#            2       0
panel1["CDM"] = panel2.loc[
    pd.MultiIndex.from_arrays(
        [panel1.index.get_level_values(0), panel1.mat.astype(str).rename("yr")]
    )
].to_numpy()
print(panel1)
#                 ytm  mat  CDM
# Dates      id
# 2017-05-02 id1    2    1    1
#            id2    1    2    2
#            id3    0    0    0
# 2017-05-03 id1    2    1    2
#            id2    0    1    2
#            id3    1    1    2
# 2017-05-04 id1    2    2    0
#            id2    2    0    1
#            id3    1    0    1

EDIT

To compare mat and yr as float and to use .reindex instead of .loc.

import pandas as pd
import numpy as np

np.random.seed(12)
idx1 = pd.MultiIndex.from_product(
    [["2017-05-02", "2017-05-03", "2017-05-04"], ["id1", "id2", "id3"]],
    names=["Dates", "id"],
)
panel1 = pd.DataFrame(
    np.random.randint(3, size=(9, 2)), index=idx1, columns=["ytm", "mat"]
)
panel1.iloc[0, 1] = np.nan
idx2 = pd.MultiIndex.from_product(
    [["2017-05-02", "2017-05-03", "2017-05-04"], ["0", "1", "2"]], names=["Dates", "yr"]
)
panel2 = pd.DataFrame(np.random.randint(3, size=9), index=idx2, columns=["curve"])
panel2 = panel2.rename(float, level=1)
print(panel1)
#                 ytm  mat
# Dates      id
# 2017-05-02 id1    2  NaN
#            id2    1  2.0
#            id3    0  0.0
# 2017-05-03 id1    2  1.0
#            id2    0  1.0
#            id3    1  1.0
# 2017-05-04 id1    2  2.0
#            id2    2  0.0
#            id3    1  0.0
print(panel2)
#                 curve
# Dates      yr
# 2017-05-02 0.0      0
#            1.0      1
#            2.0      2
# 2017-05-03 0.0      1
#            1.0      2
#            2.0      0
# 2017-05-04 0.0      1
#            1.0      2
#            2.0      0
panel1["CDM"] = panel2.reindex(
    pd.MultiIndex.from_arrays(
        [panel1.index.get_level_values(0), panel1.mat.rename("yr")]
    )
).to_numpy()
print(panel1)
#                 ytm  mat  CDM
# Dates      id
# 2017-05-02 id1    2  NaN  NaN
#            id2    1  2.0  2.0
#            id3    0  0.0  0.0
# 2017-05-03 id1    2  1.0  2.0
#            id2    0  1.0  2.0
#            id3    1  1.0  2.0
# 2017-05-04 id1    2  2.0  0.0
#            id2    2  0.0  1.0
#            id3    1  0.0  1.0

In order to generate any non-empty and repeatable result of my code, I changed a little the way how both panels are created:

np.random.seed(0)
idx1 = pd.MultiIndex.from_product([['2017-05-02', '2017-05-03', '2017-05-04'],
    ['id1', 'id2', 'id3']], names=['Dates', 'id'])
panel1 = pd.DataFrame({'ytm': np.random.randn(9),
    'mat': [0.5, 0.82, 1.06, -0.27, 1.5, 0.59, 0.62, 1.89, 2.5]}, index=idx1)
idx2 = pd.MultiIndex.from_product([['2017-05-02', '2017-05-03', '2017-05-04'],
    [0.5, 1.5, 2.5]], names=['Dates', 'yr'])
panel2 = pd.DataFrame(np.random.randn(9), index=idx2, columns=['curve'])

The changes comprise:

  • np.random.seed - te get repatable results.
  • Only ytm column of panel1 is created as random numbers. In order to have some matching values in mat, I put predefined values there, providing one match with yr for each date.
  • Level 1 of idx2 is of float type. Your sample includes strings, which obviously will not be equal to mat values.

I also suppose that for each group from panel1 the look for match should be performed in rows from panel2 with the same date (not in groups for all dates).

To generate the result (CDB column), proceed as follows:

  1. Define a function generating CDB column for the current group of rows (for each date):

     def getCDB(grp):
         cdb = panel2.xs(grp.index[0][0], level=0).reindex(grp.mat).curve
         return pd.Series(cdb.values, index=grp.index)
    
  2. Then apply it and save the result in a new column:

     panel1['CDB'] = panel1.groupby(level=0).apply(getCDB)\
         .reset_index(level=0, drop=True)
    

For my input data, the result is:

                     ytm   mat       CDB
Dates      id                           
2017-05-02 id1  1.764052  0.50  0.410599
           id2  0.400157  0.82       NaN
           id3  0.978738  1.06       NaN
2017-05-03 id1  2.240893 -0.27       NaN
           id2  1.867558  1.50  0.121675
           id3 -0.977278  0.59       NaN
2017-05-04 id1  0.950088  0.62       NaN
           id2 -0.151357  1.89       NaN
           id3 -0.103219  2.50 -0.205158
Related