Coverage for pesummary/utils/bounded_2d_kde.py: 97.2%

1import numpy as np 

2from scipy.stats import gaussian_kde as kde 

3 

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

5 

6 

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 

23 

24 @property 

25 def xlow(self): 

26 """The lower bound of the x domain 

27 """ 

28 return self._xlow 

29 

30 @property 

31 def xhigh(self): 

32 """The upper bound of the x domain 

33 """ 

34 return self._xhigh 

35 

36 @property 

37 def ylow(self): 

38 """The lower bound of the y domain 

39 """ 

40 return self._ylow 

41 

42 @property 

43 def yhigh(self): 

44 """The upper bound of the y domain 

45 """ 

46 return self._yhigh 

47 

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 

55 

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]) 

60 

61 if self.xhigh is not None: 

62 pdf += super(Bounded_2d_kde, self).evaluate([2 * self.xhigh - x, y]) 

63 

64 if self.ylow is not None: 

65 pdf += super(Bounded_2d_kde, self).evaluate([x, 2 * self.ylow - y]) 

66 

67 if self.yhigh is not None: 

68 pdf += super(Bounded_2d_kde, self).evaluate([x, 2 * self.yhigh - y]) 

69 

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]) 

74 

75 if self.yhigh is not None: 

76 pdf += super(Bounded_2d_kde, self).evaluate( 

77 [2 * self.xlow - x, 2 * self.yhigh - y]) 

78 

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 

87 

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