notebooks/geothermal_2D_gpu_kp_ad_inv.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Geothermal 2D on GPUs using kernel programming and AD for inversion\n",
    "\n",
    "Load modules"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Printf, LinearAlgebra\n",
    "using CairoMakie\n",
    "using CUDA\n",
    "using Enzyme"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Averaging and local maximum macros and support 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])))\n",
    "\n",
    "function smooth_d!(A2, A)\n",
    "    ix = (blockIdx().x - 1) * blockDim().x + threadIdx().x\n",
    "    iz = (blockIdx().y - 1) * blockDim().y + threadIdx().y\n",
    "    if (ix>1 && ix<size(A, 1) && iz>1 && iz<size(A, 2))\n",
    "        @inbounds A2[ix, iz] = A[ix, iz] + 0.2 * (A[ix+1, iz] - 2A[ix, iz] + A[ix-1, iz] + A[ix, iz+1] - 2A[ix, iz] + A[ix, iz-1])\n",
    "    end\n",
    "    return\n",
    "end\n",
    "\n",
    "function smooth!(A2, A, nthread, nblock; nsm=1)\n",
    "    for _ ∈ 1:nsm\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock smooth_d!(A2, A)\n",
    "        A, A2 = A2, A\n",
    "    end\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Forward kernels"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function residual_fluxes!()\n",
    "    #= ??? =#\n",
    "    return\n",
    "end\n",
    "\n",
    "function residual_pressure!()\n",
    "    #= ??? =#\n",
    "    return\n",
    "end\n",
    "\n",
    "function update_fluxes!()\n",
    "    #= ??? =#\n",
    "    return\n",
    "end\n",
    "\n",
    "function update_pressure!()\n",
    "    #= ??? =#\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Generic autodiff call"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "@inline ∇(fun, args...) = (Enzyme.autodiff_deferred(Enzyme.Reverse, fun, args...); return)\n",
    "const DupNN = DuplicatedNoNeed"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Forward solver"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "@views function forward_solve!(logK, fields, scalars, iter_params; visu=nothing)\n",
    "    (;Pf, qx, qz, Qf, RPf, Rqx, Rqz, K)               = fields\n",
    "    (;nx, nz, dx, dz, nthread, nblock)                = scalars\n",
    "    (;cfl, re, vdτ, lz, ϵtol, maxiter, ncheck, K_max) = iter_params\n",
    "    isnothing(visu) || ((;qx_c, qz_c, qM, fig, plt, st) = visu)\n",
    "    K .= exp.(logK)\n",
    "    # approximate diagonal (Jacobi) preconditioner\n",
    "    K_max .= K; K_max[2:end-1, 2:end-1] .= maxloc(K); K_max[:, [1, end]] .= K_max[:, [2, end-1]]\n",
    "    # iterative loop\n",
    "    iters_evo = Float64[]; errs_evo = Float64[]\n",
    "    err = 2ϵtol; iter = 1\n",
    "    while err >= ϵtol && iter <= maxiter\n",
    "        #= ??? =#\n",
    "        #= ??? =#\n",
    "        #= ??? =#\n",
    "        #= ??? =#\n",
    "        if iter % ncheck == 0\n",
    "            err = maximum(abs.(RPf))\n",
    "            push!(iters_evo, iter/nx); push!(errs_evo, err)\n",
    "            @printf(\"  #iter/nx=%.1f, max(err)=%1.3e\\n\", iter/nx, err)\n",
    "            if !isnothing(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",
    "            end\n",
    "        end\n",
    "        iter += 1\n",
    "    end\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Adjoint solver"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "@views function adjoint_solve!(logK, fwd_params, adj_params, loss_params)\n",
    "    # unpack forward\n",
    "    (;Pf, qx, qz, Qf, RPf, Rqx, Rqz, K) = fwd_params.fields\n",
    "    (;nx, nz, dx, dz, nthread, nblock)  = fwd_params.scalars\n",
    "    # unpack adjoint\n",
    "    (;P̄f, q̄x, q̄z, R̄Pf, R̄qx, R̄qz, Ψ_qx, Ψ_qz, Ψ_Pf)      = adj_params.fields\n",
    "    (;∂J_∂Pf)                                           = loss_params.fields\n",
    "    (;cfl, re_a, vdτ, lz, ϵtol, maxiter, ncheck, K_max) = adj_params.iter_params\n",
    "    # iterative loop\n",
    "    iters_evo = Float64[]; errs_evo = Float64[]\n",
    "    err = 2ϵtol; iter = 1\n",
    "    while err >= ϵtol && iter <= maxiter\n",
    "        R̄qx .= Ψ_qx\n",
    "        R̄qz .= Ψ_qz\n",
    "        P̄f  .= #= ??? =#\n",
    "        q̄x  .= 0.0\n",
    "        q̄z  .= 0.0\n",
    "        #= ??? =#\n",
    "        #= ??? =#\n",
    "        #= ??? =#\n",
    "        R̄Pf .= #= ??? =#\n",
    "        #= ??? =#\n",
    "        #= ??? =#\n",
    "        if iter % ncheck == 0\n",
    "            err = maximum(abs.(P̄f))\n",
    "            push!(iters_evo, iter/nx); push!(errs_evo, err)\n",
    "            @printf(\"  #iter/nx=%.1f, max(err)=%1.6e\\n\", iter/nx, err)\n",
    "        end\n",
    "        iter += 1\n",
    "    end\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Loss function"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "@views function loss(logK, fwd_params, loss_params; kwargs...)\n",
    "    (;Pf_obs)       = loss_params.fields\n",
    "    (;ixobs, izobs) = loss_params.scalars\n",
    "    @info \"Forward solve\"\n",
    "    #= ??? =#\n",
    "    Pf = fwd_params.fields.Pf\n",
    "    return 0.5*sum((Pf[ixobs, izobs] .- Pf_obs).^2)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Gradient of the loss function"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function ∇loss!(logK̄, logK, fwd_params, adj_params, loss_params; reg=nothing, kwargs...)\n",
    "    # unpack\n",
    "    (;R̄qx, R̄qz, Ψ_qx, Ψ_qz)    = adj_params.fields\n",
    "    (;Pf, qx, qz, Rqx, Rqz, K) = fwd_params.fields\n",
    "    (;dx, dz, nthread, nblock) = fwd_params.scalars\n",
    "    (;Pf_obs, ∂J_∂Pf)          = loss_params.fields\n",
    "    (;ixobs, izobs)            = loss_params.scalars\n",
    "    @info \"Forward solve\"\n",
    "    #= ??? =#\n",
    "    # set tangent\n",
    "    ∂J_∂Pf[ixobs, izobs] .= #= ??? =#\n",
    "    @info \"Adjoint solve\"\n",
    "    #= ??? =#\n",
    "    # evaluate gradient dJ_dK\n",
    "    R̄qx .= .-Ψ_qx\n",
    "    R̄qz .= .-Ψ_qz\n",
    "    logK̄ .= 0.0\n",
    "    #= ??? =#\n",
    "    # Tikhonov regularisation (smoothing)\n",
    "    if !isnothing(reg)\n",
    "        (;nsm, Tmp) = reg\n",
    "        Tmp .= logK̄; smooth!(logK̄, Tmp, nthread, nblock; nsm)\n",
    "    end\n",
    "    logK̄ .*= K # convert to dJ_dlogK by chain rule\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",
    "    # observations\n",
    "    xobs_rng = LinRange(-lx / 6, lx / 6, 8)\n",
    "    zobs_rng = LinRange(0.25lz, 0.85lz , 8)\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π # fwd re\n",
    "    st       = ceil(Int, nx / 30)\n",
    "    # GD params\n",
    "    ngd      = 50\n",
    "    Δγ       = 0.2\n",
    "    # preprocessing\n",
    "    re_a     = 2re  # adjoint re\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",
    "    ixobs    = floor.(Int, (xobs_rng .- xc[1]) ./ dx) .+ 1\n",
    "    izobs    = floor.(Int, (zobs_rng .- zc[1]) ./ dz) .+ 1\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",
    "    logK     = CUDA.zeros(Float64, nx, nz)\n",
    "    Tmp      = CUDA.zeros(Float64, nx, nz)\n",
    "    # init adjoint storage\n",
    "    Ψ_qx     = CUDA.zeros(Float64, nx + 1, nz)\n",
    "    q̄x       = CUDA.zeros(Float64, nx + 1, nz)\n",
    "    R̄qx      = CUDA.zeros(Float64, nx + 1, nz)\n",
    "    Ψ_qz     = CUDA.zeros(Float64, nx, nz + 1)\n",
    "    q̄z       = CUDA.zeros(Float64, nx, nz + 1)\n",
    "    R̄qz      = CUDA.zeros(Float64, nx, nz + 1)\n",
    "    Ψ_Pf     = CUDA.zeros(Float64, nx, nz)\n",
    "    P̄f       = CUDA.zeros(Float64, nx, nz)\n",
    "    R̄Pf      = CUDA.zeros(Float64, nx, nz)\n",
    "    ∂J_∂Pf   = CUDA.zeros(Float64, nx, nz)\n",
    "    dJ_dlogK = CUDA.zeros(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",
    "    logK .= log.(K)\n",
    "    K_max = copy(K)\n",
    "    # set wells location\n",
    "    x_iw, x_ew, z_w = ceil.(Int, (lx / 5 / dx, 4lx / 5 / dx, 0.45lz / dz))\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",
    "                   xobs = scatter!(ax.K , vec(Point2.(xobs_rng, zobs_rng')); color=:white),\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",
    "    # action\n",
    "    fwd_params = (\n",
    "        fields      = (;Pf, qx, qz, Qf, RPf, Rqx, Rqz, K),\n",
    "        scalars     = (;nx, nz, dx, dz, nthread, nblock),\n",
    "        iter_params = (;cfl, re, vdτ, lz, ϵtol, maxiter, ncheck, K_max),\n",
    "    )\n",
    "    fwd_visu = (;qx_c, qz_c, qM, fig, plt, st)\n",
    "    @info \"Synthetic solve\"\n",
    "    #= ??? =#\n",
    "    # store true data\n",
    "    Pf_obs = copy(Pf[ixobs, izobs])\n",
    "    adj_params = (\n",
    "        fields  = (;P̄f, q̄x, q̄z, R̄Pf, R̄qx, R̄qz, Ψ_qx, Ψ_qz, Ψ_Pf),\n",
    "        iter_params = (;cfl, re_a, vdτ, lz, ϵtol, maxiter, ncheck, K_max),\n",
    "    )\n",
    "    loss_params = (\n",
    "        fields  = (;Pf_obs, ∂J_∂Pf),\n",
    "        scalars = (;ixobs, izobs),\n",
    "    )\n",
    "    reg = (;nsm=20, Tmp)\n",
    "    # loss functions\n",
    "    J(_logK) = loss(_logK, fwd_params, loss_params)\n",
    "    ∇J!(_logK̄, _logK) = ∇loss!(_logK̄, _logK, fwd_params, adj_params, loss_params; reg)\n",
    "    @info \"Inversion for K\"\n",
    "    # initial guess\n",
    "    K    .= k0_μ\n",
    "    logK .= log.(K)\n",
    "    @info \"Gradient descent - inversion for K\"\n",
    "    cost_evo = Float64[]\n",
    "    for igd in 1:ngd\n",
    "        printstyled(\"> GD iter $igd \\n\"; bold=true, color=:green)\n",
    "        # evaluate gradient of the cost function\n",
    "        #= ??? =#\n",
    "        # update logK\n",
    "        γ = Δγ / maximum(abs.(dJ_dlogK))\n",
    "        @. logK -= #= ??? =#\n",
    "        @printf \"  min(K) = %1.2e \\n\" minimum(K)\n",
    "        # loss\n",
    "        push!(cost_evo, J(logK))\n",
    "        @printf \"  --> Loss J = %1.2e (γ = %1.2e)\\n\" last(cost_evo)/first(cost_evo) γ\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.(1:igd, cost_evo ./ 0.999cost_evo[1])\n",
    "        IJulia.clear_output(true)\n",
    "        display(fig)\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
}