add dvu_ode vector to custom quantities

Author: LasNikas

examples/fluid/hydrostatic_water_column_2d.jl

@@ -73,7 +73,7 @@ semi = Semidiscretization(fluid_system, boundary_system,
 ode = semidiscretize(semi, tspan)
 
 info_callback = InfoCallback(interval=50)
-saving_callback = SolutionSavingCallback(dt=0.02, prefix="")
+saving_callback = SolutionSavingCallback(dt=0.02, prefix="", dv=TrixiParticles.acceleration)
 
 # This is to easily add a new callback with `trixi_include`
 extra_callback = nothing

src/TrixiParticles.jl

@@ -23,7 +23,8 @@ using ReadVTK: ReadVTK
 using RecipesBase: RecipesBase, @series
 using Random: seed!
 using SciMLBase: CallbackSet, DiscreteCallback, DynamicalODEProblem, u_modified!,
-                 get_tmp_cache, set_proposed_dt!, ODESolution, ODEProblem, terminate!
+                 get_tmp_cache, set_proposed_dt!, ODESolution, ODEProblem, terminate!,
+                 get_du
 @reexport using StaticArrays: SVector
 using StaticArrays: @SMatrix, SMatrix, setindex
 using StrideArrays: PtrArray, StaticInt

src/callbacks/post_process.jl

@@ -229,6 +229,8 @@ end
 # `affect!`
 function (pp::PostprocessCallback)(integrator)
     @trixi_timeit timer() "apply postprocess cb" begin
+        dvu_ode = get_du(integrator)
+        dv_ode, du_ode = dvu_ode.x
         vu_ode = integrator.u
         v_ode, u_ode = vu_ode.x
         semi = integrator.p
@@ -252,7 +254,7 @@ function (pp::PostprocessCallback)(integrator)
             system_index = system_indices(system, semi)
 
             for (key, f) in pp.func
-                result = custom_quantity(f, system, v_ode, u_ode, semi, t)
+                result = custom_quantity(f, system, dv_ode, du_ode, v_ode, u_ode, semi, t)
                 if result !== nothing
                     add_entry!(pp, string(key), t, result, filenames[system_index])
                     new_data = true

src/callbacks/solution_saving.jl

@@ -154,6 +154,7 @@ function (solution_callback::SolutionSavingCallback)(integrator)
     (; interval, output_directory, custom_quantities, write_meta_data, git_hash,
      verbose, prefix, latest_saved_iter, max_coordinates) = solution_callback
 
+    dvu_ode = get_du(integrator)
     vu_ode = integrator.u
     semi = integrator.p
     iter = get_iter(interval, integrator)
@@ -173,7 +174,7 @@ function (solution_callback::SolutionSavingCallback)(integrator)
         println("Writing solution to $output_directory at t = $(integrator.t)")
     end
 
-    @trixi_timeit timer() "save solution" trixi2vtk(vu_ode, semi, integrator.t;
+    @trixi_timeit timer() "save solution" trixi2vtk(dvu_ode, vu_ode, semi, integrator.t;
                                                     iter, output_directory, prefix,
                                                     write_meta_data, git_hash=git_hash[],
                                                     max_coordinates, custom_quantities...)

src/general/custom_quantities.jl

@@ -3,7 +3,7 @@
 
 Returns the total kinetic energy of all particles in a system.
 """
-function kinetic_energy(system, v_ode, u_ode, semi, t)
+function kinetic_energy(system, dv_ode, du_ode, v_ode, u_ode, semi, t)
     v = wrap_v(v_ode, system, semi)
 
     # If `each_moving_particle` is empty (no moving particles), return zero
@@ -18,7 +18,7 @@ end
 
 Returns the total mass of all particles in a system.
 """
-function total_mass(system, v_ode, u_ode, semi, t)
+function total_mass(system, dv_ode, du_ode, v_ode, u_ode, semi, t)
     return sum(eachparticle(system)) do particle
         return system.mass[particle]
     end
@@ -43,7 +43,7 @@ end
 
 Returns the maximum pressure over all particles in a system.
 """
-function max_pressure(system::FluidSystem, v_ode, u_ode, semi, t)
+function max_pressure(system::FluidSystem, dv_ode, du_ode, v_ode, u_ode, semi, t)
     v = wrap_v(v_ode, system, semi)
     return maximum(current_pressure(v, system))
 end
@@ -57,7 +57,7 @@ end
 
 Returns the minimum pressure over all particles in a system.
 """
-function min_pressure(system::FluidSystem, v_ode, u_ode, semi, t)
+function min_pressure(system::FluidSystem, dv_ode, du_ode, v_ode, u_ode, semi, t)
     v = wrap_v(v_ode, system, semi)
     return minimum(current_pressure(v, system))
 end
@@ -71,13 +71,13 @@ end
 
 Returns the average pressure over all particles in a system.
 """
-function avg_pressure(system::FluidSystem, v_ode, u_ode, semi, t)
+function avg_pressure(system::FluidSystem, dv_ode, du_ode, v_ode, u_ode, semi, t)
     v = wrap_v(v_ode, system, semi)
     sum_ = sum(current_pressure(v, system))
     return sum_ / nparticles(system)
 end
 
-function avg_pressure(system, v_ode, u_ode, semi, t)
+function avg_pressure(system, dv_ode, du_ode, v_ode, u_ode, semi, t)
     return NaN
 end
 
@@ -86,12 +86,12 @@ end
 
 Returns the maximum density over all particles in a system.
 """
-function max_density(system::FluidSystem, v_ode, u_ode, semi, t)
+function max_density(system::FluidSystem, dv_ode, du_ode, v_ode, u_ode, semi, t)
     v = wrap_v(v_ode, system, semi)
     return maximum(current_density(v, system))
 end
 
-function max_density(system, v_ode, u_ode, semi, t)
+function max_density(system, dv_ode, du_ode, v_ode, u_ode, semi, t)
     return NaN
 end
 
@@ -100,12 +100,12 @@ end
 
 Returns the minimum density over all particles in a system.
 """
-function min_density(system::FluidSystem, v_ode, u_ode, semi, t)
+function min_density(system::FluidSystem, dv_ode, du_ode, v_ode, u_ode, semi, t)
     v = wrap_v(v_ode, system, semi)
     return minimum(current_density(v, system))
 end
 
-function min_density(system, v_ode, u_ode, semi, t)
+function min_density(system, dv_ode, du_ode, v_ode, u_ode, semi, t)
     return NaN
 end
 
@@ -114,7 +114,7 @@ end
 
 Returns the average_density over all particles in a system.
 """
-function avg_density(system::FluidSystem, v_ode, u_ode, semi, t)
+function avg_density(system::FluidSystem, dv_ode, du_ode, v_ode, u_ode, semi, t)
     v = wrap_v(v_ode, system, semi)
     sum_ = sum(current_density(v, system))
     return sum_ / nparticles(system)
@@ -123,3 +123,7 @@ end
 function avg_density(system, v_ode, u_ode, semi, t)
     return NaN
 end
+
+function acceleration(system, dv_ode, du_ode, v_ode, u_ode, semi, t)
+    return wrap_v(dv_ode, system, semi)
+end

src/io/read_vtk.jl

@@ -57,7 +57,8 @@ function vtk2trixi(file)
             @info "No '$field' field found in VTK file. Will be set to zero."
         end
     end
-    return InitialCondition(; coordinates, particle_spacing=results["particle_spacing"],
+    return InitialCondition(; coordinates,
+                            particle_spacing=first(results["particle_spacing"]),
                             velocity=results["velocity"],
                             mass=results["mass"],
                             density=results["density"],

src/io/write_vtk.jl

@@ -48,15 +48,15 @@ trixi2vtk(sol.u[end], semi, 0.0, iter=1, my_custom_quantity=kinetic_energy)
 
 ```
 """
-function trixi2vtk(vu_ode, semi, t; iter=nothing, output_directory="out", prefix="",
-                   write_meta_data=true, git_hash=compute_git_hash(),
+function trixi2vtk(dvu_ode, vu_ode, semi, t; iter=nothing, output_directory="out",
+                   prefix="", write_meta_data=true, git_hash=compute_git_hash(),
                    max_coordinates=Inf, custom_quantities...)
     (; systems) = semi
-    v_ode, u_ode = vu_ode.x
 
     # Update quantities that are stored in the systems. These quantities (e.g. pressure)
     # still have the values from the last stage of the previous step if not updated here.
     @trixi_timeit timer() "update systems" begin
+        v_ode, u_ode = vu_ode.x
         # Don't create sub-timers here to avoid cluttering the timer output
         @notimeit timer() update_systems_and_nhs(v_ode, u_ode, semi, t;
                                                  update_from_callback=true)
@@ -68,24 +68,26 @@ function trixi2vtk(vu_ode, semi, t; iter=nothing, output_directory="out", prefix
         system_index = system_indices(system, semi)
         periodic_box = get_neighborhood_search(system, semi).periodic_box
 
-        trixi2vtk(system, v_ode, u_ode, semi, t, periodic_box;
+        trixi2vtk(system, dvu_ode, vu_ode, semi, t, periodic_box;
                   system_name=filenames[system_index], output_directory, iter, prefix,
                   write_meta_data, git_hash, max_coordinates, custom_quantities...)
     end
 end
 
 # Convert data for a single TrixiParticle system to VTK format
-function trixi2vtk(system_, v_ode_, u_ode_, semi_, t, periodic_box; output_directory="out",
-                   prefix="", iter=nothing, system_name=vtkname(system_),
-                   write_meta_data=true, max_coordinates=Inf, git_hash=compute_git_hash(),
-                   custom_quantities...)
+function trixi2vtk(system_, dvu_ode_, vu_ode_, semi_, t, periodic_box;
+                   output_directory="out", prefix="", iter=nothing,
+                   system_name=vtkname(system_), write_meta_data=true, max_coordinates=Inf,
+                   git_hash=compute_git_hash(), custom_quantities...)
     mkpath(output_directory)
 
     # Skip empty systems
     if nparticles(system_) == 0
         return
     end
 
+    v_ode_, u_ode_ = vu_ode_.x
+
     # Transfer to CPU if data is on the GPU. Do nothing if already on CPU.
     v_ode, u_ode, system, semi = transfer2cpu(v_ode_, u_ode_, system_, semi_)
 
@@ -137,8 +139,10 @@ function trixi2vtk(system_, v_ode_, u_ode_, semi_, t, periodic_box; output_direc
         end
 
         # Extract custom quantities for this system
+        # TODO: Check if any custom quantities and then adapt from gpu
+        dv_ode, du_ode = dvu_ode_.x
         for (key, quantity) in custom_quantities
-            value = custom_quantity(quantity, system, v_ode, u_ode, semi, t)
+            value = custom_quantity(quantity, system, dv_ode, du_ode, v_ode, u_ode, semi, t)
             if value !== nothing
                 vtk[string(key)] = value
             end
@@ -164,20 +168,21 @@ function transfer2cpu(v_, u_, system_, semi_)
     return v_, u_, system_, semi_
 end
 
-function custom_quantity(quantity::AbstractArray, system, v_ode, u_ode, semi, t)
+function custom_quantity(quantity::AbstractArray, system, dv_ode, du_ode, v_ode, u_ode,
+                         semi, t)
     return quantity
 end
 
-function custom_quantity(quantity, system, v_ode, u_ode, semi, t)
+function custom_quantity(quantity, system, dv_ode, du_ode, v_ode, u_ode, semi, t)
     # Check if `quantity` is a function of `system`, `v_ode`, `u_ode`, `semi` and `t`
     if !isempty(methods(quantity,
-                        (typeof(system), typeof(v_ode), typeof(u_ode),
-                         typeof(semi), typeof(t))))
-        return quantity(system, v_ode, u_ode, semi, t)
+                        (typeof(system), typeof(dv_ode), typeof(du_ode), typeof(v_ode),
+                         typeof(u_ode), typeof(semi), typeof(t))))
+        return quantity(system, dv_ode, du_ode, v_ode, u_ode, semi, t)
     end
 
     # Assume `quantity` is a function of `data`
-    data = system_data(system, v_ode, u_ode, semi)
+    data = system_data(system, dv_ode, du_ode, v_ode, u_ode, semi)
     return quantity(system, data, t)
 end
 
@@ -301,12 +306,12 @@ function write2vtk!(vtk, v, u, t, system::FluidSystem; write_meta_data=true)
             grad_kernel = smoothing_kernel_grad(system, pos_diff, distance, particle)
 
             surface_tension[1:ndims(system),
-                            particle] .+= surface_tension_force(surface_tension_a,
-                                                                surface_tension_b,
-                                                                system, system, particle,
-                                                                neighbor, pos_diff,
-                                                                distance, rho_a, rho_b,
-                                                                grad_kernel)
+            particle] .+= surface_tension_force(surface_tension_a,
+                                                surface_tension_b,
+                                                system, system, particle,
+                                                neighbor, pos_diff,
+                                                distance, rho_a, rho_b,
+                                                grad_kernel)
         end
         vtk["surface_tension"] = surface_tension
 

src/schemes/boundary/open_boundary/system.jl

@@ -499,7 +499,7 @@ end
 # When the neighbor is an open boundary system, just use the viscosity of the fluid `system` instead
 @inline viscosity_model(system, neighbor_system::OpenBoundarySPHSystem) = system.viscosity
 
-function system_data(system::OpenBoundarySPHSystem, v_ode, u_ode, semi)
+function system_data(system::OpenBoundarySPHSystem, dv_ode, du_ode, v_ode, u_ode, semi)
     v = wrap_v(v_ode, system, semi)
     u = wrap_u(u_ode, system, semi)
 

src/schemes/boundary/system.jl

@@ -455,12 +455,12 @@ function system_correction(system::BoundarySPHSystem{<:BoundaryModelDummyParticl
     return system.boundary_model.correction
 end
 
-function system_data(system::BoundarySPHSystem, v_ode, u_ode, semi)
+function system_data(system::BoundarySPHSystem, dv_ode, du_ode, v_ode, u_ode, semi)
     v = wrap_v(v_ode, system, semi)
     u = wrap_u(u_ode, system, semi)
 
     coordinates = current_coordinates(u, system)
-    velocity = current_velocity(v, system)
+    velocity = 0 # TODO this is broken
     density = current_density(v, system)
     pressure = current_pressure(v, system)
 
@@ -471,7 +471,7 @@ function available_data(::BoundarySPHSystem)
     return (:coordinates, :velocity, :density, :pressure)
 end
 
-function system_data(system::BoundaryDEMSystem, v_ode, u_ode, semi)
+function system_data(system::BoundaryDEMSystem, dv_ode, du_ode, v_ode, u_ode, semi)
     (; coordinates, radius, normal_stiffness) = system
 
     return (; coordinates, radius, normal_stiffness)

src/schemes/fluid/fluid.jl

@@ -174,7 +174,7 @@ end
     return nothing
 end
 
-function system_data(system::FluidSystem, v_ode, u_ode, semi)
+function system_data(system::FluidSystem, dv_ode, du_ode, v_ode, u_ode, semi)
     (; mass) = system
 
     v = wrap_v(v_ode, system, semi)