特定列上的Python数据框总和

Question

My task is to calculate the cumulative principal. Using numpy.ppmt it only gives a specific month's paid principal, so I want to add columns which contains monthly principal for each record and then take the sum to get the cumulative principal.

For example, I have a dataframe looks like the following:

frame = pd.DataFrame({'rate':[0.1, 0.1], 'per':[2, 4], 'nper':[360, 360], 'pv':[100000, 200000]})

max_per = frame['per'].max()
columns = ['principal%s'%i for i in range(1, max_per + 1)]
df = pd.DataFrame(index=frame.index, columns=columns, dtype='float').fillna(0)

for index, column in enumerate(columns):
    df[column] = -np.ppmt(rate=frame['rate'] / 100 / 12, per=index + 1, nper=frame['nper'], 
                          pv=frame['pv'], when=when)
frame.join(df)

The dataframe will look like the following:

   nper  per      pv  rate          epp  principal1  principal2  principal3  \
0   360    2  100000   0.1   547.309838  273.643517  273.666321  273.689126   
1   360    4  200000   0.1  2189.421796  547.287034  547.332642  547.378253   

   principal4  
0  273.711934  
1  547.423868

The problem is that for record one, principal3 and principal4 should be 0. One workaround is to calculate the sum of principal1-principal4 based on column 'per', for example if frame.per == 2, then I only sum principal1 and principal 2, and if frame.per == 4, then I only sum principal1 through principal4. Any help to do that.

I can calculate the cumulative principal by calling apply, but I do not want to do that because it is slow.

Thanks.

Answer 1

One possible solution is set 0 by mask before join by boolean mask with comparing arange by length of column with per column for 2d numpy array :

#subtract 1 because python counts from 0
mask = np.arange(len(df.columns)) > frame['per'].values[:, None] - 1
df = frame.join(df.mask(mask, 0))
print(df)
   rate  per  nper      pv  principal1  principal2  principal3  principal4
0   0.1    2   360  100000  273.643517  273.666321    0.000000    0.000000
1   0.1    4   360  200000  547.287034  547.332642  547.378253  547.423868

Another solution with numpy.where :

mask = np.arange(len(df.columns)) > frame['per'].values[:, None] - 1
df[:] = np.where(mask, 0, df)
df = frame.join(df)
print(df)
   rate  per  nper      pv  principal1  principal2  principal3  principal4
0   0.1    2   360  100000  273.643517  273.666321    0.000000    0.000000
1   0.1    4   360  200000  547.287034  547.332642  547.378253  547.423868