def standard_normal_to_students_t(location, scale, dof):
if dof == 1:
return standard_normal_to_cauchy(location, scale)
elif dof == 2:
return ElementwiseMonotonicTransform(
forward=lambda n: location + standard_students_t_2_icdf(ndtr(n)) *
scale,
backward=lambda t: ndtri(
(standard_students_t_2_cdf(t - location) / scale)),
domain=reals,
image=reals,
)
elif dof == 4:
return ElementwiseMonotonicTransform(
forward=lambda n: location + standard_students_t_4_icdf(ndtr(n)) *
scale,
backward=lambda t: ndtri(
(standard_students_t_4_cdf(t - location) / scale)),
domain=reals,
image=reals,
)
else:
raise ValueError("Transform only defined for dof in {1, 2, 4}")
def standard_normal_to_half_normal(scale):
return ElementwiseMonotonicTransform(
forward=lambda n: ndtri((ndtr(n) + 1) / 2) * scale,
backward=lambda h: ndtri(2 * ndtr(h / scale) - 1),
domain=reals,
image=nonnegative_reals,
)
def standard_normal_to_truncated_normal(location, scale, lower, upper):
a = ndtr((lower - location) / scale)
b = ndtr((upper - location) / scale)
return ElementwiseMonotonicTransform(
forward=lambda n: ndtri(a + ndtr(n) * (b - a)) * scale + location,
backward=lambda t: ndtri((ndtr((t - location) / scale) - a) / (b - a)),
domain=reals,
image=RealInterval(lower, upper),
)
def integrate_box_1d(self, low, high):
if self.d != 1:
raise ValueError("integrate_box_1d() only handles 1D pdfs")
if jnp.ndim(low) != 0 or jnp.ndim(high) != 0:
raise ValueError(
"the limits of integration in integrate_box_1d must be scalars"
)
sigma = jnp.squeeze(jnp.sqrt(self.covariance))
low = jnp.squeeze((low - self.dataset) / sigma)
high = jnp.squeeze((high - self.dataset) / sigma)
return jnp.sum(self.weights * (special.ndtr(high) - special.ndtr(low)))
def standard_normal_to_beta(shape_a, shape_b):
if shape_b == 1:
def icdf(u):
return u**(1 / shape_a)
def cdf(x):
return x**shape_a
elif shape_a == 1:
def icdf(u):
return 1 - (1 - u)**(1 / shape_b)
def cdf(x):
return 1 - (1 - x)**shape_b
else:
raise ValueError(
"Transform only defined for shape_a == 1 or shape_b == 1")
return ElementwiseMonotonicTransform(
forward=lambda n: icdf(ndtr(n)),
backward=lambda x: ndtri(cdf(x)),
domain=reals,
image=RealInterval(0, 1),
)
def standard_normal_to_cauchy(location, scale):
return ElementwiseMonotonicTransform(
forward=lambda n: location + standard_cauchy_icdf(ndtr(n)) * scale,
backward=lambda c: ndtri(standard_cauchy_cdf((c - location) / scale)),
domain=reals,
image=reals,
)
def standard_normal_to_exponential(rate):
return ElementwiseMonotonicTransform(
forward=lambda n: -np.log(ndtr(n)) / rate,
backward=lambda e: ndtri(np.exp(-e * rate)),
domain=reals,
image=nonnegative_reals,
)
def standard_normal_to_uniform(lower, upper):
return ElementwiseMonotonicTransform(
forward=lambda n: lower + ndtr(n) * (upper - lower),
backward=lambda u: ndtri(u - lower / (upper - lower)),
domain=reals,
image=RealInterval(lower, upper),
)
def sample(self, key, sample_shape=()):
size = sample_shape + self.batch_shape
# We use inverse transform method:
# z ~ icdf(U), where U ~ Uniform(0, 1).
u = random.uniform(key, shape=size)
# Ref: https://en.wikipedia.org/wiki/Truncated_normal_distribution#Simulating
# icdf[cdf_a + u * (1 - cdf_a)] = icdf[1 - (1 - cdf_a)(1 - u)]
# = - icdf[(1 - cdf_a)(1 - u)]
return self.base_loc - ndtri(ndtr(self.base_loc) * (1 - u))
def _sf(self, x):
return ndtr(-x)
def _cdf(self, x):
return ndtr(x)
def cdf(x, loc=0, scale=1):
x, loc, scale = _promote_args_inexact("norm.cdf", x, loc, scale)
return special.ndtr(lax.div(lax.sub(x, loc), scale))
def sample(rng, shape=()):
a = ndtr((lower - location) / scale)
b = ndtr((upper - location) / scale)
return ndtri(a + rng.uniform(size=shape) * (b - a)) * scale + location
def variance(self):
low = (self.low - self.loc) / self.scale
low_prob_scaled = np.exp(self._normal.log_prob(
self.low)) * self.scale / ndtr(-low)
return self._normal.variance * (1 + low * low_prob_scaled -
low_prob_scaled**2)
def mean(self):
low = (self.low - self.loc) / self.scale
low_prob_scaled = np.exp(self._normal.log_prob(
self.low)) * self.scale / ndtr(-low)
return self.loc + low_prob_scaled * self.scale
def actout(cut, μ, σ):
avg = np.exp(μ + 0.5 * σ**2)
lcut = log(cut)
act = 1 - ndtr((lcut - μ) / σ)
out = avg * ndtr((μ + σ**2 - lcut) / σ)
return act, out
