Simple script for estimating bankruptcy probabilities during retirement

bankruptcy.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Feb  5 11:01:52 2017

@author: andyjones
"""

import scipy as sp
import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
import pandas as pd
from itertools import product

def simulate_savings(drawdown, ret, years, vol=0.15, trials=5000):
    savings = [sp.ones(trials)]
    for i in range(years):
        returns = sp.random.normal(ret, vol, (trials,))
        after_expenses = sp.clip(savings[-1] - drawdown, 0, None)
        after_returns = sp.clip(after_expenses + returns*after_expenses, 0, None)
        savings.append(after_returns)
    savings = sp.vstack(savings)
    return savings
    
    
def failure_probability(drawdown, ret, years, **kwargs):
    savings = simulate_savings(drawdown, ret, years=years, **kwargs)
    
    return (savings == 0).mean(1)[-1]

def failure_probabilities(rets):
    results = {}
    for ret in rets:
        print(ret)
        rates = sp.linspace(0, 0.1, 51)
        years = sp.linspace(5, 50, 46, dtype=int)
        results[ret] = {}
        for drawdown, year in product(rates, years):
            results[ret][(year, drawdown)] = failure_probability(drawdown, ret, years=year)
        results[ret] = pd.Series(results[ret]).unstack()
    results = pd.Panel(results)
    results.major_axis.name = 'years of retirement'
    results.minor_axis.name = 'annual expenditure'
    
    return results
    
def plot_failure_rates():
    results = failure_probabilities(sp.linspace(0, .1, 11))

    sns.set_style('white')
    fig, axes = plt.subplots(3, 3, figsize=(22, 18))
    for ax, ret in zip(axes.flatten(), results.items[1:]):
        im = ax.imshow(100*results[ret], 
                       cmap=plt.cm.viridis, 
                       interpolation='nearest',
                       extent=(100*results[ret].columns[0], 
                               100*results[ret].columns[-1], 
                               results[ret].index[0], 
                               results[ret].index[-1]),
                       aspect='auto',
                       origin='lower',
                       vmin=0, vmax=100)
        ax.grid(True, alpha=0.2)
        ax.set_title('{:.0f}% annual return'.format(100*ret))        
        ax.set_xlabel(results[ret].columns.name)
        ax.set_ylabel(results[ret].index.name)
        plt.colorbar(im, ax=ax, format='%.0f%%')
        
        ax.xaxis.set_major_formatter(mpl.ticker.FormatStrFormatter('%.0f%%'))
        
    fig.suptitle('Probability of bankruptcy at different rates of expenditure and return\n(Independent normal returns, 15% annual volatility)', fontsize=20)
    fig.savefig('failure_prob.png', bbox_inches='tight')
    
if __name__ == '__main__':
    plot_failure_rates()