I am working on a retirement simulator and trying to create a “nervous” investor option. The goal is to replace three consecutive return values from my list with zero after a negative return rate occurs. These consecutive zeros would then represent the investor pulling out of the market.

For some reason when I add in the commented out “nervous investor” section, the 1st and 3rd samples go missing and my other investor types also appear to be impacted by this zero replacement loop.

def retirement_forecast(investor, investment, years):
    sample_count = 0
    forecast_rates = []
    forecast_samples = {}
    #require the user to enter a minimum of 10 years
    if years < 10:
        return('Please enter a forecast period of at least 10 years')
    else:
        pass
    #creating a sample that follows the same patterns of positive and negative change, as well as negative returns 
    while sample_count < 5:
        forecast_rates = random.sample(investor, years)
        negative = negative_counter(forecast_rates)
        positive_change = positive_change_counter(forecast_rates)
        if investor == aggressive:
            if negative > 10 and positive_change >= 40 and positive_change <= 50:
                column_name = 'Sample ' + str(sample_count+1)
                forecast_samples[column_name] = forecast_rates
                sample_count += 1
        #trying to create a nervous investor where the rate drops to zero for 3 years after a negative rate (to simulate pulling out of the market)
        #somehow I'm dropping my samples with this and it is affecting the other investor types (dropping zeros for those lists)
        # if investor == nervous:
        #     if negative > 10 and positive_change >= 40 and positive_change <= 50:
        #         #if investor is nervous, replace the three consecutive returns after a negative rate with a zero
        #         for i in range(len(forecast_rates)-3):
        #             if forecast_rates[i] < 0:
        #                 forecast_rates[i+1] == 0
        #                 forecast_rates[i+2] == 0
        #                 forecast_rates[i+3] == 0
        #             else:
        #                 continue
        #         column_name = 'Sample ' + str(sample_count+1)
        #         forecast_samples[column_name] = forecast_rates
        #         sample_count += 1
        #conservative or moderate investor
        else:
            if positive_change >= 40 and positive_change <= 50:
                column_name = 'Sample ' + str(sample_count+1)
                forecast_samples[column_name] = forecast_rates
                sample_count += 1