Merge pull request UofUEpiBio#18 from IsaccBarker/get-hist-total

Implement `get_hist_total`

README.md

@@ -5,7 +5,7 @@
 Build](https://github.com/UofUEpiBio/epiworldpy/actions/workflows/pip.yaml/badge.svg)](https://github.com/UofUEpiBio/epiworldpy/actions/workflows/pip.yaml)
 
 This is a python wrapper of the [`epiworld c++`
-library](https://github.com/UofUEpiBio/epiworld), an ABM simulation
+library](https://github.com/UofUEpiBio/epiworld/), an ABM simulation
 engine. This is possible using the
 [`pybind11`](https://pybind11.readthedocs.io/en/stable/) library (which
 rocks!).
@@ -21,15 +21,17 @@ target="_blank">implemented in R</a>.
 
 # Examples
 
-## Base setup
+## Basic
 
 Here we show how to create a `SEIR` object and add terms to it. We will
 use the following data:
 
 ``` python
 # Loading the module
-import epiworldpy as m
-covid19 = m.ModelSEIR(
+import epiworldpy as epiworld
+
+# Create a SEIR model (susceptible, exposed, infectious, recovered), representing COVID-19.
+covid19 = epiworld.ModelSEIR(
   name              = 'covid-19',
   n                 = 10000,
   prevalence        = .01,
@@ -38,12 +40,8 @@ covid19 = m.ModelSEIR(
   incubation_days   = 7.0,
   recovery_rate     = 0.14
 )
-```
-
-We can now take a look at the model
 
-``` python
-# Creating the object
+# Taking a look
 covid19.print(False)
 ```
 
@@ -60,8 +58,8 @@ covid19.print(False)
     Last run elapsed t  : -
     Rewiring            : off
 
-    Global actions:
-     (none)
+    Global events:
+     - Update infected individuals (runs daily)
 
     Virus(es):
      - covid-19 (baseline prevalence: 1.00%)
@@ -75,12 +73,15 @@ covid19.print(False)
      - Prob. Recovery       : 0.1400
      - Prob. Transmission   : 0.1000
 
-    <epiworldpy._core.ModelSEIRCONN at 0x10522f7f0>
+    <epiworldpy._core.ModelSEIRCONN at 0x105d14930>
 
-And run it and see what we get
+Let’s run it and to see what we get:
 
 ``` python
+# Run for 100 days with a seed of 223.
 covid19.run(100, 223)
+
+# Print an overview.
 covid19.print(False)
 ```
 
@@ -98,12 +99,12 @@ covid19.print(False)
     Number of entities  : 0
     Days (duration)     : 100 (of 100)
     Number of viruses   : 1
-    Last run elapsed t  : 49.00ms
-    Last run speed      : 20.34 million agents x day / second
+    Last run elapsed t  : 13.00ms
+    Last run speed      : 73.78 million agents x day / second
     Rewiring            : off
 
-    Global actions:
-     (none)
+    Global events:
+     - Update infected individuals (runs daily)
 
     Virus(es):
      - covid-19 (baseline prevalence: 1.00%)
@@ -118,17 +119,59 @@ covid19.print(False)
      - Prob. Transmission   : 0.1000
 
     Distribution of the population at time 100:
-      - (0) Susceptible :  9900 -> 7500
-      - (1) Exposed     :   100 -> 261
-      - (2) Infected    :     0 -> 238
-      - (3) Recovered   :     0 -> 2001
+      - (0) Susceptible :  9900 -> 7275
+      - (1) Exposed     :   100 -> 269
+      - (2) Infected    :     0 -> 292
+      - (3) Recovered   :     0 -> 2164
 
     Transition Probabilities:
      - Susceptible  1.00  0.00  0.00  0.00
-     - Exposed      0.00  0.86  0.14  0.00
+     - Exposed      0.00  0.85  0.15  0.00
      - Infected     0.00  0.00  0.86  0.14
      - Recovered    0.00  0.00  0.00  1.00
 
-    <epiworldpy._core.ModelSEIRCONN at 0x10522f7f0>
+    <epiworldpy._core.ModelSEIRCONN at 0x105d14930>
+
+We can know visualize the resulting time series:
+
+``` python
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Get the data from the database
+history = covid19.get_db().get_hist_total()
+
+# Extract unique states and dates
+unique_states = np.unique(history['states'])
+unique_dates = np.unique(history['dates'])
+
+# Remove some data that will mess with scaling
+unique_states = np.delete(unique_states, np.where(unique_states == 'Susceptible'))
+
+# Initialize a dictionary to store time series data for each state
+time_series_data = {state: [] for state in unique_states}
+
+# Populate the time series data for each state
+for state in unique_states:
+  for date in unique_dates:
+    # Get the count for the current state and date
+    mask = (history['states'] == state) & (history['dates'] == date)
+    count = history['counts'][mask][0]
+    time_series_data[state].append(count)
+
+# Start the plotting!
+plt.figure(figsize=(10, 6))
+
+for state in unique_states:
+  plt.plot(unique_dates, time_series_data[state], marker='o', label=state)
+
+plt.xlabel('Day')
+plt.ylabel('Count')
+plt.title('COVID-19 SEIR Model Data')
+plt.legend()
+plt.grid(True)
+plt.show()
+```
 
-# Acknowledgements
+![The data resulting from the COVID-19 SEIR model
+run](README_files/figure-commonmark/series-visualization-output-1.png)

README.qmd

@@ -10,40 +10,26 @@ This is a python wrapper of the [`epiworld c++` library][epiworld-git], an ABM
 simulation engine. This is possible using the
 [`pybind11`][] library (which rocks!).
 
-[`pybind11`]: https://pybind11.readthedocs.io/en/stable/
-
 The `epiworld` module is already
 [implemented in R](https://github.com/UofUEpiBio/epiworldR){target="_blank"}.
 
-
-[epiworld-git]: https://github.com/UofUEpiBio/epiworld
-[gitter-badge]:            https://badges.gitter.im/pybind/Lobby.svg
-[gitter-link]:             https://gitter.im/pybind/Lobby
-[actions-badge]:           https://github.com/UofUEpiBio/epiworldpy/workflows/Tests/badge.svg
-[actions-conda-link]:      https://github.com/UofUEpiBio/epiworldpy/actions?query=workflow%3AConda
-[actions-conda-badge]:     https://github.com/UofUEpiBio/epiworldpy/actions/workflows/conda.yml/badge.svg
-[actions-pip-link]:        https://github.com/UofUEpiBio/epiworldpy/actions?query=workflow%3APip
-[actions-pip-badge]:       https://github.com/UofUEpiBio/epiworldpy/actions/workflows/pip.yml/badge.svg
-[actions-wheels-link]:     https://github.com/UofUEpiBio/epiworldpy/actions?query=workflow%3AWheels
-[actions-wheels-badge]:    https://github.com/UofUEpiBio/epiworldpy/workflows/Wheels/badge.svg
-
 # Installation
 
-
 - clone this repository
 - `pip install ./epiworldpy`
 
-
 # Examples 
 
-## Base setup
+## Basic
 
 Here we show how to create a `SEIR` object and add terms to it. We will use the following data:
 
 ```{python}
 # Loading the module
-import epiworldpy as m
-covid19 = m.ModelSEIR(
+import epiworldpy as epiworld
+
+# Create a SEIR model (susceptible, exposed, infectious, recovered), representing COVID-19.
+covid19 = epiworld.ModelSEIR(
   name              = 'covid-19',
   n                 = 10000,
   prevalence        = .01,
@@ -53,24 +39,63 @@ covid19 = m.ModelSEIR(
   recovery_rate     = 0.14
 )
 
+# Taking a look
+covid19.print(False)
 ```
 
-We can now take a look at the model
+Let's run it and to see what we get:
 
 ```{python}
-# Creating the object
+# Run for 100 days with a seed of 223.
+covid19.run(100, 223)
+
+# Print an overview.
 covid19.print(False)
 ```
 
-And run it and see what we get
+We can know visualize the resulting time series:
 
 ```{python}
-covid19.run(100, 223)
-covid19.print(False)
-```
+#| label: series-visualization
+#| fig-cap: "The data resulting from the COVID-19 SEIR model run"
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Get the data from the database
+history = covid19.get_db().get_hist_total()
+
+# Extract unique states and dates
+unique_states = np.unique(history['states'])
+unique_dates = np.unique(history['dates'])
 
-[`cibuildwheel`]:          https://cibuildwheel.readthedocs.io
+# Remove some data that will mess with scaling
+unique_states = np.delete(unique_states, np.where(unique_states == 'Susceptible'))
 
+# Initialize a dictionary to store time series data for each state
+time_series_data = {state: [] for state in unique_states}
 
-# Acknowledgements
+# Populate the time series data for each state
+for state in unique_states:
+  for date in unique_dates:
+    # Get the count for the current state and date
+    mask = (history['states'] == state) & (history['dates'] == date)
+    count = history['counts'][mask][0]
+    time_series_data[state].append(count)
+
+# Start the plotting!
+plt.figure(figsize=(10, 6))
+
+for state in unique_states:
+  plt.plot(unique_dates, time_series_data[state], marker='o', label=state)
+
+plt.xlabel('Day')
+plt.ylabel('Count')
+plt.title('COVID-19 SEIR Model Data')
+plt.legend()
+plt.grid(True)
+plt.show()
+```
 
+[epiworld-git]: https://github.com/UofUEpiBio/epiworld/
+[`pybind11`]: https://pybind11.readthedocs.io/en/stable/

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["scikit-build-core>=0.3.3", "pybind11"]
+requires = ["scikit-build-core>=0.3.3", "pybind11", "matplotlib>=3.5.0"]
 build-backend = "scikit_build_core.build"
 
 

src/main.cpp

@@ -6,6 +6,7 @@
 #include "epiworld-common.hpp"
 
 using namespace epiworld;
+using namespace pybind11::literals;
 
 namespace py = pybind11;
 
@@ -52,7 +53,32 @@ PYBIND11_MODULE(_core, m) {
            ModelSEIR
     )pbdoc";
 
-    // Only this is necesary to expose the class
+    // Only this is necessary to expose the class
+    py::class_<Model<int>, std::shared_ptr<Model<int>>>(m, "Model")
+        .def("get_name", &Model<int>::get_name);
+
+    // Only this is necessary to expose the class
+    py::class_<DataBase<int>, std::shared_ptr<DataBase<int>>>(m, "DataBase")
+        .def("get_hist_total", [](DataBase<int> &self) {
+            std::vector<std::string> states;
+            std::vector<int> dates;
+            std::vector<int> counts;
+
+            self.get_hist_total(&dates, &states, &counts);
+
+            py::array py_states = py::array(py::cast(states));
+            py::array_t<int> py_dates(dates.size(), dates.data());
+            py::array_t<int> py_counts(counts.size(), counts.data());
+
+            py::dict ret(
+                "dates"_a=py_dates,
+                "states"_a=py_states,
+                "counts"_a=py_counts);
+
+            return ret;
+        });
+
+    // Only this is necessary to expose the class
     py::class_<epimodels::ModelSEIRCONN<int>, std::shared_ptr<epimodels::ModelSEIRCONN<int>>>(m, "ModelSEIRCONN")
         // .def(py::init<>())
         .def("print", &epimodels::ModelSEIRCONN<int>::print,
@@ -69,7 +95,11 @@ PYBIND11_MODULE(_core, m) {
         )pbdoc",
             py::arg("ndays"),
             py::arg("seed") = 1u
-            );
+            )
+        .def("get_db", [](epimodels::ModelSEIRCONN<int> &self) {
+            //std::cout << "!!! " << self.get_db().get_model()->get_name() << std::endl;
+            return std::shared_ptr<DataBase<int>>(&self.get_db(), [](DataBase<int>*){ /* do nothing, no delete */ });
+        }, py::return_value_policy::reference);
 
     // Example with shared_ptr
     m.def("ModelSEIR", &ModelSEIR, R"pbdoc(

tests/test_db.py

@@ -0,0 +1,30 @@
+import epiworldpy as epiworld
+
+def test_db_simple():
+    DAYS = 100
+
+    covid19 = epiworld.ModelSEIR(
+        name              = 'covid-19',
+        n                 = 10000,
+        prevalence        = .01,
+        contact_rate      = 2.0,
+        transmission_rate = .1,
+        incubation_days   = 7.0,
+        recovery_rate     = 0.14
+    )
+
+    covid19.run(DAYS, 223)
+
+    history = covid19.get_db().get_hist_total()
+    dates = history['dates']
+    states = history['states']
+    counts = history['counts']
+
+    # Considering that the SEIR model has four states (susceptible, exposed, infected, and
+    # recovered), we expect DAYS + 1 * 4 (we do the plus one since the resulting time series
+    # starts at 0 and the upper bound is treated as inclusive by epiworld).
+    EXPECTED_ENTRIES = (DAYS + 1) * 4
+
+    assert len(dates) == EXPECTED_ENTRIES
+    assert len(states) == EXPECTED_ENTRIES
+    assert len(counts) == EXPECTED_ENTRIES

tests/test_seir.py

@@ -11,6 +11,5 @@ def test_seir_simple():
         recovery_rate     = 0.14
     )
 
-    covid19.print(False)
     covid19.run(100, 223)