[codex] Fix density reinitialization callback

examples/fluid/dam_break_2d.jl

@@ -111,8 +111,10 @@ extra_callback = nothing
 extra_callback2 = nothing
 
 use_reinit = false
+# The callback stores the system index, so it remains valid if `semidiscretize`
+# replaces systems internally.
 density_reinit_cb = use_reinit ?
-                    DensityReinitializationCallback(semi.systems[1], interval=10) :
+                    DensityReinitializationCallback(fluid_system, semi; interval=10) :
                     nothing
 stepsize_callback = StepsizeCallback(cfl=0.9)
 

src/callbacks/density_reinit.jl

@@ -1,45 +1,50 @@
-"""
-    DensityReinitializationCallback(; interval::Integer=0, dt=0.0)
-
-Callback to reinitialize the density field when using [`ContinuityDensity`](@ref) [Panizzo2007](@cite).
-
-# Keywords
-- `interval=0`:              Reinitialize the density every `interval` time steps.
-- `dt`:                      Reinitialize the density in regular intervals of `dt` in terms
-                             of integration time. This callback does not add extra time
-                             steps / `tstops`; instead, reinitialization is triggered at
-                             the first solver step after each `dt` interval has elapsed.
-- `reinit_initial_solution`: Reinitialize the initial solution (default=true)
-"""
-mutable struct DensityReinitializationCallback{I}
+mutable struct DensityReinitializationCallbackAffect{I}
     interval::I
+    system_index::Int
     last_t::Float64
     reinit_initial_solution::Bool
 end
 
-function Base.show(io::IO, cb::DiscreteCallback{<:Any, <:DensityReinitializationCallback})
+function Base.show(io::IO,
+                   cb::DiscreteCallback{<:Any, <:DensityReinitializationCallbackAffect})
     @nospecialize cb # reduce precompilation time
     callback = cb.affect!
     print(io, "DensityReinitializationCallback(interval=", callback.interval,
           ", reinit_initial_solution=", callback.reinit_initial_solution, ")")
 end
 
 function Base.show(io::IO, ::MIME"text/plain",
-                   cb::DiscreteCallback{<:Any, <:DensityReinitializationCallback})
+                   cb::DiscreteCallback{<:Any, <:DensityReinitializationCallbackAffect})
     @nospecialize cb # reduce precompilation time
     if get(io, :compact, false)
         show(io, cb)
     else
         callback = cb.affect!
-        setup = [
-            "interval" => callback.interval,
-            "reinit_initial_solution" => callback.reinit_initial_solution
-        ]
+        setup = Pair{String, Any}["interval" => callback.interval,
+                                  "reinit_initial_solution" => callback.reinit_initial_solution]
         summary_box(io, "DensityReinitializationCallback", setup)
     end
 end
 
