-
Notifications
You must be signed in to change notification settings - Fork 12
/
TruncatedNormal.py
166 lines (137 loc) · 6.08 KB
/
TruncatedNormal.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
import math
from numbers import Number
import torch
from torch.distributions import Distribution, constraints
from torch.distributions.utils import broadcast_all
CONST_SQRT_2 = math.sqrt(2)
CONST_INV_SQRT_2PI = 1 / math.sqrt(2 * math.pi)
CONST_INV_SQRT_2 = 1 / math.sqrt(2)
CONST_LOG_INV_SQRT_2PI = math.log(CONST_INV_SQRT_2PI)
CONST_LOG_SQRT_2PI_E = 0.5 * math.log(2 * math.pi * math.e)
class TruncatedStandardNormal(Distribution):
"""
Truncated Standard Normal distribution
https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
"""
arg_constraints = {
'a': constraints.real,
'b': constraints.real,
}
support = constraints.real
has_rsample = True
def __init__(self, a, b, eps=1e-8, validate_args=None):
self.a, self.b = broadcast_all(a, b)
if isinstance(a, Number) and isinstance(b, Number):
batch_shape = torch.Size()
else:
batch_shape = self.a.size()
super(TruncatedStandardNormal, self).__init__(batch_shape, validate_args=validate_args)
if self.a.dtype != self.b.dtype:
raise ValueError('Truncation bounds types are different')
if any((self.a >= self.b).view(-1,).tolist()):
raise ValueError('Incorrect truncation range')
self._dtype_min_gt_0 = torch.tensor(torch.finfo(self.a.dtype).eps, dtype=self.a.dtype)
self._dtype_max_lt_1 = torch.tensor(1 - torch.finfo(self.a.dtype).eps, dtype=self.a.dtype)
self._little_phi_a = self._little_phi(self.a)
self._little_phi_b = self._little_phi(self.b)
self._big_phi_a = self._big_phi(self.a)
self._big_phi_b = self._big_phi(self.b)
self._Z = (self._big_phi_b - self._big_phi_a).clamp_min(eps)
self._log_Z = self._Z.log()
self._lpbb_m_lpaa_d_Z = (self._little_phi_b * self.b - self._little_phi_a * self.a) / self._Z
self._mean = -(self._little_phi_b - self._little_phi_a) / self._Z
self._variance = 1 - self._lpbb_m_lpaa_d_Z - ((self._little_phi_b - self._little_phi_a) / self._Z) ** 2
self._entropy = CONST_LOG_SQRT_2PI_E + self._log_Z - 0.5 * self._lpbb_m_lpaa_d_Z
@property
def mean(self):
return self._mean
@property
def variance(self):
return self._variance
@property
def entropy(self):
return self._entropy
@property
def auc(self):
return self._Z
@staticmethod
def _little_phi(x):
return (-(x ** 2) * 0.5).exp() * CONST_INV_SQRT_2PI
@staticmethod
def _big_phi(x):
return 0.5 * (1 + (x * CONST_INV_SQRT_2).erf())
@staticmethod
def _inv_big_phi(x):
return CONST_SQRT_2 * (2 * x - 1).erfinv()
def cdf(self, value):
if self._validate_args:
self._validate_sample(value)
return ((self._big_phi(value) - self._big_phi_a) / self._Z).clamp(0, 1)
def icdf(self, value):
if self._validate_args:
self._validate_sample(value)
return self._inv_big_phi(self._big_phi_a + value * self._Z)
def log_prob(self, value):
if self._validate_args:
self._validate_sample(value)
return CONST_LOG_INV_SQRT_2PI - self._log_Z - (value ** 2) * 0.5
def rsample(self, sample_shape=torch.Size()):
shape = self._extended_shape(sample_shape)
p = torch.empty(shape).uniform_(self._dtype_min_gt_0, self._dtype_max_lt_1)
return self.icdf(p)
def expand(self, batch_shape, _instance=None):
# TODO: it is likely that keeping temporary variables in private attributes violates the logic of this method
raise NotImplementedError
class TruncatedNormal(TruncatedStandardNormal):
"""
Truncated Normal distribution
https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
"""
arg_constraints = {
'loc': constraints.real,
'scale': constraints.positive,
'a': constraints.real,
'b': constraints.real,
}
support = constraints.real
has_rsample = True
def __init__(self, loc, scale, a, b, eps=1e-8, validate_args=None):
self.loc, self.scale, self.a, self.b = broadcast_all(loc, scale, a, b)
a_standard = (a - self.loc) / self.scale
b_standard = (b - self.loc) / self.scale
super(TruncatedNormal, self).__init__(a_standard, b_standard, eps=eps, validate_args=validate_args)
self._log_scale = self.scale.log()
self._mean = self._mean * self.scale + self.loc
self._variance = self._variance * self.scale ** 2
self._entropy += self._log_scale
def _to_std_rv(self, value):
if self._validate_args:
self._validate_sample(value)
return (value - self.loc) / self.scale
def _from_std_rv(self, value):
if self._validate_args:
self._validate_sample(value)
return value * self.scale + self.loc
def cdf(self, value):
return super(TruncatedNormal, self).cdf(self._to_std_rv(value))
def icdf(self, value):
if self._validate_args:
self._validate_sample(value)
return self._from_std_rv(super(TruncatedNormal, self).icdf(value))
def log_prob(self, value):
if self._validate_args:
self._validate_sample(value)
return super(TruncatedNormal, self).log_prob(self._to_std_rv(value)) - self._log_scale
if __name__ == '__main__':
from scipy.stats import truncnorm
loc, scale, a, b = 1., 2., 1., 2.
tn_pt = TruncatedNormal(loc, scale, a, b)
mean_pt, var_pt = tn_pt.mean.item(), tn_pt.variance.item()
alpha, beta = (a - loc) / scale, (b - loc) / scale
mean_sp, var_sp = truncnorm.stats(alpha, beta, loc=loc, scale=scale, moments='mv')
print('mean', mean_pt, mean_sp)
print('var', var_pt, var_sp)
print('cdf', tn_pt.cdf(1.4).item(), truncnorm.cdf(1.4, alpha, beta, loc=loc, scale=scale))
print('icdf', tn_pt.icdf(0.333).item(), truncnorm.ppf(0.333, alpha, beta, loc=loc, scale=scale))
print('logpdf', tn_pt.log_prob(1.5).item(), truncnorm.logpdf(1.5, alpha, beta, loc=loc, scale=scale))
print('entropy', tn_pt.entropy.item(), truncnorm.entropy(alpha, beta, loc=loc, scale=scale))