diff --git a/pyrenew/convolve.py b/pyrenew/convolve.py index 5ab8ae44..6b8c2c0c 100755 --- a/pyrenew/convolve.py +++ b/pyrenew/convolve.py @@ -217,3 +217,77 @@ def compute_delay_ascertained_incidence( mode="valid", ) return delay_obs_incidence + + +def daily_to_weekly( + daily_values: ArrayLike, + input_data_first_dow: int = 0, + week_start_dow: int = 0, +) -> ArrayLike: + """ + Aggregate daily values (e.g. + incident hospital admissions) into weekly total values. + + Parameters + ---------- + daily_values : ArrayLike + Daily timeseries values (e.g. incident infections or incident ed visits). + 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 `input_data_first_dow` does not match `week_start_dow`, the incomplete first + week is ignored and weekly values starting + from the second week are 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 weekly values starting + with the first full week available. + """ + if input_data_first_dow < 0 or input_data_first_dow > 6: + raise ValueError( + "First day of the week for input timeseries must be between 0 and 6." + ) + + if week_start_dow < 0 or week_start_dow > 6: + raise ValueError( + "Week start date for output aggregated values must be between 0 and 6." + ) + + 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] + + return weekly_values + + +def daily_to_mmwr_epiweekly( + daily_values: ArrayLike, input_data_first_dow: int = 0 +) -> ArrayLike: + """ + 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 + ) diff --git a/test/test_daily_to_weekly.py b/test/test_daily_to_weekly.py new file mode 100644 index 00000000..a4fb403a --- /dev/null +++ b/test/test_daily_to_weekly.py @@ -0,0 +1,101 @@ +# numpydoc ignore=GL08 + +import jax.numpy as jnp +import pytest + +from pyrenew.convolve import daily_to_mmwr_epiweekly, daily_to_weekly + + +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 in the final week. + """ + 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 for input timeseries must be between 0 and 6.", + ): + daily_to_weekly(daily_values, input_data_first_dow=-1) + + with pytest.raises( + ValueError, + match="Week start date for output aggregated values 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)