Coverage for pesummary/utils/bounded_2d_kde.py: 97.2%
72 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
1import numpy as np
2from scipy.stats import gaussian_kde as kde
4__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
7class Bounded_2d_kde(kde):
8 """Class to generate a two-dimensional KDE for a probability distribution
9 functon that exists on a bounded domain
10 """
11 def __init__(self, pts, xlow=None, xhigh=None, ylow=None, yhigh=None,
12 transform=None, *args, **kwargs):
13 pts = np.atleast_2d(pts)
14 assert pts.ndim == 2, 'Bounded_kde can only be two-dimensional'
15 self._transform = transform
16 if transform is not None:
17 pts = transform(pts)
18 super(Bounded_2d_kde, self).__init__(pts, *args, **kwargs)
19 self._xlow = xlow
20 self._xhigh = xhigh
21 self._ylow = ylow
22 self._yhigh = yhigh
24 @property
25 def xlow(self):
26 """The lower bound of the x domain
27 """
28 return self._xlow
30 @property
31 def xhigh(self):
32 """The upper bound of the x domain
33 """
34 return self._xhigh
36 @property
37 def ylow(self):
38 """The lower bound of the y domain
39 """
40 return self._ylow
42 @property
43 def yhigh(self):
44 """The upper bound of the y domain
45 """
46 return self._yhigh
48 def evaluate(self, pts):
49 """Return an estimate of the density evaluated at the given
50 points."""
51 _pts = np.atleast_2d(pts)
52 assert _pts.ndim == 2, 'points must be two-dimensional'
53 if _pts.shape[0] != 2 and _pts.shape[1] == 2:
54 _pts = _pts.T
56 x, y = _pts
57 pdf = super(Bounded_2d_kde, self).evaluate(_pts)
58 if self.xlow is not None:
59 pdf += super(Bounded_2d_kde, self).evaluate([2 * self.xlow - x, y])
61 if self.xhigh is not None:
62 pdf += super(Bounded_2d_kde, self).evaluate([2 * self.xhigh - x, y])
64 if self.ylow is not None:
65 pdf += super(Bounded_2d_kde, self).evaluate([x, 2 * self.ylow - y])
67 if self.yhigh is not None:
68 pdf += super(Bounded_2d_kde, self).evaluate([x, 2 * self.yhigh - y])
70 if self.xlow is not None:
71 if self.ylow is not None:
72 pdf += super(Bounded_2d_kde, self).evaluate(
73 [2 * self.xlow - x, 2 * self.ylow - y])
75 if self.yhigh is not None:
76 pdf += super(Bounded_2d_kde, self).evaluate(
77 [2 * self.xlow - x, 2 * self.yhigh - y])
79 if self.xhigh is not None:
80 if self.ylow is not None:
81 pdf += super(Bounded_2d_kde, self).evaluate(
82 [2 * self.xhigh - x, 2 * self.ylow - y])
83 if self.yhigh is not None:
84 pdf += super(Bounded_2d_kde, self).evaluate(
85 [2 * self.xhigh - x, 2 * self.yhigh - y])
86 return pdf
88 def __call__(self, pts):
89 _pts = np.atleast_2d(pts)
90 if _pts.shape[0] != 2 and _pts.shape[1] == 2:
91 _pts = _pts.T
92 if self._transform is not None:
93 _pts = self._transform(_pts)
94 transpose = _pts.T
95 out_of_bounds = np.zeros(transpose.shape[0], dtype='bool')
96 if self.xlow is not None:
97 out_of_bounds[transpose[:, 0] < self.xlow] = True
98 if self.xhigh is not None:
99 out_of_bounds[transpose[:, 0] > self.xhigh] = True
100 if self.ylow is not None:
101 out_of_bounds[transpose[:, 1] < self.ylow] = True
102 if self.yhigh is not None:
103 out_of_bounds[transpose[:, 1] > self.yhigh] = True
104 results = self.evaluate(_pts)
105 results[out_of_bounds] = 0.
106 return results