def to perfect

FRBs · Aug 16, 2024 · c47dd0f · c47dd0f
1 parent 705477a
commit c47dd0f
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 17 deletions.
diff --git a/docs/scripts.rst b/docs/scripts.rst
@@ -110,7 +110,10 @@ input).  For an (optionally input; tuple) confidence interval,
 it reports back the putative redshift range for the FRB. It also 
 allows for plotting the host redshift range on the magnitude vs redshift 
 evolution and setting a title for the figure.  These calculations can be 
-done assuming a few different telescope models or a perfect telescope model.
+done assuming a few different telescope models (CHIME, DSA, Parkes, FAST,
+CRAFT, CRAFT_ICS_892/1300/1632) or a perfect telescope model (default). 
+The telescope models are used to determine the DM-z grids that have been 
+computed with the zdm code/repository.
 
 Here is the usage::
 
@@ -140,7 +143,7 @@ Here is the usage::
 
     --telescope TELESCOPE  telescope model for the DM-z grid: CHIME, DSA, Parkes,
                             FAST, CRAFT, CRAFT_ICS_892/1300/1632, perfect. Default
-                            = CHIME
+                            = perfect
 
 
 

diff --git a/frb/scripts/pz_dm.py b/frb/scripts/pz_dm.py
@@ -1,7 +1,7 @@
 
 """
 Estimate p(z|DM) for an assumed location on the sky and DM_FRB
-Defaults to using the CHIME telescope model for the DM-z grid
+Defaults to using a perfect telescope model for the DM-z grid
 """
 from IPython import embed
 
@@ -17,8 +17,8 @@ def parser(options=None):
                         help="Confidence limits for the z estimate [default is a 95 percent c.l., (2.5,97.5)]")
     parser.add_argument("--magdm_plot", default=False, action='store_true', 
                         help="Plot the host redshift range given DM on the magnitude vs redshift evolution")
-    parser.add_argument("--telescope", type=str, default='CHIME', help="telescope model for the DM-z grid: CHIME, DSA, Parkes, FAST, CRAFT, \
-                        CRAFT_ICS_892/1300/1632, perfect. Default = CHIME")
+    parser.add_argument("--telescope", type=str, default='perfect', help="telescope model for the DM-z grid: CHIME, DSA, Parkes, FAST, CRAFT, \
+                        CRAFT_ICS_892/1300/1632, perfect. Default = perfect")
     parser.add_argument("--fig_title", type=str,  help="title for the figure; e.g., FRBXXXXX")
 
     if options is None:
@@ -40,7 +40,7 @@ def main(pargs):
     from frb.dm import prob_dmz
     from frb.galaxies import mag_dm
 
- 
+
     # Deal with coord
     icoord = ltu.radec_to_coord(coord_arg_to_coord(pargs.coord))
 
@@ -68,30 +68,36 @@ def main(pargs):
         'perfect': 'PDM_z.npz'
     }
 
-    # Get z and DM arrays from CHIME
-    sdict = prob_dmz.grab_repo_grid(telescope_dict['CHIME'])
-    PDM_z = sdict['pzdm']
+    # Get the perfect telescope grid (default)
+    sdict = prob_dmz.grab_repo_grid(telescope_dict['perfect'])
+    PDM_z = sdict['PDM_z']
     z = sdict['z']
     DM = sdict['DM']
+    iDM = np.argmin(np.abs(DM - DM_cosmic))
+    PzDM = PDM_z[iDM, :] / np.sum(PDM_z[iDM, :])
+
+    print (len(z), len(PzDM), PzDM.shape)   
+    aslkdj
 
 
     # Get the telescope specific PZDM grid
     if pargs.telescope and pargs.telescope != 'CHIME' and pargs.telescope != 'perfect':
         if pargs.telescope not in telescope_dict:
             raise ValueError(f"Unknown telescope: {pargs.telescope}")
+        zdict = prob_dmz.grab_repo_grid(telescope_dict['CHIME'])
+        z = zdict['z']
         sdict = prob_dmz.grab_repo_grid(telescope_dict[pargs.telescope])
         PDM_z = sdict
+        iDM = np.argmin(np.abs(DM - DM_cosmic))
+        PzDM = PDM_z[:,iDM] / np.sum(PDM_z[:,iDM])
 
-    iDM = np.argmin(np.abs(DM - DM_cosmic))
-    PzDM = PDM_z[:,iDM] / np.sum(PDM_z[:,iDM])
-
-    if pargs.telescope and pargs.telescope == 'perfect':
-        sdict = prob_dmz.grab_repo_grid(telescope_dict[pargs.telescope])
-        PDM_z = sdict['PDM_z']
+    if pargs.telescope and pargs.telescope == 'CHIME':
+        sdict = prob_dmz.grab_repo_grid(telescope_dict['CHIME'])
+        PDM_z = sdict['pzdm']
         z = sdict['z']
         DM = sdict['DM']
         iDM = np.argmin(np.abs(DM - DM_cosmic))
-        PzDM = PDM_z[iDM, :] / np.sum(PDM_z[iDM, :])
+        PzDM = PDM_z[:,iDM] / np.sum(PDM_z[:,iDM])
 
     cum_sum = np.cumsum(PzDM)
     limits = pargs.cl
@@ -113,7 +119,7 @@ def main(pargs):
     print(f"The redshift range for your confidence interval {pargs.cl} is:")
     print(f"z = [{z_min:.3f}, {z_max:.3f}]")
     print("")
-    if pargs.telescope == 'perfect':
+    if not pargs.telescope or pargs.telescope == 'perfect':
         print("WARNING: This all assumes a perfect telescope and a model of the scatter in DM_cosmic (Macqurt+2020)")
     else:
         print("This assumes the "+(str(pargs.telescope))+" telescope and a model of the scatter in DM_cosmic (Macqurt+2020)")

diff --git a/frb/tests/test_frbgalaxies.py b/frb/tests/test_frbgalaxies.py
@@ -165,6 +165,19 @@ def test_pzdm_telescopes():
     assert len(DM) == 1400
     assert PDM_z.shape == (500, 1400)
 
+    # Load the perfect grid
+    sdict = prob_dmz.grab_repo_grid('PDM_z.npz')
+    PDM_z = sdict['PDM_z']
+    z = sdict['z']
+    DM = sdict['DM']
+
+    # Test      
+    assert len(z) == 500    
+    assert len(DM) == 1400
+    assert PDM_z.shape == (500, 1400)
+
+
+