def guide(): loc = numpyro.param("loc", np.zeros(3)) cov = numpyro.param("cov", np.eye(3), constraint=constraints.positive_definite) x = numpyro.sample("x", dist.MultivariateNormal(loc, cov)) with numpyro.plate("plate", len(data)): with handlers.mask(mask=np.invert(mask)): numpyro.sample("y_unobserved", dist.MultivariateNormal(x, np.eye(3)))
def model(): x = numpyro.sample("x", dist.MultivariateNormal(np.zeros(3), np.eye(3))) with numpyro.plate("plate", len(data)): y = numpyro.sample("y", dist.MultivariateNormal(x, np.eye(3)), obs=data, obs_mask=mask) if not_jax_tracer(y): assert ((y == data).all(-1) == mask).all()
def model(): x0 = numpyro.sample( "x0", dist.MultivariateNormal(loc=jnp.zeros(D), covariance_matrix=cov00) ) posterior_loc1 = jnp.matmul(cov_10_cov00_inv, x0) numpyro.sample( "x1", dist.MultivariateNormal( loc=posterior_loc1, covariance_matrix=posterior_cov1 ), )
def kernel(x, y): if self.precond_mode == 'const': wxs = jnp.array([1.]) wys = jnp.array([1.]) else: wxs = jax.nn.softmax( jax.vmap(lambda z, q_inv: dist.MultivariateNormal(z, posdef(q_inv)).log_prob(x))(particles, qs_inv)) wys = jax.nn.softmax( jax.vmap(lambda z, q_inv: dist.MultivariateNormal(z, posdef(q_inv)).log_prob(y))(particles, qs_inv)) return jnp.sum( jax.vmap(lambda qs, qis, wx, wy: wx * wy * (qis @ inner_kernel(qs @ x, qs @ y) @ qis.transpose()))( qs_sqrt, qs_inv_sqrt, wxs, wys), axis=0)
def model(cov): w = numpyro.sample("w", dist.Normal(0, 1000).expand([2]).to_event(1)) x = numpyro.sample("x", dist.Normal(0, 1000).expand([1]).to_event(1)) y = numpyro.sample("y", dist.Normal(0, 1000).expand([1]).to_event(1)) z = numpyro.sample("z", dist.Normal(0, 1000).expand([1]).to_event(1)) wxyz = jnp.concatenate([w, x, y, z]) numpyro.sample("obs", dist.MultivariateNormal(jnp.zeros(5), cov), obs=wxyz)
def twoh_c_kf(T=None, T_forecast=15, obs=None): """Define Kalman Filter with two hidden variates.""" T = len(obs) if T is None else T # Define priors over beta, tau, sigma, z_1 (keep the shapes in mind) #W = numpyro.sample(name="W", fn=dist.Normal(loc=jnp.zeros((2,4)), scale=jnp.ones((2,4)))) beta = numpyro.sample(name="beta", fn=dist.Normal(loc=jnp.array([0.,0.]), scale=jnp.ones(2))) tau = numpyro.sample(name="tau", fn=dist.HalfCauchy(scale=jnp.ones(2))) sigma = numpyro.sample(name="sigma", fn=dist.HalfCauchy(scale=.1)) z_prev = numpyro.sample(name="z_1", fn=dist.Normal(loc=jnp.zeros(2), scale=jnp.ones(2))) # Define LKJ prior L_Omega = numpyro.sample("L_Omega", dist.LKJCholesky(2, 10.)) Sigma_lower = jnp.matmul(jnp.diag(jnp.sqrt(tau)), L_Omega) # lower cholesky factor of the covariance matrix noises = numpyro.sample("noises", fn=dist.MultivariateNormal(loc=jnp.zeros(2), scale_tril=Sigma_lower), sample_shape=(T+T_forecast,)) # Propagate the dynamics forward using jax.lax.scan carry = (beta, z_prev, tau) z_collection = [z_prev] carry, zs_exp = lax.scan(f, carry, noises, T+T_forecast) z_collection = jnp.concatenate((jnp.array(z_collection), zs_exp), axis=0) c = numpyro.sample(name="c", fn=dist.Normal(loc=jnp.array([[0.], [0.]]), scale=jnp.ones((2,1)))) obs_mean = jnp.dot(z_collection[:T,:], c).squeeze() pred_mean = jnp.dot(z_collection[T:,:], c).squeeze() # Sample the observed y (y_obs) numpyro.sample(name="y_obs", fn=dist.Normal(loc=obs_mean, scale=sigma), obs=obs) numpyro.sample(name="y_pred", fn=dist.Normal(loc=pred_mean, scale=sigma), obs=None)
def get_posterior(self, params): """ Returns a multivariate Normal posterior distribution. """ transform = self.get_transform(params) return dist.MultivariateNormal(transform.loc, scale_tril=transform.scale_tril)
def sample_kernel(sa_state, model_args=(), model_kwargs=None): pe_fn = potential_fn if potential_fn_gen: pe_fn = potential_fn_gen(*model_args, **model_kwargs) zs, pes, loc, scale = sa_state.adapt_state # we recompute loc/scale after each iteration to avoid precision loss # XXX: consider to expose a setting to do this job periodically # to save some computations loc = jnp.mean(zs, 0) if scale.ndim == 2: cov = jnp.cov(zs, rowvar=False, bias=True) if cov.shape == (): # JAX returns scalar for 1D input cov = cov.reshape((1, 1)) cholesky = jnp.linalg.cholesky(cov) scale = jnp.where(jnp.any(jnp.isnan(cholesky)), scale, cholesky) else: scale = jnp.std(zs, 0) rng_key, rng_key_z, rng_key_reject, rng_key_accept = random.split(sa_state.rng_key, 4) _, unravel_fn = ravel_pytree(sa_state.z) z = loc + _sample_proposal(scale, rng_key_z) pe = pe_fn(unravel_fn(z)) pe = jnp.where(jnp.isnan(pe), jnp.inf, pe) diverging = (pe - sa_state.potential_energy) > max_delta_energy # NB: all terms having the pattern *s will have shape N x ... # and all terms having the pattern *s_ will have shape (N + 1) x ... locs, scales = _get_proposal_loc_and_scale(zs, loc, scale, z) zs_ = jnp.concatenate([zs, z[None, :]]) pes_ = jnp.concatenate([pes, pe[None]]) locs_ = jnp.concatenate([locs, loc[None, :]]) scales_ = jnp.concatenate([scales, scale[None, ...]]) if scale.ndim == 2: # dense_mass log_weights_ = dist.MultivariateNormal(locs_, scale_tril=scales_).log_prob(zs_) + pes_ else: log_weights_ = dist.Normal(locs_, scales_).log_prob(zs_).sum(-1) + pes_ # mask invalid values (nan, +inf) by -inf log_weights_ = jnp.where(jnp.isfinite(log_weights_), log_weights_, -jnp.inf) # get rejecting index j = random.categorical(rng_key_reject, log_weights_) zs = _numpy_delete(zs_, j) pes = _numpy_delete(pes_, j) loc = locs_[j] scale = scales_[j] adapt_state = SAAdaptState(zs, pes, loc, scale) # NB: weights[-1] / sum(weights) is the probability of rejecting the new sample `z`. accept_prob = 1 - jnp.exp(log_weights_[-1] - logsumexp(log_weights_)) itr = sa_state.i + 1 n = jnp.where(sa_state.i < wa_steps, itr, itr - wa_steps) mean_accept_prob = sa_state.mean_accept_prob + (accept_prob - sa_state.mean_accept_prob) / n # XXX: we make a modification of SA sampler in [1] # in [1], each MCMC state contains N points `zs` # here we do resampling to pick randomly a point from those N points k = random.categorical(rng_key_accept, jnp.zeros(zs.shape[0])) z = unravel_fn(zs[k]) pe = pes[k] return SAState(itr, z, pe, accept_prob, mean_accept_prob, diverging, adapt_state, rng_key)
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 likelihood(self, X, Y): """ The likelihood model for Gaussian Process :param X: Features :param Y: Targets :return: Samples from the a Gaussian Process model. """ params_samples = {} for param in self.params_priors: if "distribution" in str(type(self.params_priors[param])): params_samples[param] = numpyro.sample(param, self.params_priors[param]) if type(self.params_priors[param]) == float: params_samples[param] = numpyro.param(param, self.params_priors[param]) noise = params_samples.get("noise", 0) k = self.kernel(X, X, params_samples) + (noise + self.jitter) * np.eye( X.shape[0] ) # 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 obyo(y, tag, nusdx, nus, mdbCO, mdbH2O, cdbH2H2, cdbH2He): #CO SijM_CO, ngammaLM_CO, nsigmaDl_CO = exomol(mdbCO, Tarr, Parr, R_CO, molmassCO) xsm_CO = xsmatrix(cnu_CO, indexnu_CO, R_CO, pmarray_CO, nsigmaDl_CO, ngammaLM_CO, SijM_CO, nus, dgm_ngammaL_CO) dtaumCO = dtauM(dParr, jnp.abs(xsm_CO), MMR_CO * ONEARR, molmassCO, g) #H2O SijM_H2O, ngammaLM_H2O, nsigmaDl_H2O = exomol(mdbH2O, Tarr, Parr, R_H2O, molmassH2O) xsm_H2O = xsmatrix(cnu_H2O, indexnu_H2O, R_H2O, pmarray_H2O, nsigmaDl_H2O, ngammaLM_H2O, SijM_H2O, nus, dgm_ngammaL_H2O) dtaumH2O = dtauM(dParr, jnp.abs(xsm_H2O), MMR_H2O * ONEARR, molmassH2O, g) #CIA dtaucH2H2=dtauCIA(nus,Tarr,Parr,dParr,vmrH2,vmrH2,\ mmw,g,cdbH2H2.nucia,cdbH2H2.tcia,cdbH2H2.logac) dtaucH2He=dtauCIA(nus,Tarr,Parr,dParr,vmrH2,vmrHe,\ mmw,g,cdbH2He.nucia,cdbH2He.tcia,cdbH2He.logac) dtau = dtaumCO + dtaumH2O + dtaucH2H2 + dtaucH2He sourcef = planck.piBarr(Tarr, nus) Ftoa = Fref / Rp**2 F0 = rtrun(dtau, sourcef) / baseline / Ftoa Frot = response.rigidrot(nus, F0, vsini, u1, u2) mu = response.ipgauss_sampling(nusdx, nus, Frot, beta, RV) cov = gpkernel_RBF(nu1, tau, a, e1) sample(tag, dist.MultivariateNormal(loc=mu, covariance_matrix=cov), obs=y)
def model(X, Y, hypers): S, P, N = hypers['expected_sparsity'], X.shape[1], X.shape[0] sigma = numpyro.sample("sigma", dist.HalfNormal(hypers['alpha3'])) phi = sigma * (S / np.sqrt(N)) / (P - S) eta1 = numpyro.sample("eta1", dist.HalfCauchy(phi)) msq = numpyro.sample("msq", dist.InverseGamma(hypers['alpha1'], hypers['beta1'])) xisq = numpyro.sample("xisq", dist.InverseGamma(hypers['alpha2'], hypers['beta2'])) eta2 = np.square(eta1) * np.sqrt(xisq) / msq lam = numpyro.sample("lambda", dist.HalfCauchy(np.ones(P))) kappa = np.sqrt(msq) * lam / np.sqrt(msq + np.square(eta1 * lam)) # sample observation noise var_obs = numpyro.sample( "var_obs", dist.InverseGamma(hypers['alpha_obs'], hypers['beta_obs'])) # compute kernel kX = kappa * X k = kernel(kX, kX, eta1, eta2, hypers['c']) + var_obs * np.eye(N) assert k.shape == (N, N) # 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 model(self, home_team, away_team): sigma_a = pyro.sample("sigma_a", dist.HalfNormal(1.0)) sigma_b = pyro.sample("sigma_b", dist.HalfNormal(1.0)) mu_b = pyro.sample("mu_b", dist.Normal(0.0, 1.0)) rho_raw = pyro.sample("rho_raw", dist.Beta(2, 2)) rho = 2.0 * rho_raw - 1.0 log_gamma = pyro.sample("log_gamma", dist.Normal(0, 1)) with pyro.plate("teams", self.n_teams): abilities = pyro.sample( "abilities", dist.MultivariateNormal( np.array([0.0, mu_b]), covariance_matrix=np.array([ [sigma_a**2.0, rho * sigma_a * sigma_b], [rho * sigma_a * sigma_b, sigma_b**2.0], ]), ), ) log_a = abilities[:, 0] log_b = abilities[:, 1] 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.exp(log_a[home_inds] + log_b[away_inds] + log_gamma) away_rate = np.exp(log_a[away_inds] + log_b[home_inds]) pyro.sample("home_goals", dist.Poisson(home_rate).to_event(1)) pyro.sample("away_goals", dist.Poisson(away_rate).to_event(1))
def sample_posterior(self, rng_key, params, sample_shape=()): transform = self._get_transform(params) loc, scale_tril = transform.loc, transform.scale_tril latent_sample = dist.MultivariateNormal(loc, scale_tril=scale_tril).sample( rng_key, sample_shape) return self._unpack_and_constrain(latent_sample, params)
def rethinking_model(B, M, K): # priors a = numpyro.sample("a", dist.Normal(0, 0.5)) muB = numpyro.sample("muB", dist.Normal(0, 0.5)) muM = numpyro.sample("muM", dist.Normal(0, 0.5)) bB = numpyro.sample("bB", dist.Normal(0, 0.5)) bM = numpyro.sample("bM", dist.Normal(0, 0.5)) sigma = numpyro.sample("sigma", dist.Exponential(1)) Rho_BM = numpyro.sample("Rho_BM", dist.LKJ(2, 2)) Sigma_BM = numpyro.sample("Sigma_BM", dist.Exponential(1).expand([2])) # define B_merge as mix of observed and imputed values B_impute = numpyro.sample( "B_impute", dist.Normal(0, 1).expand([int(np.isnan(B).sum())]).mask(False)) B_merge = ops.index_update(B, np.nonzero(np.isnan(B))[0], B_impute) # M and B correlation MB = jnp.stack([M, B_merge], axis=1) cov = jnp.outer(Sigma_BM, Sigma_BM) * Rho_BM numpyro.sample("MB", dist.MultivariateNormal(jnp.stack([muM, muB]), cov), obs=MB) # K as function of B and M mu = a + bB * B_merge + bM * M numpyro.sample("K", dist.Normal(mu, sigma), obs=K)
def model(X, Y, hypers): S, P, N = hypers["expected_sparsity"], X.shape[1], X.shape[0] sigma = numpyro.sample("sigma", dist.HalfNormal(hypers["alpha3"])) phi = sigma * (S / jnp.sqrt(N)) / (P - S) eta1 = numpyro.sample("eta1", dist.HalfCauchy(phi)) msq = numpyro.sample("msq", dist.InverseGamma(hypers["alpha1"], hypers["beta1"])) xisq = numpyro.sample("xisq", dist.InverseGamma(hypers["alpha2"], hypers["beta2"])) eta2 = jnp.square(eta1) * jnp.sqrt(xisq) / msq lam = numpyro.sample("lambda", dist.HalfCauchy(jnp.ones(P))) kappa = jnp.sqrt(msq) * lam / jnp.sqrt(msq + jnp.square(eta1 * lam)) # compute kernel kX = kappa * X k = kernel(kX, kX, eta1, eta2, hypers["c"]) + sigma**2 * jnp.eye(N) assert k.shape == (N, N) # 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 __call__(self): assignment = numpyro.sample("assignment", dist.Categorical(self.weights)) loc = self.loc[assignment] cov = self.cov[assignment] nu_max = numpyro.sample("nu_max", self.nu_max) log_nu_max = jnp.log10(nu_max) teff = numpyro.sample("teff", self.teff) loc0101 = loc[0:2] cov0101 = jnp.array([[cov[0, 0], cov[0, 1]], [cov[1, 0], cov[1, 1]]]) L = jax.scipy.linalg.cho_factor(cov0101, lower=True) A = jax.scipy.linalg.cho_solve(L, jnp.array([log_nu_max, teff]) - loc0101) loc2323 = loc[2:] cov2323 = jnp.array([[cov[2, 2], cov[2, 3]], [cov[3, 2], cov[3, 3]]]) cov0123 = jnp.array([[cov[0, 2], cov[1, 2]], [cov[0, 3], cov[1, 3]]]) v = jax.scipy.linalg.cho_solve(L, cov0123.T) cond_loc = loc2323 + jnp.dot(cov0123, A) cond_cov = ( cov2323 - jnp.dot(cov0123, v) + self.noise * jnp.eye(2) # Add white noise ) numpyro.sample("log_tau", dist.MultivariateNormal(cond_loc, cond_cov))
def parametric_draws(subposteriors, num_draws, diagonal=False, rng_key=None): """ Merges subposteriors following (embarrassingly parallel) parametric Monte Carlo algorithm. **References:** 1. *Asymptotically Exact, Embarrassingly Parallel MCMC*, Willie Neiswanger, Chong Wang, Eric Xing :param list subposteriors: a list in which each element is a collection of samples. :param int num_draws: number of draws from the merged posterior. :param bool diagonal: whether to compute weights using variance or covariance, defaults to `False` (using covariance). :param jax.random.PRNGKey rng_key: source of the randomness, defaults to `jax.random.PRNGKey(0)`. :return: a collection of `num_draws` samples with the same data structure as each subposterior. """ rng_key = random.PRNGKey(0) if rng_key is None else rng_key if diagonal: mean, var = parametric(subposteriors, diagonal=True) samples_flat = dist.Normal(mean, np.sqrt(var)).sample( rng_key, (num_draws, )) else: mean, cov = parametric(subposteriors, diagonal=False) samples_flat = dist.MultivariateNormal(mean, cov).sample( rng_key, (num_draws, )) _, unravel_fn = ravel_pytree(tree_map(lambda x: x[0], subposteriors[0])) return vmap(lambda x: unravel_fn(x))(samples_flat)
def mll(ds: Dataset): x, y = ds.X, ds.y params = {} for iname, iparam in numpyro_params.items(): if iparam["param_type"] == "prior": params[iname] = numpyro.sample(name=iname, fn=iparam["prior"]) else: params[iname] = numpyro.param( name=iname, init_value=iparam["init_value"], constraint=iparam["constraint"], ) # get mean function mu = gp.prior.mean_function(x) # covariance function gram_matrix = gram(gp.prior.kernel, x, params) gram_matrix += params["obs_noise"] * I(x.shape[0]) # scale triangular matrix L = cholesky(gram_matrix, lower=True) return numpyro.sample( "y", dist.MultivariateNormal(loc=mu, scale_tril=L), obs=y.squeeze(), )
def _get_posterior(self): loc = numpyro.param('{}_loc'.format(self.prefix), self._init_latent) scale_tril = numpyro.param('{}_scale_tril'.format(self.prefix), jnp.identity(self.latent_dim) * self._init_scale, constraint=constraints.lower_cholesky) return dist.MultivariateNormal(loc, scale_tril=scale_tril)
def gaussian_gibbs_fn(rng_key, hmc_sites, gibbs_sites): x1 = hmc_sites['x1'] posterior_loc0 = jnp.matmul(cov_01_cov11_inv, x1) x0_proposal = dist.MultivariateNormal( loc=posterior_loc0, covariance_matrix=posterior_cov0).sample(rng_key) return {'x0': x0_proposal}
def model(self, *views: np.ndarray): n = views[0].shape[0] p = [view.shape[1] for view in views] # mean of column in each view of data (p_1,) mu = [ numpyro.sample("mu_" + str(i), dist.MultivariateNormal(0., 10 * jnp.eye(p_))) for i, p_ in enumerate(p) ] """ Generates cholesky factors of correlation matrices using an LKJ prior. The expected use is to combine it with a vector of variances and pass it to the scale_tril parameter of a multivariate distribution such as MultivariateNormal. E.g., if theta is a (positive) vector of covariances with the same dimensionality as this distribution, and Omega is sampled from this distribution, scale_tril=torch.mm(torch.diag(sqrt(theta)), Omega) """ psi = [ numpyro.sample("psi_" + str(i), dist.LKJCholesky(p_)) for i, p_ in enumerate(p) ] # sample weights to get from latent to data space (k,p) with numpyro.plate("plate_views", self.latent_dims): self.weights_list = [ numpyro.sample( "W_" + str(i), dist.MultivariateNormal(0., jnp.diag(jnp.ones(p_)))) for i, p_ in enumerate(p) ] with numpyro.plate("plate_i", n): # sample from latent z - normally disributed (n,k) z = numpyro.sample( "z", dist.MultivariateNormal(0., jnp.diag(jnp.ones(self.latent_dims)))) # sample from multivariate normal and observe data [ numpyro.sample("obs" + str(i), dist.MultivariateNormal((z @ W_) + mu_, scale_tril=psi_), obs=X_) for i, ( X_, psi_, mu_, W_) in enumerate(zip(views, psi, mu, self.weights_list)) ]
def marginal_likelihood(self, X, y, noise, is_observed=True, **kwargs): R""" Returns the marginal likelihood distribution, given the input locations `X` and the data `y`. This is integral over the product of the GP prior and a normal likelihood. .. math:: y \mid X,\theta \sim \int p(y \mid f,\, X,\, \theta) \, p(f \mid X,\, \theta) \, df Parameters ---------- name: string Name of the random variable X: array-like Function input values. If one-dimensional, must be a column vector with shape `(n, 1)`. y: array-like Data that is the sum of the function with the GP prior and Gaussian noise. Must have shape `(n, )`. noise: scalar, Variable, or Covariance Standard deviation of the Gaussian noise. Can also be a Covariance for non-white noise. is_observed: bool Whether to set `y` as an `observed` variable in the `model`. Default is `True`. **kwargs Extra keyword arguments that are passed to `MvNormal` distribution constructor. """ if not isinstance(noise, Covariance): self.noise = WhiteNoise(noise) else: self.noise = noise mu, cov = self._build_marginal_likelihood(X) _ = npy.deterministic(f"{self.name}_y", y) if is_observed: return npy.sample( f"{self.name}", dist.MultivariateNormal(loc=mu, covariance_matrix=cov), obs=y, ) else: shape = infer_shape(X, kwargs.pop("shape", None)) return npy.sample( f"{self.name}", dist.MultivariateNormal(loc=mu, covariance_matrix=cov))
def _get_posterior(self): loc = numpyro.param("{}_loc".format(self.prefix), self._init_latent) scale_tril = numpyro.param( "{}_scale_tril".format(self.prefix), jnp.identity(self.latent_dim) * self._init_scale, constraint=self.scale_tril_constraint, ) return dist.MultivariateNormal(loc, scale_tril=scale_tril)
def proposal_dist(z, g): g = -self._preconditioner.flatten(g) dim = jnp.size(g) rho2 = jnp.clip(jnp.dot(g, g), a_min=1.0) covar = (self._mu2 * jnp.eye(dim) + self._lam2_minus_mu2 * jnp.outer(g, g) / jnp.dot(g, g)) / rho2 return dist.MultivariateNormal(loc=self._preconditioner.flatten(z), covariance_matrix=covar)
def multivariate_kf(T=None, T_forecast=15, obs=None): """Define Kalman Filter in a multivariate fashion. The "time-series" are correlated. To define these relationships in a efficient manner, the covarianze matrix of h_t (or, equivalently, the noises) is drown from a Cholesky decomposed matrix. Parameters ---------- T: int T_forecast: int obs: np.array observed variable (infected, deaths...) """ T = len(obs) if T is None else T beta = numpyro.sample( name="beta", fn=dist.Normal(loc=jnp.zeros(2), scale=jnp.ones(2)) ) tau = numpyro.sample(name="tau", fn=dist.HalfCauchy(scale=jnp.ones(2))) sigma = numpyro.sample(name="sigma", fn=dist.HalfCauchy(scale=0.1)) z_prev = numpyro.sample( name="z_1", fn=dist.Normal(loc=jnp.zeros(2), scale=jnp.ones(2)) ) # Define LKJ prior L_Omega = numpyro.sample("L_Omega", dist.LKJCholesky(2, 10.0)) Sigma_lower = jnp.matmul( jnp.diag(jnp.sqrt(tau)), L_Omega ) # lower cholesky factor of the covariance matrix noises = numpyro.sample( "noises", fn=dist.MultivariateNormal(loc=jnp.zeros(2), scale_tril=Sigma_lower), sample_shape=(T + T_forecast - 1,), ) # Propagate the dynamics forward using jax.lax.scan carry = (beta, z_prev, tau) z_collection = [z_prev] carry, zs_exp = lax.scan(f, carry, noises, T + T_forecast - 1) z_collection = jnp.concatenate((jnp.array(z_collection), zs_exp), axis=0) # Sample the observed y (y_obs) and missing y (y_mis) numpyro.sample( name="y_obs1", fn=dist.Normal(loc=z_collection[:T, 0], scale=sigma), obs=obs[:, 0], ) numpyro.sample( name="y_pred1", fn=dist.Normal(loc=z_collection[T:, 0], scale=sigma), obs=None ) numpyro.sample( name="y_obs2", fn=dist.Normal(loc=z_collection[:T, 1], scale=sigma), obs=obs[:, 1], ) numpyro.sample( name="y_pred2", fn=dist.Normal(loc=z_collection[T:, 1], scale=sigma), obs=None )
def model(): a = numpyro.sample("a", dist.Normal(0, 1)) b = numpyro.sample("b", dist.Normal(a[..., None], jnp.ones(3)).to_event(1)) c = numpyro.sample( "c", dist.MultivariateNormal(jnp.zeros(3) + a[..., None], jnp.eye(3)) ) with numpyro.plate("i", 2): d = numpyro.sample("d", dist.Dirichlet(jnp.exp(b + c))) numpyro.sample("e", dist.Categorical(logits=d), obs=jnp.array([0, 0])) return a, b, c, d
def model_c(nu1, y1): Rp = numpyro.sample('Rp', dist.Uniform(0.4, 1.2)) RV = numpyro.sample('RV', dist.Uniform(5.0, 15.0)) MMR_CO = numpyro.sample('MMR_CO', dist.Uniform(0.0, 0.015)) vsini = numpyro.sample('vsini', dist.Uniform(15.0, 25.0)) g = 2478.57730044555 * Mp / Rp**2 # gravity u1 = 0.0 u2 = 0.0 # Layer-by-layer T-P model// lnsT = 6.0 # lnsT = numpyro.sample('lnsT', dist.Uniform(3.0,5.0)) sT = 10**lnsT lntaup = 0.5 # lntaup = numpyro.sample('lntaup', dist.Uniform(0,1)) taup = 10**lntaup cov = modelcov(lnParr, taup, sT) # T0=numpyro.sample('T0', dist.Uniform(1000.0,1100.0)) T0 = numpyro.sample('T0', dist.Uniform(1000, 2000)) Tarr = numpyro.sample( 'Tarr', dist.MultivariateNormal(loc=ONEARR, covariance_matrix=cov)) + T0 # line computation CO qt_CO = vmap(mdbCO.qr_interp)(Tarr) def obyo(y, tag, nusd, nus, numatrix_CO, mdbCO, cdbH2H2): # CO SijM_CO = jit(vmap(SijT, (0, None, None, None, 0)))(Tarr, mdbCO.logsij0, mdbCO.dev_nu_lines, mdbCO.elower, qt_CO) gammaLMP_CO = jit(vmap(gamma_exomol, (0, 0, None, None)))(Parr, Tarr, mdbCO.n_Texp, mdbCO.alpha_ref) gammaLMN_CO = gamma_natural(mdbCO.A) gammaLM_CO = gammaLMP_CO + gammaLMN_CO[None, :] sigmaDM_CO = jit(vmap(doppler_sigma, (None, 0, None)))(mdbCO.dev_nu_lines, Tarr, molmassCO) xsm_CO = xsmatrix(numatrix_CO, sigmaDM_CO, gammaLM_CO, SijM_CO) dtaumCO = dtauM(dParr, xsm_CO, MMR_CO * ONEARR, molmassCO, g) # CIA dtaucH2H2 = dtauCIA(nus, Tarr, Parr, dParr, vmrH2, vmrH2, mmw, g, cdbH2H2.nucia, cdbH2H2.tcia, cdbH2H2.logac) dtau = dtaumCO + dtaucH2H2 sourcef = planck.piBarr(Tarr, nus) F0 = rtrun(dtau, sourcef) / norm Frot = response.rigidrot(nus, F0, vsini, u1, u2) mu = response.ipgauss_sampling(nusd, nus, Frot, beta, RV) numpyro.sample(tag, dist.Normal(mu, sigmain), obs=y) obyo(y1, 'y1', nu1, nus, numatrix_CO, mdbCO, cdbH2H2)
def model_c(nu1, y1, e1): Rp = sample('Rp', dist.Uniform(0.5, 1.5)) Mp = sample('Mp', dist.Normal(33.5, 0.3)) RV = sample('RV', dist.Uniform(26.0, 30.0)) MMR_CO = sample('MMR_CO', dist.Uniform(0.0, maxMMR_CO)) MMR_H2O = sample('MMR_H2O', dist.Uniform(0.0, maxMMR_H2O)) T0 = sample('T0', dist.Uniform(1000.0, 1700.0)) alpha = sample('alpha', dist.Uniform(0.05, 0.15)) vsini = sample('vsini', dist.Uniform(10.0, 20.0)) # Kipping Limb Darkening Prior q1 = sample('q1', dist.Uniform(0.0, 1.0)) q2 = sample('q2', dist.Uniform(0.0, 1.0)) u1, u2 = ld_kipping(q1, q2) #GP logtau = sample('logtau', dist.Uniform(-1.5, 0.5)) #tau=1 <=> 5A tau = 10**(logtau) loga = sample('loga', dist.Uniform(-4.0, -2.0)) a = 10**(loga) #gravity g = getjov_gravity(Rp, Mp) #T-P model// Tarr = T0 * (Parr / Pref)**alpha #CO SijM_CO, gammaLM_CO, sigmaDM_CO = exomol(mdbCO, Tarr, Parr, molmassCO) xsm_CO = xsmatrix(numatrix_CO, sigmaDM_CO, gammaLM_CO, SijM_CO) dtaumCO = dtauM(dParr, xsm_CO, MMR_CO * ONEARR, molmassCO, g) #H2O SijM_H2O, gammaLM_H2O, sigmaDM_H2O = exomol(mdbH2O, Tarr, Parr, molmassH2O) xsm_H2O = xsmatrix(numatrix_H2O, sigmaDM_H2O, gammaLM_H2O, SijM_H2O) dtaumH2O = dtauM(dParr, xsm_H2O, MMR_H2O * ONEARR, molmassH2O, g) #CIA dtaucH2H2=dtauCIA(nus,Tarr,Parr,dParr,vmrH2,vmrH2,\ mmw,g,cdbH2H2.nucia,cdbH2H2.tcia,cdbH2H2.logac) dtaucH2He=dtauCIA(nus,Tarr,Parr,dParr,vmrH2,vmrHe,\ mmw,g,cdbH2He.nucia,cdbH2He.tcia,cdbH2He.logac) dtau = dtaumCO + dtaumH2O + dtaucH2H2 + dtaucH2He sourcef = planck.piBarr(Tarr, nus) Ftoa = Fref / Rp**2 F0 = rtrun(dtau, sourcef) / baseline / Ftoa Frot = response.rigidrot(nus, F0, vsini, u1, u2) mu = response.ipgauss_sampling(nu1, nus, Frot, beta, RV) cov = gpkernel_RBF(nu1, tau, a, e1) sample("y1", dist.MultivariateNormal(loc=mu, covariance_matrix=cov), obs=y1)
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)