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 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_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)
pymc_logp = model_1.logp()
np.testing.assert_allclose(manual_logp, pymc_logp)
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_scalar_ode_2_param(self):
"""Test running model for a scalar ODE with 2 parameters"""
def system(y, t, p):
return p[0] * np.exp(-p[0] * t) - p[1] * 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=2)
with pm.Model() as model:
alpha = pm.HalfCauchy("alpha", 1)
beta = pm.HalfCauchy("beta", 1)
y0 = pm.LogNormal("y0", 0, 1)
sigma = pm.HalfCauchy("sigma", 1)
forward = ode_model(theta=[alpha, beta], 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["beta"].shape == (1, 100)
assert idata.posterior["y0"].shape == (1, 100)
assert idata.posterior["sigma"].shape == (1, 100)
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), sigma=sigma, observed=yobs)
with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
idata = pm.sample(50, tune=0, chains=1)
assert idata.posterior["alpha"].shape == (1, 50)
assert idata.posterior["y0"].shape == (1, 50)
assert idata.posterior["sigma"].shape == (1, 50)
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=(IS_FLOAT32 and IS_WINDOWS),
reason="Fails on float32 on Windows")
def test_too_many_params(self):
with pytest.raises(pm.ShapeError):
self.ode_model(theta=[1, 1], y0=[0])
@pytest.mark.xfail(condition=(IS_FLOAT32 and IS_WINDOWS),
reason="Fails on float32 on Windows")
def test_too_many_y0(self):
with pytest.raises(pm.ShapeError):
self.ode_model(theta=[1], y0=[0, 0])
@pytest.mark.xfail(condition=(IS_FLOAT32 and IS_WINDOWS),
reason="Fails on float32 on Windows")
def test_too_few_params(self):
with pytest.raises(pm.ShapeError):
self.ode_model(theta=[], y0=[1])
@pytest.mark.xfail(condition=(IS_FLOAT32 and IS_WINDOWS),
reason="Fails on float32 on Windows")
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_number_of_states(self):
with pytest.raises(ValueError,
match="Argument n_states must be at least 1."):
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,
match="Argument n_theta must be positive"):
DifferentialEquation(func=self.system,
t0=0,
times=self.times,
n_states=1,
n_theta=0)
def test_func_callable(self):
with pytest.raises(ValueError,
match="Argument func must be callable."):
DifferentialEquation(func=1,
t0=0,
times=self.times,
n_states=1,
n_theta=1)
def setup_method(self, method):
def system(y, t, p):
return np.exp(-t) - p[0] * y[0]
self.system = system
self.times = np.arange(0, 9)
self.ode_model = DifferentialEquation(
func=system, t0=0, times=self.times, n_states=1, n_theta=1
)
def test_unexpected_return_type_dict(self):
with pytest.raises(
TypeError, match="Unexpected type, <class 'dict'>, returned by ode_func."
):
def system_dict(y, t, p):
return {"rhs": np.exp(-t) - p[0] * y[0]}
DifferentialEquation(func=system_dict, t0=0, times=self.times, n_states=4, n_theta=1)
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_list_shape(self):
with pytest.raises(ValueError, match="returned a 2-dimensional tensor"):
def system_2d_list(y, t, p):
s0 = np.exp(-t) - p[0] * y[0]
s1 = np.exp(-t) - p[0] * y[1]
s2 = np.exp(-t) - p[0] * y[2]
s3 = np.exp(-t) - p[0] * y[3]
return [[s0, s1], [s2, s3]]
DifferentialEquation(func=system_2d_list, t0=0, times=self.times, n_states=4, n_theta=1)
def test_tensor_shape(self):
with pytest.raises(ValueError, match="returned a 2-dimensional tensor"):
def system_2d_tensor(y, t, p):
s0 = np.exp(-t) - p[0] * y[0]
s1 = np.exp(-t) - p[0] * y[1]
s2 = np.exp(-t) - p[0] * y[2]
s3 = np.exp(-t) - p[0] * y[3]
return at.stack((s0, s1, s2, s3)).reshape((2, 2))
DifferentialEquation(
func=system_2d_tensor, t0=0, times=self.times, n_states=4, n_theta=1
)
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_sens_ic_vector_2_param_tensor(self):
# Vector ODE 2 Param with return type at.TensorVariable
def ode_func_4_t(y, t, p):
# Make sure that ds and di are vectors by slicing
ds = -p[0:1] * y[0:1] * y[1:]
di = p[0:1] * y[0:1] * y[1:] - p[1:] * y[1:]
return at.concatenate([ds, di], axis=0)
# Instantiate ODE model
model4_t = DifferentialEquation(
func=ode_func_4_t, t0=0, times=self.t, n_states=2, n_theta=2
)
model4_sens_ic_t = np.array([1, 0, 0, 0, 0, 1, 0, 0])
np.testing.assert_array_equal(model4_sens_ic_t, model4_t._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, initval=1)
gamma = pm.HalfCauchy("gamma", 1, initval=1)
sigma = pm.HalfCauchy("sigma", 1, shape=2, initval=[1, 1])
forward = ode_model(theta=[beta, gamma], y0=[0.99, 0.01])
y = pm.LogNormal("y",
mu=pm.math.log(forward),
sigma=sigma,
observed=yobs)
with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
idata = pm.sample(50, tune=0, chains=1)
assert idata.posterior["beta"].shape == (1, 50)
assert idata.posterior["gamma"].shape == (1, 50)
assert idata.posterior["sigma"].shape == (1, 50, 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, initval=1)
sigma = pm.HalfCauchy("sigma", 1, shape=2, initval=[0.5, 0.5])
forward = ode_model(theta=[R], 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["R"].shape == (1, 100)
assert idata.posterior["sigma"].shape == (1, 100, 2)