Integration tests for PDE solved added (#18)

The test suite is starting to become a little bit more complete.
ODINN-SciML · Jul 10, 2023 · ce6bee5 · ce6bee5
1 parent 2c61077
commit ce6bee5
Show file tree

Hide file tree

Showing 12 changed files with 187 additions and 36 deletions.
diff --git a/Project.toml b/Project.toml
@@ -19,19 +19,19 @@ Sleipnir = "f5e6c550-199f-11ee-3608-394420200519"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
-[extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Test"]
-
 [compat]
+Infiltrator = "1"
+JLD2 = "0.4"
 OrdinaryDiffEq = "6"
-ProgressMeter = "1"
 PlotThemes = "3"
-Infiltrator = "1"
 Plots = "1"
-Tullio = "0.3"
-JLD2 = "0.4"
+ProgressMeter = "1"
+PyCall = "1"
 Reexport = "1"
-PyCall = "1"
+Tullio = "0.3"
+
+[extras]
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test"]
diff --git a/src/parameters/SolverParameters.jl b/src/parameters/SolverParameters.jl
@@ -38,6 +38,10 @@ function SolverParameters(;
     return solver_parameters
 end
 
+Base.:(==)(a::SolverParameters, b::SolverParameters) = a.solver == b.solver && a.reltol == b.reltol && a.step == b.step && 
+                                      a.tstops == b.tstops && a.save_everystep == b.save_everystep && a.progress == b.progress &&
+                                      a.progress_steps == b.progress_steps
+
 """
 Parameters(;
         simulation::SimulationParameters = SimulationParameters()
@@ -66,7 +70,6 @@ function Parameters(;
     return parameters
 end
 
-# Base.:(==)(a::Parameters, b::Parameters) = a.physical == b.physical && a.simulation == b.simulation && a.OGGM == b.OGGM && a.solver == b.solver
 
 """
     define_callback_steps(tspan::Tuple{Float64, Float64}, step::Float64)

diff --git a/src/simulations/predictions/Prediction.jl b/src/simulations/predictions/Prediction.jl
@@ -21,9 +21,9 @@ Keyword arguments
 """
 function Prediction(
     model::Sleipnir.Model,
-    glaciers::Vector{Sleipnir.AbstractGlacier},
+    glaciers::Vector{G},
     parameters::Sleipnir.Parameters
-    )
+    ) where {G <: Sleipnir.AbstractGlacier}
 
     # Build the results struct based on input values
     prediction = Prediction(model,

diff --git a/src/simulations/predictions/prediction_utils.jl b/src/simulations/predictions/prediction_utils.jl
@@ -23,7 +23,7 @@ function batch_iceflow_PDE(glacier_idx::Int, simulation::Prediction)
     params = simulation.parameters
     glacier = simulation.glaciers[glacier_idx]
 
-    glacier_id = isnothing(glacier.gdir) ? "unnamed" : glacier.gdir
+    glacier_id = isnothing(glacier.gdir) ? "unnamed" : glacier.rgi_id
     println("Processing glacier: ", glacier_id)
 
     # Initialize glacier ice flow model

diff --git a/test/PDE_UDE_solve.jl b/test/PDE_UDE_solve.jl
@@ -0,0 +1,102 @@
+
+
+function pde_solve_test(atol::F, save_refs::Bool = false; MB::Bool=false, fast::Bool=true) where {F <: AbstractFloat}
+    println("PDE solving with MB = $MB")
+    working_dir = joinpath(homedir(), "OGGM/Huginn_tests")
+
+    params = Parameters(OGGM = OGGMparameters(working_dir=working_dir,
+                                              multiprocessing=true,
+                                              workers=2,
+                                              ice_thickness_source = "Farinotti19"),
+                        simulation = SimulationParameters(use_MB=MB,
+                                                        velocities=false,
+                                                        tspan=(2010.0, 2015.0),
+                                                        working_dir = Huginn.root_dir)
+                        ) 
+
+    ## Retrieving gdirs and climate for the following glaciers
+    ## Fast version includes less glacier to reduce the amount of downloaded files and computation time on GitHub CI  
+    if fast
+        rgi_ids = ["RGI60-11.03638", "RGI60-11.01450", "RGI60-08.00213", "RGI60-04.04351", "RGI60-01.02170"]
+    else
+        rgi_ids = ["RGI60-11.03638", "RGI60-11.01450", "RGI60-08.00213", "RGI60-04.04351", "RGI60-01.02170",
+        "RGI60-02.05098", "RGI60-01.01104", "RGI60-01.09162", "RGI60-01.00570", "RGI60-04.07051",                	
+        "RGI60-07.00274", "RGI60-07.01323",  "RGI60-01.17316"]
+    end
+
+    # Load reference values for the simulation
+    if MB
+        PDE_refs = load(joinpath(Huginn.root_dir, "test/data/PDE/PDE_refs_MB.jld2"))["results"]
+    else
+        PDE_refs = load(joinpath(Huginn.root_dir, "test/data/PDE/PDE_refs_noMB.jld2"))["results"]
+    end
+
+    model = Model(iceflow = SIA2Dmodel(params))
+
+    # We retrieve some glaciers for the simulation
+    glaciers = initialize_glaciers(rgi_ids, params)
+
+    # We create an ODINN prediction
+    prediction = Prediction(model, glaciers, params)
+
+    # We run the simulation
+    @time run!(prediction)
+
+    # /!\ Saves current run as reference values
+    if save_refs
+        if MB
+            jldsave(joinpath(Huginn.root_dir, "test/data/PDE/PDE_refs_MB.jld2"); prediction.results)
+        else
+            jldsave(joinpath(Huginn.root_dir, "test/data/PDE/PDE_refs_noMB.jld2"); prediction.results)
+        end
+    end
+
+    # # Run one epoch of the UDE training
+    # θ = zeros(10) # dummy data for the NN
+    # UA_f = zeros(10)
+    # UDE_settings, train_settings = Huginn.get_default_training_settings!(gdirs)
+    # context_batches = Huginn.get_UDE_context(gdirs, tspan; testmode=true, velocities=false)
+    # H_V_preds = @time predict_iceflow(θ, UA_f, gdirs, context_batches, UDE_settings, mb_model) # Array{(H_pred, V̄x_pred, V̄y_pred, rgi_id)}
+
+    let results=prediction.results
+    for result in results
+        let result=result, test_ref=nothing, UDE_pred=nothing
+        for PDE_ref in PDE_refs
+            if result.rgi_id == PDE_ref.rgi_id
+                test_ref = PDE_ref
+            end
+
+            # if result.rgi_id == H_V_pred[4]
+            #     UDE_pred = H_V_pred
+            # end
+        end
+
+        ##############################
+        #### Make plots of errors ####
+        ##############################
+        test_plot_path = joinpath(Huginn.root_dir, "test/plots")
+        if !isdir(test_plot_path)
+            mkdir(test_plot_path)
+        end
+        MB ? vtol = 30.0*atol : vtol = 12.0*atol # a little extra tolerance for UDE surface velocities
+        ### PDE ###
+        plot_test_error(result, test_ref, "H", result.rgi_id, atol, MB)
+        plot_test_error(result, test_ref, "Vx", result.rgi_id, vtol, MB)
+        plot_test_error(result, test_ref, "Vy", result.rgi_id, vtol, MB)
+        ### UDE ###
+        # Huginn.plot_test_error(UDE_pred, test_ref, "H", rgi_id, atol, MB)
+        # Huginn.plot_test_error(UDE_pred, test_ref, "Vx", rgi_id, vtol, MB)
+        # Huginn.plot_test_error(UDE_pred, test_ref, "Vy", rgi_id, vtol, MB)
+
+        # Test that the PDE simulations are correct
+        @test all(isapprox.(result.H[end], test_ref.H[end], atol=atol))
+        @test all(isapprox.(result.Vx, test_ref.Vx, atol=vtol))
+        @test all(isapprox.(result.Vy, test_ref.Vy, atol=vtol))
+        # Test that the UDE simulations are correct
+        # @test all(isapprox.(UDE_pred[1], test_ref["H"], atol=atol))
+        # @test all(isapprox.(UDE_pred[2], test_ref["Vx"], atol=vtol)) 
+        # @test all(isapprox.(UDE_pred[3], test_ref["Vy"], atol=vtol))
+        end # let
+    end
+    end
+end
diff --git a/test/data/PDE/PDE_refs_noMB.jld2 b/test/data/PDE/PDE_refs_noMB.jld2
diff --git a/test/data/params/solver_params_default.jld2 b/test/data/params/solver_params_default.jld2
diff --git a/test/data/params/solver_params_specified.jld2 b/test/data/params/solver_params_specified.jld2
diff --git a/test/halfar.jl b/test/halfar.jl
@@ -1,16 +1,4 @@
 
-using Revise
-using Infiltrator
-using Test 
-using Plots
-using Reexport
-
-# @reexport using Sleipnir 
-using Huginn
-
-# Include utils for halfar 
-include("./utils_test.jl")
-
 """
     unit_halfar_test(; A, t₀, t₁, Δx, Δy, nx, ny, h₀, r₀, rtol=0.02, atol=1.0, distance_to_border=3, save_plot=false)
 
@@ -36,7 +24,8 @@ function unit_halfar_test(; A, t₀, t₁, Δx, Δy, nx, ny, h₀, r₀, rtol=0.
                                                             use_MB=false,
                                                             use_iceflow=true,
                                                             multiprocessing=true,
-                                                            workers=1),
+                                                            workers=1,
+                                                            working_dir=Huginn.root_dir),
                             physical=PhysicalParameters(A=A),
                             solver=SolverParameters(reltol=1e-12))
 

diff --git a/test/params_construction.jl b/test/params_construction.jl
@@ -10,13 +10,25 @@ function params_constructor_specified(save_refs::Bool = false)
                                     progress_steps = 10)
 
     if save_refs
