add LogitNormalVariable as pre-specified RandomVariable (#864)

* added logit normal RV, tests, updated docs

* cleanup

* changes per code review

* changes per code review

Author: cdc-mitzimorris

docs/tutorials/random_variables.qmd

@@ -46,6 +46,7 @@ from pyrenew.randomvariable import (
     DistributionalVariable,
     StaticDistributionalVariable,
     DynamicDistributionalVariable,
+    LogitNormalVariable,
     TransformedVariable,
 )
 from pyrenew.observation import PopulationCounts, NegativeBinomialNoise
@@ -150,6 +151,29 @@ with numpyro.handlers.seed(rng_seed=0):
     print(f"Sum: {effect.sum():.1f}")
 ```
 
+`LogitNormalVariable` is a convenience factory for a common transformed prior: a Normal distribution on the log-odds scale, transformed to a probability in `(0, 1)`.
+The `median` is specified on the probability scale, while `scale` is the standard deviation on the log-odds scale.
+It returns a `TransformedVariable`, rather than introducing another `RandomVariable` implementation.
+
+```{python}
+#| label: logit-normal-example
+
+iedr = LogitNormalVariable(
+    name="iedr",
+    base_name="logit_iedr",
+    median=0.004,
+    scale=0.3,
+)
+
+with numpyro.handlers.seed(rng_seed=0):
+    value = iedr()
+    print(f"IEDR: {value:.4%}")
+    print(f"Latent sample name: {iedr.base_rv.name}")
+```
+
+The underlying Normal sample is named `logit_iedr`; the transformed variable is named `iedr`.
+An optional NumPyro reparameterizer can be supplied with `reparam` when a hierarchical model would benefit from different sampling coordinates.
+
 ### Interchangeability
 
 Because all implementations share the `sample()` interface, you can swap them freely.
@@ -478,6 +502,7 @@ The right panel shows the additional weekly oscillation introduced by day-of-wee
   | Fixed known PMF                                                       | `DeterministicPMF`                                  |
   | Sample from a fixed distribution                                      | `DistributionalVariable` (static)                   |
   | Sample from a distribution parameterized at sample time               | `DistributionalVariable` (dynamic, pass a callable) |
+  | Probability with a Normal prior on its log-odds                       | `LogitNormalVariable`                               |
   | Deterministic transformation of another RV                            | `TransformedVariable`                               |
   | Multiple sample statements, custom validation, or derived computation | Custom `RandomVariable` subclass                    |
 

pyrenew/randomvariable/__init__.py

@@ -5,13 +5,15 @@
     DynamicDistributionalVariable,
     StaticDistributionalVariable,
 )
+from pyrenew.randomvariable.logitnormalvariable import LogitNormalVariable
 from pyrenew.randomvariable.transformedvariable import TransformedVariable
 from pyrenew.randomvariable.vectorizedvariable import VectorizedVariable
 
 __all__ = [
     "DistributionalVariable",
     "StaticDistributionalVariable",
     "DynamicDistributionalVariable",
+    "LogitNormalVariable",
     "TransformedVariable",
     "VectorizedVariable",
 ]

pyrenew/randomvariable/logitnormalvariable.py

@@ -0,0 +1,76 @@
+"""Factory for logit-normal random variables."""
+
+import jax.numpy as jnp
+import numpyro.distributions as dist
+from jax.typing import ArrayLike
+from numpyro.infer.reparam import Reparam
+from numpyro.util import not_jax_tracer
+
+import pyrenew.transformation as transformation
+from pyrenew.randomvariable.distributionalvariable import DistributionalVariable
+from pyrenew.randomvariable.transformedvariable import TransformedVariable
+
+
+def LogitNormalVariable(
+    name: str,
+    median: ArrayLike,
+    scale: ArrayLike,
+    base_name: str | None = None,
+    reparam: Reparam = None,
+) -> TransformedVariable:
+    """Create a logit-normal distributed random variable.
+
+    Parameters
+    ----------
+    name
+        Name of the random variable.
+    median
+        Median of the random variable on the probability scale. All values
+        must be finite and strictly between 0 and 1.
+    scale
+        Standard deviation of the Normal distribution on the logit scale. All
+        values must be finite and strictly positive.
+    base_name
+        Name of the underlying Normal random variable. Defaults to
+        ``f"logit_{name}"``.
+    reparam
+        If not None, reparameterize sampling from the underlying Normal
+        distribution according to the given NumPyro reparameterizer.
+
+    Returns
+    -------
+    TransformedVariable
+        A transformed variable whose underlying Normal random variable has
+        location ``logit(median)`` and the specified scale.
+
+    Raises
+    ------
+    ValueError
+        If any statically available median is not finite and strictly between
+        0 and 1, or any statically available scale is not finite and positive.
+    """
+    median = jnp.asarray(median)
+    invalid_median = jnp.any(~jnp.isfinite(median) | (median <= 0) | (median >= 1))
+    if not_jax_tracer(invalid_median) and bool(invalid_median):
+        raise ValueError(
+            "median must contain only finite values strictly between 0 and 1"
+        )
+
+    scale = jnp.asarray(scale)
+    invalid_scale = jnp.any(~jnp.isfinite(scale) | (scale <= 0))
+    if not_jax_tracer(invalid_scale) and bool(invalid_scale):
+        raise ValueError("scale must contain only finite positive values")
+
+    sigmoid_transform = transformation.SigmoidTransform()
+    return TransformedVariable(
+        name=name,
+        base_rv=DistributionalVariable(
+            name=base_name or f"logit_{name}",
+            distribution=dist.Normal(
+                sigmoid_transform.inv(median),
+                scale,
+            ),
+            reparam=reparam,
+        ),
+        transforms=sigmoid_transform,
+    )

test/test_logit_normal_rv.py

@@ -0,0 +1,132 @@
+"""Tests for the LogitNormalVariable factory."""
+
+import jax
+import jax.numpy as jnp
+import numpyro
+import numpyro.distributions as dist
+import pytest
+from numpy.testing import assert_allclose
+from numpyro.infer import Predictive
+from numpyro.infer.reparam import LocScaleReparam
+
+import pyrenew.transformation as transformation
+from pyrenew.randomvariable import (
+    LogitNormalVariable,
+    StaticDistributionalVariable,
+    TransformedVariable,
+)
+
+
+def test_logit_normal_variable_construction():
+    """The factory constructs the expected transformed Normal variable."""
+    median = 0.004
+    scale = 0.3
+
+    rv = LogitNormalVariable(name="iedr", median=median, scale=scale)
+
+    assert isinstance(rv, TransformedVariable)
+    assert isinstance(rv.base_rv, StaticDistributionalVariable)
+    assert isinstance(rv.base_rv.distribution, dist.Normal)
+    assert_allclose(
+        rv.base_rv.distribution.loc,
+        transformation.SigmoidTransform().inv(median),
+    )
+    assert_allclose(rv.base_rv.distribution.scale, scale)
+    assert len(rv.transforms) == 1
+    assert isinstance(rv.transforms[0], transformation.SigmoidTransform)
+
+
+def test_logit_normal_variable_samples_are_probabilities():
+    """Samples from a logit-normal variable lie strictly between zero and one."""
+    rv = LogitNormalVariable(name="iedr", median=0.004, scale=0.3)
+
+    def model():  # numpydoc ignore=GL08
+        return rv.sample(record=True)
+
+    samples = Predictive(model, num_samples=100)(jax.random.key(0))
+    values = samples["iedr"][0]
+
+    assert (values > 0).all()
+    assert (values < 1).all()
+
+
+@pytest.mark.parametrize(
+    ("base_name", "expected_base_name"),
+    [
+        (None, "logit_iedr"),
+        ("iedr_unconstrained", "iedr_unconstrained"),
+    ],
+)
+def test_logit_normal_variable_names(base_name, expected_base_name):
+    """Default and explicit names are used for object and trace sites."""
+    rv = LogitNormalVariable(
+        name="iedr",
+        median=0.004,
+        scale=0.3,
+        base_name=base_name,
+    )
+
+    assert rv.name == "iedr"
+    assert rv.base_rv.name == expected_base_name
+
+    with numpyro.handlers.seed(rng_seed=0), numpyro.handlers.trace() as trace:
+        rv.sample(record=True)
+
+    assert trace[expected_base_name]["type"] == "sample"
+    assert trace["iedr"]["type"] == "deterministic"
+
+
+def test_logit_normal_variable_reparameterization():
+    """A reparameterizer is applied to the underlying Normal sample site."""
+    reparam = LocScaleReparam(0)
+    rv = LogitNormalVariable(
+        name="iedr",
+        median=0.004,
+        scale=0.3,
+        reparam=reparam,
+    )
+
+    assert rv.base_rv.reparam_dict == {"logit_iedr": reparam}
+
+    with numpyro.handlers.seed(rng_seed=0), numpyro.handlers.trace() as trace:
+        value = rv.sample()
+
+    assert trace["logit_iedr_decentered"]["type"] == "sample"
+    assert trace["logit_iedr"]["type"] == "deterministic"
+    assert 0 < value < 1
+
+
+@pytest.mark.parametrize(
+    "median",
+    [
+        -0.1,
+        0.0,
+        1.0,
+        1.1,
+        float("-inf"),
+        float("inf"),
+        float("nan"),
+        jnp.array([0.2, 1.0]),
+    ],
+)
+def test_logit_normal_variable_rejects_invalid_median(median):
+    """Medians must be finite and strictly inside the unit interval."""
+    with pytest.raises(ValueError, match="median must contain only finite values"):
+        LogitNormalVariable(name="invalid", median=median, scale=0.3)
+
+
+@pytest.mark.parametrize(
+    "scale",
+    [
+        -0.1,
+        0.0,
+        float("-inf"),
+        float("inf"),
+        float("nan"),
+        jnp.array([0.3, 0.0]),
+    ],
+)
+def test_logit_normal_variable_rejects_invalid_scale(scale):
+    """Scales must be finite and strictly positive."""
+    with pytest.raises(ValueError, match="scale must contain only finite positive"):
+        LogitNormalVariable(name="invalid", median=0.4, scale=scale)