def test_op_equality(self): """Tests that the equality of mathematically identical Ops evaluates True""" # Create ODE to test with def ode_func(y, t, p): return np.exp(-t) - p[0] * y[0] t = np.linspace(0, 2, 12) # Instantiate two Ops op_1 = DifferentialEquation(func=ode_func, t0=0, times=t, n_states=1, n_theta=1) op_2 = DifferentialEquation(func=ode_func, t0=0, times=t, n_states=1, n_theta=1) op_other = DifferentialEquation(func=ode_func, t0=0, times=np.linspace(0, 2, 16), n_states=1, n_theta=1) assert op_1 == op_2 assert op_1 != op_other return
def test_number_of_params(self): with pytest.raises(ValueError): DifferentialEquation(func=self.system, t0=0, times=self.times, n_states=1, n_theta=0)
def test_func_callable(self): with pytest.raises(ValueError): DifferentialEquation(func=1, t0=0, times=self.times, n_states=1, n_theta=1)
class TestErrors: """Test running model for a scalar ODE with 1 parameter""" def system(y, t, p): return np.exp(-t) - p[0] * y[0] times = np.arange(0, 9) ode_model = DifferentialEquation(func=system, t0=0, times=times, n_states=1, n_theta=1) @pytest.mark.xfail(condition=(aesara.config.floatX == "float32"), reason="Fails on float32") def test_too_many_params(self): with pytest.raises(pm.ShapeError): self.ode_model(theta=[1, 1], y0=[0]) @pytest.mark.xfail(condition=(aesara.config.floatX == "float32"), reason="Fails on float32") def test_too_many_y0(self): with pytest.raises(pm.ShapeError): self.ode_model(theta=[1], y0=[0, 0]) @pytest.mark.xfail(condition=(aesara.config.floatX == "float32"), reason="Fails on float32") def test_too_few_params(self): with pytest.raises(pm.ShapeError): self.ode_model(theta=[], y0=[1]) @pytest.mark.xfail(condition=(aesara.config.floatX == "float32"), reason="Fails on float32") def test_too_few_y0(self): with pytest.raises(pm.ShapeError): self.ode_model(theta=[1], y0=[]) def test_func_callable(self): with pytest.raises(ValueError): DifferentialEquation(func=1, t0=0, times=self.times, n_states=1, n_theta=1) def test_number_of_states(self): with pytest.raises(ValueError): DifferentialEquation(func=self.system, t0=0, times=self.times, n_states=0, n_theta=1) def test_number_of_params(self): with pytest.raises(ValueError): DifferentialEquation(func=self.system, t0=0, times=self.times, n_states=1, n_theta=0)
def test_scalar_ode_1_param(self): """Test running model for a scalar ODE with 1 parameter""" def system(y, t, p): return np.exp(-t) - p[0] * y[0] times = np.array([ 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0, 7.5 ]) yobs = np.array([ 0.31, 0.57, 0.51, 0.55, 0.47, 0.42, 0.38, 0.3, 0.26, 0.22, 0.22, 0.14, 0.14, 0.09, 0.1 ])[:, np.newaxis] ode_model = DifferentialEquation(func=system, t0=0, times=times, n_states=1, n_theta=1) with pm.Model() as model: alpha = pm.HalfCauchy("alpha", 1) y0 = pm.Lognormal("y0", 0, 1) sigma = pm.HalfCauchy("sigma", 1) forward = ode_model(theta=[alpha], y0=[y0]) y = pm.Lognormal("y", mu=pm.math.log(forward), sd=sigma, observed=yobs) idata = pm.sample(100, tune=0, chains=1) assert idata.posterior["alpha"].shape == (1, 100) assert idata.posterior["y0"].shape == (1, 100) assert idata.posterior["sigma"].shape == (1, 100)
def test_logp_scalar_ode(): """Test the computation of the log probability for these models""" # Differential equation def system_1(y, t, p): return np.exp(-t) - p[0] * y[0] # Parameters and inital condition alpha = 0.4 y0 = 0.0 times = np.arange(0.5, 8, 0.5) yobs = np.array( [0.30, 0.56, 0.51, 0.55, 0.47, 0.42, 0.38, 0.30, 0.26, 0.21, 0.22, 0.13, 0.13, 0.09, 0.09] )[:, np.newaxis] ode_model = DifferentialEquation(func=system_1, t0=0, times=times, n_theta=1, n_states=1) integrated_solution, *_ = ode_model._simulate([y0], [alpha]) assert integrated_solution.shape == yobs.shape # compare automatic and manual logp values manual_logp = norm.logpdf(x=np.ravel(yobs), loc=np.ravel(integrated_solution), scale=1).sum() with pm.Model() as model_1: forward = ode_model(theta=[alpha], y0=[y0]) y = pm.Normal("y", mu=forward, sd=1, observed=yobs) pymc3_logp = model_1.logp() np.testing.assert_allclose(manual_logp, pymc3_logp)
def test_simulate(): """Tests the integration in DifferentialEquation""" # Create an ODe to integrate def ode_func(y, t, p): return np.exp(-t) - p[0] * y[0] # Evaluate exact solution y0 = 0 t = np.arange(0, 12, 0.25).reshape(-1, 1) a = 0.472 y = 1.0 / (a - 1) * (np.exp(-t) - np.exp(-a * t)) # Instantiate ODE model ode_model = DifferentialEquation(func=ode_func, t0=0, times=t, n_states=1, n_theta=1) simulated_y, sens = ode_model._simulate([y0], [a]) assert simulated_y.shape == (len(t), 1) assert sens.shape == (len(t), 1, 1 + 1) np.testing.assert_allclose(y, simulated_y, rtol=1e-5)
def inferenceCases(self,S0,I0,R0,observedCumConfirmedCases): import pymc3 as pm import theano import theano.tensor as tt # <- from theano import printing from pymc3.ode import DifferentialEquation from scipy.integrate import odeint import arviz as az N = S0+I0+R0 def odem(y,t,p): # y= [S,I,R,c] # p= [beta,gamma,delta] dS = -1.*y[0]*y[1]*p[0] dI = y[0]*y[1]*p[0] - ( p[1]*y[1] ) # infected dR = p[1]*y[1] dc = y[0]*y[1]*p[0] # cumulative cases return [dS,dI,dR,dc] # last state is cumulative cases dta = observedCumConfirmedCases.reshape(-1,) #dta = O[:,-1].reshape(-1,) timesteps = len(dta) times = np.arange(0,timesteps) sir_model = DifferentialEquation( func = odem, times = times, n_states = 4, n_theta = 2, t0 = 0, ) with pm.Model() as model: beta = pm.Normal('beta' ,2.0,1) gamma = pm.Normal('gamma' ,2.0,1) tt.printing.Print('beta')(beta) tt.printing.Print('gamma')(gamma) sigmas = pm.Gamma('sigmas', 0.5,0.5, shape=1) sir_curves = sir_model(y0=[S0,I0,R0,0], theta=[beta,gamma]) # maybe?? obs = pm.Normal("obs", observed = dta , mu = sir_curves[:,-1], sd=sigmas) # maybe? with model: MAP = pm.find_MAP() self.MAP = MAP self.beta = MAP['beta'] self.gamma = MAP['gamma']
def inferenceCasesAndDeaths(self, S0, I0, R0, D0, observedNewConfirmedCases, observedNewDeaths): import pymc3 as pm import theano import theano.tensor as tt from theano import printing from pymc3.ode import DifferentialEquation from scipy.integrate import odeint import arviz as az def odem(y, t, p): # y= [S,I,R,D,c,d,r] # p= [beta,gamma,delta] dS = -1. * y[0] * y[1] * p[0] dI = y[0] * y[1] * p[0] - (p[1] * y[1] + p[2] * y[1]) dR = p[1] * y[1] dD = p[2] * y[1] dc = y[0] * y[1] * p[0] return [dS, dI, dR, dD, dc] dta = np.hstack((observedNewDeaths.reshape(-1, 1), observedNewConfirmedCases.reshape(-1, 1))) timesteps = len(dta) times = np.arange(0, timesteps) sir_model = DifferentialEquation( func=odem, times=times, n_states=5, n_theta=3, t0=0, ) with pm.Model() as model: beta = pm.Uniform('beta', 0.0, 10.0) gamma = pm.Uniform('gamma', 0.0, 10.0) delta = pm.Uniform('delta', 0.0, 10.0) sigmas = pm.Gamma('sigmas', 0.5, 0.5, shape=2) sir_curves = sir_model(y0=[S0, I0, R0, D0, 0], theta=[beta, gamma, delta]) obs = pm.Normal("obs", observed=dta, mu=sir_curves[:, [3, 4]], sd=sigmas) with model: MAP = pm.find_MAP() self.MAP = MAP
def test_sens_ic_scalar_2_param(self): # Scalar ODE 2 Param def ode_func_2(y, t, p): return p[0] * np.exp(-p[0] * t) - p[1] * y[0] # Instantiate ODE model model2 = DifferentialEquation(func=ode_func_2, t0=0, times=self.t, n_states=1, n_theta=2) model2_sens_ic = np.array([1, 0, 0]) np.testing.assert_array_equal(model2_sens_ic, model2._sens_ic)
def test_sens_ic_vector_2_param(self): # Vector ODE 2 Param def ode_func_4(y, t, p): ds = -p[0] * y[0] * y[1] di = p[0] * y[0] * y[1] - p[1] * y[1] return [ds, di] # Instantiate ODE model model4 = DifferentialEquation(func=ode_func_4, t0=0, times=self.t, n_states=2, n_theta=2) model4_sens_ic = np.array([1, 0, 0, 0, 0, 1, 0, 0]) np.testing.assert_array_equal(model4_sens_ic, model4._sens_ic)
def test_sens_ic_scalar_1_param(self): """Tests the creation of the initial condition for the sensitivities""" # Scalar ODE 1 Param # Create an ODe to integrate def ode_func_1(y, t, p): return np.exp(-t) - p[0] * y[0] # Instantiate ODE model # Instantiate ODE model model1 = DifferentialEquation(func=ode_func_1, t0=0, times=self.t, n_states=1, n_theta=1) # Sensitivity initial condition for this model should be 1 by 2 model1_sens_ic = np.array([1, 0]) np.testing.assert_array_equal(model1_sens_ic, model1._sens_ic)
def test_vector_ode_2_param(self): """Test running model for a vector ODE with 2 parameters""" def system(y, t, p): ds = -p[0] * y[0] * y[1] di = p[0] * y[0] * y[1] - p[1] * y[1] return [ds, di] times = np.array( [0.0, 0.8, 1.6, 2.4, 3.2, 4.0, 4.8, 5.6, 6.4, 7.2, 8.0]) yobs = np.array([ [1.02, 0.02], [0.86, 0.12], [0.43, 0.37], [0.14, 0.42], [0.05, 0.43], [0.03, 0.14], [0.02, 0.08], [0.02, 0.04], [0.02, 0.01], [0.02, 0.01], [0.02, 0.01], ]) ode_model = DifferentialEquation(func=system, t0=0, times=times, n_states=2, n_theta=2) with pm.Model() as model: beta = pm.HalfCauchy("beta", 1) gamma = pm.HalfCauchy("gamma", 1) sigma = pm.HalfCauchy("sigma", 1, shape=2) forward = ode_model(theta=[beta, gamma], y0=[0.99, 0.01]) y = pm.Lognormal("y", mu=pm.math.log(forward), sd=sigma, observed=yobs) idata = pm.sample(100, tune=0, chains=1) assert idata.posterior["beta"].shape == (1, 100) assert idata.posterior["gamma"].shape == (1, 100) assert idata.posterior["sigma"].shape == (1, 100, 2)
def test_sens_ic_vector_3_params(self): # Big System with Many Parameters def ode_func_5(y, t, p): dx = p[0] * (y[1] - y[0]) ds = y[0] * (p[1] - y[2]) - y[1] dz = y[0] * y[1] - p[2] * y[2] return [dx, ds, dz] # Instantiate ODE model model5 = DifferentialEquation(func=ode_func_5, t0=0, times=self.t, n_states=3, n_theta=3) # First three columns are derivatives with respect to ode parameters # Last three coluimns are derivatives with repsect to initial condition # So identity matrix should appear in last 3 columns model5_sens_ic = np.array([[1, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0]]) np.testing.assert_array_equal(np.ravel(model5_sens_ic), model5._sens_ic)
def test_vector_ode_1_param(self): """Test running model for a vector ODE with 1 parameter""" def system(y, t, p): ds = -p[0] * y[0] * y[1] di = p[0] * y[0] * y[1] - y[1] return [ds, di] times = np.array([0.0, 0.8, 1.6, 2.4, 3.2, 4.0, 4.8, 5.6, 6.4, 7.2, 8.0]) yobs = np.array( [ [1.02, 0.02], [0.86, 0.12], [0.43, 0.37], [0.14, 0.42], [0.05, 0.43], [0.03, 0.14], [0.02, 0.08], [0.02, 0.04], [0.02, 0.01], [0.02, 0.01], [0.02, 0.01], ] ) ode_model = DifferentialEquation(func=system, t0=0, times=times, n_states=2, n_theta=1) with pm.Model() as model: R = pm.Lognormal("R", 1, 5) sigma = pm.HalfCauchy("sigma", 1, shape=2) forward = ode_model(theta=[R], y0=[0.99, 0.01]) y = pm.Lognormal("y", mu=pm.math.log(forward), sd=sigma, observed=yobs) trace = pm.sample(100, tune=0, chains=1) assert trace["R"].size > 0 assert trace["sigma"].size > 0
def get_SIR(x, y, y0, country, forecast_len=0, load_post=False): ''' If 'forecast_len' is nonzero, attempts to load a trace corresponding to the country of interest from the directory 'traces' and retrieves predicted numbers of infected and susceptible patients 'forecast_len' days into the future after the 1st case is detected in the country. ''' # If in 'prediction mode', modify x, y to reflect forecast length if forecast_len != 0: ext = np.arange(1, forecast_len + 1).astype(float) ext += x[-1] x = np.append(x, ext) y = np.empty((x.shape[0], y.shape[1])) # SIR Model # p[0]: beta, p[1]: lambda def SIR(y, t, p): ds = -p[0] * y[0] * y[1] # Susceptible differential di = p[0] * y[0] * y[1] - p[1] * y[1] # Infected differential return [ds, di] # Initialize ODE sir_ode = DifferentialEquation(func=SIR, times=x, n_states=2, n_theta=2, t0=0) load_dir = osp.join('traces', country.lower()) with pm.Model() as model: sigma = pm.HalfNormal('sigma', 3, shape=2) # R0 is bounded below by 1 because we see an epidemic has occured R0 = pm.Normal('R0', 2, 3) lmbda = pm.Normal('lambda', 0.1, 0.1) beta = pm.Deterministic('beta', lmbda * R0) print('Setting up model for ' + country) sir_curves = sir_ode(y0=y0, theta=[beta, lmbda]) y_obs = pm.Normal('y_obs', mu=sir_curves, sigma=sigma, observed=y) if forecast_len == 0: trace = pm.sample(2000, tune=1000, cores=2, chains=2, progressbar=True) # Save trace pm.save_trace(trace, load_dir, overwrite=True) # Get the posterior post = pm.sample_posterior_predictive(trace, progressbar=True) out_post = post else: # Load trace print('Loading trace') trace = pm.load_trace(load_dir) print('Computing posterior') #Get posterior if not load_post: post = pm.sample_posterior_predictive(trace[500:], progressbar=True) out_post = post with open(country + '_post.pkl', 'wb') as buff: pickle.dump({'post': post}, buff) else: with open(country + '_post.pkl', 'rb') as buff: data = pickle.load(buff) out_post = data['post'] print('Done') return trace, out_post, x
# Define ODE def freefall(y, t, p): return 2.0 * p[1] - p[0] * y[0] # Get data to fit times = np.arange(0, 10, 0.5) gamma, g, y0, sigma = 0.4, 9.8, -2, 2 y = odeint(freefall, t=times, y0=y0, args=tuple([[gamma, g]])) yobs = np.random.normal(y, sigma) # Define inference model ode_model = DifferentialEquation(func=freefall, times=times, n_states=1, n_theta=2, t0=0) with pm.Model() as model: # priors sigma = pm.HalfCauchy('sigma', 1) g = pm.HalfCauchy('g', 1) gamma = pm.Lognormal('gamma', 0, 1) # observed ode_solution = ode_model(y0=[0], theta=[gamma, g]) '''The ode_solution has a shape of (n_times, n_states)''' Y = pm.Normal('Y', mu=ode_solution, sd=sigma, observed=yobs) # inference trace = pm.sample(1000, tune=1000, cores=1) # prior = pm.sample_prior_predictive() # posterior_predictive = pm.sample_posterior_predictive(trace)
# s: y[0], e: y[1], i_1: y[2], i_2: y[3], a: y[4] def SEIR(y, t, p): ds = -p[0] * y[0] * (y[3] + alpha_1 * y[2] + alpha_2 * y[4]) / N de = p[0] * y[0] * (y[3] + alpha_1 * y[2] + alpha_2 * y[4]) / N - (1-r) * y[1] / D_e - r * y[1] / D_e di_1 = (1-r) * y[1] / D_e - y[2] / D_i di_2 = y[2] / D_i - y[3] / D_r da = r * y[1] / D_e - y[4] / D_r return [ds, de, di_1, di_2, da] seir_model = DifferentialEquation( func=SEIR, times=one_estimation_period, n_states=5, # S, E, I, A, H. (no need to have R) n_theta=1, t0=0, ) with pm.Model() as modelSEIR: beta1 = pm.Uniform('beta1', 0, 5) beta2 = pm.Uniform('beta2', 0, 5) beta3 = pm.Uniform('beta3', 0, 5) beta4 = pm.Uniform('beta4', 0, 5) beta5 = pm.Uniform('beta5', 0, 5) beta6 = pm.Uniform('beta6', 0, 5) seir_curves1 = seir_model(y0 = [sus0, exp0, inf0, asy0, hos0], theta=[beta1,])
otypes=[t.dmatrix]) def seirdaq_ode_wrapper(time_exp, initial_conditions, beta, gamma_I): time_span = (time_exp.min(), time_exp.max()) args = [beta, gamma_I] y_model = seirdaq_ode_solver(initial_conditions, time_span, time_exp, *args) simulated_time = y_model.t simulated_ode_solution = y_model.y return simulated_ode_solution seir_jia_model_diffeq = DifferentialEquation(func=seir_jia_model_pymc3, times=time_range, n_states=1, n_theta=13, t0=0) # %% # beta_deterministic, gamma_I_deterministic, gamma_A_deterministic, gamma_D_deterministic, d_I_deterministic, d_D_deterministic = result_seirdaq.x beta_deterministic, gamma_I_deterministic = result_seirdaq.x print( f"parameters: beta = {beta_deterministic}, gamma_I = {gamma_I_deterministic}" ) # %% number_of_cores = 24 population_uncertain_variance = 0.05 * np.max(dead_individuals) variance = population_uncertain_variance * population_uncertain_variance
# Observed system response Xobs = np.random.normal(X, 0.02) # Plot the actual and observed system response plt.plot(timerange, Xobs, marker='o', linestyle='none') plt.plot(timerange, X, color='C0', alpha=0.5) plt.legend() plt.show() # Define the diferential equation for the pymc3odint module lmm_model = DifferentialEquation( func=X_dot, # The DE function, first order times=np.linspace(0, 0.1, 10), # time range (why not use already defined time? n_states=4, # Degrees of freedom n_theta=1, # Model parameters t0=0, # Start at t=0 ) with pm.Model(): sigma = pm.HalfNormal('sigma', 0.05, shape=4) #m1 is bounded because mass cannot take negative value m1 = pm.Bound(pm.Normal, lower=0)('m1', 0.5, 3) # Set intitial conditions to known values, only parameter is mass lmm_sol = lmm_model(y0=X0, theta=[m1]) Y = pm.Lognormal('Y', mu=pm.math.log(lmm_sol), sd=sigma, observed=Xobs)
X=np.array([np.array(data['s'])[7:14], np.array(data['i'])[7:14]] ) #X=X/pop #code to normalize data to 1 time_range = np.arange(0,len(X[0])) #%% # Create differential equation models model1 = DifferentialEquation( func = SIR_diffeq, # what is the differential equation? times = time_range, # what is the time grid for numerical integration? n_states = 3, # what is the dimensionality of the system? n_theta = 2, # how many parameters are there? # t0 = 0 # are we starting at the beginning? ) model2 = DifferentialEquation( func = SIR_wDR_diffeq, # what is the differential equation? times = time_range, # what is the time grid for numerical integration? n_states = 2, # what is the dimensionality of the system? n_theta = 3, # how many parameters are there? # t0 = 0 # are we starting at the beginning? ) model3=DifferentialEquation( func = SI_diffeq, # what is the differential equation?
plt.legend() # plt.show() # exit() survivor = ([0] + observations[ObsEnum.RSURVIVOR.value]) - ( observations[ObsEnum.RSURVIVOR.value] + [0]) print(observations[ObsEnum.RSURVIVOR.value]) print("Making the model...") # This is a DiffEq definition for pymc3. pm_ode_model = DifferentialEquation( func=ode_model, times=[x for x in range(STEPS)], n_states=10, # dimension of the return value of 'func' n_theta=10, # dimension of p (additional parameters) t0=0) # Now we use some declarations to se up the random # variables and how they are tied together basic_model = pm.Model() with basic_model: # See : https://docs.pymc.io/notebooks/getting_started.html # Prior belief (because it's not based on data). # Bound = put bounds around the parameter we're modelling # Normal = the parameter is ditributed around some normal distribution gamma1 = pm.Bound(pm.Normal, lower=1 / 10,
# Observed system response yobs = np.random.lognormal(mean=np.log(y[1::]), sigma=[0.2, 0.3]) # Plot the actual and observed system response plt.plot(times[1::], yobs, marker='o', linestyle='none') plt.plot(times, y[:, 0], color='C0', alpha=0.5, label=f'$S(t)$') plt.plot(times, y[:, 1], color='C1', alpha=0.5, label=f'$I(t)$') plt.legend() plt.show() # Define the diferential equation for the pymc3odint module sir_model = DifferentialEquation( func=SIR, # The DE function, first order times=np.arange(0.25, 5, 0.25), # time range (why not use already defined time? n_states=2, # Degrees of freedom n_theta=2, # Model parameters t0=0, # Start at t=0 ) with pm.Model(): sigma = pm.HalfCauchy('sigma', 1, shape=2) # R0 is bounded below by 1 because we see an epidemic has occured R0 = pm.Bound(pm.Normal, lower=1)('R0', 2, 3) lam = pm.Lognormal('lambda', pm.math.log(2), 2) beta = pm.Deterministic('beta', lam * R0) sir_curves = sir_model(y0=[0.99, 0.01], theta=[beta, lam]) Y = pm.Lognormal('Y',
ds = -p[0] * y[0] * y[1] di = p[0] * y[0] * y[1] - p[1] * y[1] return [ds, di] # Get data to fit times = np.arange(0, 5, 0.25) beta, gamma = 4, 1.0 y = odeint(SIR, t=times, y0=[0.99, 0.01], args=((beta, gamma), ), rtol=1e-8) yobs = np.random.lognormal(mean=np.log(y[1::]), sigma=[0.2, 0.3]) # Define inference model sir_model = DifferentialEquation( func=SIR, times=np.arange(0.25, 5, 0.25), n_states=2, n_theta=2, t0=0, ) with pm.Model() as model: # priors sigma = pm.HalfCauchy('sigma', 1, shape=2) R0 = pm.Bound(pm.Normal, lower=1)( 'R0', 2, 3) # R0 is bounded below by 1 because we see an epidemic has occured lam = pm.Lognormal('lambda', pm.math.log(2), 2) beta = pm.Deterministic('beta', lam * R0) # observed sir_curves = sir_model(y0=[0.99, 0.01], theta=[beta, lam]) Y = pm.Lognormal('Y', mu=pm.math.log(sir_curves), sd=sigma, observed=yobs) # inference