{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solve a Lotka-Volterra based ODE Optimal Control Problem\n",
    "\n",
    "## Notebook to demonstrate how to solve a Lotka-Volterra optimal control via oppy\n",
    "\n",
    "We consider the Lotka-Volterra optimal control problem \n",
    "\\begin{align*}\n",
    "&\\min_{y,u}J(y,u)\\\\\n",
    "&\\mathrm{s.t.}\\begin{cases}\\, &\\dot{y}_1(t)&=ay_1(t)-dy_1(t)y_2(t) +u(t)&\\quad \\mathrm{in}\\,t\\in(0,T)\\quad\\mathrm{with}\\quad y_1(0)=y^0_1>0,\\\\\n",
    "&\\dot{y}_2(t)&=-by_2(t)+cy_1(t)y_2(t)&\\quad \\mathrm{in}\\,t\\in(0,T)\\quad \\mathrm{with}\\quad y_2(0)=y^0_2>0,\\end{cases}\n",
    "\\end{align*}\n",
    "\n",
    "\n",
    "where $J(y,u)=-(y_1(T)+y_2(T))+\\frac{\\nu}{2}\\int_0^T|u(t)|^2dt,$ is the costfunctional, $u\\in U_{ad}:=\\{x\\in L^2(0,T;\\mathbb{R}) \\,: \\,0\\leq x(t)\\leq 1\\}$ is the control function, $y= [y_1,y_2]^T$ is the solution of the state equation, $a,b,c$ are model parameters and $\\nu$ a regularization parameter.\n",
    "\n",
    "For implementation purpose we use the following equivalent formulation of the state equation:\n",
    "\\begin{align*}\n",
    "\\dot{y}=Ay+F(y)+uB,\n",
    "\\end{align*}\n",
    "where $A=\\left(\\begin{array}{ c c}\n",
    "a&0\\\\\n",
    "0&-b\\end{array}\\right)$, $F=\\left(\\begin{array}{c}\n",
    "-dy_1y_2\\\\\n",
    "cy_1y_2\\end{array}\\right)$ and $B=\\left(\\begin{array}{ c }\n",
    "1\\\\\n",
    "0\\end{array}\\right)$\n",
    "\n",
    "To approach this problem we define the reduced cost functional:\n",
    "\\begin{align*}\n",
    "\\hat{J}(u)=J(\\mathcal{S}(u),u),\n",
    "\\end{align*}\n",
    "where $\\mathcal{S}$ is the solution operator which means $y=\\mathcal{S}(u)$ solves the differntial equation in the constraints with the corresponding control $u$. One can show that solving the reduced problem is equivalent to the original problem (see [1], Chapter 3.1). \n",
    "\n",
    "The idea is to use an constrained optimization method on the reduced problem. \n",
    "Therefore, following the Lagrange method (see [1], Chapter 4.1) we derive the gradient condition for computing the optimal solution $(\\hat{y},\\hat{u})$ of the linear quadratic optimal control problem:\n",
    "\\begin{align*}\n",
    "\\nabla \\hat{J}(\\hat{u})= \\nu \\hat{u}+B^T\\hat{p}=0,\n",
    "\\end{align*}\n",
    "where $\\hat{p}$ solves the adjoint equation:\n",
    "\\begin{align*}\n",
    "\\begin{array}{r l r}\n",
    "-\\dot{p}(t)&=A^Tp(t)+J_F(\\hat{y}(t))^Tp(t),&\\mathrm{for}\\,t\\in[0,T],\\\\\n",
    "p(T)&=-(y_1(T)+y_2(T))&\\end{array},\n",
    "\\end{align*}\n",
    "where $J_F(\\hat{y}(t))$ is the Jacobian matrix of $F$ evaluated at $\\hat{y}(t)$.\n",
    "\n",
    "Moreover for evaluating the reduced cost functional $\\hat{J}(u)$ we solve the differential equation in the constraints by integrating forward and for computing the reduced gradient (if required) $\\nabla\\hat{J}(u)$ we integrate the adjoint equation backwards. These algorithms are written in a seperate file called $\\small{\\textbf{functions.py}}$. \n",
    "\n",
    "The parameter data for the Lotka-Volterra problem as well as the option parameters for the linesearch and the options of the optimization methods are written in the $\\small{\\textbf{problem_data.py}}$ and thus have to be imported.\n",
    "\n",
    "[1] Gabriele Ciaramella, Script lecture: $\\scriptsize{\\textit{Optimal Control of Ordinary Differential Equations}}$, July 21, 2019.\n",
    "\n",
    "\n",
    "First let's import all dependencies needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from oppy.options import Options\n",
    "from oppy.conOpt import projected_gradient\n",
    "\n",
    "import functions_lotka_volterra as fun\n",
    "import problem_data_lotka_volterra as data\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define the initial control function, the lower and upper boundaries, the reduced cost functional and its derivatives for the optimization method. Furthermore the Armijo Goldstein and Projected Armijo linesearch are defined with optional options. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lower boundaries:\n",
    "a = np.array([0])\n",
    "\n",
    "# Upper boundaries:\n",
    "b = np.array([1])\n",
    "\n",
    "# Define cost function and derivatives\n",
    "def fhandle(x):\n",
    "    f = fun.FG(x,data.time,data,1)\n",
    "    return f\n",
    "\n",
    "def dfhandle(x):\n",
    "    f, grad, y, p = fun.FG(x,data.time,data.data,4)\n",
    "    return grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we set the optional options for the optimization methods. Here we define the options for the Gradient, L-BFGS and the Newton method, where we also use different norm formulas. For comparison purposes we define for each optimization methods a Armijo Goldstein and a projected Armijo linesearch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define norm for projected Gradient\n",
    "def norm(x,d):\n",
    "    nor = np.linalg.norm(x-np.clip(x-d, a, b))**2*data.diff[0]\n",
    "    return nor\n",
    "\n",
    "def normNewton(x,x_1):\n",
    "    return data.diff[0]*np.abs(np.sum((x-x_1)))**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we can run the optimization methods with an random initial control $u\\in\\mathbb{R}^{N}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Random initial guess\n",
    "u0 = np.random.rand(data.N)\n",
    "\n",
    "# Options for Projected Gradient\n",
    "options_proj_grad_armijo_gold = Options(norm = norm, disp=True)\n",
    "\n",
    "# Projected gradient scheme\n",
    "result_proj_grad = projected_gradient(fhandle, dfhandle, u0, a, b, options_proj_grad_armijo_gold)\n",
    "residual_proj_grad = result_proj_grad.res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The projected gradient terminates at 5 iterations. To reach a faster convergence we consider the L-BFGS-B method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from oppy.conOpt import l_bfgs_b\n",
    "\n",
    "# Set options for L-BFGS-B scheme\n",
    "options_l_bfgs_b = Options(norm = norm, disp=True)\n",
    "\n",
    "# L-BFGS-B scheme\n",
    "result_l_bfgs_b = l_bfgs_b(fhandle, dfhandle, u0, a, b, options_l_bfgs_b)\n",
    "residual_l_bfgs_b = result_l_bfgs_b.res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The L-BFGS-B method terminates at 6 iterations, which means that this method with default options is slightly slower than the projected gradient method. We try another method which is called the projected Newton method. Therefore we define the action of the hessian of the costfunction. Thereafter we define the options of the projected Newton method and execute the algorithm. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from oppy.conOpt import projected_newton\n",
    "\n",
    "def ddfhandle(x, u, Active, Inactive):\n",
    "    J, grad, y, p = fun.FG(u,data.time,data.data,4)\n",
    "    RAct = fun.Hessian(x, u, y, p, data.time, Active, Inactive, data.data, 1)\n",
    "    return RAct\n",
    "\n",
    "# Set options for Projected Newton scheme\n",
    "options_proj_newton = Options(norm = normNewton, disp=True)\n",
    "\n",
    "# Projected Newton scheme\n",
    "result_proj_newton = projected_newton(fhandle, dfhandle, ddfhandle, u0, a, b, options_proj_newton)\n",
    "residual_proj_newton = result_proj_newton.res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This method needs 5 iterations steps to terminate like the projected gradient method. An error plot with respect to iteration numbers of the different optimization methods with a semi y-axis log scale is used to compare the different schemes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot convergence rate\n",
    "plt.semilogy(residual_proj_grad, label='Projected Gradient')\n",
    "plt.semilogy(residual_l_bfgs_b, label='L_BFGS_B')\n",
    "plt.semilogy(residual_proj_newton, label='Projected Newton' )\n",
    "plt.title('Norm of the gradient, different schemes')\n",
    "plt.grid(True)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All methods converge within 6 iterations. In this Lotka Volterra optimal problem the projected Newton method with default options provides the best solution, since it terminates at 5 iterations with the lowest residual value."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For better understanding of the problem, we plot the states $y$ without a control function $u$ with respect to the time $t$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot uncontrolled state\n",
    "u0 = np.zeros((data.N,))\n",
    "f_0 , grad_0 , y_0 , p_0 = fun.FG( u0 , data.time, data.data , 4 )\n",
    "\n",
    "xAxe = np.linspace(0,data.T, num=data.N)\n",
    "plt.plot(xAxe,y_0[0,:], label='$y_1(t)$')\n",
    "plt.plot(xAxe,y_0[1,:], label='$y_2(t)$')\n",
    "plt.title('Solution of the state with no control')\n",
    "plt.grid(True)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now extract the optimal solution of the results of the quasi Newton method and solve the differential equation, which is the constraint of the optimal control problem. Therefore we use the forward solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot optimal solution\n",
    "f1 , grad , y , p = fun.FG( result_proj_newton.x , data.time, data.data , 4 )\n",
    "xAxe = np.linspace(0,data.T, num=data.N)\n",
    "plt.plot(xAxe,y[0,:], label='$y_1$')\n",
    "plt.plot(xAxe,y[1,:], label='$y_2$')\n",
    "plt.plot(xAxe, result_proj_newton.x, label='$u(t)$')\n",
    "plt.xlabel('time')\n",
    "plt.ylabel('prey $y_1(t)$ and predators $y_2(t)$')\n",
    "plt.title('Optimal Solution')\n",
    "plt.grid(True)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The state outcome with optimal control aligns with the optimal control problem since the costfunction aims to maximize the states in each variable."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}