Mass balance mask fixed + multiprocessing implemented (#38)

* Integration tests for PDE solved added

The test suite is starting to become a little bit more complete.

* Fix merge conflicts

* Muninn added as dependency

Now we can run simulations with MB, run by Muninn. New tests have been added covering those usescases. This means that Huginn should be ready to run simulations for most cases. Now it needs to be extensively tested for different regions and finish polishing the remaining issues.

Project.toml

@@ -1,7 +1,7 @@
 name = "Huginn"
 uuid = "9cb796e9-d0ea-421b-b37d-eb97bc1add55"
 authors = ["Jordi Bolibar <[email protected]>", "Facundo Sapienza <[email protected]>"]
-version = "0.2.1"
+version = "0.3.1"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
@@ -25,7 +25,7 @@ Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 BenchmarkTools = "1"
 Infiltrator = "1"
 JLD2 = "0.4"
-Muninn = "0.2.4"
+Muninn = "0.2.5"
 OrdinaryDiffEq = "6"
 PlotThemes = "3"
 Plots = "1"

src/Huginn.jl

@@ -35,22 +35,23 @@ const global root_plots::String = joinpath(root_dir, "plots")
 # ############ PYTHON LIBRARIES ##############
 # ##############################################
 
-const netCDF4::PyObject = PyNULL()
-const cfg::PyObject = PyNULL()
-const utils::PyObject = PyNULL()
-const workflow::PyObject = PyNULL()
-const tasks::PyObject = PyNULL()
-const global_tasks::PyObject = PyNULL()
-const graphics::PyObject = PyNULL()
-const bedtopo::PyObject = PyNULL()
-const millan22::PyObject = PyNULL()
-const MBsandbox::PyObject = PyNULL()
-const salem::PyObject = PyNULL()
+# We either retrieve the reexported Python libraries from Sleipnir or we start from scratch
+const netCDF4::PyObject = isdefined(Sleipnir, :netCDF4) ? Sleipnir.netCDF4 : PyNULL()
+const cfg::PyObject = isdefined(Sleipnir, :cfg) ? Sleipnir.cfg : PyNULL()
+const utils::PyObject = isdefined(Sleipnir, :utils) ? Sleipnir.utils : PyNULL()
+const workflow::PyObject = isdefined(Sleipnir, :workflow) ? Sleipnir.workflow : PyNULL()
+const tasks::PyObject = isdefined(Sleipnir, :tasks) ? Sleipnir.tasks : PyNULL()
+const global_tasks::PyObject = isdefined(Sleipnir, :global_tasks) ? Sleipnir.global_tasks : PyNULL()
+const graphics::PyObject = isdefined(Sleipnir, :graphics) ? Sleipnir.graphics : PyNULL()
+const bedtopo::PyObject = isdefined(Sleipnir, :bedtopo) ? Sleipnir.bedtopo : PyNULL()
+const millan22::PyObject = isdefined(Sleipnir, :millan22) ? Sleipnir.millan22 : PyNULL()
+const MBsandbox::PyObject = isdefined(Sleipnir, :MBsandbox) ? Sleipnir.MBsandbox : PyNULL()
+const salem::PyObject = isdefined(Sleipnir, :salem) ? Sleipnir.salem : PyNULL()
 
 # Essential Python libraries
-const xr::PyObject = PyNULL()
-const rioxarray::PyObject = PyNULL()
-const pd::PyObject = PyNULL()
+const xr::PyObject = isdefined(Sleipnir, :xr) ? Sleipnir.xr : PyNULL()
+const rioxarray::PyObject = isdefined(Sleipnir, :rioxarray) ? Sleipnir.rioxarray : PyNULL()
+const pd::PyObject = isdefined(Sleipnir, :pd) ? Sleipnir.pd : PyNULL()
 
 # ##############################################
 # ############ HUGINN LIBRARIES ##############

src/models/iceflow/SIA2D/SIA2D.jl

@@ -39,8 +39,8 @@ mutable struct SIA2Dmodel{F <: AbstractFloat, I <: Integer} <: SIAmodel
 end
 
 function SIA2Dmodel(params::Sleipnir.Parameters;
-                    A::Union{Ref{F}, Nothing} = nothing,
-                    n::Union{Ref{F}, Nothing} = nothing,
+                    A::Union{F, Nothing} = nothing,
+                    n::Union{F, Nothing} = nothing,
                     H₀::Union{Matrix{F}, Nothing} = nothing,
                     H::Union{Matrix{F}, Nothing} = nothing,
                     H̄::Union{Matrix{F}, Nothing} = nothing,
@@ -62,14 +62,27 @@ function SIA2Dmodel(params::Sleipnir.Parameters;
                     V::Union{Matrix{F}, Nothing} = nothing,
                     Vx::Union{Matrix{F}, Nothing} = nothing,
                     Vy::Union{Matrix{F}, Nothing} = nothing,
-                    Γ::Union{Ref{F}, Nothing} = nothing,
+                    Γ::Union{F, Nothing} = nothing,
                     MB::Union{Matrix{F}, Nothing} = nothing,
                     MB_mask::Union{BitMatrix, Nothing} = nothing,
                     MB_total::Union{Matrix{F}, Nothing} = nothing,
-                    glacier_idx::Union{Ref{I}, Nothing} = nothing) where {F <: AbstractFloat, I <: Integer}
+                    glacier_idx::Union{I, Nothing} = nothing) where {F <: AbstractFloat, I <: Integer}
 
     ft = params.simulation.float_type
     it = params.simulation.int_type
+    if !isnothing(A)
+        A = Ref{F}(A)
+    end
+    if !isnothing(n)
+        n = Ref{F}(n)
+    end
+    if !isnothing(Γ)
+        Γ = Ref{F}(Γ)
+    end
+    if !isnothing(glacier_idx)
+        glacier_idx = Ref{I}(glacier_idx)
+    end
+
     SIA2D_model = SIA2Dmodel{ft,it}(A, n, H₀, H, H̄, S, dSdx, dSdy, D, Dx, Dy, dSdx_edges, dSdy_edges,
                             ∇S, ∇Sx, ∇Sy, Fx, Fy, Fxx, Fyy, V, Vx, Vy, Γ, MB, MB_mask, MB_total, glacier_idx)
 
@@ -99,8 +112,8 @@ function initialize_iceflow_model!(iceflow_model::IF,
                                    ) where {IF <: IceflowModel, I <: Int, G <: Sleipnir.AbstractGlacier}
     nx, ny = glacier.nx, glacier.ny
     F = params.simulation.float_type
-    iceflow_model.A = Ref{F}(glacier.A)
-    iceflow_model.n = Ref{F}(glacier.n)
+    iceflow_model.A = isnothing(iceflow_model.A) ? Ref{F}(glacier.A) : iceflow_model.A
+    iceflow_model.n = isnothing(iceflow_model.n) ? Ref{F}(glacier.n) : iceflow_model.n
     iceflow_model.H₀ = deepcopy(glacier.H₀)::Matrix{F}
     iceflow_model.H = deepcopy(glacier.H₀)::Matrix{F}
     iceflow_model.H̄ = zeros(F,nx-1,ny-1)
@@ -122,7 +135,7 @@ function initialize_iceflow_model!(iceflow_model::IF,
     iceflow_model.V = zeros(F,nx-1,ny-1)
     iceflow_model.Vx = zeros(F,nx-1,ny-1)
     iceflow_model.Vy = zeros(F,nx-1,ny-1)
-    iceflow_model.Γ = Ref{F}(0.0)
+    iceflow_model.Γ = isnothing(iceflow_model.Γ) ? Ref{F}(0.0) : iceflow_model.Γ
     iceflow_model.MB = zeros(F,nx,ny)
     iceflow_model.MB_mask= zeros(F,nx,ny)
     iceflow_model.MB_total = zeros(F,nx,ny)

src/setup/config.jl

@@ -6,24 +6,25 @@ function __init__()
         mkpath(OGGM_path)
     end
 
-    println("Initializing Python libraries...")
-
-    # Load Python packages
     try
-        copy!(netCDF4, pyimport("netCDF4"))
-        copy!(cfg, pyimport("oggm.cfg"))
-        copy!(utils, pyimport("oggm.utils"))
-        copy!(workflow, pyimport("oggm.workflow"))
-        copy!(tasks, pyimport("oggm.tasks"))
-        copy!(global_tasks, pyimport("oggm.global_tasks"))
-        copy!(graphics, pyimport("oggm.graphics"))
-        copy!(bedtopo, pyimport("oggm.shop.bedtopo"))
-        copy!(millan22, pyimport("oggm.shop.millan22"))
-        copy!(MBsandbox, pyimport("MBsandbox.mbmod_daily_oneflowline"))
-        copy!(salem, pyimport("salem"))
-        copy!(pd, pyimport("pandas"))
-        copy!(xr, pyimport("xarray"))
-        copy!(rioxarray, pyimport("rioxarray"))
+        # Only load Python packages if not previously loaded by Sleipnir
+        if cfg == PyNULL() && workflow == PyNULL() && utils == PyNULL() && MBsandbox == PyNULL() 
+            println("Initializing Python libraries in Huginn...")
+            copy!(netCDF4, pyimport("netCDF4"))
+            copy!(cfg, pyimport("oggm.cfg"))
+            copy!(utils, pyimport("oggm.utils"))
+            copy!(workflow, pyimport("oggm.workflow"))
+            copy!(tasks, pyimport("oggm.tasks"))
+            copy!(global_tasks, pyimport("oggm.global_tasks"))
+            copy!(graphics, pyimport("oggm.graphics"))
+            copy!(bedtopo, pyimport("oggm.shop.bedtopo"))
+            copy!(millan22, pyimport("oggm.shop.millan22"))
+            copy!(MBsandbox, pyimport("MBsandbox.mbmod_daily_oneflowline"))
+            copy!(salem, pyimport("salem"))
+            copy!(pd, pyimport("pandas"))
+            copy!(xr, pyimport("xarray"))
+            copy!(rioxarray, pyimport("rioxarray"))
+        end
     catch e
         @warn "It looks like you have not installed and/or activated the virtual Python environment. \n 
         Please follow the guidelines in: https://github.com/ODINN-SciML/ODINN.jl#readme"
@@ -38,3 +39,24 @@ function clean()
     end
     exit()
  end
+
+ function enable_multiprocessing(procs::Int)
+    if procs > 0
+        if nprocs() < procs
+            @eval begin
+            addprocs($procs - nprocs(); exeflags="--project")
+            println("Number of cores: ", nprocs())
+            println("Number of workers: ", nworkers())
+            @everywhere using Huginn
+            end # @eval
+        elseif nprocs() != procs && procs == 1
+            @eval begin
+            rmprocs(workers(), waitfor=0)
+            println("Number of cores: ", nprocs())
+            println("Number of workers: ", nworkers())
+            end # @eval
+        end
+    end
+    return nworkers()
+end
+

src/simulations/predictions/prediction_utils.jl

@@ -8,6 +8,8 @@ In-place run of the model.
 """
 function run!(simulation::Prediction)
 
+    enable_multiprocessing(simulation.parameters.simulation.workers)
+
     println("Running forward in-place PDE ice flow model...\n")
     results_list = @showprogress pmap((glacier_idx) -> batch_iceflow_PDE!(glacier_idx, simulation), 1:length(simulation.glaciers))
 
@@ -209,7 +211,7 @@ end
 function apply_MB_mask!(H::Matrix{F}, glacier::G, ifm::IceflowModel) where {F <: AbstractFloat, G <: Sleipnir.AbstractGlacier}
     # Appy MB only over ice, and avoid applying it to the borders in the accummulation area to avoid overflow
     MB::Matrix{F}, MB_mask::BitMatrix, MB_total::Matrix{F} = ifm.MB, ifm.MB_mask, ifm.MB_total
-    MB_mask .= ((H .> 0.0) .&& (MB .< 0.0)) .|| ((H .> 0.0) .&& (glacier.dist_border .> 1.0) .&& (MB .>= 0.0)) 
+    MB_mask .= ((H .> 0.0) .&& (MB .< 0.0)) .|| ((H .> 10.0) .&& (MB .>= 0.0)) 
     H[MB_mask] .+= MB[MB_mask]
     MB_total[MB_mask] .+= MB[MB_mask]
 end

test/PDE_UDE_solve.jl

@@ -72,10 +72,6 @@ function pde_solve_test(; rtol::F, atol::F, save_refs::Bool=false, MB::Bool=fals
             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
 
         ##############################

test/runtests.jl

@@ -31,7 +31,7 @@ ENV["GKSwstype"]="nul"
 
 @testset "PDE solving integration tests w/ MB" pde_solve_test(; rtol=0.01, atol=0.01, save_refs=false, MB=true, fast=true)
 
-@testset "Run TI models in place" TI_run_test!()
+@testset "Run TI models in place" TI_run_test!(false)
 
 @testset "Conservation of Mass - Flat Bed" unit_mass_flatbed_test(; rtol=1.0e-7)