|
| 1 | +# numpydoc ignore=GL08 |
| 2 | + |
| 3 | +import jax.numpy as jnp |
| 4 | +import pytest |
| 5 | + |
| 6 | +from pyrenew.convolve import daily_to_mmwr_epiweekly, daily_to_weekly |
| 7 | + |
| 8 | + |
| 9 | +def test_daily_to_weekly_no_offset(): |
| 10 | + """ |
| 11 | + Tests that the function correctly aggregates |
| 12 | + daily values into weekly totals when there |
| 13 | + is no offset both input and output start dow on Monday. |
| 14 | + """ |
| 15 | + daily_values = jnp.arange(1, 15) |
| 16 | + result = daily_to_weekly(daily_values) |
| 17 | + expected = jnp.array([28, 77]) |
| 18 | + assert jnp.array_equal(result, expected) |
| 19 | + |
| 20 | + |
| 21 | +def test_daily_to_weekly_with_input_data_offset(): |
| 22 | + """ |
| 23 | + Tests that the function correctly aggregates |
| 24 | + daily values into weekly totals with dow |
| 25 | + offset in the input data. |
| 26 | + """ |
| 27 | + daily_values = jnp.arange(1, 15) |
| 28 | + result = daily_to_weekly(daily_values, input_data_first_dow=2) |
| 29 | + expected = jnp.array([63]) |
| 30 | + assert jnp.array_equal(result, expected) |
| 31 | + |
| 32 | + |
| 33 | +def test_daily_to_weekly_with_different_week_start(): |
| 34 | + """ |
| 35 | + Tests aggregation when the desired week start |
| 36 | + differs from the input data start. |
| 37 | + """ |
| 38 | + daily_values = jnp.arange(1, 15) |
| 39 | + result = daily_to_weekly( |
| 40 | + daily_values, input_data_first_dow=2, week_start_dow=5 |
| 41 | + ) |
| 42 | + expected = jnp.array([49]) |
| 43 | + assert jnp.array_equal(result, expected) |
| 44 | + |
| 45 | + |
| 46 | +def test_daily_to_weekly_incomplete_week(): |
| 47 | + """ |
| 48 | + Tests that the function raises a |
| 49 | + ValueError when there are |
| 50 | + insufficient daily values to |
| 51 | + form a complete week. |
| 52 | + """ |
| 53 | + daily_values = jnp.arange(1, 5) |
| 54 | + with pytest.raises( |
| 55 | + ValueError, match="No complete weekly values available" |
| 56 | + ): |
| 57 | + daily_to_weekly(daily_values, input_data_first_dow=0) |
| 58 | + |
| 59 | + |
| 60 | +def test_daily_to_weekly_missing_daily_values(): |
| 61 | + """ |
| 62 | + Tests that the function correctly |
| 63 | + aggregates the available daily values |
| 64 | + into weekly values when there are |
| 65 | + fewer daily values than required for |
| 66 | + complete weekly totals in the final week. |
| 67 | + """ |
| 68 | + daily_values = jnp.arange(1, 10) |
| 69 | + result = daily_to_weekly(daily_values, input_data_first_dow=0) |
| 70 | + expected = jnp.array([28]) |
| 71 | + assert jnp.array_equal(result, expected) |
| 72 | + |
| 73 | + |
| 74 | +def test_daily_to_weekly_invalid_offset(): |
| 75 | + """ |
| 76 | + Tests that the function raises a |
| 77 | + ValueError when the offset is |
| 78 | + outside the valid range (0-6). |
| 79 | + """ |
| 80 | + daily_values = jnp.arange(1, 15) |
| 81 | + with pytest.raises( |
| 82 | + ValueError, |
| 83 | + match="First day of the week for input timeseries must be between 0 and 6.", |
| 84 | + ): |
| 85 | + daily_to_weekly(daily_values, input_data_first_dow=-1) |
| 86 | + |
| 87 | + with pytest.raises( |
| 88 | + ValueError, |
| 89 | + match="Week start date for output aggregated values must be between 0 and 6.", |
| 90 | + ): |
| 91 | + daily_to_weekly(daily_values, week_start_dow=7) |
| 92 | + |
| 93 | + |
| 94 | +def test_daily_to_mmwr_epiweekly(): |
| 95 | + """ |
| 96 | + Tests aggregation for MMWR epidemiological week. |
| 97 | + """ |
| 98 | + daily_values = jnp.arange(1, 15) |
| 99 | + result = daily_to_mmwr_epiweekly(daily_values) |
| 100 | + expected = jnp.array([70]) |
| 101 | + assert jnp.array_equal(result, expected) |
0 commit comments