Skip to content

Commit

Permalink
docs: added SKA tutorial
Browse files Browse the repository at this point in the history
  • Loading branch information
@DanielaBreitman
DanielaBreitman committed Feb 24, 2025
1 parent 4761e84 commit 3889854
Showing 1 changed file with 41 additions and 41 deletions.
82 changes: 41 additions & 41 deletions docs/tutorials/SKA_forecast.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,21 +33,15 @@
"source": [
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import py21cmsense as p21sense\n",
"from astropy import units as un\n",
"from astropy.cosmology import Planck15\n",
"from matplotlib import colors, rcParams\n",
"\n",
"import py21cmsense as p21sense\n",
"from py21cmsense.observatory import Observatory\n",
"from py21cmsense.sensitivity import PowerSpectrum\n",
"from py21cmsense.theory import EOS2021\n",
"\n",
"rcParams.update({\"font.size\": 20})\n",
"\n",
"import random\n",
"\n",
"random.seed(0)\n",
"np.random.seed(0)"
"rcParams.update({\"font.size\": 20})"
]
},
{
Expand All @@ -64,47 +58,50 @@
"outputs": [],
"source": [
"def freq2z(f):\n",
" \"\"\"Convert frequency in MHz to redshift.\"\"\"\n",
" return 1420.4 / f - 1.0\n",
"\n",
"\n",
"def z2freq(z):\n",
" \"\"\"Convert redshift to frequency in MHz.\"\"\"\n",
" return 1420.4 / (z + 1.0)\n",
"\n",
"\n",
"def get_senses(\n",
" observation,\n",
" freq_bands,\n",
" kperp_edges_hMpc,\n",
" kpar_edges_hMpc,\n",
" calc_2D=False,\n",
" kperp_edges_hmpc,\n",
" kpar_edges_hmpc,\n",
" calc_2d=False,\n",
" foreground_model=\"moderate\",\n",
" theory_model=EOS2021,\n",
" **kwargs,\n",
"):\n",
" \"\"\"Calculate the sensitivity for a set of frequency bands.\"\"\"\n",
" h = observation.cosmo.H0.value / 100.0\n",
" redshifts = kwargs[\"redshifts\"]\n",
" mock_senses = {}\n",
" mock_senses[\"kperp_edges_hMpc\"] = kperp_edges_hMpc\n",
" mock_senses[\"kpar_edges_hMpc\"] = kpar_edges_hMpc\n",
" mock_senses[\"kperp_edges_hMpc\"] = kperp_edges_hmpc\n",
" mock_senses[\"kpar_edges_hMpc\"] = kpar_edges_hmpc\n",
" mock_senses[\"redshifts\"] = redshifts\n",
" mock_senses[\"freq_bands\"] = freq_bands\n",
" mock_senses[\"h\"] = h\n",
"\n",
" for i, band, zval in zip(range(len(freq_bands)), freq_bands, redshifts):\n",
" for _i, band, zval in zip(range(len(freq_bands)), freq_bands, redshifts):\n",
" this_z = {}\n",
" band_name = str(np.round(band, 1)) + \" MHz\"\n",
"\n",
" sense = PowerSpectrum(\n",
" observation=observation, foreground_model=foreground_model, theory_model=theory_model()\n",
" ).at_frequency(band * un.MHz)\n",
" if calc_2D:\n",
" if calc_2d:\n",
" sense2d_sample = sense.calculate_sensitivity_2d_grid(\n",
" kperp_edges=kperp_edges_hMpc, kpar_edges=kpar_edges_hMpc, thermal=False, sample=True\n",
" kperp_edges=kperp_edges_hmpc, kpar_edges=kpar_edges_hmpc, thermal=False, sample=True\n",
" )\n",
" this_z[\"sample_2D_mK2\"] = sense2d_sample.value\n",
"\n",
" sense2d_thermal = sense.calculate_sensitivity_2d_grid(\n",
" kperp_edges=kperp_edges_hMpc, kpar_edges=kpar_edges_hMpc, thermal=True, sample=False\n",
" kperp_edges=kperp_edges_hmpc, kpar_edges=kpar_edges_hmpc, thermal=True, sample=False\n",
" )\n",
" this_z[\"thermal_2D_mK2\"] = sense2d_thermal.value\n",
"\n",
Expand Down Expand Up @@ -226,7 +223,8 @@
}
],
"source": [
"# aa4.baselines_metres has shape (Nants, Nants, 3): it gives us the basline vectors between all pairs of antennas in x,y,z.\n",
"# aa4.baselines_metres has shape (Nants, Nants, 3):\n",
"# it gives us the basline vectors between all pairs of antennas in x,y,z.\n",
"# We ignore the z coordinate and plot the x and y baselines for all antenna pairs.\n",
"fig, ax = plt.subplots(1, 2, figsize=(14, 6), sharey=True, gridspec_kw={\"wspace\": 0.1})\n",
"im = ax[0].hexbin(\n",
Expand Down Expand Up @@ -379,27 +377,27 @@
"# Only needed if we want to calculate the 2D sensitivity\n",
"# Set the kperp and kpar bins at which we want to calculate the 2D sensitivity\n",
"sense_params = {}\n",
"sense_params[\"calc_2D\"] = True\n",
"sense_params[\"calc_2d\"] = True\n",
"sense_params[\"foreground_model\"] = \"moderate\"\n",
"sense_params[\"redshifts\"] = zs\n",
"sense_params[\"freq_bands\"] = freq_bands\n",
"sense_params[\"kperp_edges_hMpc\"] = un.Quantity(np.logspace(-1.5, 0.3, 15), \"littleh/Mpc\")\n",
"sense_params[\"kpar_edges_hMpc\"] = un.Quantity(np.logspace(-1.5, 0.3, 15), \"littleh/Mpc\")\n",
"sense_params[\"kpar_Mpc\"] = (\n",
"sense_params[\"kperp_edges_hmpc\"] = un.Quantity(np.logspace(-1.5, 0.3, 15), \"littleh/Mpc\")\n",
"sense_params[\"kpar_edges_hmpc\"] = un.Quantity(np.logspace(-1.5, 0.3, 15), \"littleh/Mpc\")\n",
"sense_params[\"kpar_mpc\"] = (\n",
" np.exp(\n",
" (\n",
" np.log(sense_params[\"kpar_edges_hMpc\"].value[:-1])\n",
" + np.log(sense_params[\"kpar_edges_hMpc\"].value[1:])\n",
" np.log(sense_params[\"kpar_edges_hmpc\"].value[:-1])\n",
" + np.log(sense_params[\"kpar_edges_hmpc\"].value[1:])\n",
" )\n",
" / 2\n",
" )\n",
" * observation_params[\"h\"]\n",
")\n",
"sense_params[\"kperp_Mpc\"] = (\n",
"sense_params[\"kperp_mpc\"] = (\n",
" np.exp(\n",
" (\n",
" np.log(sense_params[\"kperp_edges_hMpc\"].value[:-1])\n",
" + np.log(sense_params[\"kperp_edges_hMpc\"].value[1:])\n",
" np.log(sense_params[\"kperp_edges_hmpc\"].value[:-1])\n",
" + np.log(sense_params[\"kperp_edges_hmpc\"].value[1:])\n",
" )\n",
" / 2\n",
" )\n",
Expand All @@ -422,7 +420,7 @@
},
{
"cell_type": "code",
"execution_count": 53,
"execution_count": 84,
"metadata": {},
"outputs": [
{
Expand All @@ -435,8 +433,8 @@
}
],
"source": [
"print(\"AA* beam area:%.5f rad2\" % aaast.beam.area.value) # for 150 MHz (default value)\n",
"print(\"AA4 beam area:%.5f rad2\" % aa4.beam.area.value)"
"print(f\"AA* beam area:{aaast.beam.area.value:.5f} rad2\") # for 150 MHz (default value)\n",
"print(f\"AA4 beam area:{aa4.beam.area.value:.5f} rad2\")"
]
},
{
Expand Down Expand Up @@ -482,7 +480,7 @@
" n_days=observation_params[\"ndays\"],\n",
" cosmo=observation_params[\"cosmo\"],\n",
" coherent=True,\n",
" # phase_center_dec=-30.0*un.deg, # For track scan mode\n",
" # To point to dish off-zenith, use: phase_center_dec=-30.0*un.deg,\n",
" # max baseline is 73km and there are ~260k bls so the gridding takes a very long time.\n",
" # we add a baseline filer to only keep baselines up to 1km.\n",
" baseline_filters=p21sense.BaselineRange(bl_max=1e3 * un.m),\n",
Expand Down Expand Up @@ -741,29 +739,30 @@
},
{
"cell_type": "code",
"execution_count": 82,
"execution_count": 85,
"metadata": {},
"outputs": [],
"source": [
"def compare_senses(\n",
" senses1,\n",
" senses2,\n",
" redshift,\n",
" kperp_Mpc,\n",
" kpar_Mpc,\n",
" kperp_mpc,\n",
" kpar_mpc,\n",
" label1=\"AA*\",\n",
" label2=\"AA4\",\n",
" plot_1d=True,\n",
" **kwargs,\n",
"):\n",
" \"\"\"Plot the 2D sensitivity of two observations at a given redshift.\"\"\"\n",
" # We assume both senses have the same redshifts / freq bands\n",
" if np.all(senses1[\"redshifts\"] != senses2[\"redshifts\"]):\n",
" raise ValueError(\"Redshifts do not match\")\n",
" closest_redshift = np.argmin(np.abs(senses1[\"redshifts\"] - redshift))\n",
" band_name = str(np.round(z2freq(senses1[\"redshifts\"][closest_redshift]), 1)) + \" MHz\"\n",
" if (\n",
" \"sample_and_thermal_2D_mK2\" in senses1[band_name].keys()\n",
" and \"sample_and_thermal_2D_mK2\" in senses2[band_name].keys()\n",
" \"sample_and_thermal_2D_mK2\" in senses1[band_name]\n",
" and \"sample_and_thermal_2D_mK2\" in senses2[band_name]\n",
" ):\n",
" fig, ax = plt.subplots(1, 2, sharey=True, figsize=(12, 6))\n",
" mask1 = ~np.logical_or(\n",
Expand Down Expand Up @@ -791,24 +790,24 @@
" ]\n",
" )\n",
" ax[0].pcolormesh(\n",
" kperp_Mpc,\n",
" kpar_Mpc,\n",
" kperp_mpc,\n",
" kpar_mpc,\n",
" senses1[band_name][\"sample_and_thermal_2D_mK2\"].T,\n",
" norm=colors.LogNorm(vmin=vmin, vmax=vmax),\n",
" )\n",
" ax[0].set_title(label1)\n",
" ax[0].loglog()\n",
"\n",
" im = ax[1].pcolormesh(\n",
" kperp_Mpc,\n",
" kpar_Mpc,\n",
" kperp_mpc,\n",
" kpar_mpc,\n",
" senses2[band_name][\"sample_and_thermal_2D_mK2\"].T,\n",
" norm=colors.LogNorm(vmin=vmin, vmax=vmax),\n",
" )\n",
" ax[1].set_title(label2)\n",
" ax[1].loglog()\n",
" plt.suptitle(\"Redshift z = \" + str(np.round(redshift, 1)))\n",
" cbar = fig.colorbar(im, label=r\"sensitivity [mK$^2$]\")\n",
" fig.colorbar(im, label=r\"sensitivity [mK$^2$]\")\n",
" ax[0].set_xlabel(r\"$k_\\perp$ [Mpc$^{-1}$]\")\n",
" ax[1].set_xlabel(r\"$k_\\perp$ [Mpc$^{-1}$]\")\n",
" ax[0].set_ylabel(r\"$k_\\parallel$ [Mpc$^{-1}$]\")\n",
Expand All @@ -833,6 +832,7 @@
"\n",
"\n",
"def compare_all_senses(senses, redshift=None, k=None, labels=None, colors=None, lss=None):\n",
" \"\"\"Plot the sensitivity of observations at a given z vs k and at a given k vs z.\"\"\"\n",
" # We assume all senses have the same redshifts / freq bands and k bins\n",
" closest_redshift = np.argmin(np.abs(senses[0][\"redshifts\"] - redshift))\n",
" band_name = str(np.round(z2freq(senses[0][\"redshifts\"][closest_redshift]), 1)) + \" MHz\"\n",
Expand Down

0 comments on commit 3889854

Please sign in to comment.