Test/fix PeriodicBroadcaster

[damonbayer]

We do not currently have any tests for PeriodicBroadcaster, and I don't think it works how we want it to.

It seems like there is one too many inputs into these functions, leading to strange behavior. I think everything could be done without a period_size argument.

from pyrenew.arrayutils import PeriodicBroadcaster
import jax.numpy as jnp


base_array = jnp.array([1, 2, 3])

tile_broadcaster = PeriodicBroadcaster(
    offset=0, period_size=7, broadcast_type="tile"
)

tile_broadcaster(base_array, 10)
# Array([1, 2, 3, 1, 2, 3], dtype=int32)

repeat_broadcaster = PeriodicBroadcaster(
    offset=0, period_size=7, broadcast_type="repeat"
)

repeat_broadcaster(base_array, 10)
# Array([1, 1, 1, 1, 1, 1, 1, 2, 2, 2], dtype=int32)

[gvegayon]

There's one test using the PeriodicBroadcaster implicitly: https://github.com/CDCgov/multisignal-epi-inference/blob/64e0ed832166d6c5ec2d0ca817c3cfdd1eeb0e0d/model/src/test/test_periodiceffect.py. I'll take a look at it.

[gvegayon]

See if this works better:

dow_effects[pred_hosp_day_index] pred_hosp_day_index is an array of the same length as pred_hosp that indexes the days of the week (entries are 0,...6)

[damonbayer]

After discussing in standup, the plan is to

• Use indexing instead of creating new arrays
• Probably not have a common interface for the "tile" and "repeat" operations.