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 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 forward(self, U, pi, mu, Gamma, **kwargs):
j = jnp.argmax(U[0] * jnp.sum(pi) <= jnp.cumsum(pi))
Gamma = Gamma[j, ...]
mu = mu[j, ...]
if self._ill_cond:
L = msqrt(Gamma)
else:
L = jnp.linalg.cholesky(Gamma)
return L @ ndtri(U[1:]) + mu
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 icdf(self, q):
return self.loc + self.scale * ndtri(q)
def _ppf(self, q):
return ndtri(q)
def _isf(self, q):
return -ndtri(q)
def forward(self, U, mu, Gamma, **kwargs):
if self._ill_cond:
L = msqrt(Gamma)
else:
L = jnp.linalg.cholesky(Gamma)
return L @ ndtri(U) + mu
def ppf(q, loc=0, scale=1):
return jnp.array(special.ndtri(q) * scale + loc, 'float64')
def forward(self, U, mu, gamma, **kwargs):
return ndtri(U) * gamma + mu
def ndtri_(u):
return ndtri(u)
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 forward(self, U, x0, omega, **kwargs):
return x0 + omega * jnp.cumsum(ndtri(U).reshape((self.T, -1)), axis=0)
def forward(self, U, pi, mu, gamma, **kwargs):
j = jnp.argmax(U[0] <= jnp.cumsum(pi) / jnp.sum(pi))
gamma = gamma[j, ...]
mu = mu[j, ...]
return gamma * ndtri(U[1:]) + mu
