UDEs working with experiment with different MB (#122)

Fixed a few bugs in the UDE workflow. The UDEs work suprisingly well with a different MB model. This means that using surface velocities is really a well proxy, which is robust to changes in surface elevation.

Project.toml

@@ -7,6 +7,7 @@ version = "0.1.0"
 AbbreviatedStackTraces = "ac637c84-cc71-43bf-9c33-c1b4316be3d4"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
@@ -38,7 +39,6 @@ SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 AbbreviatedStackTraces = "0.1.7"
@@ -66,10 +66,10 @@ PyCall = "1.9"
 PyPlot = "2.11"
 Revise = "3.1"
 SciMLSensitivity = "7.20"
+ChainRules = "1.50"
 SnoopPrecompile = "1.0.3"
 TimerOutputs = "0.5.22"
 Tullio = "0.3"
-Zygote = "0.6"
 julia = "1.7"
 
 [extras]

scripts/dhdt_plots.jl

@@ -8,15 +8,17 @@ import ODINN: fillZeros
 function make_plots()
 
  # plot_type = "only_H" # plot final H
- # plot_type = "MB_diff" # differences between runs with different MB models
- plot_type = "H_diff" # H - H₀
+ plot_type = "MB_diff" # differences between runs with different MB models
+ # plot_type = "H_diff" # H - H₀
  tspan = (2010.0, 2015.0) # period in years for simulation
 
  root_dir = dirname(Base.current_project())
 
  # Load forward simulations with different surface MB
  grefs = load(joinpath(root_dir, "data/gdir_refs_$tspan.jld2"))["gdir_refs"]
- grefs_MBu1 = load(joinpath(root_dir, "data/gdir_refs_updatedMB1.jld2"))["gdir_refs"]
+ grefs_MBu1 = load(joinpath(root_dir, "data/gdir_refs_$(tspan)_MB1.jld2"))["gdir_refs"]
+ grefs_MBu2 = load(joinpath(root_dir, "data/gdir_refs_$(tspan)_MB2.jld2"))["gdir_refs"]
+ grefs_MBu3 = load(joinpath(root_dir, "data/gdir_refs_$(tspan)_MB3.jld2"))["gdir_refs"]
 
  n=4
  m=3
@@ -33,16 +35,15 @@ function make_plots()
  H = reverse(grefs[i]["H"]', dims=2)
  H₀ = reverse(grefs[i]["H₀"]', dims=2)
  H_MBu1 = reverse(grefs_MBu1[i]["H"]', dims=2)
- # H = reverse(grefs[i]["H"])
- # H_MBu1 = reverse(grefs_MBu1[i]["H"])
+ H_MBu2 = reverse(grefs_MBu3[i]["H"]', dims=2)
  if plot_type == "only_H"
  H_plot = H
  label = "Predicted H (m)"
  elseif plot_type == "H_diff"
  H_plot = H .- H₀
  label = "H - H₀ (m)"
  elseif plot_type == "MB_diff"
- H_plot = H .- H_MBu1
+ H_plot = H_MBu1 .- H_MBu2
  label="Surface mass balance difference (m)"
  end
  push!(MBdiffs, H_plot)
@@ -66,35 +67,6 @@ function make_plots()
 
  end # let
 
- # hms = []
- # for (gref, gref_MBu1) in zip(grefs, grefs_MBu1)
- # H = reverse(gref["H"], dims=1)
- # H_MBu1 = reverse(gref_MBu1["H"], dims=1)
- # # H = gref["H"]
- # # H_MBu1 = gref_MBu1["H"]
- # push!(hms, heatmap(H .- H_MBu1, 
- # clims=(0.0,5.0),
- # ylimits=(0, size(H)[1]),
- # xlimits=(0, size(H)[2]),
- # colorbar = false)
- # )
- # end
-
- # h2 = scatter([0,0], [0,1], clims=(0.0,5.0),
- # xlims=(1,1.1), xshowaxis=false, yshowaxis=false, label="", colorbar_title="cbar", grid=false)
-
-
- # l = @layout [grid(6,5) a{0.01w}]
-
- # # Create the combined plot with the subplots and shared colormap
- # p_dhdt = plot(hms..., h2,
- # size=(1800, 1200),
- # layout=l,
- # link=:all,
- # aspect_ratio=:equal)
-
- # savefig(p_dhdt, joinpath(root_dir, "plots/MB/dhdt_MB_1"))
-
 end
 
 make_plots()

scripts/toy_model.jl

@@ -24,7 +24,7 @@ ENV["GKSwstype"]="nul"
 
 function run_toy_model()
 
- processes = 1
+ processes = 18
 
  # Enable multiprocessing
  ODINN.enable_multiprocessing(processes)
@@ -35,9 +35,9 @@ function run_toy_model()
  ODINN.set_run_spinup(false) # Run the spin-up random_MB = generate_random_MB(gdirs_climate, tspan; plot=false)n
  ODINN.set_use_spinup(false) # Use the updated spinup
  # Reference simulations
- ODINN.set_create_ref_dataset(true) # Generate reference data for UDE training
+ ODINN.set_create_ref_dataset(false) # Generate reference data for UDE training
  # UDE training
- ODINN.set_train(false) # Train UDE
+ ODINN.set_train(true) # Train UDE
  ODINN.set_retrain(false) # Re-use previous NN weights to continue training
  # Optimization method for differentiating the model
  ODINN.set_optimization_method("AD+AD")
@@ -69,7 +69,7 @@ function run_toy_model()
  if ODINN.create_ref_dataset[]
  println("Generating reference dataset for training...")
  # Compute reference dataset in parallel
- mb_model = ODINN.TI_model_1(DDF=8.0/1000.0, acc_factor=1.0/1000.0) # in m.w.e.
+ mb_model = ODINN.TI_model_1(DDF=4.0/1000.0, acc_factor=2.0/1000.0) # in m.w.e.
  gdir_refs = @time "PDEs" generate_ref_dataset(gdirs, tspan, solver=RDPK3Sp35(), mb_model)
 
  println("Saving reference data")
@@ -89,7 +89,7 @@ function run_toy_model()
  if ODINN.train[]
  # Train iceflow UDE in parallel
  n_ADAM = 20
- n_BFGS = 100
+ n_BFGS = 50
  batch_size = length(gdir_refs)
  # batch_size = 9
  UDE_settings = Dict("reltol"=>1e-7,
@@ -103,7 +103,7 @@ function run_toy_model()
  # optimizer = BFGS(initial_stepnorm=0.0001)
  optimizer = Adam(0.001)
  train_settings = (optimizer, n_ADAM, batch_size) # optimizer, epochs, batch_size
- mb_model = ODINN.TI_model_1(DDF=6.0/1000.0, acc_factor=1.0/1000.0) # in m.w.e.
+ mb_model = ODINN.TI_model_1(DDF=8.0/1000.0, acc_factor=1.0/1000.0) # in m.w.e.
  iceflow_trained, UA_f, loss_history = @time "UDEs" train_iceflow_UDE(gdirs, gdir_refs,
  tspan, train_settings, θ_trained; 
  UDE_settings=UDE_settings,

src/ODINN.jl

@@ -12,7 +12,7 @@ using SciMLSensitivity
 using Optimization, Optim, OptimizationOptimJL
 import OptimizationOptimisers.Adam
 using IterTools: ncycle
-using Zygote: @ignore 
+using ChainRules: @ignore_derivatives
 using Base: @kwdef
 using Flux
 using Tullio

src/helpers/climate.jl

@@ -160,7 +160,7 @@ function trim_period(period, climate)
  return period
 end
 
-function partial_year(period::Type{<:Period}, float::AbstractFloat)
+function partial_year(period::Type{<:Period}, float)
  _year, Δ = divrem(float, 1)
  year_start = Date(_year)
  year = period((year_start + Year(1)) - year_start)
@@ -172,7 +172,7 @@ partial_year(float::AbstractFloat) = partial_year(Day, float)
 
 function get_longterm_temps(gdir::PyObject, tspan)
  climate = xr.open_dataset(joinpath(gdir.dir, "raw_climate_$tspan.nc")) # load only once at the beginning
- dem = xr.open_rasterio(gdir.get_filepath("dem"))
+ dem = rioxarray.open_rasterio(gdir.get_filepath("dem"))
  apply_t_grad!(climate, dem)
  longterm_temps = climate.groupby("time.year").mean().temp.data
  return longterm_temps

src/helpers/iceflow.jl

@@ -33,7 +33,7 @@ function generate_ref_dataset(gdirs, tspan, mb_model; solver = RDPK3Sp35(), velo
  println("Running forward PDE ice flow model...\n")
  # Run batches in parallel
  A_noises = randn(rng_seed(), length(gdirs)) .* noise_A_magnitude
- refs = @showprogress pmap((gdir, A_noise) -> batch_iceflow_PDE(gdir, A_noise, tspan, solver, mb_model; run_spinup=false, velocities=velocities), gdirs, A_noises)
+ refs = @showprogress map((gdir, A_noise) -> batch_iceflow_PDE(gdir, A_noise, tspan, solver, mb_model; run_spinup=false, velocities=velocities), gdirs, A_noises)
 
  # Gather information per gdir
  gdir_refs = get_gdir_refs(refs, gdirs)
@@ -359,10 +359,10 @@ function batch_iceflow_UDE(θ, UA_f, context, gdir, UDE_settings, mb_model)
  mb_model = TI_model_1(DDF=5.0/1000.0, acc_factor=1.2/1000.0)
  end
  # Initialize climate dataset
- tstops, step = @ignore define_callback_steps(tspan)
+ tstops, step = @ignore_derivatives define_callback_steps(tspan)
  S_coords = context[13]
  dist_border = context[17]
- climate = @ignore init_climate(gdir, tspan, step, context[9], S_coords)
+ climate = @ignore_derivatives init_climate(gdir, tspan, step, context[9], S_coords)
  # Callback 
  # Define stop times every one month
  stop_condition(u,t,integrator) = stop_condition_tstops(u,t,integrator, tstops) #closure
@@ -371,7 +371,7 @@ function batch_iceflow_UDE(θ, UA_f, context, gdir, UDE_settings, mb_model)
  S::Matrix{Float64} = context[9]
  S_coords = context[13]
  MB = context[15]
- @ignore MB_timestep!(MB, mb_model, climate, S, S_coords, integrator.t, step)
+ @ignore_derivatives MB_timestep!(MB, mb_model, climate, S, S_coords, integrator.t, step)
  apply_MB_mask!(integrator.u, MB, MB_total, dist_border)
  end
  # Recompute A value
@@ -401,12 +401,12 @@ function batch_iceflow_UDE(θ, UA_f, context, gdir, UDE_settings, mb_model)
  testmode = context[16]
  V̄x_pred::Matrix{Float64}, V̄y_pred::Matrix{Float64} = avg_surface_V(context, H_pred, mean(climate.longterm_temps), "UDE", θ, UA_f; 
  testmode=testmode) # Average velocity with average temperature
- rgi_id::String = @ignore gdir.rgi_id
+ rgi_id::String = @ignore_derivatives gdir.rgi_id
  H_V_pred = (H_pred, V̄x_pred, V̄y_pred, rgi_id)
 
  # println("Total MB for $rgi_id: ", sum(MB_total))
 
- @ignore GC.gc() # Run the garbage collector to tame the RAM
+ @ignore_derivatives GC.gc() # Run the garbage collector to tame the RAM
 
  return H_V_pred
 end
@@ -418,11 +418,11 @@ function batch_iceflow_UDE_inplace(θ, UA_f, gdir, tspan; solver = RDPK3Sp35())
  mb_model = TI_model_1(DDF=5.0/1000.0, acc_factor=1.2/1000.0)
 
  # Initialize climate dataset
- _, step = @ignore define_callback_steps(tspan)
+ _, step = @ignore_derivatives define_callback_steps(tspan)
  S = context[3]
  S_coords = context[32]
  dist_border = context[33]
- climate = @ignore init_climate(gdir, tspan, step, S, S_coords)
+ climate = @ignore_derivatives init_climate(gdir, tspan, step, S, S_coords)
  # Define stop times every one month
  tstops, step = define_callback_steps(tspan)
  stop_condition(u,t,integrator) = stop_condition_tstops(u,t,integrator, tstops) #closure
@@ -432,7 +432,7 @@ function batch_iceflow_UDE_inplace(θ, UA_f, gdir, tspan; solver = RDPK3Sp35())
  MB = context[25]
  S = context[3]
  S_coords = context[32]
- MB_timestep!(MB, mb_model, climate, S, S_coords, integrator.t, step)
+ @ignore_derivatives MB_timestep!(MB, mb_model, climate, S, S_coords, integrator.t, step)
  apply_MB_mask!(integrator.u, MB, MB_total, context)
  end
  # Recompute A value

src/helpers/initial_conditions.jl

@@ -137,7 +137,7 @@ function build_UDE_context_inv(gdir, gdir_ref, tspan; run_spinup=false)
  # Get evolved tickness and surface
  H = gdir_ref["H"]
 
- @ignore begin
+ @ignore_derivatives begin
  glacier_gd = xr.open_dataset(gdir.get_filepath("gridded_data"))
  H₁ = glacier_gd.consensus_ice_thickness.data
  fillNaN!(H₀) # Fill NaNs with 0s to have real boundary conditions

src/helpers/inversions.jl

@@ -145,7 +145,7 @@ function loss_iceflow_inversion(θ, UD, gdirs_climate, gdir_refs, context_batche
  # l_V = (l_Vx + l_Vy)/2
 
  # Plot V diffs to understand training
- @ignore begin
+ @ignore_derivatives begin
  plot_V_diffs(gdirs_climate, gdir_refs, V_preds)
  end
 

src/helpers/oggm.jl

@@ -24,8 +24,10 @@ function oggm_config(working_dir=joinpath(homedir(), "Python/OGGM_data"); oggm_p
 
  # Multiprocessing 
  multiprocessing = $oggm_processes > 1 ? true : false
- PARAMS["mp_processes"] = $oggm_processes
  PARAMS["use_multiprocessing"] = multiprocessing # Let's use multiprocessing for OGGM
+ if multiprocessing
+ PARAMS["mp_processes"] = $oggm_processes
+ end
 
  end # @eval ODINN
  end # @everywhere

src/helpers/utils.jl

@@ -230,7 +230,7 @@ function get_NN(θ_trained)
  Dense(10,3, x->softplus.(x)),
  Dense(3,1, sigmoid_A)
  )
- Flux.f64(UA)
+ UA = Flux.f64(UA)
  # See if parameters need to be retrained or not
  θ, UA_f = Flux.destructure(UA)
  if !isempty(θ_trained)

test/test.jl

@@ -14,7 +14,7 @@ using Test
 # using HDF5 
 # using JLD2
 # using OrdinaryDiffEq, DiffEqFlux
-# using Zygote: @ignore
+# using Zygote: @ignore_derivatives
 # using Flux
 # using Tullio, RecursiveArrayTools
 # using Infiltrator