Merge pull request #184 from FRBs/mag_vs_dm

CHIME telescope added to p(z|DM)

docs/scripts.rst

@@ -107,11 +107,18 @@ This script takes as input the FRB DM and its coordinates (approximate
 are fine) and then estimates the redshift range assuming 
 the Macquart relation (and simple host + MW contributions, optionally 
 input).  For an (optionally input; tuple) confidence interval, 
-it reports back the putative redshift range for the FRB.
+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 (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::
 
-    usage: frb_pz_dm [-h] [--dm_hostmw DM_HOSTMW] [--cl CL] coord DM_FRB
+    usage: frb_pz_dm [-h] [--dm_hostmw DM_HOSTMW] [--cl CL] coord DM_FRB 
+                    [--magdm_plot] [--fig_title FIG_TITLE] [--telescope TELESCOPE]
 
     Script to print a summary of an FRB to the screen [v1.0]
 
@@ -129,6 +136,16 @@ Here is the usage::
     --cl CL               Confidence limits for the z estimate [default is a 95
                             percent c.l., (2.5,97.5)]
 
+    --magdm_plot          Plot the host redshift range given DM on the magnitude
+                            vs redshift evolution. Default=False.   
+
+    --fig_title FIG_TITLE  title for the figure; e.g., FRBXXXXX
+
+    --telescope TELESCOPE  telescope model for the DM-z grid: CHIME, DSA, Parkes,
+                            FAST, CRAFT, CRAFT_ICS_892/1300/1632, perfect. Default
+                            = perfect
+
+
 
 frb_sightline
 =============

frb/dm/prob_dmz.py

@@ -1,7 +1,9 @@
 """ Code for calculations of P(DM|z) and P(z|DM)"""
 import numpy as np
 import os
-from pkg_resources import resource_filename
+
+
+from importlib.resources import files
 
 from scipy.stats import norm, lognorm
 from scipy.interpolate import interp1d, interp2d
@@ -150,32 +152,35 @@ def build_grid_for_repo(outfile:str):
     print("This will be used going forth")
 
 
-def grab_repo_grid():
+
+def grab_repo_grid(grid_name):
     """
-    Grab the grid from the Repository
-    This may require the code to build it first!
+    Grab the grid from the Repository based on the given grid name
+
+    Args:
+        grid_name (str): Name of the grid to grab
 
     Returns:
-        dict: Numpy dict from the npz save file
+        dict: Numpy dict from the npz or npy save file
     """
 
     # File
-    PDM_z_grid_file = os.path.join(
-        resource_filename('frb', 'data'), 'DM',
-        'PDM_z.npz')
-
+    grid_file = files('frb.data.DM').joinpath( grid_name)
+
     # Build?
-    if not os.path.isfile(PDM_z_grid_file):
-        build_grid_for_repo(PDM_z_grid_file)
-
+    if grid_name == 'PDM_z.npz':
+        if not os.path.isfile(grid_file):
+            build_grid_for_repo(grid_file)
+
     # Load
-    print(f"Loading P(DM,z) grid from {PDM_z_grid_file}")
-    sdict = np.load(PDM_z_grid_file)
+    print(f"Loading P(DM,z) grid from {grid_file}")
+    sdict = np.load(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, logz_scale=False):
     """
     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.3 * (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 logz_scale == True:
+        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/galaxies/utils.py

@@ -70,7 +70,7 @@ def load_specdb(specdb_file=None):
             specdb_file = specdb_files[0]
             print("Loading spectra from {:s}".format(specdb_file))
         else:
-            raise IOError("There are no FRB_specdb.hdf5 files in your SPECDB folder")
+            raise IOError("There are no FRB_specDB_*.hdf5 files in your SPECDB folder")
     # Load it up
     specDB = SpecDB(db_file=specdb_file)
     # Return

frb/scripts/pz_dm.py

@@ -1,21 +1,26 @@
-#!/usr/bin/env python
+
 """
 Estimate p(z|DM) for an assumed location on the sky and DM_FRB
-  Warning:  This assumes a *perfect* telescope
+Defaults to using a perfect telescope model for the DM-z grid
 """
 from IPython import embed
 
+
 def parser(options=None):
     import argparse
     # Parse
     parser = argparse.ArgumentParser(description='Script to print a summary of an FRB to the screen [v1.0]')
     parser.add_argument("coord", type=str, help="Coordinates, e.g. J081240.7+320809 or 122.223,-23.2322 or 07:45:00.47,34:17:31.1 or FRB name (FRB180924)")
     parser.add_argument("DM_FRB", type=float, help="FRB DM (pc/cm^3)")
-    parser.add_argument("--dm_hostmw", type=float, default=100., help="Assumed DM contribution from the Milky Way Halo (ISM is calculated from NE2001) and Host. Default = 100")
+    parser.add_argument("--dm_host", type=float, default=50., help="Assumed DM contribution from the Host. Default = 50")
+    parser.add_argument("--dm_mwhalo", type=float, default=50., help="Assumed DM contribution from the MW halo. Default = 50")
     parser.add_argument("--cl", type=tuple, default=(2.5,97.5), 
                         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='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:
         pargs = parser.parse_args()
@@ -36,7 +41,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))
 
@@ -46,21 +51,56 @@ def main(pargs):
     print("-----------------------------------------------------")
     print(f"NE2001 = {DM_ISM:.2f}")
 
-    # DM Cosmic
-    DM_cosmic = pargs.DM_FRB - DM_ISM.value - pargs.dm_hostmw
+    # DM cosmic and EG
+    DM_cosmic = pargs.DM_FRB - DM_ISM.value - pargs.dm_mwhalo
+    DM_extragalactic = DM_cosmic + pargs.dm_host  
 
     # Redshift estimates
 
-    # Load
-    sdict = prob_dmz.grab_repo_grid()
+    # Load the telescope specific grid
+    telescope_dict = {
+        'CHIME': 'CHIME_pzdm.npz',
+        'DSA': 'DSA_pzdm.npy',
+        'Parkes': 'parkes_mb_class_I_and_II_pzdm.npy',
+        'CRAFT': 'CRAFT_class_I_and_II_pzdm.npy',
+        'CRAFT_ICS_1300': 'CRAFT_ICS_1300_pzdm.npy',
+        'CRAFT_ICS_892': 'CRAFT_ICS_892_pzdm.npy',
+        'CRAFT_ICS_1632': 'CRAFT_ICS_1632_pzdm.npy',
+        'FAST': 'FAST_pzdm.npy',
+        'perfect': 'PDM_z.npz'
+    }
+
+    # 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']
 
-    # Do it
+    # Grab the right entry
     iDM = np.argmin(np.abs(DM - DM_cosmic))
     PzDM = PDM_z[iDM, :] / np.sum(PDM_z[iDM, :])
 
+
+    # 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']
+        DM = zdict['DM']
+        PDM_z = prob_dmz.grab_repo_grid(telescope_dict[pargs.telescope])
+        iDM = np.argmin(np.abs(DM - DM_extragalactic))
+        PzDM = PDM_z[:,iDM] / np.sum(PDM_z[:,iDM])
+
+
+    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_extragalactic))
+        PzDM = PDM_z[:,iDM] / np.sum(PDM_z[:,iDM])
+
     cum_sum = np.cumsum(PzDM)
     limits = pargs.cl
 
@@ -72,7 +112,8 @@ def main(pargs):
 
     # Finish
     print("")
-    print(f"Allowing for the MW and Host, DM_cosmic = {int(DM_cosmic)} pc/cm^3")
+    print(f"Allowing for the MW halo, DM_MW_halo = {int(pargs.dm_mwhalo))} pc/cm^3")
+    print(f"Allowing for the Host, DM_host = {int(pargs.dm_host)} pc/cm^3")
     print("")
     print("")
     print(f"The mean redshift value is: {z_50:.3f}")
@@ -81,12 +122,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 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)")
     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, logz_scale=False)
 
     return z_min, z_max, z_50, z_mode

frb/tests/test_frbgalaxies.py

@@ -150,3 +150,51 @@ def test_mag_dm_figure():
                flipy=True, known_hosts = False)
     assert os.path.exists('./temp_fig.png') 
     os.system('rm ./temp_fig.png')
+
+
+def test_pzdm_telescopes():
+
+    telescope_dict = {
+        'DSA': 'DSA_pzdm.npy',
+        'Parkes': 'parkes_mb_class_I_and_II_pzdm.npy',
+        'CRAFT': 'CRAFT_class_I_and_II_pzdm.npy',
+        'CRAFT_ICS_1300': 'CRAFT_ICS_1300_pzdm.npy',
+        'CRAFT_ICS_892': 'CRAFT_ICS_892_pzdm.npy',
+        'CRAFT_ICS_1632': 'CRAFT_ICS_1632_pzdm.npy',
+        'FAST': 'FAST_pzdm.npy',
+    }
+
+    # Load the CHIME grid
+    from frb.dm import prob_dmz
+    sdict = prob_dmz.grab_repo_grid('CHIME_pzdm.npz')
+    PDM_z = sdict['pzdm']
+    z = sdict['z']
+    DM = sdict['DM']
+
+    # Test
+    assert len(z) == 500
+    assert len(DM) == 1400
+    assert PDM_z.shape == (500, 1400)
+
+    # check 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) == 200    
+    assert len(DM) == 1000
+    assert PDM_z.shape == (1000, 200)
+
+    # check the other grids
+    for key in telescope_dict.keys():
+        PDM_z = prob_dmz.grab_repo_grid(telescope_dict[key])
+        # Test
+        assert PDM_z.shape == (500, 1400)
+
+
+
+
+
+