<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Dependence on primary cosmic ray flux
Step1:  Create an instance of an MCEqRun class. Most options are defined in the mceq_config module, and do not require change. Look into mceq_config.py or use the documentation.
If the initialization succeeds it will print out some information according to the debug level.
Step2:  Solve and store results
This code below computes fluxes of neutrinos, averaged over all directions, for different primary models.
Step3:  Plot with matplotlib
Step4:  Save as in ASCII file for other types of processing
<ASSISTANT_TASK:>
Python Code:
import matplotlib.pyplot as plt
import numpy as np
#import solver related modules
from MCEq.core import MCEqRun
import mceq_config as config
#import primary model choices
import crflux.models as pm
Explanation: Dependence on primary cosmic ray flux
End of explanation
mceq_run = MCEqRun(
#provide the string of the interaction model
interaction_model='SIBYLL2.3c',
#primary cosmic ray flux model
#support a tuple (primary model class (not instance!), arguments)
primary_model=(pm.HillasGaisser2012, "H3a"),
# Zenith angle in degrees. 0=vertical, 90=horizontal
theta_deg=0.0
)
Explanation: Create an instance of an MCEqRun class. Most options are defined in the mceq_config module, and do not require change. Look into mceq_config.py or use the documentation.
If the initialization succeeds it will print out some information according to the debug level.
End of explanation
# Bump up the debug level to see what the calculation is doing
config.debug_level = 2
#Define equidistant grid in cos(theta)
angles = np.arccos(np.linspace(1,0,11))*180./np.pi
#Power of energy to scale the flux
mag = 3
#obtain energy grid (nver changes) of the solution for the x-axis of the plots
e_grid = mceq_run.e_grid
p_spectrum_flux = []
#Initialize empty grid
for  pmcount, pmodel in enumerate([(pm.HillasGaisser2012,'H3a'),
                                   (pm.HillasGaisser2012,'H4a'),
                                   (pm.GaisserStanevTilav,'3-gen'),
                                   (pm.GaisserStanevTilav,'4-gen')]):
    
    mceq_run.set_primary_model(*pmodel)
    
    flux = {}
    for frac in ['mu_conv','mu_pr','mu_total',
                 'numu_conv','numu_pr','numu_total',
                 'nue_conv','nue_pr','nue_total','nutau_pr']:
        flux[frac] = np.zeros_like(e_grid)
        
    #Sum fluxes, calculated for different angles
    for theta in angles:
        mceq_run.set_theta_deg(theta)
        mceq_run.solve()
        #_conv means conventional (mostly pions and kaons)
        flux['mu_conv'] += (mceq_run.get_solution('conv_mu+', mag)
                           + mceq_run.get_solution('conv_mu-', mag))
        # _pr means prompt (the mother of the muon had a critical energy
        # higher than a D meson. Includes all charm and direct resonance
        # contribution)
        flux['mu_pr'] += (mceq_run.get_solution('pr_mu+', mag)
                         + mceq_run.get_solution('pr_mu-', mag))
        # total means conventional + prompt
        flux['mu_total'] += (mceq_run.get_solution('total_mu+', mag)
                            + mceq_run.get_solution('total_mu-', mag))
        # same meaning of prefixes for muon neutrinos as for muons
        flux['numu_conv'] += (mceq_run.get_solution('conv_numu', mag)
                             + mceq_run.get_solution('conv_antinumu', mag))
        flux['numu_pr'] += (mceq_run.get_solution('pr_numu', mag)
                           + mceq_run.get_solution('pr_antinumu', mag))
        flux['numu_total'] += (mceq_run.get_solution('total_numu', mag)
                              + mceq_run.get_solution('total_antinumu', mag))
        # same meaning of prefixes for electron neutrinos as for muons
        flux['nue_conv'] += (mceq_run.get_solution('conv_nue', mag)
                            + mceq_run.get_solution('conv_antinue', mag))
        flux['nue_pr'] += (mceq_run.get_solution('pr_nue', mag)
                          + mceq_run.get_solution('pr_antinue', mag))
        flux['nue_total'] += (mceq_run.get_solution('total_nue', mag)
                             + mceq_run.get_solution('total_antinue', mag))
        # since there are no conventional tau neutrinos, prompt=total
        flux['nutau_pr'] += (mceq_run.get_solution('total_nutau', mag)
                            + mceq_run.get_solution('total_antinutau', mag))
    #average the results
    for frac in ['mu_conv','mu_pr','mu_total',
                 'numu_conv','numu_pr','numu_total',
                 'nue_conv','nue_pr','nue_total','nutau_pr']:
        flux[frac] = flux[frac]/float(len(angles))
        
    p_spectrum_flux.append((flux,mceq_run.pmodel.sname,mceq_run.pmodel.name)) 
Explanation: Solve and store results
This code below computes fluxes of neutrinos, averaged over all directions, for different primary models.
End of explanation
#get path of the home directory + Desktop
desktop = os.path.join(os.path.expanduser("~"),'Desktop')
for pref, lab in [('numu_',r'\nu_\mu'), 
                  ('mu_',r'\mu'), 
                  ('nue_',r'\nu_e')
                 ]:
    plt.figure(figsize=(4.5, 3.5))
    for (flux, p_sname, p_name), col in zip(p_spectrum_flux,['k','r','g','b','c']):
        
        plt.loglog(e_grid, flux[pref + 'total'], color=col, ls='-', lw=2.5,
                  label=p_sname, alpha=0.4)
        plt.loglog(e_grid, flux[pref + 'conv'], color=col, ls='--', lw=1,
                   label='_nolabel_')
        plt.loglog(e_grid, flux[pref + 'pr'], color=col,ls='-', lw=1, 
                   label='_nolabel_')
    plt.xlim(50,1e9)
    plt.ylim(1e-5,1)
    plt.xlabel(r"$E_{{{0}}}$ [GeV]".format(lab))
    plt.ylabel(r"$\Phi_{" + lab + "}$ (E/GeV)$^{" + str(mag) +" }$" + 
               "(cm$^{2}$ s sr GeV)$^{-1}$")
    plt.legend(loc='upper right')
    plt.tight_layout()
    # Uncoment if you want to save the plot
    # plt.savefig(os.path.join(desktop,pref + 'flux.pdf'))
Explanation: Plot with matplotlib
End of explanation
for (flux, p_sname, p_name) in p_spectrum_flux:
    np.savetxt(open(os.path.join(desktop, 'numu_flux_' + p_sname + '.txt'),'w'),
    zip(e_grid, 
        flux['mu_conv'],flux['mu_pr'],flux['mu_total'],
        flux['numu_conv'],flux['numu_pr'],flux['numu_total'],
        flux['nue_conv'],flux['nue_pr'],flux['nue_total'],
        flux['nutau_pr']),
    fmt='%6.5E',
    header=('lepton flux scaled with E**{0}. Order (E, mu_conv, mu_pr, mu_total, ' +
            'numu_conv, numu_pr, numu_total, nue_conv, nue_pr, nue_total, ' +
            'nutau_pr').format(mag)
    )
Explanation: Save as in ASCII file for other types of processing
End of explanation
<END_TASK>