def model( covariates: jnp.ndarray, x: Optional[jnp.ndarray] = None, x_dim: int = 1, ) -> None: if x is not None: x_dim = x.shape[-1] seq_len, batch, c_dim = covariates.shape weight_var = numpyro.sample( "weight_var", dist.LogNormal(-5 * jnp.ones((c_dim, x_dim)), 5 * jnp.ones( (c_dim, x_dim)))) sigma = numpyro.sample( "sigma", dist.LogNormal(-20 * jnp.ones(x_dim), 20 * jnp.ones(x_dim))) def transition_fn( carry: Tuple[jnp.ndarray, jnp.ndarray], t: jnp.ndarray, ) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], jnp.ndarray]: z_prev, w_prev = carry z = numpyro.sample("z", dist.Normal(z_prev, 1)) weight = numpyro.sample("weight", dist.Normal(w_prev, weight_var)) numpyro.sample( "x", dist.Normal(z + jnp.matmul(covariates[t], weight), sigma)) return (z, weight), None z_init = jnp.zeros((batch, x_dim)) w_init = jnp.zeros((c_dim, x_dim)) with numpyro.handlers.condition(data={"x": x}): scan(transition_fn, (z_init, w_init), jnp.arange(seq_len))
def model(self, batch): XL, XH = batch['XL'], batch['XH'] y = batch['y'] NL, NH = XL.shape[0], XH.shape[0] D = XH.shape[1] # set uninformative log-normal priors for low-fidelity kernel var_L = sample('kernel_var_L', dist.LogNormal(0.0, 1.0), sample_shape = (1,)) length_L = sample('kernel_length_L', dist.LogNormal(0.0, 1.0), sample_shape = (D,)) theta_L = np.concatenate([var_L, length_L]) # set uninformative log-normal priors for high-fidelity kernel var_H = sample('kernel_var_H', dist.LogNormal(0.0, 1.0), sample_shape = (1,)) length_H = sample('kernel_length_H', dist.LogNormal(0.0, 1.0), sample_shape = (D,)) theta_H = np.concatenate([var_H, length_H]) # prior for rho rho = sample('rho', dist.Normal(0.0, 10.0), sample_shape = (1,)) # Compute kernels K_LL = self.kernel(XL, XL, theta_L) + np.eye(NL)*1e-8 K_LH = rho*self.kernel(XL, XH, theta_L) K_HH = rho**2 * self.kernel(XH, XH, theta_L) + \ self.kernel(XH, XH, theta_H) + np.eye(NH)*1e-8 K = np.vstack((np.hstack((K_LL,K_LH)), np.hstack((K_LH.T,K_HH)))) L = cholesky(K, lower=True) # Generate latent function beta_L = sample('beta_L', dist.Normal(0.0, 1.0)) beta_H = sample('beta_H', dist.Normal(0.0, 1.0)) eta_L = sample('eta_L', dist.Normal(0.0, 1.0), sample_shape=(NL,)) eta_H = sample('eta_H', dist.Normal(0.0, 1.0), sample_shape=(NH,)) beta = np.concatenate([beta_L*np.ones(NL), beta_H*np.ones(NH)]) eta = np.concatenate([eta_L, eta_H]) f = np.matmul(L, eta) + beta # Bernoulli likelihood sample('y', dist.Bernoulli(logits=f), obs=y)
def model(N, y=None): """ :param int N: number of measurement times :param numpy.ndarray y: measured populations with shape (N, 2) """ # initial population z_init = numpyro.sample("z_init", dist.LogNormal(jnp.log(10), 1).expand([2])) # measurement times ts = jnp.arange(float(N)) # parameters alpha, beta, gamma, delta of dz_dt theta = numpyro.sample( "theta", dist.TruncatedNormal( low=0.0, loc=jnp.array([1.0, 0.05, 1.0, 0.05]), scale=jnp.array([0.5, 0.05, 0.5, 0.05]), ), ) # integrate dz/dt, the result will have shape N x 2 z = odeint(dz_dt, z_init, ts, theta, rtol=1e-6, atol=1e-5, mxstep=1000) # measurement errors sigma = numpyro.sample("sigma", dist.LogNormal(-1, 1).expand([2])) # measured populations numpyro.sample("y", dist.LogNormal(jnp.log(z), sigma), obs=y)
def model( covariates: jnp.ndarray, x: Optional[jnp.ndarray] = None, x_dim: int = 1, z_dim: int = 1, seasonality: int = 7, ) -> None: seq_len, batch, c_dim = covariates.shape if x is not None: x_dim = x.shape[-1] season_trans = jax.ops.index_add(jnp.eye(seasonality - 1, k=-1), 0, -1) season_var = numpyro.sample("season_var", dist.LogNormal(-5, 5)) trend_trans = jnp.array([[1, 1], [0, 1]]) trend_var = numpyro.sample( "trend_var", dist.LogNormal(jnp.array([-5, -5]), jnp.array([5, 5]))) weight_var = numpyro.sample( "weight_var", dist.LogNormal(-5 * jnp.ones((c_dim, x_dim)), 5 * jnp.ones( (c_dim, x_dim)))) sigma = numpyro.sample( "sigma", dist.LogNormal(-5 * np.ones(x_dim), 5 * np.ones(x_dim))) def transition_fn( carry: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray], t: jnp.ndarray ) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray], jnp.ndarray]: z_prev, s_prev, t_prev, w_prev = carry z = numpyro.sample("z", dist.Normal(z_prev, jnp.ones(z_dim))) s = jnp.matmul(s_prev, season_trans.T) s0 = numpyro.sample("s0", dist.Normal(s[:, 0], season_var)) s = jax.ops.index_update(s, jax.ops.index[:, 0], s0) trend_mu = jnp.matmul(t_prev, trend_trans.T) trend = numpyro.sample("trend", dist.Normal(trend_mu, trend_var)) weight = numpyro.sample("weight", dist.Normal(w_prev, weight_var)) exogenous = jnp.matmul(covariates[t], weight) numpyro.sample( "x", dist.Normal(z.sum(-1) + s0 + trend[:, 0] + exogenous, sigma)) return (z, s, trend, weight), None z_init = jnp.zeros((batch, z_dim)) s_init = jnp.zeros((batch, seasonality - 1)) t_init = jnp.zeros((batch, 2)) w_init = jnp.zeros((c_dim, x_dim)) with numpyro.handlers.condition(data={"x": x}): scan(transition_fn, (z_init, s_init, t_init, w_init), jnp.arange(seq_len))
def model(data): xc = numpyro.sample('xc', dist.Normal(0, 10)) yc = numpyro.sample('yc', dist.Normal(0, 10)) w = numpyro.sample('w', dist.LogNormal(0, 10)) h = numpyro.sample('h', dist.LogNormal(0, 10)) phi = numpyro.sample('phi', dist.Uniform(0, np.pi)) px, py = forward(xc, yc, w, h, phi) sigma = numpyro.sample('sigma', dist.Exponential(1.)) return numpyro.sample('obs', dist.Normal(np.vstack([px, py]), sigma), obs=data)
def model(X, Y): # set uninformative log-normal priors on our three kernel hyperparameters var = numpyro.sample("kernel_var", dist.LogNormal(0.0, 10.0)) noise = numpyro.sample("kernel_noise", dist.LogNormal(0.0, 10.0)) length = numpyro.sample("kernel_length", dist.LogNormal(0.0, 10.0)) # compute kernel k = kernel(X, X, var, length, noise) # sample Y according to the standard gaussian process formula numpyro.sample("Y", dist.MultivariateNormal(loc=jnp.zeros(X.shape[0]), covariance_matrix=k), obs=Y)
def model(self, batch): X = batch['X'] y = batch['y'] N, D = X.shape # set uninformative log-normal priors var = sample('kernel_var', dist.LogNormal(0.0, 10.0)) length = sample('kernel_length', dist.LogNormal(np.zeros(D), 10.0*np.ones(D))) noise = sample('noise_var', dist.LogNormal(0.0, 10.0)) theta = np.concatenate([np.array([var]), np.array(length)]) # compute kernel K = self.kernel(X, X, theta) + np.eye(N)*(noise + 1e-8) # sample Y according to the standard gaussian process formula sample("y", dist.MultivariateNormal(loc=np.zeros(N), covariance_matrix=K), obs=y)
def GPC(X, y): N = y.shape[0] # Priors. alpha = numpyro.sample('alpha', dist.LogNormal(0, 1)) rho = numpyro.sample('rho', dist.LogNormal(0, 1)) beta = numpyro.sample('beta', dist.Normal(0, 1)) eta = numpyro.sample('eta', dist.Normal(np.zeros(N), 1)) # Latent function. f = compute_f(alpha, rho, beta, eta, X, 1e-3) # Likelihood. numpyro.sample('obs', dist.Bernoulli(logits=f), obs=y)
def GP(X, y): # Set informative log-normal priors on kernel hyperparameters. variance = numpyro.sample("kernel_var", dist.LogNormal(0.0, 0.1)) lengthscale = numpyro.sample("kernel_length", dist.LogNormal(0.0, 1.0)) sigma = numpyro.sample("sigma", dist.LogNormal(0.0, 1.0)) # Compute kernel K = squared_exp_cov_1D(X, variance, lengthscale) K += np.eye(X.shape[0]) * np.power(sigma, 2) # Sample y according to the standard gaussian process formula numpyro.sample("y", dist.MultivariateNormal(loc=np.zeros(X.shape[0]), covariance_matrix=K), obs=y)
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 test_numpyro_dict_priors_defaults_numpyro(): demo_priors = { "lengthscale": dist.LogNormal(loc=0.0, scale=1.0), "variance": dist.LogNormal(loc=0.0, scale=1.0), "obs_noise": dist.LogNormal(loc=0.0, scale=1.0), } numpyro_params = numpyro_dict_params(demo_priors) assert set(numpyro_params) == set(demo_priors.keys()) for ikey, iparam in demo_priors.items(): # check keys exist for param assert set(numpyro_params[ikey].keys()) == set(("prior", "param_type")) # check init value is the same as initial value chex.assert_equal(numpyro_params[ikey]["prior"], iparam)
def test_model_with_transformed_distribution(): x_prior = dist.HalfNormal(2) y_prior = dist.LogNormal(scale=3.) # transformed distribution def model(): numpyro.sample('x', x_prior) numpyro.sample('y', y_prior) params = {'x': jnp.array(-5.), 'y': jnp.array(7.)} model = handlers.seed(model, random.PRNGKey(0)) inv_transforms = { 'x': biject_to(x_prior.support), 'y': biject_to(y_prior.support) } expected_samples = partial(transform_fn, inv_transforms)(params) expected_potential_energy = (-x_prior.log_prob(expected_samples['x']) - y_prior.log_prob(expected_samples['y']) - inv_transforms['x'].log_abs_det_jacobian( params['x'], expected_samples['x']) - inv_transforms['y'].log_abs_det_jacobian( params['y'], expected_samples['y'])) reparam_model = handlers.reparam(model, {'y': TransformReparam()}) base_params = {'x': params['x'], 'y_base': params['y']} actual_samples = constrain_fn(handlers.seed(reparam_model, random.PRNGKey(0)), (), {}, base_params, return_deterministic=True) actual_potential_energy = potential_energy(reparam_model, (), {}, base_params) assert_allclose(expected_samples['x'], actual_samples['x']) assert_allclose(expected_samples['y'], actual_samples['y']) assert_allclose(actual_potential_energy, expected_potential_energy)
def test_lift(): def model(): loc1 = numpyro.param("loc1", 0.0) scale1 = numpyro.param("scale1", 1.0, constraint=constraints.positive) numpyro.sample("latent1", dist.Normal(loc1, scale1)) loc2 = numpyro.param("loc2", 1.0) scale2 = numpyro.param("scale2", 2.0, constraint=constraints.positive) latent2 = numpyro.sample("latent2", dist.Normal(loc2, scale2)) return latent2 loc1_prior = dist.Normal() scale1_prior = dist.LogNormal() prior = {"loc1": loc1_prior, "scale1": scale1_prior} with handlers.trace() as tr: with handlers.seed(rng_seed=1): model() with handlers.trace() as lifted_tr: with handlers.seed(rng_seed=2): with handlers.lift(prior=prior): model() for name in tr.keys(): assert name in lifted_tr if name in prior: assert lifted_tr[name]["fn"] is prior[name] assert lifted_tr[name]["type"] == "sample" assert lifted_tr[name]["value"] not in (0.0, 1.0) elif name in ("loc2", "scale2"): assert lifted_tr[name]["type"] == "param"
def model(T, q=1, r=1, phi=0., beta=0.): x = 0. mu = 0. for i in range(T): x = numpyro.sample(f'x_{i}', dist.LogNormal(phi * x, q)) mu = beta * mu + x numpyro.sample(f'y_{i}', dist.Normal(mu, r))
def model( covariates: jnp.ndarray, x: Optional[jnp.ndarray] = None, x_dim: int = 1, z_dim: int = 1, ) -> None: if x is not None: x_dim = x.shape[-1] seq_len, batch, c_dim = covariates.shape weight = numpyro.sample( "weight", dist.Normal(np.zeros((c_dim, x_dim)), np.ones((c_dim, x_dim)) * 0.1)) bias = numpyro.sample("bias", dist.Normal(np.zeros(x_dim), np.ones(x_dim) * 10)) sigma = numpyro.sample( "sigma", dist.LogNormal(-5 * np.ones(x_dim), 5 * np.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(z_prev, jnp.ones(z_dim))) numpyro.sample( "x", dist.Cauchy(z + jnp.matmul(covariates[t], weight) + bias, sigma)) return (z, ), None with numpyro.handlers.condition(data={"x": x}): scan(transition_fn, (jnp.zeros((batch, z_dim)), ), jnp.arange(seq_len))
def test_model_with_transformed_distribution(): x_prior = dist.HalfNormal(2) y_prior = dist.LogNormal(scale=3.) # transformed distribution def model(): sample('x', x_prior) sample('y', y_prior) params = {'x': np.array(-5.), 'y': np.array(7.)} model = seed(model, random.PRNGKey(0)) inv_transforms = { 'x': biject_to(x_prior.support), 'y': biject_to(y_prior.support) } expected_samples = partial(transform_fn, inv_transforms)(params) expected_potential_energy = (-x_prior.log_prob(expected_samples['x']) - y_prior.log_prob(expected_samples['y']) - inv_transforms['x'].log_abs_det_jacobian( params['x'], expected_samples['x']) - inv_transforms['y'].log_abs_det_jacobian( params['y'], expected_samples['y'])) base_inv_transforms = { 'x': biject_to(x_prior.support), 'y': biject_to(y_prior.base_dist.support) } actual_samples = constrain_fn(seed(model, random.PRNGKey(0)), (), {}, base_inv_transforms, params) actual_potential_energy = potential_energy(model, (), {}, base_inv_transforms, params) assert_allclose(expected_samples['x'], actual_samples['x']) assert_allclose(expected_samples['y'], actual_samples['y']) assert_allclose(actual_potential_energy, expected_potential_energy)
def test_numpyro_marginal_ll_numpyro_priors_type(n_samples, n_features, n_latents, dtype): # create sample data ds = _gen_training_data(n_samples, n_features, n_latents) # convert to tyle ds = jax.tree_util.tree_map(lambda x: x.astype(dtype), ds) # initialize parameters params, posterior = _get_conjugate_posterior_params() # convert to numpyro-style params numpyro_params = numpyro_dict_params(params) # convert to priors numpyro_params = add_priors(numpyro_params, dist.LogNormal(0.0, 10.0)) # initialize numpyro-style GP model npy_model = numpyro_marginal_ll(posterior, numpyro_params) # do one forward pass with context with numpyro.handlers.seed(rng_seed=KEY): pred = npy_model(ds) chex.assert_equal(pred.dtype, ds.y.dtype)
def test_lift(): def model(): loc1 = numpyro.param("loc1", 0.) scale1 = numpyro.param("scale1", 1., constraint=constraints.positive) numpyro.sample("latent1", dist.Normal(loc1, scale1)) loc2 = numpyro.param("loc2", 1.) scale2 = numpyro.param("scale2", 2., constraint=constraints.positive) latent2 = numpyro.sample("latent2", dist.Normal(loc2, scale2)) return latent2 loc1_prior = dist.Normal() scale1_prior = dist.LogNormal() prior = {"loc1": loc1_prior, "scale1": scale1_prior} with handlers.trace() as tr: with handlers.seed(rng_seed=1): model() with handlers.trace() as lifted_tr: with handlers.seed(rng_seed=2): with handlers.lift(prior=prior): model() for name in tr.keys(): assert name in lifted_tr if name in prior: assert lifted_tr[name]['fn'] is prior[name] assert lifted_tr[name]['type'] == 'sample' assert lifted_tr[name]['value'] not in (0., 1.) elif name in ('loc2', 'scale2'): assert lifted_tr[name]['type'] == 'param'
def model(T, q=1, r=1, phi=0.0, beta=0.0): x = 0.0 mu = 0.0 for i in range(T): x = numpyro.sample(f"x_{i}", dist.LogNormal(phi * x, q)) mu = beta * mu + x numpyro.sample(f"y_{i}", dist.Normal(mu, r))
def create_model(yT, yC, num_components): # Cosntants nC = yC.shape[0] nT = yT.shape[0] zC = jnp.isinf(yC).sum().item() zT = jnp.isinf(yT).sum().item() yT_finite = yT[jnp.isinf(yT) == False] yC_finite = yC_finite = yC[jnp.isinf(yC) == False] K = num_components p = numpyro.sample('p', dist.Beta(.5, .5)) gammaC = numpyro.sample('gammaC', dist.Beta(1, 1)) gammaT = numpyro.sample('gammaT', dist.Beta(1, 1)) etaC = numpyro.sample('etaC', dist.Dirichlet(jnp.ones(K) / K)) etaT = numpyro.sample('etaT', dist.Dirichlet(jnp.ones(K) / K)) with numpyro.plate('mixutre_components', K): nu = numpyro.sample('nu', dist.LogNormal(3.5, 0.5)) mu = numpyro.sample('mu', dist.Normal(0, 3)) sigma = numpyro.sample('sigma', dist.LogNormal(0, .5)) phi = numpyro.sample('phi', dist.Normal(0, 3)) gammaT_star = simulate_data.compute_gamma_T_star(gammaC, gammaT, p) etaT_star = simulate_data.compute_eta_T_star(etaC, etaT, p, gammaC, gammaT, gammaT_star) with numpyro.plate('y_C', nC - zC): numpyro.sample('finite_obs_C', Mix(nu[None, :], mu[None, :], sigma[None, :], phi[None, :], etaC[None, :]), obs=yC_finite[:, None]) with numpyro.plate('y_T', nT - zT): numpyro.sample('finite_obs_T', Mix(nu[None, :], mu[None, :], sigma[None, :], phi[None, :], etaT_star[None, :]), obs=yT_finite[:, None]) numpyro.sample('N_C', dist.Binomial(nC, gammaC), obs=zC) numpyro.sample('N_T', dist.Binomial(nT, gammaT_star), obs=zT)
def LogNormalFromInterval(low, high): """This assumes a centered 90% confidence interval, i.e. the left endpoint marks 0.05% on the CDF, the right 0.95%.""" loghigh = math.log(high) loglow = math.log(low) mean = (loghigh + loglow) / 2 stdev = (loghigh - loglow) / (2 * 1.645) return dist.LogNormal(mean, stdev)
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 model(self, batch): X = batch['X'] y = batch['y'] N, D = X.shape # set uninformative log-normal priors var = sample('kernel_var', dist.LogNormal(0.0, 1.0), sample_shape = (1,)) length = sample('kernel_length', dist.LogNormal(0.0, 1.0), sample_shape = (D,)) theta = np.concatenate([var, length]) # compute kernel K = self.kernel(X, X, theta) + np.eye(N)*1e-8 L = cholesky(K, lower=True) # Generate latent function beta = sample('beta', dist.Normal(0.0, 1.0)) eta = sample('eta', dist.Normal(0.0, 1.0), sample_shape=(N,)) f = np.matmul(L, eta) + beta # Bernoulli likelihood sample('y', dist.Bernoulli(logits=f), obs=y)
def model(self, batch): X = batch['X'] y = batch['y'] N = y.shape[0] dim_X = X.shape[1] dim_H = self.dim_H D = dim_X + dim_H # Generate latent inputs H = sample('H', dist.Normal(np.zeros((N, dim_H)), np.ones((N, dim_H)))) X = np.concatenate([X, H], axis = 1) # set uninformative log-normal priors on GP hyperparameters var = sample('kernel_var', dist.LogNormal(0.0, 10.0)) length = sample('kernel_length', dist.LogNormal(np.zeros(D), 10.0*np.ones(D))) noise = sample('noise_var', dist.LogNormal(0.0, 10.0)) theta = np.concatenate([np.array([var]), np.array(length)]) # compute kernel K = self.kernel(X, X, theta) + np.eye(N)*(noise + 1e-8) # sample Y according to the GP likelihood sample("y", dist.MultivariateNormal(loc=np.zeros(N), covariance_matrix=K), obs=y)
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 model(data): with numpyro.plate("states", dim): transition = numpyro.sample("transition", dist.Dirichlet(jnp.ones(dim))) emission_loc = numpyro.sample("emission_loc", dist.Normal(0, 1)) emission_scale = numpyro.sample("emission_scale", dist.LogNormal(0, 1)) trans_prob = numpyro.sample("initialize", dist.Dirichlet(jnp.ones(dim))) for t, y in markov(enumerate(data)): x = numpyro.sample("x_{}".format(t), dist.Categorical(trans_prob)) numpyro.sample("y_{}".format(t), dist.Normal(emission_loc[x], emission_scale[x]), obs=y) trans_prob = transition[x]
def model(batch, subsample, full_size): drift = numpyro.sample("drift", dist.LogNormal(-1, 0.5)) with handlers.substitute(data={"data": subsample}): plate = numpyro.plate("data", full_size, subsample_size=len(subsample)) assert plate.size == 50 def transition_fn(z_prev, y_curr): with plate: z_curr = numpyro.sample("state", dist.Normal(z_prev, drift)) y_curr = numpyro.sample("obs", dist.Bernoulli(logits=z_curr), obs=y_curr) return z_curr, y_curr _, result = scan(transition_fn, jnp.zeros(len(subsample)), batch, length=num_time_steps) return result
def model2(): data = [np.array([-1.0, -1.0, 0.0]), np.array([-1.0, 1.0])] p = numpyro.param("p", np.array([0.25, 0.75])) loc = numpyro.sample("loc", dist.Normal(0, 1).expand([2]).to_event(1)) # FIXME results in infinite loop in transformeddist_to_funsor. # scale = numpyro.sample("scale", dist.LogNormal(0, 1)) z1 = numpyro.sample("z1", dist.Categorical(p)) scale = numpyro.sample("scale", dist.LogNormal(jnp.array([0.0, 1.0])[z1], 1)) with numpyro.plate("data[0]", 3): numpyro.sample("x1", dist.Normal(loc[z1], scale), obs=data[0]) with numpyro.plate("data[1]", 2): z2 = numpyro.sample("z2", dist.Categorical(p)) numpyro.sample("x2", dist.Normal(loc[z2], scale), obs=data[1])
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 model(N, y=None): """ :param int N: number of measurement times :param numpy.ndarray y: measured populations with shape (N, 2) """ # initial population z_init = numpyro.sample("z_init", dist.LogNormal(jnp.log(10), 1), sample_shape=(2,)) # measurement times ts = jnp.arange(float(N)) # parameters alpha, beta, gamma, delta of dz_dt theta = numpyro.sample( "theta", dist.TruncatedNormal(low=0., loc=jnp.array([0.5, 0.05, 1.5, 0.05]), scale=jnp.array([0.5, 0.05, 0.5, 0.05]))) # integrate dz/dt, the result will have shape N x 2 z = odeint(dz_dt, z_init, ts, theta, rtol=1e-5, atol=1e-3, mxstep=500) # measurement errors, we expect that measured hare has larger error than measured lynx sigma = numpyro.sample("sigma", dist.Exponential(jnp.array([1, 2]))) # measured populations (in log scale) numpyro.sample("y", dist.Normal(jnp.log(z), sigma), obs=y)