create and use default rt across tests

.github/workflows/deptry.yml

@@ -29,4 +29,4 @@ jobs:
 
       - name: run deptry
         run: |
-          poetry run deptry . --per-rule-ignores "DEP001=pyrenew,DEP001=pytest,DEP003=pytest"
+          poetry run deptry . --per-rule-ignores "DEP001=pyrenew,DEP001=pytest,DEP003=pytest,DEP001=test"

model/src/test/test_forecast.py

@@ -1,21 +1,21 @@
 # numpydoc ignore=GL08
 
+from test.utils import simple_rt
+
 import jax.numpy as jnp
 import jax.random as jr
 import numpyro
 import numpyro.distributions as dist
-import pyrenew.transformation as t
 from numpy.testing import assert_array_equal
 from pyrenew.deterministic import DeterministicPMF
 from pyrenew.latent import (
     InfectionInitializationProcess,
     Infections,
     InitializeInfectionsZeroPad,
 )
-from pyrenew.metaclass import DistributionalRV, TransformedRandomVariable
+from pyrenew.metaclass import DistributionalRV
 from pyrenew.model import RtInfectionsRenewalModel
 from pyrenew.observation import PoissonObservation
-from pyrenew.process import SimpleRandomWalkProcess
 
 
 def test_forecast():
@@ -30,19 +30,7 @@ def test_forecast():
     )
     latent_infections = Infections()
     observed_infections = PoissonObservation(name="poisson_rv")
