def model(self, home_team, away_team, gameweek):
n_gameweeks = max(gameweek) + 1
sigma_0 = pyro.sample("sigma_0", dist.HalfNormal(5))
sigma_b = pyro.sample("sigma_b", dist.HalfNormal(5))
gamma = pyro.sample("gamma", dist.LogNormal(0, 1))
b = pyro.sample("b", dist.Normal(0, 1))
loc_mu_b = pyro.sample("loc_mu_b", dist.Normal(0, 1))
scale_mu_b = pyro.sample("scale_mu_b", dist.HalfNormal(1))
with pyro.plate("teams", self.n_teams):
log_a0 = pyro.sample("log_a0", dist.Normal(0, sigma_0))
mu_b = pyro.sample(
"mu_b",
dist.TransformedDistribution(
dist.Normal(0, 1),
dist.transforms.AffineTransform(loc_mu_b, scale_mu_b),
),
)
sigma_rw = pyro.sample("sigma_rw", dist.HalfNormal(0.1))
with pyro.plate("random_walk", n_gameweeks - 1):
diffs = pyro.sample(
"diff",
dist.TransformedDistribution(
dist.Normal(0, 1),
dist.transforms.AffineTransform(0, sigma_rw)),
)
diffs = np.vstack((log_a0, diffs))
log_a = np.cumsum(diffs, axis=-2)
with pyro.plate("weeks", n_gameweeks):
log_b = pyro.sample(
"log_b",
dist.TransformedDistribution(
dist.Normal(0, 1),
dist.transforms.AffineTransform(
mu_b + b * log_a, sigma_b),
),
)
pyro.sample("log_a", dist.Delta(log_a), obs=log_a)
home_inds = np.array([self.team_to_index[team] for team in home_team])
away_inds = np.array([self.team_to_index[team] for team in away_team])
home_rate = np.clip(
log_a[gameweek, home_inds] - log_b[gameweek, away_inds] + gamma,
-7, 2)
away_rate = np.clip(
log_a[gameweek, away_inds] - log_b[gameweek, home_inds], -7, 2)
pyro.sample("home_goals", dist.Poisson(np.exp(home_rate)))
pyro.sample("away_goals", dist.Poisson(np.exp(away_rate)))
def guide(X: DeviceArray):
n_stores, n_days, n_features = X.shape
n_features -= 1 # remove one dim for target
plate_features = numpyro.plate(Plate.features, n_features, dim=-1)
plate_stores = numpyro.plate(Plate.stores, n_stores, dim=-2)
numpyro.sample(
Site.disp_param_mu,
dist.Normal(loc=model_params[Param.loc_disp_param_mu],
scale=model_params[Param.scale_disp_param_mu]))
numpyro.sample(
Site.disp_param_sigma,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_disp_param_logsigma],
scale=model_params[Param.scale_disp_param_logsigma]),
transforms=dist.transforms.ExpTransform()))
with plate_stores:
numpyro.sample(
Site.disp_param_offsets,
dist.Normal(loc=numpyro.param(Param.loc_disp_param_offsets,
jnp.zeros((n_stores, 1))),
scale=numpyro.param(
Param.scale_disp_param_offsets,
0.1 * jnp.ones((n_stores, 1)),
constraint=dist.constraints.positive)))
with plate_features:
numpyro.sample(
Site.coef_mus,
dist.Normal(loc=model_params[Param.loc_coef_mus],
scale=model_params[Param.scale_coef_mus]))
numpyro.sample(
Site.coef_sigmas,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_coef_logsigmas],
scale=model_params[Param.scale_coef_logsigmas]),
transforms=dist.transforms.ExpTransform()))
with plate_stores:
numpyro.sample(
Site.coef_offsets,
dist.Normal(loc=numpyro.param(
Param.loc_coef_offsets,
jnp.zeros((n_stores, n_features))),
scale=numpyro.param(
Param.scale_coef_offsets,
0.5 * jnp.ones((n_stores, n_features)),
constraint=dist.constraints.positive)))
def test_transformed_distributions():
from tensorflow_probability.substrates.jax import bijectors as tfb
from numpyro.contrib.tfp import distributions as tfd
d = dist.TransformedDistribution(dist.Normal(0, 1), dist.transforms.ExpTransform())
d1 = tfd.TransformedDistribution(tfd.Normal(0, 1), tfb.Exp())
d2 = dist.TransformedDistribution(dist.Normal(0, 1), tfd.BijectorTransform(tfb.Exp()))
x = random.normal(random.PRNGKey(0), (1000,))
d_x = d.log_prob(x).sum()
d1_x = d1.log_prob(x).sum()
d2_x = d2.log_prob(x).sum()
assert_allclose(d_x, d1_x)
assert_allclose(d_x, d2_x)
def test_transformed_transformed_distribution():
loc, scale = -2, 3
dist1 = dist.TransformedDistribution(dist.Normal(2, 3), transforms.PowerTransform(2.))
dist2 = dist.TransformedDistribution(dist1, transforms.AffineTransform(-2, 3))
assert isinstance(dist2.base_dist, dist.Normal)
assert len(dist2.transforms) == 2
assert isinstance(dist2.transforms[0], transforms.PowerTransform)
assert isinstance(dist2.transforms[1], transforms.AffineTransform)
rng_key = random.PRNGKey(0)
assert_allclose(loc + scale * dist1.sample(rng_key), dist2.sample(rng_key))
intermediates = dist2.sample_with_intermediates(rng_key)
assert len(intermediates) == 2
def reparam_model(dim=10):
y = numpyro.sample('y', dist.Normal(0, 3))
with numpyro.handlers.reparam(config={'x': TransformReparam()}):
numpyro.sample(
'x',
dist.TransformedDistribution(dist.Normal(jnp.zeros(dim - 1), 1),
AffineTransform(0, jnp.exp(y / 2))))
def model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
with handlers.reparam(config={'loc': TransformReparam()}):
loc = numpyro.sample('loc', dist.TransformedDistribution(
dist.Uniform(0, 1).mask(False),
AffineTransform(0, alpha)))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def test_elbo_dynamic_support():
x_prior = dist.TransformedDistribution(
dist.Normal(),
[AffineTransform(0, 2),
SigmoidTransform(),
AffineTransform(0, 3)])
x_guide = dist.Uniform(0, 3)
def model():
numpyro.sample('x', x_prior)
def guide():
numpyro.sample('x', x_guide)
adam = optim.Adam(0.01)
# set base value of x_guide is 0.9
x_base = 0.9
guide = substitute(guide, base_param_map={'x': x_base})
svi = SVI(model, guide, elbo, adam)
svi_state = svi.init(random.PRNGKey(0), (), ())
actual_loss = svi.evaluate(svi_state)
assert np.isfinite(actual_loss)
x, _ = x_guide.transform_with_intermediates(x_base)
expected_loss = x_guide.log_prob(x) - x_prior.log_prob(x)
assert_allclose(actual_loss, expected_loss)
def _sample_latent(self, base_dist, *args, **kwargs):
sample_shape = kwargs.pop('sample_shape', ())
transform = self._get_transform()
posterior = dist.TransformedDistribution(base_dist, transform)
return numpyro.sample("_{}_latent".format(self.prefix),
posterior,
sample_shape=sample_shape)
def model():
fn = dist.TransformedDistribution(
dist.Normal(jnp.zeros_like(loc), jnp.ones_like(scale)),
[AffineTransform(loc, scale), ExpTransform()]).expand(shape)
if event_shape:
fn = fn.to_event(len(event_shape)).expand_by([100000])
with numpyro.plate_stack("plates", batch_shape):
with numpyro.plate("particles", 100000):
return numpyro.sample("x", fn)
def model(data):
alpha = numpyro.sample("alpha", dist.Uniform(0, 1))
with numpyro.handlers.reparam(config={"loc": TransformReparam()}):
loc = numpyro.sample(
"loc",
dist.TransformedDistribution(
dist.Uniform(0, 1).mask(False), AffineTransform(0, alpha)),
)
numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
def model():
with numpyro.plate_stack("plates", shape):
with numpyro.plate("particles", 100000):
return numpyro.sample(
"x",
dist.TransformedDistribution(
dist.Normal(jnp.zeros_like(loc), jnp.ones_like(scale)),
[AffineTransform(loc, scale),
ExpTransform()]).expand_by([100000]))
def test_transformed_potential_energy():
beta_dist = dist.Beta(np.ones(5), np.ones(5))
transform = transforms.AffineTransform(3, 4)
inv_transform = transforms.AffineTransform(-0.75, 0.25)
z = random.normal(random.PRNGKey(0), (5,))
pe_expected = -dist.TransformedDistribution(beta_dist, transform).log_prob(z)
potential_fn = lambda x: -beta_dist.log_prob(x) # noqa: E731
pe_actual = transformed_potential_energy(potential_fn, inv_transform, z)
assert_allclose(pe_actual, pe_expected)
def actual_model(data):
alpha = numpyro.sample("alpha", dist.Uniform(0, 1))
with numpyro.handlers.reparam(config={"loc": TransformReparam()}):
loc = numpyro.sample(
"loc",
dist.TransformedDistribution(
dist.Uniform(0, 1), transforms.AffineTransform(0, alpha)),
)
with numpyro.plate("N", len(data)):
numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
def get_posterior(self, params):
"""
Returns the posterior distribution.
:param dict params: Current parameters of model and autoguide.
The parameters can be obtained using :meth:`~numpyro.infer.svi.SVI.get_params`
method from :class:`~numpyro.infer.svi.SVI`.
"""
base_dist = self.get_base_dist()
transform = self.get_transform(params)
return dist.TransformedDistribution(base_dist, transform)
def _get_posterior(self):
if self.latent_dim == 1:
raise ValueError('latent dim = 1. Consider using AutoDiagonalNormal instead')
flows = []
for i in range(self.num_flows):
if i > 0:
flows.append(PermuteTransform(jnp.arange(self.latent_dim)[::-1]))
residual = "gated" if i < (self.num_flows - 1) else None
arn = BlockNeuralAutoregressiveNN(self.latent_dim, self._hidden_factors, residual)
arnn = numpyro.module('{}_arn__{}'.format(self.prefix, i), arn, (self.latent_dim,))
flows.append(BlockNeuralAutoregressiveTransform(arnn))
return dist.TransformedDistribution(self.get_base_dist(), flows)
def BinomialApprox(n, p, conc=None):
'''
Return distribution that is a continuous approximation to
Binomial(n, p); allows overdispersion by setting conc < n
'''
if conc is None:
conc = n
a = conc * p
b = conc * (1 - p)
# This is the distribution of n * Beta(a, b)
return dist.TransformedDistribution(
dist.Beta(a, b), dist.transforms.AffineTransform(loc=0, scale=n))
def _get_posterior(self):
if self.latent_dim == 1:
raise ValueError('latent dim = 1. Consider using AutoDiagonalNormal instead')
hidden_dims = [self.latent_dim, self.latent_dim] if self._hidden_dims is None else self._hidden_dims
flows = []
for i in range(self.num_flows):
if i > 0:
flows.append(PermuteTransform(jnp.arange(self.latent_dim)[::-1]))
arn = AutoregressiveNN(self.latent_dim, hidden_dims,
permutation=jnp.arange(self.latent_dim),
skip_connections=self._skip_connections,
nonlinearity=self._nonlinearity)
arnn = numpyro.module('{}_arn__{}'.format(self.prefix, i), arn, (self.latent_dim,))
flows.append(InverseAutoregressiveTransform(arnn))
return dist.TransformedDistribution(self.get_base_dist(), flows)
def model_noncentered(num: int,
sigma: np.ndarray,
y: Optional[np.ndarray] = None) -> None:
mu = numpyro.sample("mu", dist.Normal(0, 5))
tau = numpyro.sample("tau", dist.HalfCauchy(5))
with numpyro.plate("num", num):
with numpyro.handlers.reparam(config={"theta": TransformReparam()}):
theta = numpyro.sample(
"theta",
dist.TransformedDistribution(
dist.Normal(0.0, 1.0),
dist.transforms.AffineTransform(mu, tau)),
)
numpyro.sample("obs", dist.Normal(theta, sigma), obs=y)
def __call__(self, *args, **kwargs):
"""
An automatic guide with the same ``*args, **kwargs`` as the base ``model``.
:return: A dict mapping sample site name to sampled value.
:rtype: dict
"""
if self.prototype_trace is None:
# run model to inspect the model structure
self._setup_prototype(*args, **kwargs)
plates = self._create_plates(*args, **kwargs)
result = {}
for name, site in self.prototype_trace.items():
if site["type"] != "sample" or site["is_observed"]:
continue
event_dim = self._event_dims[name]
init_loc = self._init_locs[name]
with ExitStack() as stack:
for frame in site["cond_indep_stack"]:
stack.enter_context(plates[frame.name])
site_loc = numpyro.param("{}_{}_loc".format(name, self.prefix),
init_loc,
event_dim=event_dim)
site_scale = numpyro.param(
"{}_{}_scale".format(name, self.prefix),
jnp.full(jnp.shape(init_loc), self._init_scale),
constraint=self.scale_constraint,
event_dim=event_dim,
)
site_fn = dist.Normal(site_loc, site_scale).to_event(event_dim)
if site["fn"].support is constraints.real or (
isinstance(site["fn"].support, constraints.independent)
and site["fn"].support is constraints.real):
result[name] = numpyro.sample(name, site_fn)
else:
transform = biject_to(site["fn"].support)
guide_dist = dist.TransformedDistribution(
site_fn, transform)
result[name] = numpyro.sample(name, guide_dist)
return result
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
def test_elbo_dynamic_support():
x_prior = dist.TransformedDistribution(
dist.Normal(), [AffineTransform(0, 2), SigmoidTransform(), AffineTransform(0, 3)])
x_guide = dist.Uniform(0, 3)
def model():
numpyro.sample('x', x_prior)
def guide():
numpyro.sample('x', x_guide)
adam = optim.Adam(0.01)
x = 2.
guide = substitute(guide, data={'x': x})
svi = SVI(model, guide, adam, Trace_ELBO())
svi_state = svi.init(random.PRNGKey(0))
actual_loss = svi.evaluate(svi_state)
assert jnp.isfinite(actual_loss)
expected_loss = x_guide.log_prob(x) - x_prior.log_prob(x)
assert_allclose(actual_loss, expected_loss)
def ExponentialRandomWalk(loc=1., scale=1e-2, drift=0., num_steps=100):
'''
Return distrubtion of exponentiated Gaussian random walk
Variables are x_0, ..., x_{T-1}
Dynamics in log-space are random walk with drift:
log(x_0) := log(loc)
log(x_t) := log(x_{t-1}) + drift + eps_t, eps_t ~ N(0, scale)
==> Dynamics in non-log space are:
x_0 := loc
x_t := x_{t-1} * exp(drift + eps_t), eps_t ~ N(0, scale)
'''
log_loc = np.log(loc) + drift * (np.arange(num_steps) + 0.)
return dist.TransformedDistribution(
dist.GaussianRandomWalk(scale=scale, num_steps=num_steps), [
dist.transforms.AffineTransform(loc=log_loc, scale=1.),
dist.transforms.ExpTransform()
])
def ar_k(n_coefs, obs=None, X=None):
beta = numpyro.sample(name="beta",
sample_shape=(n_coefs, ),
fn=dist.TransformedDistribution(
dist.Normal(loc=0., scale=1),
transforms=dist.transforms.AffineTransform(
loc=0,
scale=1,
domain=dist.constraints.interval(-1, 1))))
tau = numpyro.sample(name="tau", fn=dist.HalfCauchy(scale=1))
z_init = numpyro.sample(name='z_init',
fn=dist.Normal(0, 1),
sample_shape=(n_coefs, ))
obs_init = z_init[:n_coefs - 1]
(beta, obs_last), zs_exp = scan_fn(n_coefs, beta, obs_init, obs)
Z_exp = np.concatenate((z_init, zs_exp), axis=0)
Z = numpyro.sample(name="Z", fn=dist.Normal(loc=Z_exp, scale=tau), obs=obs)
return Z_exp, obs_last
def guide(X: DeviceArray):
"""Guide with parameters of the posterior
Args:
X: input data
"""
n_stores, n_days, n_features = X.shape
n_features -= 1 # remove one dim for target
plate_features = numpyro.plate(Plate.features, n_features, dim=-1)
plate_stores = numpyro.plate(Plate.stores, n_stores, dim=-2)
numpyro.sample(
Site.disp_param_mu,
dist.Normal(loc=numpyro.param(Param.loc_disp_param_mu,
4. * jnp.ones(1)),
scale=numpyro.param(Param.scale_disp_param_mu,
1. * jnp.ones(1),
constraint=dist.constraints.positive)))
numpyro.sample(
Site.disp_param_sigma,
dist.TransformedDistribution(
dist.Normal(loc=numpyro.param(Param.loc_disp_param_logsigma,
1.0 * jnp.ones(1)),
scale=numpyro.param(
Param.scale_disp_param_logsigma,
0.1 * jnp.ones(1),
constraint=dist.constraints.positive)),
transforms=dist.transforms.ExpTransform()))
with plate_stores:
numpyro.sample(
Site.disp_param_offsets,
dist.Normal(loc=numpyro.param(Param.loc_disp_param_offsets,
jnp.zeros((n_stores, 1))),
scale=numpyro.param(
Param.scale_disp_param_offsets,
0.1 * jnp.ones((n_stores, 1)),
constraint=dist.constraints.positive)))
with plate_features:
numpyro.sample(
Site.coef_mus,
dist.Normal(loc=numpyro.param(Param.loc_coef_mus,
jnp.ones(n_features)),
scale=numpyro.param(
Param.scale_coef_mus,
0.5 * jnp.ones(n_features),
constraint=dist.constraints.positive)))
numpyro.sample(
Site.coef_sigmas,
dist.TransformedDistribution(
dist.Normal(loc=numpyro.param(Param.loc_coef_logsigmas,
jnp.zeros(n_features)),
scale=numpyro.param(
Param.scale_coef_logsigmas,
0.5 * jnp.ones(n_features),
constraint=dist.constraints.positive)),
transforms=dist.transforms.ExpTransform()))
with plate_stores:
numpyro.sample(
Site.coef_offsets,
dist.Normal(
loc=numpyro.param(Param.loc_coef_offsets,
jnp.zeros((n_stores, n_features))),
scale=numpyro.param(Param.scale_coef_offsets,
0.5 * jnp.ones((n_stores, n_features)),
constraint=dist.constraints.positive)))
def reparam_model(dim=10):
y = numpyro.sample('y', dist.Normal(0, 3))
numpyro.sample(
'x',
dist.TransformedDistribution(dist.Normal(np.zeros(dim - 1), 1),
AffineTransform(0, np.exp(y / 2))))
def model(X: DeviceArray):
n_stores, n_days, n_features = X.shape
n_features -= 1 # remove one dim for target
plate_features = numpyro.plate(Plate.features, n_features, dim=-1)
plate_stores = numpyro.plate(Plate.stores, n_stores, dim=-2)
plate_days = numpyro.plate(Plate.days, n_days, dim=-1)
disp_param_mu = numpyro.sample(
Site.disp_param_mu,
dist.Normal(
loc=model_params[Param.loc_disp_param_mu],
scale=model_params[Param.scale_disp_param_mu],
),
)
disp_param_sigma = numpyro.sample(
Site.disp_param_sigma,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_disp_param_logsigma],
scale=model_params[Param.scale_disp_param_logsigma],
),
transforms=dist.transforms.ExpTransform(),
),
)
with plate_stores:
with numpyro.handlers.reparam(
config={Site.disp_params: TransformReparam()}):
disp_params = numpyro.sample(
Site.disp_params,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_disp_params],
scale=model_params[Param.scale_disp_params],
),
dist.transforms.AffineTransform(
disp_param_mu, disp_param_sigma),
),
)
with plate_features:
coef_mus = numpyro.sample(
Site.coef_mus,
dist.Normal(
loc=model_params[Param.loc_coef_mus],
scale=model_params[Param.scale_coef_mus],
),
)
coef_sigmas = numpyro.sample(
Site.coef_sigmas,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_coef_logsigmas],
scale=model_params[Param.scale_coef_logsigmas],
),
transforms=dist.transforms.ExpTransform(),
),
)
with plate_stores:
with numpyro.handlers.reparam(
config={Site.coefs: TransformReparam()}):
coefs = numpyro.sample(
Site.coefs,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_coefs],
scale=model_params[Param.scale_coefs],
),
dist.transforms.AffineTransform(
coef_mus, coef_sigmas),
),
)
with plate_days, plate_stores:
features = jnp.nan_to_num(X[..., :-1])
means = jnp.exp(
jnp.sum(jnp.expand_dims(coefs, axis=1) * features, axis=2))
betas = jnp.exp(-disp_params)
alphas = means * betas
return numpyro.sample(Site.days, dist.GammaPoisson(alphas, betas))
assert np.shape(actual) == batch_shape
if len(shape) == transform.event_dim:
if isinstance(transform, PermuteTransform):
expected = onp.linalg.slogdet(jax.jacobian(transform)(x))[1]
inv_expected = onp.linalg.slogdet(jax.jacobian(
transform.inv)(y))[1]
else:
expected = np.log(np.abs(grad(transform)(x)))
inv_expected = np.log(np.abs(grad(transform.inv)(y)))
assert_allclose(actual, expected, atol=1e-6)
assert_allclose(actual, -inv_expected, atol=1e-6)
@pytest.mark.parametrize('transformed_dist', [
dist.TransformedDistribution(dist.Normal(np.array([2., 3.]), 1.),
constraints.ExpTransform()),
dist.TransformedDistribution(dist.Exponential(np.ones(2)), [
constraints.PowerTransform(0.7),
constraints.AffineTransform(0.,
np.ones(2) * 3)
]),
])
def test_transformed_distribution_intermediates(transformed_dist):
sample, intermediates = transformed_dist.sample_with_intermediates(
random.PRNGKey(1))
assert_allclose(transformed_dist.log_prob(sample, intermediates),
transformed_dist.log_prob(sample))
def test_transformed_transformed_distribution():
loc, scale = -2, 3
def guide(X: DeviceArray):
n_stores, n_days, n_features = X.shape
n_features -= 1 # remove one dim for target
plate_features = numpyro.plate(Plate.features, n_features, dim=-1)
plate_stores = numpyro.plate(Plate.stores, n_stores, dim=-2)
disp_param_mu = numpyro.sample(
Site.disp_param_mu,
dist.Normal(
loc=model_params[Param.loc_disp_param_mu],
scale=model_params[Param.scale_disp_param_mu],
),
)
disp_param_sigma = numpyro.sample(
Site.disp_param_sigma,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_disp_param_logsigma],
scale=model_params[Param.scale_disp_param_logsigma],
),
transforms=dist.transforms.ExpTransform(),
),
)
with plate_stores:
with numpyro.handlers.reparam(
config={Site.disp_params: TransformReparam()}):
numpyro.sample(
Site.disp_params,
dist.TransformedDistribution(
dist.Normal(
loc=numpyro.param(Param.loc_disp_params,
jnp.zeros((n_stores, 1))),
scale=numpyro.param(
Param.scale_disp_params,
0.1 * jnp.ones((n_stores, 1)),
constraint=dist.constraints.positive,
),
),
dist.transforms.AffineTransform(
disp_param_mu, disp_param_sigma),
),
)
with plate_features:
coef_mus = numpyro.sample(
Site.coef_mus,
dist.Normal(
loc=model_params[Param.loc_coef_mus],
scale=model_params[Param.scale_coef_mus],
),
)
coef_sigmas = numpyro.sample(
Site.coef_sigmas,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_coef_logsigmas],
scale=model_params[Param.scale_coef_logsigmas],
),
transforms=dist.transforms.ExpTransform(),
),
)
with plate_stores:
with numpyro.handlers.reparam(
config={Site.coefs: TransformReparam()}):
numpyro.sample(
Site.coefs,
dist.TransformedDistribution(
dist.Normal(
loc=numpyro.param(
Param.loc_coefs,
jnp.zeros((n_stores, n_features))),
scale=numpyro.param(
Param.scale_coefs,
0.5 * jnp.ones((n_stores, n_features)),
constraint=dist.constraints.positive,
),
),
dist.transforms.AffineTransform(
coef_mus, coef_sigmas),
),
)
def model(X: DeviceArray) -> DeviceArray:
"""Gamma-Poisson hierarchical model for daily sales forecasting
Args:
X: input data
Returns:
output data
"""
n_stores, n_days, n_features = X.shape
n_features -= 1 # remove one dim for target
eps = 1e-12 # epsilon
plate_features = numpyro.plate(Plate.features, n_features, dim=-1)
plate_stores = numpyro.plate(Plate.stores, n_stores, dim=-2)
plate_days = numpyro.plate(Plate.days, n_days, dim=-1)
disp_param_mu = numpyro.sample(Site.disp_param_mu,
dist.Normal(loc=4.0, scale=1.0))
disp_param_sigma = numpyro.sample(Site.disp_param_sigma,
dist.HalfNormal(scale=1.0))
with plate_stores:
with numpyro.handlers.reparam(
config={Site.disp_params: TransformReparam()}):
disp_params = numpyro.sample(
Site.disp_params,
dist.TransformedDistribution(
dist.Normal(loc=jnp.zeros((n_stores, 1)), scale=0.1),
dist.transforms.AffineTransform(disp_param_mu,
disp_param_sigma),
),
)
with plate_features:
coef_mus = numpyro.sample(
Site.coef_mus,
dist.Normal(loc=jnp.zeros(n_features), scale=jnp.ones(n_features)),
)
coef_sigmas = numpyro.sample(
Site.coef_sigmas,
dist.HalfNormal(scale=2.0 * jnp.ones(n_features)))
with plate_stores:
with numpyro.handlers.reparam(
config={Site.coefs: TransformReparam()}):
coefs = numpyro.sample(
Site.coefs,
dist.TransformedDistribution(
dist.Normal(loc=jnp.zeros((n_stores, n_features)),
scale=1.0),
dist.transforms.AffineTransform(coef_mus, coef_sigmas),
),
)
with plate_days, plate_stores:
targets = X[..., -1]
features = jnp.nan_to_num(X[..., :-1]) # padded features to 0
is_observed = jnp.where(jnp.isnan(targets), jnp.zeros_like(targets),
jnp.ones_like(targets))
not_observed = 1 - is_observed
means = (is_observed * jnp.exp(
jnp.sum(jnp.expand_dims(coefs, axis=1) * features, axis=2)) +
not_observed * eps)
betas = is_observed * jnp.exp(-disp_params) + not_observed
alphas = means * betas
return numpyro.sample(Site.days,
dist.GammaPoisson(alphas, betas),
obs=jnp.nan_to_num(targets))
def model(X: DeviceArray):
n_stores, n_days, n_features = X.shape
n_features -= 1 # remove one dim for target
plate_features = numpyro.plate(Plate.features, n_features, dim=-1)
plate_stores = numpyro.plate(Plate.stores, n_stores, dim=-2)
plate_days = numpyro.plate(Plate.days, n_days, dim=-1)
disp_param_mu = numpyro.sample(
Site.disp_param_mu,
dist.Normal(loc=model_params[Param.loc_disp_param_mu],
scale=model_params[Param.scale_disp_param_mu]))
disp_param_sigma = numpyro.sample(
Site.disp_param_sigma,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_disp_param_logsigma],
scale=model_params[Param.scale_disp_param_logsigma]),
transforms=dist.transforms.ExpTransform()))
with plate_stores:
disp_param_offsets = numpyro.sample(
Site.disp_param_offsets,
dist.Normal(
loc=model_params[Param.loc_disp_param_offsets],
scale=model_params[Param.scale_disp_param_offsets]),
)
disp_params = disp_param_mu + disp_param_offsets * disp_param_sigma
disp_params = numpyro.sample(Site.disp_params,
dist.Delta(disp_params),
obs=disp_params)
with plate_features:
coef_mus = numpyro.sample(
Site.coef_mus,
dist.Normal(loc=model_params[Param.loc_coef_mus],
scale=model_params[Param.scale_coef_mus]))
coef_sigmas = numpyro.sample(
Site.coef_sigmas,
dist.TransformedDistribution(
dist.Normal(
loc=model_params[Param.loc_coef_logsigmas],
scale=model_params[Param.scale_coef_logsigmas]),
transforms=dist.transforms.ExpTransform()))
with plate_stores:
coef_offsets = numpyro.sample(
Site.coef_offsets,
dist.Normal(loc=model_params[Param.loc_coef_offsets],
scale=model_params[Param.scale_coef_offsets]))
coefs = coef_mus + coef_offsets * coef_sigmas
coefs = numpyro.sample(Site.coefs,
dist.Delta(coefs),
obs=coefs)
with plate_days, plate_stores:
features = jnp.nan_to_num(X[..., :-1])
means = jnp.exp(
jnp.sum(jnp.expand_dims(coefs, axis=1) * features, axis=2))
betas = jnp.exp(-disp_params)
alphas = means * betas
return numpyro.sample(Site.days, dist.GammaPoisson(alphas, betas))