-function DensityReinitializationCallback(particle_system; interval::Integer=0, dt=0.0,
+"""
+    DensityReinitializationCallback(system, semi; interval::Integer=0, dt=0.0,
+                                    reinit_initial_solution=true)
+
+Callback to reinitialize the density field when using [`ContinuityDensity`](@ref) [Panizzo2007](@cite).
+
+Pass `system` and the [`Semidiscretization`](@ref) containing it. The callback stores
+the system index and uses the corresponding system from the integrator semidiscretization
+at runtime, which remains valid if [`semidiscretize`](@ref) replaces systems internally.
+
+# Keywords
+- `interval=0`: Reinitialize the density every `interval` time steps.
+- `dt`:         Reinitialize the density in regular intervals of `dt` in terms
+                of integration time. This callback does not add extra time
+                steps / `tstops`; instead, reinitialization is triggered at
+                the first solver step after each `dt` interval has elapsed.
+- `reinit_initial_solution`: Reinitialize the initial solution (default=true).
+"""
+function DensityReinitializationCallback(system, semi; interval::Integer=0, dt=0.0,
                                          reinit_initial_solution=true)
     if dt > 0 && interval > 0
         error("Setting both interval and dt is not supported!")
@@ -49,13 +54,13 @@ function DensityReinitializationCallback(particle_system; interval::Integer=0, d
         interval = Float64(dt)
     end
 
-    if particle_system.density_calculator isa SummationDensity
-        throw(ArgumentError("density reinitialization doesn't provide any advantage for summation density"))
-    end
+    check_density_reinit_system(system)
 
+    system_index = system_indices(system, semi)
     last_t = -Inf
 
-    reinit_cb = DensityReinitializationCallback(interval, last_t, reinit_initial_solution)
+    reinit_cb = DensityReinitializationCallbackAffect(interval, system_index, last_t,
+                                                      reinit_initial_solution)
 
     return DiscreteCallback(reinit_cb, reinit_cb, save_positions=(false, false),
                             initialize=(initialize_reinit_cb!))
@@ -65,11 +70,12 @@ function initialize_reinit_cb!(cb, u, t, integrator)
     initialize_reinit_cb!(cb.affect!, u, t, integrator)
 end
 
-function initialize_reinit_cb!(cb::DensityReinitializationCallback, u, t, integrator)
-    # Reinitialize initial solution
+function initialize_reinit_cb!(cb::DensityReinitializationCallbackAffect, u, t, integrator)
+    semi = integrator.p.semi
+    check_density_reinit_system(current_reinit_system(cb.system_index, semi))
+
     if cb.reinit_initial_solution
         # Update systems to compute quantities like density and pressure.
-        semi = integrator.p.semi
         v_ode, u_ode = u.x
         update_systems_and_nhs(v_ode, u_ode, semi, t)
 
@@ -83,27 +89,68 @@ function initialize_reinit_cb!(cb::DensityReinitializationCallback, u, t, integr
 end
 
 # condition with interval
-function (reinit_callback::DensityReinitializationCallback{Int})(u, t, integrator)
+function (reinit_callback::DensityReinitializationCallbackAffect{<:Integer})(u, t,
+                                                                             integrator)
     (; interval) = reinit_callback
 
     return condition_integrator_interval(integrator, interval, save_final_solution=false)
 end
 
 # condition with dt
-function (reinit_callback::DensityReinitializationCallback)(u, t, integrator)
+function (reinit_callback::DensityReinitializationCallbackAffect)(u, t, integrator)
     (; interval, last_t) = reinit_callback
 
     return (t - last_t) > interval
 end
 
 # affect!
-function (reinit_callback::DensityReinitializationCallback)(integrator)
+function (reinit_callback::DensityReinitializationCallbackAffect)(integrator)
     vu_ode = integrator.u
     semi = integrator.p.semi
 
-    @trixi_timeit timer() "reinit density" reinit_density!(vu_ode, semi)
+    @trixi_timeit timer() "reinit density" reinitialize_density!(reinit_callback, vu_ode,
+                                                                 semi)
 
     reinit_callback.last_t = integrator.t
 
+    # Tell OrdinaryDiffEq that `integrator.u` has been modified
     u_modified!(integrator, true)
+
+    return integrator
+end
+
+function reinitialize_density!(reinit_callback::DensityReinitializationCallbackAffect,
+                               vu_ode, semi)
+    v_ode, u_ode = vu_ode.x
+
+    particle_system = current_reinit_system(reinit_callback.system_index, semi)
+    check_density_reinit_system(particle_system)
+
+    v = wrap_v(v_ode, particle_system, semi)
+    u = wrap_u(u_ode, particle_system, semi)
+
+    reinit_density!(particle_system, v, u, v_ode, u_ode, semi)
+
+    return reinit_callback
+end
+
+function current_reinit_system(system_index, semi)
+    if !(1 <= system_index <= length(semi.systems))
+        throw(ArgumentError("system index $system_index is out of bounds for a " *
+                            "semidiscretization with $(length(semi.systems)) systems"))
+    end
+
+    return semi.systems[system_index]
+end
+
+function check_density_reinit_system(particle_system)
+    if !hasproperty(particle_system, :density_calculator)
+        throw(ArgumentError("density reinitialization requires a system with a density calculator"))
+    end
+
+    if particle_system.density_calculator isa SummationDensity
+        throw(ArgumentError("density reinitialization doesn't provide any advantage for summation density"))
+    end
+
+    return particle_system
 end

test/callbacks/callbacks.jl

@@ -56,6 +56,7 @@
     include("update.jl")
     include("sorting.jl")
     include("solution_saving.jl")
+    include("density_reinit.jl")
     include("steady_state_reached.jl")
     include("mechanical_work_calculator.jl")
 end

test/callbacks/density_reinit.jl

@@ -0,0 +1,160 @@
+@testset verbose=true "DensityReinitializationCallback" begin
+    struct MockDensityReinitSystem
+        density_calculator::Any
+        name::Symbol
+    end
+
+    struct MockNoDensityReinitSystem
+        name::Symbol
+    end
+
+    struct MockDensityReinitIntegrator
+        p::Any
+        u::Any
+        t::Float64
+    end
+
+    density_reinit_calls = Symbol[]
+
+    TrixiParticles.wrap_v(v_ode, system::MockDensityReinitSystem, semi) = (:v, system.name)
+    TrixiParticles.wrap_u(u_ode, system::MockDensityReinitSystem, semi) = (:u, system.name)
+
+    function TrixiParticles.reinit_density!(system::MockDensityReinitSystem, v, u,
+                                            v_ode, u_ode, semi)
+        push!(density_reinit_calls, system.name)
+        return system
+    end
+
+    @testset verbose=true "show" begin
+        system = MockDensityReinitSystem(nothing, :fluid)
+        semi = (; systems=(system,))
+        callback = DensityReinitializationCallback(system, semi; interval=10)
+
+        show_compact = "DensityReinitializationCallback(interval=10, reinit_initial_solution=true)"
+        @test repr(callback) == show_compact
+
+        show_box = """
+        ┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
+        │ DensityReinitializationCallback                                                                  │
+        │ ═══════════════════════════════                                                                  │
+        │ interval: ……………………………………………………… 10                                                               │
+        │ reinit_initial_solution: ……………… true                                                             │
+        └──────────────────────────────────────────────────────────────────────────────────────────────────┘"""
+        @test repr("text/plain", callback) == show_box
+    end
+
+    @testset verbose=true "dt condition" begin
+        system = MockDensityReinitSystem(nothing, :fluid)
+        semi = (; systems=(system,))
+        callback = DensityReinitializationCallback(system, semi; dt=0.1).affect!
+        callback.last_t = 0.0
+
+        @test !callback(nothing, 0.099, nothing)
+        @test !callback(nothing, 0.1, nothing)
+        @test callback(nothing, 0.101, nothing)
+    end
+
+    @testset verbose=true "reinit initial solution" begin
+        system = MockDensityReinitSystem(nothing, :fluid)
+        semi = (; systems=(system,))
+        callback = DensityReinitializationCallback(system, semi; interval=1,
+                                                   reinit_initial_solution=false).affect!
+
+        @test !callback.reinit_initial_solution
+    end
+
+    @testset verbose=true "initialize respects reinit_initial_solution" begin
+        empty!(density_reinit_calls)
+
+        system = MockDensityReinitSystem(nothing, :fluid)
+        vu_ode = (; x=(:v_ode, :u_ode))
+        semi = (; systems=(system,))
+        integrator = MockDensityReinitIntegrator((; semi), vu_ode, 0.0)
+
+        TrixiParticles.get_neighborhood_search(system::MockDensityReinitSystem,
+                                               neighbor::MockDensityReinitSystem,
+                                               semi) = nothing
+        TrixiParticles.update_nhs!(neighborhood_search::Nothing,
+                                   system::MockDensityReinitSystem,
+                                   neighbor::MockDensityReinitSystem,
+                                   u_system, u_neighbor, semi) = nothing
+        TrixiParticles.u_modified!(integrator::MockDensityReinitIntegrator,
+                                   is_modified) = nothing
+
+        callback = DensityReinitializationCallback(system, semi; interval=1,
+                                                   reinit_initial_solution=false).affect!
+        TrixiParticles.initialize_reinit_cb!(callback, vu_ode, 0.0, integrator)
+        @test isempty(density_reinit_calls)
+
+        callback = DensityReinitializationCallback(system, semi; interval=1,
+                                                   reinit_initial_solution=true).affect!
+        TrixiParticles.initialize_reinit_cb!(callback, vu_ode, 0.0, integrator)
+        @test density_reinit_calls == [:fluid]
+    end
+
+    @testset verbose=true "selected system" begin
+        empty!(density_reinit_calls)
+
+        system1 = MockDensityReinitSystem(nothing, :fluid1)
+        system2 = MockDensityReinitSystem(nothing, :fluid2)
+        vu_ode = (; x=(:v_ode, :u_ode))
+        semi = (; systems=(system1, system2))
+        callback = DensityReinitializationCallback(system1, semi; interval=1).affect!
+
+        @test callback.system_index == 1
+
+        TrixiParticles.reinitialize_density!(callback, vu_ode, semi)
+
+        @test density_reinit_calls == [:fluid1]
+    end
+
+    @testset verbose=true "selected semidiscretized system" begin
+        empty!(density_reinit_calls)
+
+        original_system1 = MockDensityReinitSystem(nothing, :original1)
+        original_system2 = MockDensityReinitSystem(nothing, :original2)
+        semidiscretized_system1 = MockDensityReinitSystem(nothing, :semidiscretized1)
+        semidiscretized_system2 = MockDensityReinitSystem(nothing, :semidiscretized2)
+        vu_ode = (; x=(:v_ode, :u_ode))
+        semi = (; systems=(original_system1, original_system2))
+        semi_replaced = (; systems=(semidiscretized_system1, semidiscretized_system2))
+
+        # Simulate the case where `semidiscretize` creates a copy of the system.
+        callback = DensityReinitializationCallback(original_system2, semi;
+                                                   interval=1).affect!
+        @test callback.system_index == 2
+
+        TrixiParticles.reinitialize_density!(callback, vu_ode, semi_replaced)
+        @test density_reinit_calls == [:semidiscretized2]
+
+        @test_throws ArgumentError TrixiParticles.reinitialize_density!(callback, vu_ode,
+                                                                        (; systems=()))
+    end
+
+    @testset verbose=true "system validation" begin
+        empty!(density_reinit_calls)
+
+        system = MockDensityReinitSystem(nothing, :fluid)
+        semi = (; systems=(system,))
+        other_semi = (; systems=(MockDensityReinitSystem(nothing, :other_fluid),))
+
+        @test_throws MethodError DensityReinitializationCallback(; interval=1)
+        @test_throws MethodError DensityReinitializationCallback(system; interval=1)
+        @test_throws ArgumentError DensityReinitializationCallback(MockDensityReinitSystem(SummationDensity(),
+                                                                                           :fluid),
+                                                                   semi;
+                                                                   interval=1)
+        @test_throws ArgumentError DensityReinitializationCallback(MockNoDensityReinitSystem(:boundary),
+                                                                   semi;
+                                                                   interval=1)
+        @test_throws ArgumentError DensityReinitializationCallback(system, other_semi;
+                                                                   interval=1)
+
+        callback = DensityReinitializationCallback(system, semi; interval=1).affect!
+        vu_ode = (; x=(:v_ode, :u_ode))
+
+        @test_throws ArgumentError TrixiParticles.reinitialize_density!(callback, vu_ode,
+                                                                        (;
+                                                                         systems=(MockNoDensityReinitSystem(:boundary),)))
+    end
+end