def model(home_id, away_id, score1_obs=None, score2_obs=None):
# priors
alpha = numpyro.sample("alpha", dist.Normal(0.0, 1.0))
sd_att = numpyro.sample(
"sd_att",
dist.FoldedDistribution(dist.StudentT(3.0, 0.0, 2.5)),
)
sd_def = numpyro.sample(
"sd_def",
dist.FoldedDistribution(dist.StudentT(3.0, 0.0, 2.5)),
)
home = numpyro.sample("home", dist.Normal(0.0, 1.0)) # home advantage
nt = len(np.unique(home_id))
# team-specific model parameters
with numpyro.plate("plate_teams", nt):
attack = numpyro.sample("attack", dist.Normal(0, sd_att))
defend = numpyro.sample("defend", dist.Normal(0, sd_def))
# likelihood
theta1 = jnp.exp(alpha + home + attack[home_id] - defend[away_id])
theta2 = jnp.exp(alpha + attack[away_id] - defend[home_id])
with numpyro.plate("data", len(home_id)):
numpyro.sample("s1", dist.Poisson(theta1), obs=score1_obs)
numpyro.sample("s2", dist.Poisson(theta2), obs=score2_obs)
def model(loc, scale):
with numpyro.plate_stack("plates", shape[:len(shape) - event_dim]):
with numpyro.plate("particles", 10000):
if "dist_type" == "Normal":
numpyro.sample("x", dist.Normal(loc, scale).to_event(event_dim))
else:
numpyro.sample("x", dist.StudentT(10.0, loc, scale).to_event(event_dim))
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 transition_fn(
carry: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], t: jnp.ndarray
) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], jnp.ndarray]:
level, s, moving_sum = carry
season = s[0] * level**pow_season
exp_val = level + coef_trend * level**pow_trend + season
exp_val = jnp.clip(exp_val, a_min=0)
y_t = jnp.where(t >= num_lim, exp_val, y[t])
moving_sum = moving_sum + y[t] - jnp.where(t >= seasonality,
y[t - seasonality], 0.0)
level_p = jnp.where(t >= seasonality, moving_sum / seasonality,
y_t - season)
level = level_sm * level_p + (1 - level_sm) * level
level = jnp.clip(level, a_min=0)
new_s = (s_sm * (y_t - level) / season + (1 - s_sm)) * s[0]
new_s = jnp.where(t >= num_lim, s[0], new_s)
s = jnp.concatenate([s[1:], new_s[None]], axis=0)
omega = sigma * exp_val**powx + offset_sigma
y_ = numpyro.sample("y", dist.StudentT(nu, exp_val, omega))
return (level, s, moving_sum), y_
def model(returns):
step_size = numpyro.sample("sigma", dist.Exponential(50.0))
s = numpyro.sample(
"s", dist.GaussianRandomWalk(scale=step_size, num_steps=jnp.shape(returns)[0])
)
nu = numpyro.sample("nu", dist.Exponential(0.1))
return numpyro.sample(
"r", dist.StudentT(df=nu, loc=0.0, scale=jnp.exp(s)), obs=returns
)
def model(returns):
step_size = numpyro.sample('sigma', dist.Exponential(50.))
s = numpyro.sample(
's',
dist.GaussianRandomWalk(scale=step_size,
num_steps=np.shape(returns)[0]))
nu = numpyro.sample('nu', dist.Exponential(.1))
return numpyro.sample('r',
dist.StudentT(df=nu, loc=0., scale=np.exp(-2 * s)),
obs=returns)
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 model(N=n):
plate1 = numpyro.plate("aa", N - 3, dim=-2)
plate2 = numpyro.plate("coord", 3, dim=-1)
with plate1, plate2:
M_last = numpyro.sample("M", dist.StudentT(1, 0, 50))
#M_last = numpyro.sample('M', dist.Normal(0, 10).expand_by([N-3,3]).to_event(1))
# Stack fixed and moving coordinates
M = np.concatenate((M_first, M_last))
# Make sure bond distances are around 3.8 Å
bonds = M[0:-1] - M[1:]
Bonds = (bonds[:, 0]**2 + bonds[:, 1]**2 + bonds[:, 2]**2)**(1 / 2)
# for i in pyro.plate('bonds', N-1):
# bond=Bonds[i]
# bond_obs = pyro.sample('bond_%i' % i, dist.Normal(bond, 0.001), obs=torch.tensor(3.8))
i = 0
with numpyro.plate("Bonds", 13):
bond_obs = numpyro.sample("Bonds_%i" % i,
dist.Normal(Bonds, 0.001),
obs=3.8)
i += 1
# Add a distance restraint between first and last point
D = M[0] - M[-1]
d = (D[0]**2 + D[1]**2 + D[2]**2)**(1 / 2)
d_obs = numpyro.sample("d_obs",
dist.Normal(d, 0.001),
obs=(dist_nr[0].item()))
for i in range(1, distances):
D = (M[points[i][0]] - M[points[i][1]])
d = (D[0]**2 + D[1]**2 + D[2]**2)**(1 / 2)
d_obs = numpyro.sample('d%s_obs' % i,
dist.Normal(d, s_d),
obs=(dist_nr[i].item()))
def skew_t_lpdf(x, nu, loc, scale, skew, clip_min=-100):
z = (jnp.clip(x, clip_min, jnp.inf) - loc) / scale
u = skew * z * jnp.sqrt((nu + 1) / (nu + z * z));
kernel = (dist.StudentT(nu).log_prob(z) +
dist.StudentT(nu + 1).log_cdf(u)) # not implemented...
return kernel + jnp.log(2/scale)
def log_prob(self, x):
z = (x - self.loc) / self.scale
u = self.skew * z * jnp.sqrt((self.df + 1) / (self.df + z ** 2));
kernel = (dist.StudentT(self.df).log_prob(z) +
std_studentt_lcdf(z=u, df=self.df + 1))
return kernel + jnp.log(2/self.scale)
def model_hierarchical(y, condition=None, group=None, treatment=None, dist_y='normal', add_group_slope=False,
add_group_intercept=True,
add_condition_slope=True, group2=None, add_group2_slope=False,
center_intercept=True, center_slope=False, robust_slopes=False,
add_condition_intercept=False,
dist_slope=dist.Normal
):
n_subjects = np.unique(group).shape[0]
n_conditions = np.unique(condition).shape[0]
if (condition is None) or not n_conditions:
add_condition_slope = False # Override
add_condition_intercept = False
if center_intercept:
intercept = numpyro.sample('intercept', dist.Normal(0, 100))
else:
intercept = 0
if (group is not None) and add_group_intercept:
if dist_y == 'poisson':
print('poisson intercepts')
a_group = numpyro.sample(f'mu_intercept_per_group', dist.Poisson(jnp.tile(0, n_subjects)))
intercept += a_group
else:
sigma_a_group = numpyro.sample('sigma_intercept_per_group', dist.HalfNormal(100))
a_group = numpyro.sample(f'mu_intercept_per_group', dist.Normal(jnp.tile(0, n_subjects), 10))
intercept += (a_group[group] * sigma_a_group)
if add_condition_intercept:
intercept_per_condition = numpyro.sample('intercept_per_condition',
dist.Normal(jnp.tile(0, n_conditions), 100))
sigma_intercept_per_condition = numpyro.sample('sigma_intercept_per_condition', dist.HalfNormal(100))
intercept += intercept_per_condition[condition] * sigma_intercept_per_condition
if not add_condition_slope:
center_slope = True # override center_slope
if center_slope:
slope = numpyro.sample('slope', dist_slope(0, 100))
else:
slope = 0
if add_condition_slope:
if robust_slopes:
# Robust slopes:
b_per_condition = numpyro.sample('slope_per_condition',
dist.StudentT(1, jnp.tile(0, n_conditions), 100))
else:
b_per_condition = numpyro.sample('slope_per_condition',
dist_slope(jnp.tile(0, n_conditions), 100))
sigma_b_condition = numpyro.sample('sigma_slope_per_condition', dist.HalfNormal(100))
slope = slope + b_per_condition[condition] * sigma_b_condition
if (group is not None) and add_group_slope:
sigma_b_group = numpyro.sample('sigma_slope_per_group', dist.HalfNormal(100))
b_per_group = numpyro.sample('slope_per_group', dist_slope(jnp.tile(0, n_subjects), 100))
slope = slope + b_per_group[group] * sigma_b_group
if (group2 is not None) and add_group2_slope:
sigma_b_group = numpyro.sample('sigma_slope_per_group2', dist.HalfNormal(100))
b_per_group = numpyro.sample('slope_per_group2', dist_slope(jnp.tile(0, n_subjects), 100))
slope = slope + b_per_group[group] * sigma_b_group
if type(intercept) is int:
print('Caution: No intercept')
if treatment is not None:
slope = slope * treatment
sample_y(dist_y=dist_y, theta=intercept + slope, y=y)
