Neural Parameter Estimation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bf9e3c06-83b8-475b-a47c-2800bb21b148",
   "metadata": {},
   "source": [
    "# Neural Parameter Estimation for SIR models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e87468-5d75-47bd-94ec-8feb58966910",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('---------------------------------')\n",
    "import timeit\n",
    "import torch\n",
    "use_cuda = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "print(\"Device: \",device)\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import random as rd\n",
    "import torch.nn.functional as F\n",
    "from scipy.stats import bernoulli\n",
    "import numpy as np\n",
    "from scipy.integrate import odeint\n",
    "import matplotlib.pyplot as plt\n",
    "import copy\n",
    "\n",
    "\n",
    "plt.rcParams['figure.figsize'] = [16, 10]\n",
    "\n",
    "\n",
    "# Define the class for the neural network\n",
    "class NN(nn.Module):\n",
    "    def __init__(self,\n",
    "                 num_inputs = 1,\n",
    "                 num_outputs = 1,\n",
    "                 num_hidden = 3,\n",
    "                 hidden_neurons = 40,\n",
    "                 activation = nn.Tanh):\n",
    "        super().__init__()\n",
    "        # self.num_inputs = num_inputs\n",
    "        # self.num_outputs = num_outputs\n",
    "        # self.num_hidden = num_hidden\n",
    "        # self.activation = activation\n",
    "\n",
    "        # Input layer\n",
    "        self.fci = nn.Sequential(*[\n",
    "                        nn.Linear(num_inputs, hidden_neurons),\n",
    "                        activation()])\n",
    "\n",
    "        # Hidden layers\n",
    "        self.fch = nn.Sequential(*[\n",
    "                        nn.Sequential(*[\n",
    "                            nn.Linear(hidden_neurons, hidden_neurons),\n",
    "                            activation()]) for _ in range(num_hidden)])\n",
    "\n",
    "        # Output layer\n",
    "        self.fco = nn.Sequential(*[nn.Linear(hidden_neurons, num_outputs), activation()])\n",
    "        # self.fco = nn.Linear(hidden_neurons, num_outputs)\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = self.fci(x)\n",
    "        x = self.fch(x)\n",
    "        x = self.fco(x)\n",
    "        return x\n",
    "\n",
    "# Dynamics informed Neural Network\n",
    "class DNet():\n",
    "    def __init__(self, t_data, data_list, minlist, maxlist, t_physics, param, nndata):\n",
    "\n",
    "        # Time for data points and residual of physics\n",
    "        self.t_data = torch.tensor(t_data, requires_grad=True).view(-1,1).float().to(device=device)\n",
    "        self.t_physics = torch.tensor(t_physics, requires_grad=True).view(-1,1).float().to(device=device)\n",
    "\n",
    "        # Data points\n",
    "        self.u_data = []\n",
    "        for elem in data_list:\n",
    "            self.u_data.append(torch.tensor(elem, requires_grad=True).view(-1,1).float().to(device=device))\n",
    "\n",
    "        # minmax scaling settings\n",
    "        self.minlist = minlist\n",
    "        self.maxlist = maxlist\n",
    "\n",
    "        self.beta = torch.tensor([param[0]], requires_grad=True).to(device=device)\n",
    "        self.gamma = torch.tensor([param[1]], requires_grad=True).to(device=device)\n",
    "        # self.gamma = torch.nn.Parameter(self.gamma)\n",
    "        # Randomize beta\n",
    "        # Initialize beta with a random value (you can also use custom initialization)\n",
    "        self.beta = torch.tensor([rd.uniform(0.05,0.95)], requires_grad=True).to(device=device)\n",
    "        print('Initial random value of beta: %.3f' %self.beta)\n",
    "        # Load the neural network(s)\n",
    "        self.models = []\n",
    "        self.optimizers = []\n",
    "        for dim in range(len(self.u_data)):\n",
    "                network = NN(\n",
    "                    num_inputs = nndata['num_inputs'],\n",
    "                    num_outputs = nndata['num_outputs'],\n",
    "                    num_hidden =  nndata['num_hidden'],\n",
    "                    hidden_neurons = nndata['hidden_neurons'],\n",
    "                    activation = nndata['activation'],\n",
    "                ).to(device=device)\n",
    "                self.models.append(network)\n",
    "\n",
    "\n",
    "        # Register parameters for Neural network\n",
    "        self.beta = torch.nn.Parameter(self.beta)\n",
    "#         for network in self.models:\n",
    "#             network.register_parameter('b', self.beta)\n",
    "\n",
    "        # Define optimizer (ADAM)\n",
    "        # self.optimizer = torch.optim.Adam([{'params': network.parameters()} for network in self.models], lr=1e-7)\n",
    "        self.optimizers = []\n",
    "        for network in self.models:\n",
    "            self.optimizers.append(torch.optim.Adam([\n",
    "                {'params': network.parameters()},  # For other parameters (e.g., weights and biases)\n",
    "                {'params': self.beta, 'lr': 1e-5}  # For the \"beta\" parameter\n",
    "            ], lr=1e-7))\n",
    "            # self.optimizers.append(torch.optim.Adam(network.parameters(), lr=1e-5))\n",
    "\n",
    "    #----------------------------------------------------------------------------\n",
    "    # DYNAMICS\n",
    "    #---------------------------------------------------------------------------\n",
    "    def SIR_dynamics(self,t):\n",
    "        # Apply model\n",
    "        u_list = []\n",
    "        for network in self.models:\n",
    "            u_list.append(network(t))\n",
    "\n",
    "        # Inverse Scalings (Inv. minmax-> Inv. Log -> Normal data)\n",
    "        inv_results = inverse_minmax(u_list,self.minlist,self.maxlist)\n",
    "        u = []\n",
    "        for tensor in inv_results:\n",
    "            u.append(torch.exp(tensor))\n",
    "\n",
    "        # Compartments and derivatives\n",
    "        du = []\n",
    "        for tensor in u:\n",
    "            du.append(torch.autograd.grad(tensor, t, torch.ones_like(tensor), create_graph=True)[0])\n",
    "\n",
    "        # Parameters\n",
    "        beta = self.beta\n",
    "        gamma = self.gamma\n",
    "\n",
    "        # SIR Residuals\n",
    "        f0 = du[0] + beta*u[0]*u[1]\n",
    "        f1 = du[1] - beta*u[0]*u[1] + gamma*u[1]\n",
    "        return f0**2+f1**2\n",
    "    #----------------------------------------------------------------------------\n",
    "\n",
    "    #----------------------------------------------------------------------------\n",
    "    # DATA LOSS\n",
    "    #----------------------------------------------------------------------------\n",
    "    def data_loss(self,t):\n",
    "        # Apply model\n",
    "        sc_u = []\n",
    "        for network in self.models:\n",
    "            sc_u.append(network(t))\n",
    "\n",
    "        # Compartments\n",
    "        f = []\n",
    "        for i in range(len(sc_u)):\n",
    "            f.append(sc_u[i]-self.u_data[i])\n",
    "        # Data loss\n",
    "        f = f[0]**2 + f[1]**2\n",
    "        return f\n",
    "\n",
    "    def batched_data_loss(self,t,u):\n",
    "        sc_u = []\n",
    "        for network in self.models:\n",
    "            sc_u.append(network(t))\n",
    "\n",
    "        # Dloss compartments\n",
    "        f = []\n",
    "        for i in range(len(sc_u)):\n",
    "            f.append(sc_u[i]-u[i])\n",
    "\n",
    "        # Data loss\n",
    "        f = f[0]**2 + f[1]**2\n",
    "        return f\n",
    "    #----------------------------------------------------------------------------\n",
    "\n",
    "    #----------------------------------------------------------------------------\n",
    "    # PLOT TEST FUNCTION\n",
    "    #----------------------------------------------------------------------------\n",
    "    def plot_test(self):\n",
    "        tarray = torch.linspace(0,t_final,600).view(-1,1)\n",
    "\n",
    "        # Apply model\n",
    "        u_list = []\n",
    "        for network in self.models:\n",
    "            u_list.append(network(tarray))\n",
    "\n",
    "        # Inverse Scalings (Inv. minmax-> Inv. Log -> Normal data)\n",
    "        inv_results = inverse_minmax(u_list,self.minlist,self.maxlist)\n",
    "        u = []\n",
    "        for tensor in inv_results:\n",
    "            u.append(torch.exp(tensor))\n",
    "        return\n",
    "\n",
    "    #----------------------------------------------------------------------------\n",
    "    # PLOT LOSS HISTORY\n",
    "    #----------------------------------------------------------------------------\n",
    "    def loss_history(self, lu_list, lf_list):\n",
    "        plt.plot(range(len(lu_list)),np.log(lu_list))\n",
    "        plt.plot(range(len(lf_list)),np.log(lf_list), c='orange')\n",
    "        plt.plot(range(len(lu_list)), np.log(np.array(lu_list)+np.array(lf_list)), c='red', ls='dashed', alpha=0.66)\n",
    "        plt.title('Training loss functions')\n",
    "        plt.xlabel('Epochs')\n",
    "        plt.ylabel('log loss')\n",
    "        plt.legend(['data loss', 'physics loss', 'total loss'])\n",
    "        plt.show()\n",
    "        return\n",
    "    #----------------------------------------------------------------------------\n",
    "    # TRAINING\n",
    "    #----------------------------------------------------------------------------\n",
    "    def only_data_training(self,nIter,dataloader,total_loss):\n",
    "        lu_list = []\n",
    "        lf_list = []\n",
    "        for epoch in range(nIter):\n",
    "            loss1 = 0\n",
    "            for inputs, targets1, targets2 in dataloader:\n",
    "\n",
    "                for optimizer in self.optimizers:\n",
    "                    optimizer.zero_grad()\n",
    "\n",
    "                # Data\n",
    "                inputs = torch.stack(inputs).to(device)\n",
    "                targets1 = torch.stack(targets1).to(device)\n",
    "                targets2 = torch.stack(targets2).to(device)\n",
    "\n",
    "                # Data Loss\n",
    "                loss_u = torch.mean(self.batched_data_loss(inputs, [targets1,targets2]))\n",
    "                loss1 = loss1 + loss_u.item()\n",
    "\n",
    "                # Total Loss\n",
    "                loss = loss_u\n",
    "                loss.backward()\n",
    "\n",
    "                for optimizer in self.optimizers:\n",
    "                    optimizer.step()\n",
    "\n",
    "            lu_list.append(loss1/len(dataloader))\n",
    "            lf_list.append(0)\n",
    "            total_loss.append(lu_list[-1])\n",
    "        return lu_list, lf_list, total_loss\n",
    "\n",
    "\n",
    "    def train(self, nIter):\n",
    "        # self.model.train()\n",
    "        start_time = timeit.default_timer()\n",
    "\n",
    "        # Create a combined training dataset\n",
    "        dataset = TensorDataset(self.t_data, self.u_data[0], self.u_data[1])\n",
    "\n",
    "        # Loss History\n",
    "        lu_list = []\n",
    "        lf_list = []\n",
    "        total_loss = []\n",
    "\n",
    "        # Stoping conditions parameters\n",
    "        kslice = 10\n",
    "        t = nIter\n",
    "\n",
    "        # Prepare alphalist\n",
    "        alist = []\n",
    "        for i in range(13):\n",
    "            alist.append(10**(-i/2))\n",
    "        # for i in range(4,8):\n",
    "        #     alist.append(10**(-i/2))\n",
    "        #     alist.append(10**(-i/2))\n",
    "\n",
    "        # Create a DataLoader with the custom collate function\n",
    "        batch_size = 1\n",
    "        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=custom_collate)\n",
    "\n",
    "        # Only data Training\n",
    "        lu_list, lf_list, total_loss = self.only_data_training(nIter,dataloader,total_loss)\n",
    "\n",
    "        # Loss Balancing Initialization\n",
    "        loss_u = torch.mean(self.data_loss(self.t_data))\n",
    "        loss_f = torch.mean(self.SIR_dynamics(self.t_physics))\n",
    "        w = loss_u.item()/loss_f.item()\n",
    "\n",
    "        # Begin Training\n",
    "        for alpha in alist:\n",
    "            for epoch in range(nIter):\n",
    "                t = t + 1\n",
    "                loss1 = 0\n",
    "                loss2 = 0\n",
    "                for inputs, targets1, targets2 in dataloader:\n",
    "\n",
    "                    for optimizer in self.optimizers:\n",
    "                        optimizer.zero_grad()\n",
    "\n",
    "                    # Data\n",
    "                    inputs = torch.stack(inputs).to(device)\n",
    "                    targets1 = torch.stack(targets1).to(device)\n",
    "                    targets2 = torch.stack(targets2).to(device)\n",
    "\n",
    "                    # Data Loss\n",
    "                    loss_u = torch.mean(self.batched_data_loss(inputs, [targets1,targets2]))\n",
    "                    loss1 = loss1 + loss_u.item()\n",
    "\n",
    "                    # Physics Loss\n",
    "                    loss_f = w*alpha*torch.mean(self.SIR_dynamics(inputs))\n",
    "                    loss2 = loss2 + loss_f.item()\n",
    "\n",
    "                    # Total Loss\n",
    "                    loss = loss_u + loss_f\n",
    "                    loss.backward()\n",
    "\n",
    "                    for optimizer in self.optimizers:\n",
    "                        optimizer.step()\n",
    "\n",
    "                # Losses\n",
    "                lu_list.append(loss1/len(dataloader))\n",
    "                lf_list.append(loss2/len(dataloader))\n",
    "                total_loss.append(loss1/len(dataloader)+loss2/len(dataloader))\n",
    "\n",
    "                # Stopping Condition: PK - https://link.springer.com/content/pdf/10.1007/3-540-49430-8_3.pdf\n",
    "                if (t) % kslice == 0:\n",
    "                    gl = 100*(total_loss[-1]/min(total_loss) - 1)\n",
    "                    pk = 1000*(sum(total_loss[-kslice:])/(kslice*min(total_loss)) - 1)\n",
    "                    # print('GL: %2.f, PK: %2.f' %\n",
    "                    #       (gl, pk))\n",
    "                    if gl/pk>0.8:\n",
    "                        print('-------Stopping condition for alpha section-------')\n",
    "                        print('Last epoch: %d, Loss: %.3e, Time: %.1f, Stopping Ratio: %2.f' %\n",
    "                          (t, loss.item(), elapsed, gl/pk))\n",
    "                        print('---------------------RESUMING---------------------')\n",
    "                        break\n",
    "                if (t) % nIter == 0:\n",
    "                    elapsed = timeit.default_timer() - start_time\n",
    "                    print('\\rEpoch: %d, Loss: %.3e, Time: %.1f, gl: %.2f, pk: %.2f, ratio:%.2f || beta:%.3f, gamma:%.2f' %\n",
    "                          (t, total_loss[-1], elapsed, gl, pk, gl/pk, self.beta, self.gamma))\n",
    "\n",
    "        print('Total Number of Epochs: %d, Final Value of beta: %.3f' %(t, self.beta) )\n",
    "        self.loss_history(lu_list, lf_list)\n",
    "        # self.plot_test()\n",
    "    #----------------------------------------------------------------------------\n",
    "\n",
    "    #----------------------------------------------------------------------------\n",
    "    # PREDICTION\n",
    "    #----------------------------------------------------------------------------\n",
    "    def predict(self, t):\n",
    "        with torch.no_grad():\n",
    "            # Apply model\n",
    "            out_data = []\n",
    "            for network in self.models:\n",
    "                network.eval()\n",
    "                out_data.append(network(t))\n",
    "        return out_data\n",
    "\n",
    "#--------------------------------------------------------\n",
    "# Additional Activation functions\n",
    "#--------------------------------------------------------\n",
    "class ReLU1(nn.Module):\n",
    "    def __init__(self, inplace=False):\n",
    "        super(ReLU1, self).__init__()\n",
    "        self.inplace = inplace\n",
    "\n",
    "    def forward(self, x):\n",
    "        f = torch.max(x, torch.zeros_like(x), out=x if self.inplace else None)\n",
    "        f = torch.clamp(f, max=1.0)\n",
    "        return f\n",
    "\n",
    "class ClampTanh(nn.Module):\n",
    "    def __init__(self, inplace=False):\n",
    "        super(ClampTanh, self).__init__()\n",
    "        self.inplace = inplace\n",
    "\n",
    "    def forward(self, x):\n",
    "        f = nn.Tanh(x)\n",
    "        f = torch.clamp(f, min=0.0, max=1.0)\n",
    "        return f\n",
    "#--------------------------------------------------------\n",
    "# Define a custom collate function to split the targets\n",
    "#--------------------------------------------------------\n",
    "def custom_collate(batch):\n",
    "    inputs = [item[0] for item in batch]\n",
    "    targets1 = [item[1] for item in batch]\n",
    "    targets2 = [item[2] for item in batch]\n",
    "    return inputs, targets1, targets2\n",
    "#--------------------------------------------------------\n",
    "\n",
    "#--------------------------------------------------------\n",
    "# Define the SIR differential equations\n",
    "#--------------------------------------------------------\n",
    "def SIR_eqs(z,t,beta,k,n):\n",
    "    dsdt = -beta*z[0]*z[1]/n\n",
    "    didt = beta*z[0]*z[1]/n - k*z[1]\n",
    "    drdt = k*z[1]\n",
    "    dzdt = [dsdt,didt,drdt]\n",
    "    return dzdt\n",
    "#--------------------------------------------------------\n",
    "\n",
    "#--------------------------------------------------------\n",
    "# Min Max Scaling and inverse Scaling\n",
    "#--------------------------------------------------------\n",
    "def minmax_scaling(data):\n",
    "    # Takes a numpy array with several columns and performs minmax for each column separately\n",
    "    # The min and max listes are here to do the inverse min max afterwards\n",
    "    minlist = []\n",
    "    maxlist = []\n",
    "    scaled_data = []\n",
    "    for column in data.T:\n",
    "        minvar = (min(column))\n",
    "        maxvar = (max(column))\n",
    "        sc_var = (column - minvar)/(maxvar-minvar)\n",
    "        minlist.append(minvar)\n",
    "        maxlist.append(maxvar)\n",
    "        scaled_data.append(sc_var)\n",
    "    return scaled_data, minlist, maxlist\n",
    "\n",
    "def inverse_minmax(sc_data, minlist, maxlist):\n",
    "    # Performs the inverse min max scaling to return to normal data\n",
    "    data = []\n",
    "    for i in range(len(sc_data)):\n",
    "        var = sc_data[i]*(maxlist[i]-minlist[i])+minlist[i]\n",
    "        data.append(var)\n",
    "    return data\n",
    "#--------------------------------------------------------\n",
    "\n",
    "# Initialization\n",
    "torch.manual_seed(42)\n",
    "gamma = 1/10\n",
    "n_pop = 10**5\n",
    "beta = round(rd.uniform(0.2, 0.95),2)\n",
    "params = [beta, gamma, n_pop]\n",
    "i0 = 1\n",
    "\n",
    "# System Solver\n",
    "t_final = 40\n",
    "Nt = t_final\n",
    "t_data = np.linspace(0, t_final, Nt+1)\n",
    "Z0 = [n_pop-i0,i0,0] # Initial condition\n",
    "sol = odeint(SIR_eqs,Z0,t_data, tuple(params))/n_pop\n",
    "\n",
    "# Physics Settings\n",
    "noise = 0.0\n",
    "t_physics = np.linspace(0,t_final,Nt*2)\n",
    "\n",
    "# Numerical solution with scalings\n",
    "u_log = np.log(sol[:,:2])\n",
    "u_scaled, minlist, maxlist = minmax_scaling(u_log)\n",
    "\n",
    "# Collocation points\n",
    "u_train = []\n",
    "for elem in u_scaled:\n",
    "    u_train.append(elem[0:32:1])\n",
    "t_train = t_data[0:32:1]\n",
    "\n",
    "# Dictionary with architecture\n",
    "nndict = {}\n",
    "nndict['num_inputs'] = 1\n",
    "nndict['num_outputs'] = 1\n",
    "nndict['num_hidden'] = 10\n",
    "nndict['hidden_neurons'] = 128\n",
    "nndict['activation'] = nn.SELU\n",
    "\n",
    "# randbeta = round(rd.uniform(0.10, 0.99),2)\n",
    "# params = [randbeta, gamma, n_pop]\n",
    "# Model initialization\n",
    "model = DNet(t_train, u_train, minlist, maxlist, t_physics, params, nndict)\n",
    "print('---------------------------------')\n",
    "print('Data Settings')\n",
    "print('Data dimensions: ', len(t_train))\n",
    "print('Physics dimensions: ', len(t_physics))\n",
    "print('Real beta of Numerical Solution: %.2f' %beta)\n",
    "print('---------------------------------')\n",
    "print('Starting training of neural networks...')\n",
    "model.train(600)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}