Skip to content

Commit

Permalink
remove Enzyme dep from docs again
Browse files Browse the repository at this point in the history
  • Loading branch information
maximilian-gelbrecht committed Feb 5, 2025
1 parent 7bda523 commit 42ce046
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 7 deletions.
1 change: 0 additions & 1 deletion docs/Project.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
11 changes: 5 additions & 6 deletions docs/src/differentiability.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -35,23 +35,22 @@ 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
```

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.

0 comments on commit 42ce046

Please sign in to comment.