{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c4ae06f2",
   "metadata": {},
   "source": [
    "This notebook illustrates the conjugate gradient method and the extension of Newton's method on systems of two equations and two variables."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "655b6425",
   "metadata": {},
   "source": [
    "# 1. The Conjugate Gradient Method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f3474223",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Printf\n",
    "using LinearAlgebra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0aa8e3bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CGM"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "    CGM(A, b, x0, maxit, tol)\n",
    "\n",
    "Applies the Conjugate Gradient Method to solve\n",
    "the linear system Ax = b starting at x0.\n",
    "\"\"\"\n",
    "function CGM(A::Array{Float64,2},b::Array{Float64,1},\n",
    "             x0::Array{Float64,1},\n",
    "             maxit::Int64=10,tol::Float64=1.0e-8,\n",
    "             verbose=true)\n",
    "    sol = deepcopy(x0)\n",
    "    r = b - A*sol\n",
    "    p = deepcopy(r)\n",
    "    if verbose\n",
    "        println(\" norm(r)    alpha       beta\")\n",
    "    end\n",
    "    for i=1:maxit\n",
    "        res = norm(r)\n",
    "        if verbose\n",
    "            sres = @sprintf(\"%.2e\", res)\n",
    "            print(\"$sres\")\n",
    "        end \n",
    "        if res < tol\n",
    "            if verbose\n",
    "                println(\"  succeeded after \", i, \" steps\")\n",
    "            end\n",
    "            return (sol, res, i, false)\n",
    "        end\n",
    "        alpha = (transpose(r)*r)/(transpose(p)*A*p)\n",
    "        if verbose\n",
    "            salpha = @sprintf(\"%.4e\", alpha)\n",
    "            print(\"  $salpha\")\n",
    "        end\n",
    "        sol = sol + alpha*p\n",
    "        r1 = r - alpha*A*p\n",
    "        beta = (transpose(r1)*r1)/(transpose(r)*r)\n",
    "        if verbose\n",
    "            sbeta = @sprintf(\"%.4e\", beta)\n",
    "            println(\"  $sbeta\")\n",
    "        end\n",
    "        p = r1 + beta*p\n",
    "        r = r1\n",
    "    end\n",
    "    return (sol, norm(r), maxit, true)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "643c709e",
   "metadata": {},
   "source": [
    "Let us run this on a simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "71404af4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " norm(r)    alpha       beta\n",
      "6.71e+00  2.3810e-01  3.2653e-01\n",
      "3.83e+00  7.0000e-01  4.2949e-31\n",
      "2.51e-15  succeeded after 3 steps\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "([3.9999999999999987, -0.9999999999999998], 2.5121479338940403e-15, 3, false)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mat = [2.0 2.0; 2.0 5.0]\n",
    "rhs = [6.0; 3.0]\n",
    "sol = [0.0; 0.0]\n",
    "res = CGM(mat, rhs, sol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca64e40f",
   "metadata": {},
   "source": [
    "# 2. Newton's Method in Two Variables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "887337a1",
   "metadata": {},
   "source": [
    "Consider a system of two equations in two variables, ``x`` and ``y``.  To compute the derivatives, we use SymPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7a87cd0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: using SymPy.rhs in module Main conflicts with an existing identifier.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(x, y)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using SymPy\n",
    "x, y = Sym(\"x, y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fed2f6e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SymPyExpressions"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Returns evaluable objects of the expressions in x and y,\n",
    "given in the strings f and g.\n",
    "\"\"\"\n",
    "function SymPyExpressions(f::String, g::String)\n",
    "    evaluatedf = sympify(f)\n",
    "    evaluatedg = sympify(g)\n",
    "    return [lambdify(evaluatedf, (x, y));\n",
    "            lambdify(evaluatedg, (x, y))]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8d8c9358",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SymPyDerivatives"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Returns all evaluable partial derivatives of the \n",
    "expression in x and y, given in the strings f and g.\n",
    "\"\"\"\n",
    "function SymPyDerivatives(f::String, g::String)\n",
    "    evaluatedf = sympify(f)\n",
    "    fx = diff(evaluatedf, x)\n",
    "    fy = diff(evaluatedf, y)\n",
    "    evaluatedg = sympify(g)\n",
    "    gx = diff(evaluatedg, x)\n",
    "    gy = diff(evaluatedg, y)\n",
    "    return [lambdify(fx, (x, y)) lambdify(fy, (x, y));\n",
    "            lambdify(gx, (x, y)) lambdify(gy, (x, y))]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f5c40831",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SymPyMatrixEvaluate"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Given a symbolic representation for a Jacobian matrix,\n",
    "and values for its arguments, returns the evaluated Jacobian matrix.\n",
    "\"\"\"\n",
    "function SymPyMatrixEvaluate(jac,xval::Float64,yval::Float64)\n",
    "    vfx = jac[1, 1](xval, yval)\n",
    "    vfy = jac[1, 2](xval, yval)\n",
    "    vgx = jac[2, 1](xval, yval)\n",
    "    vgy = jac[2, 2](xval, yval)\n",
    "    return [vfx vfy; vgx vgy]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4625fa08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SymPyFun"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Given an array of SymPy expressions,\n",
    "evaluates the expressions at (xval, yval).\n",
    "\"\"\"\n",
    "function SymPyFun(fun,xval::Float64,yval::Float64)\n",
    "    f = fun[1]\n",
    "    g = fun[2]\n",
    "    return [f(xval, yval); g(xval, yval)]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e68d6627",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewtonStep"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Runs one step with Newton's method.\n",
    "\"\"\"\n",
    "function NewtonStep(fun, jac,\n",
    "                    x0::Float64, y0::Float64)\n",
    "    valfun = -SymPyFun(fun, x0, y0)\n",
    "    nfx = norm(valfun)\n",
    "    valmat = SymPyMatrixEvaluate(jac, x0, y0)\n",
    "    update = valmat\\valfun\n",
    "    ndx = norm(update)\n",
    "    x1 = x0 + update[1]\n",
    "    y1 = y0 + update[2]\n",
    "    return [x1, y1, ndx, nfx]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "906042dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Newton (generic function with 4 methods)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function Newton(fun, jac,\n",
    "                x0::Float64, y0::Float64,\n",
    "                maxit::Int64=5,\n",
    "                dxtol::Float64=1.0e-8,\n",
    "                fxtol::Float64=1.0e-8)\n",
    "    xsol, ysol = x0, y0\n",
    "    stri = @sprintf(\"%3d\", 0)\n",
    "    sx1 = @sprintf(\"%.16e\", xsol)\n",
    "    sy1 = @sprintf(\"%.16e\", ysol)\n",
    "    print(\"step           x                        y     \")\n",
    "    println(\"        |update|   |f(x,y)|\")\n",
    "    println(\"$stri : $sx1  $sy1\")\n",
    "    for i=1:maxit\n",
    "        (xsol, ysol, ndx, nfx) = NewtonStep(fun, jac, xsol, ysol)\n",
    "        stri = @sprintf(\"%3d\", i)\n",
    "        sx1 = @sprintf(\"%.16e\", xsol)\n",
    "        sy1 = @sprintf(\"%.16e\", ysol)\n",
    "        sdx = @sprintf(\"%.2e\", ndx)\n",
    "        sfx = @sprintf(\"%.2e\", nfx)\n",
    "        println(\"$stri : $sx1  $sy1  $sdx   $sfx\")\n",
    "        if((ndx < dxtol) | (nfx < fxtol))\n",
    "            return (xsol, ysol, i, false)\n",
    "        end\n",
    "    end\n",
    "    return (xsol, ysol, maxit, true)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "677514ea",
   "metadata": {},
   "source": [
    "The example system we consider is defined by the two equations\n",
    "\n",
    "$$\n",
    "    \\exp(x) - y = 0\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "$$\n",
    "    x y - \\exp(x) = 0.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ac0507c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.9, 2.5)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "one = \"exp(x) - y\"\n",
    "two = \"x*y - exp(x)\"\n",
    "x0, y0 = 0.9, 2.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "39dd03f4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element Vector{SymPy.var\"#118#119\"}:\n",
       " #118 (generic function with 1 method)\n",
       " #118 (generic function with 1 method)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jacmat = SymPyDerivatives(one, two)\n",
    "valmat = SymPyMatrixEvaluate(jacmat, x0, y0)\n",
    "vecfun = SymPyExpressions(one, two)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "efe45657",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step           x                        y             |update|   |f(x,y)|\n",
      "  0 : 9.0000000000000002e-01  2.5000000000000000e+00\n",
      "  1 : 1.0091197782934511e+00  2.7279944573362789e+00  2.53e-01   2.13e-01\n",
      "  2 : 1.0000325513375456e+00  2.7182573929531251e+00  1.33e-02   1.80e-02\n",
      "  3 : 9.9999999970740971e-01  2.7182818262235138e+00  4.07e-05   1.16e-04\n",
      "  4 : 1.0000000000000000e+00  2.7182818284590451e+00  2.25e-09   2.66e-09\n"
     ]
    }
   ],
   "source": [
    "(xsol, ysol, nbit, fail) = Newton(vecfun, jacmat, x0, y0)\n",
    "for i=1:5\n",
    "    xsol, ysol = NewtonStep(vecfun, jacmat, x0, y0)\n",
    "    x0, y0 = xsol, ysol\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b46b70aa",
   "metadata": {},
   "source": [
    "Observe the quadratic convergence."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.8.0",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}