Minimal transmural slab with spatially varying diffusion tensor.

examples/transmural_slab_fractionation.jl

@@ -0,0 +1,89 @@
+using Thunderbolt, LinearAlgebra, SparseArrays, UnPack
+import Thunderbolt: AbstractIonicModel
+
+using TimerOutputs, BenchmarkTools
+
+
+import LinearAlgebra: mul!
+
+######################################################
+using StaticArrays
+
+mutable struct IOCallback{IO}
+    io::IO
+end
+
+function (iocb::IOCallback{ParaViewWriter{PVD}})(t, problem::Thunderbolt.SplitProblem, solver_cache) where {PVD}
+    store_timestep!(iocb.io, t, problem.A.dh.grid)
+    Thunderbolt.store_timestep_field!(iocb.io, t, problem.A.dh, solver_cache.A_solver_cache.uₙ, :φₘ)
+    Thunderbolt.store_timestep_field!(iocb.io, t, problem.A.dh, solver_cache.B_solver_cache.sₙ[1:length(solver_cache.A_solver_cache.uₙ)], :s)
+    Thunderbolt.finalize_timestep!(iocb.io, t)
+end
+
+######################################################
+function steady_state_initializer(problem::Thunderbolt.SplitProblem, t₀)
+    # TODO cleaner implementation. We need to extract this from the types or via dispatch.
+    dh = problem.A.dh
+    ionic_model = problem.B.ode
+    default_values = Thunderbolt.default_initial_state(ionic_model)
+    u₀ = ones(ndofs(dh))*default_values[1]
+    s₀ = zeros(ndofs(dh), Thunderbolt.num_states(ionic_model));
+    for i ∈ 1:Thunderbolt.num_states(ionic_model)
+        s₀[:, i] .= default_values[1+i]
+    end
+    for cell in CellIterator(dh)
+        _celldofs = celldofs(cell)
+        ϕₘ_celldofs = _celldofs[dof_range(dh, :ϕₘ)]
+        coordinates = getcoordinates(cell)
+        for (i, (x₁, x₂)) in enumerate(coordinates)
+            if x₁ <= 1.25 && x₂ <= 1.25
+                u₀[ϕₘ_celldofs[i]] = 50.0
+            end
+        end
+    end
+    return u₀, s₀
+end
+
+function varying_tensor_field(x,t)
+    λ₁ = 0.20
+    λ₂ = 0.10
+    λ₃ = 0.05
+    f₀ = Vec((1.0, 0.0, 0.0))
+    s₀ = Vec((0.0, 1.0, 0.0))
+    n₀ = Vec((0.0, 0.0, 1.0))
+
+    return λ₁ * f₀ ⊗ f₀ + λ₂ * s₀ ⊗ s₀ + λ₃ * n₀ ⊗ n₀
+end
+
+model = MonodomainModel(
+    ConstantCoefficient(1.0),
+    ConstantCoefficient(1.0),
+    AnalyticalCoefficient(varying_tensor_field, Lagrange{RefHexahedron,1}()^3),
+    NoStimulationProtocol(),
+    Thunderbolt.PCG2019()
+)
+
+mesh = generate_mesh(Hexahedron, (25, 50, 50), Vec((0.0,0.0,0.0)), Vec((4.0,8.0,8.0)))
+
+problem = semidiscretize(
+    ReactionDiffusionSplit(model),
+    FiniteElementDiscretization(Dict(:φₘ => LagrangeCollection{1}())),
+    mesh
+)
+
+solver = LTGOSSolver(
+    BackwardEulerSolver(),
+    Thunderbolt.ThreadedForwardEulerCellSolver(64)
+)
+
+
+# Idea: We want to solve a semidiscrete problem, with a given compatible solver, on a time interval, with a given initial condition
+# TODO iterator syntax
+solve(
+    problem,
+    solver,
+    0.01,
+    (0.0, 10.0),
+    steady_state_initializer,
+    IOCallback(ParaViewWriter("transmural-wave"))
+)

src/Thunderbolt.jl

@@ -56,6 +56,7 @@ export
     # Coefficients
     ConstantCoefficient,
     FieldCoefficient,
+    AnalyticalCoefficient,
     CalciumHatField,
     # Collections
     LagrangeCollection,

src/mesh/meshes.jl

@@ -148,3 +148,6 @@ elementtypes(::SimpleMesh2D{Quadrilateral}) = @SVector [Quadrilateral]
 @inline Ferrite.getnodes(mesh::Union{SimpleMesh2D,SimpleMesh3D}, setname::String) = Ferrite.getnodes(mesh.grid, setname)
 
 @inline Ferrite.vtk_grid(filename::AbstractString, mesh::Union{SimpleMesh2D,SimpleMesh3D}; kwargs...) = Ferrite.vtk_grid(filename, mesh.grid, kwargs...)
