Skip to content

Commit

Permalink
Disable reverse-mode in autodiff example
Browse files Browse the repository at this point in the history
  • Loading branch information
bgroenks96 committed Jan 1, 2025
1 parent 0ecc92f commit bc03743
Showing 1 changed file with 11 additions and 10 deletions.
21 changes: 11 additions & 10 deletions examples/heat_simple_autodiff_grad.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,8 @@ using CryoGrid
# Set up forcings and boundary conditions similarly to other examples:
forcings = loadforcings(CryoGrid.Forcings.Samoylov_ERA_obs_fitted_1979_2014_spinup_extended_2044);
soilprofile, tempprofile = CryoGrid.SamoylovDefault
soilprofile = SoilProfile(0.0u"m" => SimpleSoil())
freezecurve = PainterKarra(swrc=VanGenuchten())
soilprofile = SoilProfile(0.0u"m" => SimpleSoil(; freezecurve))
grid = CryoGrid.DefaultGrid_5cm
initT = initializer(:T, tempprofile)
tile = CryoGrid.SoilHeatTile(
Expand Down Expand Up @@ -43,10 +44,10 @@ using OrdinaryDiffEq
using Statistics
function loss(prob::CryoGridProblem, p)
newprob = remake(prob, p=p)
# Here we specify the sensitivity algorithm. Note that this is only
# necessary for reverse-mode autodiff with Zygote.
# autojacvec = true uses ForwardDiff to calculate the jacobian;
# enabling checkpointing (theoretically) reduces the memory cost of the backwards pass.
## Here we specify the sensitivity algorithm. Note that this is only
## necessary for reverse-mode autodiff with Zygote.
## autojacvec = true uses ForwardDiff to calculate the jacobian;
## enabling checkpointing (theoretically) reduces the memory cost of the backwards pass.
sensealg = InterpolatingAdjoint(autojacvec=true, checkpointing=true)
newsol = solve(newprob, Euler(), dt=300.0, sensealg=sensealg);
newout = CryoGridOutput(newsol)
Expand All @@ -57,8 +58,8 @@ end
pvec = vec(prob.p)
fd_grad = @time ForwardDiff.gradient(pᵢ -> loss(prob, pᵢ), pvec)

# We can also try with reverse-mode autodiff. This is generally slower for smaller numbers
# of parmaeters (<100) but could be worthwhile for model configurations with high-dimensional
# parameterizations.
zy_grad = @time Zygote.gradient(pᵢ -> loss(prob, pᵢ), pvec)
@assert maximum(abs.(fd_grad .- zy_grad)) .< 1e-6 "Forward and reverse gradients don't match!"
## We can also try with reverse-mode autodiff. This is generally slower for smaller numbers
## of parmaeters (<100) but could be worthwhile for model configurations with high-dimensional
## parameterizations.
## zy_grad = @time Zygote.gradient(pᵢ -> loss(prob, pᵢ), pvec)
## @assert maximum(abs.(fd_grad .- zy_grad)) .< 1e-6 "Forward and reverse gradients don't match!"

0 comments on commit bc03743

Please sign in to comment.