diff --git a/pycode/examples/simulation/sde_sir_simple.py b/pycode/examples/simulation/sde_sir_simple.py new file mode 100644 index 0000000000..1272d94f2a --- /dev/null +++ b/pycode/examples/simulation/sde_sir_simple.py @@ -0,0 +1,74 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Maximilian Betz +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# +import argparse + +import numpy as np + +from memilio.simulation import AgeGroup, Damping +from memilio.simulation.ssir import InfectionState as State +from memilio.simulation.ssir import (Model, simulate_stochastic) + + +def run_sde_sir_simulation(): + """Runs the c++ SDE SIR model""" + + # Define population of age groups + population = 10000 + + days = 5. # number of days to simulate + dt = 0.1 + + # Initialize Parameters + model = Model() + + # Compartment transition duration + model.parameters.TimeInfected.value = 10. + + # # Compartment transition propabilities + model.parameters.TransmissionProbabilityOnContact.value = 1. + + # Initial number of people in each compartment + model.populations[State.Infected] = 100 + model.populations[State.Recovered] = 1000 + model.populations.set_difference_from_total( + (State.Susceptible), population) + + model.parameters.ContactPatterns.baseline = np.ones( + (1, 1)) * 2.7 + model.parameters.ContactPatterns.minimum = np.zeros( + (1, 1)) + model.parameters.ContactPatterns.add_damping( + Damping(coeffs=np.r_[0.6], t=12.5, level=0, type=0)) + + # Check logical constraints to parameters + model.check_constraints() + + # Run Simulation + result = simulate_stochastic(0., days, dt, model) + + result.print_table(False, ["S", "I", "R"], 16, 5) + + +if __name__ == "__main__": + arg_parser = argparse.ArgumentParser( + 'sde_sir_simple', + description='Simple example demonstrating the setup and simulation of the SDE SIR model.') + args = arg_parser.parse_args() + run_sde_sir_simulation() diff --git a/pycode/examples/simulation/sde_sirs_simple.py b/pycode/examples/simulation/sde_sirs_simple.py new file mode 100644 index 0000000000..7397b80a83 --- /dev/null +++ b/pycode/examples/simulation/sde_sirs_simple.py @@ -0,0 +1,80 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Maximilian Betz +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# +import argparse + +import numpy as np + +from memilio.simulation import AgeGroup, Damping +from memilio.simulation.ssirs import InfectionState as State +from memilio.simulation.ssirs import ( + Model, simulate_stochastic, interpolate_simulation_result) + + +def run_sde_sirs_simulation(): + """Runs the c++ SDE SIRs model""" + + # Define population of age groups + population = 10000 + + days = 5. # number of days to simulate + dt = 0.001 + + # Initialize Parameters + model = Model() + + # Compartment transition duration + model.parameters.TimeInfected.value = 10. + model.parameters.TimeImmune.value = 100. + + # # Compartment transition propabilities + model.parameters.TransmissionProbabilityOnContact.value = 1. + + # Initial number of people in each compartment + model.populations[State.Infected] = 100 + model.populations[State.Recovered] = 1000 + model.populations.set_difference_from_total( + (State.Susceptible), population) + + model.parameters.ContactPatterns.baseline = np.ones( + (1, 1)) * 20.7 + model.parameters.ContactPatterns.minimum = np.zeros( + (1, 1)) + model.parameters.ContactPatterns.add_damping( + Damping(coeffs=np.r_[0.6], t=12.5, level=0, type=0)) + + # Check logical constraints to parameters + model.check_constraints() + + # Run Simulation + result = simulate_stochastic(0., days, dt, model) + + # Interpolate results + result = interpolate_simulation_result(result) + + # Print results + result.print_table(False, ["Susceptible", "Infected", "Recovered"], 16, 5) + + +if __name__ == "__main__": + arg_parser = argparse.ArgumentParser( + 'sde_sirs_simple', + description='Simple example demonstrating the setup and simulation of the SDE SIRS model.') + args = arg_parser.parse_args() + run_sde_sirs_simulation() diff --git a/pycode/memilio-simulation/CMakeLists.txt b/pycode/memilio-simulation/CMakeLists.txt index 975579d3fd..3caee964eb 100644 --- a/pycode/memilio-simulation/CMakeLists.txt +++ b/pycode/memilio-simulation/CMakeLists.txt @@ -88,3 +88,13 @@ add_pymio_module(_simulation_osecirvvs LINKED_LIBRARIES memilio ode_secirvvs SOURCES memilio/simulation/bindings/models/osecirvvs.cpp ) + +add_pymio_module(_simulation_ssir + LINKED_LIBRARIES memilio sde_sir + SOURCES memilio/simulation/bindings/models/ssir.cpp +) + +add_pymio_module(_simulation_ssirs + LINKED_LIBRARIES memilio sde_sirs + SOURCES memilio/simulation/bindings/models/ssirs.cpp +) \ No newline at end of file diff --git a/pycode/memilio-simulation/memilio/simulation/__init__.py b/pycode/memilio-simulation/memilio/simulation/__init__.py index 4ea23b166d..bd80fcca12 100644 --- a/pycode/memilio-simulation/memilio/simulation/__init__.py +++ b/pycode/memilio-simulation/memilio/simulation/__init__.py @@ -45,6 +45,13 @@ def __getattr__(attr): elif attr == "osecirvvs": import memilio.simulation.osecirvvs as osecirvvs return osecirvvs + elif attr == "ssir": + import memilio.simulation.ssir as ssir + return ssir + + elif attr == "ssirs": + import memilio.simulation.ssirs as ssirs + return ssirs raise AttributeError("module {!r} has no attribute " "{!r}".format(__name__, attr)) diff --git a/pycode/memilio-simulation/memilio/simulation/bindings/models/ssir.cpp b/pycode/memilio-simulation/memilio/simulation/bindings/models/ssir.cpp new file mode 100644 index 0000000000..4dbaa8ffed --- /dev/null +++ b/pycode/memilio-simulation/memilio/simulation/bindings/models/ssir.cpp @@ -0,0 +1,97 @@ +/* +* Copyright (C) 2020-2025 MEmilio +* +* Authors: Maximilian Betz +* +* Contact: Martin J. Kuehn +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +//Includes from pymio +#include "pybind_util.h" +#include "utils/index.h" +#include "compartments/compartmental_model.h" +#include "utils/custom_index_array.h" +#include "utils/parameter_set.h" +#include "epidemiology/populations.h" + +//Includes from MEmilio +#include "sde_sir/model.h" +#include "sde_sir/infection_state.h" +#include "memilio/compartments/stochastic_simulation.h" +#include "memilio/compartments/stochastic_model.h" +#include "memilio/data/analyze_result.h" + +#include "pybind11/pybind11.h" +#include +#include + +namespace py = pybind11; + +namespace pymio +{ +//specialization of pretty_name +template <> +inline std::string pretty_name() +{ + return "InfectionState"; +} + +} // namespace pymio + +PYBIND11_MODULE(_simulation_ssir, m) +{ + m.def("interpolate_simulation_result", + static_cast (*)(const mio::TimeSeries&, const double)>( + &mio::interpolate_simulation_result), + py::arg("ts"), py::arg("abs_tol") = 1e-14); + + m.def("interpolate_simulation_result", + static_cast (*)(const mio::TimeSeries&, const std::vector&)>( + &mio::interpolate_simulation_result), + py::arg("ts"), py::arg("interpolation_times")); + + m.def("interpolate_ensemble_results", &mio::interpolate_ensemble_results>); + + pymio::iterable_enum(m, "InfectionState") + .value("Susceptible", mio::ssir::InfectionState::Susceptible) + .value("Infected", mio::ssir::InfectionState::Infected) + .value("Recovered", mio::ssir::InfectionState::Recovered); + + pymio::bind_ParameterSet(m, "ParametersBase"); + + pymio::bind_class(m, + "Parameters") + .def(py::init<>()) + .def("check_constraints", &mio::ssir::Parameters::check_constraints) + .def("apply_constraints", &mio::ssir::Parameters::apply_constraints); + + using Populations = mio::Populations; + pymio::bind_Population(m, "Populations", mio::Tag{}); + pymio::bind_CompartmentalModel(m, "ModelBase"); + pymio::bind_class>( + m, "Model") + .def(py::init<>()); + + m.def( + "simulate_stochastic", + [](double t0, double tmax, double dt, mio::ssir::Model const& model) { + return mio::simulate_stochastic(t0, tmax, dt, model); + }, + "Simulates an SDE SIR model from t0 to tmax.", py::arg("t0"), py::arg("tmax"), py::arg("dt"), py::arg("model")); + + m.attr("__version__") = "dev"; +} diff --git a/pycode/memilio-simulation/memilio/simulation/bindings/models/ssirs.cpp b/pycode/memilio-simulation/memilio/simulation/bindings/models/ssirs.cpp new file mode 100644 index 0000000000..6554c5d025 --- /dev/null +++ b/pycode/memilio-simulation/memilio/simulation/bindings/models/ssirs.cpp @@ -0,0 +1,97 @@ +/* +* Copyright (C) 2020-2025 MEmilio +* +* Authors: Maximilian Betz +* +* Contact: Martin J. Kuehn +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +//Includes from pymio +#include "pybind_util.h" +#include "utils/index.h" +#include "compartments/compartmental_model.h" +#include "utils/custom_index_array.h" +#include "utils/parameter_set.h" +#include "epidemiology/populations.h" + +//Includes from MEmilio +#include "sde_sirs/model.h" +#include "sde_sirs/infection_state.h" +#include "memilio/compartments/stochastic_simulation.h" +#include "memilio/compartments/stochastic_model.h" +#include "memilio/data/analyze_result.h" + +#include "pybind11/pybind11.h" +#include +#include + +namespace py = pybind11; + +namespace pymio +{ +//specialization of pretty_name +template <> +inline std::string pretty_name() +{ + return "InfectionState"; +} + +} // namespace pymio + +PYBIND11_MODULE(_simulation_ssirs, m) +{ + m.def("interpolate_simulation_result", + static_cast (*)(const mio::TimeSeries&, const double)>( + &mio::interpolate_simulation_result), + py::arg("ts"), py::arg("abs_tol") = 1e-14); + + m.def("interpolate_simulation_result", + static_cast (*)(const mio::TimeSeries&, const std::vector&)>( + &mio::interpolate_simulation_result), + py::arg("ts"), py::arg("interpolation_times")); + + m.def("interpolate_ensemble_results", &mio::interpolate_ensemble_results>); + + pymio::iterable_enum(m, "InfectionState") + .value("Susceptible", mio::ssirs::InfectionState::Susceptible) + .value("Infected", mio::ssirs::InfectionState::Infected) + .value("Recovered", mio::ssirs::InfectionState::Recovered); + + pymio::bind_ParameterSet(m, "ParametersBase"); + + pymio::bind_class(m, + "Parameters") + .def(py::init<>()) + .def("check_constraints", &mio::ssirs::Parameters::check_constraints) + .def("apply_constraints", &mio::ssirs::Parameters::apply_constraints); + + using Populations = mio::Populations; + pymio::bind_Population(m, "Populations", mio::Tag{}); + pymio::bind_CompartmentalModel(m, "ModelBase"); + pymio::bind_class>( + m, "Model") + .def(py::init<>()); + + m.def( + "simulate_stochastic", + [](double t0, double tmax, double dt, mio::ssirs::Model const& model) { + return mio::simulate_stochastic(t0, tmax, dt, model); + }, + "Simulates an SDE SIR model from t0 to tmax.", py::arg("t0"), py::arg("tmax"), py::arg("dt"), py::arg("model")); + + m.attr("__version__") = "dev"; +} diff --git a/pycode/memilio-simulation/memilio/simulation/ssir.py b/pycode/memilio-simulation/memilio/simulation/ssir.py new file mode 100644 index 0000000000..d4e1e1be6b --- /dev/null +++ b/pycode/memilio-simulation/memilio/simulation/ssir.py @@ -0,0 +1,25 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Maximilian Betz +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + +""" +Python bindings for MEmilio SDE SIR model. +""" + +from memilio.simulation._simulation_ssir import * diff --git a/pycode/memilio-simulation/memilio/simulation/ssirs.py b/pycode/memilio-simulation/memilio/simulation/ssirs.py new file mode 100644 index 0000000000..91e9df249d --- /dev/null +++ b/pycode/memilio-simulation/memilio/simulation/ssirs.py @@ -0,0 +1,25 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: Maximilian Betz +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# + +""" +Python bindings for MEmilio SDE SIRS model. +""" + +from memilio.simulation._simulation_ssirs import * diff --git a/pycode/memilio-simulation/memilio/simulation_test/test_ssir.py b/pycode/memilio-simulation/memilio/simulation_test/test_ssir.py new file mode 100644 index 0000000000..7abc0eb55b --- /dev/null +++ b/pycode/memilio-simulation/memilio/simulation_test/test_ssir.py @@ -0,0 +1,127 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# +import unittest + +import numpy as np + +from memilio.simulation import Damping, LogLevel, set_log_level +from memilio.simulation.ssir import InfectionState as State +from memilio.simulation.ssir import ( + Model, + simulate_stochastic, + interpolate_simulation_result, +) + + +class Test_ssir_integration(unittest.TestCase): + """Integration tests for the SDE SIR Python bindings.""" + + set_log_level(LogLevel.Off) + + def setUp(self): + # Initialize model similar to the example + self.population = 10000 + self.days = 5.0 + self.dt = 0.1 + + model = Model() + + # Set parameters + model.parameters.TimeInfected.value = 10.0 + model.parameters.TransmissionProbabilityOnContact.value = 1.0 + + # Initial compartment counts + model.populations[State.Infected] = 100 + model.populations[State.Recovered] = 1000 + model.populations.set_difference_from_total( + State.Susceptible, self.population) + + # Contacts and damping + model.parameters.ContactPatterns.baseline = np.ones((1, 1)) * 2.7 + model.parameters.ContactPatterns.minimum = np.zeros((1, 1)) + model.parameters.ContactPatterns.add_damping( + Damping(coeffs=np.r_[0.6], t=12.5, level=0, type=0) + ) + + model.check_constraints() + self.model = model + + def test_simulate_simple(self): + result = simulate_stochastic(0.0, self.days, self.dt, self.model) + self.assertAlmostEqual(result.get_time(0), 0.0) + self.assertAlmostEqual(result.get_time(1), self.dt) + self.assertAlmostEqual(result.get_last_time(), self.days) + self.assertEqual(len(result.get_last_value()), 3) + + def test_check_constraints_parameters(self): + model = Model() + # valid + model.parameters.TimeInfected.value = 6.0 + model.parameters.TransmissionProbabilityOnContact.value = 1.0 + self.assertEqual(model.parameters.check_constraints(), 0) + # invalid: time <= 0.1 + model.parameters.TimeInfected.value = 0.0 + self.assertEqual(model.parameters.check_constraints(), 1) + # invalid: transmission < 0 + model.parameters.TimeInfected.value = 6.0 + model.parameters.TransmissionProbabilityOnContact.value = -1.0 + self.assertEqual(model.parameters.check_constraints(), 1) + # invalid: transmission > 1 + model.parameters.TransmissionProbabilityOnContact.value = 2.0 + self.assertEqual(model.parameters.check_constraints(), 1) + + def test_apply_constraints_parameters(self): + model = Model() + # set invalid values + model.parameters.TimeInfected.value = 0.0 + model.parameters.TransmissionProbabilityOnContact.value = -0.5 + corrected = model.parameters.apply_constraints() + self.assertTrue(corrected) + # after apply_constraints(), values should be set according to constraints + self.assertAlmostEqual( + model.parameters.TimeInfected.value, 0.1, places=6) + self.assertAlmostEqual( + model.parameters.TransmissionProbabilityOnContact.value, 0.0, + places=6) + + def test_population_conservation(self): + result = simulate_stochastic(0.0, self.days, self.dt, self.model) + total0 = self.population + atol = 1e-10 * total0 + for i in range(int(self.days / self.dt) + 1): + vals = result[i] + self.assertEqual(len(vals), 3) + # Non-negativity + self.assertTrue(np.all(np.array(vals) >= -1e-10)) + # Conservation of total population + self.assertAlmostEqual(np.sum(vals), total0, delta=atol) + + def test_interpolate_times(self): + result = simulate_stochastic(0.0, self.days, self.dt, self.model) + times = [0.0, 1.0, 2.0, 3.0, self.days] + interp = interpolate_simulation_result(result, times) + self.assertEqual(interp.get_time(0), times[0]) + self.assertEqual(interp.get_last_time(), times[-1]) + self.assertEqual(len(interp.get_value(0)), 3) + self.assertEqual(len(interp.get_value(len(times) - 1)), 3) + + +if __name__ == '__main__': + unittest.main() diff --git a/pycode/memilio-simulation/memilio/simulation_test/test_ssirs.py b/pycode/memilio-simulation/memilio/simulation_test/test_ssirs.py new file mode 100644 index 0000000000..0fe3b7d15a --- /dev/null +++ b/pycode/memilio-simulation/memilio/simulation_test/test_ssirs.py @@ -0,0 +1,136 @@ +############################################################################# +# Copyright (C) 2020-2025 MEmilio +# +# Authors: +# +# Contact: Martin J. Kuehn +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################# +import unittest + +import numpy as np + +from memilio.simulation import Damping, LogLevel, set_log_level +from memilio.simulation.ssirs import InfectionState as State +from memilio.simulation.ssirs import ( + Model, + simulate_stochastic, + interpolate_simulation_result, +) + + +class Test_ssirs_integration(unittest.TestCase): + """Integration tests for the SDE SIRS Python bindings.""" + + set_log_level(LogLevel.Off) + + def setUp(self): + # Initialize model similar to the sde_sirs_simple example + self.population = 10000 + self.days = 5.0 + self.dt = 0.1 + + model = Model() + + # Set parameters + model.parameters.TimeInfected.value = 10.0 + model.parameters.TimeImmune.value = 100.0 + model.parameters.TransmissionProbabilityOnContact.value = 1.0 + + # Initial compartment counts + model.populations[State.Infected] = 100 + model.populations[State.Recovered] = 1000 + model.populations.set_difference_from_total( + State.Susceptible, self.population) + + # Contacts and damping + model.parameters.ContactPatterns.baseline = np.ones((1, 1)) * 2.7 + model.parameters.ContactPatterns.minimum = np.zeros((1, 1)) + model.parameters.ContactPatterns.add_damping( + Damping(coeffs=np.r_[0.6], t=12.5, level=0, type=0) + ) + + model.check_constraints() + self.model = model + + def test_simulate_simple(self): + result = simulate_stochastic(0.0, self.days, self.dt, self.model) + self.assertAlmostEqual(result.get_time(0), 0.0) + self.assertAlmostEqual(result.get_time(1), self.dt) + self.assertAlmostEqual(result.get_last_time(), self.days) + self.assertEqual(len(result.get_last_value()), 3) + + def test_check_constraints_parameters(self): + model = Model() + # valid + model.parameters.TimeInfected.value = 6.0 + model.parameters.TimeImmune.value = 100.0 + model.parameters.TransmissionProbabilityOnContact.value = 1.0 + self.assertEqual(model.parameters.check_constraints(), 0) + # invalid: TimeInfected <= 0.1 + model.parameters.TimeInfected.value = 0.0 + self.assertEqual(model.parameters.check_constraints(), 1) + # invalid: TimeImmune <= 0.1 + model.parameters.TimeInfected.value = 6.0 + model.parameters.TimeImmune.value = 0.0 + self.assertEqual(model.parameters.check_constraints(), 1) + # invalid: transmission < 0 + model.parameters.TimeImmune.value = 100.0 + model.parameters.TransmissionProbabilityOnContact.value = -1.0 + self.assertEqual(model.parameters.check_constraints(), 1) + # invalid: transmission > 1 + model.parameters.TransmissionProbabilityOnContact.value = 2.0 + self.assertEqual(model.parameters.check_constraints(), 1) + + def test_apply_constraints_parameters(self): + model = Model() + # intentionally set invalid values + model.parameters.TimeInfected.value = 0.0 + model.parameters.TimeImmune.value = 0.0 + model.parameters.TransmissionProbabilityOnContact.value = -0.5 + corrected = model.parameters.apply_constraints() + self.assertTrue(corrected) + # after apply_constraints(), boundary values should be set + self.assertAlmostEqual( + model.parameters.TimeInfected.value, 0.1, places=6) + self.assertAlmostEqual( + model.parameters.TimeImmune.value, 0.1, places=6) + self.assertAlmostEqual( + model.parameters.TransmissionProbabilityOnContact.value, 0.0, + places=6) + + def test_population_conservation(self): + result = simulate_stochastic(0.0, self.days, self.dt, self.model) + total0 = self.population + atol = 1e-4 * total0 + for i in range(int(self.days / self.dt) + 1): + vals = result[i] + self.assertEqual(len(vals), 3) + # Non-negativity + self.assertTrue(np.all(np.array(vals) >= -1e-8)) + # Conservation of total population + self.assertAlmostEqual(np.sum(vals), total0, delta=atol) + + def test_interpolate_times(self): + result = simulate_stochastic(0.0, self.days, self.dt, self.model) + times = [0.0, 1.0, 2.0, 3.0, self.days] + interp = interpolate_simulation_result(result, times) + self.assertEqual(interp.get_time(0), times[0]) + self.assertEqual(interp.get_last_time(), times[-1]) + self.assertEqual(len(interp.get_value(0)), 3) + self.assertEqual(len(interp.get_value(len(times) - 1)), 3) + + +if __name__ == '__main__': + unittest.main()