Fix docs

docs/Project.toml

@@ -2,12 +2,16 @@
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Examples = "318dbb63-4243-420f-99f2-d56058123f9d"
 IncompressibleNavierStokes = "5e318141-6589-402b-868d-77d7df8c442e"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 NeuralClosure = "099dac27-d7f2-4047-93d5-0baee36b9c25"
 
+[sources]
+Examples = {path = "../examples"}
+IncompressibleNavierStokes = {path = ".."}
+NeuralClosure = {path = "../lib/NeuralClosure"}
+
 [compat]
 Documenter = "1"
 DocumenterCitations = "1"
@@ -17,12 +21,3 @@ IncompressibleNavierStokes = "2"
 Literate = "2"
 NeuralClosure = "1"
 julia = "1.9"
-
-[sources.Examples]
-path = "../examples"
-
-[sources.IncompressibleNavierStokes]
-path = ".."
-
-[sources.NeuralClosure]
-path = "../lib/NeuralClosure"

docs/make.jl

@@ -2,6 +2,7 @@
 localdev = haskey(ENV, "LOCALDEV")
 
 # Get access to example dependencies
+# Those are required for running the examples and generating output
 push!(LOAD_PATH, joinpath(@__DIR__, "..", "examples"))
 
 using IncompressibleNavierStokes
