CHIME telescope added to p(z|DM)

frb/dm/prob_dmz.py

@@ -176,6 +176,32 @@ def grab_repo_grid():
     return sdict
 
 
+
+def grab_chime_repo_grid():
+    """
+    Grab the CHIME grid from the Repository
+
+    Returns:
+        dict: Numpy dict from the npz save file
+    """
+
+    # File
+    CHIME_PDM_z_grid_file = os.path.join(
+        resource_filename('frb', 'data'), 'DM',
+        'CHIME_pzdm.npz')
+
+    # assert the file exists
+    if not os.path.isfile(CHIME_PDM_z_grid_file):
+        raise IOError("You need to download the CHIME grid from the FRB repo")
+
+    # Load
+    print(f"Loading P(DM,z) grid from {CHIME_PDM_z_grid_file}")
+    sdict = np.load(CHIME_PDM_z_grid_file)
+
+    # Return
+    return sdict
+
+
 def get_DMcosmic_from_z(zarray, perc=0.5, redo_pdmz_grid=False, DMevals=np.linspace(1.,2000.,1000), beta=3., F=0.31, cosmo=defs.frb_cosmo):
     """
     Gives DMcosmic values of zarray, considering the percentile.

frb/galaxies/mag_dm.py

@@ -20,7 +20,7 @@
 
 
 def r_vs_dm_figure(z_min, z_max, z, PzDM, outfile='fig_r_vs_z.png',
-               flipy=True, known_hosts = False):
+               flipy=True, known_hosts = False, title=None):
     """
     Plots the intersection of galaxy apparent magnitude evolution with redshift 
     and redshift distribution of the FRB and saves it. 
@@ -40,6 +40,10 @@ def r_vs_dm_figure(z_min, z_max, z, PzDM, outfile='fig_r_vs_z.png',
     f_mL = frb_gal_u.load_f_mL()
 
 
+
+    # Increase default font size
+    plt.rcParams.update({'font.size': 12})
+
     # set up the figure
     plt.figure(figsize=(6, 5))
     gs = gridspec.GridSpec(1,1)
@@ -75,7 +79,9 @@ def r_vs_dm_figure(z_min, z_max, z, PzDM, outfile='fig_r_vs_z.png',
 
     # plot
     # L curves
-    zvals = np.linspace(0.021, 4.0, 200)
+
+
+    zvals = np.linspace(0.021, z_max+0.5, 200)
     m_Lstar = f_mL(zvals)
     ax.plot(zvals, m_Lstar, '-r', label='L*')
 
@@ -87,13 +93,14 @@ def r_vs_dm_figure(z_min, z_max, z, PzDM, outfile='fig_r_vs_z.png',
 
 
     # Add P(z|DM)
-    xmnx = (0., 4.)
-    ymnx = (14, 33.)
+    xmnx = (0., z_max+0.5)
+    ymnx = (14, max(15, m_001Lstar[-1]+0.5))
     if flipy:
         ymnx = (ymnx[1], ymnx[0])
 
     y_range = np.linspace(ymnx[0], ymnx[1], 300) 
     x_range = np.linspace(z_min,z_max,300)
+
     pzd = np.interp(x_range,z,PzDM)
     X,Y = np.meshgrid(x_range,y_range)
     Z = X/X*pzd
@@ -102,20 +109,28 @@ def r_vs_dm_figure(z_min, z_max, z, PzDM, outfile='fig_r_vs_z.png',
     #zlbl = 'P(z|DM) [95% c.l.]'
     c=ax.pcolor(X, Y, Z*1000, cmap='Blues')
     zlbl = '95% c.l. FRB Redshift \n estimated from DM'
-    ax.text(np.mean([z_min,z_max]), 20.5, zlbl,
+    text_x = 0.35 * (xmnx[1] - xmnx[0]) + xmnx[0]
+    text_y = 0.5 * (ymnx[1] - ymnx[0]) + ymnx[0]
+    ax.text(text_x, text_y, zlbl,
             color='k', 
             size='large', ha='center')
 
-    ax.set_xlabel(r'$z$')
-    ax.set_ylabel(r'$m_r$')
+    ax.set_xlabel(r'Redshift ($z$)')
+    ax.set_ylabel(r'Apparent Magnitude ($m_r$)')
     ax.set_xlim(xmnx)
     ax.set_ylim(ymnx)
 
-    ax.xaxis.set_major_locator(plt.MultipleLocator(0.5))
+    if z_max <= 0.5:
+        ax.set_xscale('log')
+        ax.set_xlim(1e-2,z_max+0.5)
+        ax.xaxis.set_major_locator(plt.LogLocator(base=10, numticks=12))
+    else:
+        ax.xaxis.set_major_locator(plt.MultipleLocator(0.5))
     plt.colorbar(c,label='p(z|DM) [a.u.]',ax=ax)
 
-    ax.legend()
-
+    ax.legend(loc='upper right', frameon=True, shadow=True)
+    # set the title of the figure
+    ax.set_title(title)
     frb_fig_u.set_fontsize(ax, 15.)
 
     # End

frb/scripts/pz_dm.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+
 """
 Estimate p(z|DM) for an assumed location on the sky and DM_FRB
   Warning:  This assumes a *perfect* telescope
@@ -16,6 +16,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("--perfect_telescope", default=False, action='store_true', help="Use a perfect telescope model for the DM-z grid instead of CHIME")
+    parser.add_argument("--fig_title", type=str,  help="title for the figure; e.g., FRBXXXXX")
 
     if options is None:
         pargs = parser.parse_args()
@@ -52,14 +54,26 @@ def main(pargs):
     # Redshift estimates
 
     # Load
-    sdict = prob_dmz.grab_repo_grid()
-    PDM_z = sdict['PDM_z']
+    # consider CHIME telescope
+    sdict = prob_dmz.grab_chime_repo_grid() 
+    PDM_z = sdict['pzdm']
     z = sdict['z']
     DM = sdict['DM']
-
-    # Do it
     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])
+
+    # perfect telescope case
+    if pargs.perfect_telescope:
+        sdict = prob_dmz.grab_repo_grid()
+        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, :])
+
+
+
+
 
     cum_sum = np.cumsum(PzDM)
     limits = pargs.cl
@@ -81,12 +95,17 @@ 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("")
-    print("WARNING: This all assumes a perfect telescope and a model of the scatter in DM_cosmic (Macqurt+2020)")
+    if pargs.perfect_telescope:
+        print("WARNING: This all assumes a perfect telescope and a model of the scatter in DM_cosmic (Macqurt+2020)")
+    else:
+        print("This assumes the CHIME telescope and a model of the scatter in DM_cosmic (Macqurt+2020)")
     print("-----------------------------------------------------")
 
+
+
     # make the magnitude vs redshift plot with z-range if requested
     if pargs.magdm_plot:
         mag_dm.r_vs_dm_figure(z_min, z_max, z, PzDM, outfile='fig_r_vs_z.png',
-               flipy=True, known_hosts = False)
+               flipy=True, known_hosts=False, title=pargs.fig_title)
 
     return z_min, z_max, z_50, z_mode