Improve visualization tools with a video generation (#79)

* add complete_report flag and filter solution by unique time step in create_results * add make_thickness_video function * update jld2 files and bump version * remove complete_report attribute * downgrade version to 0.8 * remove get_glathida_path_and_IDs and one of the implementation of filter_missing_glaciers * add make video test in Sleipnir * put back glaciers2D_plots tests * remove previously added jld2 files
ODINN-SciML · Feb 26, 2025 · b76862b · b76862b
1 parent 0b3536a
commit b76862b
Show file tree

Hide file tree

Showing 19 changed files with 151 additions and 76 deletions.
diff --git a/Project.toml b/Project.toml
@@ -13,6 +13,7 @@ DateFormats = "44557152-fe0a-4de1-8405-416d90313ce6"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+GR = "28b8d3ca-fb5f-59d9-8090-bfdbd6d07a71"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -44,7 +45,7 @@ Revise = "3"
 Statistics = "1"
 Unitful = "1"
 Tar = "1.10.0, 1"
-julia = "1.9"
+julia = "1.10"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

diff --git a/src/Sleipnir.jl b/src/Sleipnir.jl
@@ -25,6 +25,7 @@ using Unitful: m, rad, °
 using CoordRefSystems
 using Dates, DateFormats
 using NetCDF
+using GR
 
 # ##############################################
 # ############    PARAMETERS     ###############

diff --git a/src/glaciers/glacier/glacier2D_utils.jl b/src/glaciers/glacier/glacier2D_utils.jl
@@ -227,56 +227,9 @@ function get_glathida_glacier(glacier::Glacier2D, params::Parameters, force)
     return gtd_grid
 end
 
-function get_glathida_path_and_IDs()
-    gtd_file = Downloads.download("https://cluster.klima.uni-bremen.de/~oggm/glathida/glathida-v3.1.0/data/TTT_per_rgi_id.h5")
-    glathida = h5open(gtd_file, "r")
-    rgi_ids = keys(glathida)
-    rgi_ids = String[id[2:end] for id in rgi_ids]
-    return gtd_file, rgi_ids
-end
 # [End] Glathida Utilities
 
 
-function filter_missing_glaciers!(glaciers::Vector{Glacier2D}, params::Parameters)
-    task_log = CSV.File(joinpath(params.simulation.working_dir, "task_log.csv"))
-    if params.simulation.velocities & params.simulation.use_glathida_data
-        glacier_filter = (task_log.velocity_to_gdir .!= "SUCCESS") .&& (task_log.gridded_attributes .!= "SUCCESS") .&& (task_log.thickness_to_gdir .!= "SUCCESS")
-    elseif params.simulation.use_glathida_data
-        glacier_filter = (task_log.gridded_attributes .!= "SUCCESS") .&& (task_log.thickness_to_gdir .!= "SUCCESS")
-    else
-        glacier_filter = (task_log.gridded_attributes .!= "SUCCESS")
-    end
-
-    glacier_ids = Vector{String}([])
-
-    for id in task_log["index"]
-        push!(glacier_ids, id)
-    end
-    missing_glaciers = glacier_ids[glacier_filter]
-
-    try
-        missing_glaciers_old = load(joinpath(params.simulation.working_dir, "data/missing_glaciers.jld2"))["missing_glaciers"]
-        for missing_glacier in missing_glaciers_old
-            @show missing_glacier
-            if all(missing_glacier .!= missing_glaciers) # if the glacier is already not present, let's add it
-                push!(missing_glaciers, missing_glacier)
-            end
-        end
-    catch error
-        @warn "$error: No missing_glaciers.jld file available. Skipping..."
-    end
-
-    for id in missing_glaciers
-        deleteat!(glaciers, findall(x->x.rgi_id==id, glaciers))
-    end
-
-    # Save missing glaciers in a file
-    jldsave(joinpath(params.simulation.working_dir, "data/missing_glaciers.jld2"); missing_glaciers)
-    #@warn "Filtering out these glaciers from gdir list: $missing_glaciers"
-
-    return missing_glaciers
-end
-
 function filter_missing_glaciers!(rgi_ids::Vector{String}, params::Parameters) # TODO: see if this is necessary, otherwise remove
 
     # Check which glaciers we can actually process

diff --git a/src/simulations/Simulation.jl b/src/simulations/Simulation.jl
@@ -13,4 +13,4 @@ include("results/Results.jl")
 include("simulation_utils.jl")
 include("results/results_utils.jl")
 include("results/results_plotting_utils.jl")
-
+include("results/results_plotting_video_utils.jl")
diff --git a/src/simulations/results/results_plotting_utils.jl b/src/simulations/results/results_plotting_utils.jl
@@ -309,7 +309,7 @@ function plot_glacier_statistics_evolution(results::Results, variables::Vector{S
     for metric in metrics
         if metric == "average"
             if "std" in metrics
-                band!(ax, t, avg_vals .- std_vals, avg_vals .+ std_vals, fillalpha=0.1, label="Std Dev", color=:lightgray)
+                band!(ax, t, avg_vals .- std_vals, avg_vals .+ std_vals, alpha=0.1, label="Std Dev", color=:lightgray)
             end
             lines!(ax, t, avg_vals, label="Average")
         elseif metric == "median"

diff --git a/src/simulations/results/results_plotting_video_utils.jl b/src/simulations/results/results_plotting_video_utils.jl
@@ -0,0 +1,83 @@
+export plot_glacier_vid
+
+function make_thickness_video(H::Vector{Matrix{Float64}}, glacier::Glacier2D, simuparams::SimulationParameters, pathVideo::String; framerate::Int=24, baseTitle::String="")
+    lat = glacier.Coords["lat"]
+    lon = glacier.Coords["lon"]
+    mask = glacier.H₀ .!= 0
+    X, Y = GR.meshgrid(lon,lat)
+
+    # Create a figure
+    fig = CairoMakie.Figure(; size = (600, 600))
+
+    # Create an axis
+    ax = CairoMakie.Axis(fig[1, 1])
+
+    # Number of frames
+    nFrames = size(H,1)
+
+    maxH = maximum(maximum.(H))
+    hm = CairoMakie.heatmap!(ax, reshape(X[mask],:), reshape(Y[mask],:), reshape(H[1][mask],:), colorrange=(0, maxH))
+    CairoMakie.Colorbar(fig[1, 2], hm, label = "Thickness (m)")
+
+    years = simuparams.tspan[1].+simuparams.step*collect(1:size(H,1))
+
+    # Function to update the heatmap for each frame
+    function _update_heatmap(frame_nb)
+        hm[1] = reshape(X[mask],:)
+        hm[2] = reshape(Y[mask],:)
+        hm[3] = reshape(H[frame_nb][mask],:)
+        year = Int(floor(years[frame_nb]))
+        ax.title = baseTitle*" (t = $year)"
+    end
+
+    # Record the animation
+    record(fig, pathVideo, 1:nFrames; framerate = framerate) do frame
+        _update_heatmap(frame)
+    end
+end
+
+
+"""
+    plot_glacier_vid(
+        plot_type::String,
+        H::Vector{Matrix{Float64}},
+        glacier::Glacier2D,
+        simuparams::SimulationParameters,
+        pathVideo::String;
+        framerate::Int=24,
+        baseTitle::String=""
+    )
+
+Generate various types of videos for glacier data.
+
+# Arguments
+- `plot_type`: Type of plot to generate. Options are:
+  * "thickness": Heatmap of the glacier thickness.
+- `H`: A vector of matrices containing the ice thickness over time. This should be
+    replaced by a Results instance in the future once Results no longer depends on
+    an iceflow model.
+- `glacier`: A glacier instance.
+- `simuparams`: The simulation parameters.
+- `pathVideo`: Path of the mp4 file to generate.
+
+# Optional Keyword Arguments
+- `framerate`: The framerate to use for the video generation.
+- `baseTitle`: The prefix to use in the title of the frames. In each frame it is
+    concatenated with the value of the year in the form " (t=XXXX)".
+"""
+function plot_glacier_vid(
+        plot_type::String,
+        H::Vector{Matrix{Float64}},
+        glacier::Glacier2D,
+        simuparams::SimulationParameters,
+        pathVideo::String;
+        framerate::Int=24,
+        baseTitle::String=""
+    )
+
+    if plot_type == "thickness"
+        make_thickness_video(H, glacier, simuparams, pathVideo; framerate=framerate, baseTitle=baseTitle)
+    else
+        error("Invalid plot_type: $plot_type")
+    end
+end
diff --git a/src/simulations/results/results_utils.jl b/src/simulations/results/results_utils.jl
@@ -5,8 +5,14 @@
 
 Store the results of a simulation of a single glacier into a `Results`.
 """
-function create_results(simulation::SIM, glacier_idx::I, solution, loss=nothing; light=false, batch_id::Union{Nothing, I}=nothing) where {SIM <: Simulation, I <: Integer} 
-    H = light ? [solution.u[begin],solution.u[end]] : solution.u 
+function create_results(simulation::SIM, glacier_idx::I, solution, loss=nothing; light=false, batch_id::Union{Nothing, I}=nothing) where {SIM <: Simulation, I <: Integer}
+    # The solution contains all the steps including the intermediate ones
+    # This results in solution having multiple values for a given time step, we select the last one of each time step
+    nSteps = (simulation.parameters.simulation.tspan[2]-simulation.parameters.simulation.tspan[1])/simulation.parameters.simulation.step
+    timeSteps = simulation.parameters.simulation.tspan[1].+collect(0:nSteps).*simulation.parameters.simulation.step
+    solStepIndices = [findlast(==(val), solution.t) for val in timeSteps]
+
+    H = light ? [solution.u[begin],solution.u[end]] : solution.u[solStepIndices]
     # Simulations using Reverse Diff require an iceflow model per glacier
     if isnothing(batch_id)
         iceflow_model = simulation.model.iceflow
@@ -18,15 +24,15 @@ function create_results(simulation::SIM, glacier_idx::I, solution, loss=nothing;
     else
         θ = nothing
     end
-    
+
     results = Results(simulation.glaciers[glacier_idx], iceflow_model;
                       H = H,
                       S = iceflow_model.S,
                       B = simulation.glaciers[glacier_idx].B,
                       V = iceflow_model.V,
                       Vx = iceflow_model.Vx,
                       Vy = iceflow_model.Vy,
-                      Δx = simulation.glaciers[glacier_idx].Δx,              
+                      Δx = simulation.glaciers[glacier_idx].Δx,
                       Δy = simulation.glaciers[glacier_idx].Δy,
                       lon = simulation.glaciers[glacier_idx].cenlon,
                       lat = simulation.glaciers[glacier_idx].cenlat,
@@ -35,28 +41,31 @@ function create_results(simulation::SIM, glacier_idx::I, solution, loss=nothing;
                       tspan = simulation.parameters.simulation.tspan,
                       θ = θ,
                       loss = loss
-                    )              
-                      
+                    )
+
     return results
 end
 
 """
-    save_results_file(simulation::Prediction)
+    save_results_file!(results_list::Vector{Results{F}}, simulation::SIM; path::Union{String,Nothing}=nothing) where {F <: AbstractFloat, SIM <: Simulation}
 
-Save simulation `Results` into a `.jld2` file.
+Save simulation results which are provided as a list of `Results` into a `.jld2` file.
+This function also overrides the `results`` attribute of `simulation`.
 """
 function save_results_file!(results_list::Vector{Results{F, I}}, simulation::SIM; path::Union{String,Nothing}=nothing) where {F <: AbstractFloat, I <: Int, SIM <: Simulation}
     # Create path for simulation results
-    predictions_path = joinpath(dirname(Base.current_project()), "data/results/predictions")
+    if isnothing(path)
+        predictions_path = joinpath(dirname(Base.current_project()), "data/results/predictions")
+    else
+        predictions_path = path
+    end
     if !ispath(predictions_path)
         mkpath(predictions_path)
     end
 
     simulation.results = results_list
 
-    if isnothing(path)
-        tspan = simulation.parameters.simulation.tspan
-        nglaciers = length(simulation.glaciers)
-        jldsave(joinpath(predictions_path, "prediction_$(nglaciers)glaciers_$tspan.jld2"); simulation.results)
-    end
-end
+    tspan = simulation.parameters.simulation.tspan
+    nglaciers = length(simulation.glaciers)
+    jldsave(joinpath(predictions_path, "prediction_$(nglaciers)glaciers_$tspan.jld2"); simulation.results)
+end
diff --git a/test/data/glaciers/glaciers2D_test_temporal.jld2 b/test/data/glaciers/glaciers2D_test_temporal.jld2
diff --git a/test/data/glaciers/glaciers2D_w_glathida.jld2 b/test/data/glaciers/glaciers2D_w_glathida.jld2
diff --git a/test/data/glaciers/glaciers2D_wo_glathida.jld2 b/test/data/glaciers/glaciers2D_wo_glathida.jld2
diff --git a/test/data/params/params_default.jld2 b/test/data/params/params_default.jld2
diff --git a/test/data/params/params_specified.jld2 b/test/data/params/params_specified.jld2
diff --git a/test/data/params/physical_params_default.jld2 b/test/data/params/physical_params_default.jld2
diff --git a/test/data/params/physical_params_specified.jld2 b/test/data/params/physical_params_specified.jld2
diff --git a/test/data/params/simulation_params_default.jld2 b/test/data/params/simulation_params_default.jld2
diff --git a/test/data/params/simulation_params_specified.jld2 b/test/data/params/simulation_params_specified.jld2
diff --git a/test/params_construction.jl b/test/params_construction.jl
@@ -9,7 +9,7 @@ function params_constructor_specified(; save_refs::Bool = false)
     physical_params = PhysicalParameters(ρ = 900.0,
                                         g = 9.81,
                                         ϵ = 1e-3,
-                                        η₀ = 1.0, 
+                                        η₀ = 1.0,
                                         maxA = 8e-17,
                                         minA = 8.5e-20,
                                         maxTlaw = 1.0,

diff --git a/test/plot_utils.jl b/test/plot_utils.jl
@@ -7,7 +7,7 @@ function glaciers2D_plots()
     @testset "plot_glacier tests" begin
         @testset "Heatmaps" begin
             try
-                plot_glacier(results[2], "heatmaps", [:H,:B])
+                plot_glacier(results[1], "heatmaps", [:H,:B])
                 @test true
             catch
                 @test false
@@ -16,7 +16,7 @@ function glaciers2D_plots()
 
         @testset "Statistics Evolution" begin
             try
-                plot_glacier(results[2], "evolution_statistics", [:H], tspan=(2010.0,2015.0), metrics=["average","std","max","median","min"])
+                plot_glacier(results[1], "evolution statistics", [:H], tspan=(2010.0,2015.0), metrics=["average","std","max","median","min"])
                 @test true
             catch
                 @test false
@@ -25,7 +25,7 @@ function glaciers2D_plots()
 
         @testset "Difference Evolution" begin
             try
-                plot_glacier(results[2], "evolution_difference", [:H], tspan=(2010.0,2015.0), metrics=["difference","hist"])
+                plot_glacier(results[1], "evolution difference", [:H], tspan=(2010.0,2015.0), metrics=["difference","hist"])
                 @test true
             catch
                 @test false
@@ -34,7 +34,7 @@ function glaciers2D_plots()
 
         @testset "Integrated Volume" begin
             try
-                plot_glacier(results[2], "integrated_volume", [:H], tspan=(2010.0,2015.0))
+                plot_glacier(results[1], "integrated volume", [:H], tspan=(2010.0,2015.0))
                 @test true
             catch
                 @test false
@@ -43,7 +43,7 @@ function glaciers2D_plots()
 
         @testset "Bias" begin
             try
-                plot_glacier(results[2], "bias", [:B,:S])
+                plot_glacier(results[1], "bias", [:B,:S])
                 @test true
             catch
                 @test false
@@ -53,7 +53,33 @@ function glaciers2D_plots()
 end
 
 
+function make_thickness_video_test()
+    rgi_ids = ["RGI60-11.03646"]
+    rgi_paths = get_rgi_paths()
+    working_dir = joinpath(Sleipnir.root_dir, "test/data")
 
+    params = Parameters(
+        simulation = SimulationParameters(
+            use_MB = true,
+            use_iceflow = true,
+            velocities = true,
+            use_glathida_data = false,
+            tspan = (2014.0, 2015.0),
+            working_dir = working_dir,
+            multiprocessing = true,
+            workers = 1,
+            rgi_paths = rgi_paths,
+            ice_thickness_source = "Farinotti19",
+        ),
+    )
 
+    glaciers = initialize_glaciers(rgi_ids, params)
 
+    nSteps = (params.simulation.tspan[2]-params.simulation.tspan[1])/params.simulation.step
+    timeSteps = params.simulation.tspan[1].+collect(0:nSteps).*params.simulation.step
+    H = [rand(glaciers[1].nx, glaciers[1].ny) for t in timeSteps]
 
+    tempPath = mktempdir()*".mp4"
+
+    plot_glacier_vid("thickness", H, glaciers[1], params.simulation, tempPath; baseTitle="Bossons glacier")
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -17,16 +17,18 @@ include("plot_utils.jl")
 # Activate to avoid GKS backend Plot issues in the JupyterHub
 ENV["GKSwstype"]="nul"
 
-@testset "Run all tests" begin 
+@testset "Run all tests" begin
 
-@testset "Parameters constructors with specified values" params_constructor_specified()
+@testset "Parameters constructors with specified values" params_constructor_specified(save_refs=false)
 
-@testset "Parameters constructors by default" params_constructor_default()
+@testset "Parameters constructors by default" params_constructor_default(save_refs=false)
 
 @testset "Glaciers 2D constructors w/o glathida data" glaciers2D_constructor(use_glathida_data=false, save_refs=false)
 
 @testset "Glaciers 2D constructors w/ glathida data" glaciers2D_constructor(use_glathida_data=true, save_refs=false)
 
-#@testset "Glaciers 2D plots" glaciers2D_plots()
+@testset "Glaciers 2D plots" glaciers2D_plots()
 
-end
+@testset "Video plot test" make_thickness_video_test()
+
+end