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 get_default_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 = get_default_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 get_default_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 = get_default_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 get_default_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 = get_default_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 get_default_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():

model/src/test/test_model_hosp_admissions.py

@@ -2,14 +2,15 @@
 # numpydoc ignore=GL08
 
 
+from test.utils import get_default_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

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 get_default_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 = get_default_rt()
 
 model = RtInfectionsRenewalModel(
     I0_rv=I0,

model/src/test/test_random_key.py

@@ -5,22 +5,22 @@
 with different random keys behave appropriately.
 """
 
+from test.utils import get_default_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, assert_raises
 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 create_test_model():  # numpydoc ignore=GL08
@@ -34,19 +34,7 @@ def create_test_model():  # numpydoc ignore=GL08
     )
     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 = get_default_rt()
     model = RtInfectionsRenewalModel(
         I0_rv=I0,
         gen_int_rv=gen_int,

model/src/test/utils.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+
+"""
+test utilities
+"""
+
+import numpyro.distributions as dist
+import pyrenew.transformation as t
+from pyrenew.metaclass import DistributionalRV, TransformedRandomVariable
+from pyrenew.process import SimpleRandomWalkProcess
+
+
+def get_default_rt():
+    """
+    Helper function to create a default Rt
+    RandomVariable for testing.
+
+    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(),
+    )