notebooks_solutions/s_geothermal_2D_gpu_kp_ad_inv_optim.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Solution: Geothermal 2D on GPUs using kernel programming and AD for inversion with Optim.jl\n",
    "\n",
    "Load modules"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Printf, LinearAlgebra\n",
    "using CairoMakie\n",
    "using CUDA\n",
    "using Enzyme\n",
    "using Optim"
   ],
   "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!(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": [
    "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",
    "        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",
    "            @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  .= .-∂J_∂Pf\n",
    "        q̄x  .= 0.0\n",
    "        q̄z  .= 0.0\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock ∇(residual_fluxes!,\n",
    "            DupNN(Rqx, R̄qx),\n",
    "            DupNN(Rqz, R̄qz),\n",
    "            DupNN(qx, q̄x),\n",
    "            DupNN(qz, q̄z),\n",
    "            DupNN(Pf, P̄f),\n",
    "            Const(K), Const(dx), Const(dz))\n",
    "        P̄f[[1, end], :] .= 0.0; P̄f[:, [1, end]] .= 0.0\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock update_pressure!(Ψ_Pf, P̄f, K_max, vdτ, lz, re_a)\n",
    "        R̄Pf .= Ψ_Pf\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock ∇(residual_pressure!,\n",
    "            DupNN(RPf, R̄Pf),\n",
    "            DupNN(qx, q̄x),\n",
    "            DupNN(qz, q̄z),\n",
    "            Const(Qf), Const(dx), Const(dz))\n",
    "        CUDA.@sync @cuda threads=nthread blocks=nblock update_fluxes!(Ψ_qx, Ψ_qz, q̄x, q̄z, cfl, nx, nz, re_a)\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",
    "    forward_solve!(logK, fwd_params...; kwargs...)\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",
    "    forward_solve!(logK, fwd_params...; kwargs...)\n",
    "    # set tangent\n",
    "    ∂J_∂Pf[ixobs, izobs] .= Pf[ixobs, izobs] .- Pf_obs\n",
    "    @info \"Adjoint solve\"\n",
    "    adjoint_solve!(logK, fwd_params, adj_params, loss_params)\n",
    "    # evaluate gradient dJ_dK\n",
    "    R̄qx .= .-Ψ_qx\n",
    "    R̄qz .= .-Ψ_qz\n",
    "    logK̄ .= 0.0\n",
    "    CUDA.@sync @cuda threads=nthread blocks=nblock ∇(residual_fluxes!,\n",
    "        DupNN(Rqx, R̄qx),\n",
    "        DupNN(Rqz, R̄qz),\n",
    "        Const(qx), Const(qz), Const(Pf),\n",
    "        DupNN(logK, logK̄),\n",
    "        Const(dx), Const(dz))\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!(1) # 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",
    "    # 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",
    "    logK    = CUDA.zeros(Float64, nx, nz)\n",
    "    K       = k0_μ .* CUDA.ones(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",
    "    # 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",
    "    forward_solve!(logK, fwd_params...; visu=fwd_visu)\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=50, 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",
    "    # Optim\n",
    "    opt = Optim.Options(\n",
    "        f_tol      = 1e-2,\n",
    "        g_tol      = 1e-6,\n",
    "        iterations = 20,\n",
    "        store_trace=true, show_trace=true,\n",
    "    )\n",
    "    result = optimize(J, ∇J!, logK, LBFGS(), opt)\n",
    "    K .= exp.(Optim.minimizer(result))\n",
    "    @show result\n",
    "    # visu\n",
    "    errs_evo = Optim.f_trace(result)\n",
    "    errs_evo ./= errs_evo[1]\n",
    "    iters_evo = 1:length(errs_evo)\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",
    "    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
}