M03-Poisson-3D.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8b241edb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary packages\n",
    "using PyPlot\n",
    "using SparseArrays\n",
    "using LinearAlgebra\n",
    "using IterativeSolvers\n",
    "using AlgebraicMultigrid\n",
    "using Printf\n",
    "using WriteVTK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "00e46cf2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "poisson_prolongation (generic function with 1 method)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "struct Grid\n",
    "    x::Vector{Float64}\n",
    "    y::Vector{Float64}\n",
    "    z::Vector{Float64}    \n",
    "    nx::Int64\n",
    "    ny::Int64\n",
    "    nz::Int64\n",
    "    LX::Float64\n",
    "    LY::Float64\n",
    "    LZ::Float64\n",
    "    dx::Float64\n",
    "    dy::Float64\n",
    "    dz::Float64\n",
    "    solution::Array{Float64,3}\n",
    "    rhs::Array{Float64,3}\n",
    "    residual::Array{Float64,3}\n",
    "    function Grid(W,H,D,nx,ny,nz)\n",
    "        x = LinRange(0,W,nx)\n",
    "        y = LinRange(0,H,ny)\n",
    "        z = LinRange(0,D,nz)\n",
    "        dx = x[2]-x[1]\n",
    "        dy = y[2]-y[1]\n",
    "        dz = z[2]-z[1]\n",
    "        \n",
    "        return new(x,y,z,nx,ny,nz,W,H,D,dx,dy,dz,zeros(nz,ny,nx),zeros(nz,ny,nx),zeros(nz,ny,nx))\n",
    "    end\n",
    "end\n",
    "\n",
    "function poisson_smoother!(grid::Grid,cycles::Int64)\n",
    "    coef_center = -2/grid.dx^2 -2/grid.dy^2 - 2/grid.dz^2\n",
    "    coef_right = 1/grid.dx^2\n",
    "    coef_left = 1/grid.dx^2\n",
    "    coef_up = 1/grid.dy^2\n",
    "    coef_down = 1/grid.dy^2\n",
    "    coef_front = 1/grid.dz^2\n",
    "    coef_back = 1/grid.dz^2\n",
    "    theta = 1.75;\n",
    "    cycle = 1\n",
    "    while cycle <= cycles\n",
    "        #red-black update        \n",
    "        for j in 2:(grid.nx-1)\n",
    "            for i in 2:(grid.ny-1) # julia is column-major - inner loop over columns\n",
    "                for k in 2:(grid.nz-1)\n",
    "                    # in Gauss-Seidel iterations, the solution (s) is updated within each iteration.\n",
    "                    # note that the 1.0 is the (uniform) value of the right hand side, hard-coded for speed.\n",
    "                    this_residual = grid.rhs[k,i,j] - coef_left*grid.solution[k,i,j-1] - coef_right*grid.solution[k,i,j+1] - coef_up*grid.solution[k,i-1,j] - coef_down*grid.solution[k,i+1,j] - coef_center*grid.solution[k,i,j] - coef_front*grid.solution[k-1,i,j]-coef_back*grid.solution[k+1,i,j]\n",
    "                    grid.solution[k,i,j] += theta*this_residual/coef_center\n",
    "                end\n",
    "            end            \n",
    "        end\n",
    "        cycle += 1\n",
    "    end\n",
    "    Threads.@threads for j in 2:(grid.nx-1)\n",
    "         for i in 2:(grid.ny-1) # julia is column-major - inner loop over columns\n",
    "            for k in 2:(grid.nz-1)\n",
    "                 # in Gauss-Seidel iterations, the solution (s) is updated within each iteration.\n",
    "                 # note that the 1.0 is the (uniform) value of the right hand side, hard-coded for speed.\n",
    "                 grid.residual[k,i,j] = grid.rhs[k,i,j] - coef_left*grid.solution[k,i,j-1] - coef_right*grid.solution[k,i,j+1] - coef_up*grid.solution[k,i-1,j] - coef_down*grid.solution[k,i+1,j] - coef_center*grid.solution[k,i,j] - coef_front*grid.solution[k-1,i,j]-coef_back*grid.solution[k+1,i,j]\n",
    "            end\n",
    "         end\n",
    "    end\n",
    "end\n",
    "\n",
    "function poisson_restriction!(fine::Grid,coarse::Grid,mgunit::Int64)\n",
    "    # restriction (fine) -> (coarse)\n",
    "    coarse.rhs[:,:,:] = fine.residual[1:mgunit:end,1:mgunit:end,1:mgunit:end]\n",
    "end\n",
    "\n",
    "function poisson_prolongation(fine::Grid,coarse::Grid,mgunit::Int64)\n",
    "    dsol = zeros(fine.nz,fine.ny,fine.nx)\n",
    "    # use bilinear interpolation to go from coarse to fine\n",
    "    local coarse_j::Int64\n",
    "    local coarse_i::Int64\n",
    "    local coarse_k::Int64\n",
    "    local dx::Float64\n",
    "    local dy::Float64\n",
    "    local dz::Float64\n",
    "    for j in 2:fine.nx-1\n",
    "        coarse_j = Int64( floor( fine.x[j]/coarse.dx )+1 )\n",
    "        dx = (fine.x[j]-coarse.x[coarse_j])/coarse.dx\n",
    "        \n",
    "        for i in 2:fine.ny-1\n",
    "            coarse_i = Int64( floor( fine.y[i]/coarse.dy) + 1)\n",
    "            dy = (fine.y[i]-coarse.y[coarse_i]) / coarse.dy\n",
    "            \n",
    "            for k in 2:fine.nz-1\n",
    "                coarse_k = Int64( floor( fine.z[k]/coarse.dz) + 1 )\n",
    "                dz = (fine.z[k]-coarse.z[coarse_k]) / coarse.dz\n",
    "                                \n",
    "                dsol[k,i,j] += (1.0-dz)*(1.0-dx)*(1.0-dy)*coarse.solution[coarse_k,coarse_i,coarse_j]\n",
    "                dsol[k,i,j] += (1.0-dz)*(1.0-dx)*(dy)*coarse.solution[coarse_k,coarse_i+1,coarse_j]\n",
    "                dsol[k,i,j] += (1.0-dz)*(dx)*(1.0-dy)*coarse.solution[coarse_k,coarse_i,coarse_j+1]\n",
    "                dsol[k,i,j] += (1.0-dz)*(dx)*(dy)*coarse.solution[coarse_k,coarse_i+1,coarse_j+1]\n",
    "            \n",
    "                dsol[k,i,j] += (dz)*(1.0-dx)*(1.0-dy)*coarse.solution[coarse_k+1,coarse_i,coarse_j]\n",
    "                dsol[k,i,j] += (dz)*(1.0-dx)*(dy)*coarse.solution[coarse_k+1,coarse_i+1,coarse_j]\n",
    "                dsol[k,i,j] += (dz)*(dx)*(1.0-dy)*coarse.solution[coarse_k+1,coarse_i,coarse_j+1]\n",
    "                dsol[k,i,j] += (dz)*(dx)*(dy)*coarse.solution[coarse_k+1,coarse_i+1,coarse_j+1]\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    return dsol\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9978e21b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "visualization (generic function with 1 method)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# code to write a .vtr file:\n",
    "function visualization(grid::Grid; filename=\"test.vtr\")\n",
    "    # write the visualization output from the regular grid as a .vts file.\n",
    "    vtk_grid(filename, grid.x, grid.y, grid.z) do vtk\n",
    "        vtk[\"phi\"] = grid.solution\n",
    "        vtk[\"residual\"] = grid.residual\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7e91610b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "257 257 257\n",
      "129 129 129\n",
      "65 65 65\n",
      "33 33 33\n",
      "17 17 17\n",
      "9 9 9\n",
      "4072.0234527811845\n",
      "7.195751805247488e-35\n",
      "5.28826762189299e-19\n",
      "3.06277106997723e-12\n",
      "0.026064626904493494\n",
      "0.32175193122222967\n",
      "Cycle: 1 Residual: 4.495345988710504\n",
      "2.082984158763465e-34\n",
      "1.644374112008982e-18\n",
      "8.736698114538585e-12\n",
      "0.05214179195053081\n",
      "0.37807220960764754\n",
      "Cycle: 2 Residual: 3.0069695642498884\n",
      "2.454368118700532e-34\n",
      "1.6764715810331503e-18\n",
      "8.422450634259428e-12\n",
      "0.026109846207788456\n",
      "0.06106845311503858\n",
      "Cycle: 3 Residual: 0.17395308576691493\n",
      "5.829856511851134e-35\n",
      "5.365404022879292e-19\n",
      "2.788064664751448e-12\n",
      "5.151610695533636e-5\n",
      "0.0002749022816691169\n",
      "Cycle: 4 Residual: 0.002932741067690726\n",
      "2.2718652317029178e-37\n",
      "1.58392645281046e-21\n",
      "7.804441565944164e-15\n",
      "1.740227074370019e-5\n",
      "2.6717209189257918e-5\n",
      "Cycle: 5 Residual: 0.0002736212921653943\n",
      "4.941159538911312e-38\n",
      "3.5216844684654147e-22\n",
      "1.826988571798764e-15\n",
      "1.3700881206556837e-6\n",
      "3.1283771698018624e-6\n",
      "Cycle: 6 Residual: 1.064603988115246e-5\n",
      "2.6800683031591532e-39\n",
      "2.682638212050127e-23\n",
      "1.3496140363000422e-16\n",
      "9.755633151386417e-9\n",
      "7.4342236969586785e-9\n",
      "Cycle: 7 Residual: 3.3051707906067177e-7\n",
      "2.0894386099040811e-41\n",
      "1.8684942364912128e-25\n",
      "9.94749696222506e-19\n",
      "1.8330844019464207e-9\n",
      "3.4730412827743573e-9\n",
      "Cycle: 8 Residual: 2.3255896921733756e-8\n",
      "6.2268096103917e-42\n",
      "3.4810508296686164e-26\n",
      "1.7959847691678989e-19\n",
      "8.056725497860943e-11\n",
      "1.707180838259513e-10\n",
      "Cycle: 9 Residual: 8.694241405547011e-9\n",
      "2.109722687792148e-43\n",
      "1.5988681565268777e-27\n",
      "8.035383698327502e-21\n",
      "1.393985055167889e-12\n",
      "1.7145505104135283e-12\n",
      "Cycle: 10 Residual: 8.73048789952478e-9\n",
      " 84.313531 seconds (750.78 k allocations: 4.566 GiB, 0.48% gc time, 0.57% compilation time)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1-element Vector{String}:\n",
       " \"poisson_3d.vtr\""
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# define an indexing function that maps from (i,j) node numbering to an integer index\n",
    "nx = 257\n",
    "ny = 257\n",
    "nz = 257\n",
    "width = 1e3\n",
    "height = 1.5e3\n",
    "depth = 1.25e3\n",
    "\n",
    "# Define levels of grid\n",
    "nlevel=6\n",
    "mgunit=2\n",
    "grids = Vector{Grid}() # grids to use with multigrid solver\n",
    "cycles = Vector{Int64}() # number of gauss-seidel iterations to take at each multigrid level\n",
    "for level in 1:nlevel\n",
    "    level_nx = Int64((nx-1)/mgunit^(level-1)+1)\n",
    "    level_ny = Int64((ny-1)/mgunit^(level-1)+1)\n",
    "    level_nz = Int64((nz-1)/mgunit^(level-1)+1)\n",
    "    println(level_nx,' ', level_ny,' ',level_nz)\n",
    "    push!(grids, Grid(width,height,depth,level_nx,level_ny,level_nz) )\n",
    "    push!(cycles,2^(level+1)*5)\n",
    "end\n",
    "\n",
    "# set RHS for finest grid.\n",
    "level = 1\n",
    "for j in 2:grids[level].nx-1\n",
    "    for i in 2:grids[level].ny-1\n",
    "        for k in 2:grids[level].nz-1\n",
    "            grids[level].rhs[k,i,j] = 1.0        \n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "function multigrid_smoother!(grids::Vector{Grid},cycles::Vector{Int64},outer_cycles::Int64)\n",
    "    level = 1\n",
    "    cycle = 1\n",
    "    # outer solver iterations\n",
    "    poisson_smoother!(grids[1],0) # compute starting residual\n",
    "    println(norm(grids[1].residual))\n",
    "    while cycle <= outer_cycles\n",
    "        for level in 1:nlevel-1\n",
    "            if cycle==1 || level > 1 # don't re-run smoother at first level if we're past first cycle.\n",
    "                poisson_smoother!(grids[level],cycles[level])     \n",
    "            end\n",
    "            poisson_restriction!(grids[level],grids[level+1],mgunit)\n",
    "            #println(norm(grids[level].residual))\n",
    "        end\n",
    "        # end at grid level nlevel-1\n",
    "        # compute correction to current solution using a direct solver\n",
    "        # interpolate corrections onto finer grids\n",
    "        for level in nlevel:-1:2\n",
    "            poisson_smoother!(grids[level],cycles[level])\n",
    "            println(norm(grids[level].residual))\n",
    "            dsol = poisson_prolongation(grids[level-1],grids[level],mgunit)\n",
    "            grids[level-1].solution[:,:,:] += dsol\n",
    "            # interpolate correction from level+1 onto level\n",
    "        end\n",
    "        poisson_smoother!(grids[1],cycles[1]) # run smoother on final (finest) level.  \n",
    "        println(\"Cycle: \",cycle,\" Residual: \",norm(grids[1].residual))\n",
    "        cycle += 1 # increment mg cycle counter\n",
    "    end\n",
    "\n",
    "end\n",
    "\n",
    "@time multigrid_smoother!(grids,cycles,10)\n",
    "visualization(grids[1] ; filename=\"poisson_3d.vtr\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05f12366",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 20 Threads 1.7.1",
   "language": "julia",
   "name": "julia-20-threads-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}