Merge pull request #46 from FRBs/chime_fussin

CHIME fussin

zdm/scripts/load_chime_data.py → zdm/chime/grids.py

@@ -9,7 +9,7 @@
 for a dataset.
 
 """
-from pkg_resources import resource_filename
+from importlib import resources
 
 import numpy as np
 from matplotlib import pyplot as plt
@@ -20,15 +20,29 @@
 #import states as st
 
 from zdm.craco import loading
-from zdm import cosmology as cos
 from zdm import repeat_grid as rep
 from zdm import misc_functions
-from zdm import survey
 from zdm import beams
 
-def main(Nbin=6):
+from IPython import embed
+
+def load(Nbin:int=6, make_plots:bool=False):
     """
     Nbins is the number of declination bins to use
+
+    Args:
+        Nbin (int, optional): Number of declination bins to use. Defaults to 6.
+           30 is allowed too
+        make_plots (bool, optional): Whether to make plots. Defaults to False.
+
+    Returns:
+        tuple: 
+            dmvals (np.ndarray): 1D array of DM values
+            zvals (np.ndarray): 1D array of redshift values
+            all_rates (np.ndarray): 2D array of rates
+            all_singles (np.ndarray): 2D array of single FRB rates
+            all_reps (np.ndarray): 2D array of repeating FRB rates
+        
     """
 
     opdir = 'CHIME/'
@@ -40,11 +54,13 @@ def main(Nbin=6):
     #state = st.set_state(pset)
 
     # loads survey data
-    sdir = os.path.join(resource_filename('zdm','data/Surveys/'),'CHIME/')
+    sdir = os.path.join(resources.files('zdm'),
+                        'data', 'Surveys','CHIME/')
 
     # loads beam data
     bdir='Nbounds'+str(Nbin)+'/'
-    beams.beams_path = os.path.join(resource_filename('zdm','data/BeamData/CHIME/'),bdir)
+    beams.beams_path = os.path.join(resources.files('zdm'),'data',
+                                    'BeamData', 'CHIME',bdir)
     #bounds = np.load(beams.beams_path+'bounds.npy')
     #solids = np.load(beams.beams_path+'solids.npy')
 
@@ -72,23 +88,23 @@ def main(Nbin=6):
             all_singles = all_singles + rg.exact_singles
             all_reps = all_reps + rg.exact_reps
 
-    misc_functions.plot_grid_2(all_rates,g.zvals,g.dmvals,
-            name=opdir+'all_CHIME_FRBs.pdf',norm=3,log=True,
-            label='$\\log_{10} p({\\rm DM}_{\\rm EG},z)$ [a.u.]',
-            project=False,Aconts=[0.01,0.1,0.5],
-            zmax=3.0,DMmax=3000)
-
-    misc_functions.plot_grid_2(all_reps,g.zvals,g.dmvals,
-            name=opdir+'repeating_CHIME_FRBs.pdf',norm=3,log=True,
-            label='$\\log_{10} p({\\rm DM}_{\\rm EG},z)$ [a.u.]',
-            project=False,Aconts=[0.01,0.1,0.5],
-            zmax=1.0,DMmax=1000)
-
-    misc_functions.plot_grid_2(all_singles,g.zvals,g.dmvals,
-            name=opdir+'single_CHIME_FRBs.pdf',norm=3,log=True,
-            label='$\\log_{10} p({\\rm DM}_{\\rm EG},z)$ [a.u.]',
-            project=False,Aconts=[0.01,0.1,0.5],
-            zmax=3.0,DMmax=3000)
+    if make_plots:
+        misc_functions.plot_grid_2(all_rates,g.zvals,g.dmvals,
+                name=opdir+'all_CHIME_FRBs.pdf',norm=3,log=True,
+                label='$\\log_{10} p({\\rm DM}_{\\rm EG},z)$ [a.u.]',
+                project=False,Aconts=[0.01,0.1,0.5],
+                zmax=3.0,DMmax=3000)
+
+        misc_functions.plot_grid_2(all_reps,g.zvals,g.dmvals,
+                name=opdir+'repeating_CHIME_FRBs.pdf',norm=3,log=True,
+                label='$\\log_{10} p({\\rm DM}_{\\rm EG},z)$ [a.u.]',
+                project=False,Aconts=[0.01,0.1,0.5],
+                zmax=1.0,DMmax=1000)
+
+        misc_functions.plot_grid_2(all_singles,g.zvals,g.dmvals,
+                name=opdir+'single_CHIME_FRBs.pdf',norm=3,log=True,
+                label='$\\log_{10} p({\\rm DM}_{\\rm EG},z)$ [a.u.]',
+                project=False,Aconts=[0.01,0.1,0.5],
+                zmax=3.0,DMmax=3000)
 
-# Nbin = 30 is also implemented
-main(Nbin=6)
+    return g.dmvals, g.zvals, all_rates, all_singles, all_reps