diff --git a/nuVeto/__init__.py b/nuVeto/__init__.py index 973d850..3c456d7 100644 --- a/nuVeto/__init__.py +++ b/nuVeto/__init__.py @@ -1 +1,2 @@ -__all__ = ['mu', 'utils', 'nuveto', 'barr_uncertainties', 'external','examples'] +__all__ = ['mu', 'utils', 'nuveto', + 'barr_uncertainties', 'external', 'examples'] diff --git a/nuVeto/examples/plots.py b/nuVeto/examples/plots.py index 0f9c678..8498059 100644 --- a/nuVeto/examples/plots.py +++ b/nuVeto/examples/plots.py @@ -17,12 +17,12 @@ def tex(inp): if isinstance(inp, str): categ, daughter = inp.split() - daughter_trans = {'nu_mu':r'$\nu_\mu$', - 'nu_e':r'$\nu_e$', - 'nu_tau':r'$\nu_\tau$', - 'nu_mubar':r'$\overline{\nu}_\mu$', - 'nu_ebar':r'$\overline{\nu}_e$', - 'nu_taubar':r'$\overline{\nu}_\tau$'} + daughter_trans = {'nu_mu': r'$\nu_\mu$', + 'nu_e': r'$\nu_e$', + 'nu_tau': r'$\nu_\tau$', + 'nu_mubar': r'$\overline{\nu}_\mu$', + 'nu_ebar': r'$\overline{\nu}_e$', + 'nu_taubar': r'$\overline{\nu}_\tau$'} return r'{} {}'.format(categ, daughter_trans[daughter]) else: if inp > 1: @@ -35,26 +35,28 @@ def extlabel(corr_only): if corr_only: return r'$\mathcal{P}_{\rm pass}^{\rm cor, SGRS}$' else: - return r'$\mathcal{P}_{\rm pass}^{\rm cor, SGRS} \mathcal{P}_{\rm pass}^{\rm uncor, GJKvS}$' - - + return r'$\mathcal{P}_{\rm pass}^{\rm cor, SGRS} \mathcal{P}_{\rm pass}^{\rm uncor, GJKvS}$' + + # passing fraction tests def fn(slice_val): """ decide which fn to run depending on slice_val """ - return pr_enu if slice_val <=1 else pr_cth + return pr_enu if slice_val <= 1 else pr_cth def pr_enu(cos_theta=1., kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', barr_mods=(), depth=1950*Units.m, density=('CORSIKA', ('SouthPole', 'December')), accuracy=3.5, fraction=True, prpl='ice_allm97_step_1', corr_only=False, **kwargs): """ plot the passing rate (flux or fraction) """ - ens = np.logspace(3,7,100) if corr_only else np.logspace(3,7,39) - passed = [passing(en, cos_theta, kind, pmodel, hadr, barr_mods, depth, density, accuracy, fraction, prpl, corr_only) for en in ens] + ens = np.logspace(3, 7, 100) if corr_only else np.logspace(3, 7, 39) + passed = [passing(en, cos_theta, kind, pmodel, hadr, barr_mods, depth, + density, accuracy, fraction, prpl, corr_only) for en in ens] if fraction: passed_fn = interpolate.interp1d(ens, passed, kind='quadratic') else: - passed_fn = lambda es: 10**interpolate.interp1d(ens, np.log10(passed), kind='quadratic')(es) - ens_plot = np.logspace(3,7,100) + def passed_fn(es): return 10**interpolate.interp1d(ens, + np.log10(passed), kind='quadratic')(es) + ens_plot = np.logspace(3, 7, 100) if fraction: prs = plt.plot(ens_plot, passed_fn(ens_plot), **kwargs) plt.ylim(0., 1.) @@ -72,8 +74,9 @@ def pr_enu(cos_theta=1., kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a') def pr_cth(enu=1e5, kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', barr_mods=(), depth=1950*Units.m, density=('CORSIKA', ('SouthPole', 'December')), accuracy=3.5, fraction=True, prpl='ice_allm97_step_1', corr_only=False, **kwargs): """ plot the passing rate (flux or fraction) """ - cths = np.linspace(0,1,21) - passed = [passing(enu, cos_theta, kind, pmodel, hadr, barr_mods, depth, density, accuracy, fraction, prpl, corr_only) for cos_theta in cths] + cths = np.linspace(0, 1, 21) + passed = [passing(enu, cos_theta, kind, pmodel, hadr, barr_mods, depth, + density, accuracy, fraction, prpl, corr_only) for cos_theta in cths] if fraction: prs = plt.plot(cths, passed, **kwargs) plt.ylim(0., 1.) @@ -103,7 +106,8 @@ def brackets(slice_val=1., kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a uppers = [BARR[param].error for param in params] lowers = [-BARR[param].error for param in params] all_barr_mods = [tuple(zip(params, uppers)), tuple(zip(params, lowers))] - pr = fn(slice_val)(slice_val, kind, pmodel, hadr, label=f'{tex(kind)} {tex(slice_val)}') + pr = fn(slice_val)(slice_val, kind, pmodel, hadr, + label=f'{tex(kind)} {tex(slice_val)}') for barr_mods in all_barr_mods: fn(slice_val)(slice_val, kind, pmodel, hadr, barr_mods, fraction=fraction, color=pr.get_color(), alpha=1-abs(barr_mods[0][-1])) @@ -112,11 +116,13 @@ def brackets(slice_val=1., kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a def samples(slice_val=1, kind='nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', fraction=True, seed=88, nsamples=10, params='g h1 h2 i w6 y1 y2 z ch_a ch_b ch_e'): params = params.split(' ') - pr = fn(slice_val)(slice_val, kind, pmodel, hadr=hadr, label=f'{tex(kind)} {tex(slice_val)}') + pr = fn(slice_val)(slice_val, kind, pmodel, hadr=hadr, + label=f'{tex(kind)} {tex(slice_val)}') np.random.seed(seed) for i in range(nsamples-1): # max(-1, throw) prevents throws that dip below -100% - errors = [max(-1, np.random.normal(scale=BARR[param].error)) for param in params] + errors = [max(-1, np.random.normal(scale=BARR[param].error)) + for param in params] barr_mods = tuple(zip(params, errors)) fn(slice_val)(slice_val, kind, pmodel, hadr, barr_mods, color=pr.get_color(), alpha=1-min(np.mean(np.abs(errors)), 0.9)) @@ -124,7 +130,7 @@ def samples(slice_val=1, kind='nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), had def accuracy(slice_val=1., kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', fraction=True): plt.clf() - accuracies = [2,3,4] + accuracies = [2, 3, 4] for accuracy in accuracies: fn(slice_val)(slice_val, kind, pmodel=pmodel, hadr=hadr, accuracy=accuracy, fraction=fraction, @@ -141,15 +147,15 @@ def prpls(slice_val=1., kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), def elbert(slice_val=1., kind='conv nu_mu', pmodel=(pm.GaisserHonda, None), prpl='ice_allm97_step_1', corr_only=False): - hadrs=['DPMJET-III-3.0.6', 'SIBYLL2.3c'] + hadrs = ['DPMJET-III-3.0.6', 'SIBYLL2.3c'] echoice = exthp.corr if corr_only else exthp.passrates if slice_val > 1: - cths = np.linspace(0,1, 100) + cths = np.linspace(0, 1, 100) emu = extsv.minimum_muon_energy(extsv.overburden(cths)) plt.plot(cths, echoice(kind)(slice_val, emu, cths), 'k--', label=f'Analytic approx. {tex(kind)} {tex(slice_val)}') else: - ens = np.logspace(2,9, 100) + ens = np.logspace(2, 9, 100) emu = extsv.minimum_muon_energy(extsv.overburden(slice_val)) plt.plot(ens, echoice(kind)(ens, emu, slice_val), 'k--', label=f'Analytic approx. {tex(kind)} {tex(slice_val)}') @@ -168,18 +174,18 @@ def elbert_pmodels(slice_val=1., kind='conv nu_mu', hadr='DPMJET-III-3.0.6', prp echoice = exthp.corr if corr_only else exthp.passrates label = f'{extlabel(corr_only)} {tex(kind)} {tex(slice_val)}' if slice_val > 1: - cths = np.linspace(0,1, 100) + cths = np.linspace(0, 1, 100) emu = extsv.minimum_muon_energy(extsv.overburden(cths)) plt.plot(cths, echoice(kind)(slice_val, emu, cths), 'k--', label=label) else: - ens = np.logspace(2,9, 100) + ens = np.logspace(2, 9, 100) emu = extsv.minimum_muon_energy(extsv.overburden(slice_val)) plt.plot(ens, echoice(kind)(ens, emu, slice_val), 'k--', label=label) for pmodel in pmodels: pr = fn(slice_val)(slice_val, kind, pmodel[:2], hadr, prpl=prpl, corr_only=corr_only, - label=f'{pmodel[2]} {tex(kind)} {tex(slice_val)}') + label=f'{pmodel[2]} {tex(kind)} {tex(slice_val)}') plt.legend() plt.tight_layout(0.3) @@ -198,35 +204,37 @@ def pmodels(slice_val=1., kind='conv nu_mu', hadr='SIBYLL2.3c', prpl='ice_allm97 def density_models(slice_val=1., kind='conv nu_mu', hadr='SIBYLL2.3c', prpl='ice_allm97_step_1', fraction=True): models = [('CORSIKA', ('BK_USStd', None)), ('CORSIKA', ('SouthPole', 'December')), - ('MSIS00_IC',('SouthPole', 'June')), - ('MSIS00_IC',('SouthPole', 'January'))] + ('MSIS00_IC', ('SouthPole', 'June')), + ('MSIS00_IC', ('SouthPole', 'January'))] for model in models: pr = fn(slice_val)(slice_val, kind, density=model, hadr=hadr, prpl=prpl, fraction=fraction, label=f'{model[0]} {model[1][0]} {model[1][1]} {tex(kind)} {tex(slice_val)}') plt.legend() -def corsika(cos_theta_bin=-1, kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3', density=('CORSIKA',('SouthPole', 'December')), prpl='ice_allm97_step_1', corsika_file='eff_maxmu', plot_nuveto_lines = False, plot_legacy_veto_lines = False): +def corsika(cos_theta_bin=-1, kind='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3', density=('CORSIKA', ('SouthPole', 'December')), prpl='ice_allm97_step_1', corsika_file='eff_maxmu', plot_nuveto_lines=False, plot_legacy_veto_lines=False): if isinstance(cos_theta_bin, list): - [corsika(cth, kind, pmodel, hadr, density, prpl, corsika_file, plot_nuveto_lines, plot_legacy_veto_lines) for cth in cos_theta_bin] + [corsika(cth, kind, pmodel, hadr, density, prpl, corsika_file, + plot_nuveto_lines, plot_legacy_veto_lines) for cth in cos_theta_bin] return - cfile = pickle.load(open(resources.files('nuVeto') / 'data' / 'corsika' / f'{corsika_file}.pkl', 'rb'), encoding='latin1') + cfile = pickle.load(open(resources.files('nuVeto') / 'data' / + 'corsika' / f'{corsika_file}.pkl', 'rb'), encoding='latin1') fraction = 'eff' in corsika_file eff, elow, eup, xedges, yedges = cfile[mceq_categ_format(kind)] cos_theta = centers(yedges)[cos_theta_bin] - clean = eff[:,cos_theta_bin]>0 + clean = eff[:, cos_theta_bin] > 0 xcenters = 10**centers(xedges)[clean] - pr = plt.errorbar(xcenters, eff[:,cos_theta_bin][clean], - xerr=np.asarray(list(zip(xcenters-10**xedges[:-1][clean], - 10**xedges[1:][clean]-xcenters))).T, - yerr=np.asarray(list(zip(elow[:,cos_theta_bin][clean], - eup[:,cos_theta_bin][clean]))).T, - label=f'CORSIKA {tex(kind)} {tex(cos_theta)}', - fmt='.') + pr = plt.errorbar(xcenters, eff[:, cos_theta_bin][clean], + xerr=np.asarray(list(zip(xcenters-10**xedges[:-1][clean], + 10**xedges[1:][clean]-xcenters))).T, + yerr=np.asarray(list(zip(elow[:, cos_theta_bin][clean], + eup[:, cos_theta_bin][clean]))).T, + label=f'CORSIKA {tex(kind)} {tex(cos_theta)}', + fmt='.') if plot_legacy_veto_lines and fraction: - ens = np.logspace(2,9, 100) + ens = np.logspace(2, 9, 100) emu = extsv.minimum_muon_energy(extsv.overburden(cos_theta)) plt.plot(ens, exthp.passrates(kind)(ens, emu, cos_theta), 'k--', label=f'Analytic approx. {tex(kind)} {tex(cos_theta)}') @@ -243,13 +251,16 @@ def dndee(mother, daughter): # print x_range[0], x_range[-1] x_samp = np.logspace(1, -9, 5000) - c = plt.plot(x_samp, dNdEE_interp(x_samp), label = f"Interpolated {mother} to {daughter}") - plt.plot(x_range,dNdEE, '.', color=c[0].get_color(), label = f"MCEq {mother} to {daughter}") + c = plt.plot(x_samp, dNdEE_interp(x_samp), + label=f"Interpolated {mother} to {daughter}") + plt.plot(x_range, dNdEE, '.', color=c[0].get_color( + ), label=f"MCEq {mother} to {daughter}") plt.semilogx() plt.xlabel(r"$x=E_\nu/E_p$") plt.ylabel(r"$ \frac{dN}{dE_\nu} E_p$") # plt.ylim(-0.1, 5.1) - plt.axvline(1-ParticleProperties.rr(mother, daughter), linestyle='--', color=c[0].get_color()) + plt.axvline(1-ParticleProperties.rr(mother, daughter), + linestyle='--', color=c[0].get_color()) plt.legend() @@ -263,16 +274,18 @@ def plot_prpl(interp_pkl, include_mean=False, include_cbar=True): xx, yy = np.meshgrid(emui, l_ice) prpls = prplfn(list(zip(xx.flatten(), yy.flatten()))) plt.figure() - plt.pcolormesh(emui_edges, l_ice_edges/1e3, prpls.reshape(xx.shape), cmap='magma') + plt.pcolormesh(emui_edges, l_ice_edges/1e3, + prpls.reshape(xx.shape), cmap='magma') if include_cbar: plt.colorbar() if include_mean: - small_ice = l_ice[l_ice<2.7e4] + small_ice = l_ice[l_ice < 2.7e4] plt.plot(extsv.minimum_muon_energy(small_ice), small_ice/1e3, 'w--', label=r'$l_{\rm ice,\,median} (E_\mu^{\rm i}, E_\mu^{\rm th} = 1\,{\rm TeV})$') - leg = plt.legend(frameon=False, prop={'weight':'bold'}, loc='upper left') + leg = plt.legend(frameon=False, prop={ + 'weight': 'bold'}, loc='upper left') for text in leg.get_texts(): - plt.setp(text, color = 'w', fontsize='medium') + plt.setp(text, color='w', fontsize='medium') plt.xlabel(r'$E_\mu^{\rm i}$ [GeV]') plt.ylabel(r'$l_{\rm ice}$ [km]') plt.locator_params(axis='y', nbins=8) @@ -293,14 +306,16 @@ def plot_prpl(interp_pkl, include_mean=False, include_cbar=True): axr.set_xscale('log') axr.set_xlim(1e2, 1e8) axr.minorticks_off() - xlocmaj = LogLocator(base=10,numticks=12) + xlocmaj = LogLocator(base=10, numticks=12) axr.get_xaxis().set_major_locator(xlocmaj) return emui_edges, l_ice_edges, prpls.reshape(xx.shape) def plot_prpl_ratio(interp_pkl_num, interp_pkl_den, include_cbar=True): - emui_edges, l_ice_edges, prpls_num = plot_prpl(interp_pkl_num, False, False) - emui_edges, l_ice_edges, prpls_den = plot_prpl(interp_pkl_den, False, False) + emui_edges, l_ice_edges, prpls_num = plot_prpl( + interp_pkl_num, False, False) + emui_edges, l_ice_edges, prpls_den = plot_prpl( + interp_pkl_den, False, False) plt.figure() plt.pcolormesh(emui_edges, l_ice_edges/1e3, np.ma.masked_invalid(prpls_num/prpls_den), norm=colors.LogNorm(vmin=1e-2, vmax=1e2), @@ -312,7 +327,7 @@ def plot_prpl_ratio(interp_pkl_num, interp_pkl_den, include_cbar=True): # plt.yscale('log') # plt.gca().yaxis.set_major_formatter(ScalarFormatter()) plt.xscale('log') - xlocmaj = LogLocator(base=10,numticks=12) + xlocmaj = LogLocator(base=10, numticks=12) plt.gca().xaxis.set_major_locator(xlocmaj) plt.gca().minorticks_off() plt.ticklabel_format(style='plain', axis='y') @@ -321,7 +336,7 @@ def plot_prpl_ratio(interp_pkl_num, interp_pkl_den, include_cbar=True): plt.title('{}/{}'.format(os.path.splitext(os.path.basename(interp_pkl_num))[0], os.path.splitext(os.path.basename(interp_pkl_den))[0])) - + def parent_ratio(cos_theta, parents='pi+ pi-', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', ecr=None, particle=None): plt.figure() @@ -341,21 +356,22 @@ def parent_ratio(cos_theta, parents='pi+ pi-', pmodel=(pm.HillasGaisser2012, 'H3 if ecr is None or particle is None: plt.title(r'$\cos \theta = {:.2g}$'.format(cos_theta)) else: - plt.title(r'{} at {:.2g} GeV and $\cos \theta = {:.2g}$'.format(particle, ecr, cos_theta)) + plt.title(r'{} at {:.2g} GeV and $\cos \theta = {:.2g}$'.format( + particle, ecr, cos_theta)) def parent_flux(cos_theta, parent='D0', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', - density=('CORSIKA',('SouthPole', 'December')), mag=3, ecr=None, particle=None): + density=('CORSIKA', ('SouthPole', 'December')), mag=3, ecr=None, particle=None): plt.figure() sv = nuVeto(cos_theta, pmodel, hadr, density=density) gsol = sv.grid_sol(ecr, particle) - ### hack to convert list to np array + # hack to convert list to np array sv.mceq.grid_sol = np.asarray(sv.mceq.grid_sol) ### X_vec = sv.X_vec for idx, x_val in enumerate(X_vec): mceq = sv.mceq.get_solution(parent, mag, grid_idx=idx) - mceq[mceq==0] = np.nan + mceq[mceq == 0] = np.nan calc = sv.get_solution(parent, gsol, mag, grid_idx=idx) pout = plt.loglog(sv.mceq.e_grid, mceq, label='$X={:.3g}$ km'.format( @@ -370,20 +386,21 @@ def parent_flux(cos_theta, parent='D0', pmodel=(pm.HillasGaisser2012, 'H3a'), ha plt.title(f'{parent} {tex(cos_theta)}') plt.tight_layout(0.3) # plt.savefig('/Users/tianlu/Desktop/selfveto/parent_flux/combined/{}.png'.format(parent)) - -def nu_flux(cos_theta, kinds='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', mag=3, logxlim=(3,7), corr_only=False): + +def nu_flux(cos_theta, kinds='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', mag=3, logxlim=(3, 7), corr_only=False): sv = nuVeto(cos_theta, pmodel, hadr) sv.grid_sol() - fig, axs = plt.subplots(2,1) + fig, axs = plt.subplots(2, 1) for kind in [_.strip() for _ in kinds.split(',')]: plt.sca(axs[0]) - mine = np.asarray([fluxes(en, cos_theta, kind, pmodel, hadr, corr_only=corr_only)[1] for en in sv.mceq.e_grid]) + mine = np.asarray([fluxes(en, cos_theta, kind, pmodel, hadr, corr_only=corr_only)[ + 1] for en in sv.mceq.e_grid]) pr = plt.plot(sv.mceq.e_grid, mine*sv.mceq.e_grid**mag, label=f'{hadr} {tex(kind)} {tex(cos_theta)}') plt.ylabel(r'$E_\nu^{} \Phi_\nu$'.format(mag)) plt.loglog() - plt.xlim(*np.power(10,logxlim)) + plt.xlim(*np.power(10, logxlim)) plt.ylim(ymin=1e-8) plt.legend() @@ -401,7 +418,7 @@ def nu_flux(cos_theta, kinds='conv nu_mu', pmodel=(pm.HillasGaisser2012, 'H3a'), plt.xscale('log') plt.ylim(0.5, 1.9) plt.xlabel(r'$E_\nu$ [GeV]') - plt.xlim(*np.power(10,logxlim)) + plt.xlim(*np.power(10, logxlim)) def prob_nomu(cos_theta, particle=14, pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', prpl='ice_allm97_step_1'): @@ -412,7 +429,7 @@ def prob_nomu(cos_theta, particle=14, pmodel=(pm.HillasGaisser2012, 'H3a'), hadr sv = nuVeto(cos_theta, pmodel, hadr) nmu = [sv.nmu(ecr, particle, prpl) for ecr in ecrs] nmufn = interpolate.interp1d(ecrs, nmu, kind='linear', - assume_sorted=True, fill_value=(0,np.nan)) + assume_sorted=True, fill_value=(0, np.nan)) plt.semilogx(ecrs_fine, np.exp(-nmufn(ecrs_fine)), label='interpolated') plt.semilogx(ecrs, np.exp(-np.asarray(nmu)), 'ko') plt.xlabel(r'$E_{CR} [GeV]$') @@ -430,15 +447,15 @@ def elbert_only(slice_val=1., kind='conv nu_mu'): names = [extlabel(True), extlabel(False)] if slice_val > 1: - cths = np.linspace(0,1, 100) + cths = np.linspace(0, 1, 100) emu = extsv.minimum_muon_energy(extsv.overburden(cths)) - for echoice, name in zip(echoices,names): + for echoice, name in zip(echoices, names): plt.plot(cths, echoice(kind)(slice_val, emu, cths), '--', label=name) else: - ens = np.logspace(2,9, 100) + ens = np.logspace(2, 9, 100) emu = extsv.minimum_muon_energy(extsv.overburden(slice_val)) - for echoice, name in zip(echoices,names): + for echoice, name in zip(echoices, names): plt.plot(ens, echoice(kind)(ens, emu, slice_val), '--', label=name) plt.title(r'{}, {}'.format(tex(kind), tex(slice_val))) @@ -457,7 +474,7 @@ def hist_preach(infile, plotdir=None): df = pd.read_csv(infile, delim_whitespace=True, header=None, names='ei l ef'.split()) # If the muon doesn't reach, MMC saves ef as -distance traveled - df[df<0] = 0 + df[df < 0] = 0 for idx, (ei, sdf) in enumerate(df.groupby('ei')): if idx % napf == 0: if idx > 0: @@ -466,14 +483,14 @@ def hist_preach(infile, plotdir=None): if plotdir is not None: plt.savefig(os.path.join( os.path.expanduser(plotdir), f'{(idx - 1) / napf}.png')) - fig, axs = plt.subplots(6, 6, figsize=(10,10)) + fig, axs = plt.subplots(6, 6, figsize=(10, 10)) # fig.text(0.5, 0.04, r'$E_f$', ha='center', va='center') # fig.text(0.06, 0.5, r'$P(E_f|E_i, l)$', ha='center', va='center', rotation='vertical') axs = axs.flatten() - ax = axs[idx%napf] - ax.set_prop_cycle('color',plt.cm.Blues(np.linspace(0.3,1,100))) + ax = axs[idx % napf] + ax.set_prop_cycle('color', plt.cm.Blues(np.linspace(0.3, 1, 100))) plt.sca(ax) - plt.title(r'${:.2g}$ GeV'.format(ei), fontdict={'fontsize':8}) + plt.title(r'${:.2g}$ GeV'.format(ei), fontdict={'fontsize': 8}) for l, efs in sdf.groupby('l'): bins = calc_bins(efs['ef']) counts, edges = np.histogram(efs['ef'], bins=bins, density=True) diff --git a/nuVeto/nuveto.py b/nuVeto/nuveto.py index 25dba21..7532cd2 100644 --- a/nuVeto/nuveto.py +++ b/nuVeto/nuveto.py @@ -22,6 +22,7 @@ class nuVeto(object): """Class for computing the neutrino passing fraction i.e. (1-(Veto probability))""" + def __init__(self, costh, pmodel=(pm.HillasGaisser2012, 'H3a'), hadr='SIBYLL2.3c', barr_mods=(), depth=1950*Units.m, @@ -59,7 +60,7 @@ def __init__(self, costh, 211, -211, 321, -321, 130, - -2212, -2112]#, 11, 22] + -2212, -2112] # , 11, 22] config.ctau = 2.5 self.mceq = MCEqRun( # provide the string of the interaction model @@ -72,11 +73,11 @@ def __init__(self, costh, # atmospheric density model density_model=density) - if len(barr_mods) > 0: for barr_mod in barr_mods: # Modify proton-air -> mod[0] - self.mceq.set_mod_pprod(2212, BARR[barr_mod[0]].pdg, barr_unc, barr_mod) + self.mceq.set_mod_pprod( + 2212, BARR[barr_mod[0]].pdg, barr_unc, barr_mod) # Populate the modifications to the matrices by re-filling the interaction matrix self.mceq.regenerate_matrices(skip_decay_matrix=True) @@ -104,9 +105,11 @@ def categ_to_mothers(categ, daughter): if 'nu_tau' in daughter: mothers = [f"D{rcharge}", f"D_s{rcharge}"] else: - mothers = [f"D{rcharge}", f"D_s{rcharge}", f"D{rbar}0"]#, 'Lambda'+lbar+'0']#, 'Lambda_c'+rcharge] + # , 'Lambda'+lbar+'0']#, 'Lambda_c'+rcharge] + mothers = [f"D{rcharge}", f"D_s{rcharge}", f"D{rbar}0"] elif categ == 'total': - mothers = nuVeto.categ_to_mothers('conv', daughter)+nuVeto.categ_to_mothers('pr', daughter) + mothers = nuVeto.categ_to_mothers( + 'conv', daughter)+nuVeto.categ_to_mothers('pr', daughter) else: mothers = [categ,] return mothers @@ -142,9 +145,11 @@ def nbody(fpath, esamp, enu, fn, l_ice): ddec = interpolate.RegularGridInterpolator((xnus, xmus), vals, bounds_error=False, fill_value=None) - emu_mat = xmus[:,None]*esamp[None,:]*Units.GeV + emu_mat = xmus[:, None]*esamp[None, :]*Units.GeV pmu_mat = ddec(np.stack(np.meshgrid(enu/esamp, xmus), axis=-1)) - reaching = 1-np.sum(pmu_mat*fn.prpl(np.stack([emu_mat, np.ones(emu_mat.shape)*l_ice], axis=-1)), axis=0) + reaching = 1 - \ + np.sum( + pmu_mat*fn.prpl(np.stack([emu_mat, np.ones(emu_mat.shape)*l_ice], axis=-1)), axis=0) reaching[reaching < 0.] = 0. return reaching @@ -157,19 +162,23 @@ def psib(l_ice, mother, enu, accuracy, prpl): fn = MuonProb(prpl) if mother in ['D0', 'D0-bar']: reaching = nuVeto.nbody( - resources.files('nuVeto') / 'data' / 'decay_distributions' / 'D0_numu.npz', + resources.files('nuVeto') / 'data' / + 'decay_distributions' / 'D0_numu.npz', esamp, enu, fn, l_ice) elif mother in ['D+', 'D-']: reaching = nuVeto.nbody( - resources.files('nuVeto') / 'data' / 'decay_distributions' / 'D+_numu.npz', + resources.files('nuVeto') / 'data' / + 'decay_distributions' / 'D+_numu.npz', esamp, enu, fn, l_ice) elif mother in ['Ds+', 'Ds-']: reaching = nuVeto.nbody( - resources.files('nuVeto') / 'data' / 'decay_distributions' / 'Ds_numu.npz', + resources.files('nuVeto') / 'data' / + 'decay_distributions' / 'Ds_numu.npz', esamp, enu, fn, l_ice) elif mother == 'K0L': reaching = nuVeto.nbody( - resources.files('nuVeto') / 'data' / 'decay_distributions' / 'K0L_numu.npz', + resources.files('nuVeto') / 'data' / + 'decay_distributions' / 'K0L_numu.npz', esamp, enu, fn, l_ice) else: # Assuming muon energy is E_parent - E_nu @@ -196,7 +205,8 @@ def get_dNdEE(self, mother, daughter): x_low = x_range[x_range < 5e-2] dNdEE_low = np.array([dNdEE[good][-1]]*x_low.size) - dNdEE_interp = lambda x_: interpolate.pchip( + + def dNdEE_interp(x_): return interpolate.pchip( np.concatenate([[1-rr], x_range[good], x_low])[::-1], np.concatenate([[dNdEE_edge], dNdEE[good], dNdEE_low])[::-1], extrapolate=True)(x_) * np.heaviside(1-rr-x_, 1) @@ -217,7 +227,9 @@ def nmu(self, ecr, particle, prpl='ice_allm97_step_1'): """Poisson probability of getting no muons""" grid_sol = self.grid_sol(ecr, particle) l_ice = self.geom.overburden(self.costh) - mu = np.abs(self.get_solution('mu-', grid_sol)) + np.abs(self.get_solution('mu+', grid_sol)) # np.abs hack to prevent negative fluxes + # np.abs hack to prevent negative fluxes + mu = np.abs(self.get_solution('mu-', grid_sol)) + \ + np.abs(self.get_solution('mu+', grid_sol)) fn = MuonProb(prpl) coords = list(zip(self.mceq.e_grid*Units.GeV, [l_ice]*len(self.mceq.e_grid))) @@ -227,19 +239,22 @@ def nmu(self, ecr, particle, prpl='ice_allm97_step_1'): @lru_cache(maxsize=2**12) def get_rescale_phi(self, mother, ecr=None, particle=None): """Flux of the mother at all heights""" - grid_sol = self.grid_sol(ecr, particle) # MCEq solution (fluxes tabulated as a function of height) + grid_sol = self.grid_sol( + ecr, particle) # MCEq solution (fluxes tabulated as a function of height) dX = self.dX_vec*Units.gr/Units.cm**2 rho = self.mceq.density_model.X2rho(self.X_vec)*Units.gr/Units.cm**3 - inv_decay_length_array = (ParticleProperties.mass_dict[mother] / (self.mceq.e_grid[:,None] * Units.GeV)) / (ParticleProperties.lifetime_dict[mother]*rho[None,:]) - rescale_phi = dX[None,:]* inv_decay_length_array * self.get_solution(mother, grid_sol, grid_idx=False).T + inv_decay_length_array = (ParticleProperties.mass_dict[mother] / ( + self.mceq.e_grid[:, None] * Units.GeV)) / (ParticleProperties.lifetime_dict[mother]*rho[None, :]) + rescale_phi = dX[None, :] * inv_decay_length_array * \ + self.get_solution(mother, grid_sol, grid_idx=False).T return rescale_phi def get_integrand(self, categ, daughter, enu, accuracy, prpl, ecr=None, particle=None): """flux*yield""" esamp = self.esamp(enu, accuracy) mothers = self.categ_to_mothers(categ, daughter) - nums = np.zeros((len(esamp),len(self.X_vec))) - dens = np.zeros((len(esamp),len(self.X_vec))) + nums = np.zeros((len(esamp), len(self.X_vec))) + dens = np.zeros((len(esamp), len(self.X_vec))) for mother in mothers: dNdEE = self.get_dNdEE(mother, daughter)[-1] rescale_phi = self.get_rescale_phi(mother, ecr, particle) @@ -247,22 +262,22 @@ def get_integrand(self, categ, daughter, enu, accuracy, prpl, ecr=None, particle ### # TODO: optimize to only run when esamp[0] is non-zero rescale_phi = np.exp(np.array([interpolate.interp1d( - np.log(self.mceq.e_grid[rescale_phi[:,i]>0]), - np.log(rescale_phi[:,i][rescale_phi[:,i]>0]), + np.log(self.mceq.e_grid[rescale_phi[:, i] > 0]), + np.log(rescale_phi[:, i][rescale_phi[:, i] > 0]), kind='quadratic', bounds_error=False, fill_value=-np.inf)(np.log(esamp)) - if np.count_nonzero(rescale_phi[:,i] > 0) > 2 - else [-np.inf,]*esamp.shape[0] - for i in range(rescale_phi.shape[1])])).T + if np.count_nonzero(rescale_phi[:, i] > 0) > 2 + else [-np.inf,]*esamp.shape[0] + for i in range(rescale_phi.shape[1])])).T if 'nu_mu' in daughter: # muon accompanies nu_mu only pnmsib = self.psib(self.geom.overburden(self.costh), mother, enu, accuracy, prpl) else: - pnmsib = np.ones(len(esamp)) + pnmsib = np.ones(len(esamp)) dnde = dNdEE(enu/esamp)/esamp - nums += (dnde * pnmsib)[:,None]*rescale_phi - dens += (dnde)[:,None]*rescale_phi + nums += (dnde * pnmsib)[:, None]*rescale_phi + dens += (dnde)[:, None]*rescale_phi return nums, dens @@ -295,22 +310,22 @@ def get_solution(self, p_pdg = ParticleProperties.pdg_id[particle_name] reduce_res = True - if grid_idx is None: # Surface only case + if grid_idx is None: # Surface only case sol = np.array([grid_sol[-1]]) xv = np.array([self.X_vec[-1]]) - elif isinstance(grid_idx, bool) and not grid_idx: # Whole solution case + elif isinstance(grid_idx, bool) and not grid_idx: # Whole solution case sol = np.asarray(grid_sol) xv = np.asarray(self.X_vec) reduce_res = False - elif grid_idx >= len(self.mceq.grid_sol): # Surface only case + elif grid_idx >= len(self.mceq.grid_sol): # Surface only case sol = np.array([grid_sol[-1]]) xv = np.array([self.X_vec[-1]]) - else: # Particular height case + else: # Particular height case sol = np.array([grid_sol[grid_idx]]) xv = np.array([self.X_vec[grid_idx]]) # MCEq solution for particle - direct = sol[:,ref[particle_name].lidx: + direct = sol[:, ref[particle_name].lidx: ref[particle_name].uidx] res = np.zeros(direct.shape) rho_air = 1./self.mceq.density_model.r_X2rho(xv) @@ -324,37 +339,38 @@ def get_solution(self, # number of targets per cm2 ndens = rho_air*Units.Na/Units.mol_air sec = self.mceq.pman[p_pdg] - prim2mceq = {'p+-bar':'pbar-', - 'n0-bar':'nbar0', - 'D0-bar':'Dbar0', - 'Lambda0-bar':'Lambdabar0'} + prim2mceq = {'p+-bar': 'pbar-', + 'n0-bar': 'nbar0', + 'D0-bar': 'Dbar0', + 'Lambda0-bar': 'Lambdabar0'} for prim in self.projectiles(): if prim in prim2mceq: _ = prim2mceq[prim] else: _ = prim - prim_flux = sol[:,ref[_].lidx: + prim_flux = sol[:, ref[_].lidx: ref[_].uidx] proj = self.mceq.pman[ParticleProperties.pdg_id[prim]] prim_xs = proj.inel_cross_section() try: int_yields = proj.hadr_yields[sec] - res += np.sum(int_yields[None,:,:]*prim_flux[:,None,:]*prim_xs[None,None,:]*ndens[:,None,None], axis=2) + res += np.sum(int_yields[None, :, :]*prim_flux[:, None, :] + * prim_xs[None, None, :]*ndens[:, None, None], axis=2) except KeyError as e: continue - res *= decayl[None,:] + res *= decayl[None, :] # combine with direct res[direct != 0] = direct[direct != 0] if particle_name[:-1] == 'mu': for _ in [f"k_{particle_name}", f"pi_{particle_name}"]: - res += sol[:,ref[f"{_}_l"].lidx: + res += sol[:, ref[f"{_}_l"].lidx: ref[f"{_}_l"].uidx] - res += sol[:,ref[f"{_}_r"].lidx: + res += sol[:, ref[f"{_}_r"].lidx: ref[f"{_}_r"].uidx] - res *= self.mceq.e_grid[None,:] ** mag + res *= self.mceq.e_grid[None, :] ** mag if reduce_res: res = res[0] @@ -375,7 +391,8 @@ def get_fluxes(self, enu, kind='conv nu_mu', accuracy=3.5, prpl='ice_allm97_step total = 0 if corr_only: # sum performs the dX integral - nums, dens = self.get_integrand(categ, daughter, enu, accuracy, prpl) + nums, dens = self.get_integrand( + categ, daughter, enu, accuracy, prpl) num = np.sum(nums, axis=1) den = np.sum(dens, axis=1) passed = integrate.trapezoid(num, esamp) @@ -384,7 +401,7 @@ def get_fluxes(self, enu, kind='conv nu_mu', accuracy=3.5, prpl='ice_allm97_step pmodel = self.pmodel[0](self.pmodel[1]) - #loop over primary particles + # loop over primary particles for particle in pmodel.nucleus_ids: # A continuous input energy range is allowed between # :math:`50*A~ \\text{GeV} < E_\\text{nucleus} < 10^{10}*A \\text{GeV}`. @@ -401,27 +418,29 @@ def get_fluxes(self, enu, kind='conv nu_mu', accuracy=3.5, prpl='ice_allm97_step # nmufn --> fine grid interpolation of pnm nmufn = interpolate.interp1d(ecrs, nmu, kind='linear', assume_sorted=True, bounds_error=False, - fill_value=(0,np.nan)) + fill_value=(0, np.nan)) # nums --> numerator nums = [] # dens --> denominator dens = [] # istart --> integration starting point, the lowest energy index for the integral istart = max(0, np.argmax(ecrs > enu) - 1) - for ecr in ecrs[istart:]: # integral in primary energy (E_CR) + for ecr in ecrs[istart:]: # integral in primary energy (E_CR) # cr_flux --> cosmic ray flux # phim2 --> units of flux * m^2 (look it up in the units) cr_flux = pmodel.nucleus_flux(particle, ecr.item())*Units.phim2 # poisson exp(-Nmu) [last term in eq 12] pnmarr = np.exp(-nmufn(ecr-esamp)) - num_ecr = 0 # single entry in nums - den_ecr = 0 # single entry in dens + num_ecr = 0 # single entry in nums + den_ecr = 0 # single entry in dens # dEp # integral in Ep - nums_ecr, dens_ecr = self.get_integrand(categ, daughter, enu, accuracy, prpl, ecr, particle) - num_ecr = integrate.trapz(np.sum(nums_ecr, axis=1)*pnmarr, esamp) + nums_ecr, dens_ecr = self.get_integrand( + categ, daughter, enu, accuracy, prpl, ecr, particle) + num_ecr = integrate.trapz( + np.sum(nums_ecr, axis=1)*pnmarr, esamp) den_ecr = integrate.trapz(np.sum(dens_ecr, axis=1), esamp) nums.append(num_ecr*cr_flux/Units.phicm2) diff --git a/nuVeto/resources/mu/mu.py b/nuVeto/resources/mu/mu.py index cb63b9e..40bc726 100755 --- a/nuVeto/resources/mu/mu.py +++ b/nuVeto/resources/mu/mu.py @@ -19,14 +19,15 @@ def hist_preach(infile): df = pd.read_csv(infile, delim_whitespace=True, header=None, names='ei l ef'.split()) # If the muon doesn't reach, MMC saves ef as -distance traveled - df[df<0] = 0 + df[df < 0] = 0 preach = [] for (ei, l), efs in df.groupby(['ei', 'l']): bins = calc_bins(efs['ef']) histo = Hist(*np.histogram(efs['ef'], bins=bins, density=True)) - [preach.append((ei, l, ef, ew, val)) for ef,ew,val in zip(centers(histo.edges), - np.ediff1d(histo.edges), - histo.counts)] + [preach.append((ei, l, ef, ew, val)) for ef, ew, val in zip(centers(histo.edges), + np.ediff1d( + histo.edges), + histo.counts)] return np.asarray(preach) @@ -51,14 +52,14 @@ def interp(preach, plight): intg = int_ef(preach, plight) df = pd.DataFrame(intg, columns='ei l prpl'.split()) df = df.pivot_table(index='ei', columns='l', values='prpl').fillna(0) - return interpolate.RegularGridInterpolator((df.index,df.columns), df.values, bounds_error=False, fill_value=None) + return interpolate.RegularGridInterpolator((df.index, df.columns), df.values, bounds_error=False, fill_value=None) if __name__ == '__main__': parser = argparse.ArgumentParser( description='Generate muon detection probability') parser.add_argument('mmc', metavar='MMC', - help='text file or pickled histogram containing MMC simulated data') + help='text file or pickled histogram containing MMC simulated data') parser.add_argument('--plight', default='pl_step_1000', choices=[fn for fn in dir(pl) if fn.startswith('pl_')], help='choice of a plight function to apply as defined in pl.py') diff --git a/nuVeto/resources/mu/pl.py b/nuVeto/resources/mu/pl.py index 780de4b..bac0392 100644 --- a/nuVeto/resources/mu/pl.py +++ b/nuVeto/resources/mu/pl.py @@ -44,7 +44,8 @@ def pl_noearlymu(emu): '/Users/tianlu/Projects/icecube/studies/hydrangea/lh/n_str_prob_pspl_full_flat_caus_30_combined.txt', sep=' ') nstrarr = nstrarr.reshape(2, nstrarr.shape[0]/2) # take P(Nstr|Ee) likelihood from chaack - nstr = lambda ee: 10**interpolate.interp1d( + + def nstr(ee): return 10**interpolate.interp1d( np.log10(nstrarr[0]), np.log10(nstrarr[1]), kind='linear', bounds_error=False, fill_value='extrapolate')(np.log10(ee)) return 1-nstr(emu) diff --git a/nuVeto/resources/pythia/pythia_decay.py b/nuVeto/resources/pythia/pythia_decay.py index 2a1cc09..8b328af 100644 --- a/nuVeto/resources/pythia/pythia_decay.py +++ b/nuVeto/resources/pythia/pythia_decay.py @@ -9,27 +9,27 @@ pythia.readString("Random:seed = 0") pythia.init() -UnstableParents = {411:"D+",-411:"D-"} +UnstableParents = {411: "D+", -411: "D-"} for parent in UnstableParents: # make parents unstable - pythia.particleData.mayDecay(parent,True) + pythia.particleData.mayDecay(parent, True) -StableDaughters = {-13:"mu+",13:"mu-", - 15:"tau-",-15:"tau+", - 14:"numu",-14:"antinumu", - 12:"nue",-12:"antinue", - 16:"nutau",-16:"antinutau", - 211:"pi+",-211:"pi-", - 321:"k+",321:"k-", - 130:"k0l",310:"k0s"} +StableDaughters = {-13: "mu+", 13: "mu-", + 15: "tau-", -15: "tau+", + 14: "numu", -14: "antinumu", + 12: "nue", -12: "antinue", + 16: "nutau", -16: "antinutau", + 211: "pi+", -211: "pi-", + 321: "k+", 321: "k-", + 130: "k0l", 310: "k0s"} -GuysWeWant = {-13:"mu+",13:"mu-", - 14:"numu",-14:"antinumu",} +GuysWeWant = {-13: "mu+", 13: "mu-", + 14: "numu", -14: "antinumu", } for daughter in StableDaughters: # make daughters stable - pythia.particleData.mayDecay(daughter,False) + pythia.particleData.mayDecay(daughter, False) number_of_decays = 1000000 @@ -38,35 +38,36 @@ for parent in UnstableParents: mass = pythia.particleData.m0(parent) - momenta = 1.e3 # GeV + momenta = 1.e3 # GeV energy = np.sqrt(mass*mass + momenta*momenta) - p4vec = pythia8.Vec4(momenta,0.,0.,energy) # GeV + p4vec = pythia8.Vec4(momenta, 0., 0., energy) # GeV resultsNuMu[parent] = [] resultsMuons[parent] = [] - for i in range (number_of_decays): + for i in range(number_of_decays): # clean previous stuff pythia.event.reset() # pdgid, status, col, acol, p, m # status = 91 : make normal decay products # col, acol : color and anticolor indices - pythia.event.append(parent,91,0,0,p4vec,mass) + pythia.event.append(parent, 91, 0, 0, p4vec, mass) # go all the way pythia.forceHadronLevel() for j in range(pythia.event.size()): - if(not pythia.event[j].isFinal()): + if (not pythia.event[j].isFinal()): # if not done decaying continue decaying continue - if(pythia.event[j].id() in [-14,14]): # is neutrino type I care about + if (pythia.event[j].id() in [-14, 14]): # is neutrino type I care about for sister in pythia.event[j].sisterList(): - if(pythia.event[sister].id() in [-13,13]): # is muon + if (pythia.event[sister].id() in [-13, 13]): # is muon resultsNuMu[parent].append(pythia.event[j].e()/energy) - resultsMuons[parent].append(pythia.event[sister].e()/energy) + resultsMuons[parent].append( + pythia.event[sister].e()/energy) -plt.figure(figsize=(6,5)) -H, xedges, yedges, fig = plt.hist2d(np.array(resultsNuMu[421]),np.array(resultsMuons[421]), - bins=[np.arange(0.,1.001,0.01),np.arange(0.,1.001,0.01)]) +plt.figure(figsize=(6, 5)) +H, xedges, yedges, fig = plt.hist2d(np.array(resultsNuMu[421]), np.array(resultsMuons[421]), + bins=[np.arange(0., 1.001, 0.01), np.arange(0., 1.001, 0.01)]) plt.close() H_norm_rows = H / H.max(axis=1, keepdims=True) @@ -75,22 +76,23 @@ ycenters = yedges[:-1]+np.diff(yedges) histogramas = [] -plt.figure(figsize=(7,5)) -for i,row in enumerate(H_norm_rows): - #if xcenters[i] in [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]: - norm = np.sum(row) - hist, bin_edges, fig = plt.hist(xcenters,bins=xedges,weights=row/norm,histtype="step", lw = 2, label = r"$E_\nu/E_D =" + str(ycenters[i]) + "$") - histogramas.append(hist) +plt.figure(figsize=(7, 5)) +for i, row in enumerate(H_norm_rows): + # if xcenters[i] in [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]: + norm = np.sum(row) + hist, bin_edges, fig = plt.hist(xcenters, bins=xedges, weights=row/norm, + histtype="step", lw=2, label=r"$E_\nu/E_D =" + str(ycenters[i]) + "$") + histogramas.append(hist) histogramas = np.array(histogramas) plt.legend() -plt.xlim(0,1) -plt.xlabel(r"$E_\mu/E_D$", fontsize = 20) +plt.xlim(0, 1) +plt.xlabel(r"$E_\mu/E_D$", fontsize=20) plt.close() ceros = np.zeros(len(xcenters)) ceros[0] = 1. -reshape_hist = np.vstack([histogramas,ceros]) - -np.savez("DP_meson_decay_distributions_with_two_bodies", histograms = reshape_hist, xedges = xedges) +reshape_hist = np.vstack([histogramas, ceros]) +np.savez("DP_meson_decay_distributions_with_two_bodies", + histograms=reshape_hist, xedges=xedges) diff --git a/nuVeto/utils.py b/nuVeto/utils.py index 4f7b5e2..405abfc 100644 --- a/nuVeto/utils.py +++ b/nuVeto/utils.py @@ -10,12 +10,12 @@ class Units(object): # units - Na = 6.0221415e+23 # mol/cm^3 - km = 5.0677309374099995 # km to GeV^-1 value from SQuIDS + Na = 6.0221415e+23 # mol/cm^3 + km = 5.0677309374099995 # km to GeV^-1 value from SQuIDS cm = km*1.e-5 m = km*1.e-3 - gr = 5.62e23 # gr to GeV value from SQuIDS - sec = 1523000.0 #$ sec to GeV^-1 from SQuIDS + gr = 5.62e23 # gr to GeV value from SQuIDS + sec = 1523000.0 # $ sec to GeV^-1 from SQuIDS day = sec*60**2*24 yr = day*356.25 GeV = 1 @@ -30,8 +30,11 @@ class Units(object): class ParticleProperties(object): modtab = SibyllParticleTable() pd = PYTHIAParticleData() - - mass_dict = {}; lifetime_dict = {}; pdg_id = {}; sibling = {}; + + mass_dict = {} + lifetime_dict = {} + pdg_id = {} + sibling = {} for k in modtab.part_table: pdg_id[k] = modtab.modname2pdg[k] @@ -50,9 +53,8 @@ def rr(mother, daughter): mass_tot = sum([ParticleProperties.pd.mass(abs(prod)) for prod in prod])-ParticleProperties.pd.mass(daughter_pdg) other_masses.append(mass_tot) - - return (min(other_masses)/ParticleProperties.mass_dict[mother])**2 + return (min(other_masses)/ParticleProperties.mass_dict[mother])**2 @staticmethod def br_2body(mother, daughter): @@ -63,7 +65,7 @@ def br_2body(mother, daughter): brs = 0 for br, prod in ParticleProperties.pd.decay_channels(mother_pdg): if daughter_pdg in prod and len(prod) == 2: - brs += br + brs += br return brs @@ -74,8 +76,8 @@ def __init__(self, pklfile): elif os.path.isfile(pklfile): self.mu_int = pickle.load(open(pklfile, 'rb'), encoding='latin1') else: - self.mu_int = pickle.load(open(resources.files('nuVeto') / 'data' / 'prpl' / f'{pklfile}.pkl', 'rb'), encoding='latin1') - + self.mu_int = pickle.load(open(resources.files( + 'nuVeto') / 'data' / 'prpl' / f'{pklfile}.pkl', 'rb'), encoding='latin1') def median_emui(self, distance): """ @@ -88,20 +90,18 @@ def median_emui(self, distance): b, c = 2.52151, 7.13834 return 1e3 * np.exp(1e-3 * distance / (b) + 1e-8 * (distance**2) / c) - def median_approx(self, coord): coord = np.asarray(coord) - muon_energy, ice_distance = coord[...,0], coord[...,1] + muon_energy, ice_distance = coord[..., 0], coord[..., 1] min_mue = self.median_emui(ice_distance)*Units.GeV return muon_energy > min_mue - def prpl(self, coord): pdets = self.mu_int(coord) pdets[pdets > 1] = 1 return pdets - - + + class Geometry(EarthGeometry): def __init__(self, depth): """ Depth of detector and elevation of surface above sea-level @@ -114,7 +114,6 @@ def __init__(self, depth): self.r_top = self.r_E + self.h_atm self.r_obs = self.r_E + self.h_obs - def overburden(self, cos_theta): """Returns the overburden for a detector at *depth* below some surface at *elevation*. @@ -131,7 +130,6 @@ def overburden(self, cos_theta): z = r-d return (np.sqrt(z**2*cos_theta**2+d*(2*r-d))-z*cos_theta)/Units.m - def overburden_to_cos_theta(self, l): """Returns the theta for a given overburden for a detector at *depth* below some surface at *elevation*. @@ -148,7 +146,6 @@ def overburden_to_cos_theta(self, l): z = r-d return (2*d*r-d**2-l**2)/(2*l*z) - def cos_theta_eff(self, cos_theta): """ Returns the effective cos_theta relative the the normal at earth surface. @@ -158,15 +155,15 @@ def cos_theta_eff(self, cos_theta): r = self.r_E z = r-d return np.sqrt(1-(z/r)**2*(1-cos_theta**2)) - - + + def amu(particle): """ :param particle: primary particle's corsika id :returns: the atomic mass of particle """ - return 1 if particle==14 else particle/100 + return 1 if particle == 14 else particle/100 def centers(x): @@ -177,7 +174,7 @@ def calc_nbins(x): ptile = np.percentile(x, 75) - np.percentile(x, 25) if ptile == 0: return 10 - n = (np.max(x) - np.min(x)) / (2 * len(x)**(-1./3) * ptile) + n = (np.max(x) - np.min(x)) / (2 * len(x)**(-1./3) * ptile) return np.floor(n) @@ -185,8 +182,8 @@ def calc_bins(x): nbins = calc_nbins(x) edges = np.linspace(np.min(x), np.max(x)+2, num=nbins+1) # force a bin for 0 efs - if edges[0] == 0 and edges[1]>1: - edges = np.concatenate((np.asarray([0., 1.]),edges[1:])) + if edges[0] == 0 and edges[1] > 1: + edges = np.concatenate((np.asarray([0., 1.]), edges[1:])) return edges diff --git a/paper/paper.py b/paper/paper.py index 441649a..dfbb37a 100644 --- a/paper/paper.py +++ b/paper/paper.py @@ -15,10 +15,10 @@ plt.style.use('paper.mplstyle') linestyles = ['-', '--', ':', '-.'] -titling = {'conv nu_mu':r'Conventional $\nu_\mu$', - 'conv nu_e':r'Conventional $\nu_e$', - 'pr nu_mu':r'Prompt $\nu_\mu$', - 'pr nu_e':r'Prompt $\nu_e$'} +titling = {'conv nu_mu': r'Conventional $\nu_\mu$', + 'conv nu_e': r'Conventional $\nu_e$', + 'pr nu_mu': r'Prompt $\nu_\mu$', + 'pr nu_e': r'Prompt $\nu_e$'} def save(fname): @@ -31,7 +31,7 @@ def save(fname): def earth_attenuation(enu, cos_theta, kind='conv nu_mu'): nufate = os.path.expanduser('~/projects/nuFATE/') - flavor_dict = {'numu':2,'nue':1,'antinue':-1,'antinumu':-2} + flavor_dict = {'numu': 2, 'nue': 1, 'antinue': -1, 'antinumu': -2} flavor = flavor_dict[kind.split('_')[-1]] zenith = np.arccos(cos_theta) @@ -39,23 +39,27 @@ def earth_attenuation(enu, cos_theta, kind='conv nu_mu'): sys.path.append(nufate) import cascade as cas import earth - w,v,ci,energy_nodes,phi_0 = cas.get_eigs(flavor, - os.path.join(nufate, - 'data/phiHGextrap.dat'), - os.path.join(nufate, - 'data/NuFATECrossSections.h5')) + w, v, ci, energy_nodes, phi_0 = cas.get_eigs(flavor, + os.path.join(nufate, + 'data/phiHGextrap.dat'), + os.path.join(nufate, + 'data/NuFATECrossSections.h5')) sys.path.pop() t = earth.get_t_earth(zenith)*Units.Na - phisol = np.dot(v,(ci*np.exp(w*t)))/phi_0 - phisolfn = interp1d(energy_nodes, phisol, kind='quadratic', assume_sorted=True) + phisol = np.dot(v, (ci*np.exp(w*t)))/phi_0 + phisolfn = interp1d(energy_nodes, phisol, + kind='quadratic', assume_sorted=True) return phisolfn(enu) - + def fig_prpl(): - step1000 = resources.files('nuVeto') / 'data' / 'prpl' / 'ice_allm97_step_1.pkl' - step750 = resources.files('nuVeto') / 'data' / 'prpl' / 'ice_allm97_step_0.75.pkl' - sigmoid = resources.files('nuVeto') / 'data' / 'prpl' / 'ice_allm97_sigmoid_0.75_0.25.pkl' + step1000 = resources.files('nuVeto') / 'data' / \ + 'prpl' / 'ice_allm97_step_1.pkl' + step750 = resources.files('nuVeto') / 'data' / \ + 'prpl' / 'ice_allm97_step_0.75.pkl' + sigmoid = resources.files('nuVeto') / 'data' / \ + 'prpl' / 'ice_allm97_sigmoid_0.75_0.25.pkl' plt.figure() plots.plot_prpl(step1000, True, False) plt.title(r'Heaviside $\cal P_{\rm light}$') @@ -72,9 +76,9 @@ def fig_prpl(): plt.tight_layout(0.3) save('fig/prpl_sigmoid.png') - + def fig_prpl_cbar(): - plt.figure(figsize=(5,1.5)) + plt.figure(figsize=(5, 1.5)) norm = mpl.colors.Normalize(vmin=0, vmax=1) cb = mpl.colorbar.ColorbarBase(plt.gca(), cmap='magma', norm=norm, orientation='horizontal') @@ -100,7 +104,8 @@ def fig_prs_ratio(): bb = plots.fn(cos_th)(cos_th, kind, prpl='ice_bb_step_1', linestyle='--', color=allm97.get_color()) - plt.plot(bb.get_data()[0], allm97.get_data()[1]/bb.get_data()[1], ':', color=allm97.get_color()) + plt.plot(bb.get_data()[0], allm97.get_data()[ + 1]/bb.get_data()[1], ':', color=allm97.get_color()) plt.axvline(np.nan, color='k', linestyle='-', label=labels[0]) @@ -125,7 +130,8 @@ def fig_prs(): plt.title(titling[kind]) for idx, prpl in enumerate(prpls): for cos_th in cos_ths: - clabel = r'$\cos \theta_z = {}$'.format(cos_th) if idx == 0 else None + clabel = r'$\cos \theta_z = {}$'.format( + cos_th) if idx == 0 else None plots.fn(cos_th)(cos_th, kind, prpl=prpl, label=clabel, linestyle=linestyles[idx]) @@ -140,7 +146,8 @@ def fig_prs(): def fig_pls(): kinds = ['conv nu_e', 'pr nu_e', 'conv nu_mu', 'pr nu_mu'] cos_ths = [0.25, 0.85] - prpls = ['ice_allm97_step_1', 'ice_allm97_step_0.75', 'ice_allm97_sigmoid_0.75_0.25'] + prpls = ['ice_allm97_step_1', 'ice_allm97_step_0.75', + 'ice_allm97_sigmoid_0.75_0.25'] labels = [r'$\Theta(E_\mu^{\rm f} - 1\,{\rm TeV})$', r'$\Theta(E_\mu^{\rm f} - 0.75\,{\rm TeV})$', r'$\Phi\left(\frac{E_\mu^{\rm f} - 0.75\,{\rm TeV}}{0.25\,{\rm TeV}}\right)$'] @@ -149,7 +156,8 @@ def fig_pls(): plt.title(titling[kind]) for idx, prpl in enumerate(prpls): for cos_th in cos_ths: - clabel = r'$\cos \theta_z = {}$'.format(cos_th) if idx == 0 else None + clabel = r'$\cos \theta_z = {}$'.format( + cos_th) if idx == 0 else None plots.fn(cos_th)(cos_th, kind, prpl=prpl, label=clabel, linestyle=linestyles[idx]) @@ -178,7 +186,8 @@ def fig_medium(): plt.title(titling[kind]) for idx, prpl in enumerate(prpls): for cos_th in cos_ths: - clabel = r'$\cos \theta_z = {}$'.format(cos_th) if idx == 0 else None + clabel = r'$\cos \theta_z = {}$'.format( + cos_th) if idx == 0 else None plots.fn(cos_th)(cos_th, kind, prpl=prpl[0], depth=prpl[1], label=clabel, linestyle=linestyles[idx]) @@ -189,7 +198,7 @@ def fig_medium(): plt.tight_layout(0.3) save(f'fig/medium_{kind}.eps') - + def fig_hadrs(): kinds = ['conv nu_e', 'pr nu_e', 'conv nu_mu', 'pr nu_mu'] hadrs_prompt = ['SIBYLL2.3c', 'SIBYLL2.3', 'DPMJET-III-3.0.6'] @@ -201,7 +210,8 @@ def fig_hadrs(): hadrs = hadrs_conv if kind.split('_')[0] == 'conv' else hadrs_prompt for idx, hadr in enumerate(hadrs): for cos_th in cos_ths: - clabel = r'$\cos \theta_z = {}$'.format(cos_th) if idx == 0 else None + clabel = r'$\cos \theta_z = {}$'.format( + cos_th) if idx == 0 else None plots.fn(cos_th)(cos_th, kind, hadr=hadr, label=clabel, linestyle=linestyles[idx]) @@ -215,17 +225,18 @@ def fig_hadrs(): def fig_density(): kinds = ['conv nu_e', 'pr nu_e', 'conv nu_mu', 'pr nu_mu'] - dmodels = [('CORSIKA',('SouthPole', 'June'), 'MSIS-90-E SP/Jul'), - ('CORSIKA',('SouthPole', 'December'), 'MSIS-90-E SP/Dec'), - ('MSIS00',('Karlsruhe', 'July'), 'NRLMSISE-00 KR/Jul'), - ('MSIS00',('Karlsruhe', 'December'), 'NRLMSISE-00 KR/Dec')] + dmodels = [('CORSIKA', ('SouthPole', 'June'), 'MSIS-90-E SP/Jul'), + ('CORSIKA', ('SouthPole', 'December'), 'MSIS-90-E SP/Dec'), + ('MSIS00', ('Karlsruhe', 'July'), 'NRLMSISE-00 KR/Jul'), + ('MSIS00', ('Karlsruhe', 'December'), 'NRLMSISE-00 KR/Dec')] cos_ths = [0.25, 0.85] for kind in kinds: plt.figure() plt.title(titling[kind]) for idx, dmodel in enumerate(dmodels): for cos_th in cos_ths: - clabel = r'$\cos \theta_z = {}$'.format(cos_th) if idx == 0 else None + clabel = r'$\cos \theta_z = {}$'.format( + cos_th) if idx == 0 else None plots.fn(cos_th)(cos_th, kind, density=dmodel[:2], label=clabel, linestyle=linestyles[idx]) @@ -251,7 +262,8 @@ def fig_pmodels(): plt.title(titling[kind]) for idx, pmodel in enumerate(pmodels): for cos_th in cos_ths: - clabel = r'$\cos \theta_z = {}$'.format(cos_th) if idx == 0 else None + clabel = r'$\cos \theta_z = {}$'.format( + cos_th) if idx == 0 else None plots.fn(cos_th)(cos_th, kind, pmodel=pmodel[:-1], label=clabel, linestyle=linestyles[idx]) @@ -266,7 +278,7 @@ def fig_pmodels(): def fig_extsv(): kinds = ['conv nu_e', 'pr nu_e', 'conv nu_mu', 'pr nu_mu'] cos_ths = [0.25, 0.85] - ens = np.logspace(2,9, 100) + ens = np.logspace(2, 9, 100) useexts = [False, True] labels = ['This work', 'GJKvS'] for kind in kinds: @@ -276,9 +288,11 @@ def fig_extsv(): for cos_th in cos_ths: if useext: emu = extsv.minimum_muon_energy(extsv.overburden(cos_th)) - plt.plot(ens, exthp.passrates(kind)(ens, emu, cos_th), linestyles[idx]) + plt.plot(ens, exthp.passrates(kind)( + ens, emu, cos_th), linestyles[idx]) else: - clabel = r'$\cos \theta_z = {}$'.format(cos_th) if idx == 0 else None + clabel = r'$\cos \theta_z = {}$'.format( + cos_th) if idx == 0 else None plots.fn(cos_th)(cos_th, kind, label=clabel) plt.axvline(np.nan, color='k', linestyle=linestyles[idx], @@ -290,23 +304,25 @@ def fig_extsv(): def fig_flux(): - aachen8 = lambda enu: 1.01*(enu/(100*Units.TeV))**-2.19*1e-18 + def aachen8(enu): return 1.01*(enu/(100*Units.TeV))**-2.19*1e-18 kinds = ['conv nu_e', 'conv nu_mu', 'pr nu_e', 'pr nu_mu'] ens = [1e4, 1e5] - cths = np.linspace(0,1,11) + cths = np.linspace(0, 1, 11) cths_full = np.concatenate((-cths[:0:-1], cths)) - cths_plot = np.linspace(-1,1,100) + cths_plot = np.linspace(-1, 1, 100) for enu in ens: plt.figure() plt.title(r'$E_\nu = {:.0f}$ TeV'.format(enu/1e3)) astro = np.ones(cths_full.shape)*aachen8(enu*Units.GeV) - earth = np.asarray([earth_attenuation(enu, cth, 'astro_numu') for cth in cths_full]) + earth = np.asarray( + [earth_attenuation(enu, cth, 'astro_numu') for cth in cths_full]) astrofn = interp1d(cths_full, np.log10(earth*astro*enu**3), - kind='quadratic') + kind='quadratic') plt.plot(cths_plot, 10**astrofn(cths_plot)/2, color='gray', label=r'Astrophysical $\nu_\mu$') for kind in kinds: - earth = np.asarray([earth_attenuation(enu, cth, kind) for cth in cths_full]) + earth = np.asarray([earth_attenuation(enu, cth, kind) + for cth in cths_full]) passing = [] total = [] @@ -329,13 +345,15 @@ def fig_flux(): lpassing = plt.axvline(np.nan, color='k') ltotal = plt.axvline(np.nan, color='k', linestyle=':') plt.xlabel(r'$\cos \theta_z$') - plt.ylabel(r'$E_\nu^3 \Phi_\nu$ [GeV$^2$ cm$^{-2}$ s$^{-1}$ sr$^{-1}]$') - plt.xlim(-1,1) - plt.ylim(1e-6,1e-1) + plt.ylabel( + r'$E_\nu^3 \Phi_\nu$ [GeV$^2$ cm$^{-2}$ s$^{-1}$ sr$^{-1}]$') + plt.xlim(-1, 1) + plt.ylim(1e-6, 1e-1) plt.yscale('log') if enu == 1e5: leg1 = plt.legend(loc='upper left') - plt.legend([lpassing, ltotal], ['Passing', 'Total'], loc='upper right') + plt.legend([lpassing, ltotal], [ + 'Passing', 'Total'], loc='upper right') else: leg1 = plt.legend() plt.legend([lpassing, ltotal], ['Passing', 'Total']) diff --git a/tests/test_app.py b/tests/test_app.py index c052369..6aab30a 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -12,17 +12,20 @@ def test_categ(): - assert nuVeto.categ_to_mothers('conv', 'nu_mu') == ['pi+', 'K+', 'K_L0', 'mu-'] - assert nuVeto.categ_to_mothers('conv', 'nu_mubar') == ['pi-', 'K-', 'K_L0', 'mu+'] - assert nuVeto.categ_to_mothers('conv', 'nu_e') == ['pi+', 'K+', 'K_L0', 'K_S0', 'mu+'] + assert nuVeto.categ_to_mothers('conv', 'nu_mu') == [ + 'pi+', 'K+', 'K_L0', 'mu-'] + assert nuVeto.categ_to_mothers('conv', 'nu_mubar') == [ + 'pi-', 'K-', 'K_L0', 'mu+'] + assert nuVeto.categ_to_mothers('conv', 'nu_e') == [ + 'pi+', 'K+', 'K_L0', 'K_S0', 'mu+'] assert nuVeto.categ_to_mothers('pr', 'nu_mu') == ['D+', 'D_s+', 'D0'] assert nuVeto.categ_to_mothers('pr', 'nu_mubar') == ['D-', 'D_s-', 'Dbar0'] def test_costh_effective(): geom = Geometry(1950*Units.m) - cosths = np.linspace(-1,1,100) - assert np.all(geom.cos_theta_eff(cosths)>=cosths) + cosths = np.linspace(-1, 1, 100) + assert np.all(geom.cos_theta_eff(cosths) >= cosths) center = Geometry(geom.r_E) assert np.all(center.cos_theta_eff(cosths) == np.ones(100)) @@ -30,7 +33,7 @@ def test_costh_effective(): def test_overburden(): geom = Geometry(1950*Units.m) - cosths = np.linspace(-1,1,100) + cosths = np.linspace(-1, 1, 100) assert np.all(np.diff(geom.overburden(cosths)) < 0) center = Geometry(geom.r_E) @@ -41,20 +44,21 @@ def test_pdet(): l_ice = np.linspace(1000, 200000, 500) emui = np.logspace(3, 8, 500)*Units.GeV coords = np.stack(np.meshgrid(emui, l_ice), axis=-1) - root, subdir, fpaths = next(os.walk(resources.files('nuVeto') / 'data' / 'prpl')) + root, subdir, fpaths = next( + os.walk(resources.files('nuVeto') / 'data' / 'prpl')) for fpath in fpaths: muprob = MuonProb(os.path.splitext(fpath)[0]) pdets = muprob.prpl(coords) - assert np.all(pdets >=0) and np.all(pdets <=1) + assert np.all(pdets >= 0) and np.all(pdets <= 1) def test_edge(): """ Test edge case where MCEq yields are all <= 0. """ sv = nuVeto(0., pmodel=(pm.ZatsepinSokolskaya, 'pamela'), hadr='DPMJET-III-19.1', - density=('MSIS00_IC', ('SouthPole','June'))) + density=('MSIS00_IC', ('SouthPole', 'June'))) _ = sv.get_rescale_phi('D-', 508.0218046913023, 14) - assert not np.any(_[:,-1] > 0) + assert not np.any(_[:, -1] > 0) @pytest.mark.parametrize('cth', [0.1, 0.3, 0.8]) @@ -67,27 +71,27 @@ def test_pnmshower(cth): @pytest.mark.parametrize('enu,l_ice,mother', - product(np.logspace(3,7,5), + product(np.logspace(3, 7, 5), np.linspace(1500, 100000, 5), 'pi+ K+ K_L0 D+ D0 Ds+'.split())) -def test_pnmsib(enu,l_ice,mother): +def test_pnmsib(enu, l_ice, mother): psibs = nuVeto.psib(l_ice, mother, enu, 3, 'ice_allm97_step_1') assert np.all(0 <= psibs) and np.all(psibs <= 1) -@pytest.mark.parametrize('cth', [0.1,0.3,0.8]) +@pytest.mark.parametrize('cth', [0.1, 0.3, 0.8]) def test_elbert(cth): - ens = np.logspace(2,8.9,50) + ens = np.logspace(2, 8.9, 50) mine = np.asarray( [passing(en, cth, kind='conv nu_mu', hadr='DPMJET-III-3.0.6', pmodel=(pm.GaisserHonda, None), prpl=None, corr_only=True) for en in ens]) emu = extsv.minimum_muon_energy(extsv.overburden(cth)) theirs = exthp.corr('conv nu_mu')(ens, emu, cth) print('test_elbert', cth) - assert np.all(np.abs(theirs-mine)<0.022) + assert np.all(np.abs(theirs-mine) < 0.022) -@pytest.mark.parametrize('cth', [0.1,0.3,0.8]) +@pytest.mark.parametrize('cth', [0.1, 0.3, 0.8]) def test_nuflux(cth): sv = nuVeto(cth) sv.grid_sol() @@ -98,7 +102,8 @@ def test_nuflux(cth): thres = 1e7 ensel = (sv.mceq.e_grid > 1e2) & (sv.mceq.e_grid < thres) theirs = sv.mceq.get_solution(mceq_categ_format(kind))[ensel] - mine = np.asarray([fluxes(en, cth, kind, corr_only=True)[1] for en in sv.mceq.e_grid[ensel]]) + mine = np.asarray([fluxes(en, cth, kind, corr_only=True)[1] + for en in sv.mceq.e_grid[ensel]]) print(kind, cth, theirs/mine) if _c == 'conv': @@ -107,10 +112,10 @@ def test_nuflux(cth): assert np.all(np.abs(theirs/mine - 1) < 0.8) -@pytest.mark.parametrize('cth', [0.9,1]) +@pytest.mark.parametrize('cth', [0.9, 1]) def test_nonneg(cth, capsys): with capsys.disabled(): - sv = nuVeto(cth,debug_level=2) + sv = nuVeto(cth, debug_level=2) enus = [6.2e6, 1e7] kinds = ['conv nu_mu', 'conv nu_e', 'pr nu_mu', 'pr nu_e'] for enu, kind in product(enus, kinds):