1577 Add generic Time Series Plotting function

docs/source/python/m-plot.rst

@@ -19,7 +19,7 @@ Dependencies
 Required python packages:
 
 - pandas>=1.2.2
-- matplotlib
+- matplotlib>=3.6
 - numpy>=1.22, !=1.25.*
 - openpyxl
 - xlrd
@@ -32,3 +32,82 @@ Required python packages:
 - h5py
 - imageio
 - datetime
+
+Plotting simulation results
+---------------------------
+
+The module ``memilio.plot.plotTimeSeries`` provides a standard plot for ``TimeSeries`` objects as returned by the
+simulations of the :doc:`MEmilio Python bindings <m-simulation>`:
+
+.. code-block:: python
+
+    from datetime import date
+
+    import numpy as np
+
+    from memilio.plot.plotTimeSeries import plot_time_series
+    from memilio.simulation import AgeGroup, Damping
+    from memilio.simulation.osecir import InfectionState, Model, simulate
+
+    # ODE SECIR model with two age groups
+    groups = ['0-19', '20+']
+    populations = [15000, 68000]
+    num_groups = len(groups)
+
+    model = Model(num_groups)
+    for i, population in enumerate(populations):
+        group = AgeGroup(i)
+        # time spent in the compartments (days)
+        model.parameters.TimeExposed[group] = 3.2
+        model.parameters.TimeInfectedNoSymptoms[group] = 2.
+        model.parameters.TimeInfectedSymptoms[group] = 6.
+        model.parameters.TimeInfectedSevere[group] = 12.
+        model.parameters.TimeInfectedCritical[group] = 8.
+        # transmission and transition probabilities
+        model.parameters.TransmissionProbabilityOnContact[group] = 1.0
+        model.parameters.RelativeTransmissionNoSymptoms[group] = 0.67
+        model.parameters.RiskOfInfectionFromSymptomatic[group] = 0.25
+        model.parameters.MaxRiskOfInfectionFromSymptomatic[group] = 0.5
+        model.parameters.RecoveredPerInfectedNoSymptoms[group] = 0.09
+        model.parameters.SeverePerInfectedSymptoms[group] = 0.2
+        model.parameters.CriticalPerSevere[group] = 0.25
+        model.parameters.DeathsPerCritical[group] = 0.3
+        # initial populations
+        model.populations[group, InfectionState.Exposed] = 100
+        model.populations[group, InfectionState.InfectedNoSymptoms] = 50
+        model.populations[group, InfectionState.InfectedSymptoms] = 50
+        model.populations[group, InfectionState.InfectedSevere] = 20
+        model.populations[group, InfectionState.InfectedCritical] = 10
+        model.populations[group, InfectionState.Recovered] = 10
+        model.populations.set_difference_from_group_total_AgeGroup(
+            (group, InfectionState.Susceptible), population)
+
+    # contacts: one contact per person and day, reduced by 90% from day 30 on
+    model.parameters.ContactPatterns.cont_freq_mat[0].baseline = np.ones(
+        (num_groups, num_groups))
+    model.parameters.ContactPatterns.cont_freq_mat.add_damping(Damping(
+        coeffs=np.ones((num_groups, num_groups)) * 0.9, t=30., level=0, type=0))
+    model.check_constraints()
+
+    result = simulate(0, 100, 0.1, model)
+    ax = plot_time_series(
+        result, labels=InfectionState.values(), groups=groups,
+        select=['Exposed', 'InfectedSymptoms', 'Dead'],
+        start_date=date(2020, 3, 1), title='ODE SECIR simulation')
+    ax.figure.savefig('secir.pdf')
+
+- ``labels`` names the compartments. It accepts strings or the ``InfectionState.values()`` of a model. Without labels,
+  the elements are named ``C1``, ``C2``, ... .
+- ``groups`` (the number or the names of, e.g., age groups) tells the function how the elements of the ``TimeSeries``
+  are arranged. By default, the compartments are summed over all groups; with ``sum_groups=False`` one line per group
+  and compartment is drawn.
+- ``select`` restricts the plot to some compartments, ``log_scale`` switches to a logarithmic axis and ``start_date``
+  turns the time axis into calendar dates. Colors are fixed per compartment, so a selection does not change them.
+- The function returns a matplotlib ``Axes``. Pass ``ax=`` to draw into an existing figure (e.g., a subplot grid) and
+  use ``ax.figure.savefig`` to save the figure. The legend is placed right of the axes, so create your own figures with
+  ``layout='constrained'`` or save them with ``bbox_inches='tight'``.
+
+``time_series_to_dataframe`` converts a ``TimeSeries`` into a tidy pandas ``DataFrame`` with the columns ``Time``,
+``Groups``, ``Compartments`` and ``Values`` (and ``Date`` if a start date is given) for use with other libraries such as
+seaborn, plotnine or altair. A complete example is given in
+`examples/plot/plotSimulationResults.py <https://github.com/SciCompMod/memilio/blob/main/pycode/examples/plot/plotSimulationResults.py>`_.

