Skip to content

Commit

Permalink
add daily_to_epiweekly
Browse files Browse the repository at this point in the history
  • Loading branch information
sbidari committed Dec 26, 2024
1 parent a63c4ea commit 4d5cb0a
Show file tree
Hide file tree
Showing 2 changed files with 118 additions and 0 deletions.
42 changes: 42 additions & 0 deletions pyrenew/convolve.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
76 changes: 76 additions & 0 deletions test/test_daily_to_epiweekly.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 4d5cb0a

Please sign in to comment.