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