Skip to content

Commit

Permalink
working?
Browse files Browse the repository at this point in the history
  • Loading branch information
@profxj
profxj committed Jul 16, 2024
1 parent 73fa517 commit 8b17a77
Show file tree
Hide file tree
Showing 3 changed files with 226 additions and 40 deletions.
222 changes: 194 additions & 28 deletions docs/nb/CHIME_pzDM.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,26 +10,32 @@
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": 9,
"id": "971524d6-aab9-457c-b62c-7b499dd5a2e6",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/xavier/Projects/FRB_Stuff/FRB/frb/halos/hmf.py:51: UserWarning: hmf_emulator not imported. Hope you are not intending to use the hmf.py module..\n",
" warnings.warn(\"hmf_emulator not imported. Hope you are not intending to use the hmf.py module..\")\n"
]
}
],
"outputs": [],
"source": [
"# imports\n",
"import numpy as np\n",
"from matplotlib import pyplot as plt\n",
"\n",
"from zdm.chime import grids"
]
},
{
"cell_type": "code",
"execution_count": 43,
"id": "493b5f54-db48-4014-80dd-843e308fa0aa",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%matplotlib ipympl"
]
},
{
"cell_type": "markdown",
"id": "ff1a7d56-ceab-4d96-83e0-8e4e4333d940",
Expand All @@ -50,37 +56,197 @@
"name": "stdout",
"output_type": "stream",
"text": [
"Searching for survey in directory /home/xavier/Projects/FRB_Stuff/zdm/zdm/data/Surveys/CHIME/\n",
"/home/xavier/Projects/FRB_Stuff/zdm/zdm/data/Surveys/CHIME/\n",
"Loading survey: CHIME_decbin_0_of_6\n",
"Loading survey: CHIME_decbin_0_of_6 from CHIME_decbin_0_of_6.ecsv\n",
"Loaded FRB info\n",
"FRB survey sucessfully initialised with 4 FRBs starting from 0\n",
"Initializing igamma_spline for gamma=-1.01\n",
"Initialised grid\n",
"Loaded dec bin 0\n"
]
},
{
"ename": "TypeError",
"evalue": "repeat_Grid.__init__() got an unexpected keyword argument 'Tfield'",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[2], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m dmvals, zvals, all_rates, all_singles, all_reps \u001b[38;5;241m=\u001b[39m \u001b[43mgrids\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/Projects/FRB_Stuff/zdm/zdm/chime/grids.py:79\u001b[0m, in \u001b[0;36mload\u001b[0;34m(Nbin, make_plots)\u001b[0m\n\u001b[1;32m 72\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mLoaded dec bin \u001b[39m\u001b[38;5;124m\"\u001b[39m,ibin)\n\u001b[1;32m 74\u001b[0m \u001b[38;5;66;03m# The below is specific to CHIME data. For CRAFT and other\u001b[39;00m\n\u001b[1;32m 75\u001b[0m \u001b[38;5;66;03m# FRB surveys, do not use \"bmethod=2\", and you will have to\u001b[39;00m\n\u001b[1;32m 76\u001b[0m \u001b[38;5;66;03m# enter the time per field and Nfields manually.\u001b[39;00m\n\u001b[1;32m 77\u001b[0m \u001b[38;5;66;03m# Also, for other surveys, you do not need to iterate over\u001b[39;00m\n\u001b[1;32m 78\u001b[0m \u001b[38;5;66;03m# declination bins!\u001b[39;00m\n\u001b[0;32m---> 79\u001b[0m rg \u001b[38;5;241m=\u001b[39m \u001b[43mrep\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrepeat_Grid\u001b[49m\u001b[43m(\u001b[49m\u001b[43mg\u001b[49m\u001b[43m,\u001b[49m\u001b[43mTfield\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43ms\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mTOBS\u001b[49m\u001b[43m,\u001b[49m\u001b[43mNfields\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43mMC\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43mopdir\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43mbmethod\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 80\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mLoaded repeat grid\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 82\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m ibin\u001b[38;5;241m==\u001b[39m\u001b[38;5;241m0\u001b[39m:\n",
"\u001b[0;31mTypeError\u001b[0m: repeat_Grid.__init__() got an unexpected keyword argument 'Tfield'"
"Initialised surveys\n",
"Initialised grids\n",
"Loaded dec bin 0\n",
"Loaded repeat grid\n",
"Loading survey: CHIME_decbin_1_of_6\n",
"Loading survey: CHIME_decbin_1_of_6 from CHIME_decbin_1_of_6.ecsv\n",
"Loaded FRB info\n",
"FRB survey sucessfully initialised with 23 FRBs starting from 0\n",
"Initialised surveys\n",
"Initialised grids\n",
"Loaded dec bin 1\n",
"Loaded repeat grid\n",
"Loading survey: CHIME_decbin_2_of_6\n",
"Loading survey: CHIME_decbin_2_of_6 from CHIME_decbin_2_of_6.ecsv\n",
"Loaded FRB info\n",
"FRB survey sucessfully initialised with 154 FRBs starting from 0\n",
"Initialised surveys\n",
"Initialised grids\n",
"Loaded dec bin 2\n",
"Loaded repeat grid\n",
"Loading survey: CHIME_decbin_3_of_6\n",
"Loading survey: CHIME_decbin_3_of_6 from CHIME_decbin_3_of_6.ecsv\n",
"Loaded FRB info\n",
"FRB survey sucessfully initialised with 234 FRBs starting from 0\n",
"Initialised surveys\n",
"Initialised grids\n",
"Loaded dec bin 3\n",
"Loaded repeat grid\n",
"Loading survey: CHIME_decbin_4_of_6\n",
"Loading survey: CHIME_decbin_4_of_6 from CHIME_decbin_4_of_6.ecsv\n",
"Loaded FRB info\n",
"FRB survey sucessfully initialised with 50 FRBs starting from 0\n",
"Initialised surveys\n",
"Initialised grids\n",
"Loaded dec bin 4\n",
"Loaded repeat grid\n",
"Loading survey: CHIME_decbin_5_of_6\n",
"Loading survey: CHIME_decbin_5_of_6 from CHIME_decbin_5_of_6.ecsv\n",
"Loaded FRB info\n",
"FRB survey sucessfully initialised with 27 FRBs starting from 0\n",
"Initialised surveys\n",
"Initialised grids\n",
"Loaded dec bin 5\n",
"Loaded repeat grid\n"
]
}
],
"source": [
"dmvals, zvals, all_rates, all_singles, all_reps = grids.load()"
]
},
{
"cell_type": "markdown",
"id": "294d2281-e64b-467c-990e-5d89f22a0a80",
"metadata": {},
"source": [
"# Show it!"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "e4ad7724-445e-4ef4-a260-a5450f9a8cb0",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"(500, 1400)"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"all_singles.shape"
]
},
{
"cell_type": "code",
"execution_count": 37,
"id": "b464b2dc-95f3-4b4c-81b4-649d40f05559",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"all_singles /= np.sum(all_singles)"
]
},
{
"cell_type": "code",
"execution_count": 40,
"id": "4e452458-d32d-4a22-b20a-0164aa95c006",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"0.0002946589824842096"
]
},
"execution_count": 40,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"all_singles.max()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "729a8fb0-4068-46ff-8d1e-188c74de11c2",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# Assuming you have these arrays defined:\n",
"# z_array = ...\n",
"# DM_array = ...\n",
"# PDF = ... # Your 2D array representing the PDF\n",
"z_array = zvals\n",
"DM_array = dmvals\n",
"\n",
"# Create the figure and axis\n",
"fig, ax = plt.subplots(figsize=(8, 8)) # Square figure\n",
"\n",
"# Create the heatmap\n",
"im = ax.imshow(np.log10(all_singles.T), cmap='jet', aspect='auto', \n",
" extent=[z_array.min(), z_array.max(), \n",
" DM_array.min(), DM_array.max()],\n",
" origin='lower', vmin=np.log10(all_singles.max())-4.)\n",
"\n",
"# Set labels and title\n",
"ax.set_xlabel('z')\n",
"ax.set_ylabel('DM')\n",
"ax.set_title('PDF Heatmap')\n",
"\n",
"ax.set_xlim(0., 1.)\n",
"ax.set_ylim(0., 4000.)\n",
"\n",
"# Add a colorbar\n",
"cbar = fig.colorbar(im)\n",
"cbar.set_label(r'$\\log_{10} \\, p(z,DM_{\\rm EG})$')\n",
"\n",
"# Show the plot\n",
"plt.tight_layout()\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "8fed6205-86b3-498c-87dc-bb472f05cf6b",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"500"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"zvals.size"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6598fbb9-fa3b-4ea1-b806-df5645ca3c7c",
"id": "5ddfa68f-a491-42bf-864c-1e8e3fe32cf4",
"metadata": {},
"outputs": [],
"source": []
Expand Down
27 changes: 15 additions & 12 deletions zdm/chime/grids.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
#import utilities as ute
#import states as st

from zdm.craco import loading
from zdm import loading
from zdm import repeat_grid as rep
from zdm import misc_functions
from zdm import beams
Expand Down Expand Up @@ -56,55 +56,58 @@ def load(Nbin:int=6, make_plots:bool=False):
# loads survey data
sdir = os.path.join(resources.files('zdm'),
'data', 'Surveys','CHIME/')
print(sdir)

# loads beam data
bdir='Nbounds'+str(Nbin)+'/'
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')


# Loops through CHIME declination bins
for ibin in np.arange(Nbin):
name = "CHIME_decbin_"+str(ibin)+"_of_"+str(Nbin)
s,g = loading.survey_and_grid(survey_name=name,NFRB=None,sdir=sdir)#,init_state=state)
ss,rgs = loading.surveys_and_grids(survey_names=[name],sdir=sdir,repeaters=True)#,init_state=state)
s = ss[0]
rg = rgs[0]
print("Loaded dec bin ",ibin)

# The below is specific to CHIME data. For CRAFT and other
# FRB surveys, do not use "bmethod=2", and you will have to
# enter the time per field and Nfields manually.
# Also, for other surveys, you do not need to iterate over
# declination bins!
rg = rep.repeat_Grid(g,Tfield=s.TOBS,Nfields=1,MC=False,opdir=None,bmethod=2)
# rg = rep.repeat_Grid(g,Tfield=s.TOBS,Nfields=1,MC=False,opdir=None,bmethod=2)
print("Loaded repeat grid")

if ibin==0:
all_rates = g.rates
all_rates = rg.rates
all_singles = rg.exact_singles
all_reps = rg.exact_reps
else:
all_rates = all_rates + g.rates
all_rates = all_rates + rg.rates
all_singles = all_singles + rg.exact_singles
all_reps = all_reps + rg.exact_reps

if make_plots:
misc_functions.plot_grid_2(all_rates,g.zvals,g.dmvals,
misc_functions.plot_grid_2(all_rates,rg.zvals,rg.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,
misc_functions.plot_grid_2(all_reps,rg.zvals,rg.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,
misc_functions.plot_grid_2(all_singles,rg.zvals,rg.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)

return g.dmvals, g.zvals, all_rates, all_singles, all_reps
return rg.dmvals, rg.zvals, all_rates, all_singles, all_reps

load()
17 changes: 17 additions & 0 deletions zdm/tests/test_chime.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
"""
File to test CHIME
"""

import os
import pytest

from pkg_resources import resource_filename
import pandas

from zdm.tests import tstutils

from zdm.chime import grids

def test_run():
dmvals, zvals, all_rates, all_singles, all_reps = grids.load()

0 comments on commit 8b17a77

Please sign in to comment.