def test_convolve(batch_shape, m, n, mode):
signal = torch.randn(*batch_shape, m)
kernel = torch.randn(*batch_shape, n)
actual = convolve(signal, kernel, mode)
expected = torch.stack([
torch.tensor(np.convolve(s, k, mode=mode))
for s, k in zip(signal.reshape(-1, m), kernel.reshape(-1, n))
]).reshape(*batch_shape, -1)
assert_close(actual, expected)
def heuristic_init(args, data):
"""Heuristically initialize to a feasible point."""
# Start with a single infection.
S0 = args.population - 1
# Assume 50% <= response rate <= 100%.
S2I = data * min(2.0, (S0 / data.sum()).sqrt())
S_aux = (S0 - S2I.cumsum(-1)).clamp(min=0.5)
# Account for the single initial infection.
S2I[0] += 1
# Assume infection lasts less than a month.
recovery = torch.arange(30.0).div(args.recovery_time).neg().exp()
I_aux = convolve(S2I, recovery)[:len(data)].clamp(min=0.5)
return {
"R0": torch.tensor(2.0),
"rho": torch.tensor(0.5),
"S_aux": S_aux,
"I_aux": I_aux,
}
def test_convolve_shape(m, n, mode):
signal = torch.randn(m)
kernel = torch.randn(n)
actual = convolve(signal, kernel, mode)
expected = np.convolve(signal, kernel, mode=mode)
assert actual.shape == expected.shape