update infectionswithfeedback process

pyrenew/latent/infectionswithfeedback.py

@@ -167,6 +167,10 @@ def sample(
                 **kwargs,
             )
         )
+
+        if inf_feedback_strength.ndim < Rt.ndim:
+            inf_feedback_strength = jnp.expand_dims(inf_feedback_strength, 0)
+
         # Making sure inf_feedback_strength spans the Rt length
         if inf_feedback_strength.shape[0] == 1:
             inf_feedback_strength = au.pad_edges_to_match(

test/test_infection_and_infectionwithfeedback.py

@@ -0,0 +1,79 @@
+"""
+Test to verify Infection and InfectionsWithFeedback class
+return error when input array shape for I0 and Rt are invalid
+"""
+
+import jax.numpy as jnp
+import numpy as np
+import numpyro
+import pytest
+
+import pyrenew.latent as latent
+from pyrenew.deterministic import DeterministicPMF, DeterministicVariable
+
+
+def test_infections_with_feedback_invalid_inputs():
+    """
+    Test the InfectionsWithFeedback class cannot
+    be sampled when Rt and I0 have invalid input shapes
+    """
+    I0_1d = jnp.array([0.5, 0.6, 0.7, 0.8])
+    I0_2d = jnp.array(
+        np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] * 3)
+    ).reshape((7, -1))
+    Rt = jnp.ones(10)
+    gen_int = jnp.array([0.4, 0.25, 0.25, 0.1, 0.0, 0.0, 0.0])
+
+    inf_feed_strength = DeterministicVariable(
+        name="inf_feed_strength", value=0.5
+    )
+    inf_feedback_pmf = DeterministicPMF(name="inf_feedback_pmf", value=gen_int)
+
+    # Test the InfectionsWithFeedback class
+    InfectionsWithFeedback = latent.InfectionsWithFeedback(
+        infection_feedback_strength=inf_feed_strength,
+        infection_feedback_pmf=inf_feedback_pmf,
+    )
+
+    infections = latent.Infections()
+
+    with numpyro.handlers.seed(rng_seed=0):
+        with pytest.raises(
+            ValueError,
+            match="Initial infections must be at least as long as the generation interval.",
+        ):
+            InfectionsWithFeedback(
+                gen_int=gen_int,
+                Rt=Rt,
+                I0=I0_1d,
+            )
+
+        with pytest.raises(
+            ValueError,
+            match="Initial infections vector must be at least as long as the generation interval.",
+        ):
+            infections(
+                gen_int=gen_int,
+                Rt=Rt,
+                I0=I0_1d,
+            )
+
+        with pytest.raises(
+            ValueError,
+            match="Initial infections and Rt must have the same dimensions.",
+        ):
+            InfectionsWithFeedback(
+                gen_int=gen_int,
+                Rt=Rt,
+                I0=I0_2d,
+            )
+
+        with pytest.raises(
+            ValueError,
+            match="Initial infections and Rt must have the same dimensions.",
+        ):
+            infections(
+                gen_int=gen_int,
+                Rt=Rt,
+                I0=I0_2d,
+            )

test/test_infectionsrtfeedback.py

@@ -91,7 +91,7 @@ def test_infectionsrtfeedback(Rt, I0):
     # By doing the infection feedback strength 0, Rt = Rt_adjusted
     # So infection should be equal in both
     inf_feed_strength = DeterministicVariable(
-        name="inf_feed_strength", value=jnp.zeros_like(Rt)
+        name="inf_feed_strength", value=jnp.array(0)
     )
     inf_feedback_pmf = DeterministicPMF(name="inf_feedback_pmf", value=gen_int)
 
@@ -193,70 +193,3 @@ def test_infectionsrtfeedback_feedback(Rt, I0):
     assert_array_almost_equal(samp1.rt, res["rt"])
 
     return None
-
-
-def test_infections_with_feedback_invalid_inputs():
-    """
-    Test the InfectionsWithFeedback class cannot
-    be sampled when Rt and I0 have invalid input shapes
-    """
-    I0_1d = jnp.array([0.5, 0.6, 0.7, 0.8])
-    I0_2d = jnp.array(
-        np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] * 3)
-    ).reshape((7, -1))
-    Rt = jnp.ones(10)
-    gen_int = jnp.array([0.4, 0.25, 0.25, 0.1, 0.0, 0.0, 0.0])
-
-    inf_feed_strength = DeterministicVariable(
-        name="inf_feed_strength", value=0.5
-    )
-    inf_feedback_pmf = DeterministicPMF(name="inf_feedback_pmf", value=gen_int)
-
-    # Test the InfectionsWithFeedback class
-    InfectionsWithFeedback = latent.InfectionsWithFeedback(
-        infection_feedback_strength=inf_feed_strength,
-        infection_feedback_pmf=inf_feedback_pmf,
-    )
-
-    infections = latent.Infections()
-
-    with numpyro.handlers.seed(rng_seed=0):
-        with pytest.raises(
-            ValueError,
-            match="Initial infections must be at least as long as the generation interval.",
-        ):
-            InfectionsWithFeedback(
-                gen_int=gen_int,
-                Rt=Rt,
-                I0=I0_1d,
-            )
-
-        with pytest.raises(
-            ValueError,
-            match="Initial infections vector must be at least as long as the generation interval.",
-        ):
-            infections(
-                gen_int=gen_int,
-                Rt=Rt,
-                I0=I0_1d,
-            )
-
-        with pytest.raises(
-            ValueError,
-            match="Initial infections and Rt must have the same dimensions.",
-        ):
-            InfectionsWithFeedback(
-                gen_int=gen_int,
-                Rt=Rt,
-                I0=I0_2d,
-            )
-
-        with pytest.raises(
-            ValueError,
-            match="Initial infections and Rt must have the same dimensions.",
-        ):
-            infections(
-                gen_int=gen_int,
-                Rt=Rt,
-                I0=I0_2d,
-            )

test/test_infectionwithfeedback_plate_compatibility.py

@@ -0,0 +1,44 @@
+"""
+Test the InfectionsWithFeedback class works well within numpyro plate
+"""
+
+import jax.numpy as jnp
+import numpy as np
+import numpyro
+import numpyro.distributions as dist
+
+import pyrenew.latent as latent
+from pyrenew.deterministic import DeterministicPMF
+from pyrenew.randomvariable import DistributionalVariable
+
+
+def test_infections_with_feedback_plate_compatibility():
+    """
+    Test the InfectionsWithFeedback matching the Infections class.
+    """
+    I0 = jnp.array(
+        np.array([0.0, 0.0, 0.0, 0.5, 0.6, 0.7, 0.8] * 5).reshape(-1, 5)
+    )
+    Rt = jnp.ones((10, 5))
+    gen_int = jnp.array([0.4, 0.25, 0.25, 0.1])
+
+    inf_feed_strength = DistributionalVariable(
+        "inf_feed_strength", dist.Beta(1, 1)
+    )
+    inf_feedback_pmf = DeterministicPMF(name="inf_feedback_pmf", value=gen_int)
+
+    # Test the InfectionsWithFeedback class
+    InfectionsWithFeedback = latent.InfectionsWithFeedback(
+        infection_feedback_strength=inf_feed_strength,
+        infection_feedback_pmf=inf_feedback_pmf,
+    )
+
+    with numpyro.handlers.seed(rng_seed=0):
+        with numpyro.plate("test_plate", 5):
+            samp = InfectionsWithFeedback(
+                gen_int=gen_int,
+                Rt=Rt,
+                I0=I0,
+            )
+
+    assert samp.rt.shape == Rt.shape