in covid19_spread/bar.py [0:0]
def score(self, t, ys):
assert t.size(-1) == ys.size(-1), (t.size(), ys.size())
length = ys.size(-1) - self.window + 1
# beta evolution
beta = self.beta(t, ys)
Z = th.zeros(0).sum()
if self.self_correlation:
ws = F.softplus(self.z)
ws = ws.expand(self.M, self.z.size(1))
# self-correlation
Z = F.conv1d(
F.pad(ys.unsqueeze(0) if ys.ndim == 2 else ys, (self.z.size(1) - 1, 0)),
ws.unsqueeze(1),
groups=self.M,
)
Z = Z.squeeze(0)
Z = Z.div(float(self.z.size(1)))
# cross-correlation
Ys = th.zeros(0).sum(0)
W = th.zeros(1, 1)
if self.cross_correlation:
W = self.metapopulation_weights()
Ys = th.stack(
[
F.pad(ys.narrow(-1, i, length), (self.window - 1, 0))
for i in range(self.window)
]
)
orig_shape = Ys.shape
Ys = Ys.view(-1, Ys.size(-2), Ys.size(-1)) if Ys.ndim == 4 else Ys
Ys = (
th.bmm(W.unsqueeze(0).expand(Ys.size(0), self.M, self.M), Ys)
.view(orig_shape)
.mean(dim=0)
)
with th.no_grad():
self.train_stats = (Z.mean().item(), Ys.mean().item())
if self.features is not None:
Ys = Ys + F.softplus(self.w_feat(self.features))
Ys = beta * (Z + Ys) / self.neighbors
return Ys, beta, W