Source code for pesummary.utils.probability_dict

# Licensed under an MIT style license -- see LICENSE.md

from pesummary.utils.samples_dict import SamplesDict
from pesummary.utils.dict import Dict
from pesummary.utils.pdf import DiscretePDF, DiscretePDF2D
from pesummary.core.plots.latex_labels import latex_labels
from pesummary.gw.plots.latex_labels import GWlatex_labels

__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]

latex_labels.update(GWlatex_labels)


[docs]class ProbabilityDict(Dict): """Class to store integers with discrete probabilities for multiple parameters. Parameters ---------- args: dict dictionary containing the discrete probabilities for each parameter. Key should be the parameter name and value should be a 2d array of length 2. First element integers and second element probabilities corresponding to integers. See pesummary.utils.pdf.DiscretePDF for more details Methods ------- rvs: randomly draw samples for each distribution in ProbabilityDict Examples -------- >>> from pesummary.utils.probability_dict import ProbabilityDict >>> numbers = [1,2,3,4] >>> probabilities = [0.1, 0.2, 0.3, 0.4] >>> pdf = ProbabilityDict({"parameter": [numbers, probabilities]}) >>> print(type(pdf["parameter"])) <class 'pesummary.utils.pdf.DiscretePDF'> >>> print(pdf["parameter"].probs) [0.1, 0.2, 0.3, 0.4] """ def __init__(self, *args, logger_warn="warn", autoscale=False): super(ProbabilityDict, self).__init__( *args, value_class=DiscretePDF, logger_warn=logger_warn, make_dict_kwargs={"autoscale": autoscale}, latex_labels=latex_labels ) @property def plotting_map(self): existing = super(ProbabilityDict, self).plotting_map modified = existing.copy() modified.update( { "marginalized_posterior": self._marginalized_posterior, "hist": self._marginalized_posterior, "pdf": self._analytic_pdf, } ) return modified def rvs( self, size=None, parameters=None, interpolate=True, interp_kwargs={} ): """Randomly draw samples from each distribution in ProbabilityDict Parameters ---------- size: int number of samples to draw parameters: list, optional list of parameters you wish to draw samples for. If not provided draw samples for all parameters in ProbabilityDict. Default None interpolate: Bool, optional if True, interpolate the discrete probabilities and then draw samples from the continous distribution. Default True interp_kwargs: dict, optional kwargs to pass to the DiscretePDF.interpolate() method """ if size is None: raise ValueError( "Please specify the number of samples you wish to draw from " "the interpolated distributions" ) if parameters is not None: if any(param not in self.parameters for param in parameters): raise ValueError( "Unable to draw samples because not all parameters have " "probabilities stored in ProbabilityDict. The list of " "available parameters are: {}".format( ", ".join(self.parameters) ) ) else: parameters = self.parameters mydict = {} for param in parameters: if interpolate: mydict[param] = self[param].interpolate(**interp_kwargs).rvs( size=size ) else: mydict[param] = self[param].rvs(size=size) return SamplesDict(mydict) def _marginalized_posterior( self, parameter, size=10**5, module="core", **kwargs ): """Draw samples from the probability distribution and histogram them via the `pesummary.core.plots.plot._1d_histogram_plot` or `pesummary.gw.plots.plot._1d_histogram_plot` Parameters ---------- parameter: str name of the parameter you wish to plot size: int, optional number of samples to from the probability distribution to histogram module: str, optional module you wish to use for the plotting **kwargs: dict all additional kwargs are passed to the `_1d_histogram_plot` function """ samples = SamplesDict({parameter: self[parameter].rvs(size=size)}) return samples.plot(parameter, type="hist", module=module, **kwargs) def _analytic_pdf(self, parameter, **kwargs): """Wrapper for the `pesummary.core.plots.plot._1d_analytic_plot function Parameters ---------- parameter: str name of the parameter you wish to plot **kwargs: dict, optional all additional kwargs are passed ot the `_1d_analytic_plot` function """ from pesummary.core.plots.plot import _1d_analytic_plot return _1d_analytic_plot( parameter, self[parameter].x, self[parameter].probs, self.latex_labels[parameter], **kwargs )
[docs]class ProbabilityDict2D(Dict): """Class to store the 2D discrete probability distributions for multiple parameters Parameters ---------- args: dict dictionary containing the discrete probabilities for 2 or more parameters. Key should be 2 parameter names joined with `_` and value should be an array of length 3. First element integers for x, second element integers for y and third element the 2D discrete probability distribution for xy. See pesummary.utils.pdf.DiscretePDF2D for more details Methods ------- plot: generate a plot to display the stored probability density functions Examples -------- >>> from pesummary.utils.probability_dict.ProbabilityDict2D >>> x = [-1., 0., 1.] >>> y = [-1., 0., 1.] >>> prob_xy = [ ... [1./9, 1./9, 1./9], ... [1./9, 1./9, 1./9], ... [1./9, 1./9, 1./9] ... ] >>> pdf = ProbabilityDict2D({"x_y": [x, y, prob_xy]}) >>> print(type(pdf["x_y"])) <class 'pesummary.utils.pdf.DiscretePDF2D'> """ def __init__(self, *args, logger_warn="warn", autoscale=False): super(ProbabilityDict2D, self).__init__( *args, value_class=DiscretePDF2D, logger_warn=logger_warn, make_dict_kwargs={"autoscale": autoscale}, latex_labels=latex_labels ) @property def plotting_map(self): existing = super(ProbabilityDict2D, self).plotting_map modified = existing.copy() modified.update( { "2d_kde": self._2d_kde, "triangle": self._triangle, "reverse_triangle": self._reverse_triangle } ) return modified def _2d_kde(self, parameter, **kwargs): """Generate a 2d contour plot showing the probability distribution for xy. Parameters ---------- parameter: str the key you wish to plot """ from pesummary.core.plots.publication import analytic_twod_contour_plot if "levels" in kwargs.keys(): levels = self[parameter].probs_xy.level_from_confidence( kwargs["levels"] ) kwargs["levels"] = levels return analytic_twod_contour_plot( self[parameter].x, self[parameter].y, self[parameter].probs, **kwargs ) def _base_triangle(self, parameter, function, **kwargs): pdf = self[parameter].marginalize() if "levels" in kwargs.keys(): levels = pdf.probs_xy.level_from_confidence( kwargs["levels"] ) kwargs["levels"] = levels return function( pdf.x, pdf.y, pdf.probs_x.probs, pdf.probs_y.probs, pdf.probs_xy.probs, **kwargs ) def _triangle(self, parameter, **kwargs): """Generate a triangle plot showing the probability distribution for x, y and xy. The probability distributions for x and y are found through marginalization Parameters ---------- parameter: str the key you wish to plot """ from pesummary.core.plots.publication import analytic_triangle_plot return self._base_triangle( parameter, analytic_triangle_plot, **kwargs ) def _reverse_triangle(self, parameter, **kwargs): """Generate a reverse triangle plot showing the probability distribution for x, y and xy. The probability distributions for x and y are found through marginalization Parameters ---------- parameter: str the key you wish to plot """ from pesummary.core.plots.publication import ( analytic_reverse_triangle_plot ) return self._base_triangle( parameter, analytic_reverse_triangle_plot, **kwargs )