add 3rd dimension

matthieugomez · May 28, 2024 · a9caa90 · a9caa90 · matthieugomez · May 28, 2024
1 parent c3e969c
commit a9caa90
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "EconPDEs"
 uuid = "a3315474-fad9-5060-8696-cee5f38a87b7"
-version = "1.0.3"
+version = "1.1.0"
 
 [deps]
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"

diff --git a/src/finiteschemesolve.jl b/src/finiteschemesolve.jl
@@ -13,6 +13,7 @@ function implicit_timestep(G!, ypost, Δ; is_algebraic = fill(false, size(ypost)
             method == :newton
         end
         result = nlsolve(G_helper!, ypost; iterations = iterations, show_trace = verbose, ftol = maxdist, method = method, autodiff = autodiff)
+        zero, residual_norm = result.zero, result.residual_norm
     else
         if autodiff == :forward
             jac_cache = ForwardColorJacCache(G_helper!, deepcopy(ypost); colorvec = colorvec, sparsity = J0)

diff --git a/src/utils.jl b/src/utils.jl
@@ -79,3 +79,42 @@ end
     end
 end
 
+# 3 state variables
+@generated function differentiate(::Type{Tsolution}, grid::StateGrid{T1, 3, <: NamedTuple{N}}, y::AbstractArray{T}, icar, bc) where {Tsolution, T1, N, T}
+    statename1 = N[1]
+    statename2 = N[2]
+    statename3 = N[3]
+    expr = Expr[]
+    for k in 1:length(Tsolution.parameters[1])
+        solname = Tsolution.parameters[1][k]
+        push!(expr, Expr(:(=), solname, :(y[i1, i2, i3, $k])))
+        push!(expr, Expr(:(=), Symbol(solname, statename1, :_up), :((i1 < size(y, 1)) ? (y[i1+1, i2, i3, $k] - y[i1, i2, i3, $k]) / Δx1p : convert($T, bc[end, i2, i3, $k]))))
+        push!(expr, Expr(:(=), Symbol(solname, statename1, :_down), :((i1 > 1) ? (y[i1, i2, i3, $k] - y[i1-1, i2, i3, $k]) / Δx1m : convert($T, bc[1, i2, i3, $k]))))
+        push!(expr, Expr(:(=), Symbol(solname, statename2, :_up), :((i2 < size(y, 2)) ? (y[i1, i2+1, i3, $k] - y[i1, i2, i3, $k]) / Δx2p : convert($T, bc[i1, end, i3, $k]))))
+        push!(expr, Expr(:(=), Symbol(solname, statename2, :_down),  :((i2 > 1) ? (y[i1, i2, i3, $k] - y[i1, i2-1, i3, $k]) / Δx2m : convert($T, bc[i1, 1, i3, $k]))))
+        push!(expr, Expr(:(=), Symbol(solname, statename3, :_up), :((i3 < size(y, 3)) ? (y[i1, i2, i3+1, $k] - y[i1, i2, i3, $k]) / Δx3p : convert($T, bc[i1, i2, end, $k]))))
+        push!(expr, Expr(:(=), Symbol(solname, statename3, :_down),  :((i3 > 1) ? (y[i1, i2, i3, $k] - y[i1, i2, i3-1, $k]) / Δx3m : convert($T, bc[i1, i2, 1, $k]))))
+        push!(expr, Expr(:(=), Symbol(solname, statename1, statename1), :((1 < i1 < size(y, 1)) ? (y[i1 + 1, i2, i3, $k] / (Δx1p * Δx1) + y[i1 - 1, i2, i3, $k] / (Δx1m * Δx1) - 2 * y[i1, i2, i3, $k] / (Δx1p * Δx1m)) : ((i1 == 1) ? (y[2, i2, i3, $k] / (Δx1p * Δx1) + (y[1, i2, i3, $k] - bc[1, i2, $k] * Δx1m) / (Δx1m * Δx1) - 2 * y[1, i2, i3, $k] / (Δx1p * Δx1m)) : ((y[end, i2, i3, $k] + bc[end, i2, i3, $k] * Δx1p) / (Δx1p * Δx1) + y[end - 1, i2, i3, $k] / (Δx1m * Δx1) - 2 * y[end, i2, i3, $k] / (Δx1p * Δx1m))))))
+        push!(expr, Expr(:(=), Symbol(solname, statename2, statename2), :((1 < i2 < size(y, 2)) ? (y[i1, i2 + 1, i3, $k] / (Δx2p * Δx2) + y[i1, i2 - 1, i3, $k] / (Δx2m * Δx2) - 2 * y[i1, i2, i3, $k] / (Δx2p * Δx2m)) : ((i2 == 1) ? (y[i1, 2, i3, $k] / (Δx2p * Δx2) + (y[i1, 1, i3, $k] - bc[i1, 1, i3, $k] * Δx2m) / (Δx2m * Δx2) - 2 * y[i1, 1, i3, $k] / (Δx2p * Δx2m)) : ((y[i1, end, i3, $k] + bc[i1, end, i3, $k] * Δx2p) / (Δx2p * Δx2) + y[i1, end - 1, i3, $k] / (Δx2m * Δx2) - 2 * y[i1, end, i3, $k] / (Δx2p * Δx2m))))))
+        push!(expr, Expr(:(=), Symbol(solname, statename3, statename3), :((1 < i3 < size(y, 3)) ? (y[i1, i2, i3 + 1, $k] / (Δx3p * Δx3) + y[i1, i2, i3 - 1, $k] / (Δx3m * Δx3) - 2 * y[i1, i2, i3, $k] / (Δx3p * Δx3m)) : ((i3 == 1) ? (y[i1, i2, 2, $k] / (Δx3p * Δx3) + (y[i1, i2, 1, $k] - bc[i1, i2, 1, $k] * Δx3m) / (Δx3m * Δx3) - 2 * y[i1, i2, 1, $k] / (Δx3p * Δx3m)) : ((y[i1, i2, end, $k] + bc[i1, i2, end, $k] * Δx3p) / (Δx3p * Δx3) + y[i1, i2, end - 1, $k] / (Δx3m * Δx3) - 2 * y[i1, i2, end, $k] / (Δx3p * Δx3m))))))
+        push!(expr, Expr(:(=), Symbol(solname, statename1, statename2), :((y[min(i1 + 1, size(y, 1)), min(i2 + 1, size(y, 2)), $k] - y[min(i1 + 1, size(y, 1)), max(i2 - 1, 1), $k] - y[max(i1 - 1, 1), min(i2 + 1, size(y, 2)), $k] + y[max(i1 - 1, 1), max(i2 - 1, 1), $k]) / (4 * Δx1 * Δx2))))
+        push!(expr, Expr(:(=), Symbol(solname, statename1, statename3), :((y[min(i1 + 1, size(y, 1)), i2, min(i3 + 1, size(y, 3)), $k] - y[min(i1 + 1, size(y, 1)), i2, max(i3 - 1, 1), $k] - y[max(i1 - 1, 1), i2, min(i3 + 1, size(y, 3)), $k] + y[max(i1 - 1, 1), max(i3 - 1, 1), $k]) / (4 * Δx1 * Δx3))))
+        push!(expr, Expr(:(=), Symbol(solname, statename2, statename3), :((y[i1, min(i2 + 1, size(y, 2)), min(i3 + 1, size(y, 3)), $k] - y[i1, min(i2 + 1, size(y, 2)), max(i3 - 1, 1), $k] - y[i1, max(i2 - 1, 1), min(i3 + 1, size(y, 3)), $k] + y[i1, max(i2 - 1, 1), max(i3 - 1, 1), $k]) / (4 * Δx2 * Δx3))))
+    end
+    quote
+        $(Expr(:meta, :inline))
+        i1, i2, i3 = icar[1], icar[2], icar[3]
+        grid1, grid2, grid3 = grid.x[1], grid.x[2], grid.x[3]
+        @inbounds Δx1m = grid1[max(i1, 2)] - grid1[max(i1-1, 1)]
+        @inbounds Δx1p = grid1[min(i1+1, size(y, 1))] - grid1[min(i1, size(y, 1) - 1)]
+        Δx1 = (Δx1m + Δx1p) / 2
+        @inbounds Δx2m = grid2[max(i2, 2)] - grid2[max(i2-1, 1)]
+        @inbounds Δx2p = grid2[min(i2+1, size(y, 2))] - grid2[min(i2, size(y, 2) - 1)]
+        Δx2 = (Δx2m + Δx2p) / 2
+        @inbounds Δx3m = grid3[max(i3, 2)] - grid3[max(i3-1, 1)]
+        @inbounds Δx3p = grid3[min(i3+1, size(y, 3))] - grid3[min(i3, size(y, 3) - 1)]
+        Δx3 = (Δx3m + Δx3p) / 2
+        @inbounds $(Expr(:tuple, expr...))
+    end
+end
+