{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Assignment 4\n",
    "\n",
    "\n",
    "- **Zhilong Li**\n",
    "- **USC ID:** 3884-8441-02\n",
    "- **USC email:** zhilongl@usc.edu"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "from tensorflow import keras\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from IPython import display\n",
    "import pickle\n",
    "plt.style.use('ggplot')\n",
    "plt.rcParams[\"figure.figsize\"] = (12,6)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Init Plugin\n",
      "Init Graph Optimizer\n",
      "Init Kernel\n"
     ]
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T07:00:13.073742Z",
     "start_time": "2021-10-02T07:00:12.086679Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "def get_model(input_dim, output_dim, width, depth, regularization_param,\n",
    "              act_func='tanh', output_act_func=False, random_seed=None):\n",
    "    initializer = keras.initializers.RandomNormal(seed=random_seed)\n",
    "    regularizer = keras.regularizers.l2(regularization_param)\n",
    "\n",
    "    if act_func == 'sin':\n",
    "        act_func = tf.math.sin\n",
    "\n",
    "    all_layers = [keras.layers.InputLayer(input_shape=[input_dim])]\n",
    "    all_layers += [keras.layers.Dense(width, activation=act_func,\n",
    "                                      kernel_initializer=initializer,\n",
    "                                      bias_initializer=initializer,\n",
    "                                      kernel_regularizer=regularizer,\n",
    "                                      bias_regularizer=regularizer) for _ in range(depth-1)]\n",
    "    all_layers.append(keras.layers.Dense(\n",
    "        output_dim, kernel_initializer=initializer,\n",
    "        activation=act_func if output_act_func else None,\n",
    "        bias_initializer=initializer,\n",
    "        kernel_regularizer=regularizer,\n",
    "        bias_regularizer=regularizer))\n",
    "    m = keras.Sequential(all_layers)\n",
    "    return m"
   ],
   "outputs": [],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T07:00:20.353060Z",
     "start_time": "2021-10-02T07:00:20.342013Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "def fit(model, x_train, Pe, Da, lambda_b, epoch, verbose=True):\n",
    "    ans = np.load(f'ref/diffusion-advection-reaction_pe_{Pe}_da_{Da}.npy')\n",
    "    optimizer = keras.optimizers.Adam(learning_rate=1e-3)\n",
    "    with tf.device('/cpu:0'):\n",
    "        interior_losses = []\n",
    "        boundary_losses = []\n",
    "        total_losses = []\n",
    "        for ep in range(1, epoch+1):\n",
    "            with tf.GradientTape() as tape_0:\n",
    "                with tf.GradientTape() as tape_1:\n",
    "                    tape_1.watch(x_train)\n",
    "                    with tf.GradientTape() as tape_2:\n",
    "                        tape_2.watch(x_train)\n",
    "                        u = model(x_train)\n",
    "                    du_dx = tape_2.gradient(u, x_train)\n",
    "                d2u_dx2 = tape_1.gradient(du_dx, x_train)\n",
    "                interior_loss = sum(map(lambda d2, d, u: (d2-Pe*d+Da*u*(1-u))**2,\n",
    "                                   d2u_dx2, du_dx, u))/len(x_train)\n",
    "                boundary_loss = lambda_b*(u[0]**2+(u[-1]-1)**2)\n",
    "                loss = interior_loss + boundary_loss + sum(model.losses)\n",
    "                interior_losses.append(interior_loss)\n",
    "                boundary_losses.append(boundary_loss)\n",
    "                total_losses.append(loss)\n",
    "            grads = tape_0.gradient(loss, model.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
    "\n",
    "            \n",
    "            if verbose == True and not ep % 100 or ep == epoch+1:\n",
    "                plt.clf()\n",
    "                plt.plot(ans[0], ans[1], 'r', label='Ref Truth')\n",
    "                pred = model.predict(x_train)\n",
    "                plt.plot(x_train, pred, 'b--', linewidth=2,\n",
    "                         label=\"Fitting Result\")\n",
    "                plt.title(\n",
    "                    f\"Loss = {loss.numpy()[0]:.5e},   Epoch = {ep}\\nPe={Pe},   Da={Da}\")\n",
    "                plt.ylim(-0.2, 1.2)\n",
    "                plt.legend()\n",
    "                display.clear_output(wait=True)\n",
    "                plt.pause(0.001)\n",
    "    return interior_losses, boundary_losses, total_losses, u\n"
   ],
   "outputs": [],
   "metadata": {
    "ExecuteTime": {
     "start_time": "2021-10-02T08:07:24.274Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "n_train = 100\n",
    "# Pe = 0.1\n",
    "# Da = 0.1\n",
    "lambda_b = 10\n",
    "epoch = 8000\n",
    "verbose = False\n",
    "x_train = tf.reshape(tf.linspace(0.0, 1.0, n_train), [100, 1])\n",
    "Pe_list = [0.1, 10]\n",
    "Da_list = [0.1, 10]\n",
    "\n",
    "for Pe in Pe_list:\n",
    "    for Da in Da_list:\n",
    "        for run in range(1, 4):\n",
    "            tf.keras.backend.clear_session()\n",
    "            ans = np.load(\n",
    "                f'ref/diffusion-advection-reaction_pe_{Pe}_da_{Da}.npy')\n",
    "            model = get_model(1, 1, 20, 5, 1e-7, act_func='sin',\n",
    "                              output_act_func=False)\n",
    "            interior_losses, boundary_losses, total_losses, u = fit(\n",
    "                model, x_train, Pe, Da, lambda_b, epoch, verbose)\n",
    "            # Plot\n",
    "            fig, (sp1, sp2) = plt.subplots(1, 2)\n",
    "            sp1.plot(interior_losses, label='int')\n",
    "            sp1.plot(boundary_losses, label='boundary')\n",
    "            sp1.plot(total_losses, label='total')\n",
    "            sp1.set_yscale('log')\n",
    "            sp1.legend()\n",
    "            sp2.plot(ans[0], ans[1], 'r', label='Ref Truth', linewidth=4)\n",
    "            sp2.plot(x_train, u, 'b--', label='Pred', linewidth=2)\n",
    "            sp2.legend()\n",
    "            fig.suptitle(\n",
    "                f\"Loss = {total_losses[-1][0]:.5f}, Epoch = {epoch}\\nPe={Pe},   Da={Da},   lambda_b={lambda_b}\\nRUN {run}\")\n",
    "            plt.show()\n"
   ],
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 50076 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 51268 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 53036 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 54948 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 57684 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 56664 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 40760 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 42384 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 42304 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 52504 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 52748 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 52324 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can see when using $\\lambda_b = 10$, the "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "lambda_b = 1000\n",
    "epoch = 8000\n",
    "verbose = False\n",
    "Pe_list = [10]\n",
    "Da_list = [0.1, 10]\n",
    "for Pe in Pe_list:\n",
    "    for Da in Da_list:\n",
    "        for run in range(1, 4):\n",
    "            tf.keras.backend.clear_session()\n",
    "            ans = np.load(\n",
    "                f'ref/diffusion-advection-reaction_pe_{Pe}_da_{Da}.npy')\n",
    "            model = get_model(1, 1, 20, 5, 1e-7, act_func='sin',\n",
    "                              output_act_func=False)\n",
    "            interior_losses, boundary_losses, total_losses, u = fit(\n",
    "                model, x_train, Pe, Da, lambda_b, epoch, verbose)\n",
    "            # Plot\n",
    "            fig, (sp1, sp2) = plt.subplots(1, 2)\n",
    "            sp1.plot(interior_losses, label='int')\n",
    "            sp1.plot(boundary_losses, label='boundary')\n",
    "            sp1.plot(total_losses, label='total')\n",
    "            sp1.set_yscale('log')\n",
    "            sp1.legend()\n",
    "            sp2.plot(ans[0], ans[1], 'r', label='Ref Truth', linewidth=4)\n",
    "            sp2.plot(x_train, u, 'b--', label='Pred', linewidth=2)\n",
    "            sp2.legend()\n",
    "            fig.suptitle(\n",
    "                f\"Loss = {total_losses[-1][0]:.5f}, Epoch = {epoch}\\nPe={Pe},   Da={Da},   lambda_b={lambda_b}\\nRUN {run}\")\n",
    "            plt.show()\n"
   ],
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 62524 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 58556 characters]",
      "text/plain": [
       "<Figure size 864x432 with 2 Axes>"
      ]
     },
     "metadata": {}
    }
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.12 64-bit ('tf': conda)"
  },
  "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.8.12"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  },
  "interpreter": {
   "hash": "859cadc33cf9c93d9354f622d0bfdc75f7d75fc23a9b22a18dfc646072acf83e"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}