modify plotting

datasets/transform_SIR.py

@@ -1,15 +1,23 @@
 import numpy as np
 import pandas as pd
 import matplotlib.pylab as plt
+from matplotlib import rcParams
 
-dataset_path = 'Covid_19_DINN/datasets/COVID-19-Todesfaelle_in_Deutschland/'
+dataset_path = 'datasets/COVID-19-Todesfaelle_in_Deutschland/'
 
-def transform_general_data():
+FONT_COLOR = '#595959'
+SUSCEPTIBLE = '#6399f7'
+INFECTIOUS = '#f56262'
+REMOVED = '#83eb5e'
+
+def transform_general_data(plot_name: str, plot_title: str, exclude=[], plot_size=(12,6), yscale_log=False, plot_legend=True):
     """Function to generate the SIR split from the data in the COVID-19-Todesfaelle_in_Deutschland dataset.
     """
     # read the data
     df = pd.read_csv(dataset_path + 'COVID-19-Todesfaelle_Deutschland.csv')
 
+    df = df.drop(df.index[1200:])
+
     # population of germany at the end of 2019
     N = 83100000
     S, I, R = np.zeros(df.shape[0]), np.zeros(df.shape[0]), np.zeros(df.shape[0])
@@ -43,44 +51,43 @@ def transform_general_data():
 
     # plot graphs
     t = np.arange(0, df.shape[0], 1)
-    plt.plot(t, S, label='Susceptible')
-    plt.plot(t, I, label='Infectionous')
-    plt.plot(t, R, label='Removed')
-    plt.legend()
-    # plt.yscale('log')
-    plt.savefig('Covid_19_DINN/visualizations/SIR.png')
-
-def transform_state_data(state='Sachsen', N=4100000):
-    df = pd.read_csv(dataset_path + 'COVID-19-Todesfaelle_Bundeslaender.csv')
-    S, I, R = np.zeros(df.shape[0]), np.zeros(df.shape[0]), np.zeros(df.shape[0])
-    S[0] = N - df['Faelle_gesamt'][0]
-    I[0] = df['Faelle_gesamt'][0] - df['Todesfaelle_gesamt'][0]
-    R[0] = df['Todesfaelle_gesamt'][0]
-    recovery_queue = np.zeros(14)
+    rcParams['font.family'] = 'Comfortaa'
+    rcParams['font.size'] = 12
+
+    rcParams['text.color'] = FONT_COLOR
+    rcParams['axes.labelcolor'] = FONT_COLOR
+    rcParams['xtick.color'] = FONT_COLOR
+    rcParams['ytick.color'] = FONT_COLOR
+
+    slide3 = plt.figure(figsize=plot_size)
+    ax = slide3.add_subplot(111, facecolor='#dddddd', axisbelow=True)
+    ax.set_facecolor('xkcd:white')
+
+    if 'S' not in exclude:
+        ax.plot(t, S, label='Susceptible', c=SUSCEPTIBLE, lw=3)
     
-    for day in range(1, df.shape[0]):
-        infections = df['Faelle_gesamt'][day] - df['Faelle_gesamt'][day-1]
-        deaths = df['Todesfaelle_neu'][day]
-        recoveries = recovery_queue[0]
+    if 'I' not in exclude:
+        ax.plot(t, I, label='Infectious', c=INFECTIOUS, lw=3)
 
-        S[day] = S[day-1] - infections
-        I[day] = I[day-1] + infections - deaths - recoveries
-        R[day] = R[day-1] + deaths + recoveries
+    if 'R' not in exclude:
+        ax.plot(t, R, label='Removed', c=REMOVED, lw=3)
 
-        if I[day] < 0:
-            recovery_queue[-1] -= I[day] 
-            I[day] = 0
+    if yscale_log:
+        plt.yscale('log')
 
-        recovery_queue[:-1] = recovery_queue[1:]
-        recovery_queue[-1] = infections
+    plt.ylabel('amount of poeple')
+    plt.xlabel('time / days')
+    plt.title(plot_title)
+    ax.yaxis.set_tick_params(length=0)
 
-    t = np.arange(0, df.shape[0], 1)
-    plt.plot(t, S, label='Susceptible')
-    plt.plot(t, I, label='Infectionous')
-    plt.plot(t, R, label='Removed')
-    plt.legend()
-    # plt.yscale('log')
-    plt.savefig('Covid_19_DINN/visualizations/SIR.png')
+    if plot_legend:
+        plt.legend()
+
+    ax.yaxis.set_tick_params(length=0, which='both')
+    ax.xaxis.set_tick_params(length=0, which='both')
+    ax.grid(which='major', c='black', lw=0.2, ls='-')
 
-transform_general_data()
+    for spine in ('top', 'right', 'bottom', 'left'):
+        ax.spines[spine].set_visible(False)
 
+    slide3.savefig(f'visualizations/{plot_name}.png', transparent=True)