def _f(x, beta, gamma, eta, dt):
lam = beta * (x[..., 1] + eta) / x.sum(-1)
i_f = 1. - (-lam * dt).exp()
r_f = 1. - (-gamma * dt).exp()
return concater(i_f, r_f)
def f_(x, beta, gamma, alpha, rho, sigma, eps, nu):
s = -beta * x[..., 0] * x[..., 1]
r = (1 - alpha) * gamma * x[..., 1]
i = -s - r - alpha * rho * x[..., 1]
d = alpha * rho * x[..., 1]
return concater(s, i, r, d)
def f(x, beta, gamma, delta, alpha, rho, sigma, eps):
s = -beta * x[..., 0] * x[..., 2]
e = -s - delta * x[..., 1]
r = (1 - alpha) * gamma * x[..., 2]
i = delta * x[..., 1] - r - alpha * rho * x[..., 2]
d = alpha * rho * x[..., 2]
return concater(s, e, i, r, d)
def prop_state(x, beta, gamma, eta, dt):
f = _f(x, beta, gamma, eta, dt)
bins = Independent(Binomial(x[..., :-1], f), 1)
samp = bins.sample()
s = x[..., 0] - samp[..., 0]
i = x[..., 1] + samp[..., 0] - samp[..., 1]
r = x[..., 2] + samp[..., 1]
return concater(s, i, r)
def gmvn(x, alpha, sigma):
return concater(sigma, sigma)
def fmvn(x, alpha, sigma):
x1 = alpha * x[0] + x[1] / 3
x2 = x[1]
return concater(x1, x2)
def fmvn(x, a, sigma):
return concater(x[0], x[1])
def fmvn(x, a, sigma):
return concater(x[..., 0], x[..., 1])
def fmvn(x, alpha, sigma):
x1 = alpha * x[..., 0]
x2 = x[..., 1]
return concater(x1, x2)
def g(x, beta, gamma, sigma):
g1 = -sigma * x[..., 0] * x[..., 1]
g3 = torch.zeros_like(g1)
return concater(g1, -g1, g3)
def f(x, beta, gamma, sigma):
f1 = -beta * x[..., 0] * x[..., 1]
f2 = x[..., 1] * (beta * x[..., 0] - gamma)
f3 = x[..., 1] * gamma
return concater(f1, f2, f3)