Skip to content

Improve plot recipes #862

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
efaulhaber merged 2 commits into from
Jul 16, 2025
Merged

Improve plot recipes #862

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1d0a49c
Improve plot recipes
efaulhaber Jul 16, 2025
0b58805
Merge branch 'main' into improve-plot-recipes
efaulhaber Jul 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 32 additions & 42 deletions src/visualization/recipes_plots.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,39 +4,31 @@ const TrixiParticlesODESolution = ODESolution{<:Any, <:Any, <:Any, <:Any, <:Any,
<:ODEProblem{<:Any, <:Any, <:Any,
<:Semidiscretization}}

# This is the main recipe
RecipesBase.@recipe function f(sol::TrixiParticlesODESolution)
# Redirect everything to the next recipe
return sol.u[end].x..., sol.prob.p
end

# GPU version
RecipesBase.@recipe function f(sol::TrixiParticlesODESolution, semi::Semidiscretization)
RecipesBase.@recipe function f(v_ode::AbstractGPUArray, u_ode, semi::Semidiscretization)
# Move GPU data to the CPU
v_ode = Array(sol.u[end].x[1])
u_ode = Array(sol.u[end].x[2])
semi_ = Adapt.adapt(Array, sol.prob.p)
v_ode_ = Array(v_ode)
u_ode_ = Array(u_ode)
semi_ = Adapt.adapt(Array, semi)

# Redirect everything to the next recipe
return v_ode, u_ode, semi_, particle_spacings(semi)
end

RecipesBase.@recipe function f(v_ode::AbstractGPUArray, u_ode, semi::Semidiscretization;
particle_spacings=nothing,
size=(600, 400), # Default size
xlims=(-Inf, Inf), ylims=(-Inf, Inf))
throw(ArgumentError("to plot GPU data, use `plot(sol, semi)`"))
return v_ode_, u_ode_, semi_
end

RecipesBase.@recipe function f(v_ode, u_ode, semi::Semidiscretization;
particle_spacings=TrixiParticles.particle_spacings(semi),
size=(600, 400), # Default size
xlims=(-Inf, Inf), ylims=(-Inf, Inf))
return v_ode, u_ode, semi, particle_spacings
end

RecipesBase.@recipe function f(v_ode, u_ode, semi::Semidiscretization, particle_spacings;
size=(600, 400), # Default size
xlims=(-Inf, Inf), ylims=(-Inf, Inf))
# We need to split this in two recipes in order to find the minimum and maximum
# coordinates across all systems.
# In this first recipe, we collect the data for each system,
# and then pass it to the next recipe.
systems_data = map(enumerate(semi.systems)) do (i, system)
u = wrap_u(u_ode, system, semi)
periodic_box = get_neighborhood_search(system, semi).periodic_box
Expand All @@ -54,27 +46,20 @@ RecipesBase.@recipe function f(v_ode, u_ode, semi::Semidiscretization, particle_
x_min, y_min = minimum(coordinates, dims=2) .- 0.5particle_spacing
x_max, y_max = maximum(coordinates, dims=2) .+ 0.5particle_spacing

# `x_min`, `x_max`, etc. are used to automatically set the marker size.
# When `xlims` or `ylims` are passed explicitly, we have to update these to get the correct marker size.
isfinite(first(xlims)) && (x_min = xlims[1])
isfinite(last(xlims)) && (x_max = xlims[2])

isfinite(first(ylims)) && (y_min = ylims[1])
isfinite(last(ylims)) && (y_max = ylims[2])

return (; x, y, x_min, x_max, y_min, y_max, particle_spacing,
label=timer_name(system))
end

# Pass the semidiscretization and the collected data to the next recipe
return (semi, systems_data...)
end

function particle_spacings(semi::Semidiscretization)
return [system.initial_condition.particle_spacing for system in semi.systems]
end

RecipesBase.@recipe function f((initial_conditions::InitialCondition)...;
xlims=(Inf, Inf), ylims=(Inf, Inf))
RecipesBase.@recipe function f(initial_conditions::InitialCondition...)
# This recipe is similar to the one above for the semidiscretization
idx = 0
ics = map(initial_conditions) do ic
x = collect(ic.coordinates[1, :])
Expand All @@ -88,34 +73,32 @@ RecipesBase.@recipe function f((initial_conditions::InitialCondition)...;
x_min, y_min = minimum(ic.coordinates, dims=2) .- 0.5particle_spacing
x_max, y_max = maximum(ic.coordinates, dims=2) .+ 0.5particle_spacing

# `x_min`, `x_max`, etc. are used to automatically set the marker size.
# When `xlims` or `ylims` are passed explicitly, we have to update these to get the correct marker size.
isfinite(first(xlims)) && (x_min = xlims[1])
isfinite(last(xlims)) && (x_max = xlims[2])

isfinite(first(ylims)) && (y_min = ylims[1])
isfinite(last(ylims)) && (y_max = ylims[2])

idx += 1

return (; x, y, x_min, x_max, y_min, y_max, particle_spacing,
label="initial condition " * "$idx")
end

# Pass the first initial condition and the collected data to the next recipe
return (first(initial_conditions), ics...)
end

RecipesBase.@recipe function f(::Union{InitialCondition, Semidiscretization},
data...; zcolor=nothing, size=(600, 400), colorbar_title="",
xlims=(Inf, Inf), ylims=(Inf, Inf))
data...; size=(600, 400), xlims=(Inf, Inf), ylims=(Inf, Inf))
# `data` is a tuple of named tuples, passed from the recipe above.
# Each named tuple contains coordinates and metadata for a system or initial condition.
#
# First find the minimum and maximum coordinates across all systems.
x_min = minimum(obj.x_min for obj in data)
x_max = maximum(obj.x_max for obj in data)

y_min = minimum(obj.y_min for obj in data)
y_max = maximum(obj.y_max for obj in data)

# `x_min`, `x_max`, etc. are used to automatically set the marker size.
# When `xlims` or `ylims` are passed explicitly, we have to update these to get the correct marker size.
# So they need to be the minimum and maximum coordinates of the plot area.
# When `xlims` or `ylims` are passed explicitly, we have to update these
# to get the correct marker size.
isfinite(first(xlims)) && (x_min = xlims[1])
isfinite(last(xlims)) && (x_max = xlims[2])

Expand All @@ -130,14 +113,21 @@ RecipesBase.@recipe function f(::Union{InitialCondition, Semidiscretization},

xlims --> (x_min, x_max)
ylims --> (y_min, y_max)
aspect_ratio --> :equal

# Just having the kwargs `xlims` and `ylims` (without setting them)
# is enough to make `widen = :auto` fall back to no widening.
# When no explicit limits are passed, we overwrite this.
if all(!isfinite, xlims) && all(!isfinite, ylims)
widen --> true
end

aspect_ratio --> :equal
seriestype --> :scatter
# No border around the markers
markerstrokewidth --> 0
grid --> false
colorbar_title --> colorbar_title
zcolor --> zcolor

# Now plot all systems or initial conditions
for obj in data
@series begin
if obj.particle_spacing < eps()
Expand Down
Loading