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)