import numpy as np
import scipy.stats as sts
import matplotlib.pyplot as plt
n, t = 100000, 200
dn, dt = n//1000, t//10
mu, sigma = 0.05, 0.3
R = np.random.normal(mu, sigma, (1000*dn, 10*dt))
X = R.cumsum(axis=1)
Y = np.exp(X)
S = np.ones((1000*dn, 1))
Z = np.concatenate((S, S*Y), axis=1)
Z.sort(axis=0)
W = Z[0*dn:1000*dn].sum(axis=0)
Q0 = Z[0*dn:100*dn].sum(axis=0)/W
Q1 = Z[100*dn:200*dn].sum(axis=0)/W
Q2 = Z[200*dn:300*dn].sum(axis=0)/W
Q3 = Z[300*dn:400*dn].sum(axis=0)/W
Q4 = Z[400*dn:500*dn].sum(axis=0)/W
Q5 = Z[500*dn:600*dn].sum(axis=0)/W
Q6 = Z[600*dn:700*dn].sum(axis=0)/W
Q7 = Z[700*dn:800*dn].sum(axis=0)/W
Q8 = Z[800*dn:900*dn].sum(axis=0)/W
Q9 = Z[900*dn:1000*dn].sum(axis=0)/W
Q99 = Z[990*dn:1000*dn].sum(axis=0)/W
Q999 = Z[999*dn:1000*dn].sum(axis=0)/W
T = np.arange(0, 10*dt+1, 1)
P9 = sts.norm.cdf(sigma*np.sqrt(T)-sts.norm.ppf(0.9))
P99 = sts.norm.cdf(sigma*np.sqrt(T)-sts.norm.ppf(0.99))
P999 = sts.norm.cdf(sigma*np.sqrt(T)-sts.norm.ppf(0.999))
fig, ax = plt.subplots(2, 3, sharey=True, figsize=(10.0, 5.0))
fig.suptitle('Distribution of wealth over time')
I = ('0', '1', '2', '3', '4', '5', '6', '7', '8', '9')
Q = np.stack((Q0, Q1, Q2, Q3, Q4, Q5, Q6, Q7, Q8, Q9))
ax[0, 0].bar(I, Q[0:10, 0*dt])
ax[0, 1].bar(I, Q[0:10, 1*dt])
ax[0, 2].bar(I, Q[0:10, 2*dt])
ax[1, 0].bar(I, Q[0:10, 3*dt])
ax[1, 1].bar(I, Q[0:10, 4*dt])
ax[1, 2].bar(I, Q[0:10, 5*dt])
fig, ax = plt.subplots(figsize=(10.0, 5.0))
fig.suptitle('Concentration of wealth over time')
ax.plot(Q9, 'bo', T, P9, 'c', markersize=4.0, linewidth=2.0, label='Top 10%')
ax.plot(Q99, 'bo', T, P99, 'c', markersize=2.0, linewidth=1.0, label='Top 1%')
ax.plot(Q999, 'bo', T, P999, 'c', markersize=1.0, linewidth=0.5, label='Top 0.1%')
ax.set(xlabel='Time (years)', ylabel='Proportion of total wealth')
ax.set_xticks(np.arange(0, 10*dt+1, dt))
ax.legend(loc='lower right')
ax.grid()
plt.show()