@@ -92,7 +93,7 @@ makedocs(;
     # draft = true,
     # clean = !localdev,
     modules = [IncompressibleNavierStokes, NeuralClosure],
-    warnonly = true, # Reexporting KernelAbstractions.CPU fails otherwise
+    # warnonly = true, # Reexporting KernelAbstractions.CPU fails otherwise
     plugins = [bib],
     authors = "Syver Døving Agdestein, Benjamin Sanderse, and contributors",
     repo = Remotes.GitHub("agdestein", "IncompressibleNavierStokes.jl"),

docs/references.bib

@@ -22,14 +22,19 @@ @misc{Sanderse2023
     Month         = {Jul},
     Url           = {http://arxiv.org/abs/2307.10874v1},
 }
-@article{Agdestein2024,
-  archiveprefix = {arXiv},
-  author        = {Syver Døving Agdestein and Benjamin Sanderse},
-  eprint        = {2403.18088},
-  primaryclass  = {math.NA},
-  title         = {Discretize first, filter next: learning divergence-consistent closure models for large-eddy simulation},
-  year          = {2024}
+@article{Agdestein2025,
+    title = {Discretize first, filter next: Learning divergence-consistent closure models for large-eddy simulation},
+    journal = {Journal of Computational Physics},
+    volume = {522},
+    pages = {113577},
+    year = {2025},
+    issn = {0021-9991},
+    doi = {https://doi.org/10.1016/j.jcp.2024.113577},
+    url = {https://www.sciencedirect.com/science/article/pii/S0021999124008258},
+    author = {Syver Døving Agdestein and Benjamin Sanderse},
+    keywords = {Discrete filtering, Closure modeling, Divergence-consistency, Large-eddy simulation, Neural ODE, A-posteriori training}
 }
+
 @article{Beck2021,
   author   = {Beck, Andrea and Kurz, Marius},
   doi      = {10.1002/gamm.202100002},

docs/src/.vitepress/config.mts

@@ -94,6 +94,7 @@ export default defineConfig({
               { text: 'Differentiating code', link: '/manual/differentiability' },
               { text: 'Temperature equation', link: '/manual/temperature' },
               { text: 'Large eddy simulation', link: '/manual/les' },
+              { text: 'SciML', link: '/manual/sciml' },
             ],
           },
         ],
@@ -189,6 +190,7 @@ export default defineConfig({
               { text: 'Differentiating code', link: '/manual/differentiability' },
               { text: 'Temperature equation', link: '/manual/temperature' },
               { text: 'Large eddy simulation', link: '/manual/les' },
+              { text: 'SciML', link: '/manual/sciml' },
             ],
           },
         ],

docs/src/manual/closure.md

@@ -1,6 +1,6 @@
 # Neural closure models
 
-!!! note "`NeuralClorusure`"
+!!! note "`NeuralClosure`"
     These features are experimental, and require cloning
     IncompressibleNavierStokes from GitHub:
     ```sh

docs/src/manual/differentiability.md

@@ -24,30 +24,31 @@ variant (e.g. [`divergence!`](@ref).)
     To make your code differentiable, you must use the differentiable versions
     of the operators (without the exclamation marks).
 
-#### Example: Gradient of kinetic energy
+### 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
+```@example Differentiability
+using IncompressibleNavierStokes
+
 ax = range(0, 1, 101)
-setup = INS.Setup(; x = (ax, ax), Re = 500.0)
-psolver = INS.default_psolver(setup)
-method = INS.RKMethods.RK44P2()
+setup = Setup(; x = (ax, ax), Re = 500.0)
+psolver = default_psolver(setup)
+method = 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)
+    stepper = create_stepper(method; setup, psolver, u, temp = nothing, t = 0.0)
     for it = 1:nstep
-        stepper = INS.timestep(method, stepper, Δt)
+        stepper = timestep(method, stepper, Δt)
     end
     (; u) = stepper
     sum(abs2, u[Iu[1], 1]) / 2 + sum(abs2, u[Iu[2], 2]) / 2
 end
 
-u = INS.random_field(setup)
+u = random_field(setup)
 
 using Zygote
 g, = Zygote.gradient(final_energy, u)
@@ -61,7 +62,6 @@ Now `g` is the gradient of `final_energy` with respect to the initial conditions
 Note that every operation in the `final_energy` function is non-mutating and
 thus differentiable.
 
---- 
 ## Automatic differentiation with Enzyme
 
 Enzyme.jl is highly-efficient and its ability to perform AD on optimized code allows Enzyme to meet or exceed the performance of state-of-the-art AD tools.
@@ -71,45 +71,47 @@ The downside is that restricts the user's defined f function to not do things li
     Enzyme's autodiff function can only handle functions with scalar output. To implement pullbacks for array-valued functions, use a mutating function that returns `nothing` and stores its result in one of the arguments, which must be passed wrapped in a Duplicated.
     In IncompressibleNavierStokes, we provide `enzyme_wrapper` to automatically wrap the function and its arguments in the correct way.
 
-#### Example: Gradient of the right-hand side
+### Example: Gradient of the right-hand side
 
 In this example we differentiate the right-hand side of the Navier-Stokes equations with respect to the velocity field `u`:
 
-```@example
-import IncompressibleNavierStokes as INS
+```@example Differentiability
 using Enzyme
 ax = range(0, 1, 101)
-setup = INS.Setup(; x = (ax, ax), Re = 500.0)
-psolver = INS.default_psolver(setup)
-u = INS.random_field(setup)
+setup = Setup(; x = (ax, ax), Re = 500.0)
+psolver = default_psolver(setup)
+u = random_field(setup)
 dudt = similar(u)
 t = 0.0
-f! = INS.right_hand_side!
+f! = right_hand_side!
 ```
 Notice that we are using the mutating (in-place) version of the right-hand side function. This function can not be differentiate by Zygote, which requires the slower non-mutating version of the right-hand side.
 
-We then define the `Dual` part of the input and output, required to store the adjoint values: 
-```@example
+We then define the `Dual` part of the input and output, required to store the adjoint values:
+
+```@example Differentiability
 ddudt = Enzyme.make_zero(dudt) .+ 1;
 du = Enzyme.make_zero(u);
 ```
 Remember that the derivative of the output (also called the *seed*) has to be set to $1$ in order to compute the gradient. In this case the output is the force, that we store mutating the value of `dudt` inside `right_hand_side!`.
 
 Then we pack the parameters to be passed to `right_hand_side!`:
-```@example
+
+```@example Differentiability
 params = [setup, psolver];
 params_ref = Ref(params);
 ```
 Now, we call the `autodiff` function from Enzyme:
-```@example
+
+```@example Differentiability
 Enzyme.autodiff(Enzyme.Reverse, f!, Duplicated(dudt,ddudt), Duplicated(u,du), Const(params_ref), Const(t))
 ```
-Since we have passed a `Duplicated` object, the gradient of `u` is stored in `du`. 
+Since we have passed a `Duplicated` object, the gradient of `u` is stored in `du`.
 
 Finally, we can also compare its value with the one obtained by Zygote differentiating the out-of-place (non-mutating) version of the right-hand side:
-```@example
-using Zygote
+
+```@example Differentiability
 f = create_right_hand_side(setup, psolver)
-_, zpull = Zygote.pullback(f, u, nothing, T(0));
+_, zpull = Zygote.pullback(f, u, nothing, 0.0);
 @assert zpull(dudt)[1] == du
-```
+```

docs/src/manual/sciml.md

@@ -17,30 +17,46 @@ Using IncompressibleNavierStokes it is possible to write the momentum equations
 \end{align*}
 ```
 
-The derivation and the drawbacks of this approach are discussed in the [documentation](/docs/src/manual/spatial.md).
+<!-- There is an issue with linking to other MD pages: -->
+<!-- https://github.com/LuxDL/DocumenterVitepress.jl/issues/172 -->
+The derivation and the drawbacks of this approach are discussed in the [documentation](spatial.md).
 
 This projected right-hand side can be used in the SciML solvers to solve the Navier-Stokes equations. The following example shows how to use the SciML solvers to solve the ODEs obtained from the Navier-Stokes equations.
 
-```julia
-using DifferentialEquations 
-f(u, p, t) = create_right_hand_side(setup, psolver) 
-u0 = INITIAL_CONDITION
+!!! note "OrdinaryDiffEqTsit5"
+    We here use the explicit solver `Tsit5` from OrdinaryDiffEqTsit5 directly to
+    skip loading all the toolchains for implicit solvers, which takes a while to
+    install on GitHub.
+
+```@example SciML
+using OrdinaryDiffEqTsit5
+using IncompressibleNavierStokes
+ax = range(0, 1, 101)
+setup = Setup(; x = (ax, ax), Re = 500.0)
+psolver = default_psolver(setup)
+f = create_right_hand_side(setup, psolver)
+u0 = random_field(setup)
 tspan = (0.0, 1.0)     # time span where to solve.
 problem = ODEProblem(f, u0, tspan) #SciMLBase.ODEProblem
 sol = solve(problem, Tsit5(), reltol = 1e-8, abstol = 1e-8) # sol: SciMLBase.ODESolution
 ```
 
-Alternatively, it is also possible to use an [in-place formulation](https://docs.sciml.ai/DiffEqDocs/stable/basics/problem/#In-place-vs-Out-of-Place-Function-Definition-Forms) 
+Alternatively, it is also possible to use an [in-place formulation](https://docs.sciml.ai/DiffEqDocs/stable/basics/problem/#In-place-vs-Out-of-Place-Function-Definition-Forms)
 
-```julia
-f(du,u,p,t) = right_hand_side!(du, u, Ref([setup, psolver]), t)
+```@example SciML
+f!(du,u,p,t) = right_hand_side!(du, u, Ref([setup, psolver]), t)
+u = similar(u0)
+du = similar(u0)
+p = nothing
+t = 0.0
+f!(du,u,p,t)
 ```
 that is usually faster than the out-of-place formulation.
 
 You can look [here](https://docs.sciml.ai/DiffEqDocs/stable/basics/overview/) for more information on how to use the SciML solvers and all the options available.
 
-## API 
+## API
 ```@autodocs
 Modules = [IncompressibleNavierStokes]
 Pages = ["sciml.jl"]
-```
+```

examples/Project.toml

@@ -5,27 +5,33 @@ version = "1.0.0"
 
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 IncompressibleNavierStokes = "5e318141-6589-402b-868d-77d7df8c442e"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
-[sources.IncompressibleNavierStokes]
-path = ".."
+[sources]
+IncompressibleNavierStokes = {path = ".."}
 
 [compat]
 CairoMakie = "0.12"
+Enzyme = "0.13"
 FFTW = "1"
 GLMakie = "0.10"
 IncompressibleNavierStokes = "2"
 JLD2 = "0.5"
 LaTeXStrings = "1"
 LinearAlgebra = "1"
 Literate = "2"
+OrdinaryDiffEqTsit5 = "1"
 SparseArrays = "1"
+Zygote = "0.6"
 julia = "1.9"

examples/TaylorGreenVortex2D.jl

@@ -49,11 +49,8 @@ function compute_convergence(; D, nlist, lims, Re, tlims, Δt, uref, backend = C
         )
         (; u, t), outputs = solve_unsteady(; setup, ustart, tlims, Δt, psolver)
         (; Ip) = setup.grid
-        a, b = T(0), T(0)
-        for α = 1:D
-            a += sum(abs2, u[α][Ip] - ut[α][Ip])
-            b += sum(abs2, ut[α][Ip])
-        end
+        a = sum(abs2, u[Ip, :] - ut[Ip, :])
+        b = sum(abs2, ut[Ip, :])
         e[i] = sqrt(a) / sqrt(b)
     end
     e