import os
from builtins import zip
import astropy.units as u
import matplotlib.pyplot as plt
import pytest
from threeML import *
from threeML.io.calculate_flux import _calculate_point_source_flux
from threeML.io.package_data import get_path_of_data_dir
from threeML.plugins.OGIPLike import OGIPLike
from threeML.utils.fitted_objects.fitted_point_sources import InvalidUnitError
# Init some globals
datadir = os.path.abspath(
os.path.join(get_path_of_data_dir(), "datasets", "bn090217206")
)
good_d_flux_units = ["1/(cm2 s keV)", "erg/(cm2 s keV)", "erg2/(cm2 s keV)"]
good_i_flux_units = ["1/(cm2 s )", "erg/(cm2 s )", "erg2/(cm2 s )"]
good_energy_units = ["keV", "Hz", "nm"]
bad_flux_units = ["g"]
[docs]
def make_simple_model():
triggerName = "bn090217206"
ra = 204.9
dec = -8.4
powerlaw = Powerlaw()
GRB = PointSource(triggerName, ra, dec, spectral_shape=powerlaw)
model = Model(GRB)
powerlaw.index.prior = Uniform_prior(lower_bound=-5.0, upper_bound=5.0)
powerlaw.K.prior = Log_uniform_prior(lower_bound=1.0, upper_bound=10)
return model
[docs]
def make_components_model():
triggerName = "bn090217206"
ra = 204.9
dec = -8.4
powerlaw = Powerlaw() + Blackbody()
GRB = PointSource(triggerName, ra, dec, spectral_shape=powerlaw)
model = Model(GRB)
powerlaw.index_1.prior = Uniform_prior(lower_bound=-5.0, upper_bound=5.0)
powerlaw.K_1.prior = Log_uniform_prior(lower_bound=1.0, upper_bound=10)
powerlaw.K_2.prior = Uniform_prior(lower_bound=-5.0, upper_bound=5.0)
powerlaw.kT_2.prior = Log_uniform_prior(lower_bound=1.0, upper_bound=10)
return model
[docs]
def make_dless_components_model():
triggerName = "bn090217206"
ra = 204.9
dec = -8.4
powerlaw = Powerlaw() * Constant()
GRB = PointSource(triggerName, ra, dec, spectral_shape=powerlaw)
model = Model(GRB)
powerlaw.index_1.prior = Uniform_prior(lower_bound=-5.0, upper_bound=5.0)
powerlaw.K_1.prior = Log_uniform_prior(lower_bound=1.0, upper_bound=10)
powerlaw.k_2 = 1.0
powerlaw.k_2.fix = True
return model
[docs]
@pytest.fixture
def analysis_to_test(data_list_bn090217206_nai6):
simple_model = make_simple_model()
complex_model = make_components_model()
# prepare mle
dless_model = make_dless_components_model()
jl_simple = JointLikelihood(simple_model, data_list_bn090217206_nai6)
jl_simple.fit()
jl_complex = JointLikelihood(complex_model, data_list_bn090217206_nai6)
jl_complex.fit()
jl_dless = JointLikelihood(dless_model, data_list_bn090217206_nai6)
jl_dless.fit()
bayes_simple = BayesianAnalysis(simple_model, data_list_bn090217206_nai6)
bayes_simple.set_sampler("emcee")
bayes_simple.sampler.setup(n_iterations=10, n_burn_in=10, n_walkers=20)
bayes_simple.sample()
bayes_complex = BayesianAnalysis(complex_model, data_list_bn090217206_nai6)
bayes_complex.set_sampler("emcee")
bayes_complex.sampler.setup(n_iterations=10, n_burn_in=10, n_walkers=20)
bayes_complex.sample()
bayes_dless = BayesianAnalysis(dless_model, data_list_bn090217206_nai6)
bayes_dless.set_sampler("emcee")
bayes_dless.sampler.setup(n_iterations=10, n_burn_in=10, n_walkers=20)
bayes_dless.sample()
analysis_to_test = [
jl_simple.results,
jl_complex.results,
jl_dless.results,
bayes_simple.results,
bayes_complex.results,
bayes_dless.results,
]
return analysis_to_test
[docs]
def test_fitted_point_source_plotting(analysis_to_test):
plot_keywords = {
"use_components": True,
"components_to_use": ["Powerlaw", "total"],
"sources_to_use": ["bn090217206"],
"flux_unit": "erg/(cm2 s)",
"energy_unit": "keV",
"plot_style_kwargs": {},
"contour_style_kwargs": {},
"legend_kwargs": {},
"ene_min": 10,
"ene_max": 100,
"num_ene": 5,
"show_legend": False,
"fit_cmap": "jet",
"countor_cmap": "jet",
"sum_sources": True,
}
for u1, u2 in zip(good_d_flux_units, good_i_flux_units):
for e_unit in good_energy_units:
for x in analysis_to_test:
_ = plot_spectra(
x, flux_unit=u1, energy_unit=e_unit, num_ene=5
)
_ = plot_spectra(x, **plot_keywords)
with pytest.raises(InvalidUnitError):
_ = plot_spectra(x, flux_unit=bad_flux_units[0])
plt.close("all")
[docs]
def test_fitted_point_source_flux_calculations(analysis_to_test):
flux_keywords = {
"use_components": True,
"components_to_use": ["total", "Powerlaw"],
"sources_to_use": ["bn090217206"],
"flux_unit": "erg/(cm2 s)",
"energy_unit": "keV",
"sum_sources": True,
}
_calculate_point_source_flux(
1, 10, analysis_to_test[0], flux_unit=good_i_flux_units[0], energy_unit="keV"
)
_calculate_point_source_flux(1, 10, analysis_to_test[-2], **flux_keywords)
[docs]
def test_units_on_energy_range(analysis_to_test):
_ = plot_spectra(
analysis_to_test[0], ene_min=1.0 * u.keV, ene_max=1 * u.MeV
)
with pytest.raises(RuntimeError):
plot_spectra(analysis_to_test[0], ene_min=1.0, ene_max=1 * u.MeV)
with pytest.raises(RuntimeError):
plot_spectra(analysis_to_test[0], ene_min=1.0 * u.keV, ene_max=1.0)