Add weekly_to_daily functionality, move time handling functions to their own module

pyrenew/convolve.py

@@ -217,77 +217,3 @@ 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
-    )

pyrenew/time.py

@@ -0,0 +1,166 @@
+"""
+Helper functions for handling timeseries in Pyrenew
+
+Days of the week in pyrenew are 0-indexed and follow
+ISO standards, so 0 is Monday at 6 is Sunday.
+"""
+
+import jax.numpy as jnp
+from jax.typing import ArrayLike
+
+
+def _validate_dow(
+    day_of_week: int, variable_name: str
+) -> None:  # numpydoc ignore=GL08
+    if not isinstance(day_of_week, int):
+        raise ValueError(
+            "Day-of-week indices must be a integers "
+            "between 0 and 6, inclusive. "
+            f"Got {day_of_week} for {variable_name}, "
+            "which is a "
+            f"{type(day_of_week)}"
+        )
+    if day_of_week < 0 or day_of_week > 6:
+        raise ValueError(
+            "Day-of-week indices must be a integers "
+            "between 0 and 6, inclusive. "
+            f"Got {day_of_week} for {variable_name}."
+        )
+    return None
+
+
+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.
+    """
+
+    _validate_dow(input_data_first_dow, "input_data_first_dow")
+    _validate_dow(week_start_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]
+
+    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
+    )
+
+
+def weekly_to_daily(
+    weekly_values: ArrayLike,
+    output_data_first_dow: int = 0,
+    week_start_dow: int = 0,
+) -> ArrayLike:
+    """
+    Convert a weekly timeseries to a daily
+    timeseries using :func:`jnp.repeat`.
+
+    Parameters
+    ----------
+    weekly_values: ArrayLike
+        Timeseries of weekly values, where
+        (discrete) time is the first dimension of
+        the array (following Pyrenew convention).
+
+    output_data_first_dow: int
+        First day of the week in the output daily timeseries.
+        0-indexed. Default 0 (Monday).
+
+    week_start_dow: int
+        Starting day of the week for ``weekly_ts``,
+        0-indexed. Default 0 (Monday).
+
+    Returns
+    -------
+    ArrayLike
+        The daily timeseries.
+    """
+    _validate_dow(output_data_first_dow, "output_data_first_dow")
+    _validate_dow(week_start_dow, "week_start_dow")
+
+    offset = (output_data_first_dow - week_start_dow) % 7
+    return jnp.repeat(
+        weekly_values,
+        repeats=7,
+        axis=0,
+    )[offset:]
+
+
+def mmwr_epiweekly_to_daily(
+    weekly_values: ArrayLike,
+    output_data_first_dow: int = 0,
+) -> ArrayLike:
+    """
+    Convert an MMWR epiweekly timeseries to a daily
+    timeseries using :func:`weekly_to_daily`.
+
+    Parameters
+    ----------
+    weekly_values: ArrayLike
+        Timeseries of weekly values, where
+        (discrete) time is the first dimension of
+        the array (following Pyrenew convention).
+
+    output_data_first_dow: int
+        First day of the week in the output daily timeseries.
+        0-indexed. Default 0 (Monday).
+
+    Returns
+    -------
+    ArrayLike
+        The daily timeseries.
+    """
+    return weekly_to_daily(
+        weekly_values, output_data_first_dow, week_start_dow=6
+    )

test/test_daily_to_weekly.py → test/test_time.py

@@ -1,9 +1,11 @@
-# numpydoc ignore=GL08
+"""
+Tests for the pyrenew.time module.
+"""
 
 import jax.numpy as jnp
 import pytest
 
-from pyrenew.convolve import daily_to_mmwr_epiweekly, daily_to_weekly
+import pyrenew.time as ptime
 
 
 def test_daily_to_weekly_no_offset():
@@ -13,7 +15,7 @@ def test_daily_to_weekly_no_offset():
     is no offset both input and output start dow on Monday.
     """
     daily_values = jnp.arange(1, 15)
-    result = daily_to_weekly(daily_values)
+    result = ptime.daily_to_weekly(daily_values)
     expected = jnp.array([28, 77])
     assert jnp.array_equal(result, expected)
 
@@ -25,7 +27,7 @@ def test_daily_to_weekly_with_input_data_offset():
     offset in the input data.
     """
     daily_values = jnp.arange(1, 15)
-    result = daily_to_weekly(daily_values, input_data_first_dow=2)
+    result = ptime.daily_to_weekly(daily_values, input_data_first_dow=2)
     expected = jnp.array([63])
     assert jnp.array_equal(result, expected)
 
@@ -36,7 +38,7 @@ def test_daily_to_weekly_with_different_week_start():
     differs from the input data start.
     """
     daily_values = jnp.arange(1, 15)
-    result = daily_to_weekly(
+    result = ptime.daily_to_weekly(
         daily_values, input_data_first_dow=2, week_start_dow=5
     )
     expected = jnp.array([49])
@@ -54,7 +56,7 @@ def test_daily_to_weekly_incomplete_week():
     with pytest.raises(
         ValueError, match="No complete weekly values available"
     ):
-        daily_to_weekly(daily_values, input_data_first_dow=0)
+        ptime.daily_to_weekly(daily_values, input_data_first_dow=0)
 
 
 def test_daily_to_weekly_missing_daily_values():
@@ -66,7 +68,7 @@ def test_daily_to_weekly_missing_daily_values():
     complete weekly totals in the final week.
     """
     daily_values = jnp.arange(1, 10)
-    result = daily_to_weekly(daily_values, input_data_first_dow=0)
+    result = ptime.daily_to_weekly(daily_values, input_data_first_dow=0)
     expected = jnp.array([28])
     assert jnp.array_equal(result, expected)
 
@@ -80,22 +82,22 @@ def test_daily_to_weekly_invalid_offset():
     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.",
+        match="Got -1 for input_data_first_dow",
     ):
-        daily_to_weekly(daily_values, input_data_first_dow=-1)
+        ptime.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.",
+        match="Got 7 for week_start_dow",
     ):
-        daily_to_weekly(daily_values, week_start_dow=7)
+        ptime.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)
+    result = ptime.daily_to_mmwr_epiweekly(daily_values)
     expected = jnp.array([70])
     assert jnp.array_equal(result, expected)