+
+@inline Ferrite.get_coordinate_type(::SimpleMesh2D{C,T}) where {C,T} = Vec{2,T} 
+@inline Ferrite.get_coordinate_type(::SimpleMesh3D{C,T}) where {C,T} = Vec{3,T} 

src/modeling/cells/pcg2019.jl

@@ -45,19 +45,25 @@ Base.@kwdef struct ParametrizedPCG2019Model{T} <: AbstractIonicModel
     E_Ca::T = 50.0  # [mV]
 end
 
-function cell_rhs_fast!(du,φ,s,t,p::P) where {P <: ParametrizedPCG2019Model}
-    sigmoid(φ, E_Y, k_Y, sign) = 1.0 / (1.0 + exp(sign * (φ - E_Y) / k_Y))
+const PCG2019 = ParametrizedPCG2019Model{Float64};
 
+function cell_rhs_fast!(du,φ,state,x,t,p::ParametrizedPCG2019Model{T}) where T
+    sigmoid(φ, E_Y, k_Y, sign) = 1.0 / (1.0 + exp(sign * (φ - E_Y) / k_Y))
+
+    C_m = T(1.0) # TODO pass!
+
     @unpack g_Na, g_K1, g_to, g_CaL, g_Kr, g_Ks = p
-    @unpack E_K, E_h, E_m = p
-    @unpack      k_h, k_m = p
-
-    h  = s[1]
-    m  = s[2]
-    f  = s[3]
-    s  = s[4]
-    xs = s[5]
-    xr = s[6]
+    @unpack E_K, E_Na, E_Ca, E_r, E_d, E_z, E_y, E_h, E_m = p
+    @unpack                  k_r, k_d, k_z, k_y, k_h, k_m = p
+
+    @unpack τ_h0, δ_h, τ_m = p
+
+    h  = state[1]
+    m  = state[2]
+    f  = state[3]
+    s  = state[4]
+    xs = state[5]
+    xr = state[6]
 
     # Instantaneous gates
     r∞ = sigmoid(φ, E_r, k_r, -1.0)
@@ -73,7 +79,7 @@ function cell_rhs_fast!(du,φ,s,t,p::P) where {P <: ParametrizedPCG2019Model}
     I_Kr  = g_Kr * xr * y∞ * (φ - E_K)
     I_Ks  = g_Ks * xs * (φ - E_K)
 
-    I_total = I_Na + I_K1 + I_to + I_CaL + I_Kr + I_Ks + I_stim(t)
+    I_total = I_Na + I_K1 + I_to + I_CaL + I_Kr + I_Ks
 
     du[1] = -I_total/C_m
 
@@ -85,17 +91,17 @@ function cell_rhs_fast!(du,φ,s,t,p::P) where {P <: ParametrizedPCG2019Model}
     du[3] = (m∞-m)/τ_m
 end
 
-function cell_rhs_slow!(du,φ,s,t,p::ParametrizedPCG2019Model)
+function cell_rhs_slow!(du,φ,state,x,t,p::ParametrizedPCG2019Model)
     sigmoid(φ, E_Y, k_Y, sign) = 1.0 / (1.0 + exp(sign * (φ - E_Y) / k_Y))
 
-    @unpack E_K, E_h, E_m, E_s, E_xs, E_xr = p
-    @unpack      k_h, k_m, k_s, k_xs, k_xr = p
+    @unpack E_f, E_s, E_xs, E_xr = p
+    @unpack k_f, k_s, k_xs, k_xr = p
     @unpack τ_f, τ_s, τ_xs, τ_xr = p
 
-    f  = s[3]
-    s  = s[4]
-    xs = s[5]
-    xr = s[6]
+    f  = state[3]
+    s  = state[4]
+    xs = state[5]
+    xr = state[6]
 
     f∞ = sigmoid(φ, E_f, k_f, 1.0)
     du[4] = (f∞-f)/τ_f
@@ -110,6 +116,12 @@ function cell_rhs_slow!(du,φ,s,t,p::ParametrizedPCG2019Model)
     du[7] = (xr∞-xr)/τ_xr
 end
 
