{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Solution of 2-D Poisson equation: CG method"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We consider the Poisson equation\n",
"\n",
"$$\n",
" -(u_{xx} + u_{yy}) = f, \\quad\\textrm{in}\\quad \\Omega = [0,1] \\times [0,1]\n",
"$$\n",
"\n",
"with boundary condition\n",
"\n",
"$$\n",
"u = 0, \\quad\\textrm{on}\\quad \\partial\\Omega\n",
"$$\n",
"\n",
"We will take\n",
"\n",
"$$\n",
"f = 1\n",
"$$\n",
"\n",
"Make a uniform grid with $n$ points in each direction, with spacing\n",
"\n",
"$$\n",
"h = \\frac{1}{n-1}\n",
"$$\n",
"\n",
"The grid points are\n",
"\n",
"$$\n",
"x_i = ih, \\quad 0 \\le i \\le n-1\n",
"$$\n",
"\n",
"$$\n",
"y_j = jh, \\quad 0 \\le j \\le n-1\n",
"$$\n",
"\n",
"The finite difference approximation at $(i,j)$ is\n",
"\n",
"$$\n",
"-\\frac{u_{i-1,j} - 2u_{i,j} + u_{i+1,j}}{\\Delta x^2} - \\frac{u_{i,j-1} - 2u_{i,j} + u_{i,j+1}}{\\Delta y^2} = f_{i,j}, \\qquad 1 \\le i \\le n-2, \\quad 1 \\le j \\le n-2\n",
"$$\n",
"\n",
"This set of equations is of the form\n",
"\n",
"$$\n",
"A U = b\n",
"$$\n",
"\n",
"where\n",
"\n",
"$$\n",
"A \\in \\Re^{(n-2)^2 \\times (n-2)^2}, \\qquad U,b \\in \\Re^{(n-2)^2}\n",
"$$\n",
"\n",
"However, we will never form the matrix $A$ explicitly, and we do not have to renumber the unknowns $\\{ u_{i,j} \\}$ into a long vector."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## CG Algorithm\n",
"Here is the alogorithm which we repeat from the 1-D Poisson notebook.\n",
"\n",
"* Set initial guess $u_0 = 0$, $r_0 = b - A u_0$, $p_0 = 0$\n",
"* For $k=0,1,\\ldots$\n",
" * If $\\| r_k \\| < TOL \\cdot \\|r_0\\|$, then stop\n",
" * If $k=0$, $\\beta_1 = 0$\n",
" * If $k > 0$, $\\beta_{k+1} = \\frac{r_k^\\top r_k}{r_{k-1}^\\top r_{k-1}}$\n",
" * $p_{k+1} = r_k + \\beta_{k+1} p_k$\n",
" * $\\alpha_{k+1} = \\frac{r_k^\\top r_k}{p_{k+1}^\\top A p_{k+1}}$\n",
" * $u_{k+1} = u_k + \\alpha_{k+1} p_{k+1}$\n",
" * $r_{k+1} = r_k - \\alpha_{k+1} A p_{k+1}$\n",
" \n",
"The above algorithm is not useful for computer implementation, since it is written in terms of indices, but we do not need to store all the values. We modify it below, so that we store only the quantities needed for the iterations.\n",
" \n",
"* Set initial guess $u = 0$\n",
"* $r = f - A u $\n",
"* $p = 0$\n",
"* $r_{norm} = \\| f \\|$, $r_{old} = r_{new} = 0$\n",
"* For $k=0,1,\\ldots,$ itmax\n",
" * $r_{new} = \\| r \\|$\n",
" * If $r_{new} < TOL \\cdot r_{norm}$, then stop\n",
" * If $k=0$, $\\beta = 0$\n",
" * If $k > 0$, $\\beta = r_{new}^2 / r_{old}^2$\n",
" * $p = r + \\beta p$\n",
" * $A_p = A p$\n",
" * $\\alpha = r_{new}^2 / (p^\\top A_p)$\n",
" * $u = u + \\alpha p$\n",
" * $r = r - \\alpha A_p$\n",
" * $r_{old} = r_{new}$\n",
" \n",
"Written like this, each step of the algorithm requires one matrix-vector product, two dot products and three saxpy operations."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Code"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
"%matplotlib inline\n",
"%config InlineBackend.figure_format = 'svg'\n",
"import numpy as np\n",
"from matplotlib import pyplot as plt"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This function computes matrix-vector product."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"# Note that boundary values of res are always zero.\n",
"# res is computed only for interior points.\n",
"def ax(h,u):\n",
" n = u.shape[0]\n",
" res = np.zeros((n,n))\n",
" for i in range(1,n-1):\n",
" for j in range(1,n-1):\n",
" uxx = (u[i-1,j]-2*u[i,j]+u[i+1,j])/h**2\n",
" uyy = (u[i,j-1]-2*u[i,j]+u[i,j+1])/h**2\n",
" res[i,j] = -(uxx + uyy)\n",
" return res"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Note that we computed this without forming the matrix explicitly."
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"xmin,xmax = 0.0,1.0\n",
"n = 21\n",
"\n",
"h = (xmax - xmin)/(n-1)\n",
"x = np.linspace(xmin,xmax,n)\n",
"X,Y = np.meshgrid(x,x,indexing='ij')\n",
"f = np.ones((n,n))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The next function implements the CG method, it is identical to the function used for 1d BVP. Use Frobenius norm since the arrays are two dimensional."
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"def solve_cg(res, u):\n",
" r_norm = np.linalg.norm(res, 'fro')\n",
" res_old, res_new = 0.0, 0.0\n",
" p = np.zeros_like(res)\n",
" for it in range(itmax):\n",
" res_new = np.linalg.norm(res, 'fro')\n",
" if res_new < TOL * r_norm:\n",
" break\n",
" if it==0:\n",
" beta = 0.0\n",
" else:\n",
" beta = res_new**2 / res_old**2\n",
" p = res + beta * p\n",
" ap = ax(h,p)\n",
" alpha = res_new**2 / np.sum(p*ap)\n",
" u += alpha * p\n",
" res -= alpha * ap\n",
" res_old = res_new\n",
" \n",
" print(\"Number of iterations = %d\" % it)\n",
" return u"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Solve on a sample problem."
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Number of iterations = 31\n"
]
}
],
"source": [
"TOL = 1.0e-6\n",
"itmax = 2000\n",
"\n",
"u = np.zeros((n,n))\n",
"res = np.array(f)\n",
"\n",
"# On boundary, solution is fixed to zero.\n",
"# Hence we make residual zero, in which case solution\n",
"# will not change at these points.\n",
"res[0,:] = 0.0\n",
"res[n-1,:] = 0.0\n",
"res[:,0] = 0.0\n",
"res[:,n-1] = 0.0\n",
"\n",
"u = solve_cg(res, u)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Visualize the solution as a color plot."
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"plt.contourf(X,Y,u,cmap=plt.cm.jet)\n",
"plt.axis('equal')\n",
"plt.colorbar();"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The CG algorithm looks same as in the case of 1d BVP. Only the function `ax` is different.\n",
"\n",
"**Remark**: We do not have to renumber the unknown, which is a 2-D array, into a one-dimensional vector. All quantities like search direction `p`, residual `res` are stored as 2-D array. The function `ax` takes a 2-D array and performs the matrix-vector product."
]
}
],
"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.10.8"
},
"toc": {
"base_numbering": 1,
"nav_menu": {},
"number_sections": true,
"sideBar": true,
"skip_h1_title": false,
"title_cell": "Table of Contents",
"title_sidebar": "Contents",
"toc_cell": false,
"toc_position": {},
"toc_section_display": true,
"toc_window_display": false
}
},
"nbformat": 4,
"nbformat_minor": 4
}