-        jldsave(joinpath(Sleipnir.root_dir, "test/data/params/solver_params.jld2"); solver_params)
+        jldsave(joinpath(Huginn.root_dir, "test/data/params/solver_params_specified.jld2"); solver_params)
     end
 
+    solver_params_ref = load(joinpath(Huginn.root_dir, "test/data/params/solver_params_specified.jld2"))["solver_params"]
+
+    @test solver_params == solver_params_ref
 
 end
 
-function params_constructor_default()
+function params_constructor_default(save_refs::Bool = false)
+
+    solver_params = SolverParameters()
+
+    if save_refs
+        jldsave(joinpath(Huginn.root_dir, "test/data/params/solver_params_default.jld2"); solver_params)
+    end
+
+    solver_params_ref = load(joinpath(Huginn.root_dir, "test/data/params/solver_params_default.jld2"))["solver_params"]
 
+    @test solver_params == solver_params_ref
 
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -6,22 +6,27 @@ using Test
 using JLD2
 using Plots
 using Infiltrator
+using OrdinaryDiffEq
 
 using Huginn
 
-# include("PDE_UDE_solve.jl")
+include("utils_test.jl")
+include("params_construction.jl")
 include("halfar.jl")
+include("PDE_UDE_solve.jl")
 include("mass_conservation.jl")
 
 # Activate to avoid GKS backend Plot issues in the JupyterHub
 ENV["GKSwstype"]="nul"
 