+function cell_rhs!(du::TD,φₘ::TV,s::TS,x::TX,t::TT,cell_parameters::TP) where {TD,TV,TS,TX,TT,TP <: ParametrizedPCG2019Model}
+    cell_rhs_fast!(du,φₘ,s,x,t,cell_parameters)
+    cell_rhs_slow!(du,φₘ,s,x,t,cell_parameters)
+    return nothing
+end
+
 function pcg2019_rhs!(du,u,p,t)
     pcg2019_rhs_fast!(du,u,p,t)
     pcg2019_rhs_slow!(du,u,p,t)
@@ -119,8 +131,8 @@ num_states(::ParametrizedPCG2019Model) = 6
 function default_initial_state(p::ParametrizedPCG2019Model{T}) where {T}
     sigmoid(φ, E_Y, k_Y, sign) = 1.0 / (1.0 + exp(sign * (φ - E_Y) / k_Y))
 
-    @unpack E_K, E_h, E_m, E_s, E_xs, E_xr = p
-    @unpack      k_h, k_m, k_s, k_xs, k_xr = p
+    @unpack E_K, E_h, E_m, E_f, E_s, E_xs, E_xr = p
+    @unpack      k_h, k_m, k_f, k_s, k_xs, k_xr = p
 
     u₀ = zeros(T, 7)
     u₀[1] = E_K

src/modeling/coefficients.jl

@@ -48,10 +48,10 @@ struct AnalyticalCoefficient{F, IPG}
     ip_g::IPG #TODO remove this
 end
 
-function Thunderbolt.evaluate_coefficient(coeff::AnalyticalCoefficient{F, sdim}, cell_cache, ξ::Vec{rdim, T}, t::T=T(0.0)) where {F, sdim, rdim, T}
+function Thunderbolt.evaluate_coefficient(coeff::AnalyticalCoefficient{F, <: VectorizedInterpolation{sdim}}, cell_cache, ξ::Vec{rdim, T}, t::T=T(0.0)) where {F, sdim, rdim, T}
     x = zero(Vec{sdim, T})
-    for i in 1:getnbasefunctions(coeff.ip_g)
-        x += shape_value(coeff.ip_g, ξ, i) * cell_cache.coords[i]
+    for i in 1:getnbasefunctions(coeff.ip_g.ip)
+        x += shape_value(coeff.ip_g.ip, ξ, i) * cell_cache.coords[i]
     end
     return coeff.f(x, t)
 end

src/solver/operator_splitting.jl

