diff --git a/epispot/plots/native.py b/epispot/plots/native.py index 0c90331..8219f92 100644 --- a/epispot/plots/native.py +++ b/epispot/plots/native.py @@ -96,3 +96,72 @@ def model(Model, time_frame, title='Compartment Populations over Time', plt.legend() return plt + + +def stacked(Model, time_frame, title='Compartment Populations over Time', + starting_state=None, compartments=None, names=None, show_susceptible=False, + log=False, latex=True): + """ + Plots the results of one model using `matplotlib`. + The results are displayed natively via a `matplotlib` window as a stacked area chart. + There are various ways to customize the generated plots by modifying + the time frame or compartments displayed. + Additionally, `matplotlib` allows editing plots even after + they have been created to change things like colors, margins, etc. + + - Model: An `epispot.models.Model` object + - time_frame: A `range()` describing the time period to plot + - title: (`='Compartment Populations over Time`) The title of the plot + - starting_state: (default:inherited) Initial model state (see `epispot.models.Model.integrate` parameter `starting_state`) + - compartments: (default:all) The indices of the compartments in the model to plot; + all other compartments will be hidden + - names: (default:`Model.layer_names`) A list of names for each of the compartments + - show_susceptible: (`=False`) Boolean value describing whether or not to plot the Susceptible compartment.\ + **This assumes that the Susceptible compartment is the first in `Model`**\ + Note:\ + > This can potentially result in less visibility for other compartments + > since usually the Susceptible compartment comprises of many, many + > more individuals than the other compartments combined. + - log: (`=False`) Boolean value indicating whether or not to use a logarithmic scale when plotting `Model` + - latex: (`=True`) Turn off if you do not have LaTeX installed or want quicker loading times + - return: `pyplot` figure (display with `.show()`) + """ + + if latex: + plt.style.use('science') + + DataFrame = {} + System = Model.integrate(time_frame, starting_state=starting_state) + + # parameter substitutions + if compartments is None: + compartments = list(range(len(Model.layers))) + + if names is None: + names = Model.layer_names + + # setup + for name in Model.layer_names: + DataFrame[name] = [] + + for day in System: + for i, compartment in enumerate(day): + DataFrame[Model.layer_names[i]].append(compartment) + + if not show_susceptible: + for i, compartment in enumerate(compartments): + if compartment == 0: + del compartments[i] + break + + # plotting + plt.figure(figsize=(9, 5)) + plt.stackplot(time_frame, *[DataFrame[Model.layer_names[compartment]] for compartment in compartments], + labels=[names[compartment] for compartment in compartments]) + + if log: + plt.yscale('log') + plt.title(title) + plt.legend(loc='upper left', facecolor='white') + + return plt diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 9ca7350..9266c58 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -16,6 +16,9 @@ |- native |- plain |- full + |- native-stack + |- plain + |- full """ import epispot as epi @@ -115,3 +118,25 @@ def test_full_native(): show_susceptible=True, log=True) return Figure + + +def test_plain_native_stack(): + Model = epi.pre.SEIR(R_0, N, place, gamma, delta) # compile model + Figure = epi.plots.native.stacked(Model, range(120), latex=False) # `latex=False` flag speeds up testing + return Figure + + +def test_full_native_stack(): + Model = epi.pre.SEIR(R_0, N, place, gamma, delta) # compile model + #`latex=True` enabled by default + # blank strings `''` in `names=` are used to indicate that + # the corresponding compartment is not being plotted + Figure = epi.plots.native.stacked(Model, range(120), + title='SEIR Model Plot', + starting_state=[N(0) - 50, 25, 25, 0], + compartments=[0, 3], + names=['Susceptible Population', + '', '', 'Removed/Recovered'], + show_susceptible=True, + log=True) + return Figure