2-step decay chain Ni→Co→Fe

#coding=utf8
import numpy as np
from matplotlib import *
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.ticker import AutoMinorLocator
from math import exp
from distinct_colors import get_distinct

''' Decay of 56Ni '''

### PLOTTING PARAMETERS
rcParams['font.size'] = 24.
font = {'family': 'normal', 'size': 24}
rc('axes', linewidth=2)
rc("text", usetex=True)
rc('font', family='serif')
rc('font', serif='Times')
rc('legend', fontsize=22)
rc('xtick.major', size=5, width=2)
rc('ytick.major', size=5, width=2)
rc('xtick.minor', size=3, width=1)
rc('ytick.minor', size=3, width=1)
fig = plt.figure()
fig.subplots_adjust(left=.15, bottom=.15, right=.95, top=.95)
ax = plt.subplot(111)


c = get_distinct(4)
m_sun = 1.989 * 10**33
ages = np.arange(0,600)


# Lifetime (τ = 1/λ) of 56Ni and 56Co
tau_ni = 6.0749 / np.log(2.0) # 56Ni decays to 56Co. Half-life T(1/2) = 6.0749
tau_co = 77.27  / np.log(2.0) # 56Co decays to 56Fe. Half-life T(1/2) = 77.27


M_ni0 = 0.0074 * m_sun
#N_ni0 = (m_sun/56*m_u) * M_ni0/m_sun = 2.141 * 10**55 * (M_ni0/m_sun) ⇒
N_ni0 = 2.141 * 10**55 * (M_ni0/m_sun)


N_ni_array=[]
N_co_array=[]
N_fe_array=[]
N_co_wrong=[]

for t in ages:
    N_ni_array = np.append(N_ni_array,  N_ni0 * np.exp(-t/tau_ni))
    N_co_array = np.append(N_co_array,  N_ni0 * ( tau_co/(tau_co-tau_ni) ) * (   exp(-t/tau_co) - exp(-t/tau_ni)  ) )
    N_co_wrong = np.append(N_co_wrong,  N_ni0 * np.exp(-t/tau_co))
    N_fe_array = np.append(N_fe_array,  N_ni0 * ( 1 + (tau_ni/(tau_co-tau_ni))*exp(-t/tau_ni) - (tau_co/(tau_co-tau_ni))*exp(-t/tau_co) ) )


fill = 0

if fill < 1:

    ## Number of Ni nuclides vs. time
    ax.plot(ages, (N_ni_array),  
            color=c[0],
            lw=3,
            ls='-',
            label='${\\rm N_{Ni}}$')
    ## Number of Co nuclides vs. time
    ax.plot(ages, (N_co_array),  
            color=c[1],
            lw=3,
            ls='--',
            label='${\\rm N_{Co}}$')
    ## Number of Co nuclides vs. time, wrong (for illustration)
    line, = ax.plot(ages, (N_co_wrong),  
            color=c[2],
            lw=3,
            ls='--',
            label='${\\rm N_{Co}}$ --- wrong')
    line.set_dashes([3,3,3,3])
    ## Number of Fe nuclides vs. time
    line, = ax.plot(ages, (N_fe_array),  
            color=c[3],
            lw=3,
            ls='-.',
            label='${\\rm N_{Fe}}$')
    line.set_dashes([15,5,3,5])

    x1=0.17
    x2=0.93
    y1=0.875
    y2=0.185


else:
    ### or show Fill_betweein
    ax.fill_between(ages, 0, N_co_array + N_ni_array + N_fe_array, color=c[2])
    ax.fill_between(ages, 0, N_co_array + N_ni_array, color=c[1], )
    ax.fill_between(ages, 0, N_co_array, color=c[0])

    plt.figtext(0.165,0.68, r'$^{56}{\rm {Ni}}$')
    plt.figtext(0.2,0.3, r'$^{56}{\rm {Co}}$')
    plt.figtext(0.8,0.68, r'$^{56}{\rm {Fe}}$')

    x1=0.17
    x2=0.93
    y1=0.875
    y2=0.875



# The total rate of energy production «e», erg/s
# do not plot with the same ylim, no sense
es = []
for t in ages:
    e = (6.45 * 10**43 * exp(-t/8.8) + 1.45*10**43 * exp(-t/111.3)) * M_ni0/m_sun # erg/s
    es = np.append(es, e)
line, = ax.plot(ages, es,
        color='k',
        lw=3,
        ls='-.',
        label='$E_{{\\rm tot}}\;[{\\rm erg/s]}$')
line.set_dashes([15,5,15,15])


plt.figtext(x1,y1,
            r'Initial M$_{^{56}{\rm Ni}}$ = ' \
            +str("%.4f" % float(M_ni0/m_sun)) \
            +r'$\,$$\,$M$_\odot$', 
            fontsize=20)

plt.figtext(x2,y2, r'$^{56}$Ni$\,\to\,^{56}$Co$\,\to\,^{56}$Fe', ha='right', fontsize=20)

ax.legend(frameon = False, loc='upper right', handlelength=3, fontsize=17)
ax.xaxis.set_minor_locator(AutoMinorLocator(2))
ax.yaxis.set_minor_locator(AutoMinorLocator())

#ax.set_ylim([44,45])

ax.set_xlim([-0,150])
#ax.set_ylim([1e51,2.0e53])
ax.set_ylim([1e51, N_ni0*1.2])


ax.set_xlabel('Days since explosion')
ax.set_ylabel('Number of nuclides')
plt.show()

## Nadezhin 1994: Ni → Co → Fe Decay
#dN_ni/dt = -N_ni/tau_ni
#dN_co/dt = N_ni/tau_ni - N_co/tau_co
#dN_fe/dt = N_co/tau_co
#
#Initial conditions:
#N_ni = N_ni0, N_co=0, N_fe=0    at t=0
#
#Solving gives:
#N_ni = N_ni0 * exp(-t/tau_ni)
#N_co = N_ni0 * ( tau_co/(tau_co-tau_ni) ) * (   exp(-t/tau_co) - exp(-t/tau_ni)  ) 
#N_fe = N_ni0 * ( 1 + (tau_ni/(tau_co-tau_ni))*exp(-t/tau_ni) - (tau_co/(tau_co-tau_ni))*exp(-t/tau_co) )


### HOW TO «UNDECAY» 56NI
#n0 = 5     # number of particles of 56ni at specific age (e.g., 9 days)
#age = 9    # age [days]
#tau_ni = 6.0749 / np.log(2.0)
#
#n0 = n0 * np.exp(age/tau_ni) # initial mass of ni
#
#print n0