In-situ visualization #0: Make Makie available in 2D VisualizationCallback

examples/tree_2d_dgsem/elixir_advection_amr_visualization_makie.jl

@@ -0,0 +1,82 @@
+
+using OrdinaryDiffEq
+using Trixi
+using GLMakie
+
+###############################################################################
+# semidiscretization of the linear advection equation
+
+advection_velocity = (0.2, -0.7)
+equations = LinearScalarAdvectionEquation2D(advection_velocity)
+
+function initial_condition_gauss_largedomain(x, t,
+                                             equation::LinearScalarAdvectionEquation2D)
+    # Store translated coordinate for easy use of exact solution
+    domain_length = SVector(10, 10)
+    x_trans = Trixi.x_trans_periodic_2d(x - equation.advection_velocity * t, domain_length)
+
+    return SVector(exp(-(x_trans[1]^2 + x_trans[2]^2)))
+end
+initial_condition = initial_condition_gauss_largedomain
+
+solver = DGSEM(polydeg = 3, surface_flux = flux_lax_friedrichs)
+
+coordinates_min = (-5.0, -5.0)
+coordinates_max = (5.0, 5.0)
+mesh = TreeMesh(coordinates_min, coordinates_max,
+                initial_refinement_level = 3,
+                n_cells_max = 30_000)
+
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver)
+
+###############################################################################
+# ODE solvers, callbacks etc.
+
+tspan = (0.0, 20.0)
+ode = semidiscretize(semi, tspan)
+
+summary_callback = SummaryCallback()
+
+analysis_interval = 100
+analysis_callback = AnalysisCallback(semi, interval = analysis_interval,
+                                     extra_analysis_integrals = (entropy,))
+
+alive_callback = AliveCallback(analysis_interval = analysis_interval)
+
+save_solution = SaveSolutionCallback(interval = 100,
+                                     save_initial_solution = true,
+                                     save_final_solution = true,
+                                     solution_variables = cons2prim)
+
+# Enable in-situ visualization with a new plot generated every 100 time steps.
+# plot_creator is set to show_plot_makie in order to use Makie. The specific backend is set
+# via `using GLMakie` (see above).
+# Additional keyword arguments, such as colorrange, will get passed to the respective
+# plotting command.
+visualization = VisualizationCallback(interval = 100,
+                                      plot_creator = Trixi.show_plot_makie,
+                                      colorrange = (0.0, 1.0))
+
+amr_controller = ControllerThreeLevel(semi, IndicatorMax(semi, variable = first),
+                                      base_level = 3,
+                                      med_level = 4, med_threshold = 0.1,
+                                      max_level = 5, max_threshold = 0.6)
+amr_callback = AMRCallback(semi, amr_controller,
+                           interval = 5,
+                           adapt_initial_condition = true,
+                           adapt_initial_condition_only_refine = true)
+
+stepsize_callback = StepsizeCallback(cfl = 1.6)
+
+callbacks = CallbackSet(summary_callback,
+                        analysis_callback, alive_callback,
+                        save_solution, visualization,
+                        amr_callback, stepsize_callback);
+
+###############################################################################
+# run the simulation
+
+sol = solve(ode, CarpenterKennedy2N54(williamson_condition = false),
+            dt = 1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
+            save_everystep = false, callback = callbacks);
+summary_callback() # print the timer summary

ext/TrixiMakieExt.jl

@@ -9,7 +9,7 @@ using Trixi
 
 # Use additional symbols that are not exported
 using Trixi: PlotData2DTriangulated, TrixiODESolution, PlotDataSeries, ScalarData, @muladd,
-             wrap_array_native, mesh_equations_solver_cache
+             wrap_array_native, mesh_equations_solver_cache, FigureAndAxes
 
 # Import functions such that they can be extended with new methods
 import Trixi: iplot, iplot!
@@ -136,12 +136,6 @@ end
 # We set the Makie default colormap to match Plots.jl, which uses `:inferno` by default.
 default_Makie_colormap() = :inferno
 
