Covid_19_DINN/src/dinn.py

import torch
import numpy as np

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, 
                     t_init,
                     t_final) -> 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)
            
            self.__t_init = t_init
            self.__t_final = t_final

        def forward(self, t):
            # normalize input
            t_scaled = (t - self.__t_init) / (self.__t_final - self.__t_init)
            x = self.input(t_scaled)
            x = self.hidden(x)
            x = self.output(x)
            return x

    def __init__(self,
                 output_size: int,
                 data: PandemicDataset,
                 parameter_list: list,
                 problem: PandemicProblem,
                 plotter: Plotter,
                 state_variables=[],
                 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:
            output_size (int): Number of the output nodes of the NN.
            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.
            plotter (Plotter): Plotter object to plot dataset curves.
            state_variables (list, optional): List of the names of state variables. Defaults to [].
            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().
        """
        assert len(state_variables) + data.number_groups == output_size, f'The number of groups plus the number of state variable must result in the output size\nGroups:\t{data.number_groups}\nState variables:\t{len(state_variables)}\noutput_size: {output_size}\n'
        self.device = torch.device(data.device_name)
        self.device_name = data.device_name
        self.plotter = plotter

        self.model = DINN.NN(output_size, 
                             input_size, 
                             hidden_size, 
                             hidden_layers, 
                             activation_layer, 
                             data.t_init, 
                             data.t_final)
        self.model = self.model.to(self.device)
        self.data = data
        self.parameter_regulator = parameter_regulator
        self.problem = problem
        self.problem.def_grad_matrix(output_size)

        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))})

        # new model has to be configured and then trained
        self.__is_configured = False
        self.__has_trained = False

        self.__state_variables = state_variables

        self.parameters = [np.zeros(1) for _ in range(len(parameter_list))]

        self.frames = []

    @property
    def number_state_variables(self):
        return len(self.__state_variables)

    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 configure_training(self, lr:float, epochs:int, optimizer_name='Adam', scheduler_name='CyclicLR', scheduler_factor = 1, verbose=False):
        """This method sets the optimizer, scheduler, learning rate and number of epochs for the following training process.

        Args:
            lr (float): Learning rate for the optimizer.
            epochs (int): Number of epochs the NN is supposed to be trained for.
            optimizer_name (str, optional): Name of the optimizer class that is supposed to be used. Defaults to 'Adam'.
            scheduler_name (str, optional): Name of the scheduler class that is supposed to be used. Defaults to 'CyclicLR'.
            verbose (bool, optional): Controles if the configuration process, is to be verbosed. Defaults to False.
        """
        parameter_list = list(self.model.parameters()) + list(self.parameters_tilda.values())
        self.epochs = epochs
        match optimizer_name:
            case 'Adam':
                self.optimizer = torch.optim.Adam(parameter_list, lr=lr)
            case _:
                self.optimizer = torch.optim.Adam(parameter_list, lr=lr)
                if verbose:
                    print('---------------------------------')
                    print(f' Entered unknown optimizer name: {optimizer_name}\n Defaulted to Adam.')
                    print('---------------------------------')
                optimizer_name = 'Adam'

        match scheduler_name:
            case 'CyclicLR':
                self.scheduler = torch.optim.lr_scheduler.CyclicLR(self.optimizer, base_lr=lr * 10, max_lr=lr * 1e3, step_size_up=1000, mode="exp_range", gamma=0.85, cycle_momentum=False)
            case 'LinearLR':
                self.scheduler = torch.optim.lr_scheduler.LinearLR(self.optimizer, start_factor=lr, total_iters=epochs/scheduler_factor)
            case 'PolynomialLR':
                self.scheduler = torch.optim.lr_scheduler.PolynomialLR(self.optimizer, total_iters=epochs/scheduler_factor, power=1.0)
            case _:
                self.scheduler = torch.optim.lr_scheduler.CyclicLR(self.optimizer, base_lr=lr * 10, max_lr=lr * 1e3, step_size_up=1000, mode="exp_range", gamma=0.85, cycle_momentum=False)
                if verbose:
                    print('---------------------------------')
                    print(f' Entered unknown scheduler name: {scheduler_name}\n Defaulted to CyclicLR.')
                    print('---------------------------------')
                scheduler_name = 'CyclicLR'

        if verbose:
            print(f'\nLearning Rate:\t{lr}\nOptimizer:\t{optimizer_name}\nScheduler:\t{scheduler_name}\n')

        self.__is_configured = True

    
    def train(self, 
              create_animation=False,
              animation_sample_rate=500,
              verbose=False):
        """Training routine for the DINN.

        Args:
            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.
            verbose (bool, optional): Controles if the training process, is to be verbosed. Defaults to False.
        """
        assert self.__is_configured, 'The model has to be configured before training through the use of self.configure training.'
        if verbose:
            print(f'torch seed: {torch.seed()}')
        
        # arrays to hold values for plotting
        self.losses = np.zeros(self.epochs)
        self.obs_losses = np.zeros(self.epochs)
        self.physics_losses = np.zeros(self.epochs)
        self.parameters = [np.zeros(self.epochs) for _ in self.parameters]

        for epoch in range(self.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())

            self.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.group_names):
                loss_obs += torch.mean(torch.square(self.data.get_norm(group) - prediction[:, i]))
            
            loss = loss_obs + loss_physics

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

            # append values for plotting
            self.losses[epoch] = loss.item()
            self.obs_losses[epoch] = loss_obs.item()
            self.physics_losses[epoch] = loss_physics.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.data.get_denormalized_data([prediction[:, i] for i in range(self.data.number_groups)])

                plot_labels = [name + '_pred' for name in self.data.group_names] + [name + '_true' for name in self.data.group_names]
                background_list = [0 for _ in self.data.group_names] + [1 for _ in self.data.group_names]
                self.plotter.plot(t, 
                                  list(groups) + list(self.data.data), 
                                  plot_labels,
                                  'frame',
                                  f'epoch {epoch}',
                                  figure_shape=(12, 6),
                                  is_frame=True,
                                  is_background=background_list,
                                  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 and verbose:          
                print(f'\nEpoch {epoch} | LR {self.scheduler.get_last_lr()[0]}')
                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('---------------------------------')
                if len(self.parameters_tilda.items()) != 0:
                    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')
            self.plotter.reset_animation()

        self.__has_trained = True

    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.__has_trained, 'Model has to be trained, before plotting the training graphs'
        epochs = np.arange(0, self.epochs, 1)

        # plot loss
        self.plotter.plot(epochs, [self.losses, self.obs_losses, self.physics_losses], ['loss', 'observation loss', 'physics loss'], self.data.name + '_loss', 'Loss', (6, 6), y_log_scale=True, plot_legend=True, 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)

    def plot_state_variables(self):
        for i in range(self.data.number_groups, self.data.number_groups+self.number_state_variables):
            prediction = self.model(self.data.t_batch)
            groups = [prediction[:, i] for i in range(self.data.number_groups)]
            t = torch.arange(0, self.data.t_raw[-1].item(), (self.data.t_raw[-1] / self.data.t_raw.shape[0]).item())
            self.plotter.plot(t,
                              [prediction[:, i]] + groups,
                              [self.__state_variables[i-self.data.number_groups]] + self.data.group_names,
                              f'{self.data.name}_{self.__state_variables[i-self.data.number_groups]}',
                              self.__state_variables[i-self.data.number_groups],
                              is_background=[0, 1, 1],
                              figure_shape=(12, 6),
                              plot_legend=True,
                              xlabel='time / days')