Source code for threeML.io.plotting.model_plot

__author__ = "grburgess"

import warnings
from typing import List

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
from astropy.visualization import quantity_support
from threeML.config.config import threeML_config
from threeML.io.calculate_flux import (
    _collect_sums_into_dictionaries,
    _setup_analysis_dictionaries,
)
from threeML.io.logging import setup_logger
from threeML.io.package_data import get_path_of_data_file
from threeML.io.plotting.cmap_cycle import cmap_intervals
from threeML.utils.progress_bar import tqdm

if threeML_config.plotting.use_threeml_style:

    plt.style.use(str(get_path_of_data_file("threeml.mplstyle")))


log = setup_logger(__name__)


[docs] def plot_point_source_spectra(*analysis_results, **kwargs): log.error( "plot_point_source_spectra() has been replaced by plot_spectra()." ) return plot_spectra(*analysis_results, **kwargs)
[docs] def plot_spectra(*analysis_results, **kwargs) -> plt.Figure: """ plotting routine for fitted point source spectra :param analysis_results: fitted JointLikelihood or BayesianAnalysis objects :param sources_to_use: (optional) list of PointSource string names to plot from the analysis :param energy_unit: (optional) astropy energy unit in string form (can also be frequency) :param flux_unit: (optional) astropy flux unit in string form :param confidence_level: (optional) confidence level to use (default: 0.68) :param ene_min: (optional) minimum energy to plot :param ene_max: (optional) maximum energy to plot :param num_ene: (optional) number of energies to plot :param use_components: (optional) True or False to plot the spectral components :param components_to_use: (optional) list of string names of the components to plot: including 'total' will also plot the total spectrum :param sum_sources: (optional) some all the MLE and Bayesian sources :param show_contours: (optional) True or False to plot the contour region :param plot_style_kwargs: (optional) dictionary of MPL plot styling for the best fit curve :param contour_style_kwargs: (optional) dictionary of MPL plot styling for the contour regions :param fit_cmap: MPL color map to iterate over for plotting multiple analyses :param contour_cmap: MPL color map to iterate over for plotting contours for multiple analyses :param subplot: subplot to use :param xscale: 'log' or 'linear' :param yscale: 'log' or 'linear' :param include_extended: True or False, also plot extended source spectra. :return: """ # allow matplotlib to plot quantities to the access quantity_support() _sub_menu = threeML_config.model_plot.point_source_plot _defaults = { "fit_cmap": _sub_menu.fit_cmap.value, "contour_cmap": _sub_menu.contour_cmap.value, "contour_colors": None, "fit_colors": None, "confidence_level": 0.68, "equal_tailed": True, "best_fit": "median", "energy_unit": _sub_menu.ene_unit, "flux_unit": _sub_menu.flux_unit, "ene_min": _sub_menu.emin, "ene_max": _sub_menu.emax, "num_ene": _sub_menu.num_ene, "use_components": False, "components_to_use": [], "sources_to_use": [], "sum_sources": False, "show_contours": True, "plot_style_kwargs": _sub_menu.plot_style, "contour_style_kwargs": _sub_menu.contour_style, "show_legend": _sub_menu.show_legend, "legend_kwargs": _sub_menu.legend_style, "subplot": None, "xscale": "log", "yscale": "log", "include_extended": False, } for key, value in kwargs.items(): if key in _defaults: _defaults[key] = value if isinstance(_defaults["ene_min"], u.Quantity): if not isinstance(_defaults["ene_max"], u.Quantity): log.error("both energy arguments must be Quantities") raise RuntimeError() if isinstance(_defaults["ene_max"], u.Quantity): if not isinstance(_defaults["ene_min"], u.Quantity): log.error("both energy arguments must be Quantities") raise RuntimeError() if isinstance(_defaults["ene_max"], u.Quantity): energy_range = np.linspace( _defaults["ene_min"], _defaults["ene_max"], _defaults["num_ene"] ) # type: u.Quantity # _defaults["energy_unit"] = energy_range.unit if _defaults["xscale"] == "log": energy_range = ( np.logspace( np.log10(energy_range.min().value), np.log10(energy_range.max().value), _defaults["num_ene"], ) * energy_range.unit ) energy_range = energy_range.to( _defaults["energy_unit"], equivalencies=u.spectral() ) else: energy_range = np.logspace( np.log10(_defaults["ene_min"]), np.log10(_defaults["ene_max"]), _defaults["num_ene"], ) * u.Unit(_defaults["energy_unit"]) # scale the units to the defaults _defaults["ene_min"] = _defaults["ene_min"] * u.Unit( _defaults["energy_unit"] ) _defaults["ene_max"] = _defaults["ene_max"] * u.Unit( _defaults["energy_unit"] ) ( mle_analyses, bayesian_analyses, num_sources_to_plot, duplicate_keys, ) = _setup_analysis_dictionaries( analysis_results, energy_range, _defaults["energy_unit"], _defaults["flux_unit"], _defaults["use_components"], _defaults["components_to_use"], _defaults["confidence_level"], _defaults["equal_tailed"], differential=True, sources_to_use=_defaults["sources_to_use"], include_extended=_defaults["include_extended"], ) # we are now ready to plot. # all calculations have been made. # if we are not going to sum sources if not _defaults["sum_sources"]: if _defaults["fit_colors"] is None: color_fit = cmap_intervals( num_sources_to_plot + 1, _defaults["fit_cmap"] ) else: # duck typing if isinstance(_defaults["fit_colors"], (str, str)): color_fit = [_defaults["fit_colors"]] * num_sources_to_plot elif isinstance(_defaults["fit_colors"], list): assert len(_defaults["fit_colors"]) == num_sources_to_plot, ( "list of colors (%d) must be the same length as sources ot plot (%s)" % (len(_defaults["fit_colors"]), num_sources_to_plot) ) color_fit = _defaults["fit_colors"] else: raise ValueError( "Can not setup color, wrong type:", type(_defaults["fit_colors"]), ) if _defaults["contour_colors"] is None: color_contour = cmap_intervals( num_sources_to_plot + 1, _defaults["contour_cmap"] ) else: # duck typing if isinstance(_defaults["contour_colors"], (str, str)): color_contour = [ _defaults["contour_colors"] ] * num_sources_to_plot elif isinstance(_defaults["contour_colors"], list): assert ( len(_defaults["contour_colors"]) == num_sources_to_plot ), ( "list of colors (%d) must be the same length as sources ot plot (%s)" % (len(_defaults["contour_colors"]), num_sources_to_plot) ) color_contour = _defaults["fit_colors"] else: raise ValueError( "Can not setup contour color, wrong type:", type(_defaults["contour_colors"]), ) color_itr = 0 # go thru the mle analysis and plot spectra plotter = SpectralContourPlot( num_sources_to_plot, xscale=_defaults["xscale"], yscale=_defaults["yscale"], show_legend=_defaults["show_legend"], plot_kwargs=_defaults["plot_style_kwargs"], contour_kwargs=_defaults["contour_style_kwargs"], legend_kwargs=_defaults["legend_kwargs"], emin=_defaults["ene_min"], emax=_defaults["ene_max"], subplot=_defaults["subplot"], ) for key in list(mle_analyses.keys()): # we won't assume to plot the total until the end plot_total = False if _defaults["use_components"]: # if this source has no components or none that we wish to plot # then we will plot the total spectrum after this if (not list(mle_analyses[key]["components"].keys())) or ( "total" in _defaults["components_to_use"] ): plot_total = True for component in list(mle_analyses[key]["components"].keys()): positive_error = None negative_error = None # extract the information and plot it if _defaults["best_fit"] == "average": best_fit = mle_analyses[key]["components"][ component ].average else: best_fit = mle_analyses[key]["components"][ component ].median if _defaults["show_contours"]: positive_error = mle_analyses[key]["components"][ component ].upper_error negative_error = mle_analyses[key]["components"][ component ].lower_error neg_mask = negative_error <= 0 # replace with small number negative_error[neg_mask] = min(best_fit) * 0.9 label = "%s: %s" % (key, component) # this is where we keep track of duplicates if key in duplicate_keys: label = "%s: MLE" % label if mle_analyses[key]["components"][ component ].is_dimensionless: plotter.add_dimensionless_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label=label, ) else: plotter.add_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label=label, ) color_itr += 1 else: plot_total = True if plot_total: # it ends up that we need to plot the total spectrum # which is just a repeat of the process if _defaults["best_fit"] == "average": best_fit = mle_analyses[key]["fitted point source"].average else: best_fit = mle_analyses[key]["fitted point source"].median if _defaults["show_contours"]: positive_error = mle_analyses[key][ "fitted point source" ].upper_error negative_error = mle_analyses[key][ "fitted point source" ].lower_error neg_mask = negative_error <= 0 # replace with small number negative_error[neg_mask] = min(best_fit) * 0.9 else: positive_error = None negative_error = None label = "%s" % key if key in duplicate_keys: label = "%s: MLE" % label plotter.add_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label=label, ) color_itr += 1 # we will do the exact same thing for the bayesian analyses for key in list(bayesian_analyses.keys()): plot_total = False if _defaults["use_components"]: if (not list(bayesian_analyses[key]["components"].keys())) or ( "total" in _defaults["components_to_use"] ): plot_total = True for component in list( bayesian_analyses[key]["components"].keys() ): positive_error = None negative_error = None if _defaults["best_fit"] == "average": best_fit = bayesian_analyses[key]["components"][ component ].average else: best_fit = bayesian_analyses[key]["components"][ component ].median if _defaults["show_contours"]: positive_error = bayesian_analyses[key]["components"][ component ].upper_error negative_error = bayesian_analyses[key]["components"][ component ].lower_error label = "%s: %s" % (key, component) if key in duplicate_keys: label = "%s: Bayesian" % label if bayesian_analyses[key]["components"][ component ].is_dimensionless: plotter.add_dimensionless_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label=label, ) else: plotter.add_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label=label, ) color_itr += 1 else: plot_total = True if plot_total: if _defaults["best_fit"] == "average": best_fit = bayesian_analyses[key][ "fitted point source" ].average else: best_fit = bayesian_analyses[key][ "fitted point source" ].median positive_error = None negative_error = None if _defaults["show_contours"]: positive_error = bayesian_analyses[key][ "fitted point source" ].upper_error negative_error = bayesian_analyses[key][ "fitted point source" ].lower_error label = "%s" % key if key in duplicate_keys: label = "%s: Bayesian" % label plotter.add_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label=label, ) color_itr += 1 else: # now we sum sources instead # we keep MLE and Bayes apart because it makes no # sense to sum them together ( total_analysis_mle, component_sum_dict_mle, num_sources_to_plot, ) = _collect_sums_into_dictionaries( mle_analyses, _defaults["use_components"], _defaults["components_to_use"], ) ( total_analysis_bayes, component_sum_dict_bayes, num_sources_to_plot_bayes, ) = _collect_sums_into_dictionaries( bayesian_analyses, _defaults["use_components"], _defaults["components_to_use"], ) num_sources_to_plot += num_sources_to_plot_bayes plotter = SpectralContourPlot( num_sources_to_plot, xscale=_defaults["xscale"], yscale=_defaults["yscale"], show_legend=_defaults["show_legend"], plot_kwargs=_defaults["plot_style_kwargs"], contour_kwargs=_defaults["contour_style_kwargs"], legend_kwargs=_defaults["legend_kwargs"], emin=_defaults["ene_min"], emax=_defaults["ene_max"], subplot=_defaults["subplot"], ) color_fit = cmap_intervals(num_sources_to_plot, _defaults["fit_cmap"]) color_contour = cmap_intervals( num_sources_to_plot, _defaults["contour_cmap"] ) color_itr = 0 if _defaults["use_components"] and list(component_sum_dict_mle.keys()): # we have components to plot for component, values in component_sum_dict_mle.items(): summed_analysis = sum(values) if _defaults["best_fit"] == "average": best_fit = summed_analysis.average else: best_fit = summed_analysis.median positive_error = None negative_error = None if _defaults["show_contours"]: positive_error = summed_analysis.upper_error negative_error = summed_analysis.lower_error neg_mask = negative_error <= 0 # replace with small number negative_error[neg_mask] = min(best_fit) * 0.9 if np.any( [ c.is_dimensionless for c in component_sum_dict_mle[component] ] ): plotter.add_dimensionless_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label="%s: MLE" % component, ) else: plotter.add_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label="%s: MLE" % component, ) color_itr += 1 if total_analysis_mle: # we will sum and plot the total # analysis summed_analysis = sum(total_analysis_mle) if _defaults["best_fit"] == "average": best_fit = summed_analysis.average else: best_fit = summed_analysis.median positive_error = None negative_error = None if _defaults["show_contours"]: positive_error = best_fit + summed_analysis.upper_error negative_error = best_fit - summed_analysis.lower_error neg_mask = negative_error <= 0 # replace with small number negative_error[neg_mask] = min(best_fit) * 0.9 plotter.add_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label="total: MLE", ) color_itr += 1 if _defaults["use_components"] and list( component_sum_dict_bayes.keys() ): # we have components to plot for component, values in component_sum_dict_bayes.items(): summed_analysis = sum(values) if _defaults["best_fit"] == "average": best_fit = summed_analysis.average else: best_fit = summed_analysis.median positive_error = None negative_error = None if _defaults["show_contours"]: positive_error = summed_analysis.upper_error negative_error = summed_analysis.lower_error if np.any( [ c.is_dimensionless for c in component_sum_dict_bayes[component] ] ): plotter.add_dimensionless_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label="%s: Bayesian" % component, ) else: plotter.add_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label="%s: Bayesian" % component, ) color_itr += 1 if total_analysis_bayes: # we will sum and plot the total # analysis summed_analysis = sum(total_analysis_bayes) if _defaults["best_fit"] == "average": best_fit = summed_analysis.average else: best_fit = summed_analysis.median positive_error = None negative_error = None if _defaults["show_contours"]: positive_error = summed_analysis.upper_error negative_error = summed_analysis.lower_error plotter.add_model( energy_range=energy_range, best_fit=best_fit, color=color_fit[color_itr], upper_error=positive_error, lower_error=negative_error, contour_color=color_contour[color_itr], label="total: Bayesian", ) color_itr += 1 return plotter.finalize(_defaults)
[docs] class SpectralContourPlot: def __init__( self, n_total, xscale="log", yscale="log", show_legend=True, plot_kwargs=None, contour_kwargs=None, legend_kwargs=None, emin=None, emax=None, subplot=None, ): self._n_total = n_total self._show_legend = show_legend self._legend_kwargs = legend_kwargs self._emin = emin self._emax = emax self._plot_kwargs = plot_kwargs self._contour_kwargs = contour_kwargs if subplot is None: self._fig, self._ax = plt.subplots() else: self._ax = subplot self._fig = self._ax.get_figure() self._ax_right = None self._n_plotted = 0 self._xscale = xscale self._yscale = yscale
[docs] def add_model( self, energy_range, best_fit, color, upper_error=None, lower_error=None, contour_color=None, label="model", ): self._ax.plot( energy_range, best_fit, color=color, label=label, **self._plot_kwargs, ) if (upper_error is not None) and (lower_error is not None): self._ax.fill_between( energy_range, lower_error, upper_error, facecolor=contour_color, **self._contour_kwargs, )
[docs] def add_dimensionless_model( self, energy_range, best_fit, color, upper_error=None, lower_error=None, contour_color=None, label="model", ): if self._n_total > 1: if self._ax_right is None: self._ax_right = self._ax.twinx() self._ax_right.plot( energy_range, best_fit, color=color, label=label, **self._plot_kwargs, ) if (upper_error is not None) and (lower_error is not None): self._ax_right.fill_between( energy_range, lower_error, upper_error, facecolor=contour_color, **self._contour_kwargs, ) else: self.add_model( energy_range, best_fit, color, upper_error, lower_error, contour_color, label, )
[docs] def finalize(self, _defaults): self._ax.set_xscale(self._xscale) self._ax.set_yscale(self._yscale) if self._show_legend: self._ax.legend(**self._legend_kwargs) if self._ax_right is not None: self._ax_right.set_yscale(self._yscale) self._ax_right.set_ylabel("Arbitrary units") if self._show_legend: self._ax_right.legend(**self._legend_kwargs) log.debug(f'converting {self._emin.unit} to {_defaults["energy_unit"]}') try: self._ax.set_xlim( [ self._emin.to( _defaults["energy_unit"], equivalencies=u.spectral() ), self._emax.to( _defaults["energy_unit"], equivalencies=u.spectral() ), ] ) except: pass if isinstance(self._emin, u.Quantity) and self._show_legend: # This workaround is needed because of a bug in astropy that would break the plotting of the legend # (see issue #7504 in the Astropy github repo) eemin = self._emin.to( self._ax.xaxis.get_units(), equivalencies=u.spectral() ).value eemax = self._emax.to( self._ax.xaxis.get_units(), equivalencies=u.spectral() ).value self._ax.set_xlim([eemin, eemax]) self._ax.xaxis.converter = None return self._fig