-atol = 0.01
-# @testset "PDE and UDE SIA solvers without MB" pde_solve_test(atol; MB=false, fast=true)
+@testset "Solver parameters construction with specified variables" params_constructor_specified()
 
-atol = 2.0
-# @testset "PDE and UDE SIA solvers with MB" pde_solve_test(atol; MB=true, fast=true)
+@testset "Solver parameters construction with default variables" params_constructor_default()
 
 @testset "Halfar Solution" halfar_test(rtol=0.02, atol=1.0)
 
-@testset "Conservation of Mass - Flat Bed" unit_mass_flatbed_test(rtol=1.0e-7)
+atol = 0.01
+@testset "PDE solving integration tests" pde_solve_test(atol; MB=false, fast=true)
+
+@testset "Conservation of Mass - Flat Bed" unit_mass_flatbed_test(rtol=1.0e-7)
+
diff --git a/test/utils_test.jl b/test/utils_test.jl
@@ -44,3 +44,43 @@ function halfar_solution(R, t, h₀, r₀, physical_parameters::PhysicalParamete
             0.0 
             for r in R]
 end
+
+function access(x::Results, name::String)
+    s = Symbol(name)
+    return getproperty(x, s)
+end
+
+function plot_test_error(pred::Results, ref::Results, variable, rgi_id, atol, MB; path=joinpath(Huginn.root_dir, "test/plots"))
+    @assert (variable == "H") || (variable == "Vx") || (variable == "Vy") "Wrong variable for plots. Needs to be either `H`, `Vx` or `Vy`."
+    if !all(isapprox.(access(pred,variable)[end], access(ref,variable)[end], atol=atol))
+        # @warn "Error found in PDE solve! Check plots in /test/plots⁄"
+        if variable == "H"
+            colour=:ice
+        elseif variable == "Vx" || variable == "Vy"
+            colour=:speed
+        end
+        MB ? tail = "MB" : tail = ""
+        PDE_plot = Plots.heatmap(access(pred,variable)[end] .- access(ref,variable)[end], title="$(variable): PDE simulation - Reference simulation", c=colour)
+        Plots.savefig(PDE_plot,joinpath(path,"$(variable)_PDE_$rgi_id$tail.pdf"))
+    end
+end
+
+function plot_test_error(pred::Tuple, ref::Dict{String, Any}, variable, rgi_id, atol, MB; path=joinpath(Huginn.root_dir, "test/plots"))
+    @assert (variable == "H") || (variable == "Vx") || (variable == "Vy") "Wrong variable for plots. Needs to be either `H`, `Vx` or `Vy`."
+    if variable == "H"
+        idx=1
+        colour=:ice
+    elseif variable == "Vx" 
+        idx=2
+        colour=:speed
+    elseif variable == "Vy"
+        idx=3
+        colour=:speed
+    end
+    if !all(isapprox.(pred[idx], ref[variable], atol=atol))
+        # @warn "Error found in PDE solve! Check plots in /test/plots⁄"
+        UDE_plot = Plots.heatmap(pred[idx] .- ref[variable], title="$(variable): UDE simulation - Reference simulation", c=colour)
+        MB ? tail = "MB" : tail = ""
+        Plots.savefig(UDE_plot,joinpath(path,"$(variable)_UDE_$rgi_id$tail.pdf"))
+    end
+end