diff --git a/docs/tutorials/SKA_forecast.ipynb b/docs/tutorials/SKA_forecast.ipynb index 1f5e1e6..47fc36f 100644 --- a/docs/tutorials/SKA_forecast.ipynb +++ b/docs/tutorials/SKA_forecast.ipynb @@ -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})" ] }, { @@ -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", @@ -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", @@ -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", @@ -422,7 +420,7 @@ }, { "cell_type": "code", - "execution_count": 53, + "execution_count": 84, "metadata": {}, "outputs": [ { @@ -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\")" ] }, { @@ -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", @@ -741,7 +739,7 @@ }, { "cell_type": "code", - "execution_count": 82, + "execution_count": 85, "metadata": {}, "outputs": [], "source": [ @@ -749,21 +747,22 @@ " 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", @@ -791,8 +790,8 @@ " ]\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", @@ -800,15 +799,15 @@ " 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", @@ -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",