Merge pull request #82 from agdestein/missing-pullbacks

Add missing BC pullbacks

docs/src/.vitepress/config.mts

@@ -128,6 +128,7 @@ export default defineConfig({
               { text: 'GPU Support', link: '/manual/gpu' },
               { text: 'Operators', link: '/manual/operators' },
               { text: 'Temperature equation', link: '/manual/temperature' },
+              { text: 'Differentiating code', link: '/manual/differentiability' },
               { text: 'Large eddy simulation', link: '/manual/les' },
               { text: 'Neural closure models', link: '/manual/closure' },
               { text: 'API', link: '/manual/api' },
@@ -216,6 +217,7 @@ export default defineConfig({
               { text: 'GPU Support', link: '/manual/gpu' },
               { text: 'Operators', link: '/manual/operators' },
               { text: 'Temperature equation', link: '/manual/temperature' },
+              { text: 'Differentiating code', link: '/manual/differentiability' },
               { text: 'Large eddy simulation', link: '/manual/les' },
               { text: 'Neural closure models', link: '/manual/closure' },
               { text: 'API', link: '/manual/api' },

docs/src/index.md

@@ -29,7 +29,7 @@ features:
   - icon: <img width="64" height="64" src="https://raw.githubusercontent.com/JuliaDiff/ChainRulesCore.jl/fa530b9865ec0cb3acff81ddef0967fdcc8c8214/docs/src/assets/logo.svg" alt="ChainRules.jl"/>
     title: Differentiable physics
     details: Backpropagate through the solver using Zygote.jl to optimize closure models
-    link: https://github.com/FluxML/Zygote.jl
+    link: /manual/differentiability
   - icon: <img width="64" height="64" src="https://raw.githubusercontent.com/LuxDL/Lux.jl/ca2c635f9d70a3d994efab9f0116711a8cdb1a48/assets/lux-logo.svg" alt="Lux.jl"/>
     title: Neural network integration
     details: Integrate neural network closure models with Lux.jl

docs/src/manual/bc.md

@@ -19,4 +19,5 @@ PressureBC
 ```@docs
 offset_p
 offset_u
+boundary
 ```

docs/src/manual/differentiability.md

@@ -0,0 +1,57 @@
+```@meta
+CurrentModule = IncompressibleNavierStokes
+```
+
+# Differentiating through the code
+
+IncompressibleNavierStokes is
+[reverse-mode differentiable](https://juliadiff.org/ChainRulesCore.jl/stable/index.html#Reverse-mode-AD-rules-(rrules)),
+which means that you can back-propagate gradients through the code.
+This comes at a cost however, as intermediate velocity fields need to be stored
+in memory for use in the backward pass. For this reason, many of the operators
+come in two versions:oa slow differentiable allocating non-mutating variant (e.g.
+[`divergence`](@ref)) and fast non-differentiable non-allocating mutating
+variant (e.g. [`divergence!`](@ref).)
+
+!!! warning "Differentiable code"
+    To make your code differentiable, you must use the differentiable versions
+    of the operators (without the exclamation marks).
+
+To differentiate the code, use [Zygote.jl](https://github.com/FluxML/Zygote.jl).
+
+## Example: Gradient of kinetic energy
+
+To differentiate outputs of a simulation with respect to the initial conditions,
+make a time stepping loop composed of differentiable operations:
+
+```julia
+import IncompressibleNavierStokes as INS
+setup = INS.Setup(0:0.01:1, 0:0.01:1; Re = 500.0)
+psolver = INS.default_psolver(setup)
+method = INS.RKMethods.RK44P2()
+Δt = 0.01
+nstep = 100
+(; Iu) = setup.grid
+function final_energy(u)
+    stepper = INS.create_stepper(method; setup, psolver, u, temp = nothing, t = 0.0)
+    for it = 1:nstep
+        stepper = INS.timestep(method, stepper, Δt)
+    end
+    (; u) = stepper
+    sum(abs2, u[1][Iu[1]]) / 2 + sum(abs2, u[2][Iu[2]]) / 2
+end
+
+u = INS.random_field(setup)
+
+using Zygote
+g, = Zygote.gradient(final_energy, u)
+
+@show size.(u)
+@show size.(g)
+```
+
+Now `g` is the gradient of `final_energy` with respect to the initial conditions
+`u`, and consequently has the same size.
+
+Note that every operation in the `final_energy` function is non-mutating and
+thus differentiable.

examples/Project.toml

@@ -15,6 +15,9 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
+[sources]
+IncompressibleNavierStokes = {path = ".."}
+
 [compat]
 CairoMakie = "0.12"
 FFTW = "1"
@@ -26,6 +29,3 @@ LinearAlgebra = "1"
 Literate = "2"
 SparseArrays = "1"
 julia = "1.9"
-
-[sources]
-IncompressibleNavierStokes = {path = ".."}