begin
    import Pkg
    # careful: this is _not_ a reproducible environment
    # activate the local environment
    Pkg.activate(".")
    Pkg.instantiate()
    using PlutoUI, PlutoLinks
    using CairoMakie
end
@revise using Ariadne
using Krylov
function bc_periodic!(u)
    N, M = size(u)
    N = N - 2
    M = M - 2

    # Enforce boundary conditions
    # (wrap around)
    u[0, :] .= u[N, :]
    u[N + 1, :] .= u[1, :]
    u[:, 0] .= u[:, N]
    u[:, N + 1] .= u[:, 1]
    return nothing
end
bc_periodic! (generic function with 1 method)
function bc_zero!(u)
    N, M = size(u)
    N = N - 2
    M = M - 2

    u[0, :] .= 0
    u[N + 1, :] .= 0
    u[:, 0] .= 0
    u[:, N + 1] .= 0
    return nothing
end
bc_zero! (generic function with 1 method)
function diffusion!(du, u, (a, Δx, Δy, bc!), _)
    N, M = size(u)
    N = N - 2
    M = M - 2

    # Impose boundary conditions
    bc!(u)

    for i in 1:N
        for j in 1:M
            du[i, j] = a * (
                (u[i + 1, j] - 2 * u[i, j] + u[i - 1, j]) / Δx^2 +
                    (u[i, j + 1] - 2 * u[i, j] + u[i, j - 1]) / Δy^2
            )
        end
    end
    return
end
diffusion! (generic function with 1 method)
HaloVectors = @ingredients(joinpath(dirname(pathof(Ariadne)), "../examples/halovector.jl"))
Main.var"halovector.jl"
import .HaloVectors: HaloVector
Implicit = @ingredients(joinpath(dirname(pathof(Ariadne)), "../examples/implicit.jl"));
import .Implicit: jacobian, solve, G_Euler!
using OffsetArrays
function diffusion!(du::HaloVector, u::HaloVector, p, t)
    return diffusion!(du.data, u.data, p, t)
end
diffusion! (generic function with 2 methods)
begin
    a = 0.01
    N = M = 40
end
40
begin
    Δx = 1 / (N + 1)  # x-grid spacing
    Δy = 1 / (M + 1)  # y-grid spacing

    # TODO: This is a time-step from explicit
    Δt = Δx^2 * Δy^2 / (2.0 * a * (Δx^2 + Δy^2)) # Largest stable time step
end
0.0148720999405116
let
    N = M = 12
    u₀ = HaloVector(OffsetArray(zeros(N + 2, M + 2), 0:(N + 1), 0:(M + 1)))
    jacobian(G_Euler!, diffusion!, u₀, (a, Δx, Δy, bc_zero!), Δt, 0.0)
end
144×144 SparseArrays.SparseMatrixCSC{Float64, Int64} with 672 stored entries:
⎡⠻⣦⡘⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎤
⎢⠲⣌⠱⣦⡙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥
⎢⠀⠈⠳⣌⠻⣦⡙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥
⎢⠀⠀⠀⠈⠳⣌⠛⣤⡙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥
⎢⠀⠀⠀⠀⠀⠈⠳⣌⠻⣦⡙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥
⎢⠀⠀⠀⠀⠀⠀⠀⠈⠳⣌⠻⣦⡙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥
⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣌⠻⣦⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥
⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣄⠻⣦⡙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥
⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣌⠻⣦⡙⢦⡀⠀⠀⠀⠀⠀⠀⠀⎥
⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣌⠻⣦⡙⢦⡀⠀⠀⠀⠀⠀⎥
⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣌⠛⣤⡙⢦⡀⠀⠀⠀⎥
⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣌⠻⣦⡙⢦⡀⠀⎥
⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⣌⠻⢆⡙⠦⎥
⎣⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠳⡌⠻⣦⎦
begin
    xs = 0.0:Δx:1.0
    ys = 0.0:Δy:1.0
end
0.0:0.024390243902439025:1.0
length(xs)
42
function f(x, y)
    return sin(π * x) .* sin(π * y)