@@ -57,7 +57,7 @@ function setup_solver_caches(problem::SplitProblem{APT, BPT}, solver::LTGOSSolve
     return cache
 end
 
-function setup_solver_caches(problem::SplitProblem{APT, BPT}, solver::LTGOSSolver{BackwardEulerSolver,BST}, t₀) where {APT,BPT <: TransientHeatProblem,BST}
+function setup_solver_caches(problem::SplitProblem{APT, BPT}, solver::LTGOSSolver{AST,BackwardEulerSolver}, t₀) where {APT,BPT <: TransientHeatProblem,AST}
     cache = LTGOSSolverCache(
         setup_solver_caches(problem.A, solver.A_solver, t₀),
         setup_solver_caches(problem.B, solver.B_solver, t₀),
@@ -80,6 +80,8 @@ transfer_fields!(A, A_cache, B, B_cache)
 
 transfer_fields!(A, A_cache::BackwardEulerSolverCache, B, B_cache::ForwardEulerCellSolverCache) = nothing
 transfer_fields!(A, A_cache::ForwardEulerCellSolverCache, B, B_cache::BackwardEulerSolverCache) = nothing
+transfer_fields!(A, A_cache::BackwardEulerSolverCache, B, B_cache::ThreadedForwardEulerCellSolverCache) = nothing
+transfer_fields!(A, A_cache::ThreadedForwardEulerCellSolverCache, B, B_cache::BackwardEulerSolverCache) = nothing
 
 # TODO what exactly is the job here? How do we know where to write and what to iterate?
 function setup_initial_condition!(problem::SplitProblem{<:Any, <:AbstractPointwiseProblem}, cache, initial_condition, time)

src/solver/partitioned_solver.jl

@@ -107,12 +107,19 @@ function perform_step!(cell_model::ION, t::Float64, Δt::Float64, cache::Adaptiv
 end
 
 # Base.@kwdef struct ThreadedCellSolver{SolverType<:AbstractPointwiseSolver} <: AbstractPointwiseSolver
+#     inner_solver::SolverType
 #     cells_per_thread::Int = 64
 # end
 
-# Base.@kwdef struct ThreadedCellSolverCache{CacheType<:AbstractPointwiseSolverCache} <: AbstractPointwiseSolverCache
+# struct ThreadedCellSolverCache{CacheType<:AbstractPointwiseSolverCache} <: AbstractPointwiseSolverCache
 #     scratch::Vector{CacheType}
-#     cells_per_thread::Int = 64
+# end
+
+# function setup_solver_caches(problem, solver::ThreadedCellSolver{InnerSolverType}, t₀) where {InnerSolverType}
+#     @unpack npoints = problem # TODO what is a good abstraction layer over this?
+#     return ThreadedCellSolverCache(
+#         [setup_solver_caches(problem, solver.inner_solver, t₀) for i ∈ 1:solver.cells_per_thread]
+#     )
 # end
 
 # function perform_step!(uₙ::T1, sₙ::T2, cell_model::ION, t::Float64, Δt::Float64, cache::ThreadedCellSolverCache{CacheType}) where {T1, T2, CacheType, ION <: AbstractIonicModel}
@@ -121,12 +128,79 @@ end
 #         last_cell_in_thread = min((tid+cells_per_thread),length(uₙ))
 #         tuₙ = @view uₙ[tid:last_cell_in_thread]
 #         tsₙ = @view sₙ[tid:last_cell_in_thread]
-#         Threads.@threads for tid ∈ 1:cells_per_thread:length(uₙ)
+#         Threads.@threads :static for tid ∈ 1:cells_per_thread:length(uₙ)
 #             perform_step!(tuₙ, tsₙ, cell_model, t, Δt, tcache)
 #         end
 #     end
 # end
 
+# # TODO better abstraction layer
+# function setup_solver_caches(problem::SplitProblem{APT, BPT}, solver::LTGOSSolver{BackwardEulerSolver,BST}, t₀) where {APT <: TransientHeatProblem,BPT,BST}
+#     cache = LTGOSSolverCache(
+#         setup_solver_caches(problem.A, solver.A_solver, t₀),
+#         setup_solver_caches(problem.B, solver.B_solver, t₀),
+#     )
+#     cache.B_solver_cache.uₙ = cache.A_solver_cache.uₙ
+#     return cache
+# end
+
+# function setup_solver_caches(problem::SplitProblem{APT, BPT}, solver::LTGOSSolver{AST,BackwardEulerSolver}, t₀) where {APT,BPT <: TransientHeatProblem,AST}
+#     cache = LTGOSSolverCache(
+#         setup_solver_caches(problem.A, solver.A_solver, t₀),
+#         setup_solver_caches(problem.B, solver.B_solver, t₀),
+#     )
+#     cache.B_solver_cache.uₙ = cache.A_solver_cache.uₙ
+#     return cache
+# end
+
+# TODO fix the one above somehow
+
+struct ThreadedForwardEulerCellSolver <: AbstractPointwiseSolver
+    num_cells_per_batch::Int
+end
+
+mutable struct ThreadedForwardEulerCellSolverCache{VT, MT} <: AbstractPointwiseSolverCache
+    du::VT
+    uₙ::VT
+    sₙ::MT
+    num_cells_per_batch::Int
+end
+
+function perform_step!(cell_model::ION, t::Float64, Δt::Float64, solver_cache::ThreadedForwardEulerCellSolverCache{VT}) where {VT, ION <: AbstractIonicModel}
+    # Eval buffer
+    @unpack du, uₙ, sₙ = solver_cache
+
+    Threads.@threads :static for j ∈ 1:solver_cache.num_cells_per_batch:length(uₙ)
+        for i ∈ j:min(j+solver_cache.num_cells_per_batch-1, length(uₙ))
+            @inbounds φₘ_cell = uₙ[i]
+            @inbounds s_cell  = @view sₙ[i,:]
+
+            # #TODO get x and Cₘ
+            cell_rhs!(du, φₘ_cell, s_cell, nothing, t, cell_model)
+
+            @inbounds uₙ[i] = φₘ_cell + Δt*du[1]
+
+            # # Non-allocating assignment
+            @inbounds for j ∈ 1:num_states(cell_model)
+                sₙ[i,j] = s_cell[j] + Δt*du[j+1]
+            end
+        end
+    end
+
+    return true
+end
+
+# TODO decouple the concept ForwardEuler from "CellProblem"
+function setup_solver_caches(problem, solver::ThreadedForwardEulerCellSolver, t₀)
+    @unpack npoints = problem # TODO what is a good abstraction layer over this?
+    return ThreadedForwardEulerCellSolverCache(
+        zeros(1+num_states(problem.ode)),
+        zeros(npoints),
+        zeros(npoints, num_states(problem.ode)),
+        solver.num_cells_per_batch
+    )
+end
+
 struct RushLarsenSolver
 end