Run autopep to fix whitespace and indents (#27)

nuVeto/__init__.py

@@ -1 +1,2 @@
-__all__ = ['mu', 'utils', 'nuveto', 'barr_uncertainties', 'external','examples']
+__all__ = ['mu', 'utils', 'nuveto',
+           'barr_uncertainties', 'external', 'examples']

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,42 +239,45 @@ 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)
 
             ###
             # 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)

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')

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)