docs/source/python/m-simulation_model_usage.rst

@@ -259,6 +259,15 @@ pythonic interface.
 
 Now you can use the usual data handling options and make use of the easy visualization tools that are part of Python.
 Some plotting functions specific to MEmilio and created as part of the project are combined in the :doc:`MEmilio Plot Package <m-plot>`.
+For a quick look at the result, it provides a standard plot for ``TimeSeries`` objects:
+
+.. code-block:: python
+
+   from memilio.plot.plotTimeSeries import plot_time_series
+
+   ax = plot_time_series(
+       result, labels=oseir.InfectionState.values(), groups=num_groups)
+   ax.figure.savefig('result.pdf')
 
 Additional resources
 ---------------------

pycode/examples/plot/plotSimulationResults.py

@@ -0,0 +1,128 @@
+#############################################################################
+# Copyright (C) 2020-2026 MEmilio
+#
+# Authors: Kilian Volmer
+#
+# Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
+#
+# 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.
+#############################################################################
+"""
+Example demonstrating the standard TimeSeries plot of the MEmilio plot
+package on the result of an ODE SEIR simulation with two age groups.
+"""
+import argparse
+import os
+from datetime import date
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from memilio.plot.plotTimeSeries import (plot_time_series,
+                                         time_series_to_dataframe)
+from memilio.simulation import AgeGroup, Damping
+from memilio.simulation.oseir import InfectionState, Model, simulate
+
+AGE_GROUPS = ['0-19', '20+']
+START_DATE = date(2020, 3, 1)
+
+
+def run_ode_seir_simulation(days=100, dt=0.1):
+    """ Runs the ODE SEIR model with two age groups.
+
+    :param days: Number of days to simulate. (Default value = 100)
+    :param dt: Initial time step. (Default value = 0.1)
+    :returns: Simulation result as TimeSeries.
+    """
+    group_populations = [15000, 68000]
+    num_groups = len(group_populations)
+    model = Model(num_groups)
+
+    for i, population in enumerate(group_populations):
+        group = AgeGroup(i)
+        model.parameters.TimeExposed[group] = 5.2
+        model.parameters.TimeInfected[group] = 6.
+        model.parameters.TransmissionProbabilityOnContact[group] = 1. * (i+1)
+        model.populations[group, InfectionState.Exposed] = 100
+        model.populations[group, InfectionState.Infected] = 50
+        model.populations[group, InfectionState.Recovered] = 10
+        model.populations.set_difference_from_group_total_AgeGroup(
+            (group, InfectionState.Susceptible), population)
+
+    model.parameters.ContactPatterns.cont_freq_mat[0].baseline = np.ones(
+        (num_groups, num_groups))
+    model.parameters.ContactPatterns.cont_freq_mat[0].minimum = np.zeros(
+        (num_groups, num_groups))
+    model.parameters.ContactPatterns.cont_freq_mat.add_damping(Damping(
+        coeffs=np.ones((num_groups, num_groups)) * 0.9, t=30.0, level=0,
+        type=0))
+
+    model.check_constraints()
+    return simulate(0, days, dt, model)
+
+
+def plot_results(result, output_path='.', show_plot=False):
+    """ Creates the standard plots of the simulation result.
+
+    :param result: TimeSeries returned by the simulation.
+    :param output_path: Directory the figures are written to.
+        (Default value = '.')
+    :param show_plot: Whether to show the figures interactively.
+        (Default value = False)
+    """
+    # All compartments, summed over the age groups. The compartment names
+    # are taken from the InfectionState enum of the model.
+    ax = plot_time_series(
+        result, labels=InfectionState.values(), groups=AGE_GROUPS,
+        title='ODE SEIR simulation')
+    ax.figure.savefig(
+        os.path.join(output_path, 'seir_compartments.png'), dpi=150)
+
+    # Two panels in one figure: a selection of compartments per age group
+    # on a date axis, and a selection on a logarithmic axis.
+    fig, axes = plt.subplots(1, 2, figsize=(12, 4.5), layout='constrained')
+    plot_time_series(
+        result, labels=InfectionState.values(), groups=AGE_GROUPS,
+        sum_groups=False, select=['Exposed', 'Infected'],
+        start_date=START_DATE, ax=axes[0],
+        title='Exposed and infected per age group')
+    plot_time_series(
+        result, labels=InfectionState.values(), groups=AGE_GROUPS,
+        select=['Exposed', 'Infected', 'Recovered'], log_scale=True,
+        start_date=START_DATE, ax=axes[1], title='Logarithmic scale')
+    fig.savefig(os.path.join(output_path, 'seir_selection.pdf'))
+
+    # The same data as a tidy data frame, e.g. for other plotting libraries.
+    df = time_series_to_dataframe(
+        result, labels=InfectionState.values(), groups=AGE_GROUPS,
+        start_date=START_DATE)
+    print(df.head())
+    print(df.pivot_table(index='Date', columns='Compartments', values='Values',
+                         aggfunc='sum', observed=False).tail())
+
+    if show_plot:
+        plt.show()
+    plt.close('all')
+
+
+if __name__ == '__main__':
+    arg_parser = argparse.ArgumentParser(
+        'plotSimulationResults',
+        description='Plots the result of an ODE SEIR simulation with the '
+        'standard TimeSeries plot of the MEmilio plot package.')
+    arg_parser.add_argument('-p', '--show_plot', action='store_true',
+                            help='Show the figures interactively.')
+    arg_parser.add_argument('-o', '--output_path', default='.',
+                            help='Directory the figures are written to.')
+    args = arg_parser.parse_args()
+    plot_results(run_ode_seir_simulation(), args.output_path, args.show_plot)

