def frozen_random_walk(name, num_steps=100, num_frozen=10): # last random value is repeated frozen-1 times num_random = min(max(0, num_steps - num_frozen), num_steps) num_frozen = num_steps - num_random rw = numpyro.sample(name, dist.GaussianRandomWalk(num_steps=num_random)) rw = np.concatenate((rw, np.repeat(rw[-1], num_frozen))) return rw
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 ExponentialRandomWalk(loc=1., scale=1e-2, drift=0., num_steps=100): ''' Return distrubtion of exponentiated Gaussian random walk Variables are x_0, ..., x_{T-1} Dynamics in log-space are random walk with drift: log(x_0) := log(loc) log(x_t) := log(x_{t-1}) + drift + eps_t, eps_t ~ N(0, scale) ==> Dynamics in non-log space are: x_0 := loc x_t := x_{t-1} * exp(drift + eps_t), eps_t ~ N(0, scale) ''' log_loc = np.log(loc) + drift * (np.arange(num_steps) + 0.) return dist.TransformedDistribution( dist.GaussianRandomWalk(scale=scale, num_steps=num_steps), [ dist.transforms.AffineTransform(loc=log_loc, scale=1.), dist.transforms.ExpTransform() ])
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=100., gamma_shape=100., det_prob_est=0.3, det_prob_conc=50., confirmed_dispersion=0.3, death_dispersion=0.3, rw_scale=2e-1, forecast_rw_scale=0., drift_scale=None, num_frozen=0, rw_use_last=1, confirmed=None, death=None): ''' Stochastic SEIR model. Draws random parameters and runs dynamics. ''' # Sample initial number of infected individuals # Sample dispersion parameters around specified values rw = numpyro.sample("rw", dist.GaussianRandomWalk(scale=100, num_steps=T)) D0 = numpyro.sample("D0", dist.Normal(.000002 * N, 1000)) # Sample parameters if death is None: death = None death0 = None else: death = clean_daily_obs(death) death0 = death[0] # First observation z_hat = np.exp(np.cumsum(rw) + np.log(D0)) z0 = observe_normal("z0", z_hat[0], .95, death_dispersion, obs=death0) y0 = observe_normal("y0", z_hat[0], .95, death_dispersion, obs=death0) z = observe_normal("z", z_hat, .95, death_dispersion, obs=death) y = observe_normal("y", z_hat, .95, death_dispersion, obs=death) z = np.append(z0, z) y = np.append(y0, y) if (T_future > 0): z_hat_future = np.exp( np.cumsum(np.repeat(rw[-1], T_future)) + np.log(D0)) z_future = observe_normal("z_future", z_hat_future, .95, death_dispersion, obs=death) y_future = observe_normal("y_future", z_hat_future, .95, death_dispersion, obs=death) z = np.append(z, z_future) y = np.append(y, y_future) return None, None, y, z, None, None