Coverage for pesummary/utils/probability_dict.py: 93.5%
62 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-05-02 08:42 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3from pesummary.utils.samples_dict import SamplesDict
4from pesummary.utils.dict import Dict
5from pesummary.utils.pdf import DiscretePDF, DiscretePDF2D
6from pesummary.core.plots.latex_labels import latex_labels
7from pesummary.gw.plots.latex_labels import GWlatex_labels
9__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
11latex_labels.update(GWlatex_labels)
14class ProbabilityDict(Dict):
15 """Class to store integers with discrete probabilities for multiple
16 parameters.
18 Parameters
19 ----------
20 args: dict
21 dictionary containing the discrete probabilities for each parameter.
22 Key should be the parameter name and value should be a 2d array of
23 length 2. First element integers and second element probabilities
24 corresponding to integers. See pesummary.utils.pdf.DiscretePDF for more
25 details
27 Methods
28 -------
29 rvs:
30 randomly draw samples for each distribution in ProbabilityDict
32 Examples
33 --------
34 >>> from pesummary.utils.probability_dict import ProbabilityDict
35 >>> numbers = [1,2,3,4]
36 >>> probabilities = [0.1, 0.2, 0.3, 0.4]
37 >>> pdf = ProbabilityDict({"parameter": [numbers, probabilities]})
38 >>> print(type(pdf["parameter"]))
39 <class 'pesummary.utils.pdf.DiscretePDF'>
40 >>> print(pdf["parameter"].probs)
41 [0.1, 0.2, 0.3, 0.4]
42 """
43 def __init__(self, *args, logger_warn="warn", autoscale=False):
44 super(ProbabilityDict, self).__init__(
45 *args, value_class=DiscretePDF, logger_warn=logger_warn,
46 make_dict_kwargs={"autoscale": autoscale}, latex_labels=latex_labels
47 )
49 @property
50 def plotting_map(self):
51 existing = super(ProbabilityDict, self).plotting_map
52 modified = existing.copy()
53 modified.update(
54 {
55 "marginalized_posterior": self._marginalized_posterior,
56 "hist": self._marginalized_posterior,
57 "pdf": self._analytic_pdf,
58 }
59 )
60 return modified
62 def rvs(
63 self, size=None, parameters=None, interpolate=True, interp_kwargs={}
64 ):
65 """Randomly draw samples from each distribution in ProbabilityDict
67 Parameters
68 ----------
69 size: int
70 number of samples to draw
71 parameters: list, optional
72 list of parameters you wish to draw samples for. If not provided
73 draw samples for all parameters in ProbabilityDict. Default None
74 interpolate: Bool, optional
75 if True, interpolate the discrete probabilities and then draw
76 samples from the continous distribution. Default True
77 interp_kwargs: dict, optional
78 kwargs to pass to the DiscretePDF.interpolate() method
79 """
80 if size is None:
81 raise ValueError(
82 "Please specify the number of samples you wish to draw from "
83 "the interpolated distributions"
84 )
85 if parameters is not None:
86 if any(param not in self.parameters for param in parameters):
87 raise ValueError(
88 "Unable to draw samples because not all parameters have "
89 "probabilities stored in ProbabilityDict. The list of "
90 "available parameters are: {}".format(
91 ", ".join(self.parameters)
92 )
93 )
94 else:
95 parameters = self.parameters
97 mydict = {}
98 for param in parameters:
99 if interpolate:
100 mydict[param] = self[param].interpolate(**interp_kwargs).rvs(
101 size=size
102 )
103 else:
104 mydict[param] = self[param].rvs(size=size)
105 return SamplesDict(mydict)
107 def _marginalized_posterior(
108 self, parameter, size=10**5, module="core", **kwargs
109 ):
110 """Draw samples from the probability distribution and histogram them
111 via the `pesummary.core.plots.plot._1d_histogram_plot` or
112 `pesummary.gw.plots.plot._1d_histogram_plot`
114 Parameters
115 ----------
116 parameter: str
117 name of the parameter you wish to plot
118 size: int, optional
119 number of samples to from the probability distribution to histogram
120 module: str, optional
121 module you wish to use for the plotting
122 **kwargs: dict
123 all additional kwargs are passed to the `_1d_histogram_plot`
124 function
125 """
126 samples = SamplesDict({parameter: self[parameter].rvs(size=size)})
127 return samples.plot(parameter, type="hist", module=module, **kwargs)
129 def _analytic_pdf(self, parameter, **kwargs):
130 """Wrapper for the `pesummary.core.plots.plot._1d_analytic_plot
131 function
133 Parameters
134 ----------
135 parameter: str
136 name of the parameter you wish to plot
137 **kwargs: dict, optional
138 all additional kwargs are passed ot the `_1d_analytic_plot` function
139 """
140 from pesummary.core.plots.plot import _1d_analytic_plot
142 return _1d_analytic_plot(
143 parameter, self[parameter].x, self[parameter].probs,
144 self.latex_labels[parameter], **kwargs
145 )
148class ProbabilityDict2D(Dict):
149 """Class to store the 2D discrete probability distributions for multiple
150 parameters
152 Parameters
153 ----------
154 args: dict
155 dictionary containing the discrete probabilities for 2 or more
156 parameters. Key should be 2 parameter names joined with `_` and
157 value should be an array of length 3. First element integers for
158 x, second element integers for y and third element the 2D discrete
159 probability distribution for xy. See
160 pesummary.utils.pdf.DiscretePDF2D for more details
162 Methods
163 -------
164 plot:
165 generate a plot to display the stored probability density functions
167 Examples
168 --------
169 >>> from pesummary.utils.probability_dict.ProbabilityDict2D
170 >>> x = [-1., 0., 1.]
171 >>> y = [-1., 0., 1.]
172 >>> prob_xy = [
173 ... [1./9, 1./9, 1./9],
174 ... [1./9, 1./9, 1./9],
175 ... [1./9, 1./9, 1./9]
176 ... ]
177 >>> pdf = ProbabilityDict2D({"x_y": [x, y, prob_xy]})
178 >>> print(type(pdf["x_y"]))
179 <class 'pesummary.utils.pdf.DiscretePDF2D'>
180 """
181 def __init__(self, *args, logger_warn="warn", autoscale=False):
182 super(ProbabilityDict2D, self).__init__(
183 *args, value_class=DiscretePDF2D, logger_warn=logger_warn,
184 make_dict_kwargs={"autoscale": autoscale}, latex_labels=latex_labels
185 )
187 @property
188 def plotting_map(self):
189 existing = super(ProbabilityDict2D, self).plotting_map
190 modified = existing.copy()
191 modified.update(
192 {
193 "2d_kde": self._2d_kde,
194 "triangle": self._triangle,
195 "reverse_triangle": self._reverse_triangle
196 }
197 )
198 return modified
200 def _2d_kde(self, parameter, **kwargs):
201 """Generate a 2d contour plot showing the probability distribution for
202 xy.
204 Parameters
205 ----------
206 parameter: str
207 the key you wish to plot
208 """
209 from pesummary.core.plots.publication import analytic_twod_contour_plot
211 if "levels" in kwargs.keys():
212 levels = self[parameter].probs_xy.level_from_confidence(
213 kwargs["levels"]
214 )
215 kwargs["levels"] = levels
216 return analytic_twod_contour_plot(
217 self[parameter].x, self[parameter].y, self[parameter].probs,
218 **kwargs
219 )
221 def _base_triangle(self, parameter, function, **kwargs):
222 pdf = self[parameter].marginalize()
223 if "levels" in kwargs.keys():
224 levels = pdf.probs_xy.level_from_confidence(
225 kwargs["levels"]
226 )
227 kwargs["levels"] = levels
228 return function(
229 pdf.x, pdf.y, pdf.probs_x.probs, pdf.probs_y.probs,
230 pdf.probs_xy.probs, **kwargs
231 )
233 def _triangle(self, parameter, **kwargs):
234 """Generate a triangle plot showing the probability distribution for
235 x, y and xy. The probability distributions for x and y are found
236 through marginalization
238 Parameters
239 ----------
240 parameter: str
241 the key you wish to plot
242 """
243 from pesummary.core.plots.publication import analytic_triangle_plot
245 return self._base_triangle(
246 parameter, analytic_triangle_plot, **kwargs
247 )
249 def _reverse_triangle(self, parameter, **kwargs):
250 """Generate a reverse triangle plot showing the probability distribution
251 for x, y and xy. The probability distributions for x and y are found
252 through marginalization
254 Parameters
255 ----------
256 parameter: str
257 the key you wish to plot
258 """
259 from pesummary.core.plots.publication import (
260 analytic_reverse_triangle_plot
261 )
263 return self._base_triangle(
264 parameter, analytic_reverse_triangle_plot, **kwargs
265 )