pycode/examples/simulation/ode_secir_groups.py

@@ -17,33 +17,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #############################################################################
+
+# This example requires memilio.simulation and memilio.plot to be installed!
+
 import argparse
 import os
-from datetime import date, datetime
+from datetime import date
 
 import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd
 
-from memilio.simulation import AgeGroup, ContactMatrix, Damping, UncertainContactMatrix
-from memilio.simulation.osecir import Index_InfectionState
+from memilio.plot.plotTimeSeries import plot_time_series
+from memilio.simulation import AgeGroup, Damping
 from memilio.simulation.osecir import InfectionState as State
-from memilio.simulation.osecir import (Model, Simulation,
-                                       interpolate_simulation_result, simulate)
+from memilio.simulation.osecir import (Model, interpolate_simulation_result,
+                                       simulate)
 
 
 def run_ode_secir_groups_simulation(show_plot=True):
     """Runs the c++ ODE SECIHURD model using mulitple age groups
     and plots the results
 
-    :param show_plot:  (Default value = True)
+    :param show_plot: Whether to show the figures interactively.
+        (Default value = True)
 
     """
 
-    # Define Comartment names
-    compartments = [
-        'Susceptible', 'Exposed', 'InfectedNoSymptoms', 'InfectedSymptoms',
-        'InfectedSevere', 'InfectedCritical', 'Recovered', 'Dead']
     # Define age Groups
     groups = ['0-4', '5-14', '15-34', '35-59', '60-79', '80+']
     # Define population of age groups
@@ -55,11 +54,10 @@ def run_ode_secir_groups_simulation(show_plot=True):
     start_year = 2019
     dt = 0.1
     num_groups = len(groups)
