Move and rename VectorizedVariable (#740)

Author: dylanhmorris

docs/tutorials/observation_processes_measurements.qmd

@@ -24,9 +24,8 @@ from _tutorial_theme import theme_tutorial
 from pyrenew.observation import (
     Measurements,
     HierarchicalNormalNoise,
-    VectorizedRV,
 )
-from pyrenew.randomvariable import DistributionalVariable
+from pyrenew.randomvariable import DistributionalVariable, VectorizedVariable
 from pyrenew.deterministic import DeterministicVariable, DeterministicPMF
 ```
 
@@ -83,18 +82,18 @@ Measurement data typically exhibits **sensor-level variability**: different inst
 observed ~ Normal(predicted + sensor_mode[sensor], sensor_sd[sensor])
 ```
 
-The sensor-level RVs must implement `sample(n_groups=...)`. Use `VectorizedRV` to wrap simple distributions:
+The sensor-level RVs must implement `sample(n_groups=...)`. Use `VectorizedVariable` to wrap simple distributions:
 
 ```{python}
 # | label: noise-model-general
 # Sensor modes: zero-centered, allowing positive or negative bias
-sensor_mode_rv = VectorizedRV(
+sensor_mode_rv = VectorizedVariable(
     "vec_sensor_mode",
     DistributionalVariable("sensor_mode", dist.Normal(0, 0.5)),
 )
 
 # Sensor SDs: must be positive, truncated normal is a common choice
-sensor_sd_rv = VectorizedRV(
+sensor_sd_rv = VectorizedVariable(
     "vec_sensor_sd",
     DistributionalVariable(
         "sensor_sd", dist.TruncatedNormal(loc=0.3, scale=0.15, low=0.05)
@@ -355,13 +354,13 @@ For wastewater, a "sensor" is a WWTP/lab pair—the combination of treatment pla
 ```{python}
 # | label: ww-noise-model
 # Sensor-level mode: systematic differences between WWTP/lab pairs
-ww_sensor_mode_rv = VectorizedRV(
+ww_sensor_mode_rv = VectorizedVariable(
     "vec_ww_sensor_mode",
     DistributionalVariable("ww_sensor_mode", dist.Normal(0, 0.5)),
 )
 
 # Sensor-level SD: measurement variability within each WWTP/lab pair
-ww_sensor_sd_rv = VectorizedRV(
+ww_sensor_sd_rv = VectorizedVariable(
     "vec_ww_sensor_sd",
     DistributionalVariable(
         "ww_sensor_sd", dist.TruncatedNormal(loc=0.3, scale=0.15, low=0.10)

pyrenew/observation/__init__.py

@@ -28,7 +28,6 @@
     MeasurementNoise,
     NegativeBinomialNoise,
     PoissonNoise,
-    VectorizedRV,
 )
 from pyrenew.observation.types import ObservationSample
 
@@ -44,7 +43,6 @@
     "NegativeBinomialNoise",
     "MeasurementNoise",
     "HierarchicalNormalNoise",
-    "VectorizedRV",
     # Observation processes
     "Counts",
     "CountsBySubpop",

pyrenew/observation/noise.py

@@ -1,4 +1,3 @@
-# numpydoc ignore=GL08
 """
 Noise models for observation processes.
 
@@ -16,10 +15,6 @@
 - ``HierarchicalNormalNoise``: Normal noise with hierarchical sensor effects.
   Takes ``sensor_mode_rv`` and ``sensor_sd_rv`` for sensor-level
   bias and variability.
-
-**Utilities**
-
-- ``VectorizedRV``: Wrapper that adds ``n_groups`` support to simple RVs.
 """
 
 from __future__ import annotations
@@ -36,60 +31,6 @@
 _EPSILON = 1e-10
 
 
-class VectorizedRV(RandomVariable):
-    """
-    Wrapper that adds n_groups support to simple RandomVariables.
-
-    Uses numpyro.plate to vectorize sampling, enabling simple RVs
-    to work with noise models expecting the group-level interface.
-
-    Parameters
-    ----------
-    name
-        A name for this random variable.
-        The numpyro plate is named ``f"{name}_plate"``.
-    rv
-        The underlying RandomVariable to wrap.
-    """
-
-    def __init__(self, name: str, rv: RandomVariable) -> None:
-        """
-        Initialize VectorizedRV wrapper.
-
-        Parameters
-        ----------
-        name
-            A name for this random variable.
-            The numpyro plate is named ``f"{name}_plate"``.
-        rv
-            The underlying RandomVariable to wrap.
-        """
-        super().__init__(name=name)
-        self.rv = rv
-        self.plate_name = f"{name}_plate"
-
-    def validate(self) -> None:  # pragma: no cover
-        """Validate the underlying RV."""
-        self.rv.validate()
-
-    def sample(self, n_groups: int, **kwargs: object) -> ArrayLike:
-        """
-        Sample n_groups values using numpyro.plate.
-
-        Parameters
-        ----------
-        n_groups
-            Number of group-level values to sample.
-
-        Returns
-        -------
-        ArrayLike
-            Array of shape (n_groups,).
-        """
-        with numpyro.plate(self.plate_name, n_groups):
-            return self.rv(**kwargs)
-
-
 class CountNoise(ABC):
     """
     Abstract base for count observation noise models.
@@ -356,7 +297,8 @@ class HierarchicalNormalNoise(MeasurementNoise):
 
     Notes
     -----
-    Use ``VectorizedRV`` to wrap simple RVs that lack this interface.
+    Use [`VectorizedVariable`][pyrenew.randomvariable.VectorizedVariable]
+    to wrap simple RVs that lack this interface.
     """
 
     def __init__(

pyrenew/randomvariable/__init__.py

@@ -6,10 +6,12 @@
     StaticDistributionalVariable,
 )
 from pyrenew.randomvariable.transformedvariable import TransformedVariable
+from pyrenew.randomvariable.vectorizedvariable import VectorizedVariable
 
 __all__ = [
     "DistributionalVariable",
     "StaticDistributionalVariable",
     "DynamicDistributionalVariable",
     "TransformedVariable",
+    "VectorizedVariable",
 ]

pyrenew/randomvariable/vectorizedvariable.py

@@ -0,0 +1,55 @@
+"""
+Vectorization wrapper for simple RandomVariables
+"""
+
+import numpyro
+from jax.typing import ArrayLike
+
+from pyrenew.metaclass import RandomVariable
+
+
+class VectorizedVariable(RandomVariable):
+    """
+    Wrapper that adds n_groups support to simple RandomVariables.
+
+    Uses numpyro.plate to vectorize sampling, enabling simple RVs
+    to work with noise models expecting the group-level interface.
+    """
+
+    def __init__(self, name: str, rv: RandomVariable) -> None:
+        """
+        Initialize VectorizedVariable wrapper.
+
+        Parameters
+        ----------
+        name
+            A name for this random variable.
+            The numpyro plate used to vectorize will
+            have this name with the suffix `_plate"`.
+        rv
+            The underlying RandomVariable to wrap.
+        """
+        super().__init__(name=name)
+        self.rv = rv
+        self.plate_name = f"{name}_plate"
+
+    def validate(self) -> None:  # pragma: no cover
+        """Validate the underlying RV."""
+        self.rv.validate()
+
+    def sample(self, n_groups: int, **kwargs: object) -> ArrayLike:
+        """
+        Sample n_groups values using numpyro.plate.
+
+        Parameters
+        ----------
+        n_groups
+            Number of group-level values to sample.
+
+        Returns
+        -------
+        ArrayLike
+            Array of shape (n_groups,).
+        """
+        with numpyro.plate(self.plate_name, n_groups):
+            return self.rv(**kwargs)

test/conftest.py

@@ -15,9 +15,8 @@
     Counts,
     HierarchicalNormalNoise,
     NegativeBinomialNoise,
-    VectorizedRV,
 )
-from pyrenew.randomvariable import DistributionalVariable
+from pyrenew.randomvariable import DistributionalVariable, VectorizedVariable
 
 # =============================================================================
 # PMF Fixtures
@@ -103,18 +102,18 @@ def gen_int_rv():
 @pytest.fixture
 def hierarchical_normal_noise():
     """
-    Standard HierarchicalNormalNoise with VectorizedRV wrappers.
+    Standard HierarchicalNormalNoise with VectorizedVariable wrappers.
 
     Returns
     -------
     HierarchicalNormalNoise
         Noise model for continuous measurements.
     """
-    sensor_mode_rv = VectorizedRV(
+    sensor_mode_rv = VectorizedVariable(
         name="sensor_mode_rv",
         rv=DistributionalVariable("ww_sensor_mode", dist.Normal(0, 0.5)),
     )
-    sensor_sd_rv = VectorizedRV(
+    sensor_sd_rv = VectorizedVariable(
         name="sensor_sd_rv",
         rv=DistributionalVariable(
             "ww_sensor_sd", dist.TruncatedNormal(0.3, 0.15, low=0.10)
@@ -133,11 +132,11 @@ def hierarchical_normal_noise_tight():
     HierarchicalNormalNoise
         Noise model with very small variance.
     """
-    sensor_mode_rv = VectorizedRV(
+    sensor_mode_rv = VectorizedVariable(
         name="sensor_mode_rv",
         rv=DistributionalVariable("ww_sensor_mode", dist.Normal(0, 0.01)),
     )
-    sensor_sd_rv = VectorizedRV(
+    sensor_sd_rv = VectorizedVariable(
         name="sensor_sd_rv",
         rv=DistributionalVariable(
             "ww_sensor_sd", dist.TruncatedNormal(0.01, 0.005, low=0.001)

test/test_interface_coverage.py

@@ -38,13 +38,16 @@
     NegativeBinomialNoise,
     NegativeBinomialObservation,
     PoissonNoise,
-    VectorizedRV,
 )
 from pyrenew.process import ARProcess, DifferencedProcess
 from pyrenew.process.iidrandomsequence import IIDRandomSequence, StandardNormalSequence
 from pyrenew.process.randomwalk import RandomWalk as ProcessRandomWalk
 from pyrenew.process.randomwalk import StandardNormalRandomWalk
-from pyrenew.randomvariable import DistributionalVariable, TransformedVariable
+from pyrenew.randomvariable import (
+    DistributionalVariable,
+    TransformedVariable,
+    VectorizedVariable,
+)
 from test.test_helpers import ConcreteMeasurements
 
 # =============================================================================
@@ -92,11 +95,11 @@ def _make_measurements():
     -------
     instantiated object
     """
-    sensor_mode_rv = VectorizedRV(
+    sensor_mode_rv = VectorizedVariable(
         name="sensor_mode_rv",
         rv=DistributionalVariable("mode", dist.Normal(0, 0.5)),
     )
-    sensor_sd_rv = VectorizedRV(
+    sensor_sd_rv = VectorizedVariable(
         name="sensor_sd_rv",
         rv=DistributionalVariable("sd", dist.TruncatedNormal(0.3, 0.15, low=0.1)),
     )
@@ -115,11 +118,11 @@ def _make_hierarchical_normal_noise():
     -------
     instantiated object
     """
-    sensor_mode_rv = VectorizedRV(
+    sensor_mode_rv = VectorizedVariable(
         name="sensor_mode_rv",
         rv=DistributionalVariable("mode", dist.Normal(0, 0.5)),
     )
-    sensor_sd_rv = VectorizedRV(
+    sensor_sd_rv = VectorizedVariable(
         name="sensor_sd_rv",
         rv=DistributionalVariable("sd", dist.TruncatedNormal(0.3, 0.15, low=0.1)),
     )
@@ -455,12 +458,12 @@ def test_random_variable_rejects_invalid_name(bad_name):
         pytest.param(_make_counts_by_subpop(), "test_subpop", id="CountsBySubpop"),
         pytest.param(_make_measurements(), "test_ww", id="ConcreteMeasurements"),
         pytest.param(
-            VectorizedRV(
+            VectorizedVariable(
                 name="test_vec",
                 rv=DistributionalVariable("inner", dist.Normal(0, 1)),
             ),
             "test_vec",
-            id="VectorizedRV",
+            id="VectorizedVariable",
         ),
     ],
 )

test/test_observation_measurements.py

@@ -11,30 +11,11 @@
 import pytest
 
 from pyrenew.deterministic import DeterministicPMF
-from pyrenew.observation import (
-    HierarchicalNormalNoise,
-    VectorizedRV,
-)
-from pyrenew.randomvariable import DistributionalVariable
+from pyrenew.observation import HierarchicalNormalNoise
+from pyrenew.randomvariable import DistributionalVariable, VectorizedVariable
 from test.test_helpers import ConcreteMeasurements
 
 
-class TestVectorizedRV:
-    """Test VectorizedRV wrapper class."""
-
-    def test_init_and_sample(self):
-        """Test VectorizedRV initialization and sampling."""
-        rv = DistributionalVariable("test", dist.Normal(0, 1.0))
-        vectorized = VectorizedRV(name="test_vectorized", rv=rv)
-
-        with numpyro.handlers.seed(rng_seed=42):
-            samples = vectorized.sample(n_groups=5)
-
-        assert samples.shape == (5,)
-        # Verify samples are actually different (not degenerate)
-        assert jnp.std(samples) > 0
-
-
 class TestHierarchicalNormalNoise:
     """Test HierarchicalNormalNoise model."""
 
@@ -235,11 +216,11 @@ def test_sensor_bias_differences(self):
         shedding_pmf = jnp.array([1.0])
 
         # Use wide priors to ensure sensors get distinguishable biases
-        sensor_mode_rv = VectorizedRV(
+        sensor_mode_rv = VectorizedVariable(
             name="sensor_mode_rv",
             rv=DistributionalVariable("mode", dist.Normal(0, 2.0)),
         )
-        sensor_sd_rv = VectorizedRV(
+        sensor_sd_rv = VectorizedVariable(
             name="sensor_sd_rv",
             rv=DistributionalVariable("sd", dist.TruncatedNormal(0.1, 0.05, low=0.01)),
         )