-# convenience struct for editing Makie plots after they're created.
-struct FigureAndAxes{Axes}
-    fig::Makie.Figure
-    axes::Axes
-end
-
 # for "quiet" return arguments to Makie.plot(::TrixiODESolution) and
 # Makie.plot(::PlotData2DTriangulated)
 Base.show(io::IO, fa::FigureAndAxes) = nothing
@@ -412,6 +406,59 @@ function Makie.plot!(fig, pd::PlotData2DTriangulated;
 
     return FigureAndAxes(fig, axes)
 end
-end # @muladd
 
+# converts a single int into a tuple of ints, to get a square arrangement
+# example: f(1) = (1,1) f(2) = (2,1) f(3) = (2,2) f(4) = (1,2)
+function makie_layout_helper(n)
+    if n == 1
+        return (1, 1)
+    end
+    t = makie_layout_helper(n - 1)
+    if t[1] == 1
+        return (t[2] + 1, 1)
+    elseif t[1] > t[2]
+        return (t[1], t[2] + 1)
+    elseif t[2] >= t[1]
+        return (t[1] - 1, t[2])
+    end
+end
+
+function Trixi.show_plot_makie(visualization_callback, plot_data::PlotData2DTriangulated, variable_names;
+                                show_mesh = true, plot_arguments = Dict{Symbol, Any}(),
+                                time = nothing, timestep = nothing)
+    Makie.plot(plot_data, plot_mesh = show_mesh)
+end
+
+function Trixi.show_plot_makie(visualization_callback, plot_data, variable_names;
+                               show_mesh = true, plot_arguments = Dict{Symbol, Any}(),
+                               time = nothing, timestep = nothing)
+    nvars = size(variable_names)[1]
+    if visualization_callback.figure_axes.fig === nothing
+        @info "Creating new Makie figure"
+        fig = Makie.Figure()
+        axes = [Makie.Axis(fig[makie_layout_helper(v)...], aspect = Makie.DataAspect(),
+                           title = variable_names[v])
+                for v in 1:nvars]
+        if show_mesh
+            push!(axes,
+                  Makie.Axis(fig[makie_layout_helper(nvars + 1)...],
+                             aspect = Makie.DataAspect(), title = "mesh"))
+        end
+        visualization_callback.figure_axes = FigureAndAxes(fig, axes)
+        Makie.display(visualization_callback.figure_axes.fig)
+    end
+
+    @unpack axes = visualization_callback.figure_axes
+    for v in 1:nvars
+        Makie.heatmap!(axes[v], plot_data.x, plot_data.y,
+                       permutedims(plot_data.data[v]); plot_arguments...)
+    end
+    if show_mesh
+        empty!(axes[nvars + 1])
+        Makie.lines!(axes[nvars + 1], plot_data.mesh_vertices_x,
+                     plot_data.mesh_vertices_y,
+                     color = :black)
+    end
+end
+end # @muladd
 end

src/callbacks_step/visualization.jl

@@ -5,6 +5,12 @@
 @muladd begin
 #! format: noindent
 
+# convenience struct for editing plots after they're created.
+struct FigureAndAxes{Figure, Axes}
+    fig::Figure
+    axes::Axes
+end
+
 mutable struct VisualizationCallback{SolutionVariables, VariableNames, PlotDataCreator,
                                      PlotCreator}
     interval::Int
@@ -13,6 +19,7 @@ mutable struct VisualizationCallback{SolutionVariables, VariableNames, PlotDataC
     show_mesh::Bool
     plot_data_creator::PlotDataCreator
     plot_creator::PlotCreator
+    figure_axes::FigureAndAxes
     plot_arguments::Dict{Symbol, Any}
 end
 
@@ -99,19 +106,25 @@ function VisualizationCallback(; interval = 0,
                                                    solution_variables, variable_names,
                                                    show_mesh,
                                                    plot_data_creator, plot_creator,
+                                                   FigureAndAxes(nothing, []),
                                                    Dict{Symbol, Any}(plot_arguments))
 
-    # Warn users if they create a visualization callback without having loaded the Plots package
+    # Warn users if they create a visualization callback without having loaded a plotting
+    # package
     #
-    # Note: This warning is added for convenience, as Plots is the only "officially" supported
-    #       visualization package right now. However, in general nothing prevents anyone from using
-    #       other packages such as Makie, Gadfly etc., given that appropriate `plot_creator`s are
-    #       passed. This is also the reason why the visualization callback is not included via
-    #       Requires.jl only when Plots is present.
-    #       In the future, we should update/remove this warning if other plotting packages are
-    #       starting to be used.
-    if !(:Plots in names(@__MODULE__, all = true))
-        @warn "Package `Plots` not loaded but required by `VisualizationCallback` to visualize results"
+    # Note: This warning is added for convenience, as Plots and Makie are currently the
+    # only "officially" supported visualization packages. However, in general nothing
+    # prevents anyone from using other packages, given that appropriate `plot_creator`s are
+    # passed.
+
+    # TODO: rest of this comment?
+    # This is also the reason why the visualization callback is not included via
+    # Requires.jl only when Plots is present.
+    # In the future, we should update/remove this warning if other plotting packages are
+    # starting to be used.
+    if !(:Plots in names(@__MODULE__, all = true)) &&
+       Base.get_extension(Trixi, :TrixiMakieExt) === nothing
+        @warn "Neither `Plots` nor `Makie loaded but required by `VisualizationCallback` to visualize results"
     end
 
     DiscreteCallback(visualization_callback, visualization_callback, # the first one is the condition, the second the affect!
@@ -156,7 +169,7 @@ function (visualization_callback::VisualizationCallback)(integrator)
     end
 
     # Create plot
-    plot_creator(plot_data, variable_names;
+    plot_creator(visualization_callback, plot_data, variable_names;
                  show_mesh = show_mesh, plot_arguments = plot_arguments,
                  time = integrator.t, timestep = integrator.stats.naccept)
 
@@ -166,7 +179,7 @@ function (visualization_callback::VisualizationCallback)(integrator)
 end
 
 """
-    show_plot(plot_data, variable_names;
+    show_plot(visualization_callback, plot_data, variable_names;
               show_mesh=true, plot_arguments=Dict{Symbol,Any}(),
               time=nothing, timestep=nothing)
 
@@ -182,7 +195,7 @@ This function is the default `plot_creator` argument for the [`VisualizationCall
 
 See also: [`VisualizationCallback`](@ref), [`save_plot`](@ref)
 """
-function show_plot(plot_data, variable_names;
+function show_plot(visualization_callback, plot_data, variable_names;
                    show_mesh = true, plot_arguments = Dict{Symbol, Any}(),
                    time = nothing, timestep = nothing)
     # Gather subplots
@@ -218,7 +231,7 @@ function show_plot(plot_data, variable_names;
 end
 
 """
-    save_plot(plot_data, variable_names;
+    save_plot(visualization_callback, plot_data, variable_names;
               show_mesh=true, plot_arguments=Dict{Symbol,Any}(),
               time=nothing, timestep=nothing)
 
@@ -234,7 +247,7 @@ The `timestep` is used in the filename. `time` is currently unused by this funct
 
 See also: [`VisualizationCallback`](@ref), [`show_plot`](@ref)
 """
-function save_plot(plot_data, variable_names;
+function save_plot(visualization_callback, plot_data, variable_names;
                    show_mesh = true, plot_arguments = Dict{Symbol, Any}(),
                    time = nothing, timestep = nothing)
     # Gather subplots
@@ -258,4 +271,8 @@ function save_plot(plot_data, variable_names;
     filename = joinpath("out", @sprintf("solution_%09d.png", timestep))
     Plots.savefig(filename)
 end
+
+# Add definitions of Makie plot functions here such that they can be exported from Trixi.jl
+# and extended in the TrixiMakieExt extension
+function show_plot_makie end
 end # @muladd