add daily_to_weekly function

pyrenew/convolve.py

@@ -219,43 +219,74 @@ def compute_delay_ascertained_incidence(
     return delay_obs_incidence
 
 
-def daily_to_epiweekly(
-    daily_value: ArrayLike,
-    first_dow: int = 0,
+def daily_to_weekly(
+    daily_values: ArrayLike,
+    input_data_first_dow: int = 0,
+    week_start_dow: int = 0,
 ):
     """
     Aggregate daily values (e.g.
     incident hospital admissions) into weekly total values.
 
     Parameters
     ----------
-    daily_value : ArrayLike
+    daily_values : ArrayLike
         Daily timeseries values (e.g. incident infections or incident ed visits).
-    first_dow : int
-        First day of the week in the input timeseries `daily_values`. An integer between 0 and 6, inclusive.
-        (0 for Monday, 6 for Sunday).
+    input_data_first_dow : int
+        First day of the week in the input timeseries `daily_values`.
+        An integer between 0 and 6, inclusive (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.
+        week is ignored and weekly values starting
+        from second week is returned. Defaults to 0.
+    week_start_dow : int
+        The desired starting day of the week for the output weekly aggregation.
+        An integer between 0 and 6, inclusive. Defaults to 0 (Monday).
 
     Returns
     -------
     ArrayLike
-        Data converted to epiweekly values starting
-        with the first full epiweek available.
+        Data converted to weekly values starting
+        with the first full week available.
     """
-    if first_dow < 0 or first_dow > 6:
-        raise ValueError("First day of the week must be between 0 and 6")
+    if input_data_first_dow < 0 or input_data_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 week_start_dow < 0 or week_start_dow > 6:
+        raise ValueError(
+            "First day of aggregated data must be between 0 and 6."
+        )
+
+    if input_data_first_dow != week_start_dow:
+        offset = (week_start_dow - input_data_first_dow) % 7
+        daily_values = daily_values[offset:]
+
+    if len(daily_values) < 7:
+        raise ValueError("No complete weekly values available")
+
+    weekly_values = jnp.convolve(daily_values, jnp.ones(7), mode="valid")[::7]
 
-    if len(daily_value) < 7:
-        raise ValueError("No complete epiweekly values available")
+    return weekly_values
 
-    epiweekly_values = jnp.convolve(daily_value, jnp.ones(7), mode="valid")[
-        ::7
-    ]
 
-    return epiweekly_values
+def daily_to_mmwr_epiweekly(
+    daily_values: ArrayLike, input_data_first_dow: int = 0
+):
+    """
+    Convert daily values to MMWR epidemiological weeks.
+
+    Parameters
+    ----------
+    daily_values : ArrayLike
+        Daily timeseries values.
+    input_data_first_dow : int
+        First day of the week in the input timeseries `daily_values`.
+        Defaults to 0 (Monday).
+
+    Returns
+    -------
+    ArrayLike
+        Data converted to epiweekly values.
+    """
+    return daily_to_weekly(
+        daily_values, input_data_first_dow, week_start_dow=6
+    )

test/test_daily_to_weekly.py

@@ -0,0 +1,95 @@
+# numpydoc ignore=GL08
+
+import jax.numpy as jnp
+import pytest
+
+from pyrenew.convolve import daily_to_weekly, daily_to_mmwr_epiweekly
+
+
+def test_daily_to_weekly_no_offset():
+    """
+    Tests that the function correctly aggregates
+    daily values into weekly totals when there
+    is no offset both input and output start dow on Monday.
+    """
+    daily_values = jnp.arange(1, 15)
+    result = daily_to_weekly(daily_values)
+    expected = jnp.array([28, 77])
+    assert jnp.array_equal(result, expected)
+
+
+def test_daily_to_weekly_with_input_data_offset():
+    """
+    Tests that the function correctly aggregates
+    daily values into weekly totals with dow
+    offset in the input data.
+    """
+    daily_values = jnp.arange(1, 15)
+    result = daily_to_weekly(daily_values, input_data_first_dow=2)
+    expected = jnp.array([63])
+    assert jnp.array_equal(result, expected)
+
+
+def test_daily_to_weekly_with_different_week_start():
+    """
+    Tests aggregation when the desired week start
+    differs from the input data start.
+    """
+    daily_values = jnp.arange(1, 15)
+    result = daily_to_weekly(daily_values, input_data_first_dow=2, week_start_dow=5)
+    expected = jnp.array([49])
+    assert jnp.array_equal(result, expected)
+
+
+def test_daily_to_weekly_incomplete_week():
+    """
+    Tests that the function raises a
+    ValueError when there are
+    insufficient daily values to
+    form a complete week.
+    """
+    daily_values = jnp.arange(1, 5)
+    with pytest.raises(ValueError, match="No complete weekly values available"):
+        daily_to_weekly(daily_values, input_data_first_dow=0)
+
+
+def test_daily_to_weekly_missing_daily_values():
+    """
+    Tests that the function correctly
+    aggregates the available daily values
+    into weekly values when there are
+    fewer daily values than required for
+    complete weekly totals.
+    """
+    daily_values = jnp.arange(1, 10)
+    result = daily_to_weekly(daily_values, input_data_first_dow=0)
+    expected = jnp.array([28])
+    assert jnp.array_equal(result, expected)
+
+
+def test_daily_to_weekly_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_weekly(daily_values, input_data_first_dow=-1)
+
+    with pytest.raises(
+        ValueError, match="First day of aggregated data must be between 0 and 6."
+    ):
+        daily_to_weekly(daily_values, week_start_dow=7)
+
+
+def test_daily_to_mmwr_epiweekly():
+    """
+    Tests aggregation for MMWR epidemiological week.
+    """
+    daily_values = jnp.arange(1, 15)
+    result = daily_to_mmwr_epiweekly(daily_values)
+    expected = jnp.array([70])
+    assert jnp.array_equal(result, expected)