add daily_to_weekly function

pyrenew/convolve.py

@@ -217,3 +217,45 @@ def compute_delay_ascertained_incidence(
         mode="valid",
     )
     return delay_obs_incidence
+
+
+def daily_to_epiweekly(
+    daily_value: ArrayLike,
+    first_dow: int = 0,
+):
+    """
+    Aggregate daily values (e.g.
+    hospitalizations) into epiweekly total values.
+
+    Parameters
+    ----------
+    daily_value : ArrayLike
+        Daily infections or hospitalization values.
+    first_dow : int
+        First day of the week, values between 0-6.
+        (0 for Monday, 6 for Sunday).
+        If first_dow is not 0, the incomplete first
+        epiweek is ignored and epiweekly values
+        starting from second week is returned.
+        Defaults to 0.
+
+    Returns
+    -------
+    ArrayLike
+        Data converted to epiweekly values starting
+        with the first full epiweek available.
+    """
+    if first_dow < 0 or first_dow > 6:
+        raise ValueError("First day of the week must be between 0 and 6")
+
+    if first_dow > 0:
+        daily_value = daily_value[7 - first_dow :]
+
+    if len(daily_value) < 7:
+        raise ValueError("No complete epiweekly values available")
+
+    epiweekly_values = jnp.convolve(daily_value, jnp.ones(7), mode="valid")[
+        ::7
+    ]
+
+    return epiweekly_values

test/test_daily_to_epiweekly.py

@@ -0,0 +1,76 @@
+# numpydoc ignore=GL08
+
+import jax.numpy as jnp
+import pytest
+
+from pyrenew.convolve import daily_to_epiweekly
+
+
+def test_daily_to_epiweekly_no_offset():
+    """
+    Tests that the function correctly aggregates
+    daily values into epiweekly totals when there
+    is no offset, starting from the first day.
+    """
+    daily_values = jnp.arange(1, 15)
+    result = daily_to_epiweekly(daily_values)
+    expected = jnp.array([28, 77])
+    assert jnp.array_equal(result, expected)
+
+
+def test_daily_to_epiweekly_with_offset():
+    """
+    Tests that the function correctly aggregates
+    daily values into epiweekly totals when there
+    is a dow offset.
+    """
+    daily_values = jnp.arange(1, 15)
+    result = daily_to_epiweekly(daily_values, first_dow=2)
+    expected = jnp.array([63])
+    assert jnp.array_equal(result, expected)
+
+
+def test_daily_to_epiweekly_incomplete_epiweek():
+    """
+    Tests that the function raises a
+    ValueError when there are
+    insufficient daily values to
+    form a complete epiweek.
+    """
+    daily_values = jnp.arange(1, 5)
+    with pytest.raises(
+        ValueError, match="No complete epiweekly values available"
+    ):
+        daily_to_epiweekly(daily_values, first_dow=0)
+
+
+def test_daily_to_epiweekly_missing_daily_values():
+    """
+    Tests that the function correctly
+    aggregates the available daily values
+    into epiweekly values when there are
+    fewer daily values than required for
+    complete epiweekly totals.
+    """
+    daily_values = jnp.arange(1, 10)
+    result = daily_to_epiweekly(daily_values, first_dow=0)
+    expected = jnp.array([28])
+    assert jnp.array_equal(result, expected)
+
+
+def test_daily_to_epiweekly_invalid_offset():
+    """
+    Tests that the function raises a
+    ValueError when the offset is
+    outside the valid range (0-6).
+    """
+    daily_values = jnp.arange(1, 15)
+    with pytest.raises(
+        ValueError, match="First day of the week must be between 0 and 6"
+    ):
+        daily_to_epiweekly(daily_values, first_dow=-1)
+
+    with pytest.raises(
+        ValueError, match="First day of the week must be between 0 and 6"
+    ):
+        daily_to_epiweekly(daily_values, first_dow=7)