Covid_19_DINN/src/dinn.py

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



class DINN:
    class NN(torch.nn.Module):
        def __init__(self, 
                     output_size: int,
                     input_size: int,
                     hidden_size: int,
                     hidden_layers: int, 
                     activation_layer) -> None:
            """Neural Network

            Args:
                output_size (int): number of outputs
                input_size (int): number of inputs
                hidden_size (int): number of hidden nodes per layer
                hidden_layers (int): number of hidden layers
                activation_layer (_type_): activation layer
            """
            super(DINN.NN, self).__init__()

            self.input = torch.nn.Sequential(torch.nn.Linear(input_size, hidden_size), activation_layer)
            self.hidden = torch.nn.Sequential(*[torch.nn.Sequential(torch.nn.Linear(hidden_size, hidden_size), activation_layer) for _ in range(hidden_layers)])
            self.output = torch.nn.Linear(hidden_size, output_size)

        def forward(self, t):
            x = self.input(t)
            x = self.hidden(x)
            x = self.output(x)
            return x

    def __init__(self,
                 number_groups: int,
                 data: PandemicDataset,
                 parameter_list: list,
                 problem: PandemicProblem,
                 plotter: Plotter,
                 parameter_regulator=torch.tanh,
                 input_size=1,
                 hidden_size=20,
                 hidden_layers=7, 
                 activation_layer=torch.nn.ReLU()) -> None:
        """Desease Informed Neural Network. Uses the PandemicProblem, DINN.NN and PandemicDataset to solve Inverse Problems and find the 
        parameters of a specific mathematical model.

        Args:
            number_groups (int): The number of groups, that the population is split into.
            data (PandemicDataset): Data collected showing the course of the pandemic
            parameter_list (list): List of the parameter names(strings), that are supposed to be found.
            problem (PandemicProblem): Problem class implementing the calculation of the residuals.
            parameter_regulator (optional): Function to force the parameters to be in a certain range. Defaults to torch.tanh.
            input_size (int, optional): Number of the input nodes of the NN. Defaults to 1.
            hidden_size (int, optional): Number of the hidden nodes of the NN. Defaults to 20.
            hidden_layers (int, optional): Number of the hidden layers for the NN. Defaults to 7.
            activation_layer (optional): Class of the activation function. Defaults to torch.nn.ReLU().
        """
        
        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, 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.

        Args:
            parameter_name (str): Name of the parameter to be returned.

        Returns:
            torch.Parameter: Regulated parameter object of the search parameter.
        """
        return self.parameter_regulator(self.parameters_tilda[parameter_name])
    
    def get_parameters_tilda(self):
        """Function to get the original value (not forced into any range).

        Returns:
            torch.Parameter: Parameter object of the search parameter.
        """
        return list(self.parameters_tilda.values())

    def get_regulated_param_list(self):
        """Get the list of regulated parameters (forced into a specific range).

        Returns:
            list: list of regulated parameters
        """
        return [self.parameter_regulator(parameter) for parameter in self.get_parameters_tilda()]
    
    def train(self, 
              epochs: int, 
              lr: float, 
              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
        optimizer = optimizer_class(list(self.model.parameters()) + list(self.parameters_tilda.values()), lr=lr)
        scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=1e-5, max_lr=1e-3, step_size_up=1000, mode="exp_range", gamma=0.85, cycle_momentum=False)

        self.epochs = epochs

        # arrays to hold values for plotting
        self.losses = np.zeros(epochs)
        self.parameters = [np.zeros(epochs) for _ in self.parameters]

        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.get_regulated_param_list())

            optimizer.zero_grad()

            # calculate loss from the differential system
            loss_physics = 0
            for residual in residuals:
                loss_physics += torch.mean(torch.square(residual))

            # calculate loss from the dataset
            loss_obs = 0
            for i, group in enumerate(self.data.get_group_names()):
                loss_obs += torch.mean(torch.square(self.data.get_norm(group) - prediction[:, i]))
            
            loss = loss_physics + loss_obs

            loss.backward()
            optimizer.step()
            scheduler.step()

            # append values for plotting
            self.losses[epoch] = loss.item()
            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)
                print(f'physics loss:\t\t{loss_physics.item()}')
                print(f'observation loss:\t{loss_obs.item()}')
                print(f'loss:\t\t\t{loss.item()}')
                print('---------------------------------')
                for parameter in self.parameters_tilda.items():
                    print(f'{parameter[0]}:\t\t\t{self.parameter_regulator(parameter[1]).item()}')
                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.

        Args:
            ground_truth (list): List of the ground truth parameters
        """
        assert self.epochs != None
        epochs = np.arange(0, self.epochs, 1)

        # plot loss
        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):
            if len(ground_truth) > i:
                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)