def sample_y(dist_y, theta, y, sigma_obs=None):
if not sigma_obs:
if dist_y == 'gamma':
sigma_obs = numpyro.sample('sigma_obs', dist.Exponential(1))
else:
sigma_obs = numpyro.sample('sigma_obs', dist.HalfNormal(1))
if dist_y == 'student':
numpyro.sample('y', dist.StudentT(numpyro.sample('nu_y', dist.Gamma(1, .1)), theta, sigma_obs), obs=y)
elif dist_y == 'normal':
numpyro.sample('y', dist.Normal(theta, sigma_obs), obs=y)
elif dist_y == 'lognormal':
numpyro.sample('y', dist.LogNormal(theta, sigma_obs), obs=y)
elif dist_y == 'gamma':
numpyro.sample('y', dist.Gamma(jnp.exp(theta), sigma_obs), obs=y)
elif dist_y == 'gamma_raw':
numpyro.sample('y', dist.Gamma(theta, sigma_obs), obs=y)
elif dist_y == 'poisson':
numpyro.sample('y', dist.Poisson(theta), obs=y)
elif dist_y == 'exponential':
numpyro.sample('y', dist.Exponential(jnp.exp(theta)), obs=y)
elif dist_y == 'exponential_raw':
numpyro.sample('y', dist.Exponential(theta), obs=y)
elif dist_y == 'uniform':
numpyro.sample('y', dist.Uniform(0, 1), obs=y)
else:
raise NotImplementedError
def model_bpinn(p,
t,
Y,
F,
data_type,
D_H,
u_sigma=None,
f_sigma=None,
sigma_w=1):
m = 0.15
d = 0.15
B = 0.2
D_X, D_Y = 2, 1
# sample first layer
w1 = numpyro.sample(
"w1", dist.Normal(jnp.zeros((D_X, D_H)), sigma_w * jnp.ones(
(D_X, D_H))))
b1 = numpyro.sample(
"b1", dist.Normal(jnp.zeros((D_H, 1)), sigma_w * jnp.ones((D_H, 1))))
# sample second layer
w2 = numpyro.sample(
"w2", dist.Normal(jnp.zeros((D_H, D_H)), sigma_w * jnp.ones(
(D_H, D_H))))
b2 = numpyro.sample(
"b2", dist.Normal(jnp.zeros((D_H, 1)), sigma_w * jnp.ones((D_H, 1))))
# sample final layer
w3 = numpyro.sample(
"w3", dist.Normal(jnp.zeros((D_H, D_Y)), sigma_w * jnp.ones(
(D_H, D_Y))))
b3 = numpyro.sample(
"b3", dist.Normal(jnp.zeros((D_Y, 1)), sigma_w * jnp.ones((D_Y, 1))))
u_mu, dudt = mu_grad(p, t, w1, b1, w2, b2, w3, b3)
dudtt = second_grad(p, t, w1, b1, w2, b2, w3, b3)
# prior on the observation noise
if u_sigma is None:
prec_u = numpyro.sample("prec_u", dist.Gamma(3.0, 1.0))
u_sigma = 1.0 / jnp.sqrt(prec_u)
if f_sigma is None:
prec_f = numpyro.sample("prec_f", dist.Gamma(3.0, 1.0))
f_sigma = 1.0 / jnp.sqrt(prec_f)
# observe data
with numpyro.plate('observations', p.shape[0]):
with handlers.mask(mask=data_type):
u_hat = numpyro.sample("Y", dist.Normal(u_mu, u_sigma), obs=Y)
f_mu = m * dudtt + d * dudt + B * jnp.sin(
u_mu) - p # Forcing physics-term, always=0
f_hat = numpyro.sample("F", dist.Normal(f_mu, f_sigma), obs=F)
return u_mu, f_mu
def __init__(self, params):
'''
Pull from passed parameters and initialize numpyro samplers.
'''
# if params.R0 is not None:
# self.R0 = params.R0
# self.R0 = ny.sample('r0', dist.Normal(*self.R0)) # NB: Unused right now
# if params.INCUBATION_PERIOD is not None:
# self.INCUBATION_PERIOD = params.INCUBATION_PERIOD
# self.INCUBATION_PERIOD = ny.sample('iP', dist.Gamma(*self.INCUBATION_PERIOD)) # NB: unused right now
if params.PROPORTION_PRESYMPTOMATIC_TRANSMISSION is not None:
self.PROPORTION_PRESYMPTOMATIC_TRANSMISSION = params.PROPORTION_PRESYMPTOMATIC_TRANSMISSION
self.PROPORTION_PRESYMPTOMATIC_TRANSMISSION = ny.sample(
'pP',
dist.TruncatedNormal(*self.PROPORTION_PRESYMPTOMATIC_TRANSMISSION))
if params.SYMPTOMATIC is not None:
self.SYMPTOMATIC = params.SYMPTOMATIC
self.SYMPTOMATIC = ny.sample('rS', dist.Beta(*self.SYMPTOMATIC))
if params.PRESYMPTOMATIC_CONTAGIOUS_PERIOD is not None:
self.PRESYMPTOMATIC_CONTAGIOUS_PERIOD = params.PRESYMPTOMATIC_CONTAGIOUS_PERIOD
self.PRESYMPTOMATIC_CONTAGIOUS_PERIOD = ny.sample(
'cP', dist.Gamma(*self.PRESYMPTOMATIC_CONTAGIOUS_PERIOD))
if params.ASYMPTOMATIC_CONTAGIOUS_PERIOD is not None:
self.ASYMPTOMATIC_CONTAGIOUS_PERIOD = params.ASYMPTOMATIC_CONTAGIOUS_PERIOD
self.ASYMPTOMATIC_CONTAGIOUS_PERIOD = ny.sample(
'cA', dist.Gamma(*self.ASYMPTOMATIC_CONTAGIOUS_PERIOD))
if params.SYMPTOMATIC_CONTAGIOUS_PERIOD is not None:
self.SYMPTOMATIC_CONTAGIOUS_PERIOD = params.SYMPTOMATIC_CONTAGIOUS_PERIOD
self.SYMPTOMATIC_CONTAGIOUS_PERIOD = ny.sample(
'cS', dist.Gamma(*self.SYMPTOMATIC_CONTAGIOUS_PERIOD))
if params.SYMPTOM_TO_HOSP_PERIOD is not None:
self.SYMPTOM_TO_HOSP_PERIOD = params.SYMPTOM_TO_HOSP_PERIOD
self.SYMPTOM_TO_HOSP_PERIOD = ny.sample(
'pH', dist.TruncatedNormal(*self.SYMPTOM_TO_HOSP_PERIOD))
if params.HOSP_DEATH_PERIOD is not None:
self.HOSP_DEATH_PERIOD = params.HOSP_DEATH_PERIOD
self.HOSP_DEATH_PERIOD = ny.sample('hD',
dist.Gamma(*self.HOSP_DEATH_PERIOD))
if params.HOSP_RECOVERY_PERIOD is not None:
self.HOSP_RECOVERY_PERIOD = params.HOSP_RECOVERY_PERIOD
self.HOSP_RECOVERY_PERIOD = ny.sample(
'hR', dist.Gamma(*self.HOSP_RECOVERY_PERIOD))
# Derived variables
self.EXPOSED_PERIOD = self.INCUBATION_PERIOD - self.PRESYMPTOMATIC_CONTAGIOUS_PERIOD
self.RECOVERY_PERIOD = self.INCUBATION_PERIOD + np.where(
self.SYMPTOMATIC > 0, self.SYMPTOMATIC_CONTAGIOUS_PERIOD,
self.ASYMPTOMATIC_CONTAGIOUS_PERIOD)
def GP(X, y):
X = numpyro.deterministic("X", X)
# Set informative priors on kernel hyperparameters.
η = numpyro.sample("variance", dist.HalfCauchy(scale=5.0))
ℓ = numpyro.sample("length_scale", dist.Gamma(2.0, 1.0))
σ = numpyro.sample("obs_noise", dist.HalfCauchy(scale=5.0))
# Compute kernel
K = rbf_kernel(X, X, η, ℓ)
K = add_to_diagonal(K, σ)
K = add_to_diagonal(K, wandb.config.jitter)
# cholesky decomposition
Lff = numpyro.deterministic("Lff", cholesky(K, lower=True))
# Sample y according to the standard gaussian process formula
return numpyro.sample(
"y",
dist.MultivariateNormal(loc=jnp.zeros(X.shape[0]),
scale_tril=Lff).expand_by(
y.shape[:-1]) # for multioutput scenarios
.to_event(y.ndim - 1),
obs=y,
)
def model(X, Y, D_H, D_Y=1):
N, D_X = X.shape
# sample first layer (we put unit normal priors on all weights)
w1 = numpyro.sample(
"w1", dist.Normal(jnp.zeros((D_X, D_H)), jnp.ones((D_X, D_H))))
assert w1.shape == (D_X, D_H)
z1 = nonlin(jnp.matmul(X, w1)) # <= first layer of activations
assert z1.shape == (N, D_H)
# sample second layer
w2 = numpyro.sample(
"w2", dist.Normal(jnp.zeros((D_H, D_H)), jnp.ones((D_H, D_H))))
assert w2.shape == (D_H, D_H)
z2 = nonlin(jnp.matmul(z1, w2)) # <= second layer of activations
assert z2.shape == (N, D_H)
# sample final layer of weights and neural network output
w3 = numpyro.sample(
"w3", dist.Normal(jnp.zeros((D_H, D_Y)), jnp.ones((D_H, D_Y))))
assert w3.shape == (D_H, D_Y)
z3 = jnp.matmul(z2, w3) # <= output of the neural network
assert z3.shape == (N, D_Y)
if Y is not None:
assert z3.shape == Y.shape
# we put a prior on the observation noise
prec_obs = numpyro.sample("prec_obs", dist.Gamma(3.0, 1.0))
sigma_obs = 1.0 / jnp.sqrt(prec_obs)
# observe data
with numpyro.plate("data", N):
# note we use to_event(1) because each observation has shape (1,)
numpyro.sample("Y", dist.Normal(z3, sigma_obs).to_event(1), obs=Y)
def _model(self, X, Y, D_H, train=True):
D_X, D_Y = X.shape[1], 1#self.num_bins
# if train:
# targets = Y[:,1]
# Sample first layer (we put unit normal priors on all weights)
w1 = numpyro.sample("w1", dist.Normal(np.zeros((D_X, D_H)), np.ones((D_X, D_H)))) # D_X D_H
z1 = self._nonlin(np.matmul(X, w1)) # N D_H <= first layer of activations
# # sample second layer
# w2 = numpyro.sample("w2", dist.Normal(np.zeros((D_H, D_H)), np.ones((D_H, D_H)))) # D_H D_H
# z2 = self._nonlin(np.matmul(z1, w2)) # N D_H <= second layer of activations
# sample final layer of weights and neural network output
w3 = numpyro.sample("w3", dist.Normal(np.zeros((D_H, D_Y)), np.ones((D_H, D_Y)))) # D_H D_Y
z3 = np.matmul(z1, w3) # N D_Y <= output of the neural network
# if train: # Isolate z3 to only the relevant predictions (which action was selected) from the neural network
# filter_z3 = np.array(Y[:,0])
# red_z3 = np.take(z3,filter_z3)
# we put a prior on the observation noise
prec_obs = numpyro.sample("prec_obs", dist.Gamma(10.0, 1.0))
sigma_obs = 1.0 / np.sqrt(prec_obs)
# # observe data
# if train:
# numpyro.sample("Y", dist.Normal(red_z3, sigma_obs), obs=1.0*targets)
# else:
numpyro.sample("Y", dist.Normal(z3, sigma_obs), obs=Y)
def bnn_model(self, x, y, num_hidden):
d_x = x.shape[1] # number of features in x
d_y = 1 #
# d_x, d_y = x_train.shape[1], 1
# sample first layer (we put unit normal priors on all weights)
w1 = numpyro.sample("w1",
dist.Normal(np.zeros((d_x, num_hidden)),
np.ones((d_x, num_hidden)))) # d_x d_h
b1 = numpyro.sample("b1",
dist.Normal(0, 1),
sample_shape=(num_hidden, ))
z1 = nonlin(b1 + np.matmul(x, w1))
# z1 = nonlin(np.matmul(x_train, w1)) # N d_h <= first layer of activations
# sample second layer
w2 = numpyro.sample("w2",
dist.Normal(np.zeros((num_hidden, d_y)),
np.ones((num_hidden, d_y)))) # d_h d_h
b2 = numpyro.sample("b2", dist.Normal(0, 1), sample_shape=(d_y, ))
z2 = b2 + np.matmul(z1, w2) # N d_h <= second layer of activations
# sample final layer of weights and neural network output
# w3 = numpyro.sample("w3", dist.Normal(np.zeros((d_h, d_y)), np.ones((d_h, d_y)))) # d_h D_Y
# z3 = np.matmul(z2, w3) # N D_Y <= output of the neural network
# we put a prior on the observation noise
prec_obs = numpyro.sample("prec_obs", dist.Gamma(3.0, 1.0))
sigma_obs = 1.0 / np.sqrt(prec_obs)
# observe data
numpyro.sample("y", dist.Normal(z2, sigma_obs), obs=y)
def model(X, Y, D_H):
D_X, D_Y = X.shape[1], 1
# sample first layer (we put unit normal priors on all weights)
w1 = numpyro.sample("w1",
dist.Normal(np.zeros((D_X, D_H)), np.ones(
(D_X, D_H)))) # D_X D_H
z1 = nonlin(np.matmul(X, w1)) # N D_H <= first layer of activations
# sample second layer
w2 = numpyro.sample("w2",
dist.Normal(np.zeros((D_H, D_H)), np.ones(
(D_H, D_H)))) # D_H D_H
z2 = nonlin(np.matmul(z1, w2)) # N D_H <= second layer of activations
# sample final layer of weights and neural network output
w3 = numpyro.sample("w3",
dist.Normal(np.zeros((D_H, D_Y)), np.ones(
(D_H, D_Y)))) # D_H D_Y
z3 = np.matmul(z2, w3) # N D_Y <= output of the neural network
# we put a prior on the observation noise
prec_obs = numpyro.sample("prec_obs", dist.Gamma(3.0, 1.0))
sigma_obs = 1.0 / np.sqrt(prec_obs)
# observe data
numpyro.sample("Y", dist.Normal(z3, sigma_obs), obs=Y)
def model(
x: Optional[jnp.ndarray] = None,
seq_len: int = 0,
batch: int = 0,
x_dim: int = 1,
z_dim: int = 3,
future_steps: int = 0,
) -> None:
"""Hierarchical Kalman filter."""
if x is not None:
seq_len, batch, x_dim = x.shape
trans_mu = numpyro.sample(
"trans_mu",
dist.Uniform(-jnp.ones((z_dim, z_dim)), jnp.ones((z_dim, z_dim))))
trans_var = numpyro.sample(
"trans_var",
dist.LogNormal(-20 * jnp.ones((z_dim, z_dim)), 20 * jnp.ones(
(z_dim, z_dim))))
with numpyro.plate("batch", batch, dim=-3):
trans = numpyro.sample(
"trans",
dist.Normal(trans_mu, trans_var).expand((1, z_dim, z_dim)))
emit = numpyro.sample(
"emit", dist.Normal(jnp.zeros((z_dim, x_dim)), jnp.ones(
(z_dim, x_dim))))
z_std = numpyro.sample("z_std", dist.Gamma(jnp.ones(z_dim),
jnp.ones(z_dim)))
x_std = numpyro.sample("x_std", dist.Gamma(jnp.ones(x_dim),
jnp.ones(x_dim)))
def transition_fn(
carry: Tuple[jnp.ndarray],
t: jnp.ndarray) -> Tuple[Tuple[jnp.ndarray], jnp.ndarray]:
z_prev, *_ = carry
index = jnp.arange(batch)
z = numpyro.sample(
"z", dist.Normal(jnp.matmul(z_prev, trans)[index, index], z_std))
numpyro.sample("x", dist.Normal(jnp.matmul(z, emit), x_std))
return (z, ), None
z_init = jnp.zeros((batch, z_dim))
with numpyro.handlers.condition(data={"x": x}):
scan(transition_fn, (z_init, ), jnp.arange(seq_len + future_steps))
def __init__(self, mu, tau, validate_args=None):
self.mu, self.tau = promote_shapes(mu, tau)
# converts mean var parametrisation to r and p
self.r = tau
self.var = mu + 1 / self.r * mu**2
self.p = (self.var - mu) / self.var
self._gamma = dist.Gamma(self.r, (1 - self.p) / self.p)
super(NegativeBinomial, self).__init__(self._gamma.batch_shape,
validate_args=validate_args)
def test_model_with_lift_handler():
def model(data):
c = numpyro.param("c", jnp.array(1.), constraint=dist.constraints.positive)
x = numpyro.sample("x", dist.LogNormal(c, 1.), obs=data)
return x
nuts_kernel = NUTS(numpyro.handlers.lift(model, prior={"c": dist.Gamma(0.01, 0.01)}))
mcmc = MCMC(nuts_kernel, num_warmup=10, num_samples=10)
mcmc.run(random.PRNGKey(1), jnp.exp(random.normal(random.PRNGKey(0), (1000,))))
def model(
x: Optional[jnp.ndarray] = None,
seq_len: int = 0,
batch: int = 0,
x_dim: int = 1,
future_steps: int = 0,
z_dim: int = 3,
) -> None:
"""Simple Kalman filter model (random walk).
Args:
x: **Batch-first** data, `shape = (seq_len, batch, data_dim)`.
future_steps: Forecasting time steps.
batch: Batch size for prior sampling.
x_dim: Dimension of data for prior sampling.
"""
if x is not None:
seq_len, batch, x_dim = x.shape
trans = numpyro.sample(
"trans",
dist.Normal(jnp.zeros((z_dim, z_dim)), jnp.ones((z_dim, z_dim))))
emit = numpyro.sample(
"emit", dist.Normal(jnp.zeros((z_dim, x_dim)), jnp.ones(
(z_dim, x_dim))))
z_std = numpyro.sample("z_std", dist.Gamma(jnp.ones(z_dim),
jnp.ones(z_dim)))
x_std = numpyro.sample("x_std", dist.Gamma(jnp.ones(x_dim),
jnp.ones(x_dim)))
def transition_fn(
carry: Tuple[jnp.ndarray],
t: jnp.ndarray) -> Tuple[Tuple[jnp.ndarray], jnp.ndarray]:
z_prev, *_ = carry
z = numpyro.sample("z", dist.Normal(jnp.matmul(z_prev, trans), z_std))
numpyro.sample("x", dist.Normal(jnp.matmul(z, emit), x_std))
return (z, ), None
z_init = jnp.zeros((batch, z_dim))
with numpyro.handlers.condition(data={"x": x}):
scan(transition_fn, (z_init, ), jnp.arange(seq_len + future_steps))
def bayesian_regression(x: np.ndarray, y: Optional[np.ndarray] = None) -> None:
batch, x_dim = jnp.shape(x)
theta = numpyro.sample(
"theta", dist.Normal(jnp.zeros(x_dim),
jnp.ones(x_dim) * 100))
sigma = numpyro.sample("sigma", dist.Gamma(1.0, 1.0))
x_mu = numpyro.sample(
"x_mu", dist.Normal(jnp.nanmean(x, axis=0), jnp.nanmean(x, axis=0)))
x_std = numpyro.sample("x_std", dist.Gamma(1.0, 1.0))
with numpyro.plate("batch", batch, dim=-2):
mask = ~np.isnan(x)
numpyro.sample("x", dist.Normal(x_mu, x_std).mask(mask), obs=x)
index = (~mask).astype(int).nonzero()
x_sample = numpyro.sample("x_sample", dist.Normal(x_mu, x_std))
x_filled = ops.index_update(x, index, x_sample[index])
numpyro.sample("y", dist.Normal(jnp.matmul(x_filled, theta), sigma), obs=y)
def model(self, data: xr.Dataset):
data = data.transpose("item", "feature")
X, Y = data.X.values, data.Y.values
alpha = numpyro.sample("alpha", dist.Normal(scale=self.alpha_scale))
with numpyro.plate("K", self.k):
beta = numpyro.sample("beta", dist.Normal(self.beta_loc, self.beta_scale))
nu = numpyro.sample("nu", dist.Gamma(2.0, 0.1))
sigma = numpyro.sample("sigma", dist.Exponential(1 / self.sigma_mean))
mu = alpha + X @ beta
with numpyro.plate("N", self.n):
numpyro.sample("Y", dist.StudentT(nu, mu, sigma), obs=Y)
def _(d):
gamma_dist = dist.Gamma(d.concentration)
def transform_fn(q):
# NB: icdf is not available yet for Gamma distribution
# so this will raise an NotImplementedError for now.
# We will need scipy.special.gammaincinv, which is not available yet in JAX
# see issue: https://github.com/google/jax/issues/5350
# TODO: consider wrap jaxns GammaPrior transform implementation
gammas = uniform_reparam_transform(gamma_dist)(q)
return gammas / gammas.sum(-1, keepdims=True)
return transform_fn
def dp_sb_gmm(y, num_components):
# Cosntants
N = y.shape[0]
K = num_components
# Priors
# NOTE: In numpyro, the Gamma distribution is parameterized with shape and rate.
# Hence, Gamma(shape, rate) => mean = shape/rate
alpha = numpyro.sample('alpha', dist.Gamma(1, 10))
with numpyro.plate('mixture_weights', K - 1):
v = numpyro.sample('v', dist.Beta(1, alpha, K - 1))
eta = stickbreak(v)
with numpyro.plate('components', K):
mu = numpyro.sample('mu', dist.Normal(0., 3.))
sigma = numpyro.sample('sigma', dist.Gamma(1, 10))
with numpyro.plate('data', N):
numpyro.sample('obs', NormalMixture(mu[None, :] , sigma[None, :], eta[None, :]),
obs=y[:, None])
def gammadyn_mixture_model(beliefs, y, mask):
M, T, _ = beliefs[0].shape
tau = .5
weights = npyro.sample('weights', dist.Dirichlet(tau * jnp.ones(M)))
assert weights.shape == (M, )
gamma = npyro.sample('gamma', dist.InverseGamma(2., 2.))
mu = jnp.log(jnp.exp(gamma) - 1)
p = npyro.sample('p', dist.Dirichlet(jnp.ones(3)))
p0 = npyro.deterministic(
'p0',
jnp.concatenate([
p[..., :1] / 2, p[..., :1] / 2 + p[..., 1:2], p[..., 2:] / 2,
p[..., 2:] / 2
], -1))
scale = npyro.sample('scale', dist.Gamma(1., 1.))
rho = npyro.sample('rho', dist.Beta(1., 2.))
sigma = jnp.sqrt(-(1 - rho**2) / (2 * jnp.log(rho))) * scale
U = jnp.log(p0)
def transition_fn(carry, t):
x_prev = carry
gamma_dyn = npyro.deterministic('gamma_dyn', nn.softplus(mu + x_prev))
logs = logits((beliefs[0][:, t], beliefs[1][:, t]),
jnp.expand_dims(gamma_dyn, -1), jnp.expand_dims(U, -2))
mixing_dist = dist.CategoricalProbs(weights)
component_dist = dist.CategoricalLogits(logs).mask(mask[t])
npyro.sample('y', dist.MixtureSameFamily(mixing_dist, component_dist))
with npyro.handlers.reparam(
config={"x_next": npyro.infer.reparam.TransformReparam()}):
affine = dist.transforms.AffineTransform(rho * x_prev, sigma)
x_next = npyro.sample(
'x_next',
dist.TransformedDistribution(dist.Normal(0., 1.), affine))
return (x_next), None
x0 = jnp.zeros(1)
with npyro.handlers.condition(data={"y": y}):
scan(transition_fn, (x0), jnp.arange(T))
def model(self, t, X):
noise = sample('noise', dist.LogNormal(0.0, 1.0), sample_shape=(self.D,))
hyp = sample('hyp', dist.Gamma(1.0, 0.5), sample_shape=(self.D,))
W = sample('W', dist.LogNormal(0.0, 1.0), sample_shape=(self.D,))
J0 = sample('J0', dist.Uniform(1.0, 10.0)) # 2.5
k1 = sample('k1', dist.Uniform(80., 120.0)) # 100.
k2 = sample('k2', dist.Uniform(1., 10.0)) # 6.
k3 = sample('k3', dist.Uniform(2., 20.0)) # 16.
k4 = sample('k4', dist.Uniform(80., 120.0)) # 100.
k5 = sample('k5', dist.Uniform(0.1, 2.0)) # 1.28
k6 = sample('k6', dist.Uniform(2., 20.0)) # 12.
k = sample('k', dist.Uniform(0.1, 2.0)) # 1.8
ka = sample('ka', dist.Uniform(2., 20.0)) # 13.
q = sample('q', dist.Uniform(1., 10.0)) # 4.
KI = sample('KI', dist.Uniform(0.1, 2.0)) # 0.52
phi = sample('phi', dist.Uniform(0.05, 1.0)) # 0.1
Np = sample('Np', dist.Uniform(0.1, 2.0)) # 1.
A = sample('A', dist.Uniform(1., 10.0)) #4.
IC = sample('IC', dist.Uniform(0, 1))
# compute kernel
K_11 = W[0]*self.RBF(self.t_t[0], self.t_t[0], hyp[0]) + np.eye(self.N[0])*(noise[0] + self.jitter)
K_22 = W[1]*self.RBF(self.t_t[1], self.t_t[1], hyp[1]) + np.eye(self.N[1])*(noise[1] + self.jitter)
K_33 = W[2]*self.RBF(self.t_t[2], self.t_t[2], hyp[2]) + np.eye(self.N[2])*(noise[2] + self.jitter)
K = np.concatenate([np.concatenate([K_11, np.zeros((self.N[0], self.N[1])), np.zeros((self.N[0], self.N[2]))], axis = 1),
np.concatenate([np.zeros((self.N[1], self.N[0])), K_22, np.zeros((self.N[1], self.N[2]))], axis = 1),
np.concatenate([np.zeros((self.N[2], self.N[0])), np.zeros((self.N[2], self.N[1])), K_33], axis = 1)], axis = 0)
# compute mean
x0 = np.array([0.5, 1.9, 0.18, 0.15, IC, 0.1, 0.064])
mut = odeint(self.dxdt, x0, self.t.flatten(), J0, k1, k2, k3, k4, k5, k6, k, ka, q, KI, phi, Np, A)
mu1 = mut[self.i_t[0],ind[0]] / self.max_X[0]
mu2 = mut[self.i_t[1],ind[1]] / self.max_X[1]
mu3 = mut[self.i_t[2],ind[2]] / self.max_X[2]
mu = np.concatenate((mu1,mu2,mu3),axis=0)
# Concat data
mu = mu.flatten('F')
X = np.concatenate((self.X[0],self.X[1],self.X[2]),axis=0)
X = X.flatten('F')
# sample X according to the standard gaussian process formula
sample("X", dist.MultivariateNormal(loc=mu, covariance_matrix=K), obs=X)
def numpyro_model(X, y):
if inference == "map" or "vi_mf" or "vi_full":
# Set priors on hyperparameters.
η = numpyro.sample("variance", dist.HalfCauchy(scale=5.0))
ℓ = numpyro.sample(
"length_scale", dist.Gamma(2.0, 1.0), sample_shape=(n_features,)
)
σ = numpyro.sample("obs_noise", dist.HalfCauchy(scale=5.0))
elif inference == "mll":
# set params and constraints on hyperparams
η = numpyro.param(
"variance", init_value=1.0, constraints=dist.constraints.positive
)
ℓ = numpyro.param(
"length_scale",
init_value=jnp.ones(n_features),
constraints=dist.constraints.positive,
)
σ = numpyro.param(
"obs_noise", init_value=0.01, onstraints=dist.constraints.positive
)
else:
raise ValueError(f"Unrecognized inference scheme: {inference}")
x_u = numpyro.param("x_u", init_value=X_u_init)
# Kernel Function
rbf_kernel = RBF(variance=η, length_scale=ℓ)
# GP Model
gp_model = SGPVFE(
X=X,
X_u=x_u,
y=y,
mean=zero_mean,
kernel=rbf_kernel,
obs_noise=σ,
jitter=jitter,
)
# Sample y according SGP
return gp_model.to_numpyro(y=y)
def model(x=None, num_obs_total=None):
"""
Args:
x (jax.numpy.array): Array holding all features of a single data instance.
num_obs_total (int): Number of total instances in the data set.
Samples:
site `x` similar to input x; array holding all features of a single data instance.
"""
assert x is None or len(np.shape(x)) == 2
if x is None:
N = 1
else:
N = np.shape(x)[0]
if num_obs_total is None:
num_obs_total = N
assert isinstance(num_obs_total, int) and num_obs_total > 0
assert N <= num_obs_total
leuko_mus = sample('Leukocytes_mus', dist.Normal(0., 1.))
leuko_sig = sample('Leukocytes_sig', dist.Gamma(2., 2.))
leuko_dist = dist.Normal(leuko_mus, leuko_sig)
leuko_na_prob = sample('Leukocytes_na_prob', dist.Beta(1., 1.))
leuko_na_dist = NAModel(leuko_dist, leuko_na_prob)
rhino_test_logit = sample('Rhinovirus/Enterovirus_logit',
dist.Normal(0., 1.))
rhino_test_dist = dist.Bernoulli(logits=rhino_test_logit)
rhino_test_na_prob = sample('Rhinovirus/Enterovirus_na_prob',
dist.Beta(1., 1.))
rhino_test_na_dist = NAModel(rhino_test_dist, rhino_test_na_prob)
with plate("batch", num_obs_total, N):
x_leuko = get_feature(x, 0)
x_rhino = get_feature(x, 1)
y_leuko = sample('Leukocytes', leuko_na_dist, obs=x_leuko)
y_rhino = sample('Rhinovirus/Enterovirus',
rhino_test_na_dist,
obs=x_rhino)
y = sample_combined(y_leuko, y_rhino)
def model_bnn(X, Y, D_H, sigma_obs=None, sigma_w=1):
D_X, D_Y = X.shape[1], 1
# sample first layer (we put unit normal priors on all weights)
w1 = numpyro.sample("w1",
dist.Normal(jnp.zeros((D_X, D_H)),
sigma_w * jnp.ones((D_X, D_H)))) # D_X D_H
b1 = numpyro.sample("b1",
dist.Normal(jnp.zeros((D_H, 1)),
sigma_w * jnp.ones((D_H, 1)))) # D_H 1
z1 = nonlin(jnp.matmul(X, w1) +
jnp.transpose(b1)) # N D_H <= first layer of activations
# sample second layer
w2 = numpyro.sample("w2",
dist.Normal(jnp.zeros((D_H, D_H)),
sigma_w * jnp.ones((D_H, D_H)))) # D_H D_H
b2 = numpyro.sample("b2",
dist.Normal(jnp.zeros((D_H, 1)),
sigma_w * jnp.ones((D_H, 1)))) # D_H 1
z2 = nonlin(jnp.matmul(z1, w2) +
jnp.transpose(b2)) # N D_H <= second layer of activations
# sample final layer of weights and neural network output
w3 = numpyro.sample("w3",
dist.Normal(jnp.zeros((D_H, D_Y)),
sigma_w * jnp.ones((D_H, D_Y)))) # D_H D_Y
b3 = numpyro.sample("b3",
dist.Normal(jnp.zeros((D_Y, 1)),
sigma_w * jnp.ones((D_Y, 1)))) # D_H 1
z3 = jnp.matmul(z2, w3) + jnp.transpose(
b3) # N D_Y <= output of the neural network
# prior on the observation noise
if sigma_obs is None:
prec_obs = numpyro.sample("prec_obs", dist.Gamma(1000000.0, 1.0))
sigma_obs = 1.0 / jnp.sqrt(prec_obs)
# observe data
with numpyro.plate('observations', D_X):
numpyro.sample("Y", dist.Normal(z3, sigma_obs), obs=Y)
def generative_model(beliefs, y=None, mask=True):
# generative model
T, N = beliefs[0].shape[:2]
with npyro.plate('N', N):
gamma = npyro.sample('gamma', dist.Gamma(20., 2.))
td = TransformedDistribution(
dist.Normal(jnp.array([-1., .7]), jnp.array([1.,
0.2])).to_event(1),
transforms.OrderedTransform())
lams = npyro.sample('lams', td)
U = jnp.pad(lams, ((0, 0), (0, 2)), 'constant', constant_values=(0, ))
with npyro.plate('T', T):
logs = npyro.deterministic(
'logits',
logits(beliefs, jnp.expand_dims(gamma, -1),
jnp.expand_dims(U, -2)))
npyro.sample('y', dist.CategoricalLogits(logs).mask(mask), obs=y)
def pca_regression(
x: Optional[np.ndarray] = None,
y: Optional[np.ndarray] = None,
batch: int = 100,
x_dim: int = 50,
z_dim: int = 5,
) -> None:
if x is not None:
batch, x_dim = jnp.shape(x)
x_mu = np.nanmean(x, axis=0)
x_std = np.nanstd(x, axis=0)
mask = ~np.isnan(x)
else:
x_mu = jnp.zeros(x_dim)
x_std = jnp.ones(x_dim)
mask = False
if y is not None:
y_mu = np.mean(y) * np.ones(z_dim)
y_std = np.std(y) * np.ones(z_dim)
y_std_val = np.std(y)
else:
y_mu = np.zeros(z_dim)
y_std = np.ones(z_dim)
y_std_val = 1.0
phi = numpyro.sample(
"phi", dist.Normal(jnp.zeros((z_dim, x_dim)), jnp.ones(
(z_dim, x_dim))))
eta = numpyro.sample("eta", dist.Normal(x_mu, x_std))
theta = numpyro.sample("theta", dist.Normal(y_mu, y_std))
sigma = numpyro.sample("sigma", dist.Gamma(y_std_val, 1.0))
with numpyro.plate("batch", batch, dim=-2):
z = numpyro.sample("z", dist.Normal(jnp.zeros(z_dim), jnp.ones(z_dim)))
numpyro.sample("x",
dist.Normal(jnp.matmul(z, phi) + eta, x_std).mask(mask),
obs=x)
numpyro.sample("y", dist.Normal(jnp.matmul(z, theta), sigma), obs=y)
def gammadyn_single_model(beliefs, y, mask):
T, _ = beliefs[0].shape
gamma = npyro.sample('gamma', dist.InverseGamma(2., 2.))
mu = jnp.log(jnp.exp(gamma) - 1)
p = npyro.sample('p', dist.Dirichlet(jnp.ones(3)))
p0 = npyro.deterministic(
'p0',
jnp.concatenate([
p[..., :1] / 2, p[..., :1] / 2 + p[..., 1:2], p[..., 2:] / 2,
p[..., 2:] / 2
], -1))
scale = npyro.sample('scale', dist.Gamma(1., 1.))
rho = npyro.sample('rho', dist.Beta(1., 2.))
sigma = jnp.sqrt(-(1 - rho**2) / (2 * jnp.log(rho))) * scale
U = jnp.log(p0)
def transition_fn(carry, t):
x_prev = carry
dyn_gamma = npyro.deterministic('dyn_gamma', nn.softplus(mu + x_prev))
logs = logits((beliefs[0][t], beliefs[1][t]),
jnp.expand_dims(dyn_gamma, -1), jnp.expand_dims(U, -2))
npyro.sample('y', dist.CategoricalLogits(logs).mask(mask[t]))
noise = npyro.sample('dw', dist.Normal(0., 1.))
x_next = rho * x_prev + sigma * noise
return x_next, None
x0 = jnp.zeros(1)
with npyro.handlers.condition(data={"y": y}):
scan(transition_fn, x0, jnp.arange(T))
def model_normal_likelihood(X, Y):
D_X = X.shape[1]
# sample from horseshoe prior
lambdas = numpyro.sample("lambdas", dist.HalfCauchy(jnp.ones(D_X)))
tau = numpyro.sample("tau", dist.HalfCauchy(jnp.ones(1)))
# note that in practice for a normal likelihood we would probably want to
# integrate out the coefficients (as is done for example in sparse_regression.py).
# however, this trick wouldn't be applicable to other likelihoods
# (e.g. bernoulli, see below) so we don't make use of it here.
unscaled_betas = numpyro.sample("unscaled_betas",
dist.Normal(0.0, jnp.ones(D_X)))
scaled_betas = numpyro.deterministic("betas",
tau * lambdas * unscaled_betas)
# compute mean function using linear coefficients
mean_function = jnp.dot(X, scaled_betas)
prec_obs = numpyro.sample("prec_obs", dist.Gamma(3.0, 1.0))
sigma_obs = 1.0 / jnp.sqrt(prec_obs)
# observe data
numpyro.sample("Y", dist.Normal(mean_function, sigma_obs), obs=Y)
def __call__(self,
T=50,
N=1e5,
T_future=0,
E_duration_est=4.0,
I_duration_est=2.0,
R0_est=3.0,
beta_shape=1,
sigma_shape=5,
gamma_shape=8,
det_prob_est=0.3,
det_prob_conc=50,
det_noise_scale=0.7,
rw_scale=1e-1,
drift_scale=None,
confirmed=None,
death=None):
'''
Stochastic SEIR model. Draws random parameters and runs dynamics.
'''
# Sample initial number of infected individuals
I0 = numpyro.sample("I0", dist.Uniform(0, 0.02 * N))
E0 = numpyro.sample("E0", dist.Uniform(0, 0.02 * N))
H0 = numpyro.sample("H0", dist.Uniform(0, 0.02 * N))
D0 = numpyro.sample("D0", dist.Uniform(0, 1000))
# Sample parameters
sigma = numpyro.sample(
"sigma", dist.Gamma(sigma_shape, sigma_shape * E_duration_est))
gamma = numpyro.sample(
"gamma", dist.Gamma(gamma_shape, gamma_shape * I_duration_est))
# gamma = numpyro.sample("gamma",
# dist.TruncatedNormal(loc = 1./I_duration_est, scale = 0.25)
beta0 = numpyro.sample(
"beta0",
dist.Gamma(beta_shape, beta_shape * I_duration_est / R0_est))
det_prob = numpyro.sample(
"det_prob",
dist.Beta(det_prob_est * det_prob_conc,
(1 - det_prob_est) * det_prob_conc))
det_prob_d = numpyro.sample("det_prob_d",
dist.Beta(.9 * 100, (1 - .9) * 100))
death_prob = numpyro.sample("death_prob",
dist.Beta(.1 * 100, (1 - .1) * 100))
death_rate = numpyro.sample("death_rate", dist.Gamma(10, 10 * 10))
if drift_scale is not None:
drift = numpyro.sample("drift",
dist.Normal(loc=0, scale=drift_scale))
else:
drift = 0
x0 = SEIRDModel.seed(N=N, I=I0, E=E0, H=H0, D=D0)
numpyro.deterministic("x0", x0)
# Split observations into first and rest
confirmed0, confirmed = (None, None) if confirmed is None else (
confirmed[0], np.diff(confirmed))
death0, death = (None, None) if death is None else (death[0],
np.diff(death))
# First observation
with numpyro.handlers.scale(scale_factor=0.5):
y0 = observe("dy0",
x0[6],
det_prob,
det_noise_scale,
obs=confirmed0)
with numpyro.handlers.scale(scale_factor=2.0):
z0 = observe("dz0", x0[5], det_prob_d, det_noise_scale, obs=death0)
params = (beta0, sigma, gamma, rw_scale, drift, det_prob,
det_noise_scale, death_prob, death_rate, det_prob_d)
beta, x, y, z = self.dynamics(T,
params,
x0,
confirmed=confirmed,
death=death)
x = np.vstack((x0, x))
y = np.append(y0, y)
z = np.append(z0, z)
if T_future > 0:
params = (beta[-1], sigma, gamma, rw_scale, drift, det_prob,
det_noise_scale, death_prob, death_rate, det_prob_d)
beta_f, x_f, y_f, z_f = self.dynamics(T_future + 1,
params,
x[-1, :],
suffix="_future")
x = np.vstack((x, x_f))
y = np.append(y, y_f)
z = np.append(z, z_f)
return beta, x, y, z, det_prob, death_prob
def __call__(self,
T=50,
N=1e5,
T_future=0,
E_duration_est=4.0,
I_duration_est=2.0,
H_duration_est=10.0,
R0_est=3.0,
beta_shape=1.,
sigma_shape=100.,
gamma_shape=100.,
det_prob_est=0.3,
det_prob_conc=50.,
confirmed_dispersion=0.3,
death_dispersion=0.3,
rw_scale=2e-1,
forecast_rw_scale=0.,
drift_scale=None,
num_frozen=0,
rw_use_last=1,
confirmed=None,
death=None):
'''
Stochastic SEIR model. Draws random parameters and runs dynamics.
'''
# Sample initial number of infected individuals
I0 = numpyro.sample("I0", dist.Uniform(0, 0.02 * N))
E0 = numpyro.sample("E0", dist.Uniform(0, 0.02 * N))
H0 = numpyro.sample("H0", dist.Uniform(0, 1e-3 * N))
D0 = numpyro.sample("D0", dist.Uniform(0, 1e-3 * N))
# Sample dispersion parameters around specified values
death_dispersion = numpyro.sample(
"death_dispersion",
dist.TruncatedNormal(low=0.1, loc=death_dispersion, scale=0.15))
confirmed_dispersion = numpyro.sample(
"confirmed_dispersion",
dist.TruncatedNormal(low=0.1, loc=confirmed_dispersion,
scale=0.15))
# Sample parameters
sigma = numpyro.sample(
"sigma", dist.Gamma(sigma_shape, sigma_shape * E_duration_est))
gamma = numpyro.sample(
"gamma", dist.Gamma(gamma_shape, gamma_shape * I_duration_est))
beta0 = numpyro.sample(
"beta0",
dist.Gamma(beta_shape, beta_shape * I_duration_est / R0_est))
det_prob0 = numpyro.sample(
"det_prob0",
dist.Beta(det_prob_est * det_prob_conc,
(1 - det_prob_est) * det_prob_conc))
det_prob_d = numpyro.sample("det_prob_d",
dist.Beta(.9 * 100, (1 - .9) * 100))
death_prob = numpyro.sample("death_prob",
dist.Beta(0.01 * 100, (1 - 0.01) * 100))
#dist.Beta(0.02 * 1000, (1-0.02) * 1000))
death_rate = numpyro.sample("death_rate",
dist.Gamma(10, 10 * H_duration_est))
if drift_scale is not None:
drift = numpyro.sample("drift",
dist.Normal(loc=0, scale=drift_scale))
else:
drift = 0
x0 = SEIRDModel.seed(N=N, I=I0, E=E0, H=H0, D=D0)
numpyro.deterministic("x0", x0)
# Split observations into first and rest
if confirmed is None:
confirmed0, confirmed = (None, None)
else:
confirmed0 = confirmed[0]
confirmed = clean_daily_obs(onp.diff(confirmed))
if death is None:
death0, death = (None, None)
else:
death0 = death[0]
death = clean_daily_obs(onp.diff(death))
# First observation
with numpyro.handlers.scale(scale=0.5):
y0 = observe_nb2("dy0",
x0[6],
det_prob0,
confirmed_dispersion,
obs=confirmed0)
with numpyro.handlers.scale(scale=2.0):
z0 = observe_nb2("dz0",
x0[5],
det_prob_d,
death_dispersion,
obs=death0)
params = (beta0, sigma, gamma, rw_scale, drift, det_prob0,
confirmed_dispersion, death_dispersion, death_prob,
death_rate, det_prob_d)
beta, det_prob, x, y, z = self.dynamics(T,
params,
x0,
num_frozen=num_frozen,
confirmed=confirmed,
death=death)
x = np.vstack((x0, x))
y = np.append(y0, y)
z = np.append(z0, z)
if T_future > 0:
params = (beta[-rw_use_last:].mean(), sigma, gamma,
forecast_rw_scale, drift, det_prob[-rw_use_last:].mean(),
confirmed_dispersion, death_dispersion, death_prob,
death_rate, det_prob_d)
beta_f, det_rate_rw_f, x_f, y_f, z_f = self.dynamics(
T_future + 1, params, x[-1, :], suffix="_future")
x = np.vstack((x, x_f))
y = np.append(y, y_f)
z = np.append(z, z_f)
return beta, x, y, z, det_prob, death_prob
def SEIR_stochastic(T=50,
N=1e5,
T_future=0,
E_duration_est=4.0,
I_duration_est=2.0,
R0_est=3.0,
beta_shape=1,
sigma_shape=5,
gamma_shape=5,
det_rate_est=0.3,
det_rate_conc=50,
det_noise_scale=0.15,
rw_scale=1e-1,
drift_scale=None,
obs=None,
use_hosp=False,
hosp_rate_est=0.15,
hosp_rate_conc=30,
hosp_noise_scale=0.15,
hosp=None):
'''
Stochastic SEIR model. Draws random parameters and runs dynamics.
'''
# Sample initial number of infected individuals
I0 = numpyro.sample("I0", dist.Uniform(0, 0.02 * N))
E0 = numpyro.sample("E0", dist.Uniform(0, 0.02 * N))
# Sample parameters
sigma = numpyro.sample(
"sigma", dist.Gamma(sigma_shape, sigma_shape * E_duration_est))
gamma = numpyro.sample(
"gamma", dist.Gamma(gamma_shape, gamma_shape * I_duration_est))
beta0 = numpyro.sample(
"beta0", dist.Gamma(beta_shape, beta_shape * I_duration_est / R0_est))
det_rate = numpyro.sample(
"det_rate",
dist.Beta(det_rate_est * det_rate_conc,
(1 - det_rate_est) * det_rate_conc))
if use_hosp:
hosp_rate = det_rate * numpyro.sample(
"hosp_rate",
dist.Beta(hosp_rate_est * hosp_rate_conc,
(1 - hosp_rate_est) * hosp_rate_conc))
else:
hosp_rate = 0.0
if drift_scale is not None:
drift = numpyro.sample("drift", dist.Normal(loc=0, scale=drift_scale))
else:
drift = 0
x0 = SEIRModel.seed(N, I0, E0)
numpyro.deterministic("x0", x0)
# Split observations into first and rest
obs0, obs = (None, None) if obs is None else (obs[0], obs[1:])
if use_hosp:
hosp0, hosp = (None, None) if hosp is None else (hosp[0], hosp[1:])
# First observation
y0 = observe("y0", x0[4], det_rate, det_noise_scale, obs=obs0)
if use_hosp:
z0 = observe("z0", x0[4], hosp_rate, hosp_noise_scale, obs=hosp0)
else:
z0 = 0.
params = (beta0, sigma, gamma, rw_scale, drift, det_rate, det_noise_scale,
hosp_rate, hosp_noise_scale)
beta, x, y, z = SEIR_dynamics(T,
params,
x0,
use_hosp=use_hosp,
obs=obs,
hosp=hosp)
x = np.vstack((x0, x))
y = np.append(y0, y)
z = np.append(z0, z)
if T_future > 0:
params = (beta[-1], sigma, gamma, rw_scale, drift, det_rate,
det_noise_scale, hosp_rate, hosp_noise_scale)
beta_f, x_f, y_f, z_f = SEIR_dynamics(T_future + 1,
params,
x[-1, :],
use_hosp=use_hosp,
suffix="_future")
x = np.vstack((x, x_f))
y = np.append(y, y_f)
z = np.append(z, z_f)
return beta, x, y, z, det_rate, hosp_rate
def mix_model(data):
w = sample("w", dist.Dirichlet((1 / 3) * np.ones(3)), rng_key=rng)
mu = sample("mu", dist.Normal(np.zeros(3), np.ones(3)), rng_key=rng)
std = sample("std", dist.Gamma(np.ones(3), np.ones(3)), rng_key=rng)
sample("obs", NormalMixture(w, mu, std), rng_key=rng, obs=data)
def SparseGP(X, y):
n_samples = X.shape[0]
X = numpyro.deterministic("X", X)
# Set priors on kernel hyperparameters.
η = numpyro.sample("variance", dist.HalfCauchy(scale=5.0))
ℓ = numpyro.sample("length_scale", dist.Gamma(2.0, 1.0))
σ = numpyro.sample("obs_noise", dist.HalfCauchy(scale=5.0))
x_u = numpyro.param("x_u", init_value=X_u_init)
# η = numpyro.param("kernel_var", init_value=1.0, constraints=dist.constraints.positive)
# ℓ = numpyro.param("kernel_length", init_value=0.1, constraints=dist.constraints.positive)
# σ = numpyro.param("sigma", init_value=0.1, onstraints=dist.constraints.positive)
# ================================
# Mean Function
# ================================
f_loc = np.zeros(n_samples)
# ================================
# Qff Term
# ================================
# W = (inv(Luu) @ Kuf).T
# Qff = Kfu @ inv(Kuu) @ Kuf
# Qff = W @ W.T
# ================================
Kuu = rbf_kernel(x_u, x_u, η, ℓ)
Kuf = rbf_kernel(x_u, X, η, ℓ)
# Kuu += jnp.eye(Ninducing) * jitter
# add jitter
Kuu = add_to_diagonal(Kuu, jitter)
# cholesky factorization
Luu = cholesky(Kuu, lower=True)
Luu = numpyro.deterministic("Luu", Luu)
# W matrix
W = solve_triangular(Luu, Kuf, lower=True)
W = numpyro.deterministic("W", W).T
# ================================
# Likelihood Noise Term
# ================================
# D = noise
# ================================
D = numpyro.deterministic("G", jnp.ones(n_samples) * σ)
# ================================
# trace term
# ================================
# t = tr(Kff - Qff) / noise
# t /= - 2.0
# ================================
Kffdiag = jnp.diag(rbf_kernel(X, X, η, ℓ))
Qffdiag = jnp.power(W, 2).sum(axis=1)
trace_term = (Kffdiag - Qffdiag).sum() / σ
trace_term = jnp.clip(trace_term, a_min=0.0) # numerical errors
# add trace term to the log probability loss
numpyro.factor("trace_term", -trace_term / 2.0)
# Sample y according SGP
return numpyro.sample(
"y",
dist.LowRankMultivariateNormal(loc=f_loc, cov_factor=W,
cov_diag=D).expand_by(
y.shape[:-1]).to_event(y.ndim - 1),
obs=y,
)
