Add Stable distribution with numerically integrated log-probability calculation (StableWithLogProb).

docs/source/distributions.rst

@@ -407,6 +407,13 @@ Stable
     :undoc-members:
     :show-inheritance:
 
+StableWithLogProb
+-----------------
+.. autoclass:: pyro.distributions.StableWithLogProb
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 TruncatedPolyaGamma
 -------------------
 .. autoclass:: pyro.distributions.TruncatedPolyaGamma

pyro/distributions/__init__.py

@@ -119,7 +119,7 @@
 from pyro.distributions.sine_skewed import SineSkewed
 from pyro.distributions.softlaplace import SoftLaplace
 from pyro.distributions.spanning_tree import SpanningTree
-from pyro.distributions.stable import Stable
+from pyro.distributions.stable import Stable, StableWithLogProb
 from pyro.distributions.torch import __all__ as torch_dists
 from pyro.distributions.torch_distribution import (
     ExpandedDistribution,
@@ -234,6 +234,7 @@
     "SoftLaplace",
     "SpanningTree",
     "Stable",
+    "StableWithLogProb",
     "StudentT",
     "TorchDistribution",
     "TransformModule",

pyro/distributions/stable.py

@@ -7,6 +7,7 @@
 from torch.distributions import constraints
 from torch.distributions.utils import broadcast_all
 
+from pyro.distributions.stable_log_prob import StableLogProb
 from pyro.distributions.torch_distribution import TorchDistribution
 
 
@@ -204,3 +205,24 @@ def mean(self):
     def variance(self):
         var = self.scale * self.scale
         return var.mul(2).masked_fill(self.stability < 2, math.inf)
+
+
+class StableWithLogProb(StableLogProb, Stable):
+    r"""
+    Levy :math:`\alpha`-stable distribution that is based on
+    :class:`Stable` but with an added method for calculating the
+    log probability density using numerical integration.
+
+    This should be used in cases where reparameterization does not work
+    like when trying to estimate the skew :math:`\beta` parameter. Running
+    times are slower than with reparameterization.
+
+    The numerical integration implementation is based on the algorithm
+    proposed by Chambers, Mallows and Stuck (CMS) for simulating the
+    Levy :math:`\alpha`-stable distribution. The CMS algorithm involves a
+    nonlinear transformation of two independent random variables into
+    one stable random variable. The first random variable is uniformly
+    distributed while the second is exponentially distributed. The numerical
+    integration is performed over the first uniformly distributed random
+    variable.
+    """

pyro/distributions/stable_log_prob.py

@@ -0,0 +1,194 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from functools import partial
+
+import torch
+
+value_near_zero_tolerance = 0.01
+alpha_near_one_tolerance = 0.05
+
+
+finfo = torch.finfo(torch.float64)
+MAX_LOG = math.log10(finfo.max)
+MIN_LOG = math.log10(finfo.tiny)
+
+
+def create_integrator(num_points):
+    from scipy.special import roots_legendre
+
+    roots, weights = roots_legendre(num_points)
+    roots = torch.Tensor(roots).double()
+    weights = torch.Tensor(weights).double()
+    log_weights = weights.log()
+    half_roots = roots * 0.5
+
+    def integrate(fn, domain):
+        sl = [slice(None)] + (len(domain.shape) - 1) * [None]
+        half_roots_sl = half_roots[sl]
+        value = domain[0] * (0.5 - half_roots_sl) + domain[1] * (0.5 + half_roots_sl)
+        return (
+            torch.logsumexp(fn(value) + log_weights[sl], dim=0)
+            + ((domain[1] - domain[0]) / 2).log()
+        )
+
+    return integrate
+
+
+def set_integrator(num_points):
+    global integrate
+    integrate = create_integrator(num_points)
+
+
+# Stub which is replaced by the default integrator when called for the first time
+# if a default integrator has not already been set.
+def integrate(*args, **kwargs):
+    set_integrator(num_points=501)
+    return integrate(*args, **kwargs)
+
+
+class StableLogProb:
+    def log_prob(self, value):
+        # Undo shift and scale
+        value = (value - self.loc) / self.scale
+
+        # Use double precision math
+        alpha = self.stability.double()
+        beta = self.skew.double()
+        value = value.double()
+
+        return _stable_log_prob(alpha, beta, value, self.coords) - self.scale.log()
+
+
+def _stable_log_prob(alpha, beta, value, coords):
+    # Convert to Nolan's parametrization S^0 where samples depend
+    # continuously on (alpha,beta), allowing interpolation around the hole at
+    # alpha=1.
+    if coords == "S":
+        value = torch.where(
+            alpha == 1, value, value - beta * (math.pi / 2 * alpha)
+        ).tan()
+    elif coords != "S0":
+        raise ValueError("Unknown coords: {}".format(coords))
+
+    # Find near one alpha
+    idx = (alpha - 1).abs() < alpha_near_one_tolerance
+
+    log_prob = _unsafe_alpha_stable_log_prob_S0(
+        torch.where(idx, 1 + alpha_near_one_tolerance, alpha), beta, value
+    )
+
+    # Handle alpha near one by interpolation
+    if idx.any():
+        log_prob_pos = log_prob[idx]
+        log_prob_neg = _unsafe_alpha_stable_log_prob_S0(
+            (1 - alpha_near_one_tolerance) * log_prob_pos.new_ones(log_prob_pos.shape),
+            beta[idx],
+            value[idx],
+        )
+        weights = (alpha[idx] - 1) / (2 * alpha_near_one_tolerance) + 0.5
+        log_prob[idx] = torch.logsumexp(
+            torch.stack(
+                (log_prob_pos + weights.log(), log_prob_neg + (1 - weights).log()),
+                dim=0,
+            ),
+            dim=0,
+        )
+
+    return log_prob
+
+
+def _unsafe_alpha_stable_log_prob_S0(alpha, beta, Z):
+    # Calculate log-probability of Z in Nolan's parametrization S^0. This will fail if alpha is close to 1
+
+    # Convert from Nolan's parametrization S^0 where samples depend
+    # continuously on (alpha,beta), allowing interpolation around the hole at
+    # alpha=1.
+    Z = Z + beta * (math.pi / 2 * alpha).tan()
+
+    # Find near zero values
+    per_alpha_value_near_zero_tolerance = (
+        value_near_zero_tolerance * alpha / (1 - alpha).abs()
+    )
+    idx = Z.abs() < per_alpha_value_near_zero_tolerance
+
+    # Calculate log-prob at safe values
+    log_prob = _unsafe_stable_log_prob(
+        alpha, beta, torch.where(idx, per_alpha_value_near_zero_tolerance, Z)
+    )
+
+    # Handle near zero values by interpolation
+    if idx.any():
+        log_prob_pos = log_prob[idx]
+        log_prob_neg = _unsafe_stable_log_prob(
+            alpha[idx], beta[idx], -per_alpha_value_near_zero_tolerance[idx]
+        )
+        weights = Z[idx] / (2 * per_alpha_value_near_zero_tolerance[idx]) + 0.5
+        log_prob[idx] = torch.logsumexp(
+            torch.stack(
+                (log_prob_pos + weights.log(), log_prob_neg + (1 - weights).log()),
+                dim=0,
+            ),
+            dim=0,
+        )
+
+    return log_prob
+
+
+def _unsafe_stable_log_prob(alpha, beta, Z):
+    # Calculate log-probability of Z. This will fail if alpha is close to 1
+    # or if Z is close to 0
+    ha = math.pi / 2 * alpha
+    b = beta * ha.tan()
+    atan_b = b.atan()
+    u_zero = -alpha.reciprocal() * atan_b
+
+    # If sample should be negative calculate with flipped beta and flipped value
+    flip_beta_x = Z < 0
+    beta = torch.where(flip_beta_x, -beta, beta)
+    u_zero = torch.where(flip_beta_x, -u_zero, u_zero)
+    Z = torch.where(flip_beta_x, -Z, Z)
+
+    # Set integration domwin
+    domain = torch.stack((u_zero, 0.5 * math.pi * u_zero.new_ones(u_zero.shape)), dim=0)
+
+    integrand = partial(
+        _unsafe_stable_given_uniform_log_prob, alpha=alpha, beta=beta, Z=Z
+    )
+
+    return integrate(integrand, domain) - math.log(math.pi)
+
+
+def _unsafe_stable_given_uniform_log_prob(V, alpha, beta, Z):
+    # Calculate log-probability of Z given V. This will fail if alpha is close to 1
+    # or if Z is close to 0
+    inv_alpha_minus_one = (alpha - 1).reciprocal()
+    half_pi = math.pi / 2
+    eps = torch.finfo(V.dtype).eps
+    # make V belong to the open interval (-pi/2, pi/2)
+    V = V.clamp(min=2 * eps - half_pi, max=half_pi - 2 * eps)
+    ha = half_pi * alpha
+    b = beta * ha.tan()
+    atan_b = b.atan()
+    cos_V = V.cos()
+
+    # +/- `ha` term to keep the precision of alpha * (V + half_pi) when V ~ -half_pi
+    v = atan_b - ha + alpha * (V + half_pi)
+
+    term1_log = atan_b.cos().log() * inv_alpha_minus_one
+    term2_log = (Z * cos_V / v.sin()).log() * alpha * inv_alpha_minus_one
+    term3_log = ((v - V).cos() / cos_V).log()
+
+    W_log = term1_log + term2_log + term3_log
+
+    W = W_log.clamp(min=MIN_LOG, max=MAX_LOG).exp()
+
+    log_prob = -W + (alpha * W / Z / (alpha - 1)).abs().log()
+
+    # Infinite W means zero-probability
+    log_prob = torch.where(W == torch.inf, -torch.inf, log_prob)
+
+    log_prob = log_prob.clamp(min=MIN_LOG, max=MAX_LOG)
+
+    return log_prob

pyro/infer/reparam/stable.py

@@ -158,7 +158,11 @@ def apply(self, msg):
         is_observed = msg["is_observed"]
 
         fn, event_dim = self._unwrap(fn)
-        assert isinstance(fn, dist.Stable) and fn.coords == "S0"
+        assert (
+            isinstance(fn, dist.Stable)
+            and fn.coords == "S0"
+            and not isinstance(fn, dist.StableWithLogProb)
+        )
 
         # Strategy: Let X ~ S0(a,b,s,m) be the stable variable of interest.
         # 1. WLOG scale and shift so s=1 and m=0, additionally shifting to convert

pyro/infer/reparam/strategies.py

@@ -114,7 +114,7 @@ def _minimal_reparam(fn, is_observed):
                 return TransformReparam()  # Then reparametrize new sites.
         fn = fn.base_dist
 
-    if isinstance(fn, dist.Stable):
+    if isinstance(fn, dist.Stable) and not isinstance(fn, dist.StableWithLogProb):
         if not is_observed:
             return LatentStableReparam()
         elif fn.skew.requires_grad or fn.skew.any():

tests/distributions/test_stable_with_log_prob.py

@@ -0,0 +1,99 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import pytest
+import torch
+
+import pyro
+from pyro.distributions import StableWithLogProb as Stable
+from pyro.distributions import constraints
+from pyro.infer import SVI, Trace_ELBO
+from pyro.infer.autoguide import AutoNormal
+from tests.common import assert_close
+
+torch.set_default_dtype(torch.float64)
+
+
+@pytest.mark.parametrize(
+    "alpha, beta, c, mu",
+    [
+        (1.00, 0.8, 2.0, 3.0),
+        (1.02, -0.8, 2.0, -3.0),
+        (0.98, 0.5, 1.0, -3.0),
+        (0.95, -0.5, 1.0, 3.0),
+        (1.10, 0.0, 1.0, 0.0),
+        (1.80, -0.5, 1.0, -2.0),
+        (0.50, 0.0, 1.0, 2.0),
+    ],
+)
+@pytest.mark.parametrize(
+    "alpha_0, beta_0, c_0, mu_0",
+    [
+        (1.3, 0.0, 1.0, 0.0),
+    ],
+)
+def test_stable_with_log_prob_param_fit(alpha, beta, c, mu, alpha_0, beta_0, c_0, mu_0):
+    # Sample test data
+    n = 10000
+    pyro.set_rng_seed(20240520)
+    data = Stable(alpha, beta, c, mu).sample((n,))
+
+    def model(data):
+        alpha = pyro.param(
+            "alpha", torch.tensor(alpha_0), constraint=constraints.interval(0, 2)
+        )
+        beta = pyro.param(
+            "beta", torch.tensor(beta_0), constraint=constraints.interval(-1, 1)
+        )
+        c = pyro.param("c", torch.tensor(c_0), constraint=constraints.positive)
+        mu = pyro.param("mu", torch.tensor(mu_0), constraint=constraints.real)
+        with pyro.plate("data", data.shape[0]):
+            pyro.sample("obs", Stable(alpha, beta, c, mu), obs=data)
+
+    def train(model, guide, num_steps=400, lr=0.03):
+        pyro.clear_param_store()
+        pyro.set_rng_seed(20240520)
+
+        # set up ELBO, and optimizer
+        elbo = Trace_ELBO()
+        elbo.loss(model, guide, data=data)
+        optim = pyro.optim.Adam({"lr": lr})
+        svi = SVI(model, guide, optim, loss=elbo)
+
+        # optimize
+        for i in range(num_steps):
+            loss = svi.step(data) / data.numel()
+            if i % 10 == 0:
+                logging.info(f"step {i} loss = {loss:0.6g}")
+                log_progress()
+
+        logging.info(f"Parameter estimates (n = {n}):")
+        log_progress()
+
+    def log_progress():
+        logging.info(f"alpha: Estimate = {pyro.param('alpha')}, true = {alpha}")
+        logging.info(f"beta: Estimate = {pyro.param('beta')}, true = {beta}")
+        logging.info(f"c: Estimate = {pyro.param('c')}, true = {c}")
+        logging.info(f"mu: Estimate = {pyro.param('mu')}, true = {mu}")
+
+    # Fit model to data
+    guide = AutoNormal(model)
+    train(model, guide)
+
+    # Verify fit accuracy
+    assert_close(alpha, pyro.param("alpha").item(), atol=0.03)
+    assert_close(beta, pyro.param("beta").item(), atol=0.06)
+    assert_close(c, pyro.param("c").item(), atol=0.2)
+    assert_close(mu, pyro.param("mu").item(), atol=0.2)
+
+
+# # The below tests will be executed:
+# test_stable_with_log_prob_param_fit(1.00,  0.8,  2.0,  3.0,  1.3,  0.0,  1.0,  0.0)
+# test_stable_with_log_prob_param_fit(1.02, -0.8,  2.0, -3.0,  1.3,  0.0,  1.0,  0.0)
+# test_stable_with_log_prob_param_fit(0.98,  0.5,  1.0, -3.0,  1.3,  0.0,  1.0,  0.0)
+# test_stable_with_log_prob_param_fit(0.95, -0.5,  1.0,  3.0,  1.3,  0.0,  1.0,  0.0)
+# test_stable_with_log_prob_param_fit(1.10,  0.0,  1.0,  0.0,  1.3,  0.0,  1.0,  0.0)
+# test_stable_with_log_prob_param_fit(1.80, -0.5,  1.0, -2.0,  1.3,  0.0,  1.0,  0.0)
+# test_stable_with_log_prob_param_fit(0.50,  0.0,  1.0,  2.0,  1.3,  0.0,  1.0,  0.0)