def test_create_tables_normal(self):
method_list = methods.builtin_method_functions()
cosmology_list = cosmologies.builtin_cosmology_functions()
# Create a lookup table for each method and cosmology
for method in method_list:
for key in cosmology_list:
here = os.getcwd()
cosmo = cosmologies.builtin_cosmology_functions()[key]
filename = "_".join(["pytest_output", method, key])
table.create(method=method, filename=filename,
cosmo=cosmo, output_dir=here, zmin=0,
zmax=20, num_samples=10000)
# Compare new tables to existing tables for 4 dm values
pre_calc_fn = ".".join(["_".join([method, key]), "npz"])
new_calc_fn = "".join([filename, ".npz"])
pre_calc = table.load(pre_calc_fn)
new_calc = table.load(new_calc_fn, data_dir=here)
test_dm_list = [0, 100, 1000, 2000]
for dm in test_dm_list:
new_z = table.get_z_from_table(dm, new_calc)
pre_z = table.get_z_from_table(dm, pre_calc)
assert new_z == pre_z
def cosmology_comparison(filename="",
extension="png",
usetex=False,
passed_ax=None,
**kwargs):
"""
Create a plot comparing how the estimated redshift changes as a
function of dispersion mesure for each cosmology.
Parameters
----------
filename: string, optional
The filename of the saved figure. Default: "cosmology_comparison"
extension: string, optional
The format to save the figure. e.g "png", "pdf", "eps", etc...
Default: "png"
Generates
---------
A figure displaying how estimated redshift changes as a function of
dispersion measure for each of the different cosmologies.
"""
# Update rcParams for consistent plotting style
plt.rcParams.update(set_rc_params(usetex=usetex))
if passed_ax:
ax = passed_ax
else:
fig = plt.figure(figsize=(8, 8), constrained_layout=True)
ax = fig.add_subplot(111)
# Remove EAGLE from cosmologies since it is the same as Planck13
cosmology_list = cosmologies.builtin_cosmology_functions()
cosmology_list.pop("EAGLE")
dm_vals = np.linspace(0, 3000, 1000)
add_axin = True
try:
# Add inset plot showing the part where cosmologies diverge the most.
axin = ax.inset_axes([0.05, 0.52, 0.45, 0.45])
except Exception:
print("""Skipping inset axis in cosmology plot. Requires Python 3 and
Matplotlib 3.0""")
add_axin = False
colours = [
'#a6cee3', '#1f78b4', '#b2df8a', '#33a02c', '#fb9a99', '#e31a1c'
]
label = [
r"$\rm{WMAP5}$", r"$\rm{WMAP7}$", r"$\rm{WMAP9}$", r"$\rm{Planck13}$",
r"$\rm{Planck15}$", r"$\rm{Planck18}$"
]
for j, cosmo in enumerate(cosmology_list):
z_vals = np.zeros(len(dm_vals))
if 'method' in kwargs:
method = kwargs['method']
else:
method = 'Inoue2004'
table_name = "{}.hdf5".format(method)
table_name = utils.get_path_to_file_from_here(table_name,
subdirs=["data"])
for i, dm in enumerate(dm_vals):
z_vals[i] = table.get_z_from_table(dm, table_name, cosmo)
ax.plot(dm_vals, z_vals, colours[j], label=label[j], **kwargs)
if add_axin:
axin.plot(dm_vals, z_vals, colours[j], **kwargs)
ax.set_xlabel(r"$\rm{DM\ \left[pc \ cm^{-3}\right]}$")
if not passed_ax:
ax.set_ylabel(r"$\rm{Redshift}$")
ax.legend(loc='lower right', frameon=False)
if add_axin:
axin.set_xlim(2800, 3000)
axin.set_ylim(3.0, 3.25)
axin.xaxis.set_tick_params(labelsize=8)
axin.yaxis.set_tick_params(labelsize=8)
ax.indicate_inset_zoom(axin)
if filename != "":
plt.savefig(".".join([filename, extension]))
if passed_ax:
return ax
else:
return fig
def method_comparison(filename=None,
extension="png",
usetex=False,
passed_ax=None,
**kwargs):
"""
Create a plot comparing how estimated redshift changes as a
function of dispersion measure for each DM-z relation.
Parameters
----------
filename: string or None, optional
The filename of the saved figure. Default: *None*
extension: string, optional
The format to save the figure. e.g "png", "pdf", "eps", etc...
Default: "png"
usetex: bool, optional
Use LaTeX for for fonts.
passed_ax: or None, optional
Generates
---------
A figure displaying how estimated redshift changes as a function of
dispersion measure for each of the different cosmologies.
"""
# Update rcParams for consistent plotting style
plt.rcParams.update(set_rc_params(usetex=usetex))
if passed_ax:
ax = passed_ax
else:
fig = plt.figure(figsize=(8, 8), constrained_layout=True)
ax = fig.add_subplot(111)
method_list = methods.builtin_method_functions()
method_list.pop("Batten2021")
dm_vals = np.linspace(0, 3000, 1000)
colours = ["#1b9e77", "#d95f02", "#7570b3"]
label = [r"$\rm{Ioka 2003}$", r"$\rm{Inoue 2004}$", r"$\rm{Zhang 2018}$"]
for j, method in enumerate(method_list):
z_vals = np.zeros(len(dm_vals))
if 'cosmology' in kwargs:
cosmology = kwargs['cosmology']
else:
cosmology = 'Planck18'
table_name = "{}.hdf5".format(method)
table_name = utils.get_path_to_file_from_here(table_name,
subdirs=["data"])
for i, dm in enumerate(dm_vals):
z_vals[i] = table.get_z_from_table(dm, table_name, cosmology)
ax.plot(dm_vals, z_vals, colours[j], label=label[j], **kwargs)
if not passed_ax:
ax.set_ylabel(r"$\rm{Redshift}$")
ax.set_xlabel(r"$\rm{DM\ \left[pc \ cm^{-3}\right]}$")
ax.legend(loc='lower right', frameon=False)
if filename is not None:
plt.savefig(".".join([filename, extension]))
if passed_ax:
return ax
else:
return fig
def calc_redshift(self,
method='Batten2021',
cosmology="Planck18",
subtract_host=False,
lookup_table=None):
"""
Calculate the redshift of the FRB from its :attr:`dm`,
:attr:`dm_excess` or :attr:`dm_excess` - :attr:`dm_host_est`.
Parameters
----------
method : str, optional
The dispersion meausre -redshift relation to use when
calculating the redshift. Avaliable methods: %(meth)s.
Default: 'Inoue2004'
cosmology : str, optional
The cosmology to assume when calculating the redshift.
Avaliable cosmologies: %(cosmo)s. Default: 'Planck18'
subtract_host : bool, optional
Subtract :attr:`dm_host_est` from the :attr:`dm_excess`
before calculating the redshift. This is is used to account
for the dispersion measure that arises from the FRB host
galaxy. Default: False
lookup_table : str or None, optional
The path to the lookup table file. If ``lookup_table=None``
a table will attempted to be loaded from the data directory
based on the method name. Default: *None*
Returns
-------
float
The redshift of the FRB.
Notes
-----
The methods_ section in the documentation has a description for
each of the builtin methods.
The cosmology_ section in the documentation has a list of the
cosmological parameters for each cosmology
.. _cosmology: https://fruitbat.readthedocs.io/en/latest/user_guide/method_and_cosmology.html#cosmology
.. _methods: https://fruitbat.readthedocs.io/en/latest/user_guide/method_and_cosmology.html#methods
"""
utils.check_type("subtract_host", subtract_host, bool)
if subtract_host:
input_dm = self.dm_excess - self.dm_host_est
print(input_dm)
else:
input_dm = self.dm_excess
# If the user provides a table use that table for estimation.
if lookup_table is not None:
raise NotImplementedError("Not made yet!")
if method not in methods.available_methods():
err_msg = ("Method '{}' is not a valid method. "
"Valid methods are: {}".format(
method, methods.available_methods()))
raise ValueError(err_msg)
if method in methods.methods_hydrodynamic():
zvals, pdf, dz = self.calc_redshift_pdf(method=method)
z_median = utils.calc_median_from_pdf(zvals, pdf)
self.z = z_median
self.method = method
if method == "Batten2020":
self.cosmology = "EAGLE"
elif method in methods.methods_analytic():
if cosmology in cosmologies.available_cosmologies():
#if method in methods.builtin_method_functions().keys():
table_name = table.get_table_path(method)
self.z = table.get_z_from_table(input_dm, table_name,
cosmology)
self.cosmology = cosmology
self.method = method
else:
err_msg = ("Cosmology '{}' is not a valid cosmology. Valid "
"cosmologies are: {}".format(
cosmology, cosmologies.available_cosmologies()))
raise ValueError(err_msg)
else:
table_name = utils.get_path_to_file_from_here(
"{}.hdf5".format(method), subdirs=["data"])
self.z = table.get_z_from_table(input_dm, table_name)
self.cosmology = cosmology
self.method = method
return self.z
def calc_redshift(self,
method='Inoue2004',
cosmology="Planck18",
subtract_host=False,
lookup_table=None):
"""
Calculate the redshift of the FRB from its :attr:`dm`,
:attr:`dm_excess` or :attr:`dm_excess` - :attr:`dm_host_est`.
Parameters
----------
method : str, optional
The dispersion meausre -redshift relation to use when
calculating the redshift. Avaliable methods: %(meth)s.
Default: 'Inoue2004'
cosmology : str, optional
The cosmology to assume when calculating the redshift.
Avaliable cosmologies: %(cosmo)s. Default: 'Planck18'
subtract_host : bool, optional
Subtract :attr:`dm_host_est` from the :attr:`dm_excess`
before calculating the redshift. This is is used to account
for the dispersion measure that arises from the FRB host
galaxy. Default: False
lookup_table : str or None, optional
The path to the lookup table file. If ``lookup_table=None``
a table will attempted to be loaded from the data directory
based on the method name. Default: *None*
Returns
-------
float
The redshift of the FRB.
Notes
-----
The methods_ section in the documentation has a description for
each of the builtin methods.
The cosmology_ section in the documentation has a list of the
cosmological parameters for each cosmology
.. _cosmology: https://fruitbat.readthedocs.io/en/latest/user_guide/method_and_cosmology.html#cosmology
.. _methods: https://fruitbat.readthedocs.io/en/latest/user_guide/method_and_cosmology.html#methods
"""
if not isinstance(subtract_host, bool):
raise ValueError("subtract_host must be of type bool.")
if subtract_host:
input_dm = self.dm_excess - self.dm_host_est
else:
input_dm = self.dm_excess
if method not in methods.available_methods():
err_msg = ("Method '{}' is not a valid method. "
"Valid methods are: {}".format(
method, methods.available_methods()))
raise ValueError(err_msg)
if cosmology not in cosmologies.available_cosmologies():
err_msg = ("Cosmology '{}' is not a valid cosmology. Valid "
"cosmologies are: {}".format(
cosmology, cosmologies.available_cosmologies()))
raise ValueError(err_msg)
# If the user provides a table use that table for estimation.
if lookup_table is not None:
lookup_table = table.load(lookup_table)
else:
if method in methods.builtin_method_functions().keys():
table_name = "".join(["_".join([method, cosmology]), ".npz"])
else:
table_name = "".join(["custom_", method, ".npz"])
lookup_table = table.load(table_name)
self.z = table.get_z_from_table(input_dm, lookup_table)
lookup_table.close()
self.cosmology = cosmology
self.method = method
return self.z
method_list = methods.available_methods()
colours = ["#66c2a5", "#e7298a", "#8da0cb"]
label = [r"$\rm{Ioka\ 2003}$", r"$\rm{Inoue\ 2004}$", r"$\rm{Zhang\ 2018}$"]
lstyle = ["-", "--", "-."]
for j, method in enumerate(method_list):
z_vals = np.zeros(len(dm_vals))
cosmology = 'Planck18'
table_name = "".join(["_".join([method, cosmology]), ".npz"])
lookup_table = table.load(table_name)
for i, dm in enumerate(dm_vals):
z_vals[i] = table.get_z_from_table(dm, lookup_table)
ax.plot(z_vals,
np.log10(dm_vals),
colours[j],
linestyle=lstyle[j],
label=label[j],
linewidth=2)
#divider = make_axes_locatable(ax)
#cax = divider.append_axes("right", size="5%", pad=0.05)
legend = plt.legend(frameon=False, fontsize=20)
plt.setp(legend.get_texts(), color='w')
#plt.colorbar(im, cax=cax)
plt.savefig("postimage_magma_lines_colorbar.png", dpi=800)