def transition_fn(x, y): probs = init_probs if x is None else transition_probs[x] with numpyro.plate("D", D, dim=-1): x = numpyro.sample("x", dist.Categorical(probs)) numpyro.sample("y", dist.Normal(locs[x], 1), obs=y) return x, None
def model(data): mean = numpyro.param('mean', 0.) std = numpyro.param('std', 1., constraint=constraints.positive) return numpyro.sample('obs', dist.Normal(mean, std), obs=data)
def schools_model(): mu = numpyro.sample('mu', dist.Normal(0, 5)) tau = numpyro.sample('tau', dist.HalfCauchy(5)) theta = numpyro.sample('theta', dist.Normal(mu, tau), sample_shape=(data['J'],)) numpyro.sample('obs', dist.Normal(theta, data['sigma']), obs=data['y'])
def model_c(nu1, y1, e1): Rp = numpyro.sample('Rp', dist.Uniform(0.5, 1.5)) Mp = numpyro.sample('Mp', dist.Normal(33.5, 0.3)) sigma = numpyro.sample('sigma', dist.Exponential(1.0)) RV = numpyro.sample('RV', dist.Uniform(25.0, 35.0)) MMR_CO = numpyro.sample('MMR_CO', dist.Uniform(0.0, maxMMR_CO)) MMR_H2O = numpyro.sample('MMR_H2O', dist.Uniform(0.0, maxMMR_H2O)) T0 = numpyro.sample('T0', dist.Uniform(1000.0, 1700.0)) alpha = numpyro.sample('alpha', dist.Uniform(0.05, 0.2)) vsini = numpyro.sample('vsini', dist.Uniform(10.0, 20.0)) g = 2478.57730044555 * Mp / Rp**2 #gravity #Kipping Limb Darkening Prior arxiv:1308.0009 q1 = numpyro.sample('q1', dist.Uniform(0.0, 1.0)) q2 = numpyro.sample('q2', dist.Uniform(0.0, 1.0)) sqrtq1 = jnp.sqrt(q1) u1 = 2.0 * sqrtq1 * q2 u2 = sqrtq1 * (1.0 - 2.0 * q2) #T-P model// Tarr = T0 * (Parr / Pref)**alpha #line computation CO qt_CO = vmap(mdbCO1.qr_interp)(Tarr) qt_H2O = vmap(mdbH2O1.qr_interp)(Tarr) def obyo(y, tag, nusd, nus, numatrix_CO, numatrix_H2O, mdbCO, mdbH2O, cdbH2H2, cdbH2He): #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) #H2O SijM_H2O=jit(vmap(SijT,(0,None,None,None,0)))\ (Tarr,mdbH2O.logsij0,mdbH2O.dev_nu_lines,mdbH2O.elower,qt_H2O) gammaLMP_H2O = jit(vmap(gamma_exomol,(0,0,None,None)))\ (Parr,Tarr,mdbH2O.n_Texp,mdbH2O.alpha_ref) gammaLMN_H2O = gamma_natural(mdbH2O.A) gammaLM_H2O = gammaLMP_H2O + gammaLMN_H2O[None, :] sigmaDM_H2O=jit(vmap(doppler_sigma,(None,0,None)))\ (mdbH2O.dev_nu_lines,Tarr,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(nusd, nus, Frot, beta, RV) errall = jnp.sqrt(e1**2 + sigma**2) numpyro.sample(tag, dist.Normal(mu, errall), obs=y) obyo(y1, "y1", nusd1, nus1, numatrix_CO1, numatrix_H2O1, mdbCO1, mdbH2O1, cdbH2H21, cdbH2He1)
def __call__(self, T=50, N=1e5, T_future=0, E_duration_est=5.5, I_duration_est=3.0, R0_est=3.0, beta_shape=1, sigma_shape=8, gamma_shape=8, det_prob_est=0.15, det_prob_conc=50, det_noise_scale=0.15, rw_scale=2e-1, forecast_rw_scale=0, drift_scale=None, confirmed=None, death=None): ''' Stochastic SEIR model. Draws random parameters and runs dynamics. ''' # Sample initial number of infected individuals I0 = numpyro.sample("I0", dist.Uniform(10, 0.05 * N)) E0 = numpyro.sample("E0", dist.Uniform(10, 0.05 * N)) H0 = numpyro.sample("H0", dist.Uniform(0, 0.01 * N)) D0 = numpyro.sample("D0", dist.Uniform(0, 0.01 * N)) # Sample parameters sigma = numpyro.sample( "sigma", dist.Gamma(sigma_shape, sigma_shape * E_duration_est)) gamma = numpyro.sample( "gamma", dist.Gamma(gamma_shape, gamma_shape * I_duration_est)) # gamma = numpyro.sample("gamma", # dist.TruncatedNormal(loc = 1./I_duration_est, scale = 0.25) beta0 = numpyro.sample( "beta0", dist.Gamma(beta_shape, beta_shape * I_duration_est / R0_est)) det_prob = numpyro.sample( "det_prob", dist.Beta(det_prob_est * det_prob_conc, (1 - det_prob_est) * det_prob_conc)) det_prob_d = numpyro.sample("det_prob_d", dist.Beta(.9 * 100, (1 - .9) * 100)) death_prob = numpyro.sample("death_prob", dist.Beta(.1 * 100, (1 - .1) * 100)) death_rate = numpyro.sample("death_rate", dist.Gamma(10, 10 * 10)) if drift_scale is not None: drift = numpyro.sample("drift", dist.Normal(loc=0, scale=drift_scale)) else: drift = 0 x0 = SEIRDModel.seed(N=N, I=I0, E=E0, H=H0, D=D0) numpyro.deterministic("x0", x0) # Split observations into first and rest confirmed0, confirmed = (None, None) if confirmed is None else ( confirmed[0], confirmed[1:]) death0, death = (None, None) if death is None else (death[0], death[1:]) # First observation with numpyro.handlers.scale(scale_factor=0.5): y0 = observe("y0", x0[6], det_prob, det_noise_scale, obs=confirmed0) with numpyro.handlers.scale(scale_factor=2.0): z0 = observe("z0", x0[5], det_prob_d, det_noise_scale, obs=death0) params = (beta0, sigma, gamma, rw_scale, drift, det_prob, det_noise_scale, death_prob, death_rate, det_prob_d) beta, x, y, z = self.dynamics(T, params, x0, confirmed=confirmed, death=death) x = np.vstack((x0, x)) y = np.append(y0, y) z = np.append(z0, z) if T_future > 0: params = (beta[-1], sigma, gamma, forecast_rw_scale, drift, det_prob, det_noise_scale, death_prob, death_rate, det_prob_d) beta_f, x_f, y_f, z_f = self.dynamics(T_future + 1, params, x[-1, :], suffix="_future") x = np.vstack((x, x_f)) y = np.append(y, y_f) z = np.append(z, z_f) return beta, x, y, z, det_prob, death_prob
def test_block(): with handlers.trace() as trace: with handlers.block(hide=["x"]): with handlers.seed(rng_seed=0): numpyro.sample("x", dist.Normal()) assert "x" not in trace
def model(y): mu = numpyro.sample("mu", dist.Normal()) sigma = numpyro.sample("sigma", dist.HalfNormal()) with numpyro.plate("plate1", y.shape[0]): numpyro.sample("y", dist.Normal(mu, sigma), obs=y)
def model(data, labels): coefs = numpyro.sample('coefs', dist.Normal(np.zeros(dim), np.ones(dim))) offset = numpyro.sample('offset', dist.Uniform(-1, 1)) logits = offset + np.sum(coefs * data, axis=-1) return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
def model(): a = numpyro.param('a', a_init, constraint=constraints.greater_than(a_minval)) b = numpyro.param('b', b_init, constraint=constraints.positive) numpyro.sample('x', dist.Normal(a, b))
def transition_fn(name, probs, locs, x, y): x = numpyro.sample(name, dist.Categorical(probs[x])) numpyro.sample("y_" + name, dist.Normal(locs[x], 1), obs=y) return x, None
def model(data): y_prob = numpyro.sample("y_prob", dist.Beta(1., 1.)) with numpyro.plate("data", data.shape[0]): y = numpyro.sample("y", dist.Bernoulli(y_prob)) z = numpyro.sample("z", dist.Bernoulli(0.65 * y + 0.1)) numpyro.sample("obs", dist.Normal(2. * z, 1.), obs=data)
def model(data): x = w = 0 for i, y in markov(enumerate(data)): x = numpyro.sample(f"x_{i}", dist.Categorical(probs_x[x])) w = numpyro.sample(f"w_{i}", dist.Categorical(probs_w[w])) numpyro.sample(f"y_{i}", dist.Normal(locs[w, x], 1), obs=y)
def transition_fn(x, y): x = numpyro.sample("x", dist.Categorical(probs[w, x])) numpyro.sample("y", dist.Normal(locs[x], 1), obs=y) return x, None
def model(data): w = numpyro.sample("w", dist.Bernoulli(0.6)) x = 0 for i, y in markov(enumerate(data)): x = numpyro.sample(f"x_{i}", dist.Categorical(probs[w, x])) numpyro.sample(f"y_{i}", dist.Normal(locs[x], 1), obs=y)
def guide(subsample): scale = numpyro.param("scale", 1.0) with handlers.substitute(data={"data": subsample}): with numpyro.plate("data", len(data), subsample_size): loc = numpyro.param("loc", jnp.zeros(len(data)), event_dim=0) numpyro.sample("z", dist.Normal(loc, scale))
def actual_model(data): alpha = numpyro.sample('alpha', dist.Uniform(0, 1)) loc = numpyro.sample('loc', dist.Uniform(0, alpha)) numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def guide(): with numpyro.plate_stack("plates", shape): return numpyro.sample("x", dist.Normal(0, 1))
def expected_model(data): alpha = numpyro.sample('alpha', dist.Uniform(0, 1)) loc = numpyro.sample('loc', dist.Uniform(0, 1)) * alpha numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def fn(): return numpyro.sample("x", dist.Normal())
def model(x, y): a = numpyro.sample("a", dist.Normal(0, 10)) b = numpyro.sample("b", dist.Normal(0, 10), sample_shape=(3, )) mu = a + b[0] * x + b[1] * x**2 + b[2] * x**3 numpyro.sample("y", dist.Normal(mu, 0.001), obs=y)
def model(): alpha = numpyro.sample("alpha", dist.Normal()) numpyro.deterministic("alpha", alpha * 2)
def model(data): x = numpyro.sample("x", dist.Normal(0, 1)) with optional(use_context_manager, handlers.scale(scale=10)): numpyro.sample("obs", dist.Normal(x, 1), obs=data)
def model(): x = numpyro.sample("x", dist.Normal()) with numpyro.handlers.mask(mask=False): numpyro.sample("y", dist.Normal(x), obs=1)
def _sample(): x = numpyro.sample("x", dist.Normal(0.0, 1.0)) y = numpyro.sample("y", dist.Normal(1.0, 2.0)) return jnp.stack([x, y])
def model(data): alpha = numpyro.sample('alpha', dist.Uniform(0, 1)) loc = numpyro.param('loc', 0., constraint=constraints.interval(0., alpha)) numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def model(): x = numpyro.sample("x", dist.Normal(0.0, 1.0)) y = numpyro.sample("y", dist.Normal(0.0, 1.0)) return x + y
def model(data): loc = numpyro.sample('loc', dist.Normal(0., 1.)) numpyro.sample('obs', dist.Normal(loc, 1), obs=data)
def model(subsample): with handlers.substitute(data={"data": subsample}): with numpyro.plate("data", len(data), subsample_size) as ind: x = data[ind] z = numpyro.sample("z", dist.Normal(0, 1)) numpyro.sample("x", dist.Normal(z, 1), obs=x)
def model(labels): coefs = numpyro.sample('coefs', dist.Normal(np.zeros(dim), np.ones(dim))) logits = np.sum(coefs * data, axis=-1) return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
def likelihood_func(self, yhat): """Return a normal likelihood with fitted sigma.""" _sigma = numpyro.sample("_sigma", dist.Exponential(self.sigma_prior)) return dist.Normal(yhat, _sigma)