end
f (generic function with 1 method)
u₀ = let
    _u₀ = zeros(N + 2, M + 2)
    _u₀ .= f.(xs, ys')
    HaloVector(OffsetArray(_u₀, 0:(N + 1), 0:(M + 1)))
end
1600-element HaloVector{Float64, OffsetMatrix{Float64, Matrix{Float64}}}:
 0.005859788109825737
 0.011685188652980709
 0.01744201586270899
 0.023096486387842192
 0.028615417546937465
 0.03396642205744675
 0.03911809809716897
 ⋮
 0.0339664220574469
 0.028615417546937583
 0.02309648638784227
 0.017442015862709073
 0.01168518865298079
 0.005859788109825782
plot_state(u₀)
function plot_state(u)
    data = u.data.parent
    fig = Figure()
    ax = Axis(fig[1, 1])
    heatmap!(ax, xs, ys, data)

    ax1 = Axis(fig[2, 1])
    lines!(ax1, xs, data[:, M ÷ 2])

    ax2 = Axis(fig[2, 2])
    lines!(ax2, ys, data[N ÷ 2, :])

    return fig
end
plot_state (generic function with 1 method)
begin
    bc! = bc_zero!
    t_end = 10.0

    # TODO: With periodic boundaries -- newton_krylov solve struggles to reach target accuracy.

    # bc! = bc_periodic!
    # t_end = 3*Δt

    ts = 0:Δt:t_end
end
0.0:0.0148720999405116:9.994051160023796
hist = let
    hist = [copy(u₀)]
    solve(G_Euler!, diffusion!, copy(u₀), (a, Δx, Δy, bc!), Δt, ts; callback = (u) -> push!(hist, copy(u)), verbose = 1)
    hist
end
673-element Vector{HaloVector{Float64, OffsetMatrix{Float64, Matrix{Float64}}}}:
 [0.005859788109825737, 0.011685188652980709, 0.01744201586270899, 0.023096486387842192, 0.028615417546937465, 0.03396642205744675, 0.03911809809716897, 0.04404021358262913, 0.048703883582965325, 0.053081739828188784  …  0.053081739828188985, 0.04870388358296551, 0.0440402135826293, 0.03911809809716915, 0.0339664220574469, 0.028615417546937583, 0.02309648638784227, 0.017442015862709073, 0.01168518865298079, 0.005859788109825782]
 [0.005842644628871876, 0.01165100229583603, 0.017390987248510827, 0.02302891497278929, 0.028531699866961954, 0.033867049401183597, 0.03900365362287893, 0.04391136889599416, 0.048561394795839465, 0.05292644312143304  …  0.05292644312143323, 0.04856139479583963, 0.04391136889599432, 0.03900365362287909, 0.03386704940118373, 0.02853169986696206, 0.023028914972789374, 0.017390987248510903, 0.011651002295836098, 0.0058426446288719165]
 [0.00582555130313416, 0.011616915954792889, 0.01734010792436534, 0.022961541245646817, 0.028448227112643356, 0.03376796747100269, 0.038889543969000205, 0.0437829011592403, 0.04841932287597269, 0.05277160075297155  …  0.05277160075297171, 0.04841932287597284, 0.04378290115924044, 0.03888954396900034, 0.03376796747100281, 0.028448227112643446, 0.02296154124564689, 0.017340107924365408, 0.011616915954792948, 0.005825551303134194]
 [0.005808507985877763, 0.011582929337242734, 0.017289377453507396, 0.022894364628056163, 0.028364998567423617, 0.0336691754163516, 0.038775768155977386, 0.043654809269557525, 0.04827766660377213, 0.05261721139358561  …  0.052617211393585744, 0.04827766660377226, 0.04365480926955764, 0.0387757681559775, 0.03366917541635169, 0.028364998567423697, 0.02289436462805623, 0.017289377453507452, 0.011582929337242781, 0.005808507985877793]
 [0.005791514530797149, 0.011549042151433068, 0.017238795400449654, 0.022827384543350763, 0.028282013516841068, 0.033570672389166265, 0.03866232520712082, 0.0435270921273622, 0.04813642476321294, 0.05246327371794531  …  0.05246327371794543, 0.048136424763213045, 0.043527092127362305, 0.03866232520712091, 0.033570672389166355, 0.028282013516841134, 0.02282738454335082, 0.017238795400449706, 0.01154904215143311, 0.005791514530797172]
 [0.005774570792014812, 0.011515254106464949, 0.017188361330978858, 0.02276060041655116, 0.028199271248524265, 0.03347245754386377, 0.03854921414859826, 0.04339974863628764, 0.04799559614182792, 0.05230978640459815  …  0.052309786404598255, 0.047995596141828016, 0.04339974863628773, 0.03854921414859834, 0.03347245754386384, 0.028199271248524324, 0.022760600416551213, 0.0171883613309789, 0.011515254106464985, 0.005774570792014833]
 [0.0057576766240800305, 0.011481564912290503, 0.017138074812152067, 0.022694011674360075, 0.0281167710521859, 0.033374530037334996, 0.03843643400942656, 0.04327277770317473, 0.04785517953069707, 0.052156748135957685  …  0.052156748135957776, 0.04785517953069715, 0.043272777703174815, 0.038436434009426625, 0.033374530037335065, 0.028116771052185956, 0.02269401167436012, 0.017138074812152113, 0.011481564912290533, 0.005757676624080049]
 ⋮
 [0.0008301644255912241, 0.0016554571162327076, 0.0024710349262525534, 0.0032721117207613077, 0.004053986462594607, 0.004812070799868553, 0.0055419159922530355, 0.0062392390179489615, 0.006899947708163234, 0.007520164761582417  …  0.007520164761582417, 0.006899947708163235, 0.006239239017948961, 0.0055419159922530355, 0.004812070799868552, 0.004053986462594607, 0.0032721117207613073, 0.002471034926252553, 0.0016554571162327074, 0.000830164425591224]
 [0.0008277356845255139, 0.0016506138869196061, 0.002463805630687497, 0.003262538783325273, 0.004042126061090984, 0.004797992535849441, 0.005525702482569678, 0.006220985409923196, 0.006879761121225188, 0.007498163658651082  …  0.007498163658651083, 0.006879761121225187, 0.006220985409923197, 0.005525702482569677, 0.004797992535849441, 0.004042126061090983, 0.003262538783325273, 0.0024638056306874965, 0.001650613886919606, 0.0008277356845255138]
 [0.0008253140490198372, 0.0016457848270283219, 0.0024565974852542374, 0.0032529938526136355, 0.0040303003585497555, 0.004783955459403423, 0.005509536407365048, 0.006202785204917416, 0.006859633592458223, 0.007476226922465099  …  0.007476226922465099, 0.0068596335924582235, 0.006202785204917416, 0.005509536407365049, 0.004783955459403422, 0.0040303003585497555, 0.003252993852613635, 0.002456597485254237, 0.0016457848270283219, 0.0008253140490198371]
 [0.0008228994982860658, 0.0016409698951045883, 0.0024494104280756443, 0.0032434768466895136, 0.0040185092534551405, 0.004769959450031533, 0.00549341762786415, 0.00618463824669497, 0.006839564949080899, 0.00745435436471218  …  0.007454354364712181, 0.006839564949080899, 0.006184638246694971, 0.00549341762786415, 0.004769959450031533, 0.00401850925345514, 0.0032434768466895136, 0.0024494104280756443, 0.001640969895104588, 0.0008228994982860658]
 [0.0008204920115968877, 0.0016361690498154194, 0.0024422443974556132, 0.0032339876838557466, 0.004006752644588344, 0.00475600438758735, 0.0054773460056979865, 0.0061665443794763, 0.006819555018817268, 0.007432545797630978  …  0.007432545797630977, 0.006819555018817269, 0.0061665443794763, 0.005477346005697987, 0.00475600438758735, 0.004006752644588344, 0.003233987683855746, 0.0024422443974556132, 0.0016361690498154194, 0.0008204920115968875]
 [0.000818091568285632, 0.001631382249948753, 0.002435099331878541, 0.0032245262826541845, 0.003995030431026704, 0.0047420901522759445, 0.005461321402902379, 0.006148503447937593, 0.006799603629895397, 0.007410801034009451  …  0.007410801034009451, 0.006799603629895396, 0.006148503447937594, 0.005461321402902379, 0.004742090152275945, 0.003995030431026704, 0.0032245262826541845, 0.002435099331878541, 0.0016313822499487527, 0.000818091568285632]
plot_state(hist[i])
@bind i PlutoUI.Slider(1:length(ts))
plot_state(hist[1])
any(isnan, hist[i].data)
false