def __call__(self, T=50, N=1e5, T_future=0, E_duration_est=4.0, I_duration_est=2.0, R0_est=3.0, beta_shape=1, sigma_shape=5, gamma_shape=8, det_prob_est=0.3, det_prob_conc=50, det_noise_scale=0.15, rw_scale=2e-1, forecast_rw_scale=0, drift_scale=None, num_frozen=0, 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(0, 0.02 * N)) E0 = numpyro.sample("E0", dist.Uniform(0, 0.02 * N)) H0 = numpyro.sample("H0", dist.Uniform(0, 0.02 * N)) D0 = numpyro.sample("D0", dist.Uniform(0, 0.02 * 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)) beta0 = numpyro.sample( "beta0", dist.Gamma(beta_shape, beta_shape * I_duration_est / R0_est)) det_prob0 = numpyro.sample( "det_prob0", 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(.01 * 100, (1 - .01) * 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_prob0, 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_prob0, det_noise_scale, death_prob, death_rate, det_prob_d) beta, det_prob, x, y, z = self.dynamics(T, params, x0, num_frozen=num_frozen, 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[-1], det_noise_scale, death_prob, death_rate, det_prob_d) beta_f, det_rate_rw_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 __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 = .2, gamma_shape = .2, 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, hosp=None): ''' Stochastic SEIHR 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)) R0 = numpyro.sample("R0", dist.Uniform(0, 0.01*N)) H0 = numpyro.sample("H0", dist.Uniform(0, 0.01*N)) # Sample parameters sigma = numpyro.sample("sigma", dist.LogNormal(np.log(1/E_duration_est),sigma_shape)) gamma = numpyro.sample("gamma", dist.LogNormal(np.log(1/I_duration_est),gamma_shape)) # 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)) hosp_prob = numpyro.sample("hosp_prob", dist.Beta(.1 * 100, (1-.1) * 100)) if drift_scale is not None: drift = numpyro.sample("drift", dist.Normal(loc=0, scale=drift_scale)) else: drift = 0 x0 = SEIHRModel.seed(N=N, I=I0, E=E0, R=R0, H=H0) numpyro.deterministic("x0", x0) # Split observations into first and rest #confirmed0, confirmed = (None, None) if confirmed is None else (confirmed[0], confirmed[1:]) hosp0, hosp = (None, None) if hosp is None else (hosp[0], hosp[1:]) # First observation with numpyro.handlers.scale(scale_factor=2.0): y0 = observe_nonrandom("y0", x0[3], det_noise_scale, obs=hosp0) params = (beta0, sigma, gamma, rw_scale, drift, det_prob, det_noise_scale, hosp_prob) beta, x, y = self.dynamics(T, params, x0, hosp=hosp) 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, hosp_prob) beta_f, x_f, y_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, det_prob, hosp_prob
def model(data): alpha = np.array([1.1, 1.1]) beta = np.array([1.1, 1.1]) p_latent = sample('p_latent', dist.Beta(alpha, beta)) sample('obs', dist.Bernoulli(p_latent), obs=data) return p_latent
def model(data): f = numpyro.sample("beta", dist.Beta(1., 1.)) numpyro.sample("obs", dist.Bernoulli(f), obs=data)
def guide(data): alpha_q = numpyro.param("alpha_q", 1.0, constraint=constraints.positive) beta_q = numpyro.param("beta_q", 1.0, constraint=constraints.positive) numpyro.sample("beta", dist.Beta(alpha_q, beta_q))
def model(data=None): with numpyro.plate("dim", 2): beta = numpyro.sample("beta", dist.Beta(1., 1.)) with numpyro.plate("plate", N, dim=-2): numpyro.deterministic("beta_sq", beta**2) numpyro.sample("obs", dist.Bernoulli(beta), obs=data)
def model(data): f = numpyro.sample("beta", dist.Beta(1., 1.)) with numpyro.plate("plate", 10): numpyro.deterministic("beta_sq", f**2) numpyro.sample("obs", dist.Bernoulli(f), obs=data)
def model(): c = numpyro.sample("c", dist.Gamma(1, 1)) with handlers.collapse(): probs = numpyro.sample("probs", dist.Beta(c, 2)) with numpyro.plate("plate", len(data)): numpyro.sample("obs", dist.Binomial(10, probs), obs=data)
def model(): c = numpyro.sample("c", dist.Gamma(1, 1)) with handlers.collapse(): probs = numpyro.sample("probs", dist.Beta(c, 2)) numpyro.sample("obs", dist.Bernoulli(probs), obs=data)
def model1(): c1 = numpyro.param("c1", 0.5, constraint=dist.constraints.positive) c0 = numpyro.param("c0", 1.5, constraint=dist.constraints.positive) with handlers.collapse(): probs = numpyro.sample("probs", dist.Beta(c1, c0)) numpyro.sample("obs", dist.Binomial(total_count, probs), obs=data)
def model(data): f = numpyro.sample('beta', dist.Beta(np.ones(2), np.ones(2))) numpyro.sample('obs', dist.Bernoulli(f), obs=data)
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)