-    rt = TransformedRandomVariable(
-        name="Rt_rv",
-        base_rv=SimpleRandomWalkProcess(
-            name="log_rt",
-            step_rv=DistributionalRV(
-                name="rw_step_rv", dist=dist.Normal(0, 0.025)
-            ),
-            init_rv=DistributionalRV(
-                name="init_log_rt", dist=dist.Normal(0, 0.2)
-            ),
-        ),
-        transforms=t.ExpTransform(),
-    )
+    rt = simple_rt()
 
     model = RtInfectionsRenewalModel(
         I0_rv=I0,

model/src/test/test_latent_admissions.py

@@ -1,15 +1,15 @@
 # -*- coding: utf-8 -*-
 # numpydoc ignore=GL08
 
+from test.utils import simple_rt
+
 import jax.numpy as jnp
 import numpy.testing as testing
 import numpyro
 import numpyro.distributions as dist
-from pyrenew import transformation as t
 from pyrenew.deterministic import DeterministicPMF
 from pyrenew.latent import HospitalAdmissions, Infections
-from pyrenew.metaclass import DistributionalRV, TransformedRandomVariable
-from pyrenew.process import SimpleRandomWalkProcess
+from pyrenew.metaclass import DistributionalRV
 
 
 def test_admissions_sample():
@@ -20,19 +20,7 @@ def test_admissions_sample():
 
     # Generating Rt and Infections to compute the hospital admissions
 
-    rt = TransformedRandomVariable(
-        name="Rt_rv",
-        base_rv=SimpleRandomWalkProcess(
-            name="log_rt",
-            step_rv=DistributionalRV(
-                name="rw_step_rv", dist=dist.Normal(0, 0.025)
-            ),
-            init_rv=DistributionalRV(
-                name="init_log_rt", dist=dist.Normal(0, 0.2)
-            ),
-        ),
-        transforms=t.ExpTransform(),
-    )
+    rt = simple_rt()
 
     with numpyro.handlers.seed(rng_seed=223):
         sim_rt = rt(n_steps=30)[0].value

model/src/test/test_latent_infections.py

@@ -1,15 +1,13 @@
 # -*- coding: utf-8 -*-
 # numpydoc ignore=GL08
 
+from test.utils import simple_rt
+
 import jax.numpy as jnp
 import numpy.testing as testing
 import numpyro
-import numpyro.distributions as dist
-import pyrenew.transformation as t
 import pytest
 from pyrenew.latent import Infections
-from pyrenew.metaclass import DistributionalRV, TransformedRandomVariable
-from pyrenew.process import SimpleRandomWalkProcess
 
 
 def test_infections_as_deterministic():
@@ -18,19 +16,7 @@ def test_infections_as_deterministic():
     the same seed is used.
     """
 
-    rt = TransformedRandomVariable(
-        "Rt_rv",
-        base_rv=SimpleRandomWalkProcess(
-            name="log_rt",
-            step_rv=DistributionalRV(
-                name="rw_step_rv", dist=dist.Normal(0, 0.025)
-            ),
-            init_rv=DistributionalRV(
-                name="init_log_rt", dist=dist.Normal(0, 0.2)
-            ),
-        ),
-        transforms=t.ExpTransform(),
-    )
+    rt = simple_rt()
 
     with numpyro.handlers.seed(rng_seed=223):
         sim_rt, *_ = rt(n_steps=30)

model/src/test/test_model_basic_renewal.py

@@ -2,50 +2,24 @@
 # numpydoc ignore=GL08
 
 
+from test.utils import simple_rt
+
 import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
 import numpyro
 import numpyro.distributions as dist
 import polars as pl
-import pyrenew.transformation as t
 import pytest
 from pyrenew.deterministic import DeterministicPMF, NullObservation
 from pyrenew.latent import (
     InfectionInitializationProcess,
     Infections,
     InitializeInfectionsZeroPad,
 )
-from pyrenew.metaclass import DistributionalRV, TransformedRandomVariable
+from pyrenew.metaclass import DistributionalRV
 from pyrenew.model import RtInfectionsRenewalModel
 from pyrenew.observation import PoissonObservation
-from pyrenew.process import SimpleRandomWalkProcess
-
-
-def get_default_rt():
-    """
-    Helper function to create a default Rt
-    RandomVariable for this testing session.
-
-    Returns
-    -------
-    TransformedRandomVariable :
-       A log-scale random walk with fixed
-       init value and step size priors
-    """
-    return TransformedRandomVariable(
-        "Rt_rv",
-        base_rv=SimpleRandomWalkProcess(
-            name="log_rt",
-            step_rv=DistributionalRV(
-                name="rw_step_rv", dist=dist.Normal(0, 0.025)
-            ),
-            init_rv=DistributionalRV(
-                name="init_log_rt", dist=dist.Normal(0, 0.2)
-            ),
-        ),
-        transforms=t.ExpTransform(),
-    )
 
 
 def test_model_basicrenewal_no_timepoints_or_observations():
@@ -65,7 +39,7 @@ def test_model_basicrenewal_no_timepoints_or_observations():
 
     observed_infections = PoissonObservation("poisson_rv")
 
-    rt = get_default_rt()
+    rt = simple_rt()
 
     model1 = RtInfectionsRenewalModel(
         I0_rv=I0,
@@ -96,7 +70,7 @@ def test_model_basicrenewal_both_timepoints_and_observations():
 
     observed_infections = PoissonObservation("possion_rv")
 
-    rt = get_default_rt()
+    rt = simple_rt()
 
     model1 = RtInfectionsRenewalModel(
         I0_rv=I0,
@@ -137,7 +111,7 @@ def test_model_basicrenewal_no_obs_model():
 
     latent_infections = Infections()
 
-    rt = get_default_rt()
+    rt = simple_rt()
 
     model0 = RtInfectionsRenewalModel(
         gen_int_rv=gen_int,
@@ -210,7 +184,7 @@ def test_model_basicrenewal_with_obs_model():
 
     observed_infections = PoissonObservation("poisson_rv")
 
-    rt = get_default_rt()
+    rt = simple_rt()
 
     model1 = RtInfectionsRenewalModel(
         I0_rv=I0,
@@ -259,7 +233,7 @@ def test_model_basicrenewal_padding() -> None:  # numpydoc ignore=GL08
 
     observed_infections = PoissonObservation("poisson_rv")
 
-    rt = get_default_rt()
+    rt = simple_rt()
 
     model1 = RtInfectionsRenewalModel(
         I0_rv=I0,

model/src/test/test_model_hosp_admissions.py

@@ -2,14 +2,15 @@
 # numpydoc ignore=GL08
 
 
+from test.utils import simple_rt
+
 import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
 import numpyro
 import numpyro.distributions as dist
 import polars as pl
 import pytest
-from pyrenew import transformation as t
 from pyrenew.deterministic import (
     DeterministicPMF,
     DeterministicVariable,
@@ -21,41 +22,9 @@
     Infections,
     InitializeInfectionsZeroPad,
 )
-from pyrenew.metaclass import (
-    DistributionalRV,
-    RandomVariable,
-    SampledValue,
-    TransformedRandomVariable,
-)
+from pyrenew.metaclass import DistributionalRV, RandomVariable, SampledValue
 from pyrenew.model import HospitalAdmissionsModel
 from pyrenew.observation import PoissonObservation
-from pyrenew.process import SimpleRandomWalkProcess
-
-
-def get_default_rt():
-    """
-    Helper function to create a default Rt
-    RandomVariable for this testing session.
-
-    Returns
-    -------
-    TransformedRandomVariable :
-       A log-scale random walk with fixed
-       init value and step size priors
-    """
-    return TransformedRandomVariable(
-        "Rt_rv",
-        base_rv=SimpleRandomWalkProcess(
-            name="log_rt",
-            step_rv=DistributionalRV(
-                name="rw_step_rv", dist=dist.Normal(0, 0.025)
-            ),
-            init_rv=DistributionalRV(
-                name="init_log_rt", dist=dist.Normal(0, 0.2)
-            ),
-        ),
-        transforms=t.ExpTransform(),
-    )
 
 
 class UniformProbForTest(RandomVariable):  # numpydoc ignore=GL08
@@ -91,7 +60,7 @@ def test_model_hosp_no_timepoints_or_observations():
     I0 = DistributionalRV(name="I0", dist=dist.LogNormal(0, 1))
 
     latent_infections = Infections()
-    Rt_process = get_default_rt()
+    Rt_process = simple_rt()
 
     observed_admissions = PoissonObservation("poisson_rv")
 
@@ -156,7 +125,7 @@ def test_model_hosp_both_timepoints_and_observations():
     I0 = DistributionalRV(name="I0", dist=dist.LogNormal(0, 1))
 
     latent_infections = Infections()
-    Rt_process = get_default_rt()
+    Rt_process = simple_rt()
 
     observed_admissions = PoissonObservation("poisson_rv")
 
@@ -229,7 +198,7 @@ def test_model_hosp_no_obs_model():
     )
 
     latent_infections = Infections()
-    Rt_process = get_default_rt()
+    Rt_process = simple_rt()
 
     inf_hosp = DeterministicPMF(
         name="inf_hosp",
@@ -339,7 +308,7 @@ def test_model_hosp_with_obs_model():
     )
 
     latent_infections = Infections()
-    Rt_process = get_default_rt()
+    Rt_process = simple_rt()
     observed_admissions = PoissonObservation("poisson_rv")
 
     inf_hosp = DeterministicPMF(
@@ -426,7 +395,7 @@ def test_model_hosp_with_obs_model_weekday_phosp_2():
     )
 
     latent_infections = Infections()
-    Rt_process = get_default_rt()
+    Rt_process = simple_rt()
     observed_admissions = PoissonObservation("poisson_rv")
 
     inf_hosp = DeterministicPMF(
@@ -525,7 +494,7 @@ def test_model_hosp_with_obs_model_weekday_phosp():
     )
 
     latent_infections = Infections()
-    Rt_process = get_default_rt()
+    Rt_process = simple_rt()
 
     observed_admissions = PoissonObservation("poisson_rv")
 

model/src/test/test_predictive.py

@@ -4,20 +4,20 @@
 Ensures that posterior predictive samples are not generated when no posterior samples are available.
 """
 
+from test.utils import simple_rt
+
 import jax.numpy as jnp
 import numpyro.distributions as dist
-import pyrenew.transformation as t
 import pytest
 from pyrenew.deterministic import DeterministicPMF
 from pyrenew.latent import (
     InfectionInitializationProcess,
     Infections,
     InitializeInfectionsZeroPad,
 )
-from pyrenew.metaclass import DistributionalRV, TransformedRandomVariable
+from pyrenew.metaclass import DistributionalRV
 from pyrenew.model import RtInfectionsRenewalModel
 from pyrenew.observation import PoissonObservation
-from pyrenew.process import SimpleRandomWalkProcess
 
 pmf_array = jnp.array([0.25, 0.1, 0.2, 0.45])
 gen_int = DeterministicPMF(name="gen_int", value=pmf_array)
@@ -29,17 +29,7 @@
 )
 latent_infections = Infections()
 observed_infections = PoissonObservation("poisson_rv")
-rt = TransformedRandomVariable(
-    "Rt_rv",
-    base_rv=SimpleRandomWalkProcess(
-        name="log_rt",
-        step_rv=DistributionalRV(
-            name="rw_step_rv", dist=dist.Normal(0, 0.025)
-        ),
-        init_rv=DistributionalRV(name="init_log_rt", dist=dist.Normal(0, 0.2)),
-    ),
-    transforms=t.ExpTransform(),
-)
+rt = simple_rt()
 
 model = RtInfectionsRenewalModel(
     I0_rv=I0,