-    num_compartments = len(compartments)
 
     # set contact frequency matrix
     data_dir = os.path.join(os.path.dirname(
-        __file__), "..", "..", "..", "data")
+        __file__), "..", "..", "..", "data", "Germany")
     baseline_contact_matrix0 = os.path.join(
         data_dir, "contacts/baseline_home.txt")
     baseline_contact_matrix1 = os.path.join(
@@ -133,79 +131,38 @@ def run_ode_secir_groups_simulation(show_plot=True):
     # interpolate results
     result = interpolate_simulation_result(result)
 
-    # print(result.get_last_value())
-    num_time_points = result.get_num_time_points()
-    result_array = result.as_ndarray()
-    t = result_array[0, :]
-    group_data = np.transpose(result_array[1:, :])
-
-    # sum over all groups
-    data = np.zeros((num_time_points, num_compartments))
-    for i in range(num_groups):
-        data += group_data[:, i * num_compartments: (i + 1) * num_compartments]
-
-    # Plot Results
-    datelist = np.array(
-        pd.date_range(
-            datetime(start_year, start_month, start_day),
-            periods=days, freq='D').strftime('%m-%d').tolist())
-
-    tick_range = (np.arange(int(days / 10) + 1) * 10)
-    tick_range[-1] -= 1
-    fig, ax = plt.subplots()
-    ax.plot(t, data[:, 0], label='#Susceptible')
-    ax.plot(t, data[:, 1], label='#Exposed')
-    ax.plot(t, data[:, 2], label='#InfectedNoSymptoms')
-    ax.plot(t, data[:, 3], label='#InfectedSymptoms')
-    ax.plot(t, data[:, 4], label='#Hospitalzed')
-    ax.plot(t, data[:, 5], label='#InfectedCritical')
-    ax.plot(t, data[:, 6], label='#Recovered')
-    ax.plot(t, data[:, 7], label='#Dead')
-    ax.set_title("ODE SECIR simulation results (entire population)")
-    ax.set_xticks(tick_range)
-    ax.set_xticklabels(datelist[tick_range], rotation=45)
-    ax.legend()
-    fig.tight_layout
-    fig.savefig('osecir_by_compartments.pdf')
-
-    # plot dynamics in each comparment by age group
-    fig, ax = plt.subplots(4, 2, figsize=(12, 15))
-
-    for i, title in zip(range(num_compartments), compartments):
-
-        for j, group in enumerate(groups):
-            ax[int(np.floor(i / 2)), int(i % 2)].plot(t,
-                                                      group_data[:, j*num_compartments+i], label=group)
-
-        ax[int(np.floor(i / 2)), int(i % 2)].set_title(title, fontsize=10)
-        ax[int(np.floor(i / 2)), int(i % 2)].legend()
-
-        ax[int(np.floor(i / 2)), int(i % 2)].set_xticks(tick_range)
-        ax[int(np.floor(i / 2)), int(i % 2)
-           ].set_xticklabels(datelist[tick_range], rotation=45)
-    plt.subplots_adjust(hspace=0.5, bottom=0.1, top=0.9)
+    start_date = date(start_year, start_month, start_day)
+
+    # Plot results summed over all age groups, one line per compartment. The
+    # compartment names are taken from the InfectionState enum of the model.
+    ax = plot_time_series(
+        result, labels=State.values(), groups=groups, start_date=start_date,
+        title='ODE SECIR simulation results (entire population)')
+    ax.figure.savefig('osecir_by_compartments.pdf')
+
+    # One panel per compartment with one line per age group.
+    fig, axes = plt.subplots(5, 2, figsize=(16, 18), layout='constrained')
+    for state, ax in zip(State.values(), axes.flat):
+        plot_time_series(
+            result, labels=State.values(), groups=groups, sum_groups=False,
+            select=state, start_date=start_date, ax=ax, title=state.name)
     fig.suptitle(
         'ODE SECIR simulation results by age group in each compartment')
     fig.savefig('osecir_age_groups_in_compartments.pdf')
 
-    fig, ax = plt.subplots(4, 2, figsize=(12, 15))
-    for i, title in zip(range(num_compartments), compartments):
-        ax[int(np.floor(i / 2)), int(i % 2)].plot(t, data[:, i])
-        ax[int(np.floor(i / 2)), int(i % 2)].set_title(title, fontsize=10)
-
-        ax[int(np.floor(i / 2)), int(i % 2)].set_xticks(tick_range)
-        ax[int(np.floor(i / 2)), int(i % 2)
-           ].set_xticklabels(datelist[tick_range], rotation=45)
-    plt.subplots_adjust(hspace=0.5, bottom=0.1, top=0.9)
+    # One panel per compartment, summed over all age groups.
+    fig, axes = plt.subplots(5, 2, figsize=(16, 18), layout='constrained')
+    for state, ax in zip(State.values(), axes.flat):
+        plot_time_series(
+            result, labels=State.values(), groups=groups, select=state,
+            start_date=start_date, ax=ax, title=state.name)
     fig.suptitle(
         'ODE SECIR simulation results by compartment (entire population)')
     fig.savefig('osecir_all_parts.pdf')
 
     if show_plot:
         plt.show()
-        plt.close()
-
-    # return data
+    plt.close('all')
 
 
 if __name__ == "__main__":