extra tests

FRBs · Aug 16, 2024 · 081bf19 · 081bf19
1 parent c47dd0f
commit 081bf19
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 9 deletions.
diff --git a/frb/scripts/pz_dm.py b/frb/scripts/pz_dm.py
@@ -5,6 +5,7 @@
 """
 from IPython import embed
 
+
 def parser(options=None):
     import argparse
     # Parse
@@ -76,21 +77,19 @@ def main(pargs):
     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
+        DM = zdict['DM']
+        PDM_z = prob_dmz.grab_repo_grid(telescope_dict[pargs.telescope])
         iDM = np.argmin(np.abs(DM - DM_cosmic))
         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']

diff --git a/frb/tests/test_frbgalaxies.py b/frb/tests/test_frbgalaxies.py
@@ -153,6 +153,19 @@ def test_mag_dm_figure():
 
 
 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')
@@ -165,16 +178,22 @@ def test_pzdm_telescopes():
     assert len(DM) == 1400
     assert PDM_z.shape == (500, 1400)
 
-    # Load the perfect grid
+    # 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) == 500    
-    assert len(DM) == 1400
-    assert PDM_z.shape == (500, 1400)
+    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)