add animation function

src/dinn.py

@@ -1,15 +1,15 @@
 import torch
+import os
+import imageio
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib import rcParams
 
 from .dataset import PandemicDataset
 from .problem import PandemicProblem
+from .plotter import Plotter
+
 
-FONT_COLOR = '#595959'
-SUSCEPTIBLE = '#6399f7'
-INFECTIOUS = '#f56262'
-REMOVED = '#83eb5e'
 
 class DINN:
     class NN(torch.nn.Module):
@@ -45,6 +45,7 @@ class DINN:
                  data: PandemicDataset,
                  parameter_list: list,
                  problem: PandemicProblem,
+                 plotter: Plotter,
                  parameter_regulator=torch.tanh,
                  input_size=1,
                  hidden_size=20,
@@ -65,22 +66,28 @@ class DINN:
             activation_layer (optional): Class of the activation function. Defaults to torch.nn.ReLU().
         """
         
-        self.device = torch.device('cpu')
+        self.device = torch.device(data.device_name)
+        self.device_name = data.device_name
+        self.plotter = plotter
 
         self.model = DINN.NN(number_groups, input_size, hidden_size, hidden_layers, activation_layer)
+        self.model = self.model.to(self.device)
         self.data = data
         self.parameter_regulator = parameter_regulator
         self.problem = problem
 
         self.parameters_tilda = {}
         for parameter in parameter_list:
-            self.parameters_tilda.update({parameter : torch.nn.Parameter(torch.rand(1, requires_grad=True))})
+            self.parameters_tilda.update({parameter : torch.nn.Parameter(torch.rand(1, requires_grad=True, device=self.device_name))})
         
         self.epochs = None
 
         self.losses = np.zeros(1)
         self.parameters = [np.zeros(1) for _ in range(len(parameter_list))]
 
+        self.frames = []
+
+
     def get_regulated_param(self, parameter_name: str):
         """Function to get the searched parameters, forced into a certain range.
 
@@ -111,13 +118,17 @@ class DINN:
     def train(self, 
               epochs: int, 
               lr: float, 
-              optimizer_class=torch.optim.Adam):
+              optimizer_class=torch.optim.Adam,
+              create_animation=False,
+              animation_sample_rate=500):
         """Training routine for the DINN
 
         Args:
             epochs (int): Number of epochs the NN is supposed to be trained for.
             lr (float): Learning rate for the optimizer.
             optimizer_class (optional): Class of the optimizer. Defaults to torch.optim.Adam.
+            create_animation (boolean, optional): Decides on wether a prediction animation is supposed to be created during training. Defaults to False.
+            animation_sample_rate (int, optional): Sample rate of the prediction animation. Only used, when create_animation=True. Defaults to 500.
         """
 
         # define optimizer and scheduler
@@ -133,7 +144,7 @@ class DINN:
         for epoch in range(epochs):
             # get the prediction and the fitting residuals
             prediction = self.model(self.data.t_batch)
-            residuals = self.problem.residual(prediction, self.data, *self.get_regulated_param_list())
+            residuals = self.problem.residual(prediction, *self.get_regulated_param_list())
 
             optimizer.zero_grad()
 
@@ -158,6 +169,27 @@ class DINN:
             for i, parameter in enumerate(self.parameters_tilda.items()):
                 self.parameters[i][epoch] = self.get_regulated_param(parameter[0]).item()
 
+            # do snapshot for prediction animation
+            if epoch % animation_sample_rate == 0 and create_animation:
+                # prediction
+                prediction = self.model(self.data.t_batch)
+                t = torch.arange(0, self.data.t_raw[-1].item(), (self.data.t_raw[-1] / self.data.t_raw.shape[0]).item())
+                
+                groups = self.problem.denormalization(prediction)
+                self.plotter.plot(t, 
+                                  groups + tuple(self.data.get_data()), 
+                                  [name + '_pred' for name in self.data.get_group_names()] + [name + '_true' for name in self.data.get_group_names()],
+                                  'frame',
+                                  f'epoch {epoch}',
+                                  figure_shape=(12, 6),
+                                  is_frame=True,
+                                  is_background=[0, 0, 0, 1, 1, 1],
+                                  lw=3,
+                                  legend_loc='upper right',
+                                  ylim=(0, self.data.N), 
+                                  xlabel='time / days',
+                                  ylabel='amount of people')
+
             # print training advancements
             if epoch % 1000 == 0:          
                 print('\nEpoch ', epoch)
@@ -167,7 +199,11 @@ class DINN:
                 print('---------------------------------')
                 for parameter in self.parameters_tilda.items():
                     print(f'{parameter[0]}:\t\t\t{self.parameter_regulator(parameter[1]).item()}')
-                print('#################################') 
+                print('#################################')
+
+        # create prediction animation
+        if create_animation:
+            self.plotter.animate(self.data.name + '_animation')
 
     def plot_training_graphs(self, ground_truth=[]):
         """Plot the loss graph and the graphs of the advancements of the parameters.
@@ -178,57 +214,14 @@ class DINN:
         assert self.epochs != None
         epochs = np.arange(0, self.epochs, 1)
 
-        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
-
         # plot loss
-        plt.plot(epochs, self.losses)
-        plt.title('Loss')
-        plt.yscale('log')
-        plt.show()
-
+        self.plotter.plot(epochs, [self.losses], ['loss'], self.data.name + '_loss', 'Loss', (6, 6), y_log_scale=True, plot_legend=False, xlabel='epochs')
+        
         # plot parameters
         for i, parameter in enumerate(self.parameters):
-            figure = plt.figure(figsize=(6,6))
-            ax = figure.add_subplot(111, facecolor='#dddddd', axisbelow=True)
-            ax.set_facecolor('xkcd:white')
-
-            ax.plot(epochs, parameter, c=FONT_COLOR, lw=3, label='prediction')
             if len(ground_truth) > i:
-                ax.axhline(y=ground_truth[i], color=INFECTIOUS, linestyle='-', lw=3, label='ground truth')
-        
-            plt.xlabel('epochs')
-            plt.title(list(self.parameters_tilda.items())[i][0])
-            ax.yaxis.set_tick_params(length=0)
-
-            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='-')
-
-            plt.legend()
-
-            for spine in ('top', 'right', 'bottom', 'left'):
-                ax.spines[spine].set_visible(False)
-
-            figure.savefig(f'visualizations/{list(self.parameters_tilda.items())[i][0]}.png', transparent=True)
-            plt.show()
-
-    def to_cuda(self):
-        """Move training to cuda device
-        """
-        assert torch.cuda.is_available(), "CUDA is not available"
-        
-        self.device = torch.device('cuda')
-        self.model = self.model.to(self.device)
-        self.data.to_device(self.device)
-        self.problem.to_device(self.device)
-
-        for parameter in self.parameters_tilda:
-            self.parameters_tilda[parameter] = self.parameters_tilda[parameter].to(self.device).detach().requires_grad_()
-        
+                self.plotter.plot(epochs, [parameter, np.ones_like(epochs) * ground_truth[i]], ['prediction', 'ground truth'], self.data.name + '_' + list(self.parameters_tilda.items())[i][0], list(self.parameters_tilda.items())[i][0], (6,6), is_background=[0, 1], xlabel='epochs')
+            else:
+                self.plotter.plot(epochs, [parameter], ['prediction'], self.data.name + '_' + list(self.parameters_tilda.items())[i][0], list(self.parameters_tilda.items())[i][0], (6,6), xlabel='epochs', plot_legend=False)
 
+