<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Fast chem_evol
Created by Benoit Côté
This notebook presents and tests the pre_calculate_SSPs implementation in chem_evol.py.  When many timesteps are required in OMEGA, the computational time becomes significant for a simple one-zone chemical evolution model. This fast version pre-calculates the simple stellar populations (SSPs) and the interpolation coefficients. This means the code does not need to fold the yields with the initial mass function (IMF) at each timestep. However, this option prevents you from modifying the IMF properties during a simulation.
Step1:  Original Version
Step2:  Fast Version
Step3:  By using the dt_in_SSPs array, the OMEGA timesteps can be different from the SSP timesteps. If dt_in_SSPs is not provided when running OMEGA, each SSP will have the same timsteps as OMEGA.
Step4:  Comparison Between the Original and Fast Versions
<ASSISTANT_TASK:>
Python Code:
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from NuPyCEE import omega
from NuPyCEE import sygma
Explanation: Fast chem_evol
Created by Benoit Côté
This notebook presents and tests the pre_calculate_SSPs implementation in chem_evol.py.  When many timesteps are required in OMEGA, the computational time becomes significant for a simple one-zone chemical evolution model. This fast version pre-calculates the simple stellar populations (SSPs) and the interpolation coefficients. This means the code does not need to fold the yields with the initial mass function (IMF) at each timestep. However, this option prevents you from modifying the IMF properties during a simulation.
End of explanation
# Run original OMEGA with 1000 timestesp (this may take a minute ..)
o_ori = omega.omega(galaxy='milky_way', special_timesteps=1000)
Explanation: Original Version
End of explanation
# Let's create the timestep template for each SSP
# Here, you can decide any type of timestep option
s_dt_template = sygma.sygma(special_timesteps=1000)
# Copy the SSP timestep array
dt_in_SSPs = s_dt_template.history.timesteps
Explanation: Fast Version
End of explanation
# Let's pre-calculate the SSPs.
# Here I choose a very low number of OMEGA timesteps.
# I do this because I only want to use this instance 
# to copy the SSPs, so I won't have to recalculate
# them each time I want to run an OMEGA simulation (in the fast version).
# You can ignore the Warning notice.
o_for_SSPs = omega.omega(special_timesteps=2, pre_calculate_SSPs=True, dt_in_SSPs=dt_in_SSPs)
# Let's copy the SSPs array
SSPs_in = [o_for_SSPs.ej_SSP, o_for_SSPs.ej_SSP_coef, o_for_SSPs.dt_ssp, o_for_SSPs.t_ssp]
# SSPs_in[0] --> Mass (in log) ejected for each isotope. It's an array in the form of [nb Z][nb SSP dt][isotope]
# SSPs_in[1] --> Interpolation coefficients of each isotope
# SSPs_in[2] --> List of timesteps
# SSPs_in[3] --> List of galactic ages
# Finally, let's run the fast version (1000 timesteps).
# This should be ~3 times faster
o_fast = omega.omega(galaxy='milky_way', special_timesteps=1000, pre_calculate_SSPs=True, SSPs_in=SSPs_in)
Explanation: By using the dt_in_SSPs array, the OMEGA timesteps can be different from the SSP timesteps. If dt_in_SSPs is not provided when running OMEGA, each SSP will have the same timsteps as OMEGA.
End of explanation
%matplotlib nbagg
o_ori.plot_spectro(color='b',  label='Original')
o_fast.plot_spectro(color='g', shape='--', label='Fast')
plt.xscale('log')
%matplotlib nbagg
o_ori.plot_mass( specie='Mg', color='b', markevery=100000, label='Original')
o_fast.plot_mass(specie='Mg', color='g',  markevery=100,   shape='--', label='Fast')
plt.ylabel('Mass of Mg [Msun]')
%matplotlib nbagg
o_ori.plot_spectro( xaxis='[O/H]', yaxis='[Ca/O]', color='b', \
                    markevery=100000, label='Original')
o_fast.plot_spectro(xaxis='[O/H]', yaxis='[Ca/O]', color='g', \
                    shape='--', markevery=300, label='Fast')
plt.ylim(-2.0,1.0)
Explanation: Comparison Between the Original and Fast Versions
End of explanation
<END_TASK>