How to populate data frame based on previous row values with pandas?

Viewed 54

I have the following dataframe:

df = pd.DataFrame(data={
    "id": ['a', 'd'],
    "amount": [3000, 4000],
    "rate": [0.2, 0.3],
    "date": ["2022-07-21", "2022-08-11"],
    "months": [4, 5],
})

I want to create installments for the given trades. Output should look like this:

  id    interest   principal        date       balance  installment
0  a   50.000000  731.508250  2022-07-21  2.268492e+03   781.508250
1  a   37.808196  743.700055  2022-08-21  1.524792e+03   781.508250
2  a   25.413195  756.095055  2022-09-21  7.686966e+02   781.508250
3  a   12.811611  768.696640  2022-10-21  0.00000        781.508250
4  d  100.000000  760.987444  2022-08-11  3.239013e+03   860.987444
5  d   80.975314  780.012130  2022-09-11  2.459000e+03   860.987444
6  d   61.475011  799.512433  2022-10-11  1.659488e+03   860.987444
7  d   41.487200  819.500244  2022-11-11  8.399877e+02   860.987444
8  d   20.999694  839.987750  2022-12-11  0.00000        860.987444

Key point here is that the next row values depend on previous balance value. The first time balance is amount of the source data frame.

My current solution:

import numpy_financial as npf
import pandas as pd
from dateutil.relativedelta import relativedelta

df = pd.DataFrame(data={
    "id": ['a', 'd'],
    "amount": [3000, 4000],
    "rate": [0.2, 0.3],
    "date": ["2022-07-21", "2022-08-11"],
    "months": [4, 5],
})


def get_output_df(df):
    columns = ["id", "interest", "principal", "date", "balance"]
    output_df = pd.DataFrame(columns=columns)

    for _, row in df.iterrows():
        rate = row["rate"] / 12
        months = row["months"]
        amount = row["amount"]
        date = pd.to_datetime(row["date"]).date()
        installment_amount = npf.pmt(rate=rate, nper=months, pv=-amount)
        prior_balance = amount

        loan_installment_data = []
        for i in range(months):
            interest_amount = rate * prior_balance
            principal_amount = installment_amount - interest_amount
            balance = prior_balance - principal_amount

            loan_installment_data.append(
                {
                    "id": row["id"],
                    "interest": interest_amount,
                    "principal": principal_amount,
                    "date": date,
                    "installment": installment_amount,
                    "balance": balance
                }
            )

            prior_balance = balance
            date += relativedelta(months=1)

        output_df = output_df.append(loan_installment_data, ignore_index=True)
    return output_df


output_df = get_output_df(df)

Is there any pandas feature i could use to do the same implementation?

1 Answers

You can try using pandas shift function, which shifts row values by the value specified (so can be previous row). By using shift, you can have the previous row's amount for example available to the next row, and do your calculations without looping through each row and saving each previous value in a temporary variable.

Here is an example with your data:

import pandas as pd

df = pd.DataFrame(data={
    "id": ['a', 'd'],
    "amount": [3000, 4000],
    "rate": [0.2, 0.3],
    "date": ["2022-07-21", "2022-08-11"],
    "months": [4, 5],
})

df["prior_balance"] = df["amount"].shift(1, fill_value=0)
df

Output:

output of code above

Using shift, the first row would be a NaN, so you can use fill_value=0 to set that NaN value to 0 as there is no previous amount for example.

Related