Source code for ultranest.simbase.minisbi.logistic

"""Kumaraswamy-Logistic chained distribution defined on a unit hypercube."""

import numpy as np
import torch

# ---------------------------------------------------------------------------
# Kumaraswamy-Logistic chained distribution helpers
#
# Construction:
#   1. Draw v ~ TruncatedLogistic(loc, scale) on (0,1)   [the "base" draw]
#   2. Apply the Kumaraswamy CDF  u = 1 - (1 - v^a)^b   to map v -> u in (0,1)
#
# The joint density factors as:
#   p(u) = p_L(v(u)) * |dv/du|
# where v(u) is the inverse Kumaraswamy CDF (= Kumaraswamy quantile function):
#   v = (1 - (1 - u)^{1/b})^{1/a}
# and the log |dv/du| term comes from the Kumaraswamy log-density evaluated
# at v with the extra -log(a) -log(b) factored in correctly.
# ---------------------------------------------------------------------------


# ---------------------------------------------------------------------------
# NumPy utility functions
# ---------------------------------------------------------------------------

[docs] def sigmoid(x): """ Numerically stable sigmoid function. Parameters ---------- x : numpy.ndarray or float Returns ------- numpy.ndarray or float Sigmoid of x, clipped to avoid overflow. """ return 1.0 / (1.0 + np.exp(-np.clip(x, -500.0, 500.0)))
# Internal alias used within numpy helpers _sigmoid = sigmoid
[docs] def logit(p): """ Logit (inverse sigmoid) function. Parameters ---------- p : numpy.ndarray or float Probability value(s) in (0, 1). Returns ------- numpy.ndarray or float Log-odds of p. """ p = np.clip(p, 1e-15, 1.0 - 1e-15) return np.log(p / (1.0 - p))
# Internal alias used within numpy helpers _logit = logit # --------------------------------------------------------------------------- # Kumaraswamy distribution helpers # ---------------------------------------------------------------------------
[docs] def kuma_log_prob(u, a, b): """ Log-density of Kumaraswamy(a, b) at u in (0, 1). log p(u) = log(a) + log(b) + (a-1)*log(u) + (b-1)*log(1 - u^a) Parameters ---------- u : torch.Tensor Values in (0, 1). a : torch.Tensor Shape parameter, > 0. b : torch.Tensor Shape parameter, > 0. Returns ------- torch.Tensor Log-density at u. """ u = u.clamp(1e-7, 1.0 - 1e-7) ua = u.pow(a) log_p = (torch.log(a) + torch.log(b) + (a - 1.0) * torch.log(u) + (b - 1.0) * torch.log((1.0 - ua).clamp(min=1e-30))) return log_p
[docs] def kuma_icdf(u, a, b): """ Inverse CDF (quantile function) of Kumaraswamy(a, b). Maps u in (0, 1) to v in (0, 1) via: v = (1 - (1 - u)^{1/b})^{1/a} Parameters ---------- u : torch.Tensor Values in (0, 1). a : torch.Tensor Shape parameter, > 0. b : torch.Tensor Shape parameter, > 0. Returns ------- torch.Tensor Quantile values in (0, 1). """ u = u.clamp(1e-7, 1.0 - 1e-7) return (1.0 - (1.0 - u).pow(1.0 / b)).pow(1.0 / a)
[docs] def kuma_icdf_np(u, a, b): """ Inverse CDF (quantile function) of Kumaraswamy(a, b) — NumPy version. Maps u in (0, 1) to v in (0, 1) via: v = (1 - (1 - u)^{1/b})^{1/a} Parameters ---------- u : numpy.ndarray Values in (0, 1). a : numpy.ndarray Shape parameter, > 0. b : numpy.ndarray Shape parameter, > 0. Returns ------- numpy.ndarray Quantile values in (0, 1). """ u = np.clip(u, 1e-7, 1.0 - 1e-7) return (1.0 - (1.0 - u) ** (1.0 / b)) ** (1.0 / a)
# --------------------------------------------------------------------------- # Logistic distribution helpers (PyTorch) # --------------------------------------------------------------------------- def _logistic_log_prob(u, loc, scale): """ Log-probability of u under Logistic(loc, scale). Parameters ---------- u : torch.Tensor Shape (batch, n_params). loc : torch.Tensor Shape (batch, n_params). scale : torch.Tensor Shape (batch, n_params). Returns ------- torch.Tensor Shape (batch, n_params), per-dimension log-probability. """ z = (u - loc) / scale log_p = -z - torch.log(scale) - 2.0 * torch.nn.functional.softplus(-z) return log_p
[docs] def logistic_log_cdf(x, loc, scale): """ Log CDF of Logistic(loc, scale). Computes log sigmoid((x - loc) / scale). Parameters ---------- x : torch.Tensor loc : torch.Tensor scale : torch.Tensor Returns ------- torch.Tensor Log CDF evaluated at x. """ return torch.nn.functional.logsigmoid((x - loc) / scale)
def _logistic_log_sf(x, loc, scale): """ Log survival function of Logistic(loc, scale). Computes log(1 - sigmoid((x - loc) / scale)) = log sigmoid(-(x - loc) / scale). Parameters ---------- x : torch.Tensor loc : torch.Tensor scale : torch.Tensor Returns ------- torch.Tensor Log survival function evaluated at x. """ return torch.nn.functional.logsigmoid(-(x - loc) / scale) def _logistic_log_normaliser(loc, scale): """ Log probability mass of Logistic(loc, scale) inside (0, 1). Computes log(CDF(1) - CDF(0)) per dimension. Parameters ---------- loc : torch.Tensor Shape (batch, n_params). scale : torch.Tensor Shape (batch, n_params). Returns ------- torch.Tensor Shape (batch, n_params). """ log_cdf1 = logistic_log_cdf(torch.ones_like(loc), loc, scale) log_cdf0 = logistic_log_cdf(torch.zeros_like(loc), loc, scale) log_mass = torch.log( torch.clamp(log_cdf1.exp() - log_cdf0.exp(), min=1e-30) ) return log_mass # --------------------------------------------------------------------------- # Chained (Kumaraswamy o TruncatedLogistic) distribution — PyTorch # ---------------------------------------------------------------------------
[docs] def nll_kuma_logistic_product(u, loc, scale, a, b): """Negative log-likelihood of Kumaraswamy-Logistic unit distribution. The generative model per dimension is: v ~ TruncatedLogistic(loc, scale) on (0, 1) u = Kuma_CDF(v; a, b) = 1 - (1 - v^a)^b The density of u is obtained by the change of variables v = Kuma_ICDF(u): log p_U(u) = log p_L(v; loc, scale) - log Z(loc, scale) + log|dv/du| where log|dv/du| = -log p_Kuma(v; a, b) (the Kumaraswamy log-density evaluated at v gives the magnitude of the Jacobian of the inverse map). Parameters ---------- u : torch.Tensor Shape (batch, n_params), values in (0, 1). loc : torch.Tensor Shape (batch, n_params). scale : torch.Tensor Shape (batch, n_params), > 0. a : torch.Tensor Shape (batch, n_params), > 0. Kumaraswamy shape parameter. b : torch.Tensor Shape (batch, n_params), > 0. Kumaraswamy shape parameter. Returns ------- torch.Tensor Scalar mean NLL. """ # v = Kumaraswamy inverse CDF applied to u v = kuma_icdf(u, a, b) # (batch, n_params), in (0,1) # log p_L(v) per dimension (untruncated logistic) log_pL_v = _logistic_log_prob(v, loc, scale) # (batch, n_params) # log normaliser per dimension log_Z = _logistic_log_normaliser(loc, scale) # (batch, n_params) # log|dv/du| = -log(Kuma density at v w.r.t. Kuma(a, b)) # Kuma density: p_K(v) = a*b*v^{a-1}*(1-v^a)^{b-1} v_c = v.clamp(1e-7, 1.0 - 1e-7) va = v_c.pow(a) log_kuma_at_v = (torch.log(a) + torch.log(b) + (a - 1.0) * torch.log(v_c) + (b - 1.0) * torch.log((1.0 - va).clamp(min=1e-30))) log_abs_dv_du = -log_kuma_at_v # (batch, n_params) # log p_U(u) per dimension log_p_u = log_pL_v - log_Z + log_abs_dv_du # (batch, n_params) # sum over parameter dimensions, mean over batch nll = -log_p_u.sum(dim=-1).mean() return nll
[docs] def sample_kuma_logistic_product(loc, scale, a, b, n_samples, rng): """Sample from unit Kumaraswamy-Logistic distribution. The inverse CDF of U = Kuma_CDF(V) where V ~ TruncLogistic is: 1. Draw uniform w in (0, 1). 2. v = TruncLogistic_ICDF(w; loc, scale) via standard inversion. 3. u = Kuma_CDF(v; a, b) = 1 - (1 - v^a)^b. Parameters ---------- loc : numpy.ndarray Shape (n_params,). scale : numpy.ndarray Shape (n_params,). a : numpy.ndarray Shape (n_params,), Kumaraswamy shape parameter, > 0. b : numpy.ndarray Shape (n_params,), Kumaraswamy shape parameter, > 0. n_samples : int Number of samples to draw. rng : numpy.random.Generator Random number generator. Returns ------- numpy.ndarray Shape (n_samples, n_params), values in (0, 1). """ n_params = loc.shape[0] # Truncated-logistic CDF bounds p0 = sigmoid((0.0 - loc) / scale) # (n_params,) p1 = sigmoid((1.0 - loc) / scale) # (n_params,) # Step 1: uniform draws -> truncated-logistic samples v w = rng.uniform(0.0, 1.0, size=(n_samples, n_params)) p_draw = p0[None, :] + w * (p1 - p0)[None, :] p_draw = np.clip(p_draw, 1e-15, 1.0 - 1e-15) v = loc[None, :] + scale[None, :] * logit(p_draw) # (n_samples, n_params) v = np.clip(v, 1e-7, 1.0 - 1e-7) # Step 2: apply Kumaraswamy CDF u = 1 - (1 - v^a)^b u = 1.0 - (1.0 - v ** a[None, :]) ** b[None, :] u = np.clip(u, 0.0, 1.0) return u.astype(np.float32)
# --------------------------------------------------------------------------- # Chained (Kumaraswamy o TruncatedLogistic) distribution — NumPy # ---------------------------------------------------------------------------
[docs] def kuma_logistic_cdf(x_val, loc, scale, a, b): """ CDF of the Kumaraswamy-Logistic chained distribution at a scalar point. The generative model is: v ~ TruncatedLogistic(loc, scale) on (0, 1) u = 1 - (1 - v^a)^b The CDF of u at x is P(u <= x) = P(v <= Kuma_ICDF(x; a, b)) under the truncated logistic. Parameters ---------- x_val : float Point in [0, 1]. loc : float scale : float Must be > 0. a : float Kumaraswamy shape parameter, > 0. b : float Kumaraswamy shape parameter, > 0. Returns ------- float CDF value in [0, 1]. """ x_val = float(np.clip(x_val, 1e-7, 1.0 - 1e-7)) # v = Kuma_ICDF(x_val; a, b) v = (1.0 - (1.0 - x_val) ** (1.0 / b)) ** (1.0 / a) v = float(np.clip(v, 1e-7, 1.0 - 1e-7)) # truncated-logistic CDF normaliser p0 = sigmoid((0.0 - loc) / scale) p1 = sigmoid((1.0 - loc) / scale) mass = p1 - p0 if mass < 1e-30: mass = 1e-30 # P(V <= v) under the truncated logistic pv = sigmoid((v - loc) / scale) cdf_val = (pv - p0) / mass return float(np.clip(cdf_val, 0.0, 1.0))
[docs] def kuma_logistic_cdf_vec(u_vec, loc, scale, a, b): """Cumulative distribution function of the Kumaraswamy-Logistic distribution. The CDF is evaluated element-wise over a vector of points. Parameters ---------- u_vec : numpy.ndarray Shape (d,), points at which to evaluate the CDF, each in (0, 1). loc : numpy.ndarray Shape (d,). scale : numpy.ndarray Shape (d,), must be positive. a : numpy.ndarray Shape (d,), Kumaraswamy shape parameter, must be positive. b : numpy.ndarray Shape (d,), Kumaraswamy shape parameter, must be positive. Returns ------- numpy.ndarray Shape (d,), CDF values in [0, 1]. """ u_vec = np.clip(u_vec, 1e-7, 1.0 - 1e-7) # Step 1: map u -> v via Kumaraswamy inverse CDF # v = (1 - (1 - u)^{1/b})^{1/a} v = (1.0 - (1.0 - u_vec) ** (1.0 / b)) ** (1.0 / a) v = np.clip(v, 1e-7, 1.0 - 1e-7) # Step 2: evaluate truncated-logistic CDF at v p0 = _sigmoid((0.0 - loc) / scale) p1 = _sigmoid((1.0 - loc) / scale) mass = np.clip(p1 - p0, 1e-30, None) pv = _sigmoid((v - loc) / scale) cdf = np.clip((pv - p0) / mass, 0.0, 1.0) return cdf
[docs] def kuma_logistic_icdf_vec(t_vec, loc, scale, a, b): """Inverse CDF (quantile function) of the Kumaraswamy-Logistic distribution. Parameters ---------- t_vec : numpy.ndarray Shape (d,), quantile levels in (0, 1). loc : numpy.ndarray Shape (d,). scale : numpy.ndarray Shape (d,), must be positive. a : numpy.ndarray Shape (d,), Kumaraswamy shape parameter, must be positive. b : numpy.ndarray Shape (d,), Kumaraswamy shape parameter, must be positive. Returns ------- numpy.ndarray Shape (d,), quantile values in (0, 1). """ t_vec = np.clip(t_vec, 1e-7, 1.0 - 1e-7) # Step 1: invert the truncated logistic # t = (sigmoid((v - loc)/scale) - p0) / (p1 - p0) # => sigmoid((v - loc)/scale) = p0 + t*(p1-p0) # => v = loc + scale * logit(p0 + t*(p1-p0)) p0 = _sigmoid((0.0 - loc) / scale) p1 = _sigmoid((1.0 - loc) / scale) mass = np.clip(p1 - p0, 1e-30, None) p_draw = np.clip(p0 + t_vec * mass, 1e-15, 1.0 - 1e-15) v = loc + scale * _logit(p_draw) v = np.clip(v, 1e-7, 1.0 - 1e-7) # Step 2: apply Kumaraswamy CDF u = 1 - (1 - v^a)^b u = 1.0 - (1.0 - v ** a) ** b u = np.clip(u, 1e-7, 1.0 - 1e-7) return u
[docs] def kuma_logistic_logpdf_vec(u_vec, loc, scale, a, b): """Log-density of the Kumaraswamy-Logistic distribution, summed over dimensions. This equals the log Jacobian log|du/dt| needed to correct the nested-sampling likelihood when the prior is this distribution. Parameters ---------- u_vec : numpy.ndarray Shape (d,), values in (0, 1). loc : numpy.ndarray Shape (d,). scale : numpy.ndarray Shape (d,). a : numpy.ndarray Shape (d,), Kumaraswamy shape parameter. b : numpy.ndarray Shape (d,), Kumaraswamy shape parameter. Returns ------- float Sum of per-dimension log-densities. """ u_vec = np.clip(u_vec, 1e-7, 1.0 - 1e-7) # v = Kuma_ICDF(u) v = (1.0 - (1.0 - u_vec) ** (1.0 / b)) ** (1.0 / a) v = np.clip(v, 1e-7, 1.0 - 1e-7) # log p_L(v) -- untruncated logistic z = (v - loc) / scale log_pL_v = -z - np.log(scale) - 2.0 * np.log1p(np.exp(-np.clip(z, -500, 500))) # log normaliser log(sigmoid((1-loc)/scale) - sigmoid((0-loc)/scale)) p0 = _sigmoid((0.0 - loc) / scale) p1 = _sigmoid((1.0 - loc) / scale) log_Z = np.log(np.clip(p1 - p0, 1e-30, None)) # log|dv/du| = -log(Kuma density at v) # log p_Kuma(v; a, b) = log(a)+log(b)+(a-1)log(v)+(b-1)log(1-v^a) va = v ** a log_kuma_at_v = (np.log(a) + np.log(b) + (a - 1.0) * np.log(v) + (b - 1.0) * np.log(np.clip(1.0 - va, 1e-30, None))) log_abs_dv_du = -log_kuma_at_v log_p_per_dim = log_pL_v - log_Z + log_abs_dv_du # shape (d,) return float(np.sum(log_p_per_dim))