diff --git a/docs/Project.toml b/docs/Project.toml index 3ac8cb0b1..4d8ff194c 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,7 +1,6 @@ [deps] CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6" NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228" diff --git a/docs/src/differentiability.md b/docs/src/differentiability.md index d07ce126a..b69a25790 100644 --- a/docs/src/differentiability.md +++ b/docs/src/differentiability.md @@ -12,7 +12,7 @@ For the differentiability of our model we rely on [Enzyme.jl](https://github.com First we initialize the model as usual: -```@example autodiff +```julia using SpeedyWeather, Enzyme spectral_grid = SpectralGrid(trunc=23, nlayers=3) @@ -24,7 +24,7 @@ run!(simulation, period=Day(10)) # spin-up the model a bit Then, we get all variables we need from our `simulation` -```@example autodiff +```julia (; prognostic_variables, diagnostic_variables, model) = simulation (; Δt, Δt_millisec) = model.time_stepping dt = 2Δt @@ -35,7 +35,7 @@ diagn = diagnostic_variables Next, we will prepare to use Enzyme. Enzyme saves the gradient information in a shadow of the original input. For the inputs this shadow is initialized zero, whereas for the output the shadow is used as the seed of the AD. In other words, as we are doing reverse-mode AD, the shadow of the output is the value that is backpropageted by the reverse-mode AD. Ok, let's initialize everything: -```@example autodiff +```julia dprogn = one(progn) # shadow for the progn values ddiagn = make_zero(diagn) # shadow for the diagn values dmodel = make_zero(model) # here, we'll accumulate all parameter derivatives @@ -43,15 +43,14 @@ dmodel = make_zero(model) # here, we'll accumulate all parameter derivatives Then, we can already do the differentiation with Enzyme -```@example autodiff +```julia autodiff(Reverse, SpeedyWeather.timestep!, Const, Duplicated(progn, dprogn), Duplicated(diagn, ddiagn), Const(dt), Duplicated(model, dmodel)) ``` The derivitaves are accumulated in the `dmodel` shadow. So, if we e.g. want to know the derivative with respect to the gravity constant, we just have to inspect: -```@example autodiff +```julia dmodel.planet.gravity ``` Doing a full sensitivity analysis through a long integration is computationally much more demanding, and is something that we are currently working on. -