notebooks_solutions/s_geothermal_2D_gpu_kp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Solution: Geothermal 2D on GPUs using \"kernel\" programming\n",
    "\n",
    "Load modules"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Printf\n",
    "using CairoMakie\n",
    "using CUDA"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Averaging and local maximum macros and functions"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "macro d_xa(A) esc(:($A[ix+1, iz] - $A[ix, iz])) end\n",
    "macro d_za(A) esc(:($A[ix, iz+1] - $A[ix, iz])) end\n",
    "macro avx(A)  esc(:(0.5 * ($A[ix, iz] + $A[ix+1, iz]))) end\n",
    "macro avz(A)  esc(:(0.5 * ($A[ix, iz] + $A[ix, iz+1]))) end\n",
    "@views avx(A) = 0.5 .* (A[1:end-1, :] .+ A[2:end, :])\n",
    "@views avz(A) = 0.5 .* (A[:, 1:end-1] .+ A[:, 2:end])\n",
    "@views maxloc(A) = max.(A[2:end-1, 2:end-1], max.(max.(A[1:end-2, 2:end-1], A[3:end, 2:end-1]),\n",
    "                                                  max.(A[2:end-1, 1:end-2], A[2:end-1, 3:end])))"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Compute functions aka \"kernels\""
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function residual_fluxes!(Rqx, Rqz, qx, qz, Pf, K, dx, dz)\n",
    "    ix = (blockIdx().x - 1) * blockDim().x + threadIdx().x\n",
    "    iz = (blockIdx().y - 1) * blockDim().y + threadIdx().y\n",
    "    @inbounds if (ix<=size(Rqx, 1) - 2 && iz<=size(Rqx, 2)    ) Rqx[ix+1, iz] = qx[ix+1, iz] + @avx(K) * @d_xa(Pf) / dx end\n",
    "    @inbounds if (ix<=size(Rqz, 1)     && iz<=size(Rqz, 2) - 2) Rqz[ix, iz+1] = qz[ix, iz+1] + @avz(K) * @d_za(Pf) / dz end\n",
    "    return\n",
    "end\n",
    "\n",
    "function residual_pressure!(RPf, qx, qz, Qf, dx, dz)\n",
    "    ix = (blockIdx().x - 1) * blockDim().x + threadIdx().x\n",
    "    iz = (blockIdx().y - 1) * blockDim().y + threadIdx().y\n",
    "    @inbounds if (ix<=size(RPf, 1) && iz<=size(RPf, 2)) RPf[ix, iz] = @d_xa(qx) / dx + @d_za(qz) / dz - Qf[ix, iz] end\n",
    "    return\n",
    "end\n",
    "\n",
    "function update_fluxes!(qx, qz, Rqx, Rqz, cfl, nx, nz, re)\n",
    "    ix = (blockIdx().x - 1) * blockDim().x + threadIdx().x\n",
    "    iz = (blockIdx().y - 1) * blockDim().y + threadIdx().y\n",
    "    @inbounds if (ix<=size(qx, 1) - 2 && iz<=size(qx, 2)    ) qx[ix+1, iz] -= Rqx[ix+1, iz] / (1.0 + 2cfl * nx / re) end\n",
    "    @inbounds if (ix<=size(qz, 1)     && iz<=size(qz, 2) - 2) qz[ix, iz+1] -= Rqz[ix, iz+1] / (1.0 + 2cfl * nz / re) end\n",
    "    return\n",
    "end\n",
    "\n",
    "function update_pressure!(Pf, RPf, K_max, vdτ, lz, re)\n",
    "    ix = (blockIdx().x - 1) * blockDim().x + threadIdx().x\n",
    "    iz = (blockIdx().y - 1) * blockDim().y + threadIdx().y\n",
    "    @inbounds if (ix<=size(Pf, 1) && iz<=size(Pf, 2)) Pf[ix, iz] -= RPf[ix, iz] * (vdτ * lz / re) / K_max[ix, iz] end\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Main script"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "@views function main()\n",
    "    # CUDA.device!(0) # select your GPU\n",
    "    # physics\n",
    "    lx, lz  = 2.0, 1.0 # domain extend\n",
    "    k0_μ    = 1.0      # background permeability / fluid viscosity\n",
    "    kb_μ    = 1e-6     # barrier permeability / fluid viscosity\n",
    "    Q_in    = 1.0      # injection flux\n",
    "    b_w     = 0.02lx   # barrier width\n",
    "    b_b     = 0.3lz    # barrier bottom location\n",
    "    b_t     = 0.8lz    # barrier top location\n",
    "    # numerics\n",
    "    nz      = 255\n",
    "    nx      = ceil(Int, (nz + 1) * lx / lz) - 1\n",
    "    nthread = (16, 16)\n",
    "    nblock  = cld.((nx, nz), nthread)\n",
    "    cfl     = 1 / 2.1\n",
    "    ϵtol    = 1e-6\n",
    "    maxiter = 30nx\n",
    "    ncheck  = 2nx\n",
    "    re      = 0.8π\n",
    "    st      = ceil(Int, nx / 30)\n",
    "    # preprocessing\n",
    "    dx, dz  = lx / nx, lz / nz\n",
    "    xc, zc  = LinRange(-lx / 2 + dx / 2, lx / 2 - dx / 2, nx), LinRange(dz / 2, lz - dz / 2, nz)\n",
    "    vdτ     = cfl * min(dx, dz)\n",
    "    # init\n",
    "    Pf      = CUDA.zeros(Float64, nx, nz)\n",
    "    RPf     = CUDA.zeros(Float64, nx, nz)\n",
    "    qx      = CUDA.zeros(Float64, nx + 1, nz)\n",
    "    Rqx     = CUDA.zeros(Float64, nx + 1, nz)\n",
    "    qz      = CUDA.zeros(Float64, nx, nz + 1)\n",
    "    Rqz     = CUDA.zeros(Float64, nx, nz + 1)\n",
    "    Qf      = CUDA.zeros(Float64, nx, nz)\n",
    "    K       = k0_μ .* CUDA.ones(Float64, nx, nz)\n",
    "    # set low permeability barrier location\n",
    "    K[ceil(Int, (lx/2-b_w)/dx):ceil(Int, (lx/2+b_w)/dx), ceil(Int, b_b/dz):ceil(Int, b_t/dz)] .= kb_μ\n",
    "    # set wells location\n",
    "    x_iw, x_ew, z_w = ceil.(Int, (lx / 5 / dx, 4lx / 5 / dx, 0.45lz / dz)) # well location\n",
    "    Qf[x_iw:x_iw, z_w:z_w] .=  Q_in / dx / dz # injection\n",
    "    Qf[x_ew:x_ew, z_w:z_w] .= -Q_in / dx / dz # extraction\n",
    "    # init visu\n",
    "    iters_evo = Float64[]; errs_evo = Float64[]\n",
    "    qM, qx_c, qz_c = zeros(nx, nz), zeros(nx, nz), zeros(nx, nz)\n",
    "    fig = Figure(size=(2500, 1200), fontsize=32)\n",
    "    ax = ( Pf  = Axis(fig[1, 1][1, 1]; aspect=DataAspect(), title=\"Pf\"),\n",
    "           K   = Axis(fig[1, 2][1, 1]; aspect=DataAspect(), title=\"log10(K)\"),\n",
    "           qM  = Axis(fig[2, 1][1, 1]; aspect=DataAspect(), title=\"|q|\"),\n",
    "           err = Axis(fig[2, 2]; yscale=log10, title=\"Convergence\", xlabel=\"# iter/nx\", ylabel=\"error\"), )\n",
    "    plt = (fld = ( Pf = heatmap!(ax.Pf, xc, zc, Array(Pf); colormap=:turbo, colorrange=(-1,1)),\n",
    "                   K  = heatmap!(ax.K , xc, zc, Array(log10.(K)); colormap=:turbo, colorrange=(-6,0)),\n",
    "                   qM = heatmap!(ax.qM, xc, zc, qM; colormap=:turbo, colorrange=(0,30)),\n",
    "                   ar = arrows!(ax.Pf, xc[1:st:end], zc[1:st:end], qx_c[1:st:end, 1:st:end], qz_c[1:st:end, 1:st:end]; lengthscale=0.05, arrowsize=15), ),\n",
    "           err = scatterlines!(ax.err, Point2.(iters_evo, errs_evo), linewidth=4), )\n",
    "    Colorbar(fig[1, 1][1, 2], plt.fld.Pf)\n",
    "    Colorbar(fig[1, 2][1, 2], plt.fld.K)\n",
    "    Colorbar(fig[2, 1][1, 2], plt.fld.qM)\n",
    "    # approximate diagonal (Jacobi) preconditioner\n",
    "    K_max = copy(K); K_max[2:end-1, 2:end-1] .= maxloc(K); K_max[:, [1, end]] .= K_max[:, [2, end-1]]\n",
    "    # iterative loop\n",
    "    err = 2ϵtol; iter = 1\n",
    "    while err >= ϵtol && iter <= maxiter\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock residual_fluxes!(Rqx, Rqz, qx, qz, Pf, K, dx, dz)\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock update_fluxes!(qx, qz, Rqx, Rqz, cfl, nx, nz, re)\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock residual_pressure!(RPf, qx, qz, Qf, dx, dz)\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock update_pressure!(Pf, RPf, K_max, vdτ, lz, re)\n",
    "        if iter % ncheck == 0\n",
    "            err = maximum(abs.(RPf))\n",
    "            push!(iters_evo, iter/nx); push!(errs_evo, err)\n",
    "            # visu\n",
    "            qx_c .= Array(avx(qx)); qz_c .= Array(avz(qz)); qM .= sqrt.(qx_c.^2 .+ qz_c.^2)\n",
    "            qx_c ./= qM; qz_c ./= qM\n",
    "            plt.fld.Pf[3] = Array(Pf)\n",
    "            plt.fld.K[3]  = Array(log10.(K))\n",
    "            plt.fld.qM[3] = qM\n",
    "            plt.fld.ar[3] = qx_c[1:st:end, 1:st:end]\n",
    "            plt.fld.ar[4] = qz_c[1:st:end, 1:st:end]\n",
    "            plt.err[1] = Point2.(iters_evo, errs_evo)\n",
    "            IJulia.clear_output(true)\n",
    "            display(fig)\n",
    "            @printf(\"  #iter/nx=%.1f, max(err)=%1.3e\\n\", iter/nx, err)\n",
    "        end\n",
    "        iter += 1\n",
    "    end\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Executing the main script"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "main()"
   ],
   "metadata": {},
   "execution_count": null
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.3"
  },
  "kernelspec": {
   "name": "julia-1.10",
   "display_name": "Julia 1.10.3",
   "language": "julia"
  }
 },
